diff --git a/.dockerignore b/.dockerignore index 363b8e8..09b2cef 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,18 +1,38 @@ -# Ignore everything by default +# The build context is an explicit allowlist. Runtime stages never receive the +# upstream parity oracles; reference stages copy only their selected checkout. * -# Keep Python source files anywhere in the tree -!**/*.py +!.dockerignore +!.gitmodules +!kernels.lock +!README.md +!LICENSE +!THIRD_PARTY_NOTICES.md -# Keep dependency manifest used during build -!requirements.txt +!requirements/ +!requirements/** +!LICENSES/ +!LICENSES/** +!src/ +!src/** +!tests/ +!tests/** +!benchmarks/ +!benchmarks/** +!tools/ +!tools/** +!docker/ +!docker/** +!docs/ +!docs/** +!model_cards/ +!model_cards/** -# Keep official submodules (needed for pip install -e during build) -!official/** +!vendor/ +!vendor/README.md -# Keep non-Python files needed by models -!**/*.json -!**/*.txt -!**/*.md -!Dockerfile -!.dockerignore \ No newline at end of file +**/.git +**/.git/** +**/__pycache__ +**/*.py[cod] +**/.pytest_cache diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..990a5cd --- /dev/null +++ b/.gitattributes @@ -0,0 +1,4 @@ +* text=auto eol=lf +# Preserve byte-exact third-party legal text without treating its upstream +# trailing spaces as project defects. +LICENSES/ankh/LICENSE.md whitespace=-trailing-space diff --git a/.gitignore b/.gitignore index 99cb87b..81a81c0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,27 +1,53 @@ +# Python +__pycache__/ +*.py[cod] +*.so +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.coverage +htmlcov/ +.venv/ -*.pyc -*.png -!docs/assets/*.png -*.pth -*.pt -*.safetensors -*.json -!e1_fastplms/*.json -*.bin -.qodo +# Build and run outputs +build/ +dist/ +/artifacts/ +results/ +.cache/ *.db +*.db-shm +*.db-wal *.nbc *.nbi + +# Checkpoint and tensor payloads +*.bin *.ckpt -/results_classification_lora -/results_regression_lora -/testing/results +*.pth +*.pt +*.safetensors +!tests/goldens/**/*.safetensors + +# Local credentials and workstation state +.env +.env.* +!.env.example +*.key +*.pem +*.p12 +*.pfx +.qodo/ + +# Local planning and scratch material draft.md -/dplm -/E1 -/.cache -/esm -/testing/__pycache__ -/github_issues -/internal -/marketing +github_issues/ +internal/ +marketing/ +/.codex-* +/.codex_* +/.claude +/tmp +/output +*.log +.secrets.env diff --git a/.gitmodules b/.gitmodules index 5eccbe4..f59ed67 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,12 +1,27 @@ -[submodule "official/boltz"] - path = official/boltz +[submodule "vendor/upstream/ankh"] + path = vendor/upstream/ankh + url = https://github.com/agemagician/Ankh.git +[submodule "vendor/upstream/biohub-esm"] + path = vendor/upstream/biohub-esm + url = https://github.com/Biohub/esm.git +[submodule "vendor/upstream/biohub-transformers"] + path = vendor/upstream/biohub-transformers + url = https://github.com/Biohub/transformers.git +[submodule "vendor/upstream/boltz"] + path = vendor/upstream/boltz url = https://github.com/jwohlwend/boltz.git -[submodule "official/e1"] - path = official/e1 - url = https://github.com/Profluent-AI/E1.git -[submodule "official/dplm"] - path = official/dplm +[submodule "vendor/upstream/dplm"] + path = vendor/upstream/dplm url = https://github.com/bytedance/dplm.git -[submodule "official/esm"] - path = official/esm - url = https://github.com/Biohub/esm.git +[submodule "vendor/upstream/e1"] + path = vendor/upstream/e1 + url = https://github.com/Profluent-AI/E1.git +[submodule "vendor/upstream/fair-esm"] + path = vendor/upstream/fair-esm + url = https://github.com/facebookresearch/esm.git +[submodule "vendor/upstream/openfold"] + path = vendor/upstream/openfold + url = https://github.com/aqlaboratory/openfold.git +[submodule "vendor/upstream/protein-ttt"] + path = vendor/upstream/protein-ttt + url = https://github.com/anton-bushuiev/ProteinTTT.git diff --git a/AGENTS.md b/AGENTS.md index 469268f..44a2740 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,29 +1,157 @@ # FastPLMs -## Purpose and Sources +## Purpose -FastPLMs provides Hugging Face-compatible protein language and structure models under `fastplms/`. +FastPLMs maintains the runtime source uploaded to Hugging Face protein language +and structure model repositories. This repository is a source, test, artifact, +and dependency workspace, not an installable Python distribution. Changes must +preserve biological conventions, Transformers behavior, reproducibility, legal +provenance, and the evidence boundary of each model family. -- `docs/architecture.md` and `docs/models.md`: package and model-family contracts -- `docs/embedding_api.md`: shared embedding interface -- `docs/attention_backends.md`: backend-specific behavior -- `docs/testing.md`: per-family images, markers, and parity workflow +Start with [README.md](README.md) for user-facing behavior and +[docs/README.md](docs/README.md) for documentation routing. -## Architectural Invariants +## Sources of truth -- Model and config classes remain compatible with Transformers auto classes and `trust_remote_code=True`. -- `fastplms/embedding_mixin.py` is the shared sequence embedding API. -- Tokenizer-mode families accept token IDs and attention masks. E1 sequence mode has no tokenizer and must retain its native raw-sequence preparation path. -- `testing/conftest.py` is the authoritative model registry, and `testing/test_parity.py` is the strict parity suite. -- Use per-family Docker images for native dependency parity. Respect the `gpu`, `slow`, `large`, and `structure` markers before running expensive suites. +- `src/fastplms/models.toml`: model IDs, revisions, files, AutoClasses, + tokenizer modes, state transformations, backends, precision, licenses, and + release tiers. +- `src/fastplms/registry.py`: typed parsing and validation of the manifest. +- `docs/architecture.md` and `docs/models.md`: runtime-source and model-family + contracts. +- `docs/embedding_api.md`: shared ordered embedding interface and persistence. +- `docs/attention_backends.md`: backend names, dtype constraints, masks, and + parity boundaries. +- `docs/testing.md`: candidate/reference stages, markers, and parity workflow. +- `tests/parity/`: strict model and tokenizer comparisons. -## Canonical Commands +Do not infer a model contract from an older README, an unpinned Hub card, or an +unused code path when the manifest or current tests say otherwise. + +## Repository boundaries + +- `src/fastplms/` contains runtime source copied into Hugging Face artifacts. +- `vendor/upstream/` contains pinned official repositories used as parity + oracles. Runtime code must not import from this directory. +- `tests/` contains unit, integration, parity, structure, and release checks. +- `tools/` contains artifact, conversion, remote, and maintenance workflows. +- `examples/` contains runnable research and training examples. Keep examples + directly in this directory rather than creating a tutorial subtree. +- `model_cards/` contains generated checkpoint cards. +- `LICENSES/` contains distributable third-party legal texts and provenance. + +Do not place license files, model cards, or READMEs beside runtime model +modules. Do not hand-edit generated model cards or +`docs/generated/support.md`; update the manifest or renderer and regenerate. + +## Architectural invariants + +- Model and configuration classes remain compatible with Transformers auto + classes and `trust_remote_code=True`. +- `src/fastplms/embeddings/` is the shared sequence embedding API. +- Tokenizer-mode families accept token IDs and attention masks. E1 has no + tokenizer and must retain its native raw-sequence preparation path. +- Structure families retain native chain, residue, atom, ligand, nucleic-acid, + and MSA semantics where applicable. +- A requested attention backend either executes the named implementation or + raises. Never add a silent fallback. +- Official repositories are isolated references, not build inputs for runtime + source. Production imports must not change `sys.path`, download code, compile + a kernel, initialize a model, or mutate global Torch state. +- State transformations are named, deterministic, and covered by exact tests. +- Boltz2 remains provisional until its declared native end-to-end equivalence + limits pass. Do not broaden its claims from partial contracts. + +## Common workflows + +Initialize official sources: ```bash git submodule update --init --recursive -./build_images.sh esm2 -docker run --rm --gpus all --ipc=host -v ${PWD}:/workspace fastplms-esm2 \ - python -m pytest /workspace/testing/test_parity.py -k esm2 -v ``` -Always pass `--ipc=host` to Dockerized PyTorch runs. +Run portable release checks on the declared remote environment: + +```bash +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/key \ + --suite check + +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/key \ + --suite compliance +``` + +Build the candidate and one isolated reference image: + +```bash +sudo docker buildx bake \ + -f docker/docker-bake.hcl \ + candidate reference-esm2 \ + --load +``` + +Run focused candidate tests only when their required reference results already +exist: + +```bash +sudo docker compose -f docker/compose.yaml run --rm candidate \ + python -m pytest tests/parity -k esm2 -v +``` + +Compose supplies `ipc: host` for its services. Pass `--ipc=host` to raw +Dockerized PyTorch runs. Respect the `gpu`, `slow`, `large`, and `structure` +markers before running expensive suites. + +Regenerate and check documentation: + +```bash +PYTHONPATH=src python -m tools.artifacts.generate_docs +PYTHONPATH=src python -m tools.artifacts.generate_docs --check +python -m pytest tests/release/test_documentation.py \ + tests/release/test_model_card_licenses.py -v +``` + +Build a local Hub artifact: + +```bash +PYTHONPATH=src python -m tools.artifacts.build \ + esm2_150m \ + /cache/fast-snapshot \ + --tokenizer-dir /cache/official-tokenizer-snapshot \ + --output-root dist/hub +``` + +Preview an add-only Hub update that excludes checkpoint weights: + +```bash +PYTHONPATH=src python -m tools.artifacts.publish \ + --files-only \ + --artifact-root dist/hub \ + --dry-run +``` + +Remove `--dry-run` only after reviewing every planned path. The publisher must +select every manifest model when no positional model IDs are provided; model +IDs restrict the operation to that explicit subset. It must remain +manifest-scoped, add-only, protected by the remote parent commit, and free of +token command-line arguments. It must never upload weight-shaped paths, create +repositories, delete remote files, or publish complete-artifact attestations +during a files-only update. + +## Change policy + +- Inspect the manifest, family code, tests, and current docs before changing a + biological or model-facing contract. +- Keep changes focused and preserve unrelated work in a dirty tree. +- Add or update tests for public behavior, conversion rules, file identities, + generated output, and fail-closed paths. +- Use repository-native verification proportional to risk. Do not run large + GPU or structure suites without the requested environment and authority. +- Record measured results with the environment, dtype, backend, sequence panel, + and threshold. Do not turn an implementation detail into a parity, speed, or + biological claim. +- Never inspect or print credential files. Pass credential paths opaquely to + the trusted command that needs them. diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 2890254..0000000 --- a/Dockerfile +++ /dev/null @@ -1,54 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# NOTE: switch to cudnn-devel if you need to compile CUDA extensions (e.g. flash-attn from source) -FROM nvidia/cuda:12.8.0-cudnn-runtime-ubuntu24.04 - -ENV DEBIAN_FRONTEND=noninteractive \ - PYTHONPATH=/app \ - PATH=/opt/venv/bin:/usr/local/bin:$PATH \ - TF_CPP_MIN_LOG_LEVEL=2 \ - TF_ENABLE_ONEDNN_OPTS=0 \ - TOKENIZERS_PARALLELISM=true \ - PROJECT_ROOT=/workspace \ - HF_HUB_ENABLE_HF_TRANSFER=1 \ - DISABLE_PANDERA_IMPORT_WARNING=True \ - HF_HOME=/workspace/.cache/huggingface \ - TORCH_HOME=/workspace/.cache/torch \ - XDG_CACHE_HOME=/workspace/.cache \ - WANDB_DIR=/workspace/logs \ - TQDM_CACHE=/workspace/.cache/tqdm - -RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ - --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ - apt-get update && \ - apt-get install -y --no-install-recommends \ - build-essential curl git ca-certificates \ - python3.12 python3.12-dev python3.12-venv \ - ninja-build && \ - python3.12 -m venv /opt/venv && \ - ln -sf /opt/venv/bin/python /usr/local/bin/python && \ - ln -sf /opt/venv/bin/pip /usr/local/bin/pip - -WORKDIR /app - -COPY requirements.txt . -COPY official/ official/ - -RUN pip install --upgrade pip==26.1.1 setuptools==70.2.0 - -# Install official repos from submodules for compliance testing. -# - E1: pip install -e (needed by testing/official/e1.py) -# - DPLM: NOT installed (pins torchtext==0.17.0 which is incompatible). -# Compliance uses transformers.EsmForMaskedLM directly. Submodule is for reference only. -# - ESM (Biohub): NOT pip installed (uses the same top-level `esm` package name). -# testing/official/esm_plusplus.py adds the submodule to sys.path on demand. -RUN pip install -e /app/official/e1 - -RUN pip install -r requirements.txt -RUN pip install torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128 -RUN pip install numpy==1.26.4 - -COPY . . - -WORKDIR /workspace - -CMD ["bash"] diff --git a/Dockerfile.ankh b/Dockerfile.ankh deleted file mode 100644 index b9c2d82..0000000 --- a/Dockerfile.ankh +++ /dev/null @@ -1,8 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs ANKH. Native reference is `transformers.T5EncoderModel` -# -- already in base image. -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# docker build -f Dockerfile.ankh -t fastplms-ankh . -FROM fastplms-base diff --git a/Dockerfile.base b/Dockerfile.base deleted file mode 100644 index e72b072..0000000 --- a/Dockerfile.base +++ /dev/null @@ -1,60 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Base image for FastPLMs. Shared across every family image (Dockerfile.). -# Contains torch, transformers, pinned numpy, project source code -- everything -# that every family needs. Does NOT install any family-specific native/official -# package; those live in family Dockerfiles to avoid dependency conflicts -# (e.g. Biohub `esm` depends on a Biohub transformers fork, DPLM pins -# torchtext==0.17.0). -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# -# NOTE: switch to cudnn-devel if you need to compile CUDA extensions -# (e.g. flash-attn from source). -FROM nvidia/cuda:12.8.0-cudnn-runtime-ubuntu24.04 - -ENV DEBIAN_FRONTEND=noninteractive \ - PYTHONPATH=/app \ - PATH=/opt/venv/bin:/usr/local/bin:$PATH \ - TF_CPP_MIN_LOG_LEVEL=2 \ - TF_ENABLE_ONEDNN_OPTS=0 \ - TOKENIZERS_PARALLELISM=true \ - PROJECT_ROOT=/workspace \ - HF_HUB_ENABLE_HF_TRANSFER=1 \ - DISABLE_PANDERA_IMPORT_WARNING=True \ - HF_HOME=/workspace/.cache/huggingface \ - TORCH_HOME=/workspace/.cache/torch \ - XDG_CACHE_HOME=/workspace/.cache \ - WANDB_DIR=/workspace/logs \ - TQDM_CACHE=/workspace/.cache/tqdm - -RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ - --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ - apt-get update && \ - apt-get install -y --no-install-recommends \ - build-essential curl git ca-certificates \ - python3.12 python3.12-dev python3.12-venv \ - ninja-build && \ - python3.12 -m venv /opt/venv && \ - ln -sf /opt/venv/bin/python /usr/local/bin/python && \ - ln -sf /opt/venv/bin/pip /usr/local/bin/pip - -WORKDIR /app - -COPY requirements.txt . - -RUN pip install --upgrade pip==26.1.1 setuptools==70.2.0 -# Install cu128 torch BEFORE requirements.txt so transformers/accelerate/etc -# see torch already satisfied, then AGAIN after in case a transitive dep -# silently overwrites it with the PyPI default (CUDA 13) wheel. -RUN pip install torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128 -RUN pip install -r requirements.txt -RUN pip install --force-reinstall torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128 -RUN pip install numpy==1.26.4 - -# Source code is copied LAST so edits don't invalidate dep caches. -COPY . . - -WORKDIR /workspace - -CMD ["bash"] diff --git a/Dockerfile.dplm b/Dockerfile.dplm deleted file mode 100644 index d5fcdcf..0000000 --- a/Dockerfile.dplm +++ /dev/null @@ -1,11 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs DPLM. Native reference is `transformers.EsmForMaskedLM` -# (DPLM uses the ESM2 architecture internally) -- already in base image. -# DPLM's official package is NOT installed because it pins torchtext==0.17.0 -# which is incompatible with our torch pin. Submodule at official/dplm is for -# reference only. -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# docker build -f Dockerfile.dplm -t fastplms-dplm . -FROM fastplms-base diff --git a/Dockerfile.dplm2 b/Dockerfile.dplm2 deleted file mode 100644 index 80a9bc6..0000000 --- a/Dockerfile.dplm2 +++ /dev/null @@ -1,8 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs DPLM2. Same rationale as Dockerfile.dplm -- DPLM2 uses ESM2 -# architecture internally; native reference is `transformers.EsmForMaskedLM`. -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# docker build -f Dockerfile.dplm2 -t fastplms-dplm2 . -FROM fastplms-base diff --git a/Dockerfile.e1 b/Dockerfile.e1 deleted file mode 100644 index c2deb87..0000000 --- a/Dockerfile.e1 +++ /dev/null @@ -1,14 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs E1. Installs the official E1 package (editable) from the -# submodule at official/e1. Isolated in its own image because pip install -e -# modifies the venv with E1-specific deps that aren't needed for other families. -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# docker build -f Dockerfile.e1 -t fastplms-e1 . -FROM fastplms-base - -RUN pip install -e /app/official/e1 -# Re-pin cu128 torch in case the e1 install or one of its transitive deps -# pulled the PyPI default (CUDA 13) wheel on top of our cu128 base. -RUN pip install --force-reinstall torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128 diff --git a/Dockerfile.esm2 b/Dockerfile.esm2 deleted file mode 100644 index 65c6210..0000000 --- a/Dockerfile.esm2 +++ /dev/null @@ -1,11 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs ESM2 (also covers ESMFold, which reuses the ESM2 stack). -# Native reference is `transformers.EsmForMaskedLM` -- already in base image. -# -# Build: -# docker build -f Dockerfile.base -t fastplms-base . -# docker build -f Dockerfile.esm2 -t fastplms-esm2 . -# Run parity tests: -# docker run --gpus all --ipc=host --rm -v $(pwd):/workspace \ -# fastplms-esm2 python -m pytest /workspace/testing/test_parity.py -k esm2 -v -FROM fastplms-base diff --git a/Dockerfile.esm3 b/Dockerfile.esm3 deleted file mode 100644 index 1fca5d0..0000000 --- a/Dockerfile.esm3 +++ /dev/null @@ -1,19 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs ESM3 testing. -# Biohub `esm` is loaded from the official/esm submodule via sys.path. - -FROM fastplms-base - -RUN pip install \ - attrs==26.1.0 \ - biopython==1.87 \ - biotite==1.6.0 \ - brotli==1.2.0 \ - cloudpathlib==0.24.0 \ - httpx==0.28.1 \ - msgpack==1.1.2 \ - msgpack-numpy==0.4.8 \ - pandas==3.0.3 \ - pygtrie==2.5.0 \ - tenacity==9.1.4 \ - zstd==1.5.7.3 diff --git a/Dockerfile.esm_plusplus b/Dockerfile.esm_plusplus deleted file mode 100644 index 777fc56..0000000 --- a/Dockerfile.esm_plusplus +++ /dev/null @@ -1,28 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs ESM++ / ESMC parity testing. -# Adds Biohub `esm` package runtime dependencies on top of the base image. -# NOT using `pip install -e /app/official/esm` because Biohub/esm depends on a -# Biohub transformers fork. Instead we install the runtime deps `esm.models.esmc` -# actually imports, relying on sys.path injection via testing/official/__init__.py -# to bring in the esm package itself. -# -# Build: -# docker build -f Dockerfile.esm_plusplus -t fastplms-esm_plusplus . -# Run parity script: -# docker run --gpus all --ipc=host --rm -v $(pwd):/workspace \ -# fastplms-esm_plusplus python /workspace/testing/parity_debug_esmc.py - -FROM fastplms-base - -RUN pip install \ - attrs==26.1.0 \ - biopython==1.87 \ - biotite==1.6.0 \ - brotli==1.2.0 \ - cloudpathlib==0.24.0 \ - httpx==0.28.1 \ - msgpack-numpy==0.4.8 \ - pandas==3.0.3 \ - pygtrie==2.5.0 \ - tenacity==9.1.4 \ - zstd==1.5.7.3 diff --git a/Dockerfile.esmfold2 b/Dockerfile.esmfold2 deleted file mode 100644 index 1d02e82..0000000 --- a/Dockerfile.esmfold2 +++ /dev/null @@ -1,28 +0,0 @@ -# syntax=docker/dockerfile:1.7 -# Image for FastPLMs ESMFold2 testing. -# Includes the Biohub Transformers fork only for official parity tests. The -# FastPLMs ESMFold2 AutoModel runtime uses FastPLMs ESM++ as its LM backbone. - -FROM fastplms-base - -RUN pip install \ - abnumber \ - attrs==26.1.0 \ - biopython==1.87 \ - biotite==1.6.0 \ - brotli==1.2.0 \ - cloudpathlib==0.24.0 \ - dna_features_viewer==3.1.5 \ - httpx==0.28.1 \ - msgpack==1.1.2 \ - msgpack-numpy==0.4.8 \ - pandas==3.0.3 \ - pydssp==0.9.1 \ - pygtrie==2.5.0 \ - py3dmol==2.5.5 \ - rdkit==2026.3.2 \ - tenacity==9.1.4 \ - zstd==1.5.7.3 - -RUN pip install --force-reinstall --no-deps \ - git+https://github.com/Biohub/transformers.git@3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf diff --git a/LICENSE b/LICENSE index df9ff1a..66c354f 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -PLEASE NOTE THE APACHE LICENSE ONLY APPLIES TO THE CODE IN THE FastPLMs GITHUB AND ASSOCIATED HUGGINGFACE REPOSITORIES, NOT NECESSARILY THE MODEL WEIGHTS. THOSE LICENSES CAN BE FOUND HERE https://github.com/Synthyra/FastPLMs/tree/main/licenses +PLEASE NOTE THE APACHE LICENSE ONLY APPLIES TO THE CODE IN THE FastPLMs GITHUB AND ASSOCIATED HUGGINGFACE REPOSITORIES, NOT NECESSARILY THE MODEL WEIGHTS. THOSE LICENSES CAN BE FOUND HERE https://github.com/Synthyra/FastPLMs/tree/main/LICENSES Apache License Version 2.0, January 2004 diff --git a/LICENSES/README.md b/LICENSES/README.md new file mode 100644 index 0000000..d50d61b --- /dev/null +++ b/LICENSES/README.md @@ -0,0 +1,65 @@ +# Third-party license inventory + +This directory is the distributable legal inventory. Verbatim upstream files +are copied from the pinned repositories under `vendor/upstream/`. Supplemental +provenance and modified-file notices are maintained here because they describe +the FastPLMs distribution rather than an upstream repository. + +Keep distributable license files in this directory rather than beside runtime +modules under `src/fastplms/`. Model-specific license summaries belong in the +generated cards under `model_cards/`; the canonical legal texts and notices +listed here control. + +`src/fastplms/models.toml` records the SHA-256 digest of every legal file +required by a distributable model artifact. Artifact construction verifies the +declared upstream and distribution copies before writing output. Missing or +changed declared content is a release error. + +Digests use the UTF-8 text stored by Git with LF line endings. Validation +normalizes CRLF or LF checkouts to that canonical representation, and artifact +construction always writes LF. No other whitespace or content is normalized. + +## Required inventory + +| Source | Distribution file | Purpose | +|---|---|---| +| FastPLMs | root `LICENSE` | Apache-2.0 project code license with checkpoint caveat | +| FastPLMs | root `THIRD_PARTY_NOTICES.md` | Consolidated attribution and redistribution notice | +| ANKH | `ankh/LICENSE.md` | Verbatim CC BY-NC-SA 4.0 terms | +| Biohub ESM | `biohub-esm/LICENSE.md` | Verbatim MIT license | +| Biohub ESM | `biohub-esm/THIRD_PARTY_NOTICE.md` | Verbatim Biohub third-party notices | +| Biohub Transformers | `biohub-transformers/LICENSE` | Verbatim Apache-2.0 license | +| Boltz | `boltz/LICENSE` | Verbatim MIT license | +| DPLM | `dplm/LICENSE` | Verbatim Apache-2.0 license | +| DPLM | `dplm/PROVENANCE.md` | Pinned license scope for DPLM1 and DPLM2 checkpoint weights | +| Profluent-E1 | `e1/LICENSE` | Verbatim Profluent-E1 agreement | +| Profluent-E1 | `e1/ATTRIBUTION` | Verbatim attribution guidelines | +| Profluent-E1 | `e1/NOTICE` | Verbatim required notice | +| Profluent-E1 | `e1/Apache-2.0.txt` | Complete Apache-2.0 text for the model code | +| Profluent-E1 | `e1/BSD-3-Clause.txt` | Complete BSD-3-Clause text for the FlashAttention-derived utility identified upstream | +| Profluent-E1 | `e1/MODIFICATIONS.md` | FastPLMs modified-file notice and conversion identifier | +| Meta ESM | `fair-esm/LICENSE` | Verbatim MIT license | +| Meta ESM | `fair-esm/PROVENANCE.md` | Pinned revision and parity-oracle boundary | +| OpenFold | `openfold/LICENSE` | Verbatim Apache-2.0 license | +| OpenFold | `openfold/MODIFICATIONS.md` | FastPLMs modified-file notice | +| OpenFold | `openfold/PROVENANCE.md` | Pinned revision and parity-oracle boundary | +| ProteinTTT | `protein-ttt/LICENSE` | Verbatim MIT license | +| ProteinTTT | `protein-ttt/PROVENANCE.md` | Pinned revision and optional-workflow boundary | + +## Reference-only notices + +Some legal records apply only to isolated parity environments and are not +copied into model artifacts. `dllogger/PROVENANCE.md` records the pinned +DLLogger dependency and installed-license check used by the ESMFold reference +image. FastPLMs production code does not import or distribute DLLogger. + +The pinned E1 repository contains its agreement, attribution guidelines, and +notice, but does not include standalone Apache-2.0 or BSD-3-Clause files. The +complete standard texts are included here. The BSD notice follows the official +E1 source header that identifies `flash_attention_utils.py` as adapted from +Dao-AILab FlashAttention under BSD-3-Clause. + +FastPLMs does not enforce ANKH use restrictions in software. The pinned DPLM +repository's Apache-2.0 `LICENSE` and README scope the official release to the +pretrained DPLM1 and DPLM2 weights. The immutable evidence and FastPLMs +conversion boundary are recorded in `dplm/PROVENANCE.md`. diff --git a/LICENSES/ankh/LICENSE.md b/LICENSES/ankh/LICENSE.md new file mode 100644 index 0000000..8d7a310 --- /dev/null +++ b/LICENSES/ankh/LICENSE.md @@ -0,0 +1,353 @@ +Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International + +Creative Commons Corporation ("Creative Commons") is not a law firm and +does not provide legal services or legal advice. Distribution of +Creative Commons public licenses does not create a lawyer-client or +other relationship. Creative Commons makes its licenses and related +information available on an "as-is" basis. Creative Commons gives no +warranties regarding its licenses, any material licensed under their +terms and conditions, or any related information. Creative Commons +disclaims all liability for damages resulting from their use to the +fullest extent possible. + +Using Creative Commons Public Licenses + +Creative Commons public licenses provide a standard set of terms and +conditions that creators and other rights holders may use to share +original works of authorship and other material subject to copyright and +certain other rights specified in the public license below. The +following considerations are for informational purposes only, are not +exhaustive, and do not form part of our licenses. + +Considerations for licensors: Our public licenses are intended for use +by those authorized to give the public permission to use material in +ways otherwise restricted by copyright and certain other rights. Our +licenses are irrevocable. Licensors should read and understand the terms +and conditions of the license they choose before applying it. Licensors +should also secure all rights necessary before applying our licenses so +that the public can reuse the material as expected. Licensors should +clearly mark any material not subject to the license. This includes +other CC-licensed material, or material used under an exception or +limitation to copyright. More considerations for licensors : +wiki.creativecommons.org/Considerations\_for\_licensors + +Considerations for the public: By using one of our public licenses, a +licensor grants the public permission to use the licensed material under +specified terms and conditions. If the licensor's permission is not +necessary for any reason–for example, because of any applicable +exception or limitation to copyright–then that use is not regulated by +the license. Our licenses grant only permissions under copyright and +certain other rights that a licensor has authority to grant. Use of the +licensed material may still be restricted for other reasons, including +because others have copyright or other rights in the material. A +licensor may make special requests, such as asking that all changes be +marked or described. Although not required by our licenses, you are +encouraged to respect those requests where reasonable. More +considerations for the public : +wiki.creativecommons.org/Considerations\_for\_licensees + +Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International +Public License + +By exercising the Licensed Rights (defined below), You accept and agree +to be bound by the terms and conditions of this Creative Commons +Attribution-NonCommercial-ShareAlike 4.0 International Public License +("Public License"). To the extent this Public License may be interpreted +as a contract, You are granted the Licensed Rights in consideration of +Your acceptance of these terms and conditions, and the Licensor grants +You such rights in consideration of benefits the Licensor receives from +making the Licensed Material available under these terms and conditions. + +Section 1 – Definitions. + +- a. Adapted Material means material subject to Copyright and Similar + Rights that is derived from or based upon the Licensed Material and + in which the Licensed Material is translated, altered, arranged, + transformed, or otherwise modified in a manner requiring permission + under the Copyright and Similar Rights held by the Licensor. For + purposes of this Public License, where the Licensed Material is a + musical work, performance, or sound recording, Adapted Material is + always produced where the Licensed Material is synched in timed + relation with a moving image. +- b. Adapter's License means the license You apply to Your Copyright + and Similar Rights in Your contributions to Adapted Material in + accordance with the terms and conditions of this Public License. +- c. BY-NC-SA Compatible License means a license listed at + creativecommons.org/compatiblelicenses, approved by Creative Commons + as essentially the equivalent of this Public License. +- d. Copyright and Similar Rights means copyright and/or similar + rights closely related to copyright including, without limitation, + performance, broadcast, sound recording, and Sui Generis Database + Rights, without regard to how the rights are labeled or categorized. + For purposes of this Public License, the rights specified in Section + 2(b)(1)-(2) are not Copyright and Similar Rights. +- e. Effective Technological Measures means those measures that, in + the absence of proper authority, may not be circumvented under laws + fulfilling obligations under Article 11 of the WIPO Copyright Treaty + adopted on December 20, 1996, and/or similar international + agreements. +- f. Exceptions and Limitations means fair use, fair dealing, and/or + any other exception or limitation to Copyright and Similar Rights + that applies to Your use of the Licensed Material. +- g. License Elements means the license attributes listed in the name + of a Creative Commons Public License. The License Elements of this + Public License are Attribution, NonCommercial, and ShareAlike. +- h. Licensed Material means the artistic or literary work, database, + or other material to which the Licensor applied this Public License. +- i. Licensed Rights means the rights granted to You subject to the + terms and conditions of this Public License, which are limited to + all Copyright and Similar Rights that apply to Your use of the + Licensed Material and that the Licensor has authority to license. +- j. Licensor means the individual(s) or entity(ies) granting rights + under this Public License. +- k. NonCommercial means not primarily intended for or directed + towards commercial advantage or monetary compensation. For purposes + of this Public License, the exchange of the Licensed Material for + other material subject to Copyright and Similar Rights by digital + file-sharing or similar means is NonCommercial provided there is no + payment of monetary compensation in connection with the exchange. +- l. Share means to provide material to the public by any means or + process that requires permission under the Licensed Rights, such as + reproduction, public display, public performance, distribution, + dissemination, communication, or importation, and to make material + available to the public including in ways that members of the public + may access the material from a place and at a time individually + chosen by them. +- m. Sui Generis Database Rights means rights other than copyright + resulting from Directive 96/9/EC of the European Parliament and of + the Council of 11 March 1996 on the legal protection of databases, + as amended and/or succeeded, as well as other essentially equivalent + rights anywhere in the world. +- n. You means the individual or entity exercising the Licensed Rights + under this Public License. Your has a corresponding meaning. + +Section 2 – Scope. + +- a. License grant. + - 1. Subject to the terms and conditions of this Public License, + the Licensor hereby grants You a worldwide, royalty-free, + non-sublicensable, non-exclusive, irrevocable license to + exercise the Licensed Rights in the Licensed Material to: + - A. reproduce and Share the Licensed Material, in whole or in + part, for NonCommercial purposes only; and + - B. produce, reproduce, and Share Adapted Material for + NonCommercial purposes only. + - 2. Exceptions and Limitations. For the avoidance of doubt, where + Exceptions and Limitations apply to Your use, this Public + License does not apply, and You do not need to comply with its + terms and conditions. + - 3. Term. The term of this Public License is specified in Section + 6(a). + - 4. Media and formats; technical modifications allowed. The + Licensor authorizes You to exercise the Licensed Rights in all + media and formats whether now known or hereafter created, and to + make technical modifications necessary to do so. The Licensor + waives and/or agrees not to assert any right or authority to + forbid You from making technical modifications necessary to + exercise the Licensed Rights, including technical modifications + necessary to circumvent Effective Technological Measures. For + purposes of this Public License, simply making modifications + authorized by this Section 2(a)(4) never produces Adapted + Material. + - 5. Downstream recipients. + - A. Offer from the Licensor – Licensed Material. Every + recipient of the Licensed Material automatically receives an + offer from the Licensor to exercise the Licensed Rights + under the terms and conditions of this Public License. + - B. Additional offer from the Licensor – Adapted Material. + Every recipient of Adapted Material from You automatically + receives an offer from the Licensor to exercise the Licensed + Rights in the Adapted Material under the conditions of the + Adapter's License You apply. + - C. No downstream restrictions. You may not offer or impose + any additional or different terms or conditions on, or apply + any Effective Technological Measures to, the Licensed + Material if doing so restricts exercise of the Licensed + Rights by any recipient of the Licensed Material. + - 6. No endorsement. Nothing in this Public License constitutes or + may be construed as permission to assert or imply that You are, + or that Your use of the Licensed Material is, connected with, or + sponsored, endorsed, or granted official status by, the Licensor + or others designated to receive attribution as provided in + Section 3(a)(1)(A)(i). +- b. Other rights. + - 1. Moral rights, such as the right of integrity, are not + licensed under this Public License, nor are publicity, privacy, + and/or other similar personality rights; however, to the extent + possible, the Licensor waives and/or agrees not to assert any + such rights held by the Licensor to the limited extent necessary + to allow You to exercise the Licensed Rights, but not otherwise. + - 2. Patent and trademark rights are not licensed under this + Public License. + - 3. To the extent possible, the Licensor waives any right to + collect royalties from You for the exercise of the Licensed + Rights, whether directly or through a collecting society under + any voluntary or waivable statutory or compulsory licensing + scheme. In all other cases the Licensor expressly reserves any + right to collect such royalties, including when the Licensed + Material is used other than for NonCommercial purposes. + +Section 3 – License Conditions. + +Your exercise of the Licensed Rights is expressly made subject to the +following conditions. + +- a. Attribution. + - 1. If You Share the Licensed Material (including in modified + form), You must: + - A. retain the following if it is supplied by the Licensor + with the Licensed Material: + - i. identification of the creator(s) of the Licensed + Material and any others designated to receive + attribution, in any reasonable manner requested by the + Licensor (including by pseudonym if designated); + - ii. a copyright notice; + - iii. a notice that refers to this Public License; + - iv. a notice that refers to the disclaimer of + warranties; + - v. a URI or hyperlink to the Licensed Material to the + extent reasonably practicable; + + - B. indicate if You modified the Licensed Material and retain + an indication of any previous modifications; and + - C. indicate the Licensed Material is licensed under this + Public License, and include the text of, or the URI or + hyperlink to, this Public License. + - 2. You may satisfy the conditions in Section 3(a)(1) in any + reasonable manner based on the medium, means, and context in + which You Share the Licensed Material. For example, it may be + reasonable to satisfy the conditions by providing a URI or + hyperlink to a resource that includes the required information. + - 3. If requested by the Licensor, You must remove any of the + information required by Section 3(a)(1)(A) to the extent + reasonably practicable. +- b. ShareAlike.In addition to the conditions in Section 3(a), if You + Share Adapted Material You produce, the following conditions also + apply. + - 1. The Adapter's License You apply must be a Creative Commons + license with the same License Elements, this version or later, + or a BY-NC-SA Compatible License. + - 2. You must include the text of, or the URI or hyperlink to, the + Adapter's License You apply. You may satisfy this condition in + any reasonable manner based on the medium, means, and context in + which You Share Adapted Material. + - 3. You may not offer or impose any additional or different terms + or conditions on, or apply any Effective Technological Measures + to, Adapted Material that restrict exercise of the rights + granted under the Adapter's License You apply. + +Section 4 – Sui Generis Database Rights. + +Where the Licensed Rights include Sui Generis Database Rights that apply +to Your use of the Licensed Material: + +- a. for the avoidance of doubt, Section 2(a)(1) grants You the right + to extract, reuse, reproduce, and Share all or a substantial portion + of the contents of the database for NonCommercial purposes only; +- b. if You include all or a substantial portion of the database + contents in a database in which You have Sui Generis Database + Rights, then the database in which You have Sui Generis Database + Rights (but not its individual contents) is Adapted Material, + including for purposes of Section 3(b); and +- c. You must comply with the conditions in Section 3(a) if You Share + all or a substantial portion of the contents of the database. + For the avoidance of doubt, this Section 4 supplements and does not + replace Your obligations under this Public License where the + Licensed Rights include other Copyright and Similar Rights. + +Section 5 – Disclaimer of Warranties and Limitation of Liability. + +- a. Unless otherwise separately undertaken by the Licensor, to the + extent possible, the Licensor offers the Licensed Material as-is and + as-available, and makes no representations or warranties of any kind + concerning the Licensed Material, whether express, implied, + statutory, or other. This includes, without limitation, warranties + of title, merchantability, fitness for a particular purpose, + non-infringement, absence of latent or other defects, accuracy, or + the presence or absence of errors, whether or not known or + discoverable. Where disclaimers of warranties are not allowed in + full or in part, this disclaimer may not apply to You. +- b. To the extent possible, in no event will the Licensor be liable + to You on any legal theory (including, without limitation, + negligence) or otherwise for any direct, special, indirect, + incidental, consequential, punitive, exemplary, or other losses, + costs, expenses, or damages arising out of this Public License or + use of the Licensed Material, even if the Licensor has been advised + of the possibility of such losses, costs, expenses, or damages. + Where a limitation of liability is not allowed in full or in part, + this limitation may not apply to You. +- c. The disclaimer of warranties and limitation of liability provided + above shall be interpreted in a manner that, to the extent possible, + most closely approximates an absolute disclaimer and waiver of all + liability. + +Section 6 – Term and Termination. + +- a. This Public License applies for the term of the Copyright and + Similar Rights licensed here. However, if You fail to comply with + this Public License, then Your rights under this Public License + terminate automatically. +- b. Where Your right to use the Licensed Material has terminated + under Section 6(a), it reinstates: + + - 1. automatically as of the date the violation is cured, provided + it is cured within 30 days of Your discovery of the violation; + or + - 2. upon express reinstatement by the Licensor. + + For the avoidance of doubt, this Section 6(b) does not affect any + right the Licensor may have to seek remedies for Your violations of + this Public License. + +- c. For the avoidance of doubt, the Licensor may also offer the + Licensed Material under separate terms or conditions or stop + distributing the Licensed Material at any time; however, doing so + will not terminate this Public License. +- d. Sections 1, 5, 6, 7, and 8 survive termination of this Public + License. + +Section 7 – Other Terms and Conditions. + +- a. The Licensor shall not be bound by any additional or different + terms or conditions communicated by You unless expressly agreed. +- b. Any arrangements, understandings, or agreements regarding the + Licensed Material not stated herein are separate from and + independent of the terms and conditions of this Public License. + +Section 8 – Interpretation. + +- a. For the avoidance of doubt, this Public License does not, and + shall not be interpreted to, reduce, limit, restrict, or impose + conditions on any use of the Licensed Material that could lawfully + be made without permission under this Public License. +- b. To the extent possible, if any provision of this Public License + is deemed unenforceable, it shall be automatically reformed to the + minimum extent necessary to make it enforceable. If the provision + cannot be reformed, it shall be severed from this Public License + without affecting the enforceability of the remaining terms and + conditions. +- c. No term or condition of this Public License will be waived and no + failure to comply consented to unless expressly agreed to by the + Licensor. +- d. Nothing in this Public License constitutes or may be interpreted + as a limitation upon, or waiver of, any privileges and immunities + that apply to the Licensor or You, including from the legal + processes of any jurisdiction or authority. + +Creative Commons is not a party to its public licenses. Notwithstanding, +Creative Commons may elect to apply one of its public licenses to +material it publishes and in those instances will be considered the +"Licensor." The text of the Creative Commons public licenses is +dedicated to the public domain under the CC0 Public Domain Dedication. +Except for the limited purpose of indicating that material is shared +under a Creative Commons public license or as otherwise permitted by the +Creative Commons policies published at creativecommons.org/policies, +Creative Commons does not authorize the use of the trademark "Creative +Commons" or any other trademark or logo of Creative Commons without its +prior written consent including, without limitation, in connection with +any unauthorized modifications to any of its public licenses or any +other arrangements, understandings, or agreements concerning use of +licensed material. For the avoidance of doubt, this paragraph does not +form part of the public licenses. + +Creative Commons may be contacted at creativecommons.org. diff --git a/fastplms/esmfold2/LICENSE b/LICENSES/biohub-esm/LICENSE.md similarity index 100% rename from fastplms/esmfold2/LICENSE rename to LICENSES/biohub-esm/LICENSE.md diff --git a/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md b/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md new file mode 100644 index 0000000..56d8280 --- /dev/null +++ b/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md @@ -0,0 +1,13 @@ +The code in this repository depends on the following third-party libraries: + +| Library | License | Link | +|----------|----------|----------| +| flash-attn | BSD | https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE | +| PyTorch | BSD | https://github.com/pytorch/pytorch/blob/main/LICENSE | +| xformers | BSD | https://github.com/facebookresearch/xformers/blob/main/LICENSE | +| jaxtyping | MIT | https://github.com/patrick-kidger/jaxtyping/blob/main/LICENSE | +| einops | MIT | https://github.com/arogozhnikov/einops/blob/main/LICENSE | +| omegaconf | BSD | https://github.com/omry/omegaconf/blob/master/LICENSE | +| attrs | MIT | https://github.com/python-attrs/attrs/blob/main/LICENSE | +| scipy | BSD-3-Clause | https://github.com/scipy/scipy/blob/main/LICENSE.txt
https://github.com/scipy/scipy/blob/main/LICENSES_bundled.txt | +| lightning / torchmetrics | Apache 2.0 | https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE | diff --git a/LICENSES/biohub-transformers/LICENSE b/LICENSES/biohub-transformers/LICENSE new file mode 100644 index 0000000..68b7d66 --- /dev/null +++ b/LICENSES/biohub-transformers/LICENSE @@ -0,0 +1,203 @@ +Copyright 2018- The Hugging Face team. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/fastplms/boltz/LICENSE b/LICENSES/boltz/LICENSE similarity index 99% rename from fastplms/boltz/LICENSE rename to LICENSES/boltz/LICENSE index 2421734..a9d6575 100644 --- a/fastplms/boltz/LICENSE +++ b/LICENSES/boltz/LICENSE @@ -18,4 +18,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. \ No newline at end of file +SOFTWARE. diff --git a/LICENSES/dllogger/PROVENANCE.md b/LICENSES/dllogger/PROVENANCE.md new file mode 100644 index 0000000..d0dc4be --- /dev/null +++ b/LICENSES/dllogger/PROVENANCE.md @@ -0,0 +1,16 @@ +# NVIDIA DLLogger provenance + +- Upstream: `https://github.com/NVIDIA/dllogger` +- Revision: `0478734ff7be75adde8d160e04872664d1c62e5f` +- Installed version: `1.1.0` +- License: Apache-2.0 +- Scope: isolated native ESMFold reference image only + +The pinned OpenFold package imports DLLogger eagerly from its package +initializer. FastPLMs production code does not import or depend on DLLogger. +The historical OpenFold environment selected DLLogger from an unpinned Git +URL, so the reference image fixes the revision above for reproducibility. + +DLLogger's wheel contains its verbatim license at +`dllogger-1.1.0.dist-info/licenses/LICENSE`. The reference image preserves that +file and fails its build if the license is absent or is not the Apache License. diff --git a/LICENSES/dplm/LICENSE b/LICENSES/dplm/LICENSE new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/LICENSES/dplm/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSES/dplm/PROVENANCE.md b/LICENSES/dplm/PROVENANCE.md new file mode 100644 index 0000000..0912632 --- /dev/null +++ b/LICENSES/dplm/PROVENANCE.md @@ -0,0 +1,23 @@ +# DPLM checkpoint license provenance + +FastPLMs uses the ByteDance DPLM repository at immutable revision +`8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d` as the official source for both +DPLM1 and DPLM2. + +At that revision: + +- the repository contains the complete [Apache License 2.0](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE); and +- the [official README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview) + defines the repository release as including the pretrained weights for the + DPLM family, specifically DPLM1 and DPLM2, alongside training and inference + implementations. + +FastPLMs therefore records the official DPLM1 and DPLM2 checkpoint weights as +Apache-2.0. Converted Synthyra checkpoints retain that license and include the +verbatim upstream `LICENSE`. The deterministic conversion identifiers are +`dplm_to_fastplms_v1` and `dplm2_to_fastplms_v1`; neither adds restrictions to +the upstream terms. + +Complete publication is permitted only after the ordinary FastPLMs artifact, +state-parity, legal-inventory, and atomic-publication checks pass. This record +does not change the terms of third-party training data or downstream outputs. diff --git a/LICENSES/e1/ATTRIBUTION b/LICENSES/e1/ATTRIBUTION new file mode 100644 index 0000000..07a4a8f --- /dev/null +++ b/LICENSES/e1/ATTRIBUTION @@ -0,0 +1,34 @@ +Profluent-E1 Attribution Guidelines +----------------------------------- + +These Profluent-E1 Attribution Guidelines (the “Guidelines”) set forth certain terms, conditions and restrictions relating to how You can use, share and distribute Profluent-E1 and any Derivative Works thereof in a manner that ensures transparent acknowledgement of Profluent-E1. The terms, conditions and restrictions set forth in these Guidelines are in addition to those set forth in the Profluent-E1 Clickthrough License Agreement, available at https://github.com/Profluent-AI/E1/blob/main/LICENSE, as may be updated or amended from time to time (the “Agreement”) and in the event of any conflict between any provision of these Guidelines and any provision of the Agreement, the provision that is more protective of Profluent Bio, Inc. (“Profluent”) and Profluent-E1 shall control. Any capitalized terms used but not defined herein shall have the meanings set forth in the Agreement. Profluent may update, modify or amend these Guidelines from time to time and Profluent will use reasonable efforts to provide You with notice of any material changes that may negatively impact Your use of Profluent-E1 by posting an updated version of these Guidelines, for example. Any violation of these Guidelines may result in the suspension or termination of Your access to and use of Profluent-E1. + +1. Attribution Requirements. You may use and redistribute Profluent-E1 and Derivative Works in Source and Object forms, with or without modification, in accordance with the terms of the Agreement provided that the following conditions are met: + + 1.1 You must include a prominently displayed attribution to “Profluent-E1” in any distribution, documentation or public use of Profluent-E1 or any Derivative Works built, developed, improved, or modified through the use of Profluent-E1 by any Commercial Entity, including, without limitation, in any publication, report, white paper, article, product documentation, slide deck or other media that incorporates Profluent-E1 or any Derivative Works; and + + 1.2 You must identify any drug, drug candidate, medication or pharmaceutical of any kind and any molecular or biological target, hit or lead identification that is created, identified, discovered, developed, modified or improved using Profluent-E1 as “Built with Profluent-E1,” including, without limitation, in public disclosures to the United States Food and Drug Administration and similar regulatory authorities. + +The term “Commercial Entity” as used in this Section 1 shall mean any entity engaged in any activity intended for or directed toward commercial advantage or monetary compensation, including, without limitation, the development of any product or service intended to be sold or made available for a fee, expressly excluding any university, non-profit organization, research institute, educational or government body and any not-for-profit Legal Entity. + +2. No Endorsement. You may not use or allow the use of the Profluent name nor any Profluent trade name, trademark, service mark or product name to endorse or promote, or imply that Profluent endorses or promotes, any products or other offerings derived from Profluent-E1 without Profluent's specific prior written permission. + +3. Format of Derivative Works. All attribution notices should be easy to find and easy to understand. The attribution requirements set forth in these Guidelines applies to both textual and visual materials where Profluent-E1 contributes materially to the resulting product or output. Any uses, distributions or documentation in Source form must incorporate a prominent display of text that reads “Profluent-E1” in verifiable form, while redistributions in Object form must be accompanied by a prominent display of “Profluent-E1” text in any documentation and each time the resulting executable program or a program dependent thereon is launched, You must also launch a prominent display (e.g., splash screen or banner text) of such text. + +4. Placement and Prominence. any entity engaged in any activity intended for or directed toward commercial advantage or monetary compensation, including, without limitation, the development of any product or service intended to be sold or made available for a fee, expressly excluding any university, non-profit organization, research institute, educational or government body and any not-for-profit Legal Entity. The attribution must appear within the main body of any publication, report, white paper or article. For product documentation or software tool or interface references, the attribution must be listed on the title page or in the introductory section, near other third-party license notices. In slide decks or conference presentations, the attribution must appear on a title slide, or a slide within the first three slides of the presentation. For webpages, software interfaces, or interactive demos, attribution must be visible on-screen in a footer, “About” section, or a prominently accessible credits area without requiring extra actions like scrolling or menu navigation. + +5. Formatting. You must use the exact “Profluent-E1” model name as provided by Profluent, and the font size must be at least equal to the main body text or at least 75% of the largest text on the page where the attribution is incorporated (whichever is greater). The text must be visually legible and in a color that contrasts sufficiently with its background. Graphic or visual logos provided by the Profluent may be used in place of plain-text attribution if approved, but they must remain clearly readable at normal viewing sizes. + +6. Illustrative Examples. The following are examples of how attribution to Profluent-E1 may be implemented in accordance with the requirements of these guidelines: + + 6.1 Example 1: A research article that incorporates data derived from outputs of Profluent-E1 includes reference to “Profluent-E1” in or beneath the paragraph of text first introducing these outputs. + + 6.2 Example 2: A company releases a model fine-tuned from Profluent-E1 and the website's footer includes “Built with Profluent-E1.” + + 6.3 Example 3: A presentation at a conference includes “This work incorporates works derived from use of Profluent-E1” on the title slide. + + 6.4 Example 4: An open-source repository “README” file includes a visible statement near the top: “This project builds upon Profluent-E1 under the terms of the Profluent-E1 Clickthrough License Agreement.” + +7. Non-Obscuration. The attribution must not be hidden or disguised by other visual or textual elements. It should remain readable at standard display resolutions and must not be abbreviated, truncated or replaced with unofficial branding. + +8. Exceptions. Written exceptions to these attribution requirements may be granted by the Licensor upon request, provided that alternative acknowledgment conveys equivalent visibility and transparency. diff --git a/LICENSES/e1/Apache-2.0.txt b/LICENSES/e1/Apache-2.0.txt new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/LICENSES/e1/Apache-2.0.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSES/e1/BSD-3-Clause.txt b/LICENSES/e1/BSD-3-Clause.txt new file mode 100644 index 0000000..08dc525 --- /dev/null +++ b/LICENSES/e1/BSD-3-Clause.txt @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2022, Tri Dao. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from this + software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/fastplms/e1/LICENSE b/LICENSES/e1/LICENSE similarity index 91% rename from fastplms/e1/LICENSE rename to LICENSES/e1/LICENSE index 8e6acc3..5a4accb 100644 --- a/fastplms/e1/LICENSE +++ b/LICENSES/e1/LICENSE @@ -1,22 +1,3 @@ -Your use of the Profluent-E1 model code is governed by the Apache License 2.0, while your use of the Profluent-E1 model weights and the full release of the Profluent-E1 model is governed by a similarly permissive license with additional attribution requirements - see the NOTICE file for details. You can use, share, and modify Profluent-E1 for free, but you must follow our ATTRIBUTION guidelines to give credit, include the license when you share and follow some other basic rules. Profluent is not responsible for what you build, and may terminate your rights to use Profluent-E1 if you breach the license. - -Code in src/E1/model/flash_attention_utils.py is adapted from flash-attention project under BSD-3-Clause license. - -Profluent-E1 Notice File ------------------------- - -Copyright 2025 Profluent Bio Inc. - -Licensed under the Profluent-E1 Clickthrough License Agreement (the “Agreement”); you may not use -this file except in compliance with the Agreement. Unless required by applicable law or agreed to -in writing, software distributed under the Agreement is distributed on an "AS IS" BASIS, WITHOUT -WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the Agreement for the specific -language governing permissions and limitations under the Agreement and the Attribution Guidelines -for more information about attribution requirements for use of Profluent-E1. You may obtain a copy -of the Agreement at https://github.com/Profluent-AI/E1/blob/main/LICENSE and a copy of the Attribution -Guidelines at https://github.com/Profluent-AI/E1/blob/main/ATTRIBUTION, each as may be updated or -amended from time to time. - Profluent-E1 Clickthrough License Agreement ------------------------------------------- diff --git a/LICENSES/e1/MODIFICATIONS.md b/LICENSES/e1/MODIFICATIONS.md new file mode 100644 index 0000000..ad46b4a --- /dev/null +++ b/LICENSES/e1/MODIFICATIONS.md @@ -0,0 +1,27 @@ +# Profluent-E1 modified-file notice + +FastPLMs implements Profluent-E1 behavior against the pinned official source at +revision `bfd2620a602248499f3d2583d85a7ecddf0b6e02`. The FastPLMs files listed +below are modified or independently reorganized implementations of the +corresponding E1 interfaces. They are not byte-for-byte copies of the upstream +files. + +| FastPLMs file | Modification notice | +|---|---| +| `src/fastplms/models/e1/modeling_e1.py` | Reorganized for Transformers AutoClasses, shared attention selection, sequence and RAG preparation, task heads, and checkpoint-compatible loading. | +| `src/fastplms/models/e1/attention.py` | Isolated the declared SDPA and Flex Attention paths and their masking contracts. | +| `src/fastplms/models/e1/preparation.py` | Reorganized raw-sequence, boundary-token, and retrieval-context preparation. | +| `src/fastplms/models/e1/cache.py` | Adapted the cache interface used by the reorganized Transformers implementation. | +| `src/fastplms/models/e1/retrieval.py` | Adapted retrieval helpers and their FastPLMs model outputs. | +| `src/fastplms/models/e1/__init__.py` | Added FastPLMs package exports. | +| `tools/conversion/state_transforms.py` | Added the deterministic `e1_to_fastplms_v1` checkpoint mapping. | + +These changes were present in the FastPLMs 1.0 repository as reviewed on +2026-07-20. The conversion identifier is `e1_to_fastplms_v1`. Recipients must +retain the Profluent-E1 agreement, `ATTRIBUTION`, `NOTICE`, this modified-file +notice, and the applicable Apache-2.0 and BSD-3-Clause texts. + +The BSD-3-Clause component is the padding utility identified by the official E1 +repository as adapted from Dao-AILab FlashAttention. FastPLMs does not copy that +official utility into its production package, but preserves the notice because +the official source is the parity oracle for the E1 behavior contract. diff --git a/LICENSES/e1/NOTICE b/LICENSES/e1/NOTICE new file mode 100644 index 0000000..e8c1aaa --- /dev/null +++ b/LICENSES/e1/NOTICE @@ -0,0 +1,14 @@ +Profluent-E1 Notice File +------------------------ + +Copyright 2025 Profluent Bio Inc. + +Licensed under the Profluent-E1 Clickthrough License Agreement (the “Agreement”); you may not use +this file except in compliance with the Agreement. Unless required by applicable law or agreed to +in writing, software distributed under the Agreement is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the Agreement for the specific +language governing permissions and limitations under the Agreement and the Attribution Guidelines +for more information about attribution requirements for use of Profluent-E1. You may obtain a copy +of the Agreement at https://github.com/Profluent-AI/E1/blob/main/LICENSE and a copy of the Attribution +Guidelines at https://github.com/Profluent-AI/E1/blob/main/ATTRIBUTION, each as may be updated or +amended from time to time. diff --git a/fastplms/esm2/LICENSE b/LICENSES/fair-esm/LICENSE similarity index 92% rename from fastplms/esm2/LICENSE rename to LICENSES/fair-esm/LICENSE index 2a155d2..b93be90 100644 --- a/fastplms/esm2/LICENSE +++ b/LICENSES/fair-esm/LICENSE @@ -1,5 +1,3 @@ -License for FastESM models, from the ESM2 repo https://github.com/facebookresearch/esm - MIT License Copyright (c) Meta Platforms, Inc. and affiliates. diff --git a/LICENSES/fair-esm/PROVENANCE.md b/LICENSES/fair-esm/PROVENANCE.md new file mode 100644 index 0000000..6b2059f --- /dev/null +++ b/LICENSES/fair-esm/PROVENANCE.md @@ -0,0 +1,8 @@ +# Meta ESM provenance + +FastPLMs uses `facebookresearch/esm` revision +`2b369911bb5b4b0dda914521b9475cad1656b2ac` as the official parity oracle for +ESM2 and ESMFold. The repository is pinned at +`vendor/upstream/fair-esm/` and is not a production dependency or runtime image +component. The accompanying `LICENSE` is the verbatim MIT text from that +revision. diff --git a/LICENSES/openfold/LICENSE b/LICENSES/openfold/LICENSE new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/LICENSES/openfold/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSES/openfold/MODIFICATIONS.md b/LICENSES/openfold/MODIFICATIONS.md new file mode 100644 index 0000000..6265484 --- /dev/null +++ b/LICENSES/openfold/MODIFICATIONS.md @@ -0,0 +1,25 @@ +# OpenFold reference-build modification + +FastPLMs pins `aqlaboratory/openfold` revision +`4b41059694619831a7db195b7e0988fc4ff3a307` as an ESMFold parity oracle. +The pinned checkout under `vendor/upstream/openfold/` is not modified. + +The `reference-esmfold` image applies +`docker/constraints/openfold-sm90.patch` to the copied `setup.py` before +installing OpenFold. The patch replaces OpenFold's build-time fallback CUDA +architecture list with `sm90` and changes the extension build flag from C++14 +to C++17. Docker BuildKit cannot observe the workstation GPU, so the original +setup otherwise requests legacy architectures that CUDA 12.1 no longer +accepts. PyTorch 2.2 requires C++17, and the H100 requires an `sm90` extension. + +This is a build-only packaging change. It does not alter OpenFold model classes, +extension source, checkpoint data, or the public API used by the native oracle. +The resulting reference image is specific to the declared H100 validation +environment. + +OpenFold imports PyTorch Lightning and NVIDIA DLLogger from its package +initializers even though ESMFold inference does not use their training or +logging features. The native image therefore pins PyTorch Lightning `1.9.5`, +TorchMetrics `0.11.4`, Lightning Utilities `0.15.2`, and NVIDIA DLLogger commit +`0478734ff7be75adde8d160e04872664d1c62e5f`. These Apache-2.0 dependencies are +native-reference-only and are not FastPLMs runtime dependencies. diff --git a/LICENSES/openfold/PROVENANCE.md b/LICENSES/openfold/PROVENANCE.md new file mode 100644 index 0000000..a8ecb8c --- /dev/null +++ b/LICENSES/openfold/PROVENANCE.md @@ -0,0 +1,7 @@ +# OpenFold provenance + +FastPLMs uses `aqlaboratory/openfold` revision +`4b41059694619831a7db195b7e0988fc4ff3a307` as an official ESMFold parity +oracle. The repository is pinned at `vendor/upstream/openfold/` and is not a +production dependency or runtime image component. The accompanying `LICENSE` +is the verbatim Apache-2.0 text from that revision. diff --git a/fastplms/esm3/LICENSE b/LICENSES/protein-ttt/LICENSE similarity index 86% rename from fastplms/esm3/LICENSE rename to LICENSES/protein-ttt/LICENSE index 462eb1d..2e8ce69 100644 --- a/fastplms/esm3/LICENSE +++ b/LICENSES/protein-ttt/LICENSE @@ -1,10 +1,6 @@ -License for FastPLMs ESM3 - -This derivative is built from Biohub ESM3. Biohub ESM is released under the MIT License. - MIT License -Copyright 2026 Chan Zuckerberg Biohub, Inc. +Copyright (c) 2024 Anton Bushuiev Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/LICENSES/protein-ttt/PROVENANCE.md b/LICENSES/protein-ttt/PROVENANCE.md new file mode 100644 index 0000000..b218841 --- /dev/null +++ b/LICENSES/protein-ttt/PROVENANCE.md @@ -0,0 +1,8 @@ +# ProteinTTT provenance + +FastPLMs uses `anton-bushuiev/ProteinTTT` revision +`fde2817cd84b936167cc76ccabf31e5c0fe49962` as the official reference for the +optional protein test-time training workflow. The repository is pinned at +`vendor/upstream/protein-ttt/` and is not a production dependency or runtime +image component. The accompanying `LICENSE` is the verbatim MIT text from that +revision. diff --git a/README.md b/README.md index 71b8c36..5d0bf91 100644 --- a/README.md +++ b/README.md @@ -2,745 +2,920 @@ FastPLMs Hero Image -FastPLMs is an open-source initiative dedicated to making protein language models (pLMs) efficient and easy to use. By replacing native, often suboptimal attention implementations with **Flash Attention** or **Flex Attention**, we provide high-performance alternatives that are fully compatible with the HuggingFace `transformers` ecosystem and can easily be loaded with no extra code with `AutoModel`. +FastPLMs maintains the runtime code published with Hugging Face protein +language and structure models. This repository is the source, test, and +dependency workspace for those model repositories. It is not an installable +Python distribution. + +Published models keep the familiar Transformers interface while making +attention, embedding, generation, folding, and validation behavior explicit. +Their runtime code does not import an official model checkout. Each supported +family instead has a pinned upstream source under `vendor/upstream/`, an +immutable checkpoint identity, and a declared state transformation. Release +workflows compare the resulting Hugging Face artifact against that official +source. + +## Contents + +- [Why FastPLMs](#why-fastplms) +- [Supported models](#supported-models) +- [Dependencies](#dependencies) +- [Quick start](#quick-start) +- [Usage examples](#usage-examples) +- [Attention backends](#attention-backends) +- [Design choices](#design-choices) +- [Validation and reproducibility](#validation-and-reproducibility) +- [Files-only Hub publication](#files-only-hub-publication) +- [Documentation](#documentation) +- [Contributing and citation](#contributing-and-citation) + +## Why FastPLMs + +Protein models are often published with architecture-specific loading code, +tokenization conventions, attention implementations, and output formats. +FastPLMs separates those concerns: + +- Transformers auto classes provide a consistent loading interface. +- Model-specific adapters preserve native biological token and structure + semantics. +- A shared embedding API returns ordered, residue-aware representations. +- Attention backends are explicit capabilities rather than silent fallbacks. +- The model registry records checkpoint, source, conversion, precision, + license, and test contracts in one typed manifest. +- Official repositories are isolated parity oracles, not production + dependencies. + +A supported model is more than an architecture implementation. Its +configuration, input preparation, checkpoint conversion, representative +inference, artifact contents, and legal inventory are all defined and tested. + +## Supported models + +The generated [support matrix](docs/generated/support.md) is the authoritative +checkpoint list. It is rendered from +[`src/fastplms/models.toml`](src/fastplms/models.toml), which also defines valid +AutoClasses, attention backends, precision paths, and release tiers. + +| Family | Primary use | Typical user input | Important distinction | +| --- | --- | --- | --- | +| ESM2 | Sequence representations and masked language modeling | Amino-acid sequences tokenized to residue IDs | Preserves ESM2 encoder, MLM, contact, and classification contracts | +| ESM++ / ESMC | Sequence representations and masked language modeling | Amino-acid sequences tokenized to residue IDs | Biohub ESMC implementation and ESMFold2 language-model backbone | +| ESM3 | Multimodal protein modeling and generation | Sequence, structure, and function tracks | Retains all three tracks through its multimodal interface | +| E1 | Retrieval-augmented protein encoding | Raw amino-acid sequences | No tokenizer; native E1 preparation is preserved | +| DPLM | Discrete diffusion protein generation | Amino-acid sequences tokenized to masked residue IDs | Confidence-based iterative unmasking | +| DPLM2 | Amino-acid and structure co-generation | Amino-acid and structure token tracks | Separate structure and amino-acid boundary tokens | +| ANKH | T5 protein encoding and sequence-to-sequence modeling | Amino-acid sequences tokenized for encoder or seq2seq use | The 1.0 artifact contract is one full official-compatible encoder-decoder checkpoint with encoder-default embeddings | +| ESMFold | Sequence-to-structure inference | Raw amino-acid sequences | Meta ESMFold contract with FastPLMs ESM2 backbone | +| ESMFold2 | Sequence and complex structure prediction | Raw amino-acid sequences or complex specifications | Full variants have 48 folding blocks and optional MSA conditioning; Fast variants have 24 blocks and no MSA conditioning | +| Boltz2 | Structure prediction | Raw amino-acid sequences or prepared model features | Provisional end-to-end numerical-equivalence status | + +The model manifest, not this summary, controls support. A backend or AutoClass +that is valid for one family may be rejected by another. + +The Synthyra ANKH repositories contain the complete 1.0 encoder-decoder +checkpoints. `AutoModel` loads the encoder view, and +`AutoModelForSeq2SeqLM` loads the decoder, cross-attention, and language-model +head from the same repository. + +## Dependencies + +Published FastPLMs model repositories require Python 3.11 through 3.14, +PyTorch 2.13, and Transformers 5.13. Install those runtime dependencies +directly, then load the model from Hugging Face: ---- +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" +``` -## Table of Contents -1. [Introduction](#introduction) -2. [Documentation](#documentation) -3. [Supported Models](#supported-models) -4. [Attention Backends](#attention-backends) -5. [Embedding & Pooling](#embedding--pooling) -6. [Experimental Test-Time Training](#experimental-test-time-training) -7. [Concrete Examples](#concrete-examples) -8. [Testing & Benchmarking](#testing--benchmarking) -9. [Installation & Docker](#installation--docker) +For a source checkout used to build or validate artifacts, install the named +dependency profile and put `src` on `PYTHONPATH` when running repository code: ---- +```bash +git clone https://github.com/Synthyra/FastPLMs.git +cd FastPLMs +uv venv +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu +PYTHONPATH=src python -m pytest tests/cpu -m cpu_contract +``` -## Documentation +Official reference repositories are not runtime or routine-development +dependencies. Initialize them only for a live release-candidate compliance run, +as described under [Validation and reproducibility](#validation-and-reproducibility). -Detailed documentation is available in the [`docs/`](docs/) folder: - -- [Architecture Overview](docs/architecture.md) - How FastPLMs wraps official models, the attention backend system, Docker layout -- [Per-Model Guides](docs/models.md) - Loading, configuration, and special handling for each model family -- [Attention Backends](docs/attention_backends.md) - SDPA, Flash, Flex, Auto: how they work, when to use each, numerical properties -- [Embedding & Pooling API](docs/embedding_api.md) - Pooler strategies, `embed_dataset()` parameters, SQLite/pth storage -- [Binder Design Example](docs/binder_design.md) - FastPLMs-only ESMFold2 plus ESM++ binder optimization, CLI, metrics, and EGFR result -- [Fine-Tuning Guide](docs/finetuning.md) - LoRA, Trainer patterns, dataset classes, metrics -- [Testing & Benchmarking](docs/testing.md) - Docker commands, pytest markers, compliance architecture, throughput benchmarks -- [Contributing](docs/contributing.md) - Code style, adding new models, required tests - ---- - -## Introduction - -### What are Protein Language Models (pLMs)? -Protein Language Models are transformer-based architectures trained on massive datasets of protein sequences (such as UniProt). These models learn the "grammar" of proteins, capturing evolutionary information, structural constraints, and functional motifs. They are used for: -- **Representation Learning**: Generating high-dimensional embeddings for downstream tasks (e.g., stability, function prediction). -- **Protein Generation**: Designing novel sequences with specific properties. -- **Structure Prediction**: Mapping sequences to their 3D folds (e.g., Boltz2). - -### What is this repository? -FastPLMs provides optimized versions of these models. Our focus is on: -- **Speed**: Drastically faster inference through optimized attention kernels. -- **Memory Efficiency**: Lower VRAM usage, enabling larger batch sizes or longer sequences. -- **Seamless Integration**: Use `AutoModel.from_pretrained(..., trust_remote_code=True)` to load our optimized weights directly from HuggingFace. - ---- - -## Supported Models - -We maintain a comprehensive [HuggingFace Collection](https://huggingface.co/collections/Synthyra/pretrained-plms-675351ecc050f63baedd77de) of optimized models. Below is a summary of the supported families and their origins. - -### Model Registry Summary - -| Model Family | Organization | Official Implementation | FastPLMs Optimization | Checkpoints | -| :--- | :--- | :--- | :--- | :--- | -| **E1** | Profluent Bio | [Profluent-Bio/E1](https://github.com/Profluent-Bio/E1) | Flex Attention, Block-Causal | 150M, 300M, 600M | -| **ESM2** | Meta AI | [facebookresearch/esm](https://github.com/facebookresearch/esm) | Flash (SDPA) / Flex Attention | 8M, 35M, 150M, 650M, 3B | -| **ESM++** | Biohub | [Biohub/esm](https://github.com/Biohub/esm) | Optimized SDPA / Flex | Small (300M), Large (600M), 6B | -| **ESM3** | Biohub | [Biohub/esm](https://github.com/Biohub/esm) | HF AutoModel wrapper | Open Small | -| **ESMFold2** | Biohub | [Biohub/esm](https://github.com/Biohub/esm) | Self-contained HF AutoModel wrapper with FastPLMs ESM++ LM backbone, opt-in experimental TTT | Full, Fast, Experimental, Cutoff2025 | -| **DPLM** | ByteDance | [bytedance/dplm](https://github.com/bytedance/dplm) | Diffusion Optimized Attention | 150M, 650M, 3B | -| **DPLM2** | ByteDance | [bytedance/dplm](https://github.com/bytedance/dplm) | Multimodal Diffusion | 150M, 650M, 3B | -| **ANKH** | Elnaggar Lab | [ElnaggarLab/ankh](https://huggingface.co/ElnaggarLab/ankh-base) | T5 RPE via Flex score_mod | Base, Large, ANKH2-L, ANKH3-L, ANKH3-XL | -| **ESMFold** | Meta AI | [facebookresearch/esm](https://github.com/facebookresearch/esm) | Fast ESM2 backbone, opt-in experimental ProteinTTT | Standard | -| **Boltz2** | MIT / Various | [jwohlwend/boltz](https://github.com/jwohlwend/boltz) | Optimized Structure Prediction | Standard | - -### Full Model List - -| Model Key | Family | Parameters | Organization | FastPLMs Repo ID | Official Reference | -| :--- | :--- | :--- | :--- | :--- | :--- | -| `esm2_8m` | ESM2 | 7.5M | Meta AI | [Synthyra/ESM2-8M](https://huggingface.co/Synthyra/ESM2-8M) | [facebook/esm2_t6_8M_UR50D](https://huggingface.co/facebook/esm2_t6_8M_UR50D) | -| `esm2_35m` | ESM2 | 33.5M | Meta AI | [Synthyra/ESM2-35M](https://huggingface.co/Synthyra/ESM2-35M) | [facebook/esm2_t12_35M_UR50D](https://huggingface.co/facebook/esm2_t12_35M_UR50D) | -| `esm2_150m` | ESM2 | 148.2M | Meta AI | [Synthyra/ESM2-150M](https://huggingface.co/Synthyra/ESM2-150M) | [facebook/esm2_t30_150M_UR50D](https://huggingface.co/facebook/esm2_t30_150M_UR50D) | -| `esm2_650m` | ESM2 | 651.1M | Meta AI | [Synthyra/ESM2-650M](https://huggingface.co/Synthyra/ESM2-650M) | [facebook/esm2_t33_650M_UR50D](https://huggingface.co/facebook/esm2_t33_650M_UR50D) | -| `esm2_3b` | ESM2 | 2.84B | Meta AI | [Synthyra/ESM2-3B](https://huggingface.co/Synthyra/ESM2-3B) | [facebook/esm2_t36_3B_UR50D](https://huggingface.co/facebook/esm2_t36_3B_UR50D) | -| `esmplusplus_small` | ESM++ | 333.0M | Biohub | [Synthyra/ESMplusplus_small](https://huggingface.co/Synthyra/ESMplusplus_small) | [biohub/ESMC-300M](https://huggingface.co/biohub/ESMC-300M) | -| `esmplusplus_large` | ESM++ | 575.0M | Biohub | [Synthyra/ESMplusplus_large](https://huggingface.co/Synthyra/ESMplusplus_large) | [biohub/ESMC-600M](https://huggingface.co/biohub/ESMC-600M) | -| `esmplusplus_6b` | ESM++ | 6.35B | Biohub | [Synthyra/ESMplusplus_6B](https://huggingface.co/Synthyra/ESMplusplus_6B) | [biohub/ESMC-6B](https://huggingface.co/biohub/ESMC-6B) | -| `esm3_small` | ESM3 | 1.4B | Biohub | [Synthyra/ESM3_small](https://huggingface.co/Synthyra/ESM3_small) | [biohub/esm3-sm-open-v1](https://huggingface.co/biohub/esm3-sm-open-v1) | -| `esmfold2` | ESMFold2 | 234.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2](https://huggingface.co/Synthyra/ESMFold2) | [biohub/ESMFold2](https://huggingface.co/biohub/ESMFold2) | -| `esmfold2_fast` | ESMFold2 | 188.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2-Fast](https://huggingface.co/Synthyra/ESMFold2-Fast) | [biohub/ESMFold2-Fast](https://huggingface.co/biohub/ESMFold2-Fast) | -| `esmfold2_experimental_fast` | ESMFold2 | 188.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2-Experimental-Fast](https://huggingface.co/Synthyra/ESMFold2-Experimental-Fast) | [biohub/ESMFold2-Experimental-Fast](https://huggingface.co/biohub/ESMFold2-Experimental-Fast) | -| `esmfold2_experimental_fast_cutoff2025` | ESMFold2 | 188.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2-Experimental-Fast-Cutoff2025](https://huggingface.co/Synthyra/ESMFold2-Experimental-Fast-Cutoff2025) | [biohub/ESMFold2-Experimental-Fast-Cutoff2025](https://huggingface.co/biohub/ESMFold2-Experimental-Fast-Cutoff2025) | -| `esmfold2_experimental` | ESMFold2 | 234.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2-Experimental](https://huggingface.co/Synthyra/ESMFold2-Experimental) | [biohub/ESMFold2-Experimental](https://huggingface.co/biohub/ESMFold2-Experimental) | -| `esmfold2_experimental_cutoff2025` | ESMFold2 | 234.8M + ESM++ 6B | Biohub | [Synthyra/ESMFold2-Experimental-Cutoff2025](https://huggingface.co/Synthyra/ESMFold2-Experimental-Cutoff2025) | [biohub/ESMFold2-Experimental-Cutoff2025](https://huggingface.co/biohub/ESMFold2-Experimental-Cutoff2025) | -| `e1_150m` | E1 | 154.4M | Profluent Bio | [Synthyra/Profluent-E1-150M](https://huggingface.co/Synthyra/Profluent-E1-150M) | [Profluent-Bio/E1-150m](https://huggingface.co/Profluent-Bio/E1-150m) | -| `e1_300m` | E1 | 274.3M | Profluent Bio | [Synthyra/Profluent-E1-300M](https://huggingface.co/Synthyra/Profluent-E1-300M) | [Profluent-Bio/E1-300m](https://huggingface.co/Profluent-Bio/E1-300m) | -| `e1_600m` | E1 | 641.4M | Profluent Bio | [Synthyra/Profluent-E1-600M](https://huggingface.co/Synthyra/Profluent-E1-600M) | [Profluent-Bio/E1-600m](https://huggingface.co/Profluent-Bio/E1-600m) | -| `dplm_150m` | DPLM | 148.2M | ByteDance | [Synthyra/DPLM-150M](https://huggingface.co/Synthyra/DPLM-150M) | [airkingbd/dplm_150m](https://huggingface.co/airkingbd/dplm_150m) | -| `dplm_650m` | DPLM | 651.1M | ByteDance | [Synthyra/DPLM-650M](https://huggingface.co/Synthyra/DPLM-650M) | [airkingbd/dplm_650m](https://huggingface.co/airkingbd/dplm_650m) | -| `dplm_3b` | DPLM | 2.84B | ByteDance | [Synthyra/DPLM-3B](https://huggingface.co/Synthyra/DPLM-3B) | [airkingbd/dplm_3b](https://huggingface.co/airkingbd/dplm_3b) | -| `dplm2_150m` | DPLM2 | 158.7M | ByteDance | [Synthyra/DPLM2-150M](https://huggingface.co/Synthyra/DPLM2-150M) | [airkingbd/dplm2_150m](https://huggingface.co/airkingbd/dplm2_150m) | -| `dplm2_650m` | DPLM2 | 672.1M | ByteDance | [Synthyra/DPLM2-650M](https://huggingface.co/Synthyra/DPLM2-650M) | [airkingbd/dplm2_650m](https://huggingface.co/airkingbd/dplm2_650m) | -| `dplm2_3b` | DPLM2 | 2.88B | ByteDance | [Synthyra/DPLM2-3B](https://huggingface.co/Synthyra/DPLM2-3B) | [airkingbd/dplm2_3b](https://huggingface.co/airkingbd/dplm2_3b) | -| `ankh_base` | ANKH | 453.3M | Elnaggar Lab | [Synthyra/ANKH_base](https://huggingface.co/Synthyra/ANKH_base) | [ElnaggarLab/ankh-base](https://huggingface.co/ElnaggarLab/ankh-base) | -| `ankh_large` | ANKH | 1.15B | Elnaggar Lab | [Synthyra/ANKH_large](https://huggingface.co/Synthyra/ANKH_large) | [ElnaggarLab/ankh-large](https://huggingface.co/ElnaggarLab/ankh-large) | -| `ankh2_large` | ANKH | 1.15B | Elnaggar Lab | [Synthyra/ANKH2_large](https://huggingface.co/Synthyra/ANKH2_large) | [ElnaggarLab/ankh2-ext2](https://huggingface.co/ElnaggarLab/ankh2-ext2) | -| `ankh3_large` | ANKH | 1.15B | Elnaggar Lab | [Synthyra/ANKH3_large](https://huggingface.co/Synthyra/ANKH3_large) | [ElnaggarLab/ankh3-large](https://huggingface.co/ElnaggarLab/ankh3-large) | -| `ankh3_xl` | ANKH | 3.49B | Elnaggar Lab | [Synthyra/ANKH3_xl](https://huggingface.co/Synthyra/ANKH3_xl) | [ElnaggarLab/ankh3-xl](https://huggingface.co/ElnaggarLab/ankh3-xl) | -| `esmfold` | ESMFold | 3.53B | Meta AI | [Synthyra/FastESMFold](https://huggingface.co/Synthyra/FastESMFold) | [facebookresearch/esm](https://github.com/facebookresearch/esm) | -| `boltz2` | Boltz2 | 506.3M | MIT / Various | [Synthyra/Boltz2](https://huggingface.co/Synthyra/Boltz2) | [jwohlwend/boltz](https://github.com/jwohlwend/boltz) | - ---- - -## Experimental Test-Time Training - -FastPLMs includes experimental ProteinTTT-style test-time training utilities for -sequence PLMs and the PLM backbones used by ESMFold and ESMFold2. TTT is -disabled by default. Normal `from_pretrained`, `forward`, `embed_dataset`, -`fold_protein`, and `state_dict()` behavior is unchanged unless you explicitly -call `ttt()`, `fold_protein(..., ttt=True)`, or `fold_protein_ttt()`. - -TTT briefly adapts a model to one input protein using masked language modeling -and local LoRA adapters. It can improve predictions for difficult or -low-confidence proteins, especially structure predictions with weak baseline -pLDDT, but it is not guaranteed to help. It adds GPU memory use and runtime, -can degrade already confident predictions, and should be treated as an -experimental test-time compute option. - -Supported opt-in paths: - -| Model family | TTT API | Notes | -| :--- | :--- | :--- | -| ESM2, ESM++, ESM3, E1, DPLM, DPLM2, ANKH | `model.ttt(seq=...)` | MLM LoRA adaptation of the PLM backbone only | -| FastESMFold | `model.fold_protein(sequence, ttt=True)` or `model.fold_protein_ttt(sequence)` | Returns the best pLDDT fold across baseline and TTT steps | -| ESMFold2 | `model.fold_protein(sequence, ttt=True, ttt_config=...)` or `model.fold_protein_ttt(sequence)` | Protein-only v1 path, trains LoRA only on the ESM++ `_esmc` backbone | -| Boltz2 | Not supported | Boltz2 remains inference-only in FastPLMs | - -Sequence PLM example: +Direct dependency declarations live under `requirements/`. Profiles compose +the combinations exercised by local and container validation: -```python -from transformers import AutoModelForMaskedLM +| Dependency file or profile | Purpose | +| --- | --- | +| `profiles/runtime.in` | Core dependencies embedded model code can import | +| `profiles/cpu-validation.in` | CPU-only tests, linting, typing, structure preparation, and training contracts | +| `profiles/candidate.in` | CUDA candidate validation with FlashAttention support | +| `profiles/candidate-structure.in` | CUDA structure and binder validation | +| `profiles/candidate-fp8.in` | Experimental ESMFold2 ESMC FP8 validation | +| `features/binder.in` | AbNumber, ANARCII, pandas, and PyArrow for binder design | +| `features/cueq.in` | Optional cuEquivariance frontend and CUDA 13 kernels | +| `features/reporting.in` | Plots and statistical reports | -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-8M", - trust_remote_code=True, -).cuda().eval() +See [`requirements/README.md`](requirements/README.md) for the complete profile +layout. Official reference implementations are never runtime dependencies. -# No adapters are injected until this call. -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -``` +## Quick start -ESMFold2 example: +Load a model with its standard Transformers auto class: ```python from transformers import AutoModel model = AutoModel.from_pretrained( - "Synthyra/ESMFold2-Fast", + "Synthyra/ESM2-150M", trust_remote_code=True, - load_esmc=True, - esmc_attn_backend="flex", -).cuda().eval() - -result = model.fold_protein( - "MSTNPKPQRKTKRNT", - num_loops=1, - num_sampling_steps=10, - ttt=True, - ttt_config={"steps": 1, "ags": 1, "batch_size": 1}, + attn_implementation="sdpa", ) -print(result.ttt_metrics) +model.eval() ``` -ESMFold2 loads `Synthyra/ESMplusplus_6B` as its LM backbone by default. -Legacy `biohub/ESMC-*` IDs in older configs are accepted as migration aliases -and normalized to the matching Synthyra ESM++ checkpoint. -For FP8 LM inference, pass `esmc_precision="fp8"` and install -`transformer_engine.pytorch` in a CUDA environment with FP8-capable hardware. -FP8 is inference-only, so TTT remains a bf16/fp32 path. - -If you use the TTT functionality, cite ProteinTTT in addition to FastPLMs and -the underlying model papers. The ProteinTTT citation is listed in -[Citations](#esmfold--proteinttt). - ---- - -### License Notes - -Biohub ESM++ and ESM3 model cards include `license: mit` metadata and upload a `LICENSE` file copied from the Biohub ESM MIT license. The upstream license is linked from each model card and from the source repository at https://github.com/Biohub/esm/blob/main/LICENSE.md. - ---- - -## Attention Backends +Use the published Hub identifier for ordinary loading. Pin `revision` when an +immutable model-code snapshot is required. Contributors can instead build the +manifest-pinned artifact under `dist/hub/ESM2-150M` and load it locally before +publication. -All FastPLMs models share a common set of attention backends, controlled via `config.attn_backend`. The default is `"sdpa"`, which is safe on all hardware and numerically equivalent to standard attention. +Tokenizer-based models use the tokenizer paired with the same artifact: -### Backend Comparison - -| Backend | Key | Speed | Numerical Equivalence | Availability | -| :--- | :--- | :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Fast | Exact | Any PyTorch ≥ 2.0 | -| Flash Attention | `"kernels_flash"` | Fastest | Approximate | Requires `pip install kernels` (pre-built) | -| Flex Attention | `"flex"` | Very fast | ~Exact | Requires PyTorch ≥ 2.11 (FA4 backend on Hopper/Blackwell) | -| Auto | `"auto"` | — | — | Always (selects best available) | - -### SDPA (default) - -PyTorch's [`scaled_dot_product_attention`](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) dispatches to a fused CUDA kernel (cuDNN or efficient attention) that is faster and more memory-efficient than naive attention, while being mathematically identical to it. This is the recommended default for reproducibility and general use. It is also the only backend where `output_attentions=True` is handled natively; with other backends, attentions are computed via a separate naive matrix multiplication when requested. - -### Flash Attention (`kernels_flash`) +```python +import torch +from transformers import AutoTokenizer -Flash Attention 2 and 3 are typically the fastest options on Ampere (A100) and Hopper (H100) GPUs, often 2–4× faster than SDPA at long sequence lengths. Flash Attention achieves this by tiling the computation and applying an online softmax, which means the results are **not bitwise identical** to SDPA or naive attention. Differences are on the order of floating-point rounding and are often inconsequential for standard inference — but they are not guaranteed to be so. They can compound across layers, interact with low-precision dtypes (fp16/bf16), or affect sensitive downstream tasks. Flash Attention is standard practice in large model training and the trade-off is well understood, but it should not be treated as a drop-in numerical equivalent of SDPA. If exact reproducibility or numerical sensitivity is a concern, use `"sdpa"` instead. +tokenizer = AutoTokenizer.from_pretrained( + "Synthyra/ESM2-150M", + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) -**No compilation required.** FastPLMs uses the HuggingFace [`kernels`](https://github.com/huggingface/kernels) package to load pre-built Flash Attention 2/3 binaries at runtime — no C++ compiler, no CUDA toolkit version pinning, no waiting: +with torch.inference_mode(): + output = model(**batch) -``` -pip install kernels +print(output.last_hidden_state.shape) ``` -Building `flash-attn` from source is notoriously painful. The Ninja build system parallelizes aggressively across all available CPU cores, and each NVCC/CICC compiler process it spawns can consume **5–8 GB of RAM on its own**. On a 64-core machine this can push peak RAM usage to **~300 GB**, and even on a throttled single-threaded build (`MAX_JOBS=1 NVCC_THREADS=1`) the compile still takes **many hours** while grinding through paging. Pre-built community wheels cover 384+ version/GPU/CUDA/platform combinations and still routinely fall short of matching a user's exact environment. This is the point where most people give up and go without Flash Attention entirely. The `kernels` package sidesteps all of this by fetching a pre-compiled binary matched to your GPU architecture (SM80 for Ampere, SM90 for Hopper). If no compatible binary exists for your hardware, it gracefully falls back to `flex` or `sdpa` rather than erroring. +E1 is intentionally different. It has no tokenizer and retains its native +raw-sequence preparation path. Use `model.embed_dataset(...)` for ordinary +sequence representations or the explicit E1 preparation methods for lower +level token tensors. -### Flex Attention (`flex`) +## Usage examples -PyTorch's [`flex_attention`](https://pytorch.org/docs/stable/nn.attention.flex_attention.html) (PyTorch >= 2.11 in FastPLMs Docker images) generates a fused Triton kernel customized to the mask pattern at hand. It is numerically very close to SDPA, typically within floating-point rounding of naive computation. The primary advantage is that it can apply a **block mask** that skips padding tokens entirely, providing a meaningful speedup on batches with variable-length sequences (no compute wasted on padding). E1 uses a block-causal variant of this mask. +Calls made directly on a model loaded with `trust_remote_code=True` use the +runtime bundled in its Hugging Face repository. Examples that import +`fastplms` modules are contributor workflows and require `PYTHONPATH=src`. -The **first forward pass** triggers JIT compilation via Triton, which can take 30–120 seconds. All subsequent calls are fast. Combining with `torch.compile` yields the best sustained throughput. +### Ordered sequence embeddings -### Auto (`auto`) +The model method is available from the Hugging Face artifact. The source-level +function shown here is for repository work run with `PYTHONPATH=src`: -Automatically selects the best available backend in order of preference: `kernels_flash` → `flex` → `sdpa`. Useful when you want maximum speed without configuring the environment manually, and you accept that the resolved backend may differ across machines. +```python +from fastplms import EmbeddingInput, embed_dataset + +result = embed_dataset( + model, + [ + EmbeddingInput("protein-a", "MSTNPKPQRKTKRNT"), + EmbeddingInput("protein-a", "MKTIIALSYIFCLVFA"), + ], + batch_size=2, + pooling=("mean", "std"), + output="embeddings", +) +``` -### Setting the Backend +Insertion-ordered mappings are also accepted and preserve their keys as record +identifiers: -**At load time (every family):** ```python -from transformers import AutoConfig, AutoModel - -config = AutoConfig.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -config.attn_backend = "flex" # "sdpa", "kernels_flash", "flex", or "auto" -model = AutoModel.from_pretrained("Synthyra/ESM2-150M", config=config, trust_remote_code=True) +result = model.embed_dataset( + { + "protein-a": "MSTNPKPQRKTKRNT", + "protein-b": "MKTIIALSYIFCLVFA", + }, + batch_size=2, + pooling="mean", +) ``` -**After load time (every family):** - -Every family's `PreTrainedModel` subclass exposes a mutable `attn_backend` property whose setter propagates the change to every attention submodule in-place, so you can swap backends on a loaded model without reloading the weights: +`EmbeddingResult` preserves order, duplicate identifiers, and the original +sequence. Each record contains an identifier, sequence, and tensor: ```python -model = AutoModel.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -model.attn_backend = "flex" # every attention layer now uses flex -model.attn_backend = "kernels_flash" # flip again, no reload +for record in result.records: + tensor = record.load_tensor() + print(record.id, record.sequence, tensor.shape) ``` -This is handy for benchmarking backends on the same weights or for falling back at runtime if a backend is unavailable. The setter asserts if the requested backend isn't installed on the current GPU (e.g. `kernels_flash` without the `kernels` package). - -### Returning Attention Maps +Calling `result.as_dict(key="id")` raises when identifiers repeat unless an +explicit duplicate policy is provided. This avoids silently overwriting FASTA +records. -All backends support `output_attentions=True`. For the optimized backends (SDPA, Flash Attention, Flex), attention weights are computed via a separate naive matrix multiplication and appended to the output — so enabling this negates the memory savings of those backends. Use it only for inspection or contact prediction, not during high-throughput inference. +Safetensors output packs generation-scoped shards across batches and can resume +from the last flushed shard after interruption. An interrupted in-memory shard +is recomputed. Tensor memory is bounded by the configured shard size rather than +the full dataset size. Successful overwrites retain prior immutable generations +so already-open lazy readers remain valid. Stale generations are removed only +through explicit, dry-run-first garbage collection after the caller guarantees +there are no active readers or writers. SQLite remains available when database +transactions and queryable records are preferred. ---- +### FASTA input and in-memory output -## Embedding & Pooling +Pass a FASTA path directly. Multi-line sequences and record identifiers are +preserved: -The `EmbeddingMixin` (shared across all models) provides a standardized way to extract representations from proteins. +```python +result = model.embed_dataset( + "proteins.fasta", + batch_size=32, + batch_window_size=128, + max_tokens_per_batch=4096, + max_length=1024, + pooling=("mean", "max"), +) -### The Pooler -The `Pooler` class aggregates sequence-level residue representations into a single fixed-size vector. Supported strategies include: -- `mean`: Mask-aware average of all residues. -- `cls`: The first token's representation (Standard for classification). -- `max`: Element-wise maximum across the sequence. -- `var` / `std`: Variance or Standard Deviation of representations. -- `norm`: L2 normalization. -- `median`: Element-wise median. -- `parti`: Experimental PageRank-based attention pooling. +for record in result: + print(record.id, record.tensor.shape) +``` ---- +When `output` is omitted, tensors remain in memory. Multiple poolers are +concatenated in request order, and `result.metadata["pool_slices"]` records the +slice belonging to each transformation. FASTA is streamed line by line into an +immutable fingerprinted spool. Length bucketing is bounded by +`batch_window_size`, which defaults to sixteen times `batch_size`; output order +is restored, and `max_length` always counts biological residues rather than +tokenizer-added special tokens. SQLite prefixes commit at completed batch-window +boundaries. Safetensors prefixes publish only when a bounded shard flushes, so +an interruption replays the unflushed shard rather than necessarily one window. +Set the window equal to the batch size when per-batch SQLite checkpoint +boundaries matter. -## Concrete Examples +### Full residue and hidden-state embeddings -### 1. Batch Embedding with SQLite (Scalable) -Ideal for embedding millions of sequences where you need to stream data or avoid OOM on RAM. +Use full embeddings when a downstream task needs one vector per biological +residue: ```python -import torch -from transformers import AutoModel +residue_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) -model = AutoModel.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True).cuda() +for record in residue_result: + print(record.sequence, record.tensor.shape) +``` -sequences = ["MALWMRLLPLLALLALWGPDPAAA", "MKTIIALSYIFCLVFA", ...] +Each tensor has shape `(l, d)` after BOS, EOS, padding, chain delimiters, and +other non-biological positions are removed. To retain every returned model +state: -# Embed and store in SQLite -model.embed_dataset( - sequences=sequences, - batch_size=64, - pooling_types=['mean', 'cls'], # Concatenates both - sql=True, - sql_db_path='large_protein_db.db', - embed_dtype=torch.float32 +```python +state_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, + store_all_hidden_states=True, ) +print(state_result[0].tensor.shape) ``` -### 2. Embedding from a FASTA File -Pass a FASTA file path directly — no manual parsing required. Multi-line sequences are handled automatically. You can combine `fasta_path` with an explicit `sequences` list and the two sources are merged before embedding. +The hidden-state tensor has shape `(n, l, d)`, where `n` follows the model's +native hidden-state output order. + +ANKH selects the encoder final state by default. Decoder layers require the +full sequence-to-sequence view and explicit decoder inputs. Pass proteins as +raw residue strings without inserted spaces and keep sentinel prompts tight, +for example `M`. The model-owned and explicitly supplied tokenizer +paths apply the same ANKH normalization contract. ```python -# Embed all sequences in a FASTA file and save to SQLite -model.embed_dataset( - fasta_path='my_proteins.fasta', - batch_size=64, - pooling_types=['mean'], - sql=True, - sql_db_path='my_proteins.db', -) +from transformers import AutoModelForSeq2SeqLM -# Mix a FASTA file with an explicit list -model.embed_dataset( - sequences=["MKTIIALSYIFCLVFA"], - fasta_path='additional_proteins.fasta', - batch_size=32, - save=True, - save_path='combined_embeddings.pth', +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH_base", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, ) ``` -### 3. High-Throughput In-Memory Embedding -Perfect for medium-sized datasets that fit in memory. - -```python -# Embed and return as a dictionary -embeddings = model.embed_dataset( - sequences=sequences, - batch_size=128, - pooling_types=['mean'], - save=True, - save_path='my_embeddings.pth' -) +There is no implicit shifted-source decoder contract. Official ANKH tasks use +task-dependent prompts, sentinels, or generated tokens. Set +`hidden_state_source="encoder"` and `hidden_state_index` for any encoder layer, +or `store_all_hidden_states=True` for every layer in the selected stack. -# Access embedding -seq_vector = embeddings["MALWMRLLPLLALLALWGPDPAAA"] # torch.Tensor -``` +### Safetensors output and exact resume -### 4. Custom Pooling & Multi-Strategy -Concatenate multiple mathematical representations for richer downstream features. +Directory output uses sharded safetensors by default: ```python -# Use a variety of pooling types -embeddings = model.embed_dataset( - sequences=sequences, - pooling_types=['mean', 'max', 'std', 'var'], # All 4 concatenated - batch_size=32, - full_embeddings=False +result = model.embed_dataset( + "proteins.fasta", + batch_size=64, + pooling=("mean",), + output="artifacts/protein-embeddings", + format="safetensors", + resume=True, ) - -# Resulting vector size: 4 * hidden_size -print(embeddings[sequences[0]].shape) ``` -### 5. FastPLMs Binder Design With ESMFold2 And ESM++ -FastPLMs includes a binder design workflow that mirrors the Biohub ESMFold2 -tutorial while using only FastPLMs model repos. ESMFold2 experimental checkpoints -provide differentiable folding losses and final critics, while ESM++ provides the -masked-LM pseudoperplexity regularizer. +The directory contains tensor and descriptor shards, immutable generation +indexes, the `index.json` convenience pointer, and authoritative `run.json` +commit marker. Generation descriptors record sequence order, identifiers, +shapes, dtypes, hashes, and shard keys. The run manifest binds the input, model +state, tokenizer policy, backend, and pooling configuration. Resume is accepted +only when existing records are the exact ordered prefix of the same run. +`run.json` selects an immutable generation; older generations remain available +to readers opened before a successful overwrite. See +[Embedding API](docs/embedding_api.md#safetensors-storage) for the exclusive +garbage-collection contract. -![FastPLMs EGFR minibinder design](docs/assets/egfr_fastplms_binder_design.png) +### SQLite streaming -Run the verified EGFR 128 amino acid de novo minibinder example on a CUDA -workstation with the ESMFold2 Docker image: +SQLite commits each completed batch window and is useful for long jobs: -```bash -cd /home/ubuntu/FastPLMs - -sudo -n docker run --gpus all --rm \ - -v /home/ubuntu/FastPLMs:/app \ - -v /home/ubuntu/FastPLMs:/workspace \ - -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \ - -w /workspace fastplms-esmfold2 \ - python /app/cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli +```python +result = model.embed_dataset( + "proteins.fasta", + batch_size=64, + pooling=("mean", "std"), + output="artifacts/protein-embeddings.sqlite", + format="sqlite", + resume=True, +) ``` -The run writes `trajectory.jsonl`, `best_sequences.fasta`, `results.parquet`, -`selection.parquet`, and per-critic PDB/CIF/logit files. The verified result had -hero mean iPTM `0.913870`, hero min iPTM `0.904600`, and all four hero -ESMFold2 critics above `0.9`. +BF16 tensors are stored as raw bytes with an explicit dtype. Resume rejects a +changed model state, input order, sequence, tokenizer, pooling configuration, +or attention backend. -Binder sequence: +SQLite readers open results read-only and preserve ordered duplicate filters: -```text -SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ -``` - -See [docs/binder_design.md](docs/binder_design.md) for the full strategy, -official selection rule, Modal backend, per-critic metrics, and caveats. +```python +from fastplms.embeddings import load_sqlite_result ---- +selected = load_sqlite_result( + "artifacts/protein-embeddings.sqlite", + record_ids=["protein-b", "protein-a", "protein-b"], +) +``` -## Testing & Benchmarking +### DPLM sequence generation -FastPLMs includes a pytest-based test suite under `testing/` covering correctness, compliance, and performance. All GPU tests run inside Docker. See [docs/testing.md](docs/testing.md) for the full guide. +DPLM starts from a tokenized sequence whose biological positions define the +requested length. Iterative unmasking fills those positions: -### Test Categories +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer -| Test | What it checks | Marker | -| :--- | :--- | :--- | -| **AutoModel loading** | Every model loads via the relevant Transformers auto class with `trust_remote_code=True` and produces valid outputs | `gpu` | -| **Backend consistency** | SDPA, Flex, and Flash backends produce equivalent predictions (>= 95% agreement) | `gpu` | -| **Weight compliance** | FastPLM weights are bit-exact with the original implementations (ESM2, ESMC, ESM3, E1, DPLM) | `slow`, `gpu` | -| **Forward compliance** | Forward pass logits/predictions match the originals within tolerance | `slow`, `gpu` | -| **Rigorous parity** | Per-layer fp32 + bf16 hidden-state and last_hidden_state parity, padding-isolation, tokenizer parity, embed_dataset pipeline parity. Run per family in its own Docker image. | `gpu` | -| **NaN stability** | Batched inference with padding produces no NaN in real-token embeddings | `gpu` | -| **Batch-single match** | Batch and single-item embedding produce identical results | `gpu` | -| **Full model suite** | All of the above across every checkpoint (8M through 3B) | `gpu`, `large` | -| **Throughput benchmark** | Tokens/sec across models, backends, batch sizes, and sequence lengths | `slow`, `gpu` | -| **Structure models** | Boltz2, ESMFold, and ESMFold2 loading + forward/parity checks | `structure`, `slow`, `gpu` | +checkpoint = "Synthyra/DPLM-150M" +tokenizer = AutoTokenizer.from_pretrained(checkpoint) +generator = AutoModelForMaskedLM.from_pretrained( + checkpoint, + trust_remote_code=True, +).cuda().eval() -### Running Tests with Docker +input_ids = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() +with torch.inference_mode(): + generated_ids = generator.generate(input_ids, max_iter=100) -FastPLMs uses a per-family Docker setup. A single shared base image (`fastplms-base`) holds torch + transformers + the FastPLMs source, and one image per model family (`fastplms-esm2`, `fastplms-esm_plusplus`, `fastplms-esm3`, `fastplms-esmfold2`, `fastplms-e1`, `fastplms-dplm`, `fastplms-dplm2`, `fastplms-ankh`) layers on top with that family's native reference package. This isolates conflicting dependencies (e.g. Biohub `esm` vs `fair-esm`, DPLM's torchtext pin) and keeps each image small. +sequence = tokenizer.decode( + generated_ids[0], + skip_special_tokens=True, +).replace(" ", "") +print(sequence) +``` -```bash -# Initialize submodules (required before building Docker) -git submodule update --init --recursive +The official schedule uses 500 steps when `max_iter` is omitted. Reducing the +step count changes the sampling process. DPLM2 uses separate structure and +amino-acid tracks; see the [model guide](docs/models.md) for its explicit +boundary-token example. -# Build base + every family image -./build_images.sh +### ESM3 sequence generation -# Build a single family -./build_images.sh esm2 -./build_images.sh esm_plusplus -./build_images.sh esm3 -./build_images.sh esmfold2 -``` +ESM3 keeps multimodal helpers on the loaded model: -Run the parity / compliance tests for one family inside its image: +```python +from fastplms.models.esm3.modeling_esm3 import FastESM3GenerationConfig +from transformers import AutoModel -```bash -# ESM2 -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm2 \ - python -m pytest /workspace/testing/test_parity.py -k esm2 -v - -# ESM++ (model_key is "esmc") -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm_plusplus \ - python -m pytest /workspace/testing/test_parity.py -k esmc -v - -# ESM3, requires accepted access to biohub/esm3-sm-open-v1 for official parity -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm3 \ - python -m pytest /workspace/testing/test_parity.py -k esm3 -v - -# ESMFold2 / ESMFold2-Fast -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esmfold2 \ - python -m pytest /workspace/testing/test_esmfold2.py -v -s - -# E1, DPLM, DPLM2, ANKH -for fam in e1 dplm dplm2 ankh; do - docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-$fam \ - python -m pytest /workspace/testing/test_parity.py -k $fam -v -done +esm3 = AutoModel.from_pretrained( + "Synthyra/ESM3_small", + trust_remote_code=True, +).eval() +config = FastESM3GenerationConfig(num_steps=8, temperature=1.0, seed=7) +generated = esm3.generate("MK____A", config) +print(generated) ``` -The legacy monolithic `Dockerfile` (image tag `fastplms`) is still supported for the broader test suites that don't need native package isolation: +Underscores mark sequence positions to generate. Sequence-only forward passes +use `esm3.tokenize_sequences(...)`; structure and function tracks remain +available through the model's multimodal interface. -```bash -docker build -t fastplms . - -# Fast tests (small models, no compliance, no structure) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "gpu and not slow and not large and not structure" -v +### Boltz2 protein structure prediction -# All sequence model tests except 3B -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "not large and not structure" -v +Boltz2 exposes a protein-only convenience path in addition to its prepared +feature interface: -# Full suite including 3B models (requires 40+ GB VRAM) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "not structure" -v +```python +import torch +from transformers import AutoModel -# Structure models only (Boltz2, ESMFold, ESMFold2) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "structure" -v +boltz = AutoModel.from_pretrained( + "Synthyra/Boltz2", + trust_remote_code=True, + dtype=torch.float32, +).cuda().eval() +prediction = boltz.predict_structure( + amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG", + recycling_steps=3, + num_sampling_steps=50, + diffusion_samples=1, + seed=7, +) +boltz.save_as_cif(prediction, "prediction.cif") +print(prediction.plddt, prediction.ptm, prediction.iptm) ``` -On Windows, replace `$(pwd)` with `${PWD}`. **Always pass `--ipc=host`** with PyTorch. +Boltz2 remains provisional in FastPLMs 1.0. This interface is covered for +configuration, declared inference-core state, feature preparation, and seeded +execution, but it is not yet an official end-to-end numerical-equivalence +claim. +The helper restores caller RNG state; FP32 parameters and features execute under +the documented CUDA BF16 autocast policy. -### Compliance / Native Reference Dependencies +### ESMFold structure prediction -The parity and compliance tests compare FastPLM outputs against the original model implementations. Each per-family Docker image installs only the deps it needs; outside Docker you can install them piecewise: +ESMFold accepts raw sequences: -| Dependency | Used by | Install | -| :--- | :--- | :--- | -| `cloudpathlib`, `zstd`, `biotite` (+ `official/esm` submodule on `sys.path`) | ESM++ / ESMC, ESM3 official parity | provided by `Dockerfile.esm_plusplus` and `Dockerfile.esm3`; the Biohub `esm` package itself is **not** pip-installed because it depends on a Biohub `transformers` fork. | -| Biohub `transformers` fork, `rdkit`, `biotite`, `msgpack-numpy`, `pydssp`, `pygtrie`, `py3dmol` | ESMFold2 parity and structure export; runtime LM loading uses FastPLMs ESM++ | provided by `Dockerfile.esmfold2` | -| `E1` | E1 | `pip install -e official/e1` (or use `Dockerfile.e1`) | -| `transformers` (`EsmForMaskedLM`, `T5EncoderModel`) | ESM2, DPLM, ANKH | already in `requirements.txt` | - -If a native dep is missing in your environment, the corresponding parity tests are skipped rather than failing. +```python +import torch +from transformers import AutoModel -### Throughput Benchmarks +folder = AutoModel.from_pretrained( + "Synthyra/FastESMFold", + trust_remote_code=True, + dtype=torch.float32, +).cuda().eval() -Throughput can be measured via the pytest test (saves structured JSON/CSV/PNG results) or the standalone script (more configurable). +with torch.inference_mode(): + structure = folder.infer("MKTLLILAVVAAALA") -```bash -# Pytest (benchmarks ESM2-8M, ESMplusplus_small, DPLM-150M, DPLM2-150M across all backends) -docker run --gpus all -v $(pwd):/workspace fastplms python -m pytest /app/testing/test_throughput.py -v -s -# Output: throughput_results.json, throughput_results.csv, throughput_comparison.png - -# Standalone (fully configurable) -docker run --gpus all -v $(pwd):/workspace fastplms \ - python -m testing.throughput \ - --model_paths Synthyra/ESM2-8M Synthyra/ESMplusplus_small \ - --backends sdpa flex kernels_flash \ - --batch_sizes 2 4 8 \ - --sequence_lengths 64 128 256 512 1024 2048 \ - --output_path /workspace/throughput_comparison.png +print(structure["mean_plddt"]) ``` ---- +FastPLMs does not expose ProteinTTT for ESMFold. The pinned folding checkpoint +does not contain a trained masked-language-model head for that objective. -## Installation & Docker +### ESMFold2 folding and learned representations -### Local Installation +ESMFold2 accepts amino-acid sequences and typed molecular-complex +specifications. A target structure is not an input. Atomic coordinates and +confidence values are produced by the model. -FastPLMs is developed and tested with Python 3.12 and CUDA 12.8. For local GPU -installs, install the cu128 PyTorch wheels first, then the pinned direct -dependencies: +The two Fast checkpoints are inference-optimized for single-sequence use. They +have 24 folding blocks instead of 48 and were trained without MSA conditioning, +so they reject MSA-derived inputs. The full `ESMFold2` and +`ESMFold2-Experimental-Cutoff2025` checkpoints retain 48 blocks and optional +MSA conditioning. Fast is not necessarily single-chain-only: supported +multichain and multimolecule requests remain available, but every protein chain +uses single-sequence mode. This distinction follows the official model +description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). The +quick start below intentionally loads Fast and supplies no MSA. -```bash -git clone --recurse-submodules https://github.com/Synthyra/FastPLMs.git -cd FastPLMs -python -m pip install --upgrade pip==26.1.1 setuptools==70.2.0 -python -m pip install torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128 -python -m pip install -r requirements.txt -``` - -If you already cloned without `--recurse-submodules`, initialize submodules separately: -```bash -git submodule update --init --recursive -``` - -### Docker (Recommended for GPU Testing) - -There are two Docker layouts; pick whichever matches your task. +```python +from transformers import AutoModel -**Per-family layout (recommended for parity / compliance work).** A shared base image plus one image per model family, each with that family's native reference package isolated from the others. Build all of them once with the helper script: +folder = AutoModel.from_pretrained( + "Synthyra/ESMFold2-Fast", + trust_remote_code=True, + device_map={"": "cuda:0"}, + esmc_precision="auto", +).eval() -```bash -git submodule update --init --recursive -./build_images.sh # base + every family -./build_images.sh esm2 esm_plusplus # subset +result = folder.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = folder.result_to_pdb(result) +print(result.ptm, result.plddt.mean().item()) ``` -This produces `fastplms-base` and `fastplms-{esm2,esm_plusplus,esm3,e1,dplm,dplm2,ankh}`. Run a family's tests in its image: +Build complexes with the input types exposed by the loaded artifact: -```bash -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm2 \ - python -m pytest /workspace/testing/test_parity.py -k esm2 -v +```python +types = folder.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = folder.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +The typed interface supports RNA, ligands, modifications, covalent bonds, and +distogram conditioning. The Fast checkpoint loaded above rejects an MSA even +when it appears inside an otherwise valid typed request. To attach an MSA to a +protein input, load the full `Synthyra/ESMFold2` or +`Synthyra/ESMFold2-Experimental-Cutoff2025` checkpoint. The schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs rejects non-null pocket conditioning instead of +silently ignoring it. Prepared features contain fields such as `ref_pos`; these +are component reference geometries created during featurization, not a known +target structure. +See the offline +[`structure_preparation.py`](examples/structure_preparation.py) example for the +supported MSA, multimolecule, modification, bond, and distogram paths and the +explicit pocket rejection. Its ESMFold2 MSA branch requires one of the full +checkpoints, not a Fast checkpoint. + +Its learned sequence representation combines 81 ordered ESMC hidden states +with the folding checkpoint's projection. Use the public embedding API to +retrieve the resulting residue representation: -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace -it fastplms-esm2 bash +```python +representations = folder.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) ``` -**Monolithic layout (legacy, single image).** The original `Dockerfile` bundles every dependency that can coexist in one image. Convenient for the broad test suites and throughput benchmarks; not suitable when two families' native deps conflict (notably Biohub `esm` vs `fair-esm`). - -```bash -git submodule update --init --recursive -docker build -t fastplms . +For lower-level integrations that already hold the ordered ESMC hidden-state +stack, the projection is also exposed directly: -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -v -docker run --gpus all --ipc=host -v $(pwd):/workspace -it fastplms bash +```python +# H: (b, l, 81, 2560) +Z = folder.project_esmc_hidden_states(H) # Z: (b, l, 256) ``` -On Windows, replace `$(pwd)` with `${PWD}`. Always pass `--ipc=host` with PyTorch. +Here, `H` is the 81-state ESMC representation for a prepared sequence batch, +not a target structure. The dataset embedding API is the higher-level path for +ordinary sequence inputs. ---- +`folder.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` tensor +per single-chain sequence. The embedding path rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti`. -## Suggestions & Contributions -Found a bug or have a feature request? Please open a [GitHub Issue](https://github.com/Synthyra/FastPLMs/issues). We are actively looking for contributions to optimize more pLM architectures! +`esmc_precision="auto"` always resolves to BF16. Explicit FP8 is experimental, +inference-only, and strict: ---- +```python +folder.reload_esmc(precision="fp8", device="cuda:0") +print(folder.esmc_precision_status) +``` -## Citations +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Gradient-enabled paths reload canonical BF16 ESMC weights. -If you use FastPLMs, please cite the following along with the relevant model paper(s). +### Test-time training -### FastPLMs +Supported masked-language models expose opt-in, low-rank test-time adaptation: -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} +```python +metrics = generator.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={ + "steps": 3, + "ags": 1, + "batch_size": 1, + "seed": 7, + }, +) +generator.ttt_reset() ``` -### Flex Attention +TTT updates injected adapter parameters, not base checkpoint weights. It adds +latency and memory, can worsen a prediction, and does not establish biological +function. See the [TTT guide](docs/ttt.md) for supported families and folding +behavior. -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` +### Binder design research example -### PyTorch +The FastPLMs binder-design example optimizes a soft binder sequence against +ESMFold2 structural objectives and an ESM++ sequence prior: -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` +![FastPLMs EGFR minibinder design](docs/assets/egfr_fastplms_binder_design.png) -### ESM2 +Install the `binder` dependency profile, which includes the structure runtime +and the example-only table and antibody dependencies. The published workflow +requires Python 3.11-3.14, PyTorch 2.13, +Transformers 5.13, the verified ESMFold2 runtime assets, and a CUDA device. The +current release evidence target is the exact containerized Linux aarch64 +environment on the NVIDIA GH200 workstation. CPU-only, x86-64, Windows, macOS, +H100, and H200 binder runs do not substitute for that evidence. -```bibtex -@article{lin2023esm2, - title={Evolutionary-scale prediction of atomic-level protein structure with a language model}, - author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smestad, Nikita and Verkuil, Robert and Kabeli, Ori and Shmueli, Yaniv and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Salvatore and Rives, Alexander}, - journal={Science}, - volume={379}, - number={6637}, - pages={1123--1130}, - year={2023}, - DOI={10.1126/science.ade2574} -} +```bash +uv pip install \ + -r requirements/profiles/binder.in \ + -c requirements/constraints/validation.txt +PYTHONPATH=src python examples/binder_design_fastplms.py \ + --target-name pd-l1 \ + --binder-name minibinder \ + --batch-size 4 \ + --steps 150 \ + --output-dir artifacts/binder-design ``` -### ESM++ (ESMC) +The workflow writes optimization trajectories, ranked sequences, structures, +confidence outputs, and selection tables. These are model outputs for +prioritization, not experimental evidence of binding, specificity, expression, +or therapeutic activity. The output directory must not already exist. The +workflow publishes `run_manifest.json` atomically last, so a missing manifest +marks an incomplete run. See the [binder-design guide](docs/binder_design.md). -```bibtex -@misc{candido2026language, - title = {Language Modeling Materializes a World Model of Protein Biology}, - author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan - and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub - and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo - and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson - and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian - and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil - and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth - and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius - and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz - and Sandora, McCullen and Hla, Michael and Jia, Tina Z. - and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi - and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S. - and Sercu, Tom and Rives, Alexander}, - year = {2026}, - url = {https://biohub.ai/papers/esm_protein.pdf}, - note = {Preprint} -} -``` +### Fine-tuning -### E1 +FastPLMs follows `PreTrainedModel` conventions for Trainer, Accelerate, and +PEFT workflows. Install the `reporting` profile for training with plots: -```bibtex -@article{jain2025e1, - title={E1: Retrieval-Augmented Protein Encoder Models}, - author={Jain, Sarthak and Beazer, Joel and Ruffolo, Jeffrey A and Bhatnagar, Aadyot and Madani, Ali}, - journal={bioRxiv}, - DOI={10.1101/2025.11.12.688125}, - year={2025} -} -``` +```bash +uv pip install \ + -r requirements/profiles/reporting.in \ + -c requirements/constraints/validation.txt +PYTHONPATH=src python examples/fine_tuning.py \ + --task classification \ + --model_path Synthyra/ESM2-8M \ + --model-revision 185ecbd45665d050a8dae326d91886d330c5f9d0 \ + --classification-dataset-source GleghornLab/DL2_reg \ + --classification-dataset-revision 7e18f1b98859b0a3e3da283f63d0a153b774cf1f \ + --attn-backend sdpa \ + --output-dir artifacts/fine-tuning \ + --seed 7 \ + --full-determinism \ + --plot-results +``` + +For residue-level tasks, align labels to biological residues rather than +assuming tokenizer position equals residue position. Remote model and dataset +inputs require immutable 40-character Hub commits; existing local directory +inputs are identified by a full tree SHA-256 instead. Local datasets must use a +layout accepted by `datasets.load_dataset`; arbitrary `Dataset.save_to_disk()` +trees are not accepted. The example writes +`run_manifest.json` with ordered post-filter hashes of the rows and columns +actually consumed by training. It atomically publishes `final_model`, reloads +it against the same immutable base, and verifies the persisted adapter and +classifier tensor hashes. For the default plot-free run, install +`requirements/features/train.in` with `requirements/core.in` and omit +`--plot-results`. See the +[fine-tuning guide](docs/finetuning.md). + +## Attention backends + +Select a backend at load time or change it explicitly after loading: -### DPLM +```python +model.set_attn_implementation("sdpa") +``` + +| Backend | Use when | Main constraint | +| --- | --- | --- | +| `eager` | Attention matrices or the `parti` pooler are required | Materializes attention scores | +| `sdpa` | A stable default or official-parity path is required | Torch selects the underlying kernel | +| `flex_attention` | Padding-aware compiled attention is useful | First use compiles for the requested shape and semantics | +| `flash_attention_2` | A supported ESM2 or ESM++ BF16 CUDA path needs a precompiled kernel | BF16-only and family-limited | +| `flash_attention_3` | A supported ESM2, ESM++, or DPLM BF16 CUDA path needs a precompiled kernel | BF16-only and family-limited | + +FastPLMs does not implement an `auto` backend. An unavailable request raises. +When an optimized implementation cannot return attention matrices, +`output_attentions=True` emits one warning naming the configured backend, +effective eager backend, and reason, then runs a correctly masked eager call. +It does not change model configuration or later calls. Leaving +`attn_implementation` unspecified lets +Transformers choose its standard default. Explicit requests either run the +declared implementation or raise. + +The Flash dependency file installs Hugging Face `kernels`, not the +`flash-attn` source distribution. FastPLMs resolves immutable FlashAttention 2 +and 3 snapshots recorded +in `kernels.lock`, validates the compatible binary, and never compiles a source +fallback. See the [attention guide](docs/attention_backends.md) for exact +family, dtype, padding, and numerical contracts. + +For ESMC, SDPA is the recommended highest-fidelity path. Flex Attention and +FlashAttention 3 remain supported, non-experimental backends whose numerical +deviations are diagnostic rather than strict parity failures. The current +frozen-head release report is produced on the exact GH200/aarch64 validation +target and must publish relative L2, Q99.9, residue and pooled cosine, top-1, +and Jensen-Shannon distributions for each backend, dtype, exact hardware, and +sequence panel. H100 and H200 remain supported Hopper-class devices, but their +results are not interchangeable with or accepted as the current GH200 release +evidence. Pending measurements are +labeled pending in every ESMC card; no number is inferred from a threshold or +another checkpoint. + +## Design choices + +### Manifest-driven model support + +[`src/fastplms/models.toml`](src/fastplms/models.toml) is the source of truth +for model IDs, files, revisions, AutoClasses, tokenizer modes, transformations, +attention and precision capabilities, upstreams, licenses, and release tiers. +Support tables and model cards are generated from it. + +This avoids three common failure modes: a model appearing in documentation but +not release tooling, a checkpoint conversion with no immutable identity, and a +backend being advertised because it imports rather than because it was tested. + +### Native biological preparation + +The shared interface does not force every family through one tokenizer. E1 +keeps raw-sequence preparation. DPLM2 keeps modality-specific boundaries. +Structure models retain their native protein, nucleic-acid, ligand, and chain +representations. Full ESMFold2 checkpoints additionally retain optional MSA +conditioning; ESMFold2 Fast checkpoints reject MSA-derived inputs. Shared +pooling begins only after each model identifies its biological residue +positions. + +### Ordered embedding results + +Protein datasets routinely contain duplicate sequences and duplicate FASTA +identifiers. FastPLMs returns ordered records instead of a sequence-keyed +dictionary so those inputs are not silently lost. Persistent formats record +per-tensor hashes and complete run fingerprints for auditing and exact resume. + +### Explicit precision and backend policy + +Parameter storage, compute dtype, and attention implementation are separate +choices. DPLM and several structure paths retain FP32 parameters while using +CUDA BF16 autocast. ESMFold2 controls the ESMC backbone independently from its +folding trunk. Unsupported combinations fail before inference. + +### Official code is a parity oracle + +Pinned official repositories live under `vendor/upstream/` and are used only +in isolated reference stages. Production modules cannot import them, modify +`sys.path` to reach them, or download source code at import time. Artifacts are +self-contained and load through Transformers with `trust_remote_code=True`. + +### Fail-closed artifacts and licenses + +Artifact construction verifies required file identities, canonical legal +texts, conversion details, generated model-card metadata, and offline +loading. A missing or changed required file is a release error. Source licenses +and notices are centralized under [`LICENSES/`](LICENSES/); checkpoint-specific +terms remain distinct from the FastPLMs Apache-2.0 code license. + +## Validation and reproducibility + +All release validation is containerized. The portable runner accepts the host +and identity at invocation time: -```bibtex -@article{wang2024dplm, - title={Diffusion Language Models Are Versatile Protein Learners}, - author={Wang, Xinyou and Ye, Zaixiang and Huang, Fei and Cao, Dongyan and Liang, Shujian and Huang, Liang}, - journal={Proceedings of the 41st International Conference on Machine Learning}, - year={2024} -} +```bash +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/key \ + --suite check ``` -### DPLM2 +Release tiers cover checks, compliance, structure, features, artifacts, and +benchmarks. Missing required dependencies or declared backends fail rather than +skip. Expensive suites retain explicit `gpu`, `slow`, `large`, and `structure` +markers. -```bibtex -@article{wang2024dplm2, - title={DPLM-2: A Multimodal Diffusion Protein Language Model}, - author={Wang, Xinyou and Ye, Zaixiang and Huang, Fei and Cao, Dongyan and Liang, Shujian and Huang, Liang}, - journal={arXiv preprint arXiv:2410.13782}, - year={2024} -} -``` +Before merge, run the positive, fully offline `tests/cpu/` allowlist on the +validation workstation with Python 3.12, CPU-only Torch 2.13, and Transformers +5.13. It hides CUDA, blocks socket and Hub downloads, rejects skips and xfails, +and targets less than five minutes on four CPU cores. Live official references +remain reserved for the release-candidate `compliance` tier; routine checks +consume immutable goldens. This repository does not use GitHub Actions. -### ANKH +The canonical Docker workflow uses one candidate image and isolated official +reference images: -```bibtex -@article{elnaggar2023ankh, - title={Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling}, - author={Elnaggar, Ahmed and Essam, Hazem and Salah-Eldin, Wafaa and Moustafa, Walid and Elkerdawy, Mohamed and Rochereau, Charlotte and Rost, Burkhard}, - journal={arXiv preprint arXiv:2301.06568}, - year={2023} -} +```bash +git submodule update --init --recursive +sudo docker buildx bake -f docker/docker-bake.hcl candidate reference-esm2 --load +sudo docker compose -f docker/compose.yaml run --rm candidate \ + python -m pytest tests/parity -k esm2 -v ``` -```bibtex -@article{alsamkary2025ankh3, - title={Ankh3: Multi-Task Pretraining with Sequence Denoising and Completion Enhances Protein Representations}, - author={Alsamkary, Hazem and Elshaffei, Mohamed and Elkerdawy, Mohamed and Elnaggar, Ahmed}, - journal={arXiv preprint arXiv:2505.20052}, - year={2025} -} -``` +Always pass `--ipc=host` to Dockerized PyTorch runs. Benchmarks are separate +from correctness checks and retain raw samples, environment metadata, warm-up, +compile-time, steady-state, padding, and memory measurements. -### Boltz +Boltz2 remains provisional in FastPLMs 1.0. Configuration, declared +inference-core state, feature preparation, and seeded execution are tested, but +native-environment BF16 end-to-end inference does not yet meet the fixed +numerical-equivalence limits. FastPLMs therefore makes no official inference +equivalence claim for Boltz2. -```bibtex -@article{passaro2025boltz2, - title={Boltz-2: Exploring the Frontiers of Biomolecular Prediction}, - author={Passaro, Saro and Corso, Gabriele and Wohlwend, Jeremy and Reveiz, Mateo and Bordes, Florian and Wicky, Basile and Dayan, Peter and Jing, Bowen}, - journal={bioRxiv}, - year={2025} -} -``` +## Files-only Hub publication -```bibtex -@article{wohlwend2024boltz1, - title={Boltz-1: Democratizing Biomolecular Interaction Modeling}, - author={Wohlwend, Jeremy and Corso, Gabriele and Passaro, Saro and Reveiz, Mateo and Leidal, Ken and Swanson, Wojtek and Kher, Gilmer and Lember, Tommi and Jaakkola, Tommi}, - journal={bioRxiv}, - year={2024} -} -``` +After building and validating local artifacts, update Hub runtime files and +model cards without uploading or deleting checkpoint weights: -### ESMFold / ProteinTTT +```bash +PYTHONPATH=src python -m tools.artifacts.publish \ + --files-only \ + --artifact-root dist/hub \ + --dry-run \ + esm2_8m +``` + +`--artifact-root` is the local directory containing one built artifact +subdirectory per selected model. It tells the publisher which validated files +to upload; it is not a remote Hub path and does not change where models are +published. + +Remove `--dry-run` after reviewing the exact add-only file plan. Repository +targets come exclusively from `models.toml`. Files-only publication rejects any +ANKH selection, including the implicit all-model selection, so callers must pass +explicit non-ANKH model IDs. Authentication uses `HF_TOKEN` or the cached +Hugging Face login. See [Local Hub artifacts](docs/artifacts.md) for the full +safety contract. + +This files-only workflow is forbidden for the ANKH 1.0 migration. ANKH must +replace its encoder-only contents with the full encoder-decoder state in one +immutable commit containing every weight shard, tokenizer asset, configuration, +runtime source, card, and release record. Validate both `AutoModel` and +`AutoModelForSeq2SeqLM` from that same commit before publication is accepted. +Use the explicit `--complete ` dry-run and publication workflow. +It makes one parent-guarded atomic commit containing validated additions and +only the narrowly scoped deletions needed to replace obsolete registry-pinned +files, such as a monolithic ANKH weight file superseded by indexed shards. +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The maintained ByteDance +[license](https://github.com/bytedance/dplm/blob/main/LICENSE) +is Apache-2.0 and the [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +defines the repository release as including pretrained DPLM1 and DPLM2 weights. +Validated DPLM artifacts therefore record `weights_license_status="resolved"` +and `redistributable=true`; explicit `--complete` publication is permitted after +the same legal, parity, inventory, parent-commit, and atomic preflight checks. + +## Documentation + +- [Documentation index](docs/README.md) +- [Architecture](docs/architecture.md) +- [Models and generated support](docs/models.md) +- [Capability-to-evidence manifest](docs/generated/capability_evidence.md) +- [Embedding API](docs/embedding_api.md) +- [Attention backends](docs/attention_backends.md) +- [ESMFold2](docs/esmfold2.md) +- [Test-time training](docs/ttt.md) +- [Binder design](docs/binder_design.md) +- [Fine-tuning](docs/finetuning.md) +- [Testing and compliance](docs/testing.md) +- [Benchmarking](docs/benchmarking.md) +- [Local Hub artifacts](docs/artifacts.md) +- [Migration to 1.0](docs/migration.md) +- [Licensing](docs/licensing.md) +- [Contributing](docs/contributing.md) +- [Runnable examples](examples/README.md) + +FastPLMs 1.0 is an intentional API break. There are no legacy import, backend, +embedding-storage, or command shims. Use the migration guide when moving from +the pre-1.0 repository layout. + +## Contributing and citation + +Start with [AGENTS.md](AGENTS.md) for repository invariants and +[the contributing guide](docs/contributing.md) for the model-addition and +validation workflow. Generated support tables and model cards must be changed +through the manifest or renderer rather than edited directly. + +If FastPLMs supports your work, cite FastPLMs and the paper or model card for +the specific checkpoint family: ```bibtex -@misc{bushuiev2026proteinneed, - title={One protein is all you need}, - author={Anton Bushuiev and Roman Bushuiev and Olga Pimenova and Nikola Zadorozhny and Raman Samusevich and Elisabet Manaskova and Rachel Seongeun Kim and Hannes St\"ark and Jiri Sedlar and Martin Steinegger and Tom\'a\v{s} Pluskal and Josef Sivic}, - year={2026}, - eprint={2411.02109}, - archivePrefix={arXiv}, - primaryClass={cs.LG}, - url={https://arxiv.org/abs/2411.02109} +@misc{FastPLMs, + author = {Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, + title = {FastPLMs: Fast, efficient protein language model inference from Hugging Face AutoModel}, + year = {2024}, + url = {https://github.com/Synthyra/FastPLMs}, + doi = {10.57967/hf/3726}, + publisher = {Hugging Face} } ``` ---- diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..23006ad --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,99 @@ +# Third-party notices + +FastPLMs implements interfaces and checkpoint mappings for independently +released protein models. The pinned repositories under `vendor/upstream/` are +parity oracles. Production code does not import them, and runtime images do not +contain them. + +This notice is informational and is not legal advice. A checkpoint license can +differ from the license covering its source implementation. The typed inventory +in `src/fastplms/models.toml` and the verbatim files under `LICENSES/` are the +distribution record. + +## ANKH + +The pinned ANKH implementation and the mirrored ANKH checkpoints are identified +as CC BY-NC-SA 4.0. FastPLMs displays those terms but does not enforce them in +software. Users are responsible for determining whether their use and +redistribution comply. The complete text is in `LICENSES/ankh/LICENSE.md`. + +## Profluent-E1 + +Profluent identifies its E1 model code as Apache-2.0. The E1 weights and full +release are subject to the Profluent-E1 Clickthrough License Agreement and the +incorporated attribution requirements. Any E1 distribution must retain all of +the following files: + +- `LICENSES/e1/LICENSE`, the Profluent-E1 agreement +- `LICENSES/e1/ATTRIBUTION`, the attribution guidelines +- `LICENSES/e1/NOTICE`, the required notice +- `LICENSES/e1/Apache-2.0.txt`, the code license +- `LICENSES/e1/BSD-3-Clause.txt`, covering the FlashAttention-derived padding + utility identified by the official E1 source +- `LICENSES/e1/MODIFICATIONS.md`, the FastPLMs modified-file notice + +The exact text `Profluent-E1` must remain prominently displayed in E1 +documentation and at each launch of an executable E1 workflow, as required by +the upstream attribution guidelines. Certain commercial outputs, including +specified pharmaceutical and target-related outputs, can require the separate +`Built with Profluent-E1` statement described in `ATTRIBUTION`. + +## DPLM + +The pinned ByteDance DPLM repository is Apache-2.0. Its +[README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview) +explicitly defines the repository release as including pretrained DPLM1 and +DPLM2 weights, and the same revision carries the complete +[Apache-2.0 license](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE). +FastPLMs records both checkpoint families as Apache-2.0 and distributes the +verbatim license plus `LICENSES/dplm/PROVENANCE.md`. Converted weights retain +those terms and remain subject to the ordinary artifact and publication gates. + +## Biohub + +The pinned Biohub ESM implementation is MIT and includes a separate +`THIRD_PARTY_NOTICE.md`; both files are distributed under +`LICENSES/biohub-esm/`. The pinned Biohub Transformers fork is Apache-2.0, with +its complete text under `LICENSES/biohub-transformers/`. + +## Boltz + +The pinned Boltz source is MIT. The verbatim notice is in +`LICENSES/boltz/LICENSE`. + +## Meta ESM and OpenFold + +The pinned Meta ESM source is MIT. The pinned OpenFold source is Apache-2.0. +Their verbatim texts and revision-specific provenance notices are under +`LICENSES/fair-esm/` and `LICENSES/openfold/`. + +The native H100 ESMFold reference image applies the tracked +`docker/constraints/openfold-sm90.patch` to the copied OpenFold `setup.py`. +This build-only change restricts the CUDA extension to `sm90` and selects the +C++17 standard required by the reference PyTorch version. It leaves the pinned +submodule, extension source, model classes, checkpoint data, and public API +unchanged. The complete modified-file record is in +`LICENSES/openfold/MODIFICATIONS.md`. + +The isolated reference image also includes Apache-2.0 PyTorch Lightning, +TorchMetrics, Lightning Utilities, and NVIDIA DLLogger. Their exact versions or +revision are pinned in `docker/constraints/esmfold.txt`; OpenFold imports them +eagerly, and FastPLMs production code does not depend on them. DLLogger's exact +source identity and installed-license handling are recorded in +`LICENSES/dllogger/PROVENANCE.md`. + +## ProteinTTT + +The optional test-time training workflow is validated against the pinned +ProteinTTT repository under its MIT license. Its verbatim license and +revision-specific provenance are under `LICENSES/protein-ttt/`. + +## Conversion and packaging record + +For every supported family, `src/fastplms/models.toml` records an immutable +official checkpoint revision, an immutable FastPLMs checkpoint revision, file +digests, a named state transformation, and a mechanism-level conversion record. +Generated artifacts reproduce that record in `provenance.json`. A release or +artifact build must fail when a required file identity, legal text, attribution +notice, modified-file notice, upstream revision, or conversion record is absent +or differs from its manifest digest. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..f48c39f --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,84 @@ +# Hopper/SM90 benchmarks + +Benchmarks run independently from pytest. The steady-state path times only a +forward pass over pre-tokenized tensors already resident on the GPU. Startup, +compilation, end-to-end embedding, and the ESMFold2 projection are distinct +modes. + +Run the complete manifest-derived release matrix on the current NVIDIA GH200 +validation workstation in its exact containerized Linux aarch64 environment: + +```bash +python -m benchmarks.suite \ + --backends eager sdpa flex_attention \ + --junit-output artifacts/junit/benchmark.xml \ + --output artifacts/benchmarks/h100.json +``` + +The output name is a legacy automation identifier. Every report carries the +actual accelerator, architecture, and software fingerprint, and regression +comparisons require an exact match. H100 and H200 remain supported Hopper-class +devices, but are not interchangeable with or accepted as the current +GH200/aarch64 release evidence. +Remote orchestration binds Bake to native `linux/arm64` on the GH200 and verifies +the loaded image architecture and content digest rather than relying on +emulated `linux/amd64` images. + +For the pre-publication baseline, build and consume the manifest-selected local +Hub artifacts in the same frozen source job: + +```bash +python -m tools.artifacts.build_all \ + --benchmark-suite --source-root . --output-root dist/hub +python -m benchmarks.suite \ + --artifact-root dist/hub --local-files-only \ + --backends eager sdpa flex_attention \ + --junit-output artifacts/junit/benchmark-capture.xml \ + --output artifacts/benchmarks/h100-baseline-candidate.json +``` + +Artifact mode validates the complete selected artifacts and ESMFold2's local +ESMC dependency before loading. Reports retain registry repository/revision +case keys and path-free weights, runtime, source, canonical-state, and manifest +identities. Local filesystem paths are never baseline identities. +The GH200 release runner records FA2 as prior focused evidence and FA3 as +unavailable on linux/arm64; it never downloads, builds, or executes either +Flash kernel. Capture reports include the exact environment/artifact identities +needed for mechanical baseline promotion and keep cold compile time separate +from warm throughput blocks. + +The matrix includes startup, compilation, full embedding, `b=1, l=512` +latency, `b=8, l=1024` throughput, the fixed skewed-padding case, and BF16/FP8 +ESMFold2 projection measurements. Use `--baseline ` to apply the +one-sided 95% regression gates without modifying the baseline. + +```bash +python -m benchmarks \ + --model dist/hub/ESM2-8M \ + --backend sdpa \ + --mode steady \ + --batch-size 1 \ + --sequence-length 512 \ + --local-files-only \ + --output artifacts/benchmarks/esm2-sdpa.json +``` + +Build the manifest-pinned local artifact first. If a published Hub artifact is +used instead, pass its immutable revision and retain it in the report. + +Use `--lengths 1024 512 256 128 64 64 32 32 --batch-size 8` for the +padding-efficiency case. Every report includes raw CUDA-event samples, logical +and padded token throughput, median and P95 latency, peak allocated and +reserved memory, startup time, before/after temperatures and clocks, and the +complete accelerator/software fingerprint. + +Compare an immutable baseline with a new run using: + +```bash +python -m benchmarks.regression current.json baseline.json \ + --output artifacts/benchmarks/gate.json +``` + +The command exits nonzero when a regression is established. It never rewrites +the baseline. A speed claim is supported only when the one-sided lower +confidence bound shows at least a five-percent improvement. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e29a3a1 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1,6 @@ +"""Reproducible, exact-device Hopper/SM90 benchmarks for FastPLMs.""" + +from .regression import GateResult, GateThresholds, compare_reports + + +__all__ = ["GateResult", "GateThresholds", "compare_reports"] diff --git a/benchmarks/__main__.py b/benchmarks/__main__.py new file mode 100644 index 0000000..bdf9375 --- /dev/null +++ b/benchmarks/__main__.py @@ -0,0 +1,7 @@ +"""Run the benchmark command-line interface.""" + +from .run import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/regression.py b/benchmarks/regression.py new file mode 100644 index 0000000..28ebfce --- /dev/null +++ b/benchmarks/regression.py @@ -0,0 +1,472 @@ +"""Statistical regression gates for benchmark reports. + +The gate compares paired measurement blocks from the same benchmark case. It +keeps raw measurements in the report so a baseline change is always an +explicit, reviewable file update. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import statistics +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class GateThresholds: + """Scalar thresholds used by the exact-device Hopper/SM90 regression gate.""" + + confidence: float = 0.95 + soft_throughput_ratio: float = 0.95 + hard_throughput_ratio: float = 0.90 + claimed_improvement_ratio: float = 1.05 + memory_growth_fraction: float = 0.05 + memory_growth_bytes: int = 256 * 1024**2 + hard_memory_growth_fraction: float = 0.10 + bootstrap_samples: int = 10_000 + seed: int = 42 + + +DEFAULT_GATE_THRESHOLDS = GateThresholds() + + +@dataclass(frozen=True) +class CaseGateResult: + """Result for one matched model, backend, mode, and input shape.""" + + case: str + passed: bool + median_ratio: float + lower_confidence_bound: float + upper_confidence_bound: float + memory_growth_bytes: int + memory_growth_fraction: float + improvement_supported: bool + reasons: tuple[str, ...] + + +@dataclass(frozen=True) +class GateResult: + """Aggregate benchmark comparison.""" + + passed: bool + cases: tuple[CaseGateResult, ...] + unmatched_current: tuple[str, ...] + unmatched_baseline: tuple[str, ...] + environment_mismatches: tuple[str, ...] + artifact_mismatches: tuple[str, ...] + report_mismatches: tuple[str, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "passed": self.passed, + "cases": [asdict(case) for case in self.cases], + "unmatched_current": list(self.unmatched_current), + "unmatched_baseline": list(self.unmatched_baseline), + "environment_mismatches": list(self.environment_mismatches), + "artifact_mismatches": list(self.artifact_mismatches), + "report_mismatches": list(self.report_mismatches), + } + + +ENVIRONMENT_FIELDS = ( + "python", + "platform", + "machine", + "torch", + "cuda_runtime", + "cudnn", + "transformers", + "fastplms", + "transformer_engine", + "kernels", + "kernels_data", + "gpu", + "gpu_capability", +) +NVIDIA_SMI_IDENTITY_FIELDS = ("name", "driver_version", "memory.total") +REPORT_SCHEMA_VERSION = 3 +REPORT_IDENTITY_FIELDS = ( + "matrix_kind", + "claim_scope", + "backend_policy", + "timing_contract", + "baseline_promotion_contract", +) + + +def percentile(values: Sequence[float], probability: float) -> float: + """Return a linearly interpolated percentile without NumPy.""" + + if not values: + raise ValueError("At least one value is required") + if not 0.0 <= probability <= 1.0: + raise ValueError("probability must be between zero and one") + ordered = sorted(values) + position = probability * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def bootstrap_ratio_interval( + current: Sequence[float], + baseline: Sequence[float], + *, + confidence: float = 0.95, + samples: int = 10_000, + seed: int = 42, +) -> tuple[float, float, float]: + """Estimate a paired bootstrap interval for scalar throughput ratio ``r``.""" + + if len(current) != len(baseline): + raise ValueError("Current and baseline reports must have equally many blocks") + if not current: + raise ValueError("At least one paired block is required") + if any(not math.isfinite(value) or value <= 0.0 for value in baseline): + raise ValueError("Baseline throughput must be finite and positive") + if any(not math.isfinite(value) or value <= 0.0 for value in current): + raise ValueError("Current throughput must be finite and positive") + if not 0.0 < confidence < 1.0: + raise ValueError("confidence must be strictly between zero and one") + if samples < 1: + raise ValueError("samples must be positive") + + paired_ratios = [new / old for new, old in zip(current, baseline, strict=True)] + median_ratio = statistics.median(paired_ratios) + generator = random.Random(seed) + n = len(paired_ratios) + bootstrapped = [ + statistics.median(paired_ratios[generator.randrange(n)] for _ in range(n)) + for _ in range(samples) + ] + tail = 1.0 - confidence + # The release policy uses one-sided bounds: the upper bound detects a + # regression and the lower bound supports a speed claim. + return median_ratio, percentile(bootstrapped, tail), percentile(bootstrapped, confidence) + + +def _case_key(record: Mapping[str, Any]) -> str: + fields = ( + "model", + "revision", + "auto_class", + "backend", + "precision", + "mode", + "batch_size", + "sequence_length", + "lengths", + ) + return "|".join(f"{field}={record.get(field)}" for field in fields) + + +def _throughputs(record: Mapping[str, Any]) -> list[float]: + blocks = record.get("blocks") + if not isinstance(blocks, list) or not blocks: + raise ValueError(f"Benchmark case {_case_key(record)} has no measurement blocks") + if len(blocks) != 7: + raise ValueError( + f"Benchmark case {_case_key(record)} must contain exactly seven blocks" + ) + values: list[float] = [] + for block in blocks: + if not isinstance(block, Mapping) or "logical_tokens_per_second" not in block: + raise ValueError(f"Benchmark case {_case_key(record)} has a malformed block") + value = float(block["logical_tokens_per_second"]) + if not math.isfinite(value) or value <= 0.0: + raise ValueError( + f"Benchmark case {_case_key(record)} has non-positive/non-finite throughput" + ) + values.append(value) + return values + + +def _throughput_records(report: Mapping[str, Any]) -> dict[str, Mapping[str, Any]]: + """Index only cases containing steady-state measurement blocks. + + Startup, first-forward, compilation-only, and full-embedding records remain + descriptive. They are retained in reports but are not throughput ratios. + """ + + raw_results = report.get("results") + if not isinstance(raw_results, list): + raise ValueError("Benchmark report results must be a list") + result: dict[str, Mapping[str, Any]] = {} + for record in raw_results: + if not isinstance(record, Mapping): + raise ValueError("Benchmark report contains a non-object result") + blocks = record.get("blocks") + if isinstance(blocks, list) and blocks: + key = _case_key(record) + if key in result: + raise ValueError(f"Benchmark report contains duplicate case: {key}") + result[key] = record + return result + + +def _peak_memory(record: Mapping[str, Any]) -> int: + memory = record.get("memory") + if not isinstance(memory, Mapping): + raise ValueError(f"Benchmark case {_case_key(record)} has no memory record") + value = memory.get("peak_allocated_bytes") + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"Benchmark case {_case_key(record)} has invalid peak memory") + return value + + +def _report_mismatches( + current: Mapping[str, Any], baseline: Mapping[str, Any] +) -> tuple[str, ...]: + """Require two complete, promotion-ready reports with the same contract.""" + + mismatches: list[str] = [] + for label, report in (("current", current), ("baseline", baseline)): + if report.get("schema_version") != REPORT_SCHEMA_VERSION: + mismatches.append( + f"{label} report schema_version must be {REPORT_SCHEMA_VERSION}" + ) + if report.get("status") != "complete": + mismatches.append(f"{label} report status is not complete") + results = report.get("results") + expected = report.get("expected_case_count") + completed = report.get("completed_case_count") + if ( + not isinstance(results, list) + or isinstance(expected, bool) + or not isinstance(expected, int) + or expected <= 0 + or isinstance(completed, bool) + or not isinstance(completed, int) + or completed != expected + or len(results) != expected + ): + mismatches.append(f"{label} report case inventory is incomplete") + for field in REPORT_IDENTITY_FIELDS: + if field not in current: + mismatches.append(f"current report is missing {field}") + continue + if field not in baseline: + mismatches.append(f"baseline report is missing {field}") + continue + if current[field] != baseline[field]: + mismatches.append( + f"{field}: current={current[field]!r}, baseline={baseline[field]!r}" + ) + return tuple(mismatches) + + +def _environment_mismatches( + current: Mapping[str, Any], baseline: Mapping[str, Any] +) -> tuple[str, ...]: + """Return release-environment differences that invalidate a comparison.""" + + current_environment = current.get("environment") + baseline_environment = baseline.get("environment") + if not isinstance(current_environment, Mapping): + return ("current report has no environment fingerprint",) + if not isinstance(baseline_environment, Mapping): + return ("baseline report has no environment fingerprint",) + + mismatches: list[str] = [] + for field in ENVIRONMENT_FIELDS: + if field not in current_environment: + mismatches.append(f"environment.{field}: current report is missing the field") + continue + if field not in baseline_environment: + mismatches.append(f"environment.{field}: baseline report is missing the field") + continue + current_value = current_environment.get(field) + baseline_value = baseline_environment.get(field) + if current_value != baseline_value: + mismatches.append( + f"environment.{field}: current={current_value!r}, baseline={baseline_value!r}" + ) + + current_smi = current_environment.get("nvidia_smi") + baseline_smi = baseline_environment.get("nvidia_smi") + if not isinstance(current_smi, Mapping): + mismatches.append("environment.nvidia_smi: current report has no mapping") + if not isinstance(baseline_smi, Mapping): + mismatches.append("environment.nvidia_smi: baseline report has no mapping") + if isinstance(current_smi, Mapping) and isinstance(baseline_smi, Mapping): + for field in NVIDIA_SMI_IDENTITY_FIELDS: + if field not in current_smi: + mismatches.append( + f"environment.nvidia_smi.{field}: current report is missing the field" + ) + continue + if field not in baseline_smi: + mismatches.append( + f"environment.nvidia_smi.{field}: baseline report is missing the field" + ) + continue + current_value = current_smi.get(field) + baseline_value = baseline_smi.get(field) + if current_value != baseline_value: + mismatches.append( + f"environment.nvidia_smi.{field}: current={current_value!r}, " + f"baseline={baseline_value!r}" + ) + return tuple(mismatches) + + +def _artifact_mismatches( + current: Mapping[str, Any], baseline: Mapping[str, Any] +) -> tuple[str, ...]: + """Reject comparisons across missing or different local artifact inventories.""" + + current_has_inventory = "artifacts" in current or "artifact_load_mode" in current + baseline_has_inventory = "artifacts" in baseline or "artifact_load_mode" in baseline + if not current_has_inventory and not baseline_has_inventory: + return () + + mismatches: list[str] = [] + current_mode = current.get("artifact_load_mode") + baseline_mode = baseline.get("artifact_load_mode") + if current_mode != baseline_mode: + mismatches.append( + f"artifact_load_mode: current={current_mode!r}, baseline={baseline_mode!r}" + ) + + current_artifacts = current.get("artifacts") + baseline_artifacts = baseline.get("artifacts") + if not isinstance(current_artifacts, Mapping): + mismatches.append("current report has no artifact inventory mapping") + if not isinstance(baseline_artifacts, Mapping): + mismatches.append("baseline report has no artifact inventory mapping") + if isinstance(current_artifacts, Mapping) and isinstance(baseline_artifacts, Mapping): + for model_id in sorted(set(current_artifacts) | set(baseline_artifacts), key=str): + current_identity = current_artifacts.get(model_id) + baseline_identity = baseline_artifacts.get(model_id) + if current_identity != baseline_identity: + mismatches.append( + f"artifacts.{model_id}: current={current_identity!r}, " + f"baseline={baseline_identity!r}" + ) + return tuple(mismatches) + + +def compare_reports( + current: Mapping[str, Any], + baseline: Mapping[str, Any], + thresholds: GateThresholds = DEFAULT_GATE_THRESHOLDS, +) -> GateResult: + """Compare two reports and return a deterministic regression decision.""" + + current_cases = _throughput_records(current) + baseline_cases = _throughput_records(baseline) + environment_mismatches = _environment_mismatches(current, baseline) + artifact_mismatches = _artifact_mismatches(current, baseline) + report_mismatches = _report_mismatches(current, baseline) + shared = sorted(current_cases.keys() & baseline_cases.keys()) + case_results: list[CaseGateResult] = [] + + for index, key in enumerate(shared): + new = current_cases[key] + old = baseline_cases[key] + median_ratio, lower, upper = bootstrap_ratio_interval( + _throughputs(new), + _throughputs(old), + confidence=thresholds.confidence, + samples=thresholds.bootstrap_samples, + seed=thresholds.seed + index, + ) + baseline_memory = _peak_memory(old) + current_memory = _peak_memory(new) + memory_delta = current_memory - baseline_memory + memory_fraction = memory_delta / baseline_memory if baseline_memory else 0.0 + soft_memory_limit = max( + thresholds.memory_growth_bytes, + int(baseline_memory * thresholds.memory_growth_fraction), + ) + reasons: list[str] = [] + if upper < thresholds.soft_throughput_ratio: + reasons.append( + f"upper confidence bound {upper:.4f} is below " + f"{thresholds.soft_throughput_ratio:.4f}" + ) + if median_ratio < thresholds.hard_throughput_ratio: + reasons.append( + f"median throughput ratio {median_ratio:.4f} is below hard limit " + f"{thresholds.hard_throughput_ratio:.4f}" + ) + if memory_delta > soft_memory_limit: + reasons.append( + f"peak allocated memory grew by {memory_delta} bytes; limit is {soft_memory_limit}" + ) + if baseline_memory and memory_fraction > thresholds.hard_memory_growth_fraction: + reasons.append( + f"peak allocated memory grew by {memory_fraction:.2%}; hard limit is " + f"{thresholds.hard_memory_growth_fraction:.2%}" + ) + case_results.append( + CaseGateResult( + case=key, + passed=not reasons, + median_ratio=median_ratio, + lower_confidence_bound=lower, + upper_confidence_bound=upper, + memory_growth_bytes=memory_delta, + memory_growth_fraction=memory_fraction, + improvement_supported=lower >= thresholds.claimed_improvement_ratio, + reasons=tuple(reasons), + ) + ) + + unmatched_current = tuple(sorted(current_cases.keys() - baseline_cases.keys())) + unmatched_baseline = tuple(sorted(baseline_cases.keys() - current_cases.keys())) + passed = ( + bool(shared) + and not unmatched_current + and not unmatched_baseline + and not environment_mismatches + and not artifact_mismatches + and not report_mismatches + and all(case.passed for case in case_results) + ) + return GateResult( + passed=passed, + cases=tuple(case_results), + unmatched_current=unmatched_current, + unmatched_baseline=unmatched_baseline, + environment_mismatches=environment_mismatches, + artifact_mismatches=artifact_mismatches, + report_mismatches=report_mismatches, + ) + + +def _load(path: Path) -> Mapping[str, Any]: + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object in {path}") + return value + + +def main(argv: Iterable[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("current", type=Path) + parser.add_argument("baseline", type=Path) + parser.add_argument("--output", type=Path) + arguments = parser.parse_args(argv) + result = compare_reports(_load(arguments.current), _load(arguments.baseline)) + rendered = json.dumps(result.to_dict(), indent=2, sort_keys=True) + if arguments.output: + arguments.output.parent.mkdir(parents=True, exist_ok=True) + arguments.output.write_text(rendered + "\n", encoding="utf-8") + print(rendered) + return 0 if result.passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/run.py b/benchmarks/run.py new file mode 100644 index 0000000..75ec9f6 --- /dev/null +++ b/benchmarks/run.py @@ -0,0 +1,862 @@ +"""CUDA-event benchmark runner for FastPLMs. + +The primary steady-state path receives pre-tokenized tensors already resident +on the GPU. Startup, compilation, end-to-end embedding, the ESMFold2 learned +projection, and ESMC inference plus projection are measured separately so those +costs cannot be hidden in a single throughput number. +""" + +from __future__ import annotations + +import argparse +import contextlib +import inspect +import json +import platform +import random +import re +import statistics +import subprocess +import sys +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + + +CANONICAL_AAS = "ACDEFGHIKLMNPQRSTVWY" +HOPPER_SM90_CAPABILITY = (9, 0) +_HOPPER_PRODUCT_PATTERN = re.compile(r"(? str | None: + try: + return version(distribution_name) + except PackageNotFoundError: + return None + + +def _nvidia_smi() -> Mapping[str, str]: + fields = ( + "name", + "driver_version", + "temperature.gpu", + "clocks.sm", + "clocks.mem", + "memory.total", + ) + try: + completed = subprocess.run( + [ + "nvidia-smi", + f"--query-gpu={','.join(fields)}", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return {} + values = [value.strip() for value in completed.stdout.splitlines()[0].split(",")] + return dict(zip(fields, values, strict=False)) + + +def environment_fingerprint(torch: Any) -> dict[str, Any]: + """Return the software and accelerator metadata needed to interpret results.""" + + return { + "python": sys.version.split()[0], + "platform": platform.platform(), + "machine": platform.machine(), + "torch": str(torch.__version__), + "cuda_runtime": str(torch.version.cuda), + "cudnn": int(torch.backends.cudnn.version() or 0), + "transformers": _version("transformers"), + "fastplms": _version("fastplms"), + "transformer_engine": _version("transformer-engine"), + "kernels": _version("kernels"), + "kernels_data": _version("kernels-data"), + "gpu": torch.cuda.get_device_name(), + "gpu_capability": list(torch.cuda.get_device_capability()), + "nvidia_smi": _nvidia_smi(), + } + + +def validate_hopper_sm90_environment(environment: Mapping[str, Any]) -> None: + """Require an allowed Hopper product for a release-claim benchmark matrix.""" + + gpu = environment.get("gpu") + if not isinstance(gpu, str) or _HOPPER_PRODUCT_PATTERN.search(gpu.upper()) is None: + raise RuntimeError( + f"Release-claim benchmarks require an NVIDIA H100, H200, or GH200 GPU; got {gpu!r}." + ) + capability = environment.get("gpu_capability") + if capability != list(HOPPER_SM90_CAPABILITY): + raise RuntimeError( + f"Release-claim benchmarks require compute capability 9.0; got {capability!r}." + ) + + +def _sequence(length: int, seed: int) -> str: + generator = random.Random(seed) + return "".join(generator.choice(CANONICAL_AAS) for _ in range(length)) + + +def sequences_for_lengths(lengths: Sequence[int], *, special_tokens: int = 2) -> list[str]: + """Create deterministic proteins whose tokenized lengths approach ``lengths``.""" + + return [ + _sequence(max(1, length - special_tokens), 42 + index) + for index, length in enumerate(lengths) + ] + + +def _model_forward(model: Any, model_inputs: Mapping[str, Any]) -> Any: + arguments = dict(model_inputs) + parameters = inspect.signature(model.forward).parameters + accepts_kwargs = any( + parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ) + optional = { + "output_attentions": False, + "output_hidden_states": False, + "return_dict": True, + } + for name, value in optional.items(): + if accepts_kwargs or name in parameters: + arguments[name] = value + return model(**arguments) + + +def _uses_bf16_autocast(arguments: argparse.Namespace) -> bool: + return arguments.precision == "bf16" and arguments.bf16_execution == "fp32_parameters_autocast" + + +def _resolve_bf16_execution(arguments: argparse.Namespace) -> str: + """Resolve a registered checkpoint policy or validate a local-path override.""" + + from fastplms.registry import get_model_registry + + explicit = getattr(arguments, "bf16_execution", None) + matches = [ + spec + for spec in get_model_registry().values() + if arguments.model in {spec.fast.repo_id, spec.official.repo_id} + ] + policies = {spec.family.bf16_execution for spec in matches} + if len(policies) > 1: + raise ValueError(f"Checkpoint {arguments.model!r} has conflicting BF16 policies") + if policies: + manifest_policy = policies.pop() + if explicit is not None and explicit != manifest_policy: + raise ValueError( + f"--bf16-execution={explicit!r} conflicts with the manifest policy " + f"{manifest_policy!r} for {arguments.model!r}" + ) + return manifest_policy + if explicit is None: + raise ValueError( + "An unregistered local checkpoint requires an explicit --bf16-execution policy" + ) + return explicit + + +def _benchmark_load_dtype(arguments: argparse.Namespace, torch: Any) -> Any: + """Return the parameter-storage dtype declared by the manifest policy.""" + + return torch.float32 if _uses_bf16_autocast(arguments) else torch.bfloat16 + + +def _numeric_context(arguments: argparse.Namespace, torch: Any) -> Any: + if _uses_bf16_autocast(arguments): + return torch.autocast(device_type="cuda", dtype=torch.bfloat16) + return contextlib.nullcontext() + + +def prepare_inputs( + model: Any, + model_id: str | Path, + lengths: Sequence[int], + device: Any, + *, + revision: str | None, + local_files_only: bool, +) -> tuple[dict[str, Any], int, int, list[str]]: + """Prepare static device inputs and return biological/padded token counts.""" + + torch = _require_torch(require_cuda=False) + sequences = sequences_for_lengths(lengths) + prep_tokens = getattr(getattr(model, "model", None), "prep_tokens", None) + if prep_tokens is not None and hasattr(prep_tokens, "get_batch_kwargs"): + batch = prep_tokens.get_batch_kwargs( + sequences, device=device + ) # tensor values: model-native packed shapes + model_inputs = { + "input_ids": batch["input_ids"], # (...) + "within_seq_position_ids": batch["within_seq_position_ids"], # (...) + "global_position_ids": batch["global_position_ids"], # (...) + "sequence_ids": batch["sequence_ids"], # (...) + "attention_mask": (batch["sequence_ids"] != -1).long(), # (...) + } + else: + tokenizer = getattr(model, "tokenizer", None) + if tokenizer is None: + from transformers import AutoTokenizer + + tokenizer_kwargs: dict[str, Any] = { + "trust_remote_code": True, + "local_files_only": local_files_only, + } + if revision is not None: + tokenizer_kwargs["revision"] = revision + tokenizer = AutoTokenizer.from_pretrained(model_id, **tokenizer_kwargs) + max_length = max(lengths) + model_inputs = dict( + tokenizer( + sequences, + return_tensors="pt", + padding="max_length", + max_length=max_length, + truncation=True, + ) + ) # each tensor: (b, l) + model_inputs = { + name: value.to(device, non_blocking=True) for name, value in model_inputs.items() + } # each tensor: (b, l) + + parameters = inspect.signature(model.forward).parameters + if "sequence_id" in parameters and "sequence_id" not in model_inputs: + attention_mask = model_inputs.get("attention_mask") # (b, l) or model-native (...) + if attention_mask is None: + raise RuntimeError("A model requiring sequence_id must expose an attention mask") + model_inputs["sequence_id"] = attention_mask.to( + dtype=torch.bool + ) # (b, l) or model-native (...) + + if getattr(getattr(model, "config", None), "is_encoder_decoder", False): + decoder_start_token_id = getattr(model.config, "decoder_start_token_id", None) + if decoder_start_token_id is None: + raise RuntimeError("An encoder-decoder benchmark requires decoder_start_token_id") + batch_size = int(model_inputs["input_ids"].shape[0]) + model_inputs["decoder_input_ids"] = torch.full( + (batch_size, 1), + decoder_start_token_id, + device=device, + dtype=torch.long, + ) # (b, 1) + + # Logical throughput counts biological residues. Attention masks also include + # model-specific BOS, EOS, and other control tokens, so they cannot provide + # this count. The deterministic input strings are the shared source of truth. + logical_tokens = sum(len(sequence) for sequence in sequences) + padded_tokens = int(model_inputs["input_ids"].numel()) + return model_inputs, logical_tokens, padded_tokens, sequences + + +def cuda_sample_ms(torch: Any, operation: Callable[[], Any]) -> float: + """Time one GPU operation with CUDA events.""" + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + operation() + end.record() + end.synchronize() + return float(start.elapsed_time(end)) + + +def warm_until_stable( + torch: Any, + operation: Callable[[], Any], + *, + window: int = 10, + tolerance: float = 0.02, + minimum_samples: int = 20, + maximum_samples: int = 100, +) -> list[float]: + """Warm until adjacent timing-window medians differ by less than ``tolerance``.""" + + samples: list[float] = [] + for _ in range(maximum_samples): + samples.append(cuda_sample_ms(torch, operation)) + if len(samples) < max(minimum_samples, 2 * window): + continue + previous = statistics.median(samples[-2 * window : -window]) + current = statistics.median(samples[-window:]) + if previous > 0.0 and abs(current - previous) / previous < tolerance: + return samples + raise RuntimeError( + f"CUDA timings did not stabilize within {maximum_samples} warmup forwards; " + "inspect clocks, thermals, and background workloads" + ) + + +def measure_blocks( + torch: Any, + operation: Callable[[], Any], + *, + logical_tokens_per_forward: int, + padded_tokens_per_forward: int, + blocks: int = 7, + minimum_block_ms: float = 250.0, + minimum_forwards: int = 5, +) -> list[MeasurementBlock]: + """Collect raw samples in duration-bounded measurement blocks.""" + + output: list[MeasurementBlock] = [] + for _ in range(blocks): + samples: list[float] = [] + elapsed = 0.0 + while len(samples) < minimum_forwards or elapsed < minimum_block_ms: + duration = cuda_sample_ms(torch, operation) + samples.append(duration) + elapsed += duration + forwards = len(samples) + elapsed_seconds = elapsed / 1000.0 + logical_tokens = logical_tokens_per_forward * forwards + padded_tokens = padded_tokens_per_forward * forwards + output.append( + MeasurementBlock( + samples_ms=tuple(samples), + elapsed_ms=elapsed, + forwards=forwards, + logical_tokens=logical_tokens, + padded_tokens=padded_tokens, + logical_tokens_per_second=logical_tokens / elapsed_seconds, + padded_tokens_per_second=padded_tokens / elapsed_seconds, + ) + ) + return output + + +def _model_load_source(arguments: argparse.Namespace) -> tuple[str | Path, str | None]: + """Return physical load coordinates without changing logical report identity.""" + + return ( + getattr(arguments, "load_model", arguments.model), + getattr(arguments, "load_revision", arguments.revision), + ) + + +def _load_model(arguments: argparse.Namespace, torch: Any) -> tuple[Any, float]: + import transformers + + arguments.bf16_execution = _resolve_bf16_execution(arguments) + try: + auto_class = getattr(transformers, arguments.auto_class) + except AttributeError as error: + raise ValueError(f"Unknown Transformers AutoClass {arguments.auto_class!r}") from error + load_model, load_revision = _model_load_source(arguments) + load_kwargs: dict[str, Any] = { + "trust_remote_code": True, + "local_files_only": arguments.local_files_only, + "dtype": _benchmark_load_dtype(arguments, torch), + "device_map": torch.device("cuda"), + "attn_implementation": arguments.backend, + } + if load_revision is not None: + load_kwargs["revision"] = load_revision + if arguments.mode == "projection" and arguments.precision != "bf16": + raise ValueError( + "Learned projection consumes precomputed BF16 H; use " + "--mode esmc_projection to compare BF16 and FP8 ESMC inference." + ) + if arguments.mode == "esmfold2_embed" or arguments.precision != "bf16": + load_kwargs["esmc_precision"] = arguments.precision + esmc_load_model = getattr(arguments, "esmc_load_model", None) + if arguments.mode in {"projection", "esmc_projection"} or esmc_load_model is not None: + # Representation cases record the folding-core load and ESMC reload + # separately. Only the end-to-end case reloads ESMC before measurement. + load_kwargs["load_esmc"] = False + + torch.cuda.synchronize() + start = time.perf_counter() + model = auto_class.from_pretrained(load_model, **load_kwargs).eval() + if arguments.mode == "esmfold2_embed" and esmc_load_model is not None: + load_esmc = getattr(model, "load_esmc", None) + if load_esmc is None: + raise RuntimeError("Local ESMFold2 artifact loading requires model.load_esmc") + load_esmc( + str(esmc_load_model), + precision=arguments.precision, + device=torch.device("cuda"), + local_files_only=True, + ) + torch.cuda.synchronize() + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return model, elapsed_ms + + +def _select_backend(model: Any, backend: str) -> None: + """Change one cached model to an advertised backend without reloading it.""" + + config = getattr(model, "config", None) + current = getattr(config, "_attn_implementation", None) + if current is None: + current = getattr(config, "attn_backend", None) + if current == backend: + return + setter = getattr(model, "set_attn_implementation", None) + if setter is None: + raise RuntimeError( + f"{type(model).__name__} cannot switch from {current!r} to {backend!r}; " + "benchmark cases never reload a model to hide this contract gap." + ) + setter(backend) + + +def _measure_embedding( + torch: Any, + model: Any, + sequences: Sequence[str], + batch_size: int, + arguments: argparse.Namespace, +) -> tuple[float, Any]: + from fastplms import embed_dataset + + torch.cuda.synchronize() + start = time.perf_counter() + with _numeric_context(arguments, torch): + result = embed_dataset(model, sequences, batch_size=batch_size, pooling=("mean",)) + torch.cuda.synchronize() + return (time.perf_counter() - start) * 1000.0, result + + +def _esmfold2_residue_mask(torch: Any, lengths: Sequence[int]) -> Any: + if not lengths or any(length <= 0 for length in lengths): + raise ValueError("ESMFold2 benchmark lengths must be positive") + length_tensor = torch.tensor(lengths, device="cuda", dtype=torch.long) # (b,) + positions = torch.arange(max(lengths), device="cuda") # (l,) + residue_mask = positions.unsqueeze(0) < length_tensor.unsqueeze(1) # (b, l) + return residue_mask # (b, l) + + +def _prepare_esmfold2_inputs(torch: Any, lengths: Sequence[int]) -> tuple[dict[str, Any], int, int]: + """Create deterministic, preallocated residue tensors for ESMC inference.""" + + from fastplms.models.esmfold2.esmfold2_constants_esm3 import ( + SEQUENCE_PAD_TOKEN, + SEQUENCE_STANDARD_AA_MAX_TOKEN, + SEQUENCE_STANDARD_AA_MIN_TOKEN, + ) + + residue_mask = _esmfold2_residue_mask(torch, lengths) # (b, l) + b = len(lengths) + sequence_length = max(lengths) + input_ids = torch.full( + (b, sequence_length), SEQUENCE_PAD_TOKEN, device="cuda", dtype=torch.long + ) # (b, l) + n_residue_tokens = SEQUENCE_STANDARD_AA_MAX_TOKEN - SEQUENCE_STANDARD_AA_MIN_TOKEN + for batch_index, length in enumerate(lengths): + residues = ( + torch.arange(length, device="cuda", dtype=torch.long) + batch_index + ) % n_residue_tokens # (l_i,) + input_ids[batch_index, :length] = ( + residues + SEQUENCE_STANDARD_AA_MIN_TOKEN + ) # target slice: (l_i,) + model_inputs = { + "input_ids": input_ids, # (b, l) + "asym_id": torch.zeros_like(input_ids), # (b, l) + "residue_index": torch.arange(sequence_length, device="cuda").expand( + b, -1 + ), # (b, l) + "mol_type": torch.zeros_like(input_ids), # (b, l) + "residue_mask": residue_mask, # (b, l) + } + return model_inputs, sum(lengths), b * sequence_length + + +def _run_esmfold2_esmc_projection(model: Any, model_inputs: Mapping[str, Any]) -> Any: + """Run preallocated residues through ESMC and the learned sequence projection.""" + + # model_inputs tensor values: (b, l) + compute_hidden_states = getattr(model, "_compute_lm_hidden_states", None) + project = getattr(model, "project_esmc_hidden_states", None) + if compute_hidden_states is None or project is None: + raise RuntimeError( + "ESMC projection mode requires ESMFold2 hidden-state and projection APIs" + ) + hidden_states = compute_hidden_states( + model_inputs["input_ids"], + model_inputs["asym_id"], + model_inputs["residue_index"], + model_inputs["mol_type"], + model_inputs["residue_mask"], + ) # (b, l, 81, 2560) + projected = project( + hidden_states, residue_mask=model_inputs["residue_mask"] + ) # (b, l, 256) + return projected # (b, l, 256) + + +def _measure_projection( + torch: Any, + model: Any, + lengths: Sequence[int], +) -> tuple[float, list[float], list[MeasurementBlock]]: + project = getattr(model, "project_esmc_hidden_states", None) + if project is None: + raise RuntimeError("Projection mode requires model.project_esmc_hidden_states") + residue_mask = _esmfold2_residue_mask(torch, lengths) # (b, l) + batch_size = len(lengths) + sequence_length = max(lengths) + H = torch.randn( + (batch_size, sequence_length, 81, 2560), + device="cuda", + dtype=torch.bfloat16, + ) # (b, l, 81, 2560) + + def operation() -> Any: + with torch.inference_mode(): + projected = project(H, residue_mask=residue_mask) # (b, l, 256) + return projected # (b, l, 256) + + first_forward_ms = cuda_sample_ms(torch, operation) + warmup = warm_until_stable(torch, operation) + blocks = measure_blocks( + torch, + operation, + logical_tokens_per_forward=sum(lengths), + padded_tokens_per_forward=batch_size * sequence_length, + ) + return first_forward_ms, warmup, blocks + + +def _measure_esmc_projection( + torch: Any, + model: Any, + lengths: Sequence[int], +) -> tuple[float, list[float], list[MeasurementBlock]]: + model_inputs, logical_tokens, padded_tokens = _prepare_esmfold2_inputs( + torch, lengths + ) # each tensor: (b, l) + + def operation() -> Any: + with torch.inference_mode(): + projected = _run_esmfold2_esmc_projection( + model, model_inputs + ) # (b, l, 256) + return projected # (b, l, 256) + + first_forward_ms = cuda_sample_ms(torch, operation) + warmup = warm_until_stable(torch, operation) + blocks = measure_blocks( + torch, + operation, + logical_tokens_per_forward=logical_tokens, + padded_tokens_per_forward=padded_tokens, + ) + return first_forward_ms, warmup, blocks + + +def _precision_status_record(model: Any) -> dict[str, Any] | None: + status = getattr(model, "esmc_precision_status", None) + if status is None: + return None + if isinstance(status, Mapping): + return dict(status) + field_names = ( + "requested", + "resolved", + "reason", + "device", + "transformer_engine_version", + ) + return {name: getattr(status, name) for name in field_names if hasattr(status, name)} + + +def run_case( + arguments: argparse.Namespace, + *, + model: Any | None = None, + load_ms: float | None = None, + model_reused: bool = False, +) -> dict[str, Any]: + """Execute one benchmark case and return a JSON-serializable record.""" + + if arguments.mode == "projection" and arguments.precision != "bf16": + raise ValueError( + "Learned projection consumes precomputed BF16 H; use " + "esmc_projection for BF16 versus FP8 comparisons." + ) + arguments.bf16_execution = _resolve_bf16_execution(arguments) + torch = _require_torch() + torch.manual_seed(arguments.seed) + telemetry_before = _nvidia_smi() + if model is None: + if load_ms is not None or model_reused: + raise ValueError("load_ms and model_reused require a preloaded model") + model, load_ms = _load_model(arguments, torch) + else: + _select_backend(model, arguments.backend) + esmc_reload_ms: float | None = None + esmc_precision_status: dict[str, Any] | None = None + if arguments.mode == "esmc_projection": + reload_esmc = getattr(model, "reload_esmc", None) + if reload_esmc is None: + raise RuntimeError("ESMC projection mode requires model.reload_esmc") + status = getattr(model, "esmc_precision_status", None) + resolved = getattr(status, "resolved", None) + if resolved is None and isinstance(status, Mapping): + resolved = status.get("resolved") + if getattr(model, "_esmc", None) is None or resolved != arguments.precision: + torch.cuda.synchronize() + reload_start = time.perf_counter() + esmc_load_model = getattr(arguments, "esmc_load_model", None) + if esmc_load_model is not None and getattr(model, "_esmc", None) is None: + load_esmc = getattr(model, "load_esmc", None) + if load_esmc is None: + raise RuntimeError("Local ESMFold2 artifact loading requires model.load_esmc") + load_esmc( + str(esmc_load_model), + precision=arguments.precision, + device=torch.device("cuda"), + local_files_only=True, + ) + else: + reload_esmc( + precision=arguments.precision, + device=torch.device("cuda"), + local_files_only=arguments.local_files_only, + ) + torch.cuda.synchronize() + esmc_reload_ms = (time.perf_counter() - reload_start) * 1000.0 + status = getattr(model, "esmc_precision_status", None) + resolved = getattr(status, "resolved", None) + if resolved is None and isinstance(status, Mapping): + resolved = status.get("resolved") + if resolved != arguments.precision: + raise RuntimeError( + f"Requested ESMC precision {arguments.precision!r}, resolved {resolved!r}" + ) + esmc_precision_status = _precision_status_record(model) + elif arguments.mode == "esmfold2_embed": + esmc_precision_status = _precision_status_record(model) + resolved = None if esmc_precision_status is None else esmc_precision_status.get("resolved") + if resolved != arguments.precision: + raise RuntimeError( + f"Requested ESMC precision {arguments.precision!r}, resolved {resolved!r}" + ) + if arguments.lengths: + lengths = tuple(arguments.lengths) + if arguments.batch_size != len(lengths): + raise ValueError("--batch-size must equal the number of values passed to --lengths") + else: + lengths = (arguments.sequence_length,) * arguments.batch_size + + case = BenchmarkCase( + model=arguments.model, + revision=arguments.revision, + auto_class=arguments.auto_class, + backend=arguments.backend, + precision=arguments.precision, + bf16_execution=arguments.bf16_execution, + mode=arguments.mode, + batch_size=arguments.batch_size, + sequence_length=max(lengths), + lengths=lengths, + ) + torch.cuda.reset_peak_memory_stats() + compile_ms: float | None = None + first_forward_ms: float | None = None + embedding_ms: float | None = None + warmup_samples: list[float] = [] + blocks: list[MeasurementBlock] = [] + if arguments.mode == "startup": + pass + elif arguments.mode == "projection": + first_forward_ms, warmup_samples, blocks = _measure_projection(torch, model, lengths) + elif arguments.mode == "esmc_projection": + first_forward_ms, warmup_samples, blocks = _measure_esmc_projection(torch, model, lengths) + elif arguments.mode == "esmfold2_embed": + sequences = sequences_for_lengths(lengths, special_tokens=0) + embedding_ms, _ = _measure_embedding( + torch, model, sequences, arguments.batch_size, arguments + ) + else: + load_model, load_revision = _model_load_source(arguments) + model_inputs, logical_tokens, padded_tokens, sequences = prepare_inputs( + model, + load_model, + lengths, + torch.device("cuda"), + revision=load_revision, + local_files_only=arguments.local_files_only, + ) # tensors: tokenizer (b, l) or model-native packed shapes + + def operation() -> Any: + with torch.inference_mode(), _numeric_context(arguments, torch): + return _model_forward(model, model_inputs) + + first_forward_ms = cuda_sample_ms(torch, operation) + if arguments.mode == "compile": + torch.cuda.synchronize() + start = time.perf_counter() + model = torch.compile(model) + + def operation() -> Any: + with torch.inference_mode(), _numeric_context(arguments, torch): + return _model_forward(model, model_inputs) + + operation() + torch.cuda.synchronize() + compile_ms = (time.perf_counter() - start) * 1000.0 + elif arguments.mode == "embed": + embedding_ms, _ = _measure_embedding( + torch, model, sequences, arguments.batch_size, arguments + ) + + if arguments.mode in {"steady", "compile"}: + warmup_samples = warm_until_stable(torch, operation) + blocks = measure_blocks( + torch, + operation, + logical_tokens_per_forward=logical_tokens, + padded_tokens_per_forward=padded_tokens, + ) + + torch.cuda.synchronize() + samples = [sample for block in blocks for sample in block.samples_ms] + latency = None + if samples: + latency = { + "median_ms": statistics.median(samples), + "p95_ms": statistics.quantiles(samples, n=100, method="inclusive")[94], + } + record = asdict(case) + record.update( + { + "load_ms": load_ms, + "model_reused": model_reused, + "esmc_reload_ms": esmc_reload_ms, + "esmc_precision_status": esmc_precision_status, + "first_forward_ms": first_forward_ms, + "compile_ms": compile_ms, + "embedding_ms": embedding_ms, + "warmup_samples_ms": warmup_samples, + "blocks": [asdict(block) for block in blocks], + "latency": latency, + "memory": { + "peak_allocated_bytes": int(torch.cuda.max_memory_allocated()), + "peak_reserved_bytes": int(torch.cuda.max_memory_reserved()), + }, + "telemetry_before": telemetry_before, + "telemetry_after": _nvidia_smi(), + } + ) + return record + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Local artifact path or Hub model ID") + parser.add_argument("--revision") + parser.add_argument( + "--auto-class", + choices=("AutoModel", "AutoModelForMaskedLM", "AutoModelForSeq2SeqLM"), + default="AutoModelForMaskedLM", + ) + parser.add_argument( + "--backend", + default="sdpa", + choices=( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ), + ) + parser.add_argument("--precision", choices=("bf16", "fp8"), default="bf16") + parser.add_argument( + "--bf16-execution", + choices=("static_parameters", "fp32_parameters_autocast"), + help=( + "Required only for unregistered local checkpoints; registered IDs " + "derive and validate this policy from models.toml." + ), + ) + parser.add_argument( + "--mode", + default="steady", + choices=( + "startup", + "compile", + "steady", + "embed", + "projection", + "esmc_projection", + "esmfold2_embed", + ), + ) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--sequence-length", type=int, default=512) + parser.add_argument("--lengths", nargs="+", type=int) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: Iterable[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + torch = _require_torch() + report = { + "schema_version": 1, + "environment": environment_fingerprint(torch), + "results": [run_case(arguments)], + } + arguments.output.parent.mkdir(parents=True, exist_ok=True) + arguments.output.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(arguments.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/suite.py b/benchmarks/suite.py new file mode 100644 index 0000000..8718ec6 --- /dev/null +++ b/benchmarks/suite.py @@ -0,0 +1,1006 @@ +"""Run the manifest-declared Hopper/SM90 benchmark matrix outside pytest.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +import re +import xml.etree.ElementTree as ET +from collections.abc import Iterable, Iterator, Mapping, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +from fastplms.registry import ModelRegistry, ModelSpec, get_model_registry +from .regression import compare_reports +from .run import ( + _load_model, + _require_torch, + environment_fingerprint, + run_case, + validate_hopper_sm90_environment, +) + + +PADDING_LENGTHS = (1024, 512, 256, 128, 64, 64, 32, 32) +FIXED_SHAPES = ( + (1, 512, ()), + (8, 1024, ()), + (8, 1024, PADDING_LENGTHS), +) +EXHAUSTIVE_BATCH_SIZES = (1, 2, 4, 8) +EXHAUSTIVE_SEQUENCE_LENGTHS = (128, 256, 512, 1024) + +SEQUENCE_FORWARD_PROFILE = "sequence_forward" +ESMFOLD2_REPRESENTATION_PROFILE = "esmfold2_representation" +STRUCTURE_STARTUP_PROFILE = "structure_startup" +STRUCTURE_DEDICATED_MODE = "structure" +ESMFOLD2_DEDICATED_MODE = "representation" +FLASH_BACKEND_HISTORICAL_EVIDENCE = { + "flash_attention_2": "separate_historical_focused_evidence_only", + "flash_attention_3": "none", +} +FLASH_BACKEND_CLAIM_INELIGIBILITY_REASON = ( + "locked_release_environment_kernel_unavailable" +) +_COMMIT_PATTERN = re.compile(r"[0-9a-f]{40}") +_SHA256_PATTERN = re.compile(r"[0-9a-f]{64}") +_RUNTIME_REVISION_PATTERN = re.compile( + r"(?:[0-9a-f]{40}|source-tree-sha256:[0-9a-f]{64})" +) + + +def _load_json_object(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Unable to read benchmark artifact metadata: {path}") from error + if not isinstance(value, dict): + raise ValueError(f"Benchmark artifact metadata must be a JSON object: {path}") + return value + + +def _artifact_repository_name(spec: ModelSpec) -> str: + parts = spec.fast.repo_id.split("/") + if len(parts) != 2 or not all(parts) or parts[1] in {".", ".."}: + raise ValueError(f"Invalid registry repository ID for {spec.id}: {spec.fast.repo_id!r}") + return cast(str, parts[1]) + + +def _is_link_like(path: Path) -> bool: + is_junction = getattr(path, "is_junction", None) + return path.is_symlink() or bool(is_junction is not None and is_junction()) + + +def _validate_built_artifact( + path: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> None: + """Apply the complete Hub-artifact validator without weakening its errors.""" + + from tools.artifacts.build import ArtifactError, validate_artifact + + try: + validate_artifact(path, spec=spec, registry=registry) + except ArtifactError as error: + raise ValueError(f"Invalid benchmark artifact for {spec.id}: {error}") from error + + +def _frozen_runtime_identity( + source_root: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> tuple[str, str]: + """Return the clean tracked runtime revision and digest for one frozen source tree.""" + + from tools.artifacts.build import ArtifactError, _validated_runtime_snapshot + + try: + runtime_revision, _payloads, source_tree_sha256 = _validated_runtime_snapshot( + source_root, + registry, + spec, + ) + except ArtifactError as error: + raise ValueError( + f"Unable to validate frozen benchmark source for {spec.id}: {error}" + ) from error + return runtime_revision, source_tree_sha256 + + +def _require_identity( + value: object, + *, + name: str, + pattern: re.Pattern[str] | None = None, +) -> str: + if not isinstance(value, str) or not value or ( + pattern is not None and pattern.fullmatch(value) is None + ): + raise ValueError(f"Benchmark artifact has invalid {name}: {value!r}") + return value + + +def _artifact_path(root: Path, spec: ModelSpec) -> Path: + candidate = root / _artifact_repository_name(spec) + if _is_link_like(candidate): + raise ValueError(f"Benchmark artifact for {spec.id} may not be a link or junction") + try: + resolved = candidate.resolve(strict=True) + except OSError as error: + raise ValueError(f"Benchmark artifact is missing for {spec.id}") from error + if resolved.parent != root or not resolved.is_dir(): + raise ValueError(f"Benchmark artifact is not a contained directory for {spec.id}") + return resolved + + +def _artifact_identity( + path: Path, + spec: ModelSpec, + registry: ModelRegistry, + *, + source_root: Path, +) -> dict[str, Any]: + """Validate and return path-free identities for one locally built Hub artifact.""" + + registry.require_resolved(spec.id) + _validate_built_artifact(path, spec, registry) + expected_runtime_revision, expected_source_sha256 = _frozen_runtime_identity( + source_root, + spec, + registry, + ) + config = _load_json_object(path / "config.json") + provenance = _load_json_object(path / "provenance.json") + + expected_config: dict[str, object] = { + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "fastplms_checkpoint_revision": spec.artifact_checkpoint.revision, + "fastplms_weights_revision": spec.artifact_checkpoint.revision, + "fastplms_runtime_revision": expected_runtime_revision, + "fastplms_source_tree_sha256": expected_source_sha256, + } + config_mismatches = sorted( + name for name, expected in expected_config.items() if config.get(name) != expected + ) + if config_mismatches: + raise ValueError( + f"Benchmark artifact config for {spec.id} differs from the registry/frozen source: " + + ", ".join(config_mismatches) + ) + + expected_provenance: dict[str, object] = { + "model_id": spec.id, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": expected_runtime_revision, + "source_tree_sha256": expected_source_sha256, + } + provenance_mismatches = sorted( + name + for name, expected in expected_provenance.items() + if provenance.get(name) != expected + ) + if provenance_mismatches: + raise ValueError( + f"Benchmark artifact provenance for {spec.id} differs from the " + "registry/frozen source: " + + ", ".join(provenance_mismatches) + ) + + checkpoint = provenance.get("artifact_checkpoint") + if not isinstance(checkpoint, Mapping) or ( + checkpoint.get("repo_id") != spec.artifact_checkpoint.repo_id + or checkpoint.get("revision") != spec.artifact_checkpoint.revision + ): + raise ValueError(f"Benchmark artifact checkpoint identity differs for {spec.id}") + + runtime_revision = _require_identity( + provenance.get("runtime_revision"), + name="runtime revision", + pattern=_RUNTIME_REVISION_PATTERN, + ) + source_tree_sha256 = _require_identity( + provenance.get("source_tree_sha256"), + name="source-tree SHA-256", + pattern=_SHA256_PATTERN, + ) + runtime_bundle_sha256 = _require_identity( + provenance.get("runtime_bundle_sha256"), + name="runtime-bundle SHA-256", + pattern=_SHA256_PATTERN, + ) + if config.get("fastplms_runtime_bundle_sha256") != runtime_bundle_sha256: + raise ValueError(f"Benchmark artifact runtime-bundle identity differs for {spec.id}") + + canonical_weights = provenance.get("canonical_weights") + state_digest = ( + canonical_weights.get("state_digest") + if isinstance(canonical_weights, Mapping) + else None + ) + if not isinstance(state_digest, Mapping) or ( + state_digest.get("algorithm") != "sha256" + or state_digest.get("schema_version") != 1 + ): + raise ValueError(f"Benchmark artifact canonical-state identity is invalid for {spec.id}") + canonical_state_sha256 = _require_identity( + state_digest.get("sha256"), + name="canonical-state SHA-256", + pattern=_SHA256_PATTERN, + ) + try: + manifest_sha256 = hashlib.sha256( + (path / "artifact-manifest.json").read_bytes() + ).hexdigest() + except OSError as error: + raise ValueError(f"Benchmark artifact manifest is unavailable for {spec.id}") from error + + return { + "model_id": spec.id, + "registry_repo_id": spec.fast.repo_id, + "registry_revision": _require_identity( + spec.fast.revision, + name="registry revision", + pattern=_COMMIT_PATTERN, + ), + "checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "weights_revision": _require_identity( + provenance.get("weights_revision"), + name="weights revision", + pattern=_COMMIT_PATTERN, + ), + "runtime_revision": runtime_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + "canonical_state_sha256": canonical_state_sha256, + "artifact_manifest_sha256": manifest_sha256, + } + + +def bind_local_artifacts( + cases: Sequence[SimpleNamespace], + artifact_root: Path, + *, + source_root: Path | None = None, +) -> dict[str, dict[str, Any]]: + """Prevalidate local artifacts and bind private load-only paths to benchmark cases.""" + + if _is_link_like(artifact_root): + raise ValueError("--artifact-root may not be a link or junction") + try: + root = artifact_root.resolve(strict=True) + except OSError as error: + raise ValueError(f"--artifact-root does not exist: {artifact_root}") from error + if not root.is_dir(): + raise ValueError(f"--artifact-root is not a directory: {artifact_root}") + + registry = get_model_registry() + by_report_identity = { + (spec.fast.repo_id, spec.fast.revision): spec for spec in registry.values() + } + selected: dict[str, ModelSpec] = {} + for case in cases: + spec = by_report_identity.get((str(case.model), str(case.revision))) + if spec is None: + raise ValueError( + "Local benchmark artifacts require an exact registry model/revision identity; " + f"got {(case.model, case.revision)!r}" + ) + selected[spec.id] = spec + if spec.family.id == "esmfold2": + backbone_id = spec.family.backbone_model + if backbone_id is None: + raise ValueError(f"ESMFold2 benchmark model {spec.id} has no registry backbone") + selected[backbone_id] = registry[backbone_id] + + artifact_paths: dict[str, Path] = {} + missing: list[str] = [] + for spec in selected.values(): + try: + artifact_paths[spec.id] = _artifact_path(root, spec) + except ValueError: + missing.append(f"{spec.id} ({spec.fast.repo_id})") + if missing: + raise ValueError("Missing or invalid selected benchmark artifacts: " + ", ".join(missing)) + + path_owners: dict[Path, str] = {} + for model_id, path in artifact_paths.items(): + previous = path_owners.setdefault(path, model_id) + if previous != model_id: + raise ValueError( + f"Benchmark artifact path collision between {previous!r} and {model_id!r}" + ) + + frozen_root = (source_root or Path(__file__).resolve().parents[1]).resolve() + identities = { + model_id: _artifact_identity( + artifact_paths[model_id], + selected[model_id], + registry, + source_root=frozen_root, + ) + for model_id in sorted(selected) + } + for case in cases: + spec = by_report_identity[(str(case.model), str(case.revision))] + case.load_model = artifact_paths[spec.id] + case.load_revision = None + case.local_files_only = True + case.artifact_identity = identities[spec.id] + if spec.family.id == "esmfold2": + backbone_id = spec.family.backbone_model + if backbone_id is None: + raise ValueError(f"ESMFold2 benchmark model {spec.id} has no registry backbone") + case.esmc_load_model = artifact_paths[backbone_id] + case.artifact_dependencies = {"esmc": identities[backbone_id]} + else: + case.esmc_load_model = None + case.artifact_dependencies = {} + return identities + + +def benchmark_auto_class(spec: ModelSpec) -> str: + """Select the manifest-advertised head measured for one architecture.""" + + advertised = set(spec.auto_map) + if ( + spec.family.id == "ankh" + or spec.family.tokenizer_mode == "structure" + or "AutoModelForMaskedLM" not in advertised + ): + selected = "AutoModel" + else: + selected = "AutoModelForMaskedLM" + if selected not in advertised: + raise ValueError(f"{spec.id} does not advertise required benchmark class {selected}") + return selected + + +def benchmark_model_key(arguments: SimpleNamespace) -> tuple[str, str, str, str, str, str]: + """Return the checkpoint identity that can share one in-memory model.""" + + artifact_identity = getattr(arguments, "artifact_identity", None) + artifact_manifest = ( + artifact_identity.get("artifact_manifest_sha256") + if isinstance(artifact_identity, Mapping) + else "" + ) + return ( + str(arguments.model), + str(arguments.revision), + str(arguments.auto_class), + str(arguments.precision), + str(arguments.bf16_execution), + str(artifact_manifest), + ) + + +def _default_backend(spec: ModelSpec) -> str: + if "sdpa" in spec.family.attention: + return "sdpa" + if not spec.family.attention: + raise ValueError(f"{spec.id} does not declare a benchmark backend") + return spec.family.attention[0] + + +def _selected_backends( + spec: ModelSpec, + requested_backends: Sequence[str] | None, +) -> tuple[str, ...]: + """Select a declared backend subset while preserving manifest order.""" + + if requested_backends is None: + return cast(tuple[str, ...], spec.family.attention) + requested = set(requested_backends) + selected = tuple( + backend for backend in spec.family.attention if backend in requested + ) + if not selected: + raise ValueError( + f"Requested benchmark backends do not apply to {spec.id}: " + + ", ".join(requested_backends) + ) + return selected + + +def _arguments( + spec: ModelSpec, + *, + backend: str, + mode: str, + batch_size: int, + sequence_length: int, + lengths: tuple[int, ...] = (), + precision: str = "bf16", + local_files_only: bool, + profile: str, + dedicated_mode: str | None = None, + claim_eligible: bool | None = None, + matrix_kind: str = "fixed", +) -> SimpleNamespace: + historical_evidence = "current_release_execution" + claim_eligibility_reason = "fixed_release_benchmark_matrix" + if backend in FLASH_BACKEND_HISTORICAL_EVIDENCE: + claim_eligible = False + historical_evidence = FLASH_BACKEND_HISTORICAL_EVIDENCE[backend] + claim_eligibility_reason = FLASH_BACKEND_CLAIM_INELIGIBILITY_REASON + elif claim_eligible is None: + claim_eligible = mode in { + "compile", + "steady", + "projection", + "esmc_projection", + } + if not claim_eligible and backend not in FLASH_BACKEND_HISTORICAL_EVIDENCE: + historical_evidence = "not_applicable" + claim_eligibility_reason = "descriptive_or_startup_measurement" + return SimpleNamespace( + model=spec.fast.repo_id, + revision=spec.fast.revision, + auto_class=benchmark_auto_class(spec), + backend=backend, + precision=precision, + bf16_execution=spec.family.bf16_execution, + mode=mode, + batch_size=batch_size, + sequence_length=sequence_length, + lengths=lengths, + local_files_only=local_files_only, + seed=42, + output=None, + suite_profile=profile, + dedicated_mode=dedicated_mode, + claim_eligible=claim_eligible, + claim_eligibility_reason=claim_eligibility_reason, + historical_evidence=historical_evidence, + matrix_kind=matrix_kind, + ) + + +def _representative_specs(family: str | None) -> list[ModelSpec]: + registry = get_model_registry() + specs = [ + spec + for spec in registry.values() + if (spec.is_deep_reference or spec.family.id == "esmfold2") + and "benchmark" in spec.family.test_tiers + ] + if family is not None: + specs = [spec for spec in specs if spec.family.id == family] + if not specs: + raise ValueError(f"No benchmark representative matches family={family!r}") + return specs + + +def benchmark_artifact_model_ids() -> tuple[str, ...]: + """Return the registry models needed by the complete fixed benchmark matrix.""" + + registry = get_model_registry() + selected: dict[str, None] = {} + for spec in _representative_specs(None): + selected[spec.id] = None + if spec.family.id == "esmfold2": + backbone_id = spec.family.backbone_model + if backbone_id is None: + raise ValueError(f"ESMFold2 benchmark model {spec.id} has no registry backbone") + selected[backbone_id] = None + # Registry order is the release order and keeps artifact construction stable. + return tuple(model_id for model_id in registry if model_id in selected) + + +def _axis(values: Iterable[int], name: str) -> tuple[int, ...]: + result = tuple(values) + if not result or any(value <= 0 for value in result): + raise ValueError(f"{name} must contain positive integers") + return result + + +def benchmark_cases( + *, + family: str | None, + quick: bool, + local_files_only: bool, + backends: Sequence[str] | None = None, +) -> Iterator[SimpleNamespace]: + """Yield the fixed benchmark matrix derived from ``models.toml``.""" + + specs = _representative_specs(family) + + if quick: + spec = specs[0] + if spec.family.id == "esmfold2": + yield _arguments( + spec, + backend=_default_backend(spec), + mode="projection", + batch_size=1, + sequence_length=16, + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + ) + return + if spec.family.tokenizer_mode == "structure": + yield _arguments( + spec, + backend=_default_backend(spec), + mode="startup", + batch_size=1, + sequence_length=1, + local_files_only=local_files_only, + profile=STRUCTURE_STARTUP_PROFILE, + dedicated_mode=STRUCTURE_DEDICATED_MODE, + claim_eligible=False, + ) + return + yield _arguments( + spec, + backend=_default_backend(spec), + mode="steady", + batch_size=1, + sequence_length=128, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + ) + return + + for spec in specs: + if spec.family.id == "esmfold2": + for batch_size, sequence_length, lengths in FIXED_SHAPES: + yield _arguments( + spec, + backend=_default_backend(spec), + mode="projection", + batch_size=batch_size, + sequence_length=sequence_length, + lengths=lengths, + precision="bf16", + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + ) + for precision in ("bf16", "fp8"): + for backend in _selected_backends(spec, backends): + for batch_size, sequence_length, lengths in FIXED_SHAPES: + yield _arguments( + spec, + backend=backend, + mode="esmc_projection", + batch_size=batch_size, + sequence_length=sequence_length, + lengths=lengths, + precision=precision, + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + ) + yield _arguments( + spec, + backend=_default_backend(spec), + mode="esmfold2_embed", + batch_size=1, + sequence_length=512, + precision=precision, + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + claim_eligible=False, + ) + continue + if spec.family.tokenizer_mode == "structure": + yield _arguments( + spec, + backend=_default_backend(spec), + mode="startup", + batch_size=1, + sequence_length=1, + local_files_only=local_files_only, + profile=STRUCTURE_STARTUP_PROFILE, + dedicated_mode=STRUCTURE_DEDICATED_MODE, + claim_eligible=False, + ) + continue + + yield _arguments( + spec, + backend=_default_backend(spec), + mode="startup", + batch_size=1, + sequence_length=512, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + claim_eligible=False, + ) + yield _arguments( + spec, + backend=_default_backend(spec), + mode="embed", + batch_size=1, + sequence_length=512, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + claim_eligible=False, + ) + for backend in _selected_backends(spec, backends): + yield _arguments( + spec, + backend=backend, + mode="compile", + batch_size=1, + sequence_length=512, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + ) + yield _arguments( + spec, + backend=backend, + mode="steady", + batch_size=1, + sequence_length=512, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + ) + yield _arguments( + spec, + backend=backend, + mode="steady", + batch_size=8, + sequence_length=1024, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + ) + yield _arguments( + spec, + backend=backend, + mode="steady", + batch_size=8, + sequence_length=1024, + lengths=PADDING_LENGTHS, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + ) + + +def exhaustive_benchmark_cases( + *, + family: str | None, + batch_sizes: Iterable[int] = EXHAUSTIVE_BATCH_SIZES, + sequence_lengths: Iterable[int] = EXHAUSTIVE_SEQUENCE_LENGTHS, + local_files_only: bool, + backends: Sequence[str] | None = None, +) -> Iterator[SimpleNamespace]: + """Yield a descriptive all-checkpoint sweep that is never claim-eligible.""" + + batches = _axis(batch_sizes, "batch_sizes") + lengths = _axis(sequence_lengths, "sequence_lengths") + registry = get_model_registry() + specs = [ + spec + for spec in registry.values() + if "benchmark" in spec.family.test_tiers and (family is None or spec.family.id == family) + ] + if not specs: + raise ValueError(f"No exhaustive benchmark checkpoint matches family={family!r}") + + for spec in specs: + if spec.family.id == "esmfold2": + for batch_size in batches: + for sequence_length in lengths: + yield _arguments( + spec, + backend=_default_backend(spec), + mode="projection", + batch_size=batch_size, + sequence_length=sequence_length, + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + claim_eligible=False, + matrix_kind="exhaustive", + ) + for precision in ("bf16", "fp8"): + for backend in _selected_backends(spec, backends): + yield _arguments( + spec, + backend=backend, + mode="esmc_projection", + batch_size=batch_size, + sequence_length=sequence_length, + precision=precision, + local_files_only=local_files_only, + profile=ESMFOLD2_REPRESENTATION_PROFILE, + dedicated_mode=ESMFOLD2_DEDICATED_MODE, + claim_eligible=False, + matrix_kind="exhaustive", + ) + continue + if spec.family.tokenizer_mode == "structure": + yield _arguments( + spec, + backend=_default_backend(spec), + mode="startup", + batch_size=1, + sequence_length=1, + local_files_only=local_files_only, + profile=STRUCTURE_STARTUP_PROFILE, + dedicated_mode=STRUCTURE_DEDICATED_MODE, + claim_eligible=False, + matrix_kind="exhaustive", + ) + continue + for backend in _selected_backends(spec, backends): + for batch_size in batches: + for sequence_length in lengths: + yield _arguments( + spec, + backend=backend, + mode="steady", + batch_size=batch_size, + sequence_length=sequence_length, + local_files_only=local_files_only, + profile=SEQUENCE_FORWARD_PROFILE, + claim_eligible=False, + matrix_kind="exhaustive", + ) + + +def _write_report(path: Path, report: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def _load_report(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object in {path}") + return value + + +def _write_junit( + path: Path, + *, + suite_name: str, + failures: Sequence[str] = (), +) -> None: + """Write one atomic, dependency-free JUnit summary for orchestration.""" + + root = ET.Element( + "testsuite", + { + "name": suite_name, + "tests": "1", + "failures": "1" if failures else "0", + "errors": "0", + "skipped": "0", + }, + ) + case = ET.SubElement(root, "testcase", {"classname": "benchmarks", "name": suite_name}) + if failures: + failure = ET.SubElement(case, "failure", {"message": failures[0]}) + failure.text = "\n".join(failures) + ET.indent(root, space=" ") + payload = ET.tostring(root, encoding="utf-8", xml_declaration=True) + b"\n" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_bytes(payload) + temporary.replace(path) + + +def _gate_failures(gate: Any) -> tuple[str, ...]: + failures = [ + *gate.report_mismatches, + *gate.environment_mismatches, + *gate.artifact_mismatches, + *(f"unmatched current case: {case}" for case in gate.unmatched_current), + *(f"unmatched baseline case: {case}" for case in gate.unmatched_baseline), + ] + failures.extend( + f"{case.case}: {reason}" + for case in gate.cases + for reason in case.reasons + ) + if not gate.passed and not failures: + failures.append("Benchmark gate did not contain any comparable throughput cases") + return tuple(failures) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--baseline", type=Path) + parser.add_argument("--gate-output", type=Path) + parser.add_argument("--junit-output", type=Path) + parser.add_argument("--family") + matrix = parser.add_mutually_exclusive_group() + matrix.add_argument("--quick", action="store_true") + matrix.add_argument("--exhaustive", action="store_true") + parser.add_argument( + "--exhaustive-batch-sizes", + nargs="+", + type=int, + default=EXHAUSTIVE_BATCH_SIZES, + ) + parser.add_argument( + "--exhaustive-sequence-lengths", + nargs="+", + type=int, + default=EXHAUSTIVE_SEQUENCE_LENGTHS, + ) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument( + "--backends", + nargs="+", + choices=( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ), + help=( + "Restrict the matrix to this explicit backend subset. The GH200 release " + "runner passes eager, SDPA, and Flex and never downloads Flash kernels." + ), + ) + parser.add_argument( + "--artifact-root", + type=Path, + help=( + "Load the selected registry checkpoints from validated locally built Hub " + "artifacts while retaining registry repo/revision report identities." + ), + ) + return parser + + +def main(argv: Iterable[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + if arguments.junit_output is not None: + _write_junit( + arguments.junit_output, + suite_name="benchmark-incomplete", + failures=("Benchmark process did not complete",), + ) + if arguments.backends is not None and len(set(arguments.backends)) != len( + arguments.backends + ): + raise ValueError("--backends may not contain duplicates") + local_files_only = arguments.local_files_only or arguments.artifact_root is not None + if arguments.exhaustive: + if arguments.baseline is not None: + raise ValueError("Exhaustive sweeps are descriptive and cannot gate a baseline") + cases = list( + exhaustive_benchmark_cases( + family=arguments.family, + batch_sizes=arguments.exhaustive_batch_sizes, + sequence_lengths=arguments.exhaustive_sequence_lengths, + local_files_only=local_files_only, + backends=arguments.backends, + ) + ) + else: + cases = list( + benchmark_cases( + family=arguments.family, + quick=arguments.quick, + local_files_only=local_files_only, + backends=arguments.backends, + ) + ) + artifact_identities: dict[str, dict[str, Any]] = {} + if arguments.artifact_root is not None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + artifact_identities = bind_local_artifacts(cases, arguments.artifact_root) + + torch = _require_torch() + environment = environment_fingerprint(torch) + if not arguments.quick and not arguments.exhaustive: + validate_hopper_sm90_environment(environment) + report: dict[str, Any] = { + "schema_version": 3, + "status": "running", + "environment": environment, + "matrix_kind": ( + "exhaustive" if arguments.exhaustive else "quick" if arguments.quick else "fixed" + ), + "claim_scope": ( + "descriptive_only" + if arguments.exhaustive + else "smoke_only" + if arguments.quick + else "validated_hopper_sm90_exact_device" + ), + "artifact_load_mode": ( + "validated_local_build" if arguments.artifact_root is not None else "hub" + ), + "artifacts": artifact_identities, + "backend_policy": { + "requested": list(arguments.backends) if arguments.backends is not None else None, + "selection": ( + "explicit_subset" if arguments.backends is not None else "manifest_all" + ), + "external_kernel_downloads": False, + "external_kernel_builds": False, + }, + "timing_contract": { + "cold_compile_field": "results[].compile_ms", + "first_forward_field": "results[].first_forward_ms", + "warmup_field": "results[].warmup_samples_ms", + "warm_throughput_field": "results[].blocks", + "compile_amortized_into_throughput": False, + }, + "baseline_promotion_contract": { + "report_is_complete_when": "all cases are present and the process exits zero", + "legacy_baseline_path": "benchmarks/baselines/h100.json", + "requires_exact_environment_match": True, + "requires_exact_artifact_inventory_match": True, + }, + "expected_case_count": len(cases), + "completed_case_count": 0, + "results": [], + } + cached_key: tuple[str, str, str, str, str, str] | None = None + cached_model: Any | None = None + for case in cases: + key = benchmark_model_key(case) + reused = cached_model is not None and key == cached_key + load_ms: float | None = None + if not reused: + if cached_model is not None: + del cached_model + gc.collect() + torch.cuda.empty_cache() + cached_model, load_ms = _load_model(case, torch) + cached_key = key + result = run_case( + case, + model=cached_model, + load_ms=load_ms, + model_reused=reused, + ) + result.update( + { + "suite_profile": case.suite_profile, + "dedicated_mode": case.dedicated_mode, + "claim_eligible": case.claim_eligible, + "claim_eligibility_reason": case.claim_eligibility_reason, + "historical_evidence": case.historical_evidence, + "matrix_kind": case.matrix_kind, + "artifact": getattr(case, "artifact_identity", None), + "artifact_dependencies": getattr(case, "artifact_dependencies", {}), + } + ) + report["results"].append(result) + report["completed_case_count"] = len(report["results"]) + _write_report(arguments.output, report) + gc.collect() + torch.cuda.empty_cache() + + report["status"] = "complete" + _write_report(arguments.output, report) + if arguments.baseline is None: + if arguments.junit_output is not None: + _write_junit(arguments.junit_output, suite_name="benchmark-capture") + return 0 + gate = compare_reports(report, _load_report(arguments.baseline)) + gate_output = arguments.gate_output or arguments.output.with_suffix(".gate.json") + _write_report(gate_output, gate.to_dict()) + if arguments.junit_output is not None: + _write_junit( + arguments.junit_output, + suite_name="benchmark-regression", + failures=_gate_failures(gate), + ) + return 0 if gate.passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/build_images.sh b/build_images.sh deleted file mode 100644 index 66c5303..0000000 --- a/build_images.sh +++ /dev/null @@ -1,39 +0,0 @@ -#!/usr/bin/env bash -# Build the FastPLMs base image and every per-family image (skipping Boltz for -# now; deferred until its native deps are worked out). -# -# Usage: ./build_images.sh [family ...] -# ./build_images.sh # build base + all families -# ./build_images.sh esm_plusplus # build base + just esm_plusplus -# ./build_images.sh --no-base esm2 # skip base rebuild, just esm2 -set -euo pipefail - -FAMILIES_DEFAULT=(esm2 esm_plusplus esm3 esmfold2 e1 dplm dplm2 ankh) - -build_base=1 -families=() -for arg in "$@"; do - case "$arg" in - --no-base) build_base=0 ;; - *) families+=("$arg") ;; - esac -done - -if [ ${#families[@]} -eq 0 ]; then - families=("${FAMILIES_DEFAULT[@]}") -fi - -if [ $build_base -eq 1 ]; then - echo "==> Building fastplms-base" - docker build -f Dockerfile.base -t fastplms-base . -fi - -for family in "${families[@]}"; do - dockerfile="Dockerfile.${family}" - if [ ! -f "$dockerfile" ]; then - echo "Skipping ${family}: ${dockerfile} does not exist" - continue - fi - echo "==> Building fastplms-${family}" - docker build -f "$dockerfile" -t "fastplms-${family}" . -done diff --git a/cookbook/tutorials/binder_design_fastplms.ipynb b/cookbook/tutorials/binder_design_fastplms.ipynb deleted file mode 100644 index 4826d6c..0000000 --- a/cookbook/tutorials/binder_design_fastplms.ipynb +++ /dev/null @@ -1,320 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# FastPLMs Binder Design\n", - "\n", - "This notebook mirrors the binder design tutorial using only FastPLMs-owned models and helpers. Heavy local runs should be executed on the Linux GPU workstation or sent to Modal. The notebook assumes `cookbook/tutorials/binder_design_fastplms.py` is available from the repository root." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Install Notebook Helpers\n", - "\n", - "The script owns the model code and CLI. These packages are only for notebook orchestration, tables, and structure viewing." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "! pip install modal pandas pyarrow py3dmol tqdm biopython" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Modal Setup\n", - "\n", - "Deploy the FastPLMs binder app once, then reuse the named class for blocking or spawned jobs. Hugging Face authentication should be configured outside this notebook." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "! modal token info\n", - "! modal deploy cookbook/tutorials/binder_design_fastplms.py" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from itertools import product\n", - "import json\n", - "from pathlib import Path\n", - "\n", - "import modal\n", - "import pandas as pd\n", - "import py3Dmol\n", - "from tqdm.contrib.concurrent import thread_map\n", - "\n", - "from cookbook.tutorials.binder_design_fastplms import select_official_designs\n", - "\n", - "APP_NAME = \"fastplms-binder-design\"\n", - "FastPLMsBinderDesign = modal.Cls.from_name(APP_NAME, \"FastPLMsBinderDesignModal\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Tiny Local Smoke Run\n", - "\n", - "Use this only on a GPU workstation with the FastPLMs Docker or runtime environment available. The reduced step count is for checking wiring and outputs, not for design quality." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "! python cookbook/tutorials/binder_design_fastplms.py --backend local --target-name pd-l1 --binder-name minibinder --steps 2 --batch-size 1 --output-dir binder_design_fastplms_smoke" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Blocking Modal Run\n", - "\n", - "This waits for the remote job and returns the same artifacts as the local CLI: best sequences, optimization trajectory, critic rows, and an official-style selection table. The selection score mirrors the paper and Biohub notebook: pI-filter minibinders, average hero-critic iPTM, optionally average scaling-critic distogram proxy, then combine those two components." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "app = FastPLMsBinderDesign(\n", - " use_scaling_critics=False,\n", - " kernel_backend=None,\n", - " compile_model=False,\n", - ")\n", - "\n", - "best_sequences, trajectory, rows = app.design.remote(\n", - " target_name=\"pd-l1\",\n", - " binder_name=\"minibinder\",\n", - " seed=0,\n", - " batch_size=1,\n", - " steps=20,\n", - " output_dir=\"binder_design_fastplms_modal\",\n", - ")\n", - "\n", - "results = pd.DataFrame(rows)\n", - "selection = select_official_designs(results)\n", - "best_sequences, selection.head(), results.drop(columns=[\"pdb\", \"cif\"]).head()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Spawned Modal Run\n", - "\n", - "Spawned calls are useful for long designs and sweeps. Store the call ID so the job can be reattached later." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "future = app.design.spawn(\n", - " target_name=\"ctla4\",\n", - " binder_name=\"minibinder\",\n", - " seed=1,\n", - " batch_size=1,\n", - " steps=50,\n", - " output_dir=\"binder_design_fastplms_ctla4\",\n", - ")\n", - "\n", - "future.object_id" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "! modal app logs fastplms-binder-design -f --function-call {future.object_id}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "best_sequences, trajectory, rows = modal.FunctionCall.gather(future)[0]\n", - "results = pd.DataFrame(rows)\n", - "results.sort_values(\"final_loss\").head()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Resumable Sweep Manifest\n", - "\n", - "Each spawned call is recorded to disk immediately. If the notebook stops, recreate `modal.FunctionCall` objects from the saved IDs and gather them later." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sweep_dir = Path(\"binder_design_fastplms_sweep\")\n", - "sweep_dir.mkdir(parents=True, exist_ok=True)\n", - "manifest_path = sweep_dir / \"manifest.json\"\n", - "\n", - "axes = {\n", - " \"target_name\": [\"pd-l1\"],\n", - " \"binder_name\": [\"minibinder\", \"trastuzumab_framework_vhvl\"],\n", - " \"seed\": [0, 1],\n", - "}\n", - "jobs = [dict(zip(axes, values)) for values in product(*axes.values())]\n", - "\n", - "records = []\n", - "for index, job in enumerate(jobs):\n", - " call = app.design.spawn(\n", - " **job,\n", - " batch_size=1,\n", - " steps=50,\n", - " output_dir=str(sweep_dir / f\"job_{index:03d}\"),\n", - " )\n", - " records.append({\"index\": index, \"call_id\": call.object_id, **job})\n", - " manifest_path.write_text(json.dumps(records, indent=2))\n", - "\n", - "pd.DataFrame(records)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "records = json.loads(manifest_path.read_text())\n", - "calls = [modal.FunctionCall.from_id(record[\"call_id\"]) for record in records]\n", - "\n", - "def call_status(call):\n", - " graph = call.get_call_graph()\n", - " return graph[0].status.name\n", - "\n", - "status_table = pd.DataFrame(records)\n", - "status_table[\"status\"] = thread_map(call_status, calls)\n", - "status_table" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "outputs = modal.FunctionCall.gather(*calls)\n", - "rows = []\n", - "for record, output in zip(records, outputs):\n", - " best_sequences, trajectory, critic_rows = output\n", - " for row in critic_rows:\n", - " rows.append({**record, **row})\n", - "\n", - "sweep_results = pd.DataFrame(rows)\n", - "sweep_results.to_parquet(sweep_dir / \"results.parquet\", index=False)\n", - "selection = select_official_designs(sweep_results)\n", - "selection.to_parquet(sweep_dir / \"selection.parquet\", index=False)\n", - "selection.head()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualize A Design\n", - "\n", - "Visualize a row for the top official selection. When available, use the full Cutoff2025 hero critic structure for consistency with the Biohub tutorial visualization." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "display_rows = sweep_results if \"sweep_results\" in globals() else results\n", - "selection_rows = selection if \"selection\" in globals() else select_official_designs(display_rows)\n", - "selected_sequence = selection_rows.iloc[0][\"designed_sequence\"]\n", - "matching_rows = display_rows[display_rows[\"designed_sequence\"] == selected_sequence]\n", - "preferred_rows = matching_rows[\n", - " matching_rows[\"critic_name\"] == \"ESMFold2-Experimental-Cutoff2025\"\n", - "]\n", - "row = preferred_rows.iloc[0] if len(preferred_rows) else matching_rows.iloc[0]\n", - "\n", - "view = py3Dmol.view(width=720, height=520)\n", - "view.addModel(row[\"pdb\"], \"pdb\")\n", - "view.setStyle({\"chain\": \"A\"}, {\"cartoon\": {\"color\": \"lightgray\"}})\n", - "view.setStyle({\"chain\": \"B\"}, {\"cartoon\": {\"color\": \"deepskyblue\"}})\n", - "view.zoomTo()\n", - "view.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Local CLI Reference\n", - "\n", - "Use `--backend local` for workstation GPU runs and `--backend modal` for a deployed Modal class. The script writes `trajectory.jsonl`, `best_sequences.fasta`, `results.parquet`, per-critic structures, and final logits into the selected output directory." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "! python cookbook/tutorials/binder_design_fastplms.py --help" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/cookbook/tutorials/binder_design_fastplms.py b/cookbook/tutorials/binder_design_fastplms.py deleted file mode 100644 index 38b3ee5..0000000 --- a/cookbook/tutorials/binder_design_fastplms.py +++ /dev/null @@ -1,1421 +0,0 @@ -# /// script -# requires-python = ">=3.12" -# dependencies = [ -# "abnumber", -# "biopython", -# "modal", -# "pandas", -# "pyarrow", -# "tqdm", -# ] -# /// -"""FastPLMs binder design with local and Modal execution. - -This is a FastPLMs-only variant of the Biohub ESMFold2 binder design tutorial. -It uses FastPLMs ESMFold2 experimental checkpoints for folding and FastPLMs -ESM++ checkpoints for the masked-LM regularizer. -""" - -from __future__ import annotations - -import argparse -import inspect -import json -import logging -import math -import os -import random -from dataclasses import dataclass -from functools import cache -from pathlib import Path -from typing import Any - -import torch -import torch.nn.functional as F -import torch.optim as optim -from tqdm.auto import tqdm -from transformers import AutoModel, AutoModelForMaskedLM - -from fastplms.esm_plusplus.modeling_esm_plusplus import EsmSequenceTokenizer -from fastplms.esmfold2.esmfold2_constants import ( - ELEMENT_NUMBER_TO_SYMBOL, - PROTEIN_1TO3, - PROTEIN_3TO1, - RES_TYPE_TO_CCD, -) -from fastplms.esmfold2.modeling_esmfold2_common import _seed_context as seed_context - -try: - import modal -except ImportError: - modal = None - -os.environ["HF_XET_HIGH_PERFORMANCE"] = "1" -logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) - - -TOKENS = ["", "-"] + [RES_TYPE_TO_CCD[i] for i in range(2, 33)] -ELEMENTS = ["X"] * (max(ELEMENT_NUMBER_TO_SYMBOL) + 1) -ELEMENTS[0] = "" -for _atomic_num, _symbol in ELEMENT_NUMBER_TO_SYMBOL.items(): - ELEMENTS[_atomic_num] = _symbol[:1] + _symbol[1:].lower() -TOKEN_IDS = {token: idx for idx, token in enumerate(TOKENS)} -AA_DIMS = 20 -CYS_IDX = TOKEN_IDS[PROTEIN_1TO3["C"]] - 2 -MUTABLE_TOKEN = "#" -BinderPromptStr = str - -LOSS_WEIGHTS = {"intra_contact": 0.5, "inter_contact": 0.5, "glob": 0.2} -DEFAULT_STEPS = 150 -DEFAULT_LOG_INTERVAL = 5 -DEFAULT_LEARNING_RATE = 0.1 -DEFAULT_TEMPERATURE_MIN = 1e-2 -DEFAULT_ESMC_MASK_FRACTION = 0.15 -DEFAULT_SELECTION_TOP_K = 84 -MINIBINDER_PI_CUTOFF = 6.0 -DEFAULT_CONSENSUS_IPTM_THRESHOLD = 0.9 -SCALING_CHECKPOINT_SUBSTRING = "ESMFold2-Experimental-Fast-base" - - -def _is_scaling_critic_name(critic_name: str) -> bool: - return SCALING_CHECKPOINT_SUBSTRING in critic_name - - -@dataclass(frozen=True) -class PromptFactory: - name: str - template: str - length_ranges: dict[str, tuple[int, int]] - is_antibody: bool - - def sample(self, seed: int) -> BinderPromptStr: - random.seed(seed) - sampled_lengths = { - key: MUTABLE_TOKEN * random.randint(low, high) - for key, (low, high) in self.length_ranges.items() - } - return self.template.format(**sampled_lengths) - - -BINDER_PROMPT_FACTORIES = { - "minibinder": PromptFactory( - name="minibinder", - template="{seq}", - length_ranges={"seq": (60, 200)}, - is_antibody=False, - ), - "trastuzumab_framework_vhvl": PromptFactory( - name="trastuzumab_framework_vhvl", - template=( - "EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}YIHWVRQAPGKGLEWVARI{hcdr2}" - "TRYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCSR{hcdr3}WGQGTLVTVSS" - "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" - "{lcdr2}GVPSRFSGSRSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" - ), - length_ranges={ - "hcdr1": (7, 9), - "hcdr2": (5, 6), - "hcdr3": (9, 15), - "lcdr1": (11, 16), - "lcdr2": (7, 7), - "lcdr3": (9, 9), - }, - is_antibody=True, - ), - "atezolizumab_framework_vhvl": PromptFactory( - name="atezolizumab_framework_vhvl", - template=( - "EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}WIHWVRQAPGKGLEWVAWI{hcdr2}" - "TYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCAR{hcdr3}WGQGTLVTVSS" - "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" - "{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" - ), - length_ranges={ - "hcdr1": (7, 9), - "hcdr2": (5, 6), - "hcdr3": (9, 15), - "lcdr1": (11, 16), - "lcdr2": (7, 7), - "lcdr3": (9, 9), - }, - is_antibody=True, - ), - "ocankitug_framework_vhvl": PromptFactory( - name="ocankitug_framework_vhvl", - template=( - "QVQLVQSGAEVKKPGSSVKVSCKAS{hcdr1}WMHWVRQAPGQGLEWMGII{hcdr2}" - "TSLNQKFQGRVTITADTSTSTAYMELSSLRSEDTAVYYCAR{hcdr3}WGQGTLVTVSS" - "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" - "{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" - ), - length_ranges={ - "hcdr1": (7, 9), - "hcdr2": (5, 6), - "hcdr3": (8, 14), - "lcdr1": (11, 16), - "lcdr2": (7, 7), - "lcdr3": (9, 9), - }, - is_antibody=True, - ), -} - -TARGET_SEQUENCES = { - "cd45": ( - "GSPGEPQIIFCRSEAAHQGVITWNPPQRSFHNFTLCYIKETEKDCLNLDKNLIKYDLQNLKPYT" - "KYVLSLHAYIIAKVQRNGSAAMCHFTTKSAPPSQVWNMTVSMTSDNSMHVKCRPPRDRNGPHE" - "RYHLEVEAGNTLVRNESHKNCDFRVKDLQYSTDYTFKAYFHNGDYPGEPFILHHSTSY" - ), - "ctla4": ( - "MHVAQPAVVLASSRGIASFVCEYASPGKATEVRVTVLRQADSQVTEVCAATYMMGNELTFLDDSI" - "CTGTSSGNQVNLTIQGLRAMDTGLYICKVELMYPPPYYLGIGNGTQIYVIDPE" - ), - "egfr": ( - "RKVCNGIGIGEFKDSLSINATNIKHFKNCTSISGDLHILPVAFRGDSFTHTPPLDPQELDILKTV" - "KEITGFLLIQAWPENRTDLHAFENLEIIRGRTKQHGQFSLAVVSLNITSLGLRSLKEISDGDV" - "IISGNKNLCYANTINWKKLFGTSGQKTKIISNRGENSCKATGQVCHALCSPEGCWGPEPRDCV" - ), - "pd-l1": ( - "AFTVTVPKDLYVVEYGSNMTIECKFPVEKQLDLAALIVYWEMEDKNIIQFVHGEEDLKVQHSSYR" - "QRARLLKDQLSLGNAALQITDVKLQDAGVYRCMISYGGADYKRITVKVNA" - ), - "pdgfr": ( - "GFLPNDAEELFIFLTEITEITIPCRVTDPQLVVTLHEKKGDVALPVPYDHQRGFSGIFEDRSYIC" - "KTTIGDREVDSDAYYVYRLQVSSINVSVNAVQTVVRQGENITLMCIVIGNEVVNFEWTYPRKES" - "GRLVEPVTDFLLDMPYHIRSILHIPSAELEDSGTYTCNVTESVNDHQDEKAINITVVE" - ), -} - - -def _repo_name(name: str) -> str: - if "/" in name: - return name - return f"Synthyra/{name}" - - -def build_initial_soft_sequence_logits(sequence: str, batch_size: int) -> torch.Tensor: - if all(aa == MUTABLE_TOKEN for aa in sequence): - logits = 0.01 * torch.randn([batch_size, len(sequence), AA_DIMS]) - logits[:, :, CYS_IDX] = -1e6 - else: - logits = torch.zeros([batch_size, len(sequence), AA_DIMS]) - for i, aa in enumerate(sequence): - if aa == MUTABLE_TOKEN: - logits[:, i, :] = 0.01 * torch.randn(batch_size, AA_DIMS) - logits[:, i, CYS_IDX] = -1e6 - else: - assert aa in PROTEIN_1TO3, aa - token_id = TOKEN_IDS[PROTEIN_1TO3[aa]] - logits[:, i, token_id - 2] = 10.0 - return logits.requires_grad_(True) - - -def build_gradient_mask(sequence: str, batch_size: int) -> torch.Tensor: - mask = torch.ones([batch_size, len(sequence), AA_DIMS]) - fixed_positions = [i for i, aa in enumerate(sequence) if aa != MUTABLE_TOKEN] - mask[:, fixed_positions, :] = 0.0 - mask[:, :, CYS_IDX] = 0.0 - return mask - - -def sequence_to_one_hot(sequence: str, device: torch.device | str = "cuda") -> torch.Tensor: - target_index = [TOKEN_IDS[PROTEIN_1TO3[letter]] for letter in sequence] - one_hot = F.one_hot(torch.tensor(target_index), num_classes=len(TOKENS)) - return one_hot.to(device).unsqueeze(0).float() - - -def get_mid_points() -> torch.Tensor: - boundaries = torch.linspace(2, 52.0, 127) - lower = torch.tensor([1.0]) - upper = torch.tensor([57.0]) - exp_boundaries = torch.cat((lower, boundaries, upper)) - return (exp_boundaries[:-1] + exp_boundaries[1:]) / 2 - - -def binned_entropy( - dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float -) -> torch.Tensor: - bin_mask = ~(bin_distance < cutoff) - masked_dgram = dgram - (1e7 * bin_mask) - px = torch.softmax(masked_dgram, dim=-1) - log_px = torch.log_softmax(dgram, dim=-1) - return -(px * log_px).sum(-1) - - -def masked_min_k(x: torch.Tensor, mask: torch.Tensor, k: int) -> torch.Tensor: - mask = mask.bool() - y = torch.sort(torch.where(mask, x, float("nan")))[0] - k_mask = (torch.arange(y.shape[-1]).to(y.device) < k) & (~torch.isnan(y)) - return torch.where(k_mask, y, 0).sum(-1) / (k_mask.sum(-1) + 1e-8) - - -def masked_average(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - mask = mask.bool() - return torch.where(mask, x, 0).sum(-1) / (torch.where(mask, 1, 0).sum(-1) + 1e-8) - - -def compute_contact_loss( - distogram_logits: torch.Tensor, - bin_distance: torch.Tensor, - num_contacts: int, - min_sep: int, - cutoff: float, - chain_mask: torch.Tensor, - binder_mask: torch.Tensor, -) -> torch.Tensor: - con_loss = binned_entropy(distogram_logits, bin_distance, cutoff) - position = torch.arange(distogram_logits.shape[1]) - p_dist = position[:, None] - position[None, :] - if min_sep > 0: - separation_mask = (torch.abs(p_dist) >= min_sep).to(distogram_logits.device) - binder_mask = torch.logical_and(separation_mask, binder_mask) - per_residue = masked_min_k(con_loss, mask=binder_mask, k=num_contacts).to( - distogram_logits.device - ) - return masked_average(per_residue, mask=chain_mask).to(distogram_logits.device) - - -def compute_intra_contact_loss( - distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor -) -> torch.Tensor: - full_len = distogram_logits.shape[1] - is_binder = torch.ones(full_len, device=distogram_logits.device) - is_binder[:-binder_length] *= 0.0 - return compute_contact_loss( - distogram_logits, - bin_distance, - num_contacts=2, - min_sep=9, - cutoff=14.0, - chain_mask=is_binder, - binder_mask=is_binder, - ) - - -def compute_inter_contact_loss( - distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor -) -> torch.Tensor: - full_len = distogram_logits.shape[1] - is_binder = torch.ones(full_len, device=distogram_logits.device) - is_binder[:-binder_length] *= 0.0 - return compute_contact_loss( - distogram_logits, - bin_distance, - num_contacts=1, - min_sep=0, - cutoff=22.0, - chain_mask=1 - is_binder, - binder_mask=is_binder, - ) - - -def compute_globularity_loss( - distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor -) -> torch.Tensor: - binder_disto = distogram_logits[:, -binder_length:, -binder_length:, :] - n = binder_disto.shape[1] - disto_probs = torch.softmax(binder_disto, dim=-1) - bin_distance = bin_distance.clamp(max=27) - e_sq_dist = torch.sum(disto_probs * torch.square(bin_distance), dim=-1) - sum_sq_dist = torch.sum(torch.tril(e_sq_dist, diagonal=-1), dim=(1, 2)) - rg_term = torch.sqrt(sum_sq_dist / (n * n)) - rg_th = 2.38 * (n**0.365) - return F.elu(rg_term - rg_th) - - -def compute_structure_losses( - distogram_logits: torch.Tensor, binder_length: int -) -> dict[str, torch.Tensor]: - bin_distance = get_mid_points().to(distogram_logits.device) - losses: dict[str, torch.Tensor] = {} - losses["intra_contact_loss"] = compute_intra_contact_loss( - distogram_logits, binder_length, bin_distance - ) - losses["inter_contact_loss"] = compute_inter_contact_loss( - distogram_logits, binder_length, bin_distance - ) - losses["glob_loss"] = compute_globularity_loss( - distogram_logits, binder_length, bin_distance - ) - batch = distogram_logits.size(0) - total = torch.tensor([0.0] * batch, device=distogram_logits.device) - total = total + LOSS_WEIGHTS["intra_contact"] * losses["intra_contact_loss"] - total = total + LOSS_WEIGHTS["inter_contact"] * losses["inter_contact_loss"] - total = total + LOSS_WEIGHTS["glob"] * losses["glob_loss"] - losses["total_loss"] = total - return losses - - -def _binding_confidence_entropy( - dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float -) -> torch.Tensor: - probs = torch.softmax(dgram, dim=-1) - cutoff_mask = bin_distance < cutoff - p_cut = probs[..., cutoff_mask] - p_cut = p_cut / (p_cut.sum(-1, keepdim=True) + 1e-8) - return -(p_cut * torch.log(p_cut + 1e-10)).sum(-1) - - -def _entropy_to_confidence(mean_entropy: float) -> float: - return float(max(0.0, min(1.0, 1.0 - mean_entropy / math.log(51)))) - - -def _cdr_indices(binder_sequence: str) -> list[int]: - from abnumber import Chain - from abnumber.common import _anarci_align - - result = _anarci_align( - sequences=[binder_sequence], scheme="chothia", allowed_species=None - )[0] - chains = [ - Chain("".join(result[i][0].values()), scheme="chothia") - for i in range(len(result)) - ] - if len(chains) == 2 and not chains[0].is_heavy_chain(): - chains.reverse() - indices: list[int] = [] - for chain in chains: - for cdr in (chain.cdr1_seq, chain.cdr2_seq, chain.cdr3_seq): - start = binder_sequence.find(cdr) - assert start >= 0 - indices.extend(range(start, start + len(cdr))) - return indices - - -def compute_distogram_iptm_proxy( - distogram_logits: torch.Tensor, - target_length: int, - binder_sequence: str, - is_antibody: bool, - cdr_indices: list[int] | None = None, -) -> dict[str, float]: - if distogram_logits.ndim == 4: - distogram_logits = distogram_logits[0] - binder_length = len(binder_sequence) - assert distogram_logits.shape[0] == target_length + binder_length - - bin_distance = get_mid_points().to(distogram_logits.device) - binder_start = target_length - - def _mean_lowest_k(entropies: torch.Tensor, k: int) -> float: - sorted_entropies, _ = torch.sort(entropies.reshape(-1)) - k = min(k, sorted_entropies.numel()) - return float(sorted_entropies[:k].mean()) - - binder_to_target_entropy = _binding_confidence_entropy( - distogram_logits[binder_start:, :target_length, :], bin_distance, cutoff=22.0 - ) - distogram_iptm_proxy = _entropy_to_confidence( - _mean_lowest_k(binder_to_target_entropy, k=binder_length) - ) - - if not is_antibody: - cdr_distogram_iptm_proxy = float("nan") - else: - if cdr_indices is None: - cdr_indices = _cdr_indices(binder_sequence) - cdr_rows = [binder_start + i for i in cdr_indices] - cdr_to_target_entropy = _binding_confidence_entropy( - distogram_logits[cdr_rows, :target_length, :], bin_distance, cutoff=22.0 - ) - cdr_distogram_iptm_proxy = _entropy_to_confidence( - _mean_lowest_k(cdr_to_target_entropy, k=len(cdr_indices)) - ) - return { - "distogram_iptm_proxy": distogram_iptm_proxy, - "cdr_distogram_iptm_proxy": cdr_distogram_iptm_proxy, - } - - -_ATOM_FEATURE_DIMS = { - "ref_pos": 1, - "ref_element": 1, - "ref_charge": 1, - "ref_atom_name_chars": 1, - "ref_space_uid": 1, - "atom_attention_mask": 1, - "atom_to_token": 1, - "is_resolved": 1, - "gt_coords": 2, -} - - -def _resize_tensor(tensor: torch.Tensor, *, dim: int, size: int) -> torch.Tensor: - current = tensor.shape[dim] - if current >= size: - return tensor.narrow(dim, 0, size) - pad_shape = list(tensor.shape) - pad_shape[dim] = size - current - pad = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device) - return torch.cat((tensor, pad), dim=dim) - - -def prepare_esmfold2_tensors( - model: Any, - input_data: Any, - max_atoms: int | None = None, - seed: int | None = None, -) -> tuple[dict[str, torch.Tensor], list[Any]]: - features, chain_infos = model.prepare_structure_input(input_data, seed=seed) - if max_atoms is not None: - for key, dim in _ATOM_FEATURE_DIMS.items(): - if key in features: - features[key] = _resize_tensor(features[key], dim=dim, size=max_atoms) - return features, chain_infos - - -def _filter_model_forward_kwargs( - model: Any, kwargs: dict[str, torch.Tensor | int | bool | None] -) -> dict[str, torch.Tensor | int | bool | None]: - signature = inspect.signature(model.forward) - parameters = signature.parameters - accepts_kwargs = any( - parameter.kind == inspect.Parameter.VAR_KEYWORD - for parameter in parameters.values() - ) - if accepts_kwargs: - return kwargs - return {key: value for key, value in kwargs.items() if key in parameters} - - -def fold_and_get_distogram( - model: Any, - target_seq: str, - target_one_hot: torch.Tensor, - design: torch.Tensor, - num_loops: int = 0, - num_sampling_steps: int = 1, - calculate_confidence: bool = False, - seed: int | None = None, -) -> dict[str, Any]: - padding = (2, 11) - padded_design = F.pad(design, padding, mode="constant", value=0) - - token_lists = torch.argmax(padded_design, dim=-1) - designed_seq = [ - [PROTEIN_3TO1[TOKENS[int(tkn.item())]] for tkn in token_list] - for token_list in token_lists - ] - seq_list = [target_seq + "|" + "".join(seq) for seq in designed_seq] - max_atoms = None if len(seq_list) == 1 else ((len(seq_list[0]) - 1) * 14) // 32 * 32 - - inputs_list = [] - chain_info_list = [] - for seq in seq_list: - target, binder = seq.split("|") - input_types = model.input_types - inputs_raw = input_types.StructurePredictionInput( - sequences=[ - input_types.ProteinInput(id="A", sequence=target, msa=None), - input_types.ProteinInput(id="B", sequence=binder, msa=None), - ] - ) - features, chain_infos = prepare_esmfold2_tensors( - model, inputs_raw, max_atoms=max_atoms, seed=seed - ) - inputs_list.append(features) - chain_info_list.append(chain_infos) - - inputs = { - key: torch.cat([inp[key] for inp in inputs_list], dim=0).to(design.device) - for key in inputs_list[0] - } - inputs["res_type_soft"] = torch.cat( - (target_one_hot.repeat(design.size(0), 1, 1), padded_design), dim=1 - ) - - forward_kwargs: dict[str, torch.Tensor | int | bool | None] = dict(inputs) - forward_kwargs.update( - { - "num_diffusion_samples": 1, - "num_sampling_steps": num_sampling_steps, - "num_loops": num_loops, - "calculate_confidence": calculate_confidence, - "seed": seed, - } - ) - - with seed_context(seed): - output = model(**_filter_model_forward_kwargs(model, forward_kwargs)) - - result: dict[str, Any] = { - "distogram_logits": output["distogram_logits"], - "inputs": inputs, - "chain_info_list": chain_info_list, - "output": output, - "seq_list": seq_list, - } - if calculate_confidence: - for key in ("ptm", "iptm", "plddt"): - if key in output: - result[key] = output[key] - return result - - -@cache -def _folding_trunk_to_lm_aa_vocab_matrix(device: torch.device) -> torch.Tensor: - three_to_one_map = {v: k for k, v in PROTEIN_1TO3.items()} - ft_aas = [three_to_one_map[tok_3letter] for tok_3letter in TOKENS[2:22]] - tokenizer = EsmSequenceTokenizer() - lm_vocab = sorted(tokenizer.vocab.items(), key=lambda x: x[1]) - lm_aas = [lm_vocab[i][0] for i in range(4, 24)] - ft_to_lm_aa_matrix = torch.zeros(20, 20) - for ft_idx, ft_aa in enumerate(ft_aas): - lm_idx = lm_aas.index(ft_aa) - ft_to_lm_aa_matrix[ft_idx, lm_idx] = 1 - return ft_to_lm_aa_matrix.to(device=device) - - -def _one_hot_from_probs(probs: torch.Tensor) -> torch.Tensor: - return F.one_hot(torch.argmax(probs, dim=-1), num_classes=probs.size(-1)).to( - probs.dtype - ) - - -def _straight_through(discrete: torch.Tensor, continuous: torch.Tensor) -> torch.Tensor: - return continuous + (discrete - continuous).detach() - - -def compute_fastplms_pseudoperplexity_nll( - lm_model: Any, - binder_design: torch.Tensor, - score_mask: torch.Tensor, - batch_size: int = 4, - n_passes: int = 4, - mask_fraction: float = DEFAULT_ESMC_MASK_FRACTION, -) -> torch.Tensor: - device = binder_design.device - lm_vocab_size = lm_model.config.vocab_size - model_dtype = lm_model.embed.weight.dtype - - target_esm = binder_design @ _folding_trunk_to_lm_aa_vocab_matrix(device) - input_esm = _straight_through(_one_hot_from_probs(target_esm), target_esm) - input_ids = torch.zeros( - (binder_design.size(0), binder_design.size(1) + 2, lm_vocab_size), - dtype=model_dtype, - device=device, - ) - tokenizer = lm_model.tokenizer - input_ids[:, 0, tokenizer.cls_token_id] = 1 - input_ids[:, -1, tokenizer.eos_token_id] = 1 - input_ids[:, 1:-1, 4:24] = input_esm.to(model_dtype) - - if score_mask.ndim == 1: - score_mask = score_mask.unsqueeze(0).expand(binder_design.size(0), -1) - elif score_mask.shape != binder_design.shape[:2]: - raise ValueError( - f"Expected score_mask with shape {(binder_design.size(0), binder_design.size(1))}, " - f"got {tuple(score_mask.shape)}" - ) - score_mask = score_mask.to(device=device, dtype=torch.bool) - - mask_token = torch.zeros(lm_vocab_size, dtype=model_dtype, device=device) - mask_token[tokenizer.mask_token_id] = 1 - losses = [] - for batch_idx in range(binder_design.size(0)): - position_indices = score_mask[batch_idx].nonzero(as_tuple=False).flatten() - num_positions = int(position_indices.numel()) - if num_positions == 0: - raise ValueError("Pseudoperplexity score mask selected zero positions.") - - num_masked = max(1, math.ceil(mask_fraction * num_positions)) - random_scores = torch.rand((n_passes, num_positions), device=device) - masked_offsets = random_scores.topk(num_masked, dim=-1, largest=False).indices - pass_masks = torch.zeros( - (n_passes, binder_design.size(1)), dtype=torch.bool, device=device - ) - pass_masks[ - torch.arange(n_passes, device=device)[:, None], - position_indices[masked_offsets], - ] = True - - masked_sequences = input_ids[batch_idx : batch_idx + 1].repeat(n_passes, 1, 1) - mask_rows, mask_cols = pass_masks.nonzero(as_tuple=True) - masked_sequences[mask_rows, mask_cols + 1] = mask_token - - target_weights = target_esm[batch_idx] - masked_nlls = [] - for start in range(0, n_passes, batch_size): - stop = min(start + batch_size, n_passes) - chunk = masked_sequences[start:stop] - with torch.autocast( - device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda" - ): - hidden = lm_model.transformer( - x=chunk @ lm_model.embed.weight.to(chunk.dtype), - attention_mask=None, - output_hidden_states=False, - output_attentions=False, - ).last_hidden_state - logits = lm_model.sequence_head(hidden) - log_probs = logits.log_softmax(dim=-1)[:, 1:-1, 4:24] - nlls = -(log_probs * target_weights.to(log_probs.dtype).unsqueeze(0)).sum( - dim=-1 - ) - masked_nlls.append(nlls[pass_masks[start:stop]]) - losses.append(torch.cat(masked_nlls, dim=0).mean()) - return torch.stack(losses, dim=0) - - -def normalized_gradient_tensor( - grad: torch.Tensor, gradient_mask: torch.Tensor -) -> torch.Tensor: - masked_grad = grad * gradient_mask - index_has_nonzero_grad = torch.square(masked_grad).sum(-1) > 0 - eff_l = index_has_nonzero_grad.sum(-1) - grad_norm = torch.linalg.norm(masked_grad, axis=(-1, -2)) - normalized_grad = (masked_grad / (grad_norm[:, None, None] + 1e-7)) * torch.sqrt( - eff_l[:, None, None] - ) - return normalized_grad * gradient_mask - - -def _tensor_mean_float(tensor: torch.Tensor) -> float: - return float(tensor.detach().float().mean().cpu().item()) - - -def _metric_float(output: dict[str, Any], key: str) -> float | None: - if key not in output: - return None - value = output[key] - if value is None: - return None - if isinstance(value, torch.Tensor): - return float(value.detach().float().mean().cpu().item()) - return float(value) - - -def design_binder( - inversion_models: dict[str, Any], - critic_models: dict[str, Any], - lm_model: Any, - target_name: str | None, - target_sequence: str | None, - binder_name: str | None, - binder_sequence: str | None, - is_antibody: bool | None, - seed: int, - batch_size: int = 1, - steps: int = DEFAULT_STEPS, - log_interval: int = DEFAULT_LOG_INTERVAL, - learning_rate: float = DEFAULT_LEARNING_RATE, - temperature_min: float = DEFAULT_TEMPERATURE_MIN, - output_dir: str | Path | None = None, -) -> tuple[list[str], dict[int, dict[str, torch.Tensor]], list[dict[str, Any]]]: - assert (target_name is None) ^ ( - target_sequence is None - ), "Provide either target name or target sequence." - assert (binder_name is None) ^ ( - binder_sequence is None - ), "Provide either binder name or binder sequence." - - device = torch.device("cuda") - if target_name is not None: - assert target_name in TARGET_SEQUENCES, target_name - target_sequence = TARGET_SEQUENCES[target_name] - else: - assert target_sequence is not None - target_one_hot = sequence_to_one_hot(target_sequence, device=device) - - if binder_name is None: - assert binder_sequence is not None - if is_antibody is None: - is_antibody = False - else: - assert binder_name in BINDER_PROMPT_FACTORIES, binder_name - binder_prompt_factory = BINDER_PROMPT_FACTORIES[binder_name] - if is_antibody is not None: - assert binder_prompt_factory.is_antibody == is_antibody - is_antibody = binder_prompt_factory.is_antibody - binder_sequence = binder_prompt_factory.sample(seed=seed) - assert binder_sequence is not None - assert is_antibody is not None - mutable_binder_indices = [ - i for i, aa in enumerate(binder_sequence) if aa == MUTABLE_TOKEN - ] - binder_length = len(binder_sequence) - assert "|" not in target_sequence - assert "|" not in binder_sequence - - with seed_context(seed), torch.device(device): - logits = build_initial_soft_sequence_logits( - binder_sequence, batch_size=batch_size - ) - gradient_mask = build_gradient_mask(binder_sequence, batch_size=batch_size) - logits = logits.to(device) - gradient_mask = gradient_mask.to(device) - - trajectory: dict[int, dict[str, torch.Tensor]] = {} - optimizer = optim.SGD([logits], lr=learning_rate) - best_iptm: list[float] = [-1.0] * batch_size - best_loss: list[float] = [float("inf")] * batch_size - best_sequences: list[str] = [""] * batch_size - model_names = list(inversion_models) - - progress = tqdm(range(steps), desc="design", dynamic_ncols=True) - for step in progress: - optimizer.zero_grad() - t = (step + 1) / steps - remaining = 0.5 * (1 + math.cos(math.pi * t)) - temperature = temperature_min + (1 - temperature_min) * remaining - - random.seed(seed + step) - replicate_choice = random.randint(0, len(model_names) - 1) - inversion_model = inversion_models[model_names[replicate_choice]] - design = F.softmax(logits / temperature, dim=-1) - calculate_confidence = temperature < 0.05 - - fold_result = fold_and_get_distogram( - inversion_model, - target_sequence, - target_one_hot, - design, - num_loops=1, - num_sampling_steps=50 if calculate_confidence else 1, - calculate_confidence=calculate_confidence, - seed=seed + step, - ) - sequences: list[str] = fold_result["seq_list"] - losses = compute_structure_losses( - fold_result["distogram_logits"], binder_length - ) - structure_loss = losses["total_loss"] - structure_grad = torch.autograd.grad(structure_loss.mean(), logits)[0] - - design = F.softmax(logits / temperature, dim=-1) - score_mask = gradient_mask.sum(dim=-1) > 0 - with seed_context(seed + step): - plm_loss = compute_fastplms_pseudoperplexity_nll( - lm_model=lm_model, - binder_design=design, - score_mask=score_mask, - batch_size=4, - n_passes=4, - ) - plm_grad = torch.autograd.grad(plm_loss.mean(), logits)[0] - - logits.grad = normalized_gradient_tensor(structure_grad, gradient_mask) + ( - 0.05 if is_antibody else 0.15 - ) * normalized_gradient_tensor(plm_grad, gradient_mask) - for group in optimizer.param_groups: - group["lr"] = learning_rate * temperature - optimizer.step() - - step_losses = {key: value.detach().cpu() for key, value in losses.items()} - step_losses["plm_loss"] = plm_loss.detach().cpu() - step_losses["total_loss"] = (structure_loss + plm_loss).detach().cpu() - trajectory[step] = step_losses - - iptm = fold_result["iptm"] if "iptm" in fold_result else None - for batch_idx in range(batch_size): - current_loss = float(step_losses["total_loss"][batch_idx].item()) - if iptm is not None and iptm[batch_idx] is not None: - current_iptm = float(iptm[batch_idx].item()) - if current_iptm > best_iptm[batch_idx]: - best_iptm[batch_idx] = current_iptm - best_sequences[batch_idx] = sequences[batch_idx] - best_loss[batch_idx] = current_loss - elif current_loss < best_loss[batch_idx]: - best_sequences[batch_idx] = sequences[batch_idx] - best_loss[batch_idx] = current_loss - - if step % log_interval == 0: - loss_str = " ".join( - f"{key}={_tensor_mean_float(value):.4f}" - for key, value in step_losses.items() - ) - logger.info("step %3d | %s T=%.4f", step, loss_str, temperature) - progress.set_postfix( - loss=f"{_tensor_mean_float(step_losses['total_loss']):.3f}", - temp=f"{temperature:.3f}", - ) - - assert all(seq != "" for seq in best_sequences) - result_dir = Path(output_dir) if output_dir is not None else None - if result_dir is not None: - result_dir.mkdir(parents=True, exist_ok=True) - _write_trajectory(result_dir / "trajectory.jsonl", trajectory) - _write_fasta(result_dir / "best_sequences.fasta", best_sequences) - - critic_results: list[dict[str, Any]] = [] - target_length = len(target_sequence.replace("|", "")) - for batch_idx, best_seq in enumerate(best_sequences): - binder_seq = best_seq.split("|")[-1] - binder_design = sequence_to_one_hot(binder_seq, device=device)[..., 2:22] - for critic_name, critic_model in critic_models.items(): - is_scaling_critic = _is_scaling_critic_name(critic_name) - if is_scaling_critic: - critic_model.to(device=device) - try: - final_fold = fold_and_get_distogram( - critic_model, - target_sequence, - target_one_hot, - binder_design, - num_loops=3, - num_sampling_steps=200, - calculate_confidence=True, - seed=seed, - ) - finally: - if is_scaling_critic: - critic_model.to(device="cpu") - final_output = final_fold["output"] - final_inputs = final_fold["inputs"] - chain_infos = final_fold["chain_info_list"][0] - complex_result = critic_model.input_builder.decode( - final_output, - final_inputs, - chain_infos, - num_diffusion_samples=1, - complex_id=f"{critic_name}-{batch_idx}", - ) - cif_text = critic_model.result_to_cif(complex_result) - pdb_text = critic_model.result_to_pdb(complex_result) - iptm_proxy_scores = compute_distogram_iptm_proxy( - final_fold["distogram_logits"], - target_length, - binder_seq, - is_antibody, - cdr_indices=mutable_binder_indices if is_antibody else None, - ) - iptm_value = None - if "iptm" in final_fold: - iptm_value = float(final_fold["iptm"][0].item()) - ptm_value = _metric_float(final_fold, "ptm") - mean_plddt = _metric_float(final_fold, "plddt") - - structure_stem = f"batch{batch_idx}_{critic_name.replace('/', '_')}" - logits_path = None - if result_dir is not None: - cif_path = result_dir / f"{structure_stem}.cif" - pdb_path = result_dir / f"{structure_stem}.pdb" - logits_path_obj = result_dir / f"{structure_stem}_logits.pt" - cif_path.write_text(cif_text, encoding="utf-8") - pdb_path.write_text(pdb_text, encoding="utf-8") - torch.save(logits[batch_idx].detach().cpu(), logits_path_obj) - logits_path = str(logits_path_obj) - - row = { - "is_antibody": is_antibody, - "critic_name": critic_name, - "batch_idx": batch_idx, - "designed_sequence": best_seq, - "binder_sequence": binder_seq, - "target_length": target_length, - "binder_length": len(binder_seq), - "final_loss": float(trajectory[steps - 1]["total_loss"][batch_idx].item()), - "ptm": ptm_value, - "iptm": iptm_value, - "mean_plddt": mean_plddt, - "pdb": pdb_text, - "cif": cif_text, - "logits_path": logits_path, - } - row.update(iptm_proxy_scores) - critic_results.append(row) - - if result_dir is not None: - _write_results_table(result_dir / "results.parquet", critic_results) - _write_official_selection_table(result_dir / "selection.parquet", critic_results) - return best_sequences, trajectory, critic_results - - -def _write_trajectory( - path: Path, trajectory: dict[int, dict[str, torch.Tensor]] -) -> None: - with path.open("w", encoding="utf-8") as handle: - for step, losses in trajectory.items(): - row = {"step": step} - for key, value in losses.items(): - row[key] = [float(x) for x in value.reshape(-1).tolist()] - handle.write(json.dumps(row) + "\n") - - -def _write_fasta(path: Path, sequences: list[str]) -> None: - with path.open("w", encoding="utf-8") as handle: - for idx, sequence in enumerate(sequences): - handle.write(f">design_{idx}\n{sequence}\n") - - -def _write_results_table(path: Path, rows: list[dict[str, Any]]) -> None: - import pandas as pd - - pd.DataFrame(rows).to_parquet(path, index=False) - - -def _binder_sequence_from_designed_sequence(designed_sequence: str) -> str: - parts = designed_sequence.split("|") - assert len(parts) == 2, designed_sequence - return parts[1] - - -def _compute_isoelectric_points(sequences: list[str]) -> list[float]: - from Bio.SeqUtils.ProtParam import ProteinAnalysis - - return [float(ProteinAnalysis(sequence).isoelectric_point()) for sequence in sequences] - - -def annotate_official_selection_scores(result_df: Any) -> Any: - """Add the official binder-design selection components to critic rows. - - Mirrors the paper Appendix A.3.1.2 and the official notebook selection cell: - minibinders with pI >= 6 are filtered, hero critics contribute mean iPTM, - and optional scaling critics contribute the distogram ipTM proxy. - """ - import pandas as pd - - if isinstance(result_df, pd.DataFrame): - df = result_df.copy() - else: - df = pd.DataFrame(result_df) - required_columns = [ - "critic_name", - "designed_sequence", - "is_antibody", - "iptm", - "distogram_iptm_proxy", - "cdr_distogram_iptm_proxy", - ] - missing_columns = [column for column in required_columns if column not in df.columns] - assert not missing_columns, f"Missing selection columns: {missing_columns}" - - binder_sequences = [ - _binder_sequence_from_designed_sequence(sequence) - for sequence in df["designed_sequence"].tolist() - ] - is_antibody = df["is_antibody"].astype(bool) - is_scaling = df["critic_name"].str.contains( - SCALING_CHECKPOINT_SUBSTRING, regex=False, na=False - ) - iptm_proxy = df["distogram_iptm_proxy"].where( - ~is_antibody, df["cdr_distogram_iptm_proxy"] - ) - - df["binder_sequence"] = binder_sequences - df["isoelectric_point"] = _compute_isoelectric_points(binder_sequences) - df["passes_official_pi_filter"] = is_antibody | df["isoelectric_point"].lt( - MINIBINDER_PI_CUTOFF - ) - df["official_iptm_score_component"] = df["iptm"].where(~is_scaling) - df["official_iptm_proxy_component"] = iptm_proxy.where(is_scaling) - return df - - -def select_official_designs( - result_df: Any, - top_k: int = DEFAULT_SELECTION_TOP_K, - consensus_iptm_threshold: float = DEFAULT_CONSENSUS_IPTM_THRESHOLD, - group_columns: tuple[str, ...] = ("target_name", "binder_name"), -) -> Any: - """Rank candidates using the official ESM binder-design selection strategy.""" - df = annotate_official_selection_scores(result_df) - available_group_columns = [ - column for column in group_columns if column in df.columns - ] - selection_columns = available_group_columns + [ - "designed_sequence", - "iptm_score", - "iptm_proxy_score", - "hero_iptm_min", - "hero_iptm_median", - "hero_iptm_max", - "scaling_proxy_mean", - "critic_count", - "hero_critic_count", - "scaling_critic_count", - "batch_idx", - "binder_sequence", - "is_antibody", - "isoelectric_point", - "selection_score", - "all_hero_critics_pass", - "consensus_iptm_threshold", - ] - df = df[df["passes_official_pi_filter"]].copy() - if df.empty: - import pandas as pd - - return pd.DataFrame(columns=selection_columns) - - key_columns = available_group_columns + ["designed_sequence"] - summary_columns = [ - "batch_idx", - "binder_sequence", - "is_antibody", - "isoelectric_point", - ] - summary_aggregations = { - column: (column, "first") for column in summary_columns if column in df.columns - } - scores = df.groupby(key_columns, as_index=False).agg( - iptm_score=("official_iptm_score_component", "mean"), - iptm_proxy_score=("official_iptm_proxy_component", "mean"), - hero_iptm_min=("official_iptm_score_component", "min"), - hero_iptm_median=("official_iptm_score_component", "median"), - hero_iptm_max=("official_iptm_score_component", "max"), - scaling_proxy_mean=("official_iptm_proxy_component", "mean"), - critic_count=("critic_name", "count"), - hero_critic_count=( - "official_iptm_score_component", - lambda values: int(values.notna().sum()), - ), - scaling_critic_count=( - "official_iptm_proxy_component", - lambda values: int(values.notna().sum()), - ), - **summary_aggregations, - ) - scores["selection_score"] = 0.5 * scores["iptm_score"].fillna( - 0.0 - ) + 0.5 * scores["iptm_proxy_score"].fillna(0.0) - scores["all_hero_critics_pass"] = scores["hero_iptm_min"].gt( - consensus_iptm_threshold - ) - scores["consensus_iptm_threshold"] = consensus_iptm_threshold - - if available_group_columns: - sort_columns = available_group_columns + ["selection_score"] - ascending = [True] * len(available_group_columns) + [False] - scores = scores.sort_values(sort_columns, ascending=ascending) - return ( - scores.groupby(available_group_columns, group_keys=False, sort=False) - .head(top_k) - .reset_index(drop=True) - ) - return scores.nlargest(min(len(scores), top_k), "selection_score").reset_index( - drop=True - ) - - -def _write_official_selection_table(path: Path, rows: list[dict[str, Any]]) -> None: - selection_df = select_official_designs(rows) - selection_df.to_parquet(path, index=False) - - -def _log_official_selection_summary(rows: list[dict[str, Any]]) -> None: - selection_df = select_official_designs(rows) - if selection_df.empty: - logger.info("Official selection table is empty after pI filtering") - return - top = selection_df.iloc[0] - logger.info( - "Top official selection | score=%.4f hero_mean=%.4f proxy_mean=%.4f " - "hero_min=%.4f all_hero_pass=%s binder=%s", - float(top["selection_score"]), - float(top["iptm_score"]), - float(top["iptm_proxy_score"]) if not math.isnan(top["iptm_proxy_score"]) else 0.0, - float(top["hero_iptm_min"]), - bool(top["all_hero_critics_pass"]), - top["binder_sequence"], - ) - - -_ESMC_CACHE: Any | None = None - - -def _load_fold_model( - model_name: str, - lm_dropout: float, - cache_esmc: bool, - device: torch.device | str, - kernel_backend: str | None, - compile_model: bool, -) -> Any: - global _ESMC_CACHE - model = AutoModel.from_pretrained( - _repo_name(model_name), - trust_remote_code=True, - load_esmc=not cache_esmc, - dtype=torch.float32, - ) - if cache_esmc: - if _ESMC_CACHE is None: - model.load_esmc(model.config.esmc_id) - _ESMC_CACHE = model._esmc - else: - model._esmc = _ESMC_CACHE - model.configure_lm_dropout( - lm_dropout, force_lm_dropout_during_inference=True - ) - if kernel_backend is not None: - model.set_kernel_backend(kernel_backend) - if compile_model: - model.apply_torch_compile() - return model.to(device=device).eval().requires_grad_(False) - - -class FastPLMsBinderDesign: - lm_name = "Synthyra/ESMplusplus_6B" - inversion_model_names = [ - "ESMFold2-Experimental-Fast", - "ESMFold2-Experimental-Fast-Cutoff2025", - ] - hero_critic_model_names = [ - "ESMFold2-Experimental-Fast", - "ESMFold2-Experimental-Fast-Cutoff2025", - "ESMFold2-Experimental", - "ESMFold2-Experimental-Cutoff2025", - ] - - def load( - self, - use_scaling_critics: bool = False, - device: str = "cuda", - kernel_backend: str | None = None, - compile_model: bool = False, - ) -> None: - scaling_critic_names: list[str] = [] - if use_scaling_critics: - scaling_critic_names = [ - f"ESMFold2-Experimental-Fast-base{size}-step{step}k" - for size in ("300M", "600M", "6B") - for step in ("250", "500", "750", "1000", "1500") - ] - self.inversion_models = { - model_name: _load_fold_model( - model_name, - lm_dropout=0.5, - cache_esmc=True, - device=device, - kernel_backend=kernel_backend, - compile_model=compile_model, - ) - for model_name in self.inversion_model_names - } - self.critic_models = { - model_name: _load_fold_model( - model_name, - lm_dropout=0.25, - cache_esmc=True, - device=device, - kernel_backend=kernel_backend, - compile_model=compile_model, - ) - for model_name in self.hero_critic_model_names - } - for model_name in scaling_critic_names: - self.critic_models[model_name] = _load_fold_model( - model_name, - lm_dropout=0.25, - cache_esmc=False, - device="cpu", - kernel_backend=kernel_backend, - compile_model=False, - ) - self.lm_model = ( - AutoModelForMaskedLM.from_pretrained( - self.lm_name, trust_remote_code=True, dtype=torch.float32 - ) - .to(device=device) - .eval() - .requires_grad_(False) - ) - - def design( - self, - target_name: str | None = None, - target_sequence: str | None = None, - binder_name: str | None = None, - binder_sequence: str | None = None, - is_antibody: bool | None = None, - seed: int = 0, - batch_size: int = 1, - steps: int = DEFAULT_STEPS, - output_dir: str | None = None, - ) -> tuple[list[str], dict[int, dict[str, torch.Tensor]], list[dict[str, Any]]]: - return design_binder( - self.inversion_models, - self.critic_models, - self.lm_model, - target_name=target_name, - target_sequence=target_sequence, - binder_name=binder_name, - binder_sequence=binder_sequence, - is_antibody=is_antibody, - seed=seed, - batch_size=batch_size, - steps=steps, - output_dir=output_dir, - ) - - -def _build_modal_image(): - assert modal is not None, "Modal is not installed." - return ( - modal.Image.from_registry( - "nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04", add_python="3.12" - ) - .apt_install("git", "build-essential") - .uv_pip_install( - "torch==2.11.0", - "transformers==4.57.6", - "accelerate==1.12.0", - "hf-xet==1.5.0", - "huggingface_hub", - "numpy==1.26.4", - "einops==0.8.2", - "tokenizers", - "safetensors", - "pandas", - "pyarrow", - "biopython", - "tqdm", - "abnumber", - "biotite==1.6.0", - "rdkit==2026.3.2", - "msgpack-numpy==0.4.8", - "py3dmol==2.5.5", - index_url="https://download.pytorch.org/whl/cu128", - extra_index_url="https://pypi.org/simple", - ) - .add_local_python_source("fastplms") - .env({"HF_HOME": "/models", "HF_XET_HIGH_PERFORMANCE": "1"}) - ) - - -if modal is not None: - app = modal.App(name="fastplms-binder-design") - _MODAL_IMAGE = _build_modal_image() - _MODAL_MODEL_CACHE = modal.Volume.from_name( - "fastplms-binder-design-models", create_if_missing=True - ) - - @app.cls( - gpu="H100", - image=_MODAL_IMAGE, - volumes={"/models": _MODAL_MODEL_CACHE}, - timeout=60 * 60, - cpu=16, - memory=10 * 1024, - ) - class FastPLMsBinderDesignModal(FastPLMsBinderDesign): - use_scaling_critics: bool = modal.parameter(default=False) - kernel_backend: str | None = modal.parameter(default=None) - compile_model: bool = modal.parameter(default=False) - - @modal.enter() - def load(self) -> None: - super().load( - use_scaling_critics=self.use_scaling_critics, - kernel_backend=self.kernel_backend, - compile_model=self.compile_model, - ) - - @modal.method() - def design(self, *args, **kwargs): - return super().design(*args, **kwargs) - - @app.local_entrypoint() - def modal_main( - target_name: str | None = "pd-l1", - target_sequence: str | None = None, - binder_name: str | None = "minibinder", - binder_sequence: str | None = None, - use_scaling_critics: bool = False, - is_antibody: bool | None = None, - seed: int = 0, - batch_size: int = 1, - steps: int = DEFAULT_STEPS, - output_dir: str | None = "binder_design_out", - ) -> None: - remote_app = FastPLMsBinderDesignModal( - use_scaling_critics=use_scaling_critics - ) - best_sequences, _, results = remote_app.design.remote( - target_name=target_name, - target_sequence=target_sequence, - binder_name=binder_name, - binder_sequence=binder_sequence, - is_antibody=is_antibody, - seed=seed, - batch_size=batch_size, - steps=steps, - output_dir=output_dir, - ) - logger.info("Designed sequences: %s", best_sequences) - logger.info("Returned %d critic rows", len(results)) - _log_official_selection_summary(results) - - -def _design_kwargs_from_args(args: argparse.Namespace) -> dict[str, Any]: - return { - "target_name": args.target_name, - "target_sequence": args.target_sequence, - "binder_name": args.binder_name, - "binder_sequence": args.binder_sequence, - "is_antibody": args.is_antibody, - "seed": args.seed, - "batch_size": args.batch_size, - "steps": args.steps, - "output_dir": args.output_dir, - } - - -def run_local(args: argparse.Namespace) -> None: - runner = FastPLMsBinderDesign() - runner.load( - use_scaling_critics=args.use_scaling_critics, - kernel_backend=args.kernel_backend, - compile_model=args.compile_model, - ) - best_sequences, _, results = runner.design(**_design_kwargs_from_args(args)) - logger.info("Designed sequences: %s", best_sequences) - logger.info("Returned %d critic rows", len(results)) - _log_official_selection_summary(results) - - -def run_deployed_modal(args: argparse.Namespace) -> None: - assert modal is not None, "Modal is not installed." - remote_cls = modal.Cls.from_name(args.modal_app_name, "FastPLMsBinderDesignModal") - remote_runner = remote_cls( - use_scaling_critics=args.use_scaling_critics, - kernel_backend=args.kernel_backend, - compile_model=args.compile_model, - ) - best_sequences, _, results = remote_runner.design.remote( - **_design_kwargs_from_args(args) - ) - logger.info("Designed sequences: %s", best_sequences) - logger.info("Returned %d critic rows", len(results)) - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - _write_results_table(output_dir / "results.parquet", results) - _write_official_selection_table(output_dir / "selection.parquet", results) - _log_official_selection_summary(results) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--backend", choices=["local", "modal"], default="local") - parser.add_argument("--modal-app-name", default="fastplms-binder-design") - parser.add_argument("--target-name", default="pd-l1") - parser.add_argument("--target-sequence", default=None) - parser.add_argument("--binder-name", default="minibinder") - parser.add_argument("--binder-sequence", default=None) - parser.add_argument("--use-scaling-critics", action="store_true") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--steps", type=int, default=DEFAULT_STEPS) - parser.add_argument("--output-dir", default="binder_design_out") - parser.add_argument("--kernel-backend", default=None) - parser.add_argument("--compile-model", action="store_true") - parser.add_argument("--is-antibody", dest="is_antibody", action="store_true") - parser.add_argument("--not-antibody", dest="is_antibody", action="store_false") - parser.set_defaults(is_antibody=None) - args = parser.parse_args() - if args.target_sequence is not None: - args.target_name = None - if args.binder_sequence is not None: - args.binder_name = None - return args - - -if __name__ == "__main__": - cli_args = parse_args() - if cli_args.backend == "local": - run_local(cli_args) - else: - run_deployed_modal(cli_args) diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..3263461 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,432 @@ +# syntax=docker/dockerfile:1.19@sha256:b6afd42430b15f2d2a4c5a02b919e98a525b785b1aaff16747d2f623364e39b6 + +# Candidate validation is pinned to CUDA 13.0. The multi-platform digest is +# intentional: Buildx resolves the platform-specific image beneath it. +ARG CUDA13_IMAGE=nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04@sha256:8b2705ea7a8653ad3451b46ab835eced92d77b44e671b9cf3ad4f95fbb2efe5e +ARG CUDA12_LEGACY_IMAGE=nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04@sha256:21196d81f56b48dbee70494d5f10322e1a77cc47ffe202a3bf68eab81533c20f + +FROM ${CUDA13_IMAGE} AS python312 + +ENV DEBIAN_FRONTEND=noninteractive \ + PATH=/opt/venv/bin:$PATH \ + LD_LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/nvidia/cudnn/lib:/usr/local/cuda/lib64:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 \ + UV_CACHE_DIR=/root/.cache/uv \ + UV_COMPILE_BYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_CACHE_DIR=/root/.cache/pip \ + TOKENIZERS_PARALLELISM=true \ + HF_HOME=/cache/huggingface \ + TORCH_HOME=/cache/torch \ + XDG_CACHE_HOME=/cache/xdg \ + FASTPLMS_ARTIFACTS=/workspace/artifacts + +# PyTorch 2.13 pins its cuDNN runtime as a wheel. Keep that locked runtime ahead +# of the CUDA base image so Transformer Engine behaves identically regardless +# of whether a BF16 cuDNN kernel ran before the first FP8 reload. + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ + apt-get update && apt-get install -y --no-install-recommends \ + build-essential ca-certificates git git-lfs ninja-build \ + python3.12 python3.12-dev python3.12-venv && \ + python3.12 -m venv /opt/venv && \ + /opt/venv/bin/pip install --upgrade pip==26.1.1 uv==0.10.12 + +WORKDIR /opt/fastplms + +FROM python312 AS dependency-files + +# Keep dependency-only layers independent from runtime source, tests, docs, and +# legal text. A source edit must not rebuild the CUDA/Python environment. +COPY requirements ./requirements +COPY kernels.lock ./ + +FROM dependency-files AS source-dependencies + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/runtime.in \ + --constraint requirements/constraints/validation.txt + +FROM source-dependencies AS source + +COPY README.md LICENSE THIRD_PARTY_NOTICES.md ./ +COPY LICENSES ./LICENSES +COPY src ./src +ENV PYTHONPATH=/opt/fastplms/src + +FROM dependency-files AS runtime-dependencies + +ARG FASTPLMS_RUNTIME_PROFILE=core +RUN --mount=type=cache,target=/root/.cache/uv \ + case "${FASTPLMS_RUNTIME_PROFILE}" in \ + core) \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/runtime.in \ + --constraint requirements/constraints/validation.txt ;; \ + esmfold2-fp8) \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/runtime-fp8.in \ + --constraint requirements/constraints/validation.txt \ + --overrides requirements/overrides/cuda.txt ;; \ + *) \ + echo "Unsupported FastPLMs runtime profile: ${FASTPLMS_RUNTIME_PROFILE}" >&2; \ + exit 64 ;; \ + esac + +FROM runtime-dependencies AS runtime + +COPY README.md LICENSE THIRD_PARTY_NOTICES.md ./ +COPY LICENSES ./LICENSES +COPY src ./src +ENV PYTHONPATH=/opt/fastplms/src +WORKDIR /workspace +CMD ["python"] + +FROM dependency-files AS candidate-dependencies + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/candidate.in \ + --constraint requirements/constraints/validation.txt + +FROM candidate-dependencies AS candidate + +COPY README.md LICENSE THIRD_PARTY_NOTICES.md ./ +COPY LICENSES ./LICENSES +COPY src ./src +COPY tests /opt/fastplms/tests +COPY benchmarks /opt/fastplms/benchmarks +COPY tools /opt/fastplms/tools +ENV PYTHONPATH=/workspace/src:/workspace:/opt/fastplms/src:/opt/fastplms +WORKDIR /workspace +CMD ["python", "-m", "pytest", "tests/unit", "tests/integration", "-m", "not gpu and not slow and not structure"] + +FROM dependency-files AS candidate-structure-dependencies + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/candidate-structure.in \ + --constraint requirements/constraints/validation.txt + +FROM candidate-structure-dependencies AS candidate-structure + +COPY README.md LICENSE THIRD_PARTY_NOTICES.md ./ +COPY LICENSES ./LICENSES +COPY src ./src +COPY tests /opt/fastplms/tests +COPY benchmarks /opt/fastplms/benchmarks +COPY tools /opt/fastplms/tools +ENV PYTHONPATH=/workspace/src:/workspace:/opt/fastplms/src:/opt/fastplms +WORKDIR /workspace +CMD ["python", "-m", "pytest", "tests/structure", "-m", "structure"] + +FROM dependency-files AS candidate-fp8-dependencies + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/candidate-fp8.in \ + --constraint requirements/constraints/validation.txt \ + --overrides requirements/overrides/cuda.txt + +FROM candidate-fp8-dependencies AS candidate-fp8 + +COPY README.md LICENSE THIRD_PARTY_NOTICES.md ./ +COPY LICENSES ./LICENSES +COPY src ./src +COPY tests /opt/fastplms/tests +COPY benchmarks /opt/fastplms/benchmarks +COPY tools /opt/fastplms/tools +ENV PYTHONPATH=/workspace/src:/workspace:/opt/fastplms/src:/opt/fastplms +WORKDIR /workspace +CMD ["python", "-m", "pytest", "tests/structure", "-m", "gpu and structure"] + +# The artifact tier intentionally installs dependencies without FastPLMs. A +# mounted `dist/hub/` must therefore carry all remote-code sources it +# needs, and the probe fails if an installed `fastplms` distribution is discoverable. +FROM dependency-files AS candidate-artifact + +COPY README.md LICENSE ./ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python /opt/venv/bin/python \ + --requirement requirements/profiles/artifact.in \ + --constraint requirements/constraints/validation.txt +ENV PYTHONPATH= +WORKDIR /workspace +CMD ["python", "-m", "pytest", "tests/release/test_published_automodel.py", "-m", "artifact"] + +FROM ${CUDA12_LEGACY_IMAGE} AS python310-reference + +ENV DEBIAN_FRONTEND=noninteractive \ + PATH=/opt/venv/bin:$PATH \ + PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_CACHE_DIR=/root/.cache/pip \ + HF_HOME=/cache/huggingface \ + TORCH_HOME=/cache/torch \ + XDG_CACHE_HOME=/cache/xdg \ + FASTPLMS_ARTIFACTS=/workspace/artifacts + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ + apt-get update && apt-get install -y --no-install-recommends \ + build-essential ca-certificates git python3.10 python3.10-dev python3.10-venv && \ + python3.10 -m venv /opt/venv && \ + /opt/venv/bin/pip install --upgrade pip==26.1.1 +WORKDIR /workspace + +# The native oracle harness is independent of the FastPLMs source. Reference +# stages receive this small protocol plus only their declared upstream source. +FROM scratch AS reference-protocol + +COPY tests/__init__.py /opt/oracle/tests/__init__.py +COPY tests/parity/__init__.py /opt/oracle/tests/parity/__init__.py +COPY tests/parity/support/__init__.py /opt/oracle/tests/parity/support/__init__.py +COPY tests/parity/support/state_transforms.py /opt/oracle/tests/parity/support/state_transforms.py +COPY tests/parity/support/semantic_config.py /opt/oracle/tests/parity/support/semantic_config.py +COPY tests/parity/support/native_reference.py /opt/oracle/tests/parity/support/native_reference.py +COPY tests/parity/support/reference_adapters /opt/oracle/tests/parity/support/reference_adapters +COPY tests/structure/__init__.py /opt/oracle/tests/structure/__init__.py +COPY tests/structure/support/__init__.py /opt/oracle/tests/structure/support/__init__.py +COPY tests/structure/support/state_contract.py /opt/oracle/tests/structure/support/state_contract.py +COPY tests/structure/support/boltz2_bundle.py /opt/oracle/tests/structure/support/boltz2_bundle.py +COPY tests/structure/support/esmfold_bundle.py /opt/oracle/tests/structure/support/esmfold_bundle.py +COPY tests/structure/support/esmfold2_bundle.py /opt/oracle/tests/structure/support/esmfold2_bundle.py +COPY tools/source_provenance.py /opt/oracle/tools/source_provenance.py +COPY tools/remote/biohub_reference_environment.py /opt/oracle/tools/remote/biohub_reference_environment.py +COPY tools/remote/biohub_reference_lock.py /opt/oracle/tools/remote/biohub_reference_lock.py +COPY tools/remote/biohub_reference_requirements.py /opt/oracle/tools/remote/biohub_reference_requirements.py +COPY tools/remote/reference_source_attestation.py /opt/oracle/tools/remote/reference_source_attestation.py + +FROM python310-reference AS reference-ankh + +COPY docker/constraints/ankh.txt /tmp/constraints.txt +COPY --from=upstream_ankh --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/ankh +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130 && \ + pip install --constraint /tmp/constraints.txt /opt/oracle/vendor/upstream/ankh \ + pytest==8.4.2 safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/ankh /licenses/ankh +ENV PYTHONPATH=/opt/oracle/vendor/upstream/ankh/src:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.parity.support.reference_adapters.ankh"] + +FROM python310-reference AS reference-dplm + +COPY docker/constraints/dplm.txt /tmp/constraints.txt +COPY --from=upstream_dplm --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/dplm +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.2.0 torchtext==0.17.0 --index-url https://download.pytorch.org/whl/cu121 && \ + pip install --constraint /tmp/constraints.txt --only-binary=:all: \ + --find-links https://data.pyg.org/whl/torch-2.2.0+cu121.html \ + torch-scatter==2.1.2+pt22cu121 && \ + pip install --constraint /tmp/constraints.txt -e /opt/oracle/vendor/upstream/dplm \ + safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/dplm /licenses/dplm +ENV PYTHONPATH=/opt/oracle/vendor/upstream/dplm/src:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.parity.support.reference_adapters.dplm"] + +FROM python312 AS reference-e1 + +COPY docker/constraints/e1.txt /tmp/constraints.txt +COPY --from=upstream_e1 --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/e1 +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128 && \ + pip install --constraint /tmp/constraints.txt -e /opt/oracle/vendor/upstream/e1 \ + pytest==8.4.2 safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/e1 /licenses/e1 +ENV PYTHONPATH=/opt/oracle/vendor/upstream/e1/src:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.parity.support.reference_adapters.e1"] + +FROM python312 AS reference-biohub-esm + +ARG TARGETARCH +RUN test "${TARGETARCH}" = "arm64" && test "$(uname -m)" = "aarch64" +COPY docker/biohub-reference-lock.Dockerfile /opt/lock-root/docker/biohub-reference-lock.Dockerfile +COPY docker/constraints/biohub-reference.in /opt/lock-root/docker/constraints/biohub-reference.in +COPY docker/constraints/biohub-reference.lock.txt /opt/lock-root/docker/constraints/biohub-reference.lock.txt +COPY docker/constraints/biohub-biotraj-build.in /opt/lock-root/docker/constraints/biohub-biotraj-build.in +COPY docker/constraints/biohub-biotraj-build.lock.txt /opt/lock-root/docker/constraints/biohub-biotraj-build.lock.txt +COPY docker/constraints/biohub-reference-lock.json /opt/lock-root/docker/constraints/biohub-reference-lock.json +COPY tools/remote/biohub_reference_lock.py /opt/lock-root/biohub_reference_lock.py +COPY --from=biohub_biotraj_wheel \ + /biohub-reference-lock/biotraj-1.2.2-cp312-cp312-linux_aarch64.whl \ + /opt/wheels/biotraj-1.2.2-cp312-cp312-linux_aarch64.whl +COPY --from=upstream_biohub_transformers --exclude=.git --exclude=.git/** --exclude=**/__pycache__ --exclude=**/__pycache__/** --exclude=**/*.pyc --exclude=**/*.pyo . /opt/oracle/vendor/upstream/biohub-transformers +COPY --from=upstream_biohub_esm --exclude=.git --exclude=.git/** --exclude=**/__pycache__ --exclude=**/__pycache__/** --exclude=**/*.pyc --exclude=**/*.pyo . /opt/oracle/vendor/upstream/biohub-esm +RUN --mount=type=cache,target=/root/.cache/pip \ + python /opt/lock-root/biohub_reference_lock.py verify-contract \ + --root /opt/lock-root \ + --contract /opt/lock-root/docker/constraints/biohub-reference-lock.json && \ + python /opt/lock-root/biohub_reference_lock.py materialize-wheel-lock \ + --root /opt/lock-root \ + --contract /opt/lock-root/docker/constraints/biohub-reference-lock.json \ + --wheel /opt/wheels/biotraj-1.2.2-cp312-cp312-linux_aarch64.whl \ + --wheel-uri file:///opt/wheels/biotraj-1.2.2-cp312-cp312-linux_aarch64.whl \ + --output /tmp/biohub-reference.install.lock.txt && \ + cp -a /opt/oracle/vendor/upstream/biohub-transformers /tmp/biohub-transformers-build && \ + cp -a /opt/oracle/vendor/upstream/biohub-esm /tmp/biohub-esm-build && \ + pip install --require-hashes --only-binary=:all: --no-deps \ + --requirement /tmp/biohub-reference.install.lock.txt && \ + pip install --no-deps --no-build-isolation \ + /tmp/biohub-transformers-build /tmp/biohub-esm-build && \ + python /opt/lock-root/biohub_reference_lock.py verify-pip-check \ + --root /opt/lock-root \ + --contract /opt/lock-root/docker/constraints/biohub-reference-lock.json && \ + python /opt/lock-root/biohub_reference_lock.py verify-inventory \ + --root /opt/lock-root \ + --contract /opt/lock-root/docker/constraints/biohub-reference-lock.json \ + --profile final && \ + rm -rf -- /tmp/biohub-transformers-build /tmp/biohub-esm-build +COPY docker/constraints/biohub-esm-source.json /opt/oracle/biohub-esm-source-contract.json +COPY docker/constraints/biohub-transformers-source.json /opt/oracle/biohub-transformers-source-contract.json +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/biohub-esm /licenses/biohub-esm +COPY LICENSES/biohub-transformers /licenses/biohub-transformers +RUN test ! -e /opt/oracle/tools/remote/__init__.py && \ + PYTHONPATH=/opt/oracle python -c \ + "import importlib.util; import tools.remote.biohub_reference_environment; import tools.remote.biohub_reference_lock; import tools.remote.reference_source_attestation; missing = importlib.util.find_spec('tools.remote.run') is None; raise SystemExit(0 if missing else 'Unexpected tools.remote.run in oracle image.')" +RUN PYTHONPATH=/opt/oracle python -m tools.remote.reference_source_attestation create \ + --source-root /opt/oracle/vendor/upstream/biohub-esm \ + --contract /opt/oracle/biohub-esm-source-contract.json \ + --output /opt/oracle/biohub-esm-source-attestation.json && \ + PYTHONPATH=/opt/oracle python -m tools.remote.reference_source_attestation create \ + --source-root /opt/oracle/vendor/upstream/biohub-transformers \ + --contract /opt/oracle/biohub-transformers-source-contract.json \ + --output /opt/oracle/biohub-transformers-source-attestation.json && \ + PYTHONPATH=/opt/oracle python -m tools.remote.reference_source_attestation verify \ + --source-root /opt/oracle/vendor/upstream/biohub-esm \ + --attestation /opt/oracle/biohub-esm-source-attestation.json \ + --contract /opt/oracle/biohub-esm-source-contract.json \ + --expected-revision 82ee35553d39169d678f784c8d3f8712ffd7d2c4 && \ + PYTHONPATH=/opt/oracle python -m tools.remote.reference_source_attestation verify \ + --source-root /opt/oracle/vendor/upstream/biohub-transformers \ + --attestation /opt/oracle/biohub-transformers-source-attestation.json \ + --contract /opt/oracle/biohub-transformers-source-contract.json \ + --expected-revision 3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf && \ + PYTHONPATH=/opt/oracle python -m tools.remote.biohub_reference_lock verify-inventory \ + --root /opt/lock-root \ + --contract /opt/lock-root/docker/constraints/biohub-reference-lock.json \ + --profile final +ENV FASTPLMS_BIOHUB_ESM_REVISION=82ee35553d39169d678f784c8d3f8712ffd7d2c4 \ + FASTPLMS_BIOHUB_ESM_SOURCE=/opt/oracle/vendor/upstream/biohub-esm \ + FASTPLMS_BIOHUB_ESM_ATTESTATION=/opt/oracle/biohub-esm-source-attestation.json \ + FASTPLMS_BIOHUB_ESM_CONTRACT=/opt/oracle/biohub-esm-source-contract.json \ + FASTPLMS_BIOHUB_TRANSFORMERS_REVISION=3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf \ + FASTPLMS_BIOHUB_TRANSFORMERS_SOURCE=/opt/oracle/vendor/upstream/biohub-transformers \ + FASTPLMS_BIOHUB_TRANSFORMERS_ATTESTATION=/opt/oracle/biohub-transformers-source-attestation.json \ + FASTPLMS_BIOHUB_TRANSFORMERS_CONTRACT=/opt/oracle/biohub-transformers-source-contract.json \ + FASTPLMS_BIOHUB_LOCK_ROOT=/opt/lock-root \ + FASTPLMS_BIOHUB_LOCK_CONTRACT=/opt/lock-root/docker/constraints/biohub-reference-lock.json \ + FASTPLMS_REFERENCE_CONTAINER_IDENTITIES=/exchange/environment/container-images.json \ + FASTPLMS_REFERENCE_CONTAINER_TARGET=reference-biohub-esm \ + PYTHONPATH=/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.parity.support.reference_adapters.esm_plusplus"] + +FROM python312 AS reference-esm2 + +COPY --from=upstream_fair_esm --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/fair-esm +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130 && \ + pip install -e /opt/oracle/vendor/upstream/fair-esm huggingface-hub==0.36.2 \ + numpy==1.26.4 pytest==9.0.2 safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/fair-esm /licenses/fair-esm +ENV PYTHONPATH=/opt/oracle/vendor/upstream/fair-esm:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.parity.support.reference_adapters.esm2"] + +FROM python310-reference AS reference-boltz2 + +COPY --from=upstream_boltz --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/boltz +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cu121 && \ + pip install '/opt/oracle/vendor/upstream/boltz[test]' \ + huggingface-hub==0.36.0 safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/boltz /licenses/boltz +ENV PYTHONPATH=/opt/oracle/vendor/upstream/boltz/src:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.structure.support.boltz2_bundle", "--help"] + +# This diagnostic oracle intentionally matches the candidate numerical +# toolchain. It localizes semantic drift without replacing the upstream-native +# `reference-boltz2` release oracle above. Neither stage installs source-built +# FlashAttention; Boltz2 structure parity uses its eager attention path. +FROM python312 AS reference-boltz2-same-runtime + +COPY --from=upstream_boltz --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/boltz +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130 && \ + pip install '/opt/oracle/vendor/upstream/boltz[test]' \ + huggingface-hub==0.36.0 safetensors==0.6.2 +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/boltz /licenses/boltz +ENV PYTHONPATH=/opt/oracle/vendor/upstream/boltz/src:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.structure.support.boltz2_bundle", "--help"] + +FROM python310-reference AS reference-esmfold + +COPY docker/constraints/esmfold.txt /tmp/constraints.txt +COPY docker/constraints/openfold-sm90.patch /tmp/openfold-sm90.patch +COPY --from=upstream_openfold --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/openfold +COPY --from=upstream_fair_esm --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/fair-esm +RUN cd /opt/oracle/vendor/upstream/openfold && \ + git apply --no-index --ignore-space-change --check /tmp/openfold-sm90.patch && \ + git apply --no-index --ignore-space-change /tmp/openfold-sm90.patch && \ + grep -F 'compute_capabilities = {(9, 0)}' setup.py && \ + grep -F -- "'-std=c++17'" setup.py +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.2.2 --index-url https://download.pytorch.org/whl/cu121 && \ + pip install --constraint /tmp/constraints.txt ninja wheel && \ + pip install --no-build-isolation --no-deps /opt/oracle/vendor/upstream/openfold && \ + pip install --constraint /tmp/constraints.txt /opt/oracle/vendor/upstream/fair-esm \ + biopython dm-tree einops huggingface-hub lightning-utilities ml-collections \ + omegaconf pytorch-lightning scipy pytest safetensors torchmetrics dllogger +RUN extension="$(python -c 'import torch; import attn_core_inplace_cuda as extension; print(extension.__file__)')" && \ + cuobjdump --list-elf "$extension" | grep -F 'sm_90' && \ + test -f /opt/venv/lib/python3.10/site-packages/dllogger-1.1.0.dist-info/licenses/LICENSE && \ + grep -F 'Apache License' \ + /opt/venv/lib/python3.10/site-packages/dllogger-1.1.0.dist-info/licenses/LICENSE +COPY --from=reference-protocol /opt/oracle /opt/oracle +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/fair-esm /licenses/fair-esm +COPY LICENSES/openfold /licenses/openfold +COPY LICENSES/dllogger/PROVENANCE.md /licenses/dllogger/PROVENANCE.md +ENV PYTHONPATH=/opt/oracle/vendor/upstream/openfold:/opt/oracle/vendor/upstream/fair-esm:/opt/oracle +WORKDIR /opt/oracle +CMD ["python", "-m", "tests.structure.support.esmfold_bundle", "--help"] + +FROM reference-biohub-esm AS reference-esmfold2 + +ENV FASTPLMS_REFERENCE_CONTAINER_TARGET=reference-esmfold2 +CMD ["python", "-c", "from tests.parity.support.reference_adapters.biohub_source import reference_sources; reference_sources(); from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model; print(ESMFold2Model.__name__)"] + +FROM python310-reference AS reference-protein-ttt + +COPY --from=upstream_protein_ttt --exclude=.git --exclude=.git/** . /opt/oracle/vendor/upstream/protein-ttt +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==2.2.2 --index-url https://download.pytorch.org/whl/cu121 && \ + if [ -f /opt/oracle/vendor/upstream/protein-ttt/requirements.txt ]; then \ + pip install -r /opt/oracle/vendor/upstream/protein-ttt/requirements.txt; \ + fi +COPY THIRD_PARTY_NOTICES.md /licenses/THIRD_PARTY_NOTICES.md +COPY LICENSES/protein-ttt /licenses/protein-ttt +ENV PYTHONPATH=/opt/oracle/vendor/upstream/protein-ttt +CMD ["python", "-c", "import pathlib; print(pathlib.Path('/opt/oracle/vendor/upstream/protein-ttt').resolve())"] diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 0000000..3dbc46b --- /dev/null +++ b/docker/README.md @@ -0,0 +1,91 @@ +# Container environments + +`docker/Dockerfile` is the only maintained Dockerfile. Its candidate stages +pin Python 3.12, CUDA 13.0, PyTorch 2.13.0, and Transformers 5.13.0. The runtime +stage contains FastPLMs source but never contains `vendor/upstream/`. Candidate +dependencies are installed with `uv pip install` from the named profiles under +`requirements/profiles/` and the validation constraint. FastPLMs itself is not +installed. Candidate and runtime stages load the copied or mounted source +through `PYTHONPATH`. + +The single runtime stage accepts `FASTPLMS_RUNTIME_PROFILE=core` by default or +`FASTPLMS_RUNTIME_PROFILE=esmfold2-fp8` for ESMFold2 serving. Any other value +fails the image build. Bake exposes `runtime` and `runtime-fp8` names that both +point to this stage with the corresponding profile. The FP8 profile and the +`candidate-fp8` stage install Transformer Engine. The +`candidate-structure` and `candidate-fp8` profiles install cuEquivariance so +named ESMFold2 and Boltz2 kernel paths can be validated against the CUDA 13 +runtime. The ordinary candidate, structure, and artifact profiles retain the +canonical PyTorch CUDA dependency graph and do not install Transformer Engine. + +Reference targets are intentionally isolated because the official projects use +incompatible dependency stacks. They copy only their named submodule from the +target-specific Buildx context and act as parity oracles. The Biohub ESM +reference never resolves +the upstream package's mutable Transformers `@main` direct reference. Its +non-Transformers dependencies are derived with a fail-closed PEP 508 subset +parser, Biohub ESM is installed without dependency resolution, and the +manifest-pinned Biohub Transformers checkout is force-installed last from its +local source. The source checkout is not placed on startup `PYTHONPATH`; the +runtime gate inventories and hashes the complete tree before prioritizing its +package path. The installed wheel supplies dependency and distribution +metadata, while the independently attested source checkout is authoritative +for executed module bytes. That wheel is built from a disposable exact copy so +setuptools cannot add `build/` output to the attested tree. Image construction +runs `pip check` and validates the copied +checkout against its independent checked-in revision and tracked-tree digest. +Every ESMC, ESM3, and ESMFold2 entry path repeats that full-tree, +pre-import-origin, package-version, and post-import attestation. Native result +metadata records the source revision, source-tree and attestation hashes, +package version, and source-relative import identity under the versioned +`reference_source_attestation` field. + +Bake does not hard-code a CPU architecture. Direct invocations use the native +builder platform; `tools/remote/run.py` passes its preflight-resolved platform +explicitly and verifies each loaded image's OS, architecture, and content +digest. The GH200 oracle uses the hash-attested `linux/arm64` dependency lock, +including source-built wheels where required. Containers do not erase ABI +boundaries, so evidence is never transferred across platforms. +Biohub suites build and load the tagged `biohub-biotraj-wheel` target alongside +the reference images. Before any oracle runs, remote orchestration persists its +content digest with every other target under +`artifacts/reference/environment/container-images.json`; tags and creation +timestamps are excluded from that canonical identity. + +Production code must not import the upstream directories. + +Routine candidate checks do not require official submodules. From the repository +root, build the exact image that Compose will run and pass the complete pytest +selection explicitly: + +```bash +sudo docker buildx bake -f docker/docker-bake.hcl candidate --load +sudo docker compose -f docker/compose.yaml run --rm candidate \ + python -m pytest tests/unit tests/integration \ + -m "not gpu and not slow and not structure" +``` + +Live official parity is a compliance workflow, not a bare +`pytest tests/parity` invocation. It must prepare immutable requests, run each +official implementation in its isolated image, build candidate artifacts, and +then consume the normalized results. The remote runner performs that complete +sequence and is the canonical executable command: + +```bash +git submodule update --init --recursive +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/ssh-key \ + --suite compliance +``` + +Use `docker buildx bake -f docker/docker-bake.hcl references --load` only when +all pinned submodules are initialized and the isolated reference images are +needed for that compliance workflow. Building those images alone does not run +parity. + +Compose centralizes GPU access, `ipc: host`, source mounts, and persistent +Hugging Face and Torch caches. Reference build targets use the names recorded +for their model families in `src/fastplms/models.toml`. No image contains +checkpoint weights, SSH keys, Hub tokens, workstation addresses, or other +credentials. diff --git a/docker/biohub-reference-lock.Dockerfile b/docker/biohub-reference-lock.Dockerfile new file mode 100644 index 0000000..c9b6b89 --- /dev/null +++ b/docker/biohub-reference-lock.Dockerfile @@ -0,0 +1,94 @@ +# syntax=docker/dockerfile:1.19@sha256:b6afd42430b15f2d2a4c5a02b919e98a525b785b1aaff16747d2f623364e39b6 + +# This workflow is intentionally native-only. It builds the BioTraj wheel used +# by the Biohub reference on the GH200 Linux ARM64 target and does not claim an +# x86_64 or emulated build contract. +ARG PYTHON312_ARM64_IMAGE=python:3.12.11-slim-bookworm@sha256:9bb659dc6d5218917236f3711e866a5634bb4c2f208de9d4533aa4863f57c1d3 +ARG CUDA13_ARM64_IMAGE=nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04@sha256:6fd95f7235f228fc7cff6f45d7b2d16ed93bddff4e80a94288731c6c7cea81d7 + +FROM ${PYTHON312_ARM64_IMAGE} AS cpython312-arm64 + +FROM ${CUDA13_ARM64_IMAGE} AS biotraj-wheel-builder + +COPY --from=cpython312-arm64 /usr/local /usr/local + +ENV PATH=/opt/biotraj-build/bin:/usr/local/bin:/usr/local/cuda/bin:$PATH \ + SOURCE_DATE_EPOCH=1730547054 \ + PYTHONHASHSEED=0 \ + PYTHONDONTWRITEBYTECODE=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PYTHONPATH=/opt/contract-root \ + TZ=UTC \ + LC_ALL=C.UTF-8 \ + LANG=C.UTF-8 \ + CFLAGS="-O3 -g0 -fno-record-gcc-switches -ffile-prefix-map=/tmp/biotraj-one=/usr/src/biotraj -ffile-prefix-map=/tmp/biotraj-two=/usr/src/biotraj" \ + LDFLAGS="-Wl,--build-id=none" \ + ARFLAGS=crD \ + ZERO_AR_DATE=1 + +COPY tools/remote/biohub_reference_lock.py /opt/contract-root/tools/remote/biohub_reference_lock.py +COPY docker/biohub-reference-lock.Dockerfile /opt/contract-root/docker/biohub-reference-lock.Dockerfile +COPY docker/constraints/biohub-reference.in /opt/contract-root/docker/constraints/biohub-reference.in +COPY docker/constraints/biohub-reference.lock.txt /opt/contract-root/docker/constraints/biohub-reference.lock.txt +COPY docker/constraints/biohub-biotraj-build.in /opt/contract-root/docker/constraints/biohub-biotraj-build.in +COPY docker/constraints/biohub-biotraj-build.lock.txt /opt/contract-root/docker/constraints/biohub-biotraj-build.lock.txt +COPY docker/constraints/biohub-reference-lock.json /opt/contract-root/docker/constraints/biohub-reference-lock.json + +RUN test "$(uname -m)" = aarch64 && \ + test "$(python --version)" = "Python 3.12.11" && \ + /usr/local/bin/python -m venv /opt/biotraj-build && \ + /opt/biotraj-build/bin/python -m pip install \ + --require-hashes \ + --only-binary=:all: \ + --no-deps \ + --requirement /opt/contract-root/docker/constraints/biohub-biotraj-build.lock.txt && \ + /opt/biotraj-build/bin/python -m pip check && \ + /opt/biotraj-build/bin/python -m tools.remote.biohub_reference_lock verify-contract \ + --root /opt/contract-root \ + --contract /opt/contract-root/docker/constraints/biohub-reference-lock.json + +ADD --checksum=sha256:4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a \ + https://files.pythonhosted.org/packages/07/21/2287edfd0d2569639eea706e25c39e63b46a384cf1712db8ea05768317b0/biotraj-1.2.2.tar.gz \ + /opt/sources/biotraj-1.2.2.tar.gz + +RUN mkdir -p /tmp/biotraj-one /tmp/biotraj-two \ + /opt/wheels/one /opt/wheels/two /opt/artifact && \ + tar -xzf /opt/sources/biotraj-1.2.2.tar.gz \ + --strip-components=1 -C /tmp/biotraj-one && \ + tar -xzf /opt/sources/biotraj-1.2.2.tar.gz \ + --strip-components=1 -C /tmp/biotraj-two && \ + sed -i \ + -e 's|src/biotraj/xtc.pyx|src/biotraj/xtc.c|' \ + -e 's|src/biotraj/trr.pyx|src/biotraj/trr.c|' \ + -e 's|src/biotraj/dcd.pyx|src/biotraj/dcd.c|' \ + /tmp/biotraj-one/setup.py /tmp/biotraj-two/setup.py && \ + test "$(grep -h -E 'src/biotraj/(xtc|trr|dcd)\.c' \ + /tmp/biotraj-one/setup.py /tmp/biotraj-two/setup.py | wc -l)" = 6 && \ + ! grep -q 'src/biotraj/.*\.pyx' \ + /tmp/biotraj-one/setup.py /tmp/biotraj-two/setup.py && \ + umask 022 && \ + /opt/biotraj-build/bin/python -m pip wheel --no-deps --no-build-isolation \ + --wheel-dir /opt/wheels/one /tmp/biotraj-one && \ + /opt/biotraj-build/bin/python -m pip wheel --no-deps --no-build-isolation \ + --wheel-dir /opt/wheels/two /tmp/biotraj-two && \ + test "$(find /opt/wheels/one -maxdepth 1 -type f -name '*.whl' | wc -l)" = 1 && \ + test "$(find /opt/wheels/two -maxdepth 1 -type f -name '*.whl' | wc -l)" = 1 && \ + cmp /opt/wheels/one/*.whl /opt/wheels/two/*.whl && \ + cp /opt/wheels/one/*.whl /opt/artifact/ && \ + sha256sum /opt/artifact/*.whl > /opt/artifact/biotraj-wheel.sha256 && \ + python --version > /opt/artifact/python-version.txt && \ + gcc --version | head -n 1 > /opt/artifact/gcc-version.txt && \ + ld --version | head -n 1 > /opt/artifact/ld-version.txt && \ + uname -m > /opt/artifact/machine.txt + +FROM scratch AS biotraj-wheel-artifact + +COPY --from=biotraj-wheel-builder /opt/artifact /biohub-reference-lock + +FROM biotraj-wheel-builder AS biotraj-wheel-evidence + +ENTRYPOINT ["/opt/biotraj-build/bin/python", "-m", "tools.remote.biohub_reference_lock", \ + "write-build-evidence", "--root", "/opt/contract-root", \ + "--contract", "/opt/contract-root/docker/constraints/biohub-reference-lock.json", \ + "--wheel", "/opt/artifact/biotraj-1.2.2-cp312-cp312-linux_aarch64.whl", \ + "--output", "/output/biohub-reference-build-evidence.json"] diff --git a/docker/compose.yaml b/docker/compose.yaml new file mode 100644 index 0000000..6915d47 --- /dev/null +++ b/docker/compose.yaml @@ -0,0 +1,227 @@ +name: fastplms + +x-accelerator: &accelerator + ipc: host + init: true + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + +x-common: &common + <<: *accelerator + working_dir: /workspace + environment: + CUBLAS_WORKSPACE_CONFIG: ":4096:8" + HF_HOME: /cache/huggingface + TORCH_HOME: /cache/torch + XDG_CACHE_HOME: /cache/xdg + FASTPLMS_ARTIFACTS: /workspace/artifacts + FASTPLMS_CANDIDATE_ARTIFACTS: /workspace/dist/hub + FASTPLMS_REFERENCE_RESULTS: /workspace/artifacts/reference/results + volumes: + - ..:/workspace + - hf-cache:/cache/huggingface + - torch-cache:/cache/torch + - xdg-cache:/cache/xdg + +x-reference: &reference + <<: *accelerator + working_dir: /opt/oracle + environment: + CUBLAS_WORKSPACE_CONFIG: ":4096:8" + HF_HOME: /cache/huggingface + TORCH_HOME: /cache/torch + XDG_CACHE_HOME: /cache/xdg + volumes: + - ../artifacts/reference:/exchange + - hf-cache:/cache/huggingface + - torch-cache:/cache/torch + - xdg-cache:/cache/xdg + +services: + biohub-biotraj-wheel: + image: local/fastplms-biohub-biotraj-wheel:dev + build: + context: .. + dockerfile: docker/biohub-reference-lock.Dockerfile + target: biotraj-wheel-artifact + + candidate: + <<: *common + image: local/fastplms-candidate:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: candidate + command: + - python + - -m + - pytest + - tests/unit + - tests/integration + - -m + - not gpu and not slow and not structure + + structure: + <<: *common + image: local/fastplms-structure:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: candidate-structure + command: + - python + - -m + - pytest + - tests/structure + - tests/parity/test_boltz_source_refactor.py + - -m + - structure + + fp8: + <<: *common + image: local/fastplms-fp8:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: candidate-fp8 + command: ["python", "-m", "pytest", "tests/structure", "-m", "structure"] + + benchmark: + <<: *common + image: local/fastplms-fp8:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: candidate-fp8 + entrypoint: ["python", "-m", "benchmarks.suite"] + + artifact: + <<: *common + image: local/fastplms-artifact:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: candidate-artifact + environment: + HF_HOME: /cache/huggingface + HF_HUB_OFFLINE: "1" + TRANSFORMERS_OFFLINE: "1" + TORCH_HOME: /cache/torch + XDG_CACHE_HOME: /cache/xdg + FASTPLMS_ARTIFACTS: /workspace/artifacts + PYTHONPATH: "" + command: + - python + - -m + - pytest + - tests/release/test_published_automodel.py + - -m + - artifact + + reference-ankh: + <<: *reference + image: local/fastplms-reference-ankh:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-ankh + additional_contexts: + upstream_ankh: ../vendor/upstream/ankh + + reference-biohub-esm: + <<: *reference + image: local/fastplms-reference-biohub-esm:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-biohub-esm + additional_contexts: + biohub_biotraj_wheel: service:biohub-biotraj-wheel + upstream_biohub_esm: ../vendor/upstream/biohub-esm + upstream_biohub_transformers: ../vendor/upstream/biohub-transformers + + reference-boltz2: + <<: *reference + image: local/fastplms-reference-boltz2:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-boltz2 + additional_contexts: + upstream_boltz: ../vendor/upstream/boltz + + reference-dplm: + <<: *reference + image: local/fastplms-reference-dplm:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-dplm + additional_contexts: + upstream_dplm: ../vendor/upstream/dplm + + reference-e1: + <<: *reference + image: local/fastplms-reference-e1:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-e1 + additional_contexts: + upstream_e1: ../vendor/upstream/e1 + + reference-esm2: + <<: *reference + image: local/fastplms-reference-esm2:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-esm2 + additional_contexts: + upstream_fair_esm: ../vendor/upstream/fair-esm + + reference-esmfold: + <<: *reference + image: local/fastplms-reference-esmfold:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-esmfold + additional_contexts: + upstream_fair_esm: ../vendor/upstream/fair-esm + upstream_openfold: ../vendor/upstream/openfold + + reference-esmfold2: + <<: *reference + image: local/fastplms-reference-esmfold2:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-esmfold2 + additional_contexts: + biohub_biotraj_wheel: service:biohub-biotraj-wheel + upstream_biohub_esm: ../vendor/upstream/biohub-esm + upstream_biohub_transformers: ../vendor/upstream/biohub-transformers + + reference-protein-ttt: + <<: *reference + image: local/fastplms-reference-protein-ttt:dev + build: + context: .. + dockerfile: docker/Dockerfile + target: reference-protein-ttt + additional_contexts: + upstream_protein_ttt: ../vendor/upstream/protein-ttt + +volumes: + hf-cache: + name: fastplms-hf-cache + torch-cache: + name: fastplms-torch-cache + xdg-cache: + name: fastplms-xdg-cache diff --git a/docker/constraints/ankh.txt b/docker/constraints/ankh.txt new file mode 100644 index 0000000..df7abf1 --- /dev/null +++ b/docker/constraints/ankh.txt @@ -0,0 +1,7 @@ +torch==2.13.0 +numpy==1.26.4 +transformers==4.25.1 +tokenizers==0.13.3 +sentencepiece==0.1.99 +biopython==1.80 +datasets==2.20.0 diff --git a/docker/constraints/biohub-biotraj-build.in b/docker/constraints/biohub-biotraj-build.in new file mode 100644 index 0000000..c0a83d0 --- /dev/null +++ b/docker/constraints/biohub-biotraj-build.in @@ -0,0 +1,9 @@ +# Exact PEP 517 build inputs for the BioTraj 1.2.2 source distribution used by +# the GH200 Biohub reference. These versions satisfy the sdist's build-system +# requirements and are deliberately isolated from the runtime environment. +pip==26.1.1 +setuptools==83.0.0 +setuptools-scm==9.2.2 +wheel==0.47.0 +numpy==2.4.4 +cython==3.2.4 diff --git a/docker/constraints/biohub-biotraj-build.lock.txt b/docker/constraints/biohub-biotraj-build.lock.txt new file mode 100644 index 0000000..4f3699d --- /dev/null +++ b/docker/constraints/biohub-biotraj-build.lock.txt @@ -0,0 +1,133 @@ +# This file was autogenerated by uv via the following command: +# python -m tools.remote.biohub_reference_lock generate-build-lock +--index-url https://pypi.org/simple +--only-binary :all: + +cython==3.2.4 \ + --hash=sha256:02cb0cc0f23b9874ad262d7d2b9560aed9c7e2df07b49b920bda6f2cc9cb505e \ + --hash=sha256:03893c88299a2c868bb741ba6513357acd104e7c42265809fd58dce1456a36fc \ + --hash=sha256:14dae483ca2838b287085ff98bc206abd7a597b7bb16939a092f8e84d9062842 \ + --hash=sha256:1a64a112a34ec719b47c01395647e54fb4cf088a511613f9a3a5196694e8e382 \ + --hash=sha256:28b1e363b024c4b8dcf52ff68125e635cb9cb4b0ba997d628f25e32543a71103 \ + --hash=sha256:28e8075087a59756f2d059273184b8b639fe0f16cf17470bd91c39921bc154e0 \ + --hash=sha256:2b1f12c0e4798293d2754e73cd6f35fa5bbdf072bdc14bc6fc442c059ef2d290 \ + --hash=sha256:31a90b4a2c47bb6d56baeb926948348ec968e932c1ae2c53239164e3e8880ccf \ + --hash=sha256:35ab0632186057406ec729374c737c37051d2eacad9d515d94e5a3b3e58a9b02 \ + --hash=sha256:36bf3f5eb56d5281aafabecbaa6ed288bc11db87547bba4e1e52943ae6961ccf \ + --hash=sha256:3b6e58f73a69230218d5381817850ce6d0da5bb7e87eb7d528c7027cbba40b06 \ + --hash=sha256:3b8e62049afef9da931d55de82d8f46c9a147313b69d5ff6af6e9121d545ce7a \ + --hash=sha256:55b6c44cd30821f0b25220ceba6fe636ede48981d2a41b9bbfe3c7902ce44ea7 \ + --hash=sha256:55eb425c0baf1c8a46aa4424bc35b709db22f3c8a1de33adb3ecb8a3d54ea42a \ + --hash=sha256:64d7f71be3dd6d6d4a4c575bb3a4674ea06d1e1e5e4cd1b9882a2bc40ed3c4c9 \ + --hash=sha256:67922c9de058a0bfb72d2e75222c52d09395614108c68a76d9800f150296ddb3 \ + --hash=sha256:6d5267f22b6451eb1e2e1b88f6f78a2c9c8733a6ddefd4520d3968d26b824581 \ + --hash=sha256:72e6c0bbd978e2678b45351395f6825b9b8466095402eae293f4f7a73e9a3e85 \ + --hash=sha256:732fc93bc33ae4b14f6afaca663b916c2fdd5dcbfad7114e17fb2434eeaea45c \ + --hash=sha256:767b143704bdd08a563153448955935844e53b852e54afdc552b43902ed1e235 \ + --hash=sha256:83266c356c13c68ffe658b4905279c993d8a5337bb0160fa90c8a3e297ea9a2e \ + --hash=sha256:84226ecd313b233da27dc2eb3601b4f222b8209c3a7216d8733b031da1dc64e6 \ + --hash=sha256:869487ea41d004f8b92171f42271fbfadb1ec03bede3158705d16cd570d6b891 \ + --hash=sha256:90f43be4eaa6afd58ce20d970bb1657a3627c44e1760630b82aa256ba74b4acb \ + --hash=sha256:983f9d2bb8a896e16fa68f2b37866ded35fa980195eefe62f764ddc5f9f5ef8e \ + --hash=sha256:b362819d155fff1482575e804e43e3a8825332d32baa15245f4642022664a3f4 \ + --hash=sha256:b84d4e3c875915545f77c88dba65ad3741afd2431e5cdee6c9a20cefe6905647 \ + --hash=sha256:ca2399dc75796b785f74fb85c938254fa10c80272004d573c455f9123eceed86 \ + --hash=sha256:ca578c9cb872c7ecffbe14815dc4590a003bc13339e90b2633540c7e1a252839 \ + --hash=sha256:d4b4fd5332ab093131fa6172e8362f16adef3eac3179fd24bbdc392531cb82fa \ + --hash=sha256:e3b5ac54e95f034bc7fb07313996d27cbf71abc17b229b186c1540942d2dc28e \ + --hash=sha256:e65e4773021f8dc8532010b4fbebe782c77f9a0817e93886e518c93bd6a44e9d \ + --hash=sha256:e71efb20048358a6b8ec604a0532961c50c067b5e63e345e2e359fff72feaee8 \ + --hash=sha256:f136f379a4a54246facd0eb6f1ee15c3837cb314ce87b677582ec014db4c6845 \ + --hash=sha256:f583cad7a7eed109f0babb5035e92d0c1260598f53add626a8568b57246b62c3 \ + --hash=sha256:f81eda419b5ada7b197bbc3c5f4494090e3884521ffd75a3876c93fbf66c9ca8 \ + --hash=sha256:f8d685a70bce39acc1d62ec3916d9b724b5ef665b0ce25ae55e1c85ee09747fc \ + --hash=sha256:fdfdd753ad7e18e5092b413e9f542e8d28b8a08203126090e1c15f7783b7fe57 \ + --hash=sha256:ff9af2134c05e3734064808db95b4dd7341a39af06e8945d05ea358e1741aaed +numpy==2.4.4 \ + --hash=sha256:07077278157d02f65c43b1b26a3886bce886f95d20aabd11f87932750dfb14ed \ + --hash=sha256:08f2e31ed5e6f04b118e49821397f12767934cfdd12a1ce86a058f91e004ee50 \ + --hash=sha256:0aec54fd785890ecca25a6003fd9a5aed47ad607bbac5cd64f836ad8666f4959 \ + --hash=sha256:0d35aea54ad1d420c812bfa0385c71cd7cc5bcf7c65fed95fc2cd02fe8c79827 \ + --hash=sha256:0d4e437e295f18ec29bc79daf55e8a47a9113df44d66f702f02a293d93a2d6dd \ + --hash=sha256:0dfd3f9d3adbe2920b68b5cd3d51444e13a10792ec7154cd0a2f6e74d4ab3233 \ + --hash=sha256:1378871da56ca8943c2ba674530924bb8ca40cd228358a3b5f302ad60cf875fc \ + --hash=sha256:15716cfef24d3a9762e3acdf87e27f58dc823d1348f765bbea6bef8c639bfa1b \ + --hash=sha256:19710a9ca9992d7174e9c52f643d4272dcd1558c5f7af7f6f8190f633bd651a7 \ + --hash=sha256:23cbfd4c17357c81021f21540da84ee282b9c8fba38a03b7b9d09ba6b951421e \ + --hash=sha256:2483e4584a1cb3092da4470b38866634bafb223cbcd551ee047633fd2584599a \ + --hash=sha256:27a8d92cd10f1382a67d7cf4db7ce18341b66438bdd9f691d7b0e48d104c2a9d \ + --hash=sha256:28a650663f7314afc3e6ec620f44f333c386aad9f6fc472030865dc0ebb26ee3 \ + --hash=sha256:2aa0613a5177c264ff5921051a5719d20095ea586ca88cc802c5c218d1c67d3e \ + --hash=sha256:2c194dd721e54ecad9ad387c1d35e63dce5c4450c6dc7dd5611283dda239aabb \ + --hash=sha256:2d19e6e2095506d1736b7d80595e0f252d76b89f5e715c35e06e937679ea7d7a \ + --hash=sha256:2d390634c5182175533585cc89f3608a4682ccb173cc9bb940b2881c8d6f8fa0 \ + --hash=sha256:30caa73029a225b2d40d9fae193e008e24b2026b7ee1a867b7ee8d96ca1a448e \ + --hash=sha256:42c16925aa5a02362f986765f9ebabf20de75cdefdca827d14315c568dcab113 \ + --hash=sha256:45dbed2ab436a9e826e302fcdcbe9133f9b0006e5af7168afb8963a6520da103 \ + --hash=sha256:4636de7fd195197b7535f231b5de9e4b36d2c440b6e566d2e4e4746e6af0ca93 \ + --hash=sha256:4a19d9dba1a76618dd86b164d608566f393f8ec6ac7c44f0cc879011c45e65af \ + --hash=sha256:4bbc7f303d125971f60ec0aaad5e12c62d0d2c925f0ab1273debd0e4ba37aba5 \ + --hash=sha256:4d6d57903571f86180eb98f8f0c839fa9ebbfb031356d87f1361be91e433f5b7 \ + --hash=sha256:4e874c976154687c1f71715b034739b45c7711bec81db01914770373d125e392 \ + --hash=sha256:51fc224f7ca4d92656d5a5eb315f12eb5fe2c97a66249aa7b5f562528a3be38c \ + --hash=sha256:58c8b5929fcb8287cbd6f0a3fae19c6e03a5c48402ae792962ac465224a629a4 \ + --hash=sha256:5a285b3b96f951841799528cd1f4f01cd70e7e0204b4abebac9463eecfcf2a40 \ + --hash=sha256:5c70f1cc1c4efbe316a572e2d8b9b9cc44e89b95f79ca3331553fbb63716e2bf \ + --hash=sha256:62d6b0f03b694173f9fcb1fb317f7222fd0b0b103e784c6549f5e53a27718c44 \ + --hash=sha256:6a246d5914aa1c820c9443ddcee9c02bec3e203b0c080349533fae17727dfd1b \ + --hash=sha256:6aa3236c78803afbcb255045fbef97a9e25a1f6c9888357d205ddc42f4d6eba5 \ + --hash=sha256:6bbe4eb67390b0a0265a2c25458f6b90a409d5d069f1041e6aff1e27e3d9a79e \ + --hash=sha256:715d1c092715954784bc79e1174fc2a90093dc4dc84ea15eb14dad8abdcdeb74 \ + --hash=sha256:72944b19f2324114e9dc86a159787333b77874143efcf89a5167ef83cfee8af0 \ + --hash=sha256:81f4a14bee47aec54f883e0cad2d73986640c1590eb9bfaaba7ad17394481e6e \ + --hash=sha256:846300f379b5b12cc769334464656bc882e0735d27d9726568bc932fdc49d5ec \ + --hash=sha256:86b6f55f5a352b48d7fbfd2dbc3d5b780b2d79f4d3c121f33eb6efb22e9a2015 \ + --hash=sha256:874f200b2a981c647340f841730fc3a2b54c9d940566a3c4149099591e2c4c3d \ + --hash=sha256:8a87ec22c87be071b6bdbd27920b129b94f2fc964358ce38f3822635a3e2e03d \ + --hash=sha256:8b3b60bb7cba2c8c81837661c488637eee696f59a877788a396d33150c35d842 \ + --hash=sha256:8e3ed142f2728df44263aaf5fb1f5b0b99f4070c553a0d7f033be65338329150 \ + --hash=sha256:93e15038125dc1e5345d9b5b68aa7f996ec33b98118d18c6ca0d0b7d6198b7e8 \ + --hash=sha256:989824e9faf85f96ec9c7761cd8d29c531ad857bfa1daa930cba85baaecf1a9a \ + --hash=sha256:99d838547ace2c4aace6c4f76e879ddfe02bb58a80c1549928477862b7a6d6ed \ + --hash=sha256:9b2aec6af35c113b05695ebb5749a787acd63cafc83086a05771d1e1cd1e555f \ + --hash=sha256:9c585a1790d5436a5374bac930dad6ed244c046ed91b2b2a3634eb2971d21008 \ + --hash=sha256:a7164afb23be6e37ad90b2f10426149fd75aee07ca55653d2aa41e66c4ef697e \ + --hash=sha256:ac6b31e35612a26483e20750126d30d0941f949426974cace8e6b5c58a3657b0 \ + --hash=sha256:ad2e2ef14e0b04e544ea2fa0a36463f847f113d314aa02e5b402fdf910ef309e \ + --hash=sha256:b268594bccac7d7cf5844c7732e3f20c50921d94e36d7ec9b79e9857694b1b2f \ + --hash=sha256:b5f0362dc928a6ecd9db58868fca5e48485205e3855957bdedea308f8672ea4a \ + --hash=sha256:ba1f4fc670ed79f876f70082eff4f9583c15fb9a4b89d6188412de4d18ae2f40 \ + --hash=sha256:ba203255017337d39f89bdd58417f03c4426f12beed0440cfd933cb15f8669c7 \ + --hash=sha256:c901b15172510173f5cb310eae652908340f8dede90fff9e3bf6c0d8dfd92f83 \ + --hash=sha256:c9b39d38a9bd2ae1becd7eac1303d031c5c110ad31f2b319c6e7d98b135c934d \ + --hash=sha256:d2a8490669bfe99a233298348acc2d824d496dee0e66e31b66a6022c2ad74a5c \ + --hash=sha256:dddbbd259598d7240b18c9d87c56a9d2fb3b02fe266f49a7c101532e78c1d871 \ + --hash=sha256:df3775294accfdd75f32c74ae39fcba920c9a378a2fc18a12b6820aa8c1fb502 \ + --hash=sha256:e44319a2953c738205bf3354537979eaa3998ed673395b964c1176083dd46252 \ + --hash=sha256:e4a010c27ff6f210ff4c6ef34394cd61470d01014439b192ec22552ee867f2a8 \ + --hash=sha256:e823b8b6edc81e747526f70f71a9c0a07ac4e7ad13020aa736bb7c9d67196115 \ + --hash=sha256:e892aff75639bbef0d2a2cfd55535510df26ff92f63c92cd84ef8d4ba5a5557f \ + --hash=sha256:eea7ac5d2dce4189771cedb559c738a71512768210dc4e4753b107a2048b3d0e \ + --hash=sha256:ef4059d6e5152fa1a39f888e344c73fdc926e1b2dd58c771d67b0acfbf2aa67d \ + --hash=sha256:f169b9a863d34f5d11b8698ead99febeaa17a13ca044961aa8e2662a6c7766a0 \ + --hash=sha256:f2cf083b324a467e1ab358c105f6cad5ea950f50524668a80c486ff1db24e119 \ + --hash=sha256:f8474c4241bc18b750be2abea9d7a9ec84f46ef861dbacf86a4f6e043401f79e \ + --hash=sha256:f983334aea213c99992053ede6168500e5f086ce74fbc4acc3f2b00f5762e9db \ + --hash=sha256:f9e75681b59ddaa5e659898085ae0eaea229d054f2ac0c7e563a62205a700121 \ + --hash=sha256:fbc356aae7adf9e6336d336b9c8111d390a05df88f1805573ebb0807bd06fd1d \ + --hash=sha256:fcfe2045fd2e8f3cb0ce9d4ba6dba6333b8fa05bb8a4939c908cd43322d14c7e +packaging==26.2 \ + --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e \ + --hash=sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661 +pip==26.1.1 \ + --hash=sha256:99cb1c2899893b075ff56e4ed0af55669a955b49ad7fb8d8603ecdaf4ed653fb \ + --hash=sha256:d36762751d156a4ee895de8af39aa0abeeeb577f93a2eca6ab62467bbf0f8a78 +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 +setuptools-scm==9.2.2 \ + --hash=sha256:1c674ab4665686a0887d7e24c03ab25f24201c213e82ea689d2f3e169ef7ef57 \ + --hash=sha256:30e8f84d2ab1ba7cb0e653429b179395d0c33775d54807fc5f1dd6671801aef7 +wheel==0.47.0 \ + --hash=sha256:212281cab4dff978f6cedd499cd893e1f620791ca6ff7107cf270781e587eced \ + --hash=sha256:cc72bd1009ba0cf63922e28f94d9d83b920aa2bb28f798a31d0691b02fa3c9b3 diff --git a/docker/constraints/biohub-esm-source.json b/docker/constraints/biohub-esm-source.json new file mode 100644 index 0000000..bbb7bfb --- /dev/null +++ b/docker/constraints/biohub-esm-source.json @@ -0,0 +1,8 @@ +{ + "import_name": "esm", + "import_root": "esm", + "package_version": "3.3.0", + "schema_version": 1, + "source_revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "tree_sha256": "c5489f1fc58de200978803de2c38e1a78f769cb183a2ee90be833f0f4a0212e8" +} diff --git a/docker/constraints/biohub-reference-lock.json b/docker/constraints/biohub-reference-lock.json new file mode 100644 index 0000000..52c262c --- /dev/null +++ b/docker/constraints/biohub-reference-lock.json @@ -0,0 +1,69 @@ +{ + "schema_version": 2, + "contract": "fastplms.biohub-reference-lock", + "target": { + "hardware": "NVIDIA GH200 480GB", + "operating_system": "linux", + "architecture": "aarch64", + "container_platform": "linux/arm64", + "python_implementation": "CPython", + "python_version": "3.12", + "cuda_version": "13.0", + "torch_backend": "cu130" + }, + "container": { + "dockerfile_path": "docker/biohub-reference-lock.Dockerfile", + "dockerfile_sha256": "bfe2e238d5de518a72400233c04156d4b9fbcf509a8c8d84f021fa0827aed18d", + "dockerfile_frontend": "docker/dockerfile:1.19@sha256:b6afd42430b15f2d2a4c5a02b919e98a525b785b1aaff16747d2f623364e39b6", + "cuda_image": "nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04@sha256:6fd95f7235f228fc7cff6f45d7b2d16ed93bddff4e80a94288731c6c7cea81d7", + "cuda_image_digest": "6fd95f7235f228fc7cff6f45d7b2d16ed93bddff4e80a94288731c6c7cea81d7", + "python_image": "python:3.12.11-slim-bookworm@sha256:9bb659dc6d5218917236f3711e866a5634bb4c2f208de9d4533aa4863f57c1d3", + "python_image_digest": "9bb659dc6d5218917236f3711e866a5634bb4c2f208de9d4533aa4863f57c1d3", + "build_python_version": "3.12.11" + }, + "runtime": { + "input_path": "docker/constraints/biohub-reference.in", + "input_sha256": "78cbe5d02528baa446135400cba7bc6c28a9480b384d4f3efd05807c221fc7db", + "lock_path": "docker/constraints/biohub-reference.lock.txt", + "lock_sha256": "f87033dffffe953478b482dae82f91603fa705a68e92ee7683c1831586c94ca0", + "package_count": 108 + }, + "build": { + "input_path": "docker/constraints/biohub-biotraj-build.in", + "input_sha256": "100cb7ff60b892e66f9bf6f1c56862781997ce7932be430bc6e13086eb74f96d", + "lock_path": "docker/constraints/biohub-biotraj-build.lock.txt", + "lock_sha256": "c7864daa96028aba35081110c563b8b08c968fd312f7827449d355638f18079d", + "package_count": 7 + }, + "biotraj": { + "version": "1.2.2", + "sdist_url": "https://files.pythonhosted.org/packages/07/21/2287edfd0d2569639eea706e25c39e63b46a384cf1712db8ea05768317b0/biotraj-1.2.2.tar.gz#sha256=4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a", + "sdist_sha256": "4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a", + "wheel_filename": "biotraj-1.2.2-cp312-cp312-linux_aarch64.whl", + "wheel_sha256": "253c1354c401e97d6e951f29e0d768deb5263de6662001281870425c37719f6b", + "wheel_size": 887908 + }, + "bootstrap_inventory": { + "pip": "26.1.1" + }, + "final_inventory_overlays": { + "esm": "3.3.0", + "transformers": "4.57.6", + "uv": "0.10.12" + }, + "pip_check_platform_exceptions": [ + { + "distribution": "nvidia-cusparselt-cu13", + "version": "0.8.1", + "wheel_filename": "nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", + "wheel_sha256": "4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", + "filename_platform_tag": "py3-none-manylinux2014_aarch64", + "wheel_metadata_platform_tag": "py3-none-manylinux2014_sbsa", + "target_hardware": "NVIDIA GH200 480GB", + "target_operating_system": "linux", + "target_architecture": "aarch64", + "accepted_diagnostic": "nvidia-cusparselt-cu13 0.8.1 is not supported on this platform", + "resolution": "validated-vendor-metadata-exception-no-wheel-rewrite" + } + ] +} diff --git a/docker/constraints/biohub-reference.in b/docker/constraints/biohub-reference.in new file mode 100644 index 0000000..8d02295 --- /dev/null +++ b/docker/constraints/biohub-reference.in @@ -0,0 +1,54 @@ +# Biohub ESM reference runtime inputs. +# +# Biohub ESM source: 82ee35553d39169d678f784c8d3f8712ffd7d2c4 +# Biohub Transformers source: 3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf +# +# The mutable `transformers @ ...@main` Biohub ESM requirement is deliberately +# absent. Its pinned source checkout is installed separately without dependency +# resolution. The remaining 22 ESM requirements and the pinned fork's 10 base +# install requirements are represented below. Policy pins narrow the shared +# requirements used by the reference image. + +# Biohub ESM runtime requirements (minus the mutable Transformers direct URL). +torch==2.13.0 +ipython +einops +biotite>=1.0.0 +# BioTraj 1.2.2 publishes no CPython 3.12 Linux ARM64 wheel. Pin the one +# authoritative sdist by URL and digest; the reference build attests the wheel +# produced from it on the GH200 target. +biotraj @ https://files.pythonhosted.org/packages/07/21/2287edfd0d2569639eea706e25c39e63b46a384cf1712db8ea05768317b0/biotraj-1.2.2.tar.gz#sha256=4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a +rdkit +msgpack-numpy +biopython +scikit-learn +brotli +attrs +pandas +cloudpathlib +httpx +tenacity +zstd==1.5.6.1 +ipywidgets +py3dmol +pydssp +boto3 +pygtrie +dna_features_viewer +accelerate==1.13.0 + +# Biohub Transformers base install requirements. +filelock +huggingface-hub==0.36.2 +numpy==1.26.4 +packaging>=20.0 +pyyaml>=5.1 +regex!=2019.12.17 +requests +tokenizers>=0.22.0,<=0.23.0 +safetensors==0.5.3 +tqdm>=4.27 + +# Reference harness and deterministic wheel-install prerequisites. +pytest==9.0.2 +wheel==0.47.0 diff --git a/docker/constraints/biohub-reference.lock.txt b/docker/constraints/biohub-reference.lock.txt new file mode 100644 index 0000000..81b3ee8 --- /dev/null +++ b/docker/constraints/biohub-reference.lock.txt @@ -0,0 +1,1543 @@ +# This file was autogenerated by uv via the following command: +# python -m tools.remote.biohub_reference_lock generate +--index-url https://pypi.org/simple +--extra-index-url https://download.pytorch.org/whl/cu130 +--no-binary biotraj +--only-binary :all: + +accelerate==1.13.0 \ + --hash=sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0 \ + --hash=sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236 +anyio==4.14.2 \ + --hash=sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494 \ + --hash=sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f +asttokens==3.0.2 \ + --hash=sha256:3ecdbd8f2cc195f53ccada3a613538bb5f9ef6f6869129f13e03c30a677b8fe2 \ + --hash=sha256:9da13157f5b28becde0bd374fc677dcd3c290614264eff096f167c469cd9f933 +attrs==26.1.0 \ + --hash=sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309 \ + --hash=sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32 +biopython==1.87 \ + --hash=sha256:01ee30203bd4b2145cdfe2878499e549a7087f897a6f4d1ebd9de30790123140 \ + --hash=sha256:126e18ad44e959a1984560562fa4f295c075ce2e610e8ddbc9ec6f52e09f70d8 \ + --hash=sha256:1e951f4862ffc1dccc28e1c25245059fa653d86028a88f5dfe1b7875875f3a4f \ + --hash=sha256:331c4151608a1d8406eff0d3c52a0ff1fa3e82604fc85f11c696c562919fb161 \ + --hash=sha256:35f13796188412f135acbed196bd6cfedc1257199d50b4883e289bbd96320efc \ + --hash=sha256:3670d76759c6cb53ba617f9823d3a438c1aa5415abef6addd29cb81d61d7b312 \ + --hash=sha256:58e36efa7eaa8813cffe440af4824afa20cc3c21b1b179a95cc0fb15c4b83c01 \ + --hash=sha256:63978e1b3ae040c52369bc1bd3f4c668171f5f1f0808798eccd74b575cce455d \ + --hash=sha256:65ba69ef0273e983a9036c2a228142bc34266179a5f03660fc281d332d718630 \ + --hash=sha256:6d221b2e08e7e89713fdbfb15c8ea6744e908d59f672cd2b6fcf9ed47910d05e \ + --hash=sha256:772539297fa16a78f38651c793f53f8c11bd18317b111982e72cf30a6e57512a \ + --hash=sha256:77ccc634621904d4a8fa0a43b5e0f093fa9df8c9577ed3858af648bb3528f51e \ + --hash=sha256:8456c803459b679a9712422e5a7fd9809f2f089bf69bb085f3b077946ac9bdbf \ + --hash=sha256:856e3d64f1f27db493474ff84916ed8572731a525e001c7d0d8f41a0fd187000 \ + --hash=sha256:86596b05f39afbc5984ee6239c93fd75f6a97fffb9a630d74b45652c24aab964 \ + --hash=sha256:89ffe272517478691439a59cccd3cc2929fc8f6bfb8cbc8cc5acc103660395a7 \ + --hash=sha256:98e397096336a49804b6aaaeac8c47ad82e3e4430862f0cde37be73037f1017e \ + --hash=sha256:a3428155c3e0abbed7aad5ff08e034d435b84dfe560c8ec58e7d43abda4b6a43 \ + --hash=sha256:a6fcec2e3602ed52ced701f8f7851952383f84dbc4caeb4d202d088170e86b6d \ + --hash=sha256:b077777fd2c555434bdcee58743f6f860aa80e1e005d9671913aa73823c6a773 \ + --hash=sha256:bab39ec4108fb2542a9f1012db8269426fd6e86ed84ebc1b0bd11134ef443dd3 \ + --hash=sha256:c8da2d44a4b912c7550a051a5ff4bb72a61decc9c4b19ea92cba4c02fffb143c \ + --hash=sha256:ccf00d15e698656796ab14ff3eb175e282da7a08eedd36a29b6cacec5a33e97f \ + --hash=sha256:d4ff5369ffb7a966bcf921cae6c8a6eab5070b5df1c9f2ae8c18ad40fdcb7ec5 \ + --hash=sha256:d740c75d4bc94f9dff51719a0deda37e5e885f06ee6dfbb5e9a21bbe9de35a9c \ + --hash=sha256:db73fe16aa2b20677ac86d1997612acb0aa1a3720bc899f65d2bce5583208e1b \ + --hash=sha256:e05ef5d632c319ab3ef77705c74061190d0792b07e1f2b9eee867401b2758e7e \ + --hash=sha256:e4878a9b56775480154c686f81e98f6d907b44d87605bdc2f53538ccdfde9624 \ + --hash=sha256:efc16dd8a9312eb655fa590821495b0d8ca25d19f7b4ef4fa4da9e71d59d33e5 \ + --hash=sha256:fab1b12f6bc4646b7f56b4c390ecff685f02b5b29e3a0c10477195bb49fe62f8 \ + --hash=sha256:ff28a6f31630b3c9f52903478a2ed9dd894b07c1998e40eaeefbeefac20f2d0f +biotite==1.7.1 \ + --hash=sha256:085976013c62977eb29c5180cdffbc4b49633f35172b64c99e56eb7007d9f516 \ + --hash=sha256:1465b43485d29f5f1f80cfb43ee64e68aa0128a248b31cbc04fadfcc20f8288d \ + --hash=sha256:2ae4a5d2c2d5ba08ca5d89a647984c99f45994eb908b614e896ce9e4db3ca800 \ + --hash=sha256:327f008ec9539624d05161f34ac279b76f4ca303bd98cc9b49d6557dad7cc5c9 \ + --hash=sha256:4a6a2bdeb92adf21e2a94fe0de898aa2dc01ed8c8b8241b4b2fd646c0dabc155 \ + --hash=sha256:4f99104e7ebc1b2ab8ff7910736321e1205c85dba532826877ff263c65bd85ff \ + --hash=sha256:5212116934b9c22515f86a1c23046cf3ecf7c9706464b14f16d72cd993866095 \ + --hash=sha256:6d7dbc50b1ea807131775d38078383a9513384fad0119f519a399cf701c34cf0 \ + --hash=sha256:858b9b947933744f2251a06197cb0a15ec6b64ec0b54bae4fbafecd39725c3ce \ + --hash=sha256:965824257f716a77de3df35616cdaf302483a9ab8df98dc1dd0096c987342c49 \ + --hash=sha256:d34184f23facc4a35bc0ac7914162f6fc5b6c68bf7fbf3187c1966cf53ce9918 \ + --hash=sha256:d47443b507c6968cb6823ec0cc0bdca1baed6d11ce08bcc015b8a5409d943bf9 \ + --hash=sha256:f563bebc5f2cc754b0e3773b75d5c216cfb5e89146f0c68dfec3d1d1d50051fb +biotraj @ https://files.pythonhosted.org/packages/07/21/2287edfd0d2569639eea706e25c39e63b46a384cf1712db8ea05768317b0/biotraj-1.2.2.tar.gz#sha256=4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a \ + --hash=sha256:4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a +boto3==1.43.53 \ + --hash=sha256:5383e705d8a976a14f23bb8c113c07a396931a019db98fce4cdc68650ec6e4d8 \ + --hash=sha256:c80425acab314d7af09609562053f565139e1fe49108eacfcc1601ebfaee235b +botocore==1.43.53 \ + --hash=sha256:36d93dd8db68ee75f6b61ca9f775161b8168844e4601698701530e6efdded141 \ + --hash=sha256:b7ee9a70d187e5348883c820990ccd9436ab14e2bd6622741fc96fe561e816b8 +brotli==1.2.0 \ + --hash=sha256:022426c9e99fd65d9475dce5c195526f04bb8be8907607e27e747893f6ee3e24 \ + --hash=sha256:072e7624b1fc4d601036ab3f4f27942ef772887e876beff0301d261210bca97f \ + --hash=sha256:09ac247501d1909e9ee47d309be760c89c990defbb2e0240845c892ea5ff0de4 \ + --hash=sha256:0bbd5b5ccd157ae7913750476d48099aaf507a79841c0d04a9db4415b14842de \ + --hash=sha256:0cf8c3b8ba93d496b2fae778039e2f5ecc7cff99df84df337ca31d8f2252896c \ + --hash=sha256:14ef29fc5f310d34fc7696426071067462c9292ed98b5ff5a27ac70a200e5470 \ + --hash=sha256:15b33fe93cedc4caaff8a0bd1eb7e3dab1c61bb22a0bf5bdfdfd97cd7da79744 \ + --hash=sha256:1b1d6a4efedd53671c793be6dd760fcf2107da3a52331ad9ea429edf0902f27a \ + --hash=sha256:1b557b29782a643420e08d75aea889462a4a8796e9a6cf5621ab05a3f7da8ef2 \ + --hash=sha256:1b71754d5b6eda54d16fbbed7fce2d8bc6c052a1b91a35c320247946ee103502 \ + --hash=sha256:1ce223652fd4ed3eb2b7f78fbea31c52314baecfac68db44037bb4167062a937 \ + --hash=sha256:1e68cdf321ad05797ee41d1d09169e09d40fdf51a725bb148bff892ce04583d7 \ + --hash=sha256:260d3692396e1895c5034f204f0db022c056f9e2ac841593a4cf9426e2a3faca \ + --hash=sha256:26e8d3ecb0ee458a9804f47f21b74845cc823fd1bb19f02272be70774f56e2a6 \ + --hash=sha256:2881416badd2a88a7a14d981c103a52a23a276a553a8aacc1346c2ff47c8dc17 \ + --hash=sha256:29b7e6716ee4ea0c59e3b241f682204105f7da084d6254ec61886508efeb43bc \ + --hash=sha256:2a7f1d03727130fc875448b65b127a9ec5d06d19d0148e7554384229706f9d1b \ + --hash=sha256:2d39b54b968f4b49b5e845758e202b1035f948b0561ff5e6385e855c96625971 \ + --hash=sha256:2e1ad3fda65ae0d93fec742a128d72e145c9c7a99ee2fcd667785d99eb25a7fe \ + --hash=sha256:3173e1e57cebb6d1de186e46b5680afbd82fd4301d7b2465beebe83ed317066d \ + --hash=sha256:3219bd9e69868e57183316ee19c84e03e8f8b5a1d1f2667e1aa8c2f91cb061ac \ + --hash=sha256:350c8348f0e76fff0a0fd6c26755d2653863279d086d3aa2c290a6a7251135dd \ + --hash=sha256:35d382625778834a7f3061b15423919aa03e4f5da34ac8e02c074e4b75ab4f84 \ + --hash=sha256:3b90b767916ac44e93a8e28ce6adf8d551e43affb512f2377c732d486ac6514e \ + --hash=sha256:3e1b35d56856f3ed326b140d3c6d9db91740f22e14b06e840fe4bb1923439a18 \ + --hash=sha256:3ebe801e0f4e56d17cd386ca6600573e3706ce1845376307f5d2cbd32149b69a \ + --hash=sha256:3f3c908bcc404c90c77d5a073e55271a0a498f4e0756e48127c35d91cf155947 \ + --hash=sha256:40d918bce2b427a0c4ba189df7a006ac0c7277c180aee4617d99e9ccaaf59e6a \ + --hash=sha256:465a0d012b3d3e4f1d6146ea019b5c11e3e87f03d1676da1cc3833462e672fb0 \ + --hash=sha256:4735a10f738cb5516905a121f32b24ce196ab82cfc1e4ba2e3ad1b371085fd46 \ + --hash=sha256:4ecdb3b6dc36e6d6e14d3a1bdc6c1057c8cbf80db04031d566eb6080ce283a48 \ + --hash=sha256:50b1b799f45da91292ffaa21a473ab3a3054fa78560e8ff67082a185274431c8 \ + --hash=sha256:54a50a9dad16b32136b2241ddea9e4df159b41247b2ce6aac0b3276a66a8f1e5 \ + --hash=sha256:5732eff8973dd995549a18ecbd8acd692ac611c5c0bb3f59fa3541ae27b33be3 \ + --hash=sha256:598e88c736f63a0efec8363f9eb34e5b5536b7b6b1821e401afcb501d881f59a \ + --hash=sha256:640fe199048f24c474ec6f3eae67c48d286de12911110437a36a87d7c89573a6 \ + --hash=sha256:66c02c187ad250513c2f4fce973ef402d22f80e0adce734ee4e4efd657b6cb64 \ + --hash=sha256:67a91c5187e1eec76a61625c77a6c8c785650f5b576ca732bd33ef58b0dff49c \ + --hash=sha256:6be67c19e0b0c56365c6a76e393b932fb0e78b3b56b711d180dd7013cb1fd984 \ + --hash=sha256:6c12dad5cd04530323e723787ff762bac749a7b256a5bece32b2243dd5c27b21 \ + --hash=sha256:71a66c1c9be66595d628467401d5976158c97888c2c9379c034e1e2312c5b4f5 \ + --hash=sha256:7274942e69b17f9cef76691bcf38f2b2d4c8a5f5dba6ec10958363dcb3308a0a \ + --hash=sha256:7547369c4392b47d30a3467fe8c3330b4f2e0f7730e45e3103d7d636678a808b \ + --hash=sha256:7a47ce5c2288702e09dc22a44d0ee6152f2c7eda97b3c8482d826a1f3cfc7da7 \ + --hash=sha256:7a61c06b334bd99bc5ae84f1eeb36bfe01400264b3c352f968c6e30a10f9d08b \ + --hash=sha256:7ad8cec81f34edf44a1c6a7edf28e7b7806dfb8886e371d95dcf789ccd4e4982 \ + --hash=sha256:7e9053f5fb4e0dfab89243079b3e217f2aea4085e4d58c5c06115fc34823707f \ + --hash=sha256:7fa18d65a213abcfbb2f6cafbb4c58863a8bd6f2103d65203c520ac117d1944b \ + --hash=sha256:81da1b229b1889f25adadc929aeb9dbc4e922bd18561b65b08dd9343cfccca84 \ + --hash=sha256:82676c2781ecf0ab23833796062786db04648b7aae8be139f6b8065e5e7b1518 \ + --hash=sha256:832c115a020e463c2f67664560449a7bea26b0c1fdd690352addad6d0a08714d \ + --hash=sha256:844a8ceb8483fefafc412f85c14f2aae2fb69567bf2a0de53cdb88b73e7c43ae \ + --hash=sha256:865cedc7c7c303df5fad14a57bc5db1d4f4f9b2b4d0a7523ddd206f00c121a16 \ + --hash=sha256:88ef7d55b7bcf3331572634c3fd0ed327d237ceb9be6066810d39020a3ebac7a \ + --hash=sha256:898be2be399c221d2671d29eed26b6b2713a02c2119168ed914e7d00ceadb56f \ + --hash=sha256:8d4f47f284bdd28629481c97b5f29ad67544fa258d9091a6ed1fda47c7347cd1 \ + --hash=sha256:92edab1e2fd6cd5ca605f57d4545b6599ced5dea0fd90b2bcdf8b247a12bd190 \ + --hash=sha256:9322b9f8656782414b37e6af884146869d46ab85158201d82bab9abbcb971dc7 \ + --hash=sha256:95db242754c21a88a79e01504912e537808504465974ebb92931cfca2510469e \ + --hash=sha256:963a08f3bebd8b75ac57661045402da15991468a621f014be54e50f53a58d19e \ + --hash=sha256:96fbe82a58cdb2f872fa5d87dedc8477a12993626c446de794ea025bbda625ea \ + --hash=sha256:99cfa69813d79492f0e5d52a20fd18395bc82e671d5d40bd5a91d13e75e468e8 \ + --hash=sha256:9c79f57faa25d97900bfb119480806d783fba83cd09ee0b33c17623935b05fa3 \ + --hash=sha256:9e5825ba2c9998375530504578fd4d5d1059d09621a02065d1b6bfc41a8e05ab \ + --hash=sha256:9fe11467c42c133f38d42289d0861b6b4f9da31e8087ca2c0d7ebb4543625526 \ + --hash=sha256:a1778532b978d2536e79c05dac2d8cd857f6c55cd0c95ace5b03740824e0e2f1 \ + --hash=sha256:a387225a67f619bf16bd504c37655930f910eb03675730fc2ad69d3d8b5e7e92 \ + --hash=sha256:a56ef534b66a749759ebd091c19c03ef81eb8cd96f0d1d16b59127eaf1b97a12 \ + --hash=sha256:aa47441fa3026543513139cb8926a92a8e305ee9c71a6209ef7a97d91640ea03 \ + --hash=sha256:ac27a70bda257ae3f380ec8310b0a06680236bea547756c277b5dfe55a2452a8 \ + --hash=sha256:acec55bb7c90f1dfc476126f9711a8e81c9af7fb617409a9ee2953115343f08d \ + --hash=sha256:adedc4a67e15327dfdd04884873c6d5a01d3e3b6f61406f99b1ed4865a2f6d28 \ + --hash=sha256:af43b8711a8264bb4e7d6d9a6d004c3a2019c04c01127a868709ec29962b6036 \ + --hash=sha256:b232029d100d393ae3c603c8ffd7e3fe6f798c5e28ddca5feabb8e8fdb732997 \ + --hash=sha256:b35c13ce241abdd44cb8ca70683f20c0c079728a36a996297adb5334adfc1c44 \ + --hash=sha256:b63daa43d82f0cdabf98dee215b375b4058cce72871fd07934f179885aad16e8 \ + --hash=sha256:b908d1a7b28bc72dfb743be0d4d3f8931f8309f810af66c906ae6cd4127c93cb \ + --hash=sha256:ba76177fd318ab7b3b9bf6522be5e84c2ae798754b6cc028665490f6e66b5533 \ + --hash=sha256:bba6e7e6cfe1e6cb6eb0b7c2736a6059461de1fa2c0ad26cf845de6c078d16c8 \ + --hash=sha256:c0d6770111d1879881432f81c369de5cde6e9467be7c682a983747ec800544e2 \ + --hash=sha256:c16ab1ef7bb55651f5836e8e62db1f711d55b82ea08c3b8083ff037157171a69 \ + --hash=sha256:c1702888c9f3383cc2f09eb3e88b8babf5965a54afb79649458ec7c3c7a63e96 \ + --hash=sha256:c25332657dee6052ca470626f18349fc1fe8855a56218e19bd7a8c6ad4952c49 \ + --hash=sha256:c8565e3cdc1808b1a34714b553b262c5de5fbda202285782173ec137fd13709f \ + --hash=sha256:cf9cba6f5b78a2071ec6fb1e7bd39acf35071d90a81231d67e92d637776a6a63 \ + --hash=sha256:d206a36b4140fbb5373bf1eb73fb9de589bb06afd0d22376de23c5e91d0ab35f \ + --hash=sha256:d2d085ded05278d1c7f65560aae97b3160aeb2ea2c0b3e26204856beccb60888 \ + --hash=sha256:d8c05b1dfb61af28ef37624385b0029df902ca896a639881f594060b30ffc9a7 \ + --hash=sha256:e310f77e41941c13340a95976fe66a8a95b01e783d430eeaf7a2f87e0a57dd0a \ + --hash=sha256:e7c0af964e0b4e3412a0ebf341ea26ec767fa0b4cf81abb5e897c9338b5ad6a3 \ + --hash=sha256:e80a28f2b150774844c8b454dd288be90d76ba6109670fe33d7ff54d96eb5cb8 \ + --hash=sha256:e813da3d2d865e9793ef681d3a6b66fa4b7c19244a45b817d0cceda67e615990 \ + --hash=sha256:e85190da223337a6b7431d92c799fca3e2982abd44e7b8dec69938dcc81c8e9e \ + --hash=sha256:e99befa0b48f3cd293dafeacdd0d191804d105d279e0b387a32054c1180f3161 \ + --hash=sha256:eda5a6d042c698e28bda2507a89b16555b9aa954ef1d750e1c20473481aff675 \ + --hash=sha256:ef87b8ab2704da227e83a246356a2b179ef826f550f794b2c52cddb4efbd0196 \ + --hash=sha256:f16dace5e4d3596eaeb8af334b4d2c820d34b8278da633ce4a00020b2eac981c \ + --hash=sha256:f8d635cafbbb0c61327f942df2e3f474dde1cff16c3cd0580564774eaba1ee13 \ + --hash=sha256:fc1530af5c3c275b8524f2e24841cbe2599d74462455e9bae5109e9ff42e9361 \ + --hash=sha256:ff09cd8c5eec3b9d02d2408db41be150d8891c5566addce57513bf546e3d6c6d +certifi==2026.7.22 \ + --hash=sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775 \ + --hash=sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55 +charset-normalizer==3.4.9 \ + --hash=sha256:0327fcd59a935777d83410750c50600ee9571af2846f71ce40f25b13da1ef380 \ + --hash=sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62 \ + --hash=sha256:04ce310cb89c15df659582aee80a0603788732a5e017d5bd5c81158106ce249c \ + --hash=sha256:0d861473f743244d349b50f850d10eb87aeb22bbdcc8e64f79273c94af5a8226 \ + --hash=sha256:0e94703ec9684807f20cfb5eed95c70f67f2a8f21ad620146d7b5a13677b93e5 \ + --hash=sha256:0fa1aec2d32bcc03c8fa0f6f1712caad1adc38509f31142112e5c9daf5b9c833 \ + --hash=sha256:16b65ea0f2465b6fb52aa22de5eca612aa964ddfec00a912e26f4656cbef890b \ + --hash=sha256:16d10d789dd9bcca1173c95af82c58433122564b7bc39385124be735a35cbe99 \ + --hash=sha256:19ac87f93086ce37b86e098888555c4b4bc48102279bae3350098c0ed664b501 \ + --hash=sha256:1d22856ffbe153a602df38e4a5464f0b748a54002e0d69ac6d2ad0a197cc99ec \ + --hash=sha256:21e764fd1e70b6a3e205a0e46f3051701f98a8cb3fad66eeb80e48bb502f8698 \ + --hash=sha256:231ddcbb35e2ff8973e1365db41fe0572662893b99a05deb183b68ad4c0c8bd4 \ + --hash=sha256:253a4a220747e8b5faf57ec320c4f5efb0cef05f647420bf267143ec15dba10a \ + --hash=sha256:280081916dc341820640489a66e4696049401ef1cf6dd672f672e70ad915aca3 \ + --hash=sha256:2a441ea71902098ffe78c5abe6c494f44160b4af614ed16c3d9a3b1d17fd8ee2 \ + --hash=sha256:304b13570067b2547562e308af560b3963857b1fa90bd6afd978130130fe2d6a \ + --hash=sha256:32286a2c8d167e897177b673176c1e3e00d4057caf5d2b64eef9a3666b03018e \ + --hash=sha256:33bdcc2a32c0a0e861f60841a512c8acc658c87c2ac59d89e3a46dacf7d866e4 \ + --hash=sha256:375b83ed0aecfce76c16d198fbc21f3b11b337d68662bea0a995046682a11419 \ + --hash=sha256:3c09a49d6cde137258beb3d551994a2927fd35ad5cf96aed573f61bbd67c5f84 \ + --hash=sha256:3d92613ec25e43b05f042302531ec0f00b8445190e43325880cbd6ab7c2581da \ + --hash=sha256:40a126142a56b2dfc0aacbad1de8310cbf60da7656db0e6b16eebd48e3e93519 \ + --hash=sha256:416c229f77e5ea25b3dfd4b582f8d73d7e43c22320302b9ab128a2d3a0b38efe \ + --hash=sha256:432786d3561e69aeeae6c7e8648964ce0ad05736120135601f87ac26b9c83381 \ + --hash=sha256:43b9e366a31fdd1c87d0eb08f579b4a82b723ea54338f040d6b4e518a026ea29 \ + --hash=sha256:440eede837960000d74978f0eba527be106b5b9aee0daf779d395276ed0b0614 \ + --hash=sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0 \ + --hash=sha256:476743fe6dfe14a2da12e3ac79125dc84a3b2cf8094369a47a1529b0cd8549fe \ + --hash=sha256:4773092f8019072343a7447203308b176e10199920eb02d6195e81bbb3274c29 \ + --hash=sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0 \ + --hash=sha256:4d1c96a7a18b9690a4d46df09e3e3382406ae3213727cd1019ebade1c4a81917 \ + --hash=sha256:51307f5c71007673a2bf8232ad973483d281e74cb99c8c5a990af1eefa6277d9 \ + --hash=sha256:51447e9aa2684679af07ca5021c3db526e0284347ebf4ffcec1154c3350cfe32 \ + --hash=sha256:58150c9f9b9a552505912d182ccdf26f6396fb6094816ceebcbb20eecabaed94 \ + --hash=sha256:5b10cd92fc5c498b35a8635df6d5a100207f88b63a4dc1de7ef9a548e1e2cd63 \ + --hash=sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd \ + --hash=sha256:609b3ba8fcc0fb5ab7af00719d0fb6ad0cb518e48e7712d12fd68f1327951198 \ + --hash=sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde \ + --hash=sha256:611057cc5d5c0afc743ba8be6bd828c17e0aaa8643f9d0a9b9bb7dea80eb8012 \ + --hash=sha256:6366a16e1a25018694d6a5d784d09b046edc9eac40ea2b54065c3052672516a1 \ + --hash=sha256:65a7ff3f705e57d392f7261b6d0550fe137c3019477431f1c355e0db0a7d3e15 \ + --hash=sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b \ + --hash=sha256:67830fc78e67501f47bb950471b2dcb9b35b140084429318e862895a8e89c993 \ + --hash=sha256:68ce9f4d6b26d5ccbf7fd4459bf75f74a0a146677ebba80597df60cbdb20e6f4 \ + --hash=sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5 \ + --hash=sha256:69b157c5d3292bcd443faca052f3096f637f1e074b98212a933c074ae23dc3b8 \ + --hash=sha256:75286256590a6320cf106a0d28970d3560aad9ee09aa7b34fb40524792436d35 \ + --hash=sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642 \ + --hash=sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2 \ + --hash=sha256:79580094b00d1789d1f93ea55bc43cb2f611910c72235b7657f3482ddcc1b22d \ + --hash=sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9 \ + --hash=sha256:83aed2c10721ddd90f68140685391b50811a880af20654c59af6b6c66c40513c \ + --hash=sha256:84fd18bcc17526fc2b3c1af7d2b9217d32c9c04448c16ec693b9b4f1985c3d33 \ + --hash=sha256:871ff67ea1aad4dfd91736464934d56b32dac49f9fbe16cddba36198a7b3a0db \ + --hash=sha256:898f0e9068ca27d37f8e83a5b962821df851532e6c4a7d615c1c033f9da6eedf \ + --hash=sha256:8a79d9f4d8001473a30c163556b3c3bfebec837495a412dde78b51672f6134f9 \ + --hash=sha256:8c041122946b7ba21bb32c45b1aa57b1be35527690aeb3c5c234521085632eee \ + --hash=sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84 \ + --hash=sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44 \ + --hash=sha256:920079c3f7456fa213e0829ed2073aaa727fd39d889ead5b4f35d0de5460d04f \ + --hash=sha256:93d59d504b230e83c7a843251681959a0b6a9cd76f6e146ce1b8a80eb8739af9 \ + --hash=sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9 \ + --hash=sha256:9b8e0f3107e2200b76f6054de99016eac3ee6762713587b36baaa7e4bd2ae177 \ + --hash=sha256:9bb41182d93ea91f60b4bc8fbf4c820c69ef8a12ab2d917f3f1834f1acad07e8 \ + --hash=sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b \ + --hash=sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39 \ + --hash=sha256:a4cfde78a9f2880208d16a93b795726a3017d5977e08d1e162a7a31322479c41 \ + --hash=sha256:a4fbdde9dd4a9ce5fd52c2b3a347bb50cc89483ef783f1cb00d408c13f7a96c0 \ + --hash=sha256:aa99adc8f081b475a12843953db36831eaf83ec33eb46a90629ca6a5de45a616 \ + --hash=sha256:ac351b3b8014eead140e77e9717e2992c6bbe30b63bc3422422eb84865412e3d \ + --hash=sha256:ad41ba96094304aa090f5a30cb6e4fb3b3f1c264c523394b4c39bbacc4dc92ba \ + --hash=sha256:b5314963fce9b0b12743891de876e724997864ee22aa496f903f426c7e2fa5b2 \ + --hash=sha256:bcf74c1df76758a395bf0af608c04c82257523f55c9868b334f06270d0f2112b \ + --hash=sha256:bd47ba7fc3ca94896759ea0109775132d3e7ab921fbf54038e1bab2e46c313c9 \ + --hash=sha256:c0323c9daef75ef2e5083624b4585018a0c9d5e3b40f607eed81a311270b934b \ + --hash=sha256:c1225416b463483160e4af85d5fc3a9690ccb53fd4b1865a6437825f5ede3209 \ + --hash=sha256:c1c948747b03be832dceed96ca815cef7360de9aa19d37c730f8e3f6101aca48 \ + --hash=sha256:c25fe15c70c59eb7c5ce8c06a1f3fa1da0ecc5ea1e7a5922c40fd2fa9b0d5046 \ + --hash=sha256:cc1b0fff8ead343dae06305f954eb8468ba0ec1a97881f42489d198e4ce3c632 \ + --hash=sha256:cd6280cf040f233bd7d3407b743b4b4c74f70e8e1c4199cb112a62c941c0772a \ + --hash=sha256:cd6c3d4b783c556fa00bf540854e42f135e2f256abd29669fcd0da0f2dec79c2 \ + --hash=sha256:d4d6fcde76f94f5cb9e43e9e9a61f16dacefd228cbbf6f1a09bd9b219a92f1a1 \ + --hash=sha256:ddf4af30b417d9fe16481e9b81c27ab2a7cde1ff7ba3e85653b02db7d145dc7b \ + --hash=sha256:df115d4d83168fdf2cae48ef1ff6d1cb4c466364e30861b37121de0f3bf1b990 \ + --hash=sha256:df7276909358e5635ae203673ab7e509ddd224225a8d6b0790bf13eb2bde1cc5 \ + --hash=sha256:e4fd89cc178bced6ad29cb3e6dd4aa63fa5017c3524dbd0b25998fb64a87cc8b \ + --hash=sha256:e9701d0049d92c16703a42771b98d560b95248949f23f8cf7b4eddd201814fb9 \ + --hash=sha256:ee2f2a527e3c1a6e6411eb4209642e138b544a2d72fe5d0d76daf77b24063534 \ + --hash=sha256:f7fb7d750cfa0a070d2c24e831fd3481019a60dd317ea2b39acbcebc08b6ed81 \ + --hash=sha256:f840ed6d8ecba8255df8c42b87fadeda98ddfc6eeec05e2dc66e26d46dd6f58a \ + --hash=sha256:f86c6358749bd4fda175388691e3ba8c46e24c5347d0afd20f9b7edfc9faf07d \ + --hash=sha256:fa36ec09ef71d158186bc79e359ff5fdd6e7996fe8ab638f00d6b93139ba4fcf \ + --hash=sha256:fe2c7201c642b7c308f1675355ad7ff7b66acfe3541625efe5a3ad38f29d6115 +cloudpathlib==0.24.0 \ + --hash=sha256:b1c51e2d2ec7dc4fed6538991f4aea849d6cf11a7e6b9069f86e461aa1f9b5b4 \ + --hash=sha256:c521a984e77b47e656fe78e20a7e3e260e0ab45fc69e33ac01094227c979e34a +comm==0.2.3 \ + --hash=sha256:2dc8048c10962d55d7ad693be1e7045d891b7ce8d999c97963a5e3e99c055971 \ + --hash=sha256:c615d91d75f7f04f095b30d1c1711babd43bdc6419c1be9886a85f2f4e489417 +contourpy==1.3.3 \ + --hash=sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69 \ + --hash=sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc \ + --hash=sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880 \ + --hash=sha256:0bf67e0e3f482cb69779dd3061b534eb35ac9b17f163d851e2a547d56dba0a3a \ + --hash=sha256:0c1fc238306b35f246d61a1d416a627348b5cf0648648a031e14bb8705fcdfe8 \ + --hash=sha256:13b68d6a62db8eafaebb8039218921399baf6e47bf85006fd8529f2a08ef33fc \ + --hash=sha256:15ff10bfada4bf92ec8b31c62bf7c1834c244019b4a33095a68000d7075df470 \ + --hash=sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5 \ + --hash=sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263 \ + --hash=sha256:1fd43c3be4c8e5fd6e4f2baeae35ae18176cf2e5cced681cca908addf1cdd53b \ + --hash=sha256:22e9b1bd7a9b1d652cd77388465dc358dafcd2e217d35552424aa4f996f524f5 \ + --hash=sha256:23416f38bfd74d5d28ab8429cc4d63fa67d5068bd711a85edb1c3fb0c3e2f381 \ + --hash=sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3 \ + --hash=sha256:2a2a8b627d5cc6b7c41a4beff6c5ad5eb848c88255fda4a8745f7e901b32d8e4 \ + --hash=sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e \ + --hash=sha256:322ab1c99b008dad206d406bb61d014cf0174df491ae9d9d0fac6a6fda4f977f \ + --hash=sha256:33c82d0138c0a062380332c861387650c82e4cf1747aaa6938b9b6516762e772 \ + --hash=sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286 \ + --hash=sha256:3519428f6be58431c56581f1694ba8e50626f2dd550af225f82fb5f5814d2a42 \ + --hash=sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301 \ + --hash=sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77 \ + --hash=sha256:451e71b5a7d597379ef572de31eeb909a87246974d960049a9848c3bc6c41bf7 \ + --hash=sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411 \ + --hash=sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1 \ + --hash=sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9 \ + --hash=sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a \ + --hash=sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b \ + --hash=sha256:51e79c1f7470158e838808d4a996fa9bac72c498e93d8ebe5119bc1e6becb0db \ + --hash=sha256:556dba8fb6f5d8742f2923fe9457dbdd51e1049c4a43fd3986a0b14a1d815fc6 \ + --hash=sha256:598c3aaece21c503615fd59c92a3598b428b2f01bfb4b8ca9c4edeecc2438620 \ + --hash=sha256:5ed3657edf08512fc3fe81b510e35c2012fbd3081d2e26160f27ca28affec989 \ + --hash=sha256:626d60935cf668e70a5ce6ff184fd713e9683fb458898e4249b63be9e28286ea \ + --hash=sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67 \ + --hash=sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5 \ + --hash=sha256:66c8a43a4f7b8df8b71ee1840e4211a3c8d93b214b213f590e18a1beca458f7d \ + --hash=sha256:6afc576f7b33cf00996e5c1102dc2a8f7cc89e39c0b55df93a0b78c1bd992b36 \ + --hash=sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99 \ + --hash=sha256:709a48ef9a690e1343202916450bc48b9e51c049b089c7f79a267b46cffcdaa1 \ + --hash=sha256:70f9aad7de812d6541d29d2bbf8feb22ff7e1c299523db288004e3157ff4674e \ + --hash=sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b \ + --hash=sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8 \ + --hash=sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d \ + --hash=sha256:929ddf8c4c7f348e4c0a5a3a714b5c8542ffaa8c22954862a46ca1813b667ee7 \ + --hash=sha256:92d9abc807cf7d0e047b95ca5d957cf4792fcd04e920ca70d48add15c1a90ea7 \ + --hash=sha256:95b181891b4c71de4bb404c6621e7e2390745f887f2a026b2d99e92c17892339 \ + --hash=sha256:9e999574eddae35f1312c2b4b717b7885d4edd6cb46700e04f7f02db454e67c1 \ + --hash=sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659 \ + --hash=sha256:a22738912262aa3e254e4f3cb079a95a67132fc5a063890e224393596902f5a4 \ + --hash=sha256:ab2fd90904c503739a75b7c8c5c01160130ba67944a7b77bbf36ef8054576e7f \ + --hash=sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20 \ + --hash=sha256:afe5a512f31ee6bd7d0dda52ec9864c984ca3d66664444f2d72e0dc4eb832e36 \ + --hash=sha256:b08a32ea2f8e42cf1d4be3169a98dd4be32bafe4f22b6c4cb4ba810fa9e5d2cb \ + --hash=sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d \ + --hash=sha256:b2e8faa0ed68cb29af51edd8e24798bb661eac3bd9f65420c1887b6ca89987c8 \ + --hash=sha256:b7301b89040075c30e5768810bc96a8e8d78085b47d8be6e4c3f5a0b4ed478a0 \ + --hash=sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b \ + --hash=sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7 \ + --hash=sha256:ca658cd1a680a5c9ea96dc61cdbae1e85c8f25849843aa799dfd3cb370ad4fbe \ + --hash=sha256:cbedb772ed74ff5be440fa8eee9bd49f64f6e3fc09436d9c7d8f1c287b121d77 \ + --hash=sha256:cd5dfcaeb10f7b7f9dc8941717c6c2ade08f587be2226222c12b25f0483ed497 \ + --hash=sha256:cf9022ef053f2694e31d630feaacb21ea24224be1c3ad0520b13d844274614fd \ + --hash=sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1 \ + --hash=sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216 \ + --hash=sha256:d304906ecc71672e9c89e87c4675dc5c2645e1f4269a5063b99b0bb29f232d13 \ + --hash=sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae \ + --hash=sha256:e74a9a0f5e3fff48fb5a7f2fd2b9b70a3fe014a67522f79b7cca4c0c7e43c9ae \ + --hash=sha256:ea37e7b45949df430fe649e5de8351c423430046a2af20b1c1961cae3afcda77 \ + --hash=sha256:f64836de09927cba6f79dcd00fdd7d5329f3fccc633468507079c829ca4db4e3 \ + --hash=sha256:fd6ec6be509c787f1caf6b247f0b1ca598bef13f4ddeaa126b7658215529ba0f \ + --hash=sha256:fd907ae12cd483cd83e414b12941c632a969171bf90fc937d0c9f268a31cafff \ + --hash=sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9 \ + --hash=sha256:fde6c716d51c04b1c25d0b90364d0be954624a0ee9d60e23e850e8d48353d07a +cuda-bindings==13.3.1 \ + --hash=sha256:04436a9364059c84b8f9636f359eccda1cf814341f5b670c71d80d2f79dbc708 \ + --hash=sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86 \ + --hash=sha256:18c8c167c8907b8f02531ca810534315c458dabef31f7965095619bf647b9202 \ + --hash=sha256:1ab2f74ed65bfef4163ba07a8db16f1085e0729291db12a2423aff84ee8278b8 \ + --hash=sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7 \ + --hash=sha256:36febb7c1079d68a981dbbd8d5a67235b399802b82075c9388624719607e52b9 \ + --hash=sha256:507b0e19e7f934c5e30f30f0244ad70a75812619a7d3a0d742543caae1bd50f1 \ + --hash=sha256:61120b5e4f4a63f67efd7e7396914cb9ef871bb1f0021e990fb70277be240a4d \ + --hash=sha256:8de12ef60bf40756852cb62bbb40460609269f6ece522903d1cc93d73a3ececb \ + --hash=sha256:9851b0caa8bfd3bc6fa054eaf57bea7c8e9c3a62db2d2621224677f49f3c53d0 \ + --hash=sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf \ + --hash=sha256:b134dd8c5c66ae4c4ad814f7aee88fd215353c077010cbc47e3b55ed35ec9eff \ + --hash=sha256:c0c4b1a995098c46695c24257a342dc97d6e6d3f3050b944c9f43bd26d734051 \ + --hash=sha256:c3c772dfff49681541d59630c90f858e173ac926b9c593a2b7123f2a1043cc76 \ + --hash=sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474 \ + --hash=sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49 \ + --hash=sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a \ + --hash=sha256:efd4c814d311ec08c981f6dded1dbe7d4b371067ee4f6c14cccec4bde9590f80 +cuda-pathfinder==1.6.0 \ + --hash=sha256:1503af579d8379c24bdd65528379bc57039b0455be9f5f9686cf8e473a1fce51 +cuda-toolkit==13.0.3.0 \ + --hash=sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f +cycler==0.12.1 \ + --hash=sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30 \ + --hash=sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c +decorator==5.3.1 \ + --hash=sha256:4cbcdd55a6efadb9dbea26b858f4fb3264567b52d69ca0d25b721b553f60ea82 \ + --hash=sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c +dna-features-viewer==3.1.5 \ + --hash=sha256:4740b018f7c45427054cb1403ff8b11209bcbe4ca9d5b77f28f14ed5ecb144c4 \ + --hash=sha256:a8a383d5340d35979be50224f12f4f426ab004394d9b9a33c99d59404469fb7a +einops==0.8.2 \ + --hash=sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193 \ + --hash=sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827 +executing==2.2.1 \ + --hash=sha256:3632cc370565f6648cc328b32435bd120a1e4ebb20c77e3fdde9a13cd1e533c4 \ + --hash=sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017 +filelock==3.32.0 \ + --hash=sha256:7be2ad23a14607ccc71808e68fe30848aeace7058ace17852f68e2a68e310402 \ + --hash=sha256:d396bea984af47333ef05e50eae7eff88c84256de6112aea0ec48a233c064fe3 +fonttools==4.63.0 \ + --hash=sha256:032038247a96c1690f9f31e377c389383c902531b085aa4e4dabd6f57f870e69 \ + --hash=sha256:063e08bd17bd5a90127a14123de0d6a952dbc847695fd98b63c043d58057f90c \ + --hash=sha256:0c18358a155d75034911c5ee397a5b44cd19dd325dbb8b35fb60bf421d6a72ac \ + --hash=sha256:0eac00b9118c3c2f87d272e45341871c5b3066baa3c86897fa634a7c3fb59096 \ + --hash=sha256:1e874792a8212b44583ea02189d9e693906b2f78b261f372f95d6c563210ac1d \ + --hash=sha256:22135da48a348785c5e2d5d2d9d6bec5ed44adacbaeb9db12d9493bf6c6bfa68 \ + --hash=sha256:22693918177bd9ceabec4736d338045f357769416fc6b0b2508eefef75b08616 \ + --hash=sha256:27fdc65af8da6f88b9c6121c47a464cbe359fcfff7ff6fc2d37a1f395d755b78 \ + --hash=sha256:2b8ae05d9eacf6081414d759c0a352769ac28ce31280d6bb8e77b03f9e3c449f \ + --hash=sha256:2c14b4fd138c4bafcca294765c547914e1aa431ae1ca94ab99d8db08c958bd3b \ + --hash=sha256:308f957cdeaf8abe4e5f2f124902ef405448af92c90f80e302a3b771c2e6116b \ + --hash=sha256:37dd23e621e3b0aef1baa70a303b80aaf38449632cfc8fd2a55fb285bbccfc02 \ + --hash=sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d \ + --hash=sha256:51394295f1a51de8b5f30bdb1e1b9a4231536c7064ef5c6e211eec19fa36036f \ + --hash=sha256:58dc6bb86a78d782f00f9190ca02c119cf5bbe2807536e361e18d42019f877d8 \ + --hash=sha256:59ac449f8cca9b4ffa08d2e7bbadad87ce710d69d1eda5c3c1ce579baa987272 \ + --hash=sha256:6b2248c5decb223562f7902ff6325077a073f608ee8e33e88ad88db734eb9f49 \ + --hash=sha256:6d4741eb179121cab9eea4cb2393d24492373a260d7945006358c08cfbf45419 \ + --hash=sha256:6db5140a60a5d731d21ec076745b40a310607731b0a565b50776393188649001 \ + --hash=sha256:6e528da43bc3791085f8cb6141b1d13e459226790240340fcbb4625649238b03 \ + --hash=sha256:796f27556dbe094c4824f75ca85267e4df776c79036c8441469a4df37038c196 \ + --hash=sha256:79cdc9f567aec74a72918fd060283911406750cbc9fd28c1316023deb6ce31a9 \ + --hash=sha256:7d76edbff9014094dbf03bd2d074709dfa6ec7aba13d838c937a2b33d2d6a86e \ + --hash=sha256:7d782fac32985914c351556f68ac0855391572bcd87de50e05970d3cd4c96fc5 \ + --hash=sha256:7dd683fef0663e9f0f45cf541d788d24caa3ec9db50796b588e1757d8b3bc007 \ + --hash=sha256:85be818f5506e8a7753153def2c9550178f0ecae6a47b5e0e8dbb23f7cc90380 \ + --hash=sha256:948428a275741f0b64b113c955425a953314f4b9ab9997f73a72c83e68e569c8 \ + --hash=sha256:9ced0bd02ac751dd6319b0da88aaef24414e3b0dbc32bb4f24944821a3741a27 \ + --hash=sha256:9e12f105d2b6342c559c298afb674006bb2893afc7102dcf8a1b55b0486b4e40 \ + --hash=sha256:a8b33a82979e0a6a34ff435cc81317be1f95ec1ebb7a3a2d1c8a6a54f02ae44e \ + --hash=sha256:a9faff9e0c1f76f9fd55899d2ce785832efebab37eb8ae13995853aef178bef0 \ + --hash=sha256:af2fd1664d00a397d75f806985ddb36282091c2131a73a6485c23b4a34722263 \ + --hash=sha256:afefc1ed0a59785a7fb06ea7e1678e849c193e1e387db783579bc7b3056fcfcb \ + --hash=sha256:b1cd75a03ad8cb5bc40c90bfde68c0c47de423aa19e5c0f362b43520645eea94 \ + --hash=sha256:ba04cb5891d4c0c21b6da95eda8d7b090021508a294fff33464fc7d241e0856b \ + --hash=sha256:bf00f21eb5fb721dbaf73d1e9da6d02a1af7768f2ebcf9798be98beab8ba90f6 \ + --hash=sha256:c0425b277a59cff3d80ca42162a8de360f318438a2ac83570842a678d826d579 \ + --hash=sha256:c1aaa4b9c75798400ac043ce04d74e7830376c85095a5a6ed7cba2f17a266bf4 \ + --hash=sha256:c2a2a42198b696a6f48fad91709afb55176e66a5e566131219dba372fb7f8c59 \ + --hash=sha256:caeb583deeb5168e694b65cda8b4ee62abedfa66cf88488734466f2366b9c4e0 \ + --hash=sha256:cb014d58140a38135f16064c74c652ed57aa0b75cbf8bb59cac821f7edb5334e \ + --hash=sha256:ccf41f2efdf56994d22d73bef4ced1052161958169428d06ba9724ea9e9a64be \ + --hash=sha256:cd7e9857e5e63738b9d9fd707bc1f59c8b09e5177726d23664db393c59bb08bd \ + --hash=sha256:d76ac49f929aecaf82d83250b8347e099d7aecba0f4726c1d9b6df3b8bb5fe18 \ + --hash=sha256:d7e5c9973aa04c95650c96e5f5ad865fbf42d62079163ecfab1e01cbc2504c22 \ + --hash=sha256:dcf076a4474fe0d7367e5bbf5b052c7284fa1feca729c04176ce513521afd8a0 \ + --hash=sha256:e3297a6a4059b4acc3a1e9a8b04741f240a80044eef08ebd32e8b5bcdddce75b \ + --hash=sha256:ee08ebfa58f6e1aeff5697ab9582105bb620008c1caafb681e4c557e7483027b \ + --hash=sha256:ef3048ef05dbb552b89817713d9cac912e00d0fde4a3105c00d29e52e10c89af \ + --hash=sha256:fd1e3094f42d806d3d7c79162fc59e5910fcbe3a7360c385b8da969bc4493745 +fsspec==2026.6.0 \ + --hash=sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1 \ + --hash=sha256:f5bac145310fe30e16e1471bd6840b2d990d609e872251d7e674241822abf01a +h11==0.16.0 \ + --hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \ + --hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86 +hf-xet==1.5.2 \ + --hash=sha256:045f84440c55cdeb659cf1a1dd48c77bcd0d2e93632e2fea8f2c3bdee79f38ed \ + --hash=sha256:1da28519496eb7c8094c11e4d25509b4a468457a0302d58136099db2fd9a671d \ + --hash=sha256:4a5ecb9cda8512ba2aa8ee5d37c87a1422992165892d653098c7b90247481c3b \ + --hash=sha256:580e59e29bf37aece1f2b68537de1e3fb04f43a23d910dcf6f128280b5bfbba4 \ + --hash=sha256:6395cfe3c9cbead4f16b31808b0e67eac428b66c656f856e99636adaddea878f \ + --hash=sha256:73044bd31bae33c984af832d19c752a0dffb67518fee9ddbd91d616e1101cf47 \ + --hash=sha256:7db73c810500c54c6760be8c39d4b2e476974de85424c50063efc22fdda13025 \ + --hash=sha256:8764488197c1d7b1378c8438c18d2eea902e150dbca0b0f0d2d32603fb9b5576 \ + --hash=sha256:8d7446f72abbf7e01ca5ff131786bc2e74a56393462c17a6bf1e303fbab81db4 \ + --hash=sha256:bee28c619622d36968056532fd49cf2b35ca75099b1d616c31a618a893491380 \ + --hash=sha256:cde8cd167126bb6109b2ceb19b844433a4988643e8f3e01dd9dd0e4a34535097 \ + --hash=sha256:d6f9c58549407b84b9a5383afd68db0acc42345326a3159990b36a5ca8a20e4e \ + --hash=sha256:db78c39c83d6279daddc98e2238f373ab8980685556d42472b4ec51abcf03e8c \ + --hash=sha256:e396ab0faf6298199ad7a95305c3ca8498cb825978a6485be6d00587ee4ec577 \ + --hash=sha256:ecf63d1cb69a9a7319910f8f83fcf9b46e7a32dfcf4b8f8eeddb55f647306e65 \ + --hash=sha256:f922b8f5fb84f1dd3d7ab7a1316354a1bca9b1c73ecfc19c76e51a2a49d29799 \ + --hash=sha256:fd3add255549e8ef58fa35b2e42dc016961c050600444e7d77d030ba6b57120e +httpcore==1.0.9 \ + --hash=sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55 \ + --hash=sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8 +httpx==0.28.1 \ + --hash=sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc \ + --hash=sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad +huggingface-hub==0.36.2 \ + --hash=sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a \ + --hash=sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270 +idna==3.18 \ + --hash=sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2 \ + --hash=sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848 +iniconfig==2.3.0 \ + --hash=sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730 \ + --hash=sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12 +ipython==9.15.0 \ + --hash=sha256:515ad9c3cdf0c932a5a9f6245419e8aba706b7bd03c3e1d3a1c83d9351d6aa6e \ + --hash=sha256:da2819ce2aa83135257df830660b1176d986c3d2876db24df01974fa955b2756 +ipython-pygments-lexers==1.1.1 \ + --hash=sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81 \ + --hash=sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c +ipywidgets==8.1.8 \ + --hash=sha256:61f969306b95f85fba6b6986b7fe45d73124d1d9e3023a8068710d47a22ea668 \ + --hash=sha256:ecaca67aed704a338f88f67b1181b58f821ab5dc89c1f0f5ef99db43c1c2921e +jedi==0.20.0 \ + --hash=sha256:7bdd9c2634f56713299976f4cbd59cb3fa92165cc5e05ea811fb253480728b67 \ + --hash=sha256:c3f4ccbd276696f4b19c54618d4fb18f9fc24b0aef02acf704b23f487daa1011 +jinja2==3.1.6 \ + --hash=sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67 +jmespath==1.1.0 \ + --hash=sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d \ + --hash=sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64 +joblib==1.5.3 \ + --hash=sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713 \ + --hash=sha256:8561a3269e6801106863fd0d6d84bb737be9e7631e33aaed3fb9ce5953688da3 +jupyterlab-widgets==3.0.16 \ + --hash=sha256:423da05071d55cf27a9e602216d35a3a65a3e41cdf9c5d3b643b814ce38c19e0 \ + --hash=sha256:45fa36d9c6422cf2559198e4db481aa243c7a32d9926b500781c830c80f7ecf8 +kiwisolver==1.5.0 \ + --hash=sha256:012b1eb16e28718fa782b5e61dc6f2da1f0792ca73bd05d54de6cb9561665fc9 \ + --hash=sha256:01808c6d15f4c3e8559595d6d1fe6411c68e4a3822b4b9972b44473b24f4e679 \ + --hash=sha256:0255a027391d52944eae1dbb5d4cc5903f57092f3674e8e544cdd2622826b3f0 \ + --hash=sha256:0b85aad90cea8ac6797a53b5d5f2e967334fa4d1149f031c4537569972596cb8 \ + --hash=sha256:0bf3acf1419fa93064a4c2189ac0b58e3be7872bf6ee6177b0d4c63dc4cea276 \ + --hash=sha256:0c50b89ffd3e1a911c69a1dd3de7173c0cd10b130f56222e57898683841e4f96 \ + --hash=sha256:0cbe94b69b819209a62cb27bdfa5dc2a8977d8de2f89dfd97ba4f53ed3af754e \ + --hash=sha256:0df54df7e686afa55e6f21fb86195224a6d9beb71d637e8d7920c95cf0f89aac \ + --hash=sha256:0e3aafb33aed7479377e5e9a82e9d4bf87063741fc99fc7ae48b0f16e32bdd6f \ + --hash=sha256:12e91c215a96e39f57989c8912ae761286ac5a9584d04030ceb3368a357f017a \ + --hash=sha256:1465387ac63576c3e125e5337a6892b9e99e0627d52317f3ca79e6930d889d15 \ + --hash=sha256:16b85d37c2cbb3253226d26e64663f755d88a03439a9c47df6246b35defbdfb7 \ + --hash=sha256:1b0feb50971481a2cc44d94e88bdb02cdd497618252ae226b8eb1201b957e368 \ + --hash=sha256:1d49a49ac4cbfb7c1375301cd1ec90169dfeae55ff84710d782260ce77a75a02 \ + --hash=sha256:1d9daea4ea6b9be74fe2f01f7fbade8d6ffab263e781274cffca0dba9be9eec9 \ + --hash=sha256:1dd9b0b119a350976a6d781e7278ec7aca0b201e1a9e2d23d9804afecb6ca681 \ + --hash=sha256:1f1489f769582498610e015a8ef2d36f28f505ab3096d0e16b4858a9ec214f57 \ + --hash=sha256:2517e24d7315eb51c10664cdb865195df38ab74456c677df67bb47f12d088a27 \ + --hash=sha256:295d9ffe712caa9f8a3081de8d32fc60191b4b51c76f02f951fd8407253528f4 \ + --hash=sha256:2a075bd7bd19c70cf67c8badfa36cf7c5d8de3c9ddb8420c51e10d9c50e94920 \ + --hash=sha256:32cc0a5365239a6ea0c6ed461e8838d053b57e397443c0ca894dcc8e388d4374 \ + --hash=sha256:332b4f0145c30b5f5ad9374881133e5aa64320428a57c2c2b61e9d891a51c2f3 \ + --hash=sha256:377815a8616074cabbf3f53354e1d040c35815a134e01d7614b7692e4bf8acfa \ + --hash=sha256:38f4a703656f493b0ad185211ccfca7f0386120f022066b018eb5296d8613e23 \ + --hash=sha256:3ac2360e93cb41be81121755c6462cff3beaa9967188c866e5fce5cf13170859 \ + --hash=sha256:3c4923e404d6bcd91b6779c009542e5647fef32e4a5d75e115e3bbac6f2335eb \ + --hash=sha256:3cdcb35dc9d807259c981a85531048ede628eabcffb3239adf3d17463518992d \ + --hash=sha256:41024ed50e44ab1a60d3fe0a9d15a4ccc9f5f2b1d814ff283c8d01134d5b81bc \ + --hash=sha256:413b820229730d358efd838ecbab79902fe97094565fdc80ddb6b0a18c18a581 \ + --hash=sha256:4432b835675f0ea7414aab3d37d119f7226d24869b7a829caeab49ebda407b0c \ + --hash=sha256:4db576bb8c3ef9365f8b40fe0f671644de6736ae2c27a2c62d7d8a1b4329f099 \ + --hash=sha256:4e7f886f47ab881692f278ae901039a234e4025a68e6dfab514263a0b1c4ae05 \ + --hash=sha256:4e9750bc21b886308024f8a54ccb9a2cc38ac9fa813bf4348434e3d54f337ff9 \ + --hash=sha256:5060731cc3ed12ca3a8b57acd4aeca5bbc2f49216dd0bec1650a1acd89486bcd \ + --hash=sha256:50847dca5d197fcbd389c805aa1a1cf32f25d2e7273dc47ab181a517666b68cc \ + --hash=sha256:5092eb5b1172947f57d6ea7d89b2f29650414e4293c47707eb499ec07a0ac796 \ + --hash=sha256:5124d1ea754509b09e53738ec185584cc609aae4a3b510aaf4ed6aa047ef9303 \ + --hash=sha256:51e8c4084897de9f05898c2c2a39af6318044ae969d46ff7a34ed3f96274adca \ + --hash=sha256:530a3fd64c87cffa844d4b6b9768774763d9caa299e9b75d8eca6a4423b31314 \ + --hash=sha256:56fa888f10d0f367155e76ce849fa1166fc9730d13bd2d65a2aa13b6f5424489 \ + --hash=sha256:58f812017cd2985c21fbffb4864d59174d4903dd66fa23815e74bbc7a0e2dd57 \ + --hash=sha256:59cd8683f575d96df5bb48f6add94afc055012c29e28124fcae2b63661b9efb1 \ + --hash=sha256:5ae8e62c147495b01a0f4765c878e9bfdf843412446a247e28df59936e99e797 \ + --hash=sha256:5b233ea3e165e43e35dba1d2b8ecc21cf070b45b65ae17dd2747d2713d942021 \ + --hash=sha256:6176c1811d9d5a04fa391c490cc44f451e240697a16977f11c6f722efb9041db \ + --hash=sha256:62f59da443c4f4849f73a51a193b1d9d258dcad0c41bc4d1b8fb2bcc04bfeb22 \ + --hash=sha256:6783e069732715ad0c3ce96dbf21dbc2235ab0593f2baf6338101f70371f4028 \ + --hash=sha256:6ab8ba9152203feec73758dad83af9a0bbe05001eb4639e547207c40cfb52083 \ + --hash=sha256:70d593af6a6ca332d1df73d519fddb5148edb15cd90d5f0155e3746a6d4fcc65 \ + --hash=sha256:72ec46b7eba5b395e0a7b63025490d3214c11013f4aacb4f5e8d6c3041829588 \ + --hash=sha256:7a32f72973f0f950c1920475d5c5ea3d971b81b6f0ec53b8d0a956cc965f22e0 \ + --hash=sha256:7a4aa69609f40fce3cbc3f87b2061f042eee32f94b8f11db707b66a26461591a \ + --hash=sha256:7c60d3c9b06fb23bd9c6139281ccbdc384297579ae037f08ae90c69f6845c0b1 \ + --hash=sha256:800ee55980c18545af444d93fdd60c56b580db5cc54867d8cbf8a1dc0829938c \ + --hash=sha256:80aa065ffd378ff784822a6d7c3212f2d5f5e9c3589614b5c228b311fd3063ac \ + --hash=sha256:86e0287879f75621ae85197b0877ed2f8b7aa57b511c7331dce2eb6f4de7d476 \ + --hash=sha256:893ff3a711d1b515ba9da14ee090519bad4610ed1962fbe298a434e8c5f8db53 \ + --hash=sha256:89fc958c702ee9a745e4700378f5d23fddbc46ff89e8fdbf5395c24d5c1452a3 \ + --hash=sha256:8c63c91f95173f9c2a67c7c526b2cea976828a0e7fced9cdcead2802dc10f8a4 \ + --hash=sha256:8df31fe574b8b3993cc61764f40941111b25c2d9fea13d3ce24a49907cd2d615 \ + --hash=sha256:8f9baf6f0a6e7571c45c8863010b45e837c3ee1c2c77fcd6ef423be91b21fedb \ + --hash=sha256:9027d773c4ff81487181a925945743413f6069634d0b122d0b37684ccf4f1e18 \ + --hash=sha256:9190426b7aa26c5229501fa297b8d0653cfd3f5a36f7990c264e157cbf886b3b \ + --hash=sha256:940dda65d5e764406b9fb92761cbf462e4e63f712ab60ed98f70552e496f3bf1 \ + --hash=sha256:94eff26096eb5395136634622515b234ecb6c9979824c1f5004c6e3c3c85ccd2 \ + --hash=sha256:9eed0f7edbb274413b6ee781cca50541c8c0facd3d6fd289779e494340a2b85c \ + --hash=sha256:ad4ae4ffd1ee9cd11357b4c66b612da9888f4f4daf2f36995eda64bd45370cac \ + --hash=sha256:b0f172dc8ffaccb8522d7c5d899de00133f2f1ca7b0a49b7da98e901de87bf2d \ + --hash=sha256:b2af221f268f5af85e776a73d62b0845fc8baf8ef0abfae79d29c77d0e776aaf \ + --hash=sha256:b7d335370ae48a780c6e6a6bbfa97342f563744c39c35562f3f367665f5c1de2 \ + --hash=sha256:b83af57bdddef03c01a9138034c6ff03181a3028d9a1003b301eb1a55e161a3f \ + --hash=sha256:bb5136fb5352d3f422df33f0c879a1b0c204004324150cc3b5e3c4f310c9049f \ + --hash=sha256:bc4d8e252f532ab46a1de9349e2d27b91fce46736a9eedaa37beaca66f574ed4 \ + --hash=sha256:bdd3e53429ff02aa319ba59dfe4ceeec345bf46cf180ec2cf6fd5b942e7975e9 \ + --hash=sha256:be12f931839a3bdfe28b584db0e640a65a8bcbc24560ae3fdb025a449b3d754e \ + --hash=sha256:be4a51a55833dc29ab5d7503e7bcb3b3af3402d266018137127450005cdfe737 \ + --hash=sha256:beb7f344487cdcb9e1efe4b7a29681b74d34c08f0043a327a74da852a6749e7b \ + --hash=sha256:bf4679a3d71012a7c2bf360e5cd878fbd5e4fcac0896b56393dec239d81529ed \ + --hash=sha256:c0e1403fd7c26d77c1f03e096dc58a5c726503fa0db0456678b8668f76f521e3 \ + --hash=sha256:c31c13da98624f957b0fb1b5bae5383b2333c2c3f6793d9825dd5ce79b525cb7 \ + --hash=sha256:c438f6ca858697c9ab67eb28246c92508af972e114cac34e57a6d4ba17a3ac08 \ + --hash=sha256:c8277104ded0a51e699c8c3aff63ce2c56d4ed5519a5f73e0fd7057f959a2b9e \ + --hash=sha256:c95cab08d1965db3d84a121f1c7ce7479bdd4072c9b3dafd8fecce48a2e6b902 \ + --hash=sha256:cc0b66c1eec9021353a4b4483afb12dfd50e3669ffbb9152d6842eb34c7e29fd \ + --hash=sha256:cdee07c4d7f6d72008d3f73b9bf027f4e11550224c7c50d8df1ae4a37c1402a6 \ + --hash=sha256:ce9bf03dad3b46408c08649c6fbd6ca28a9fce0eb32fdfffa6775a13103b5310 \ + --hash=sha256:cff8e5383db4989311f99e814feeb90c4723eb4edca425b9d5d9c3fefcdd9537 \ + --hash=sha256:d168fda2dbff7b9b5f38e693182d792a938c31db4dac3a80a4888de603c99554 \ + --hash=sha256:d1ffeb80b5676463d7a7d56acbe8e37a20ce725570e09549fe738e02ca6b7e1e \ + --hash=sha256:d36ca54cb4c6c4686f7cbb7b817f66f5911c12ddb519450bbe86707155028f87 \ + --hash=sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a \ + --hash=sha256:d5cd5189fc2b6a538b75ae45433140c4823463918f7b1617c31e68b085c0022c \ + --hash=sha256:d618fd27420381a4f6044faa71f46d8bfd911bd077c555f7138ed88729bfbe79 \ + --hash=sha256:d76e2d8c75051d58177e762164d2e9ab92886534e3a12e795f103524f221dd8e \ + --hash=sha256:daae526907e262de627d8f70058a0f64acc9e2641c164c99c8f594b34a799a16 \ + --hash=sha256:db485b3847d182b908b483b2ed133c66d88d49cacf98fd278fadafe11b4478d1 \ + --hash=sha256:dd952e03bfbb096cfe2dd35cd9e00f269969b67536cb4370994afc20ff2d0875 \ + --hash=sha256:dda366d548e89a90d88a86c692377d18d8bd64b39c1fb2b92cb31370e2896bbd \ + --hash=sha256:e315e5ec90d88e140f57696ff85b484ff68bb311e36f2c414aa4286293e6dee0 \ + --hash=sha256:e4415a8db000bf49a6dd1c478bf70062eaacff0f462b92b0ba68791a905861f9 \ + --hash=sha256:e7a116ae737f0000343218c4edf5bd45893bfeaff0993c0b215d7124c9f77646 \ + --hash=sha256:e7c4c09a490dc4d4a7f8cbee56c606a320f9dc28cf92a7157a39d1ce7676a657 \ + --hash=sha256:ebae99ed6764f2b5771c522477b311be313e8841d2e0376db2b10922daebbba4 \ + --hash=sha256:ec4c85dc4b687c7f7f15f553ff26a98bfe8c58f5f7f0ac8905f0ba4c7be60232 \ + --hash=sha256:ed3a984b31da7481b103f68776f7128a89ef26ed40f4dc41a2223cda7fb24819 \ + --hash=sha256:f18c2d9782259a6dc132fdc7a63c168cbc74b35284b6d75c673958982a378384 \ + --hash=sha256:f1f9f4121ec58628c96baa3de1a55a4e3a333c5102c8e94b64e23bf7b2083309 \ + --hash=sha256:f42c23db5d1521218a3276bb08666dcb662896a0be7347cba864eca45ff64ede \ + --hash=sha256:f443b4825c50a51ee68585522ab4a1d1257fac65896f282b4c6763337ac9f5d2 \ + --hash=sha256:f6764a4ccab3078db14a632420930f6186058750df066b8ea2a7106df91d3203 \ + --hash=sha256:f7c7553b13f69c1b29a5bde08ddc6d9d0c8bfb84f9ed01c30db25944aeb852a7 \ + --hash=sha256:fa6248cd194edff41d7ea9425ced8ca3a6f838bfb295f6f1d6e6bb694a8518df \ + --hash=sha256:fa8eb9ecdb7efb0b226acec134e0d709e87a909fa4971a54c0c4f6e88635484c \ + --hash=sha256:fc20894c3d21194d8041a28b65622d5b86db786da6e3cfe73f0c762951a61167 \ + --hash=sha256:fc4d3f1fb9ca0ae9f97b095963bc6326f1dbfd3779d6679a1e016b9baaa153d3 \ + --hash=sha256:fd40bb9cd0891c4c3cb1ddf83f8bbfa15731a248fdc8162669405451e2724b09 \ + --hash=sha256:ff710414307fefa903e0d9bdf300972f892c23477829f49504e59834f4195398 +markupsafe==3.0.3 \ + --hash=sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d +matplotlib==3.11.1 \ + --hash=sha256:0c1f44890d435c1b4ef52f701ad5828cb450ea97bcc83918fda6be74965d6cd2 \ + --hash=sha256:11664c551345553db92e61cae6cf1376f138f8c47cafdf13b64b18f3e3e9e464 \ + --hash=sha256:1524e2bdd48a93557aa47ddcfe9c225dfdd57d5a01a5c49128c20f0632980ee1 \ + --hash=sha256:191163532cdefcb1571ca38a6d7e6474baccde64495783e6ba47aa07ec4b9bbb \ + --hash=sha256:1ac697e591c11b6ad04679a73c2d2f9980fe9d9f0311fb414a2e329706343dfb \ + --hash=sha256:216fbb93a74add02ddb4cb38ef5348f59ac00b3e84567eaf16598772d40e150a \ + --hash=sha256:21a67b961a6d597bca54fae826cd20695ba4a6e4d05424a08da6e13e3176fd6b \ + --hash=sha256:2abdee5ffa2fe11b2d19f7a5c63b785fb7c28cc46c7bc1814156341d9d1a33e1 \ + --hash=sha256:30c492d4ba9448595b6fd8708c6725963f8148e25c0d8842948da5b05f0ee8d3 \ + --hash=sha256:3d3fd84082b1afbd9398466c81309e20045be20d48fe0fb18c43504d164cbbb2 \ + --hash=sha256:427258425f9a3fc4ed79a91f9e9b9aaf5a82cb6571e85dc14063cc6fbb993741 \ + --hash=sha256:480194afceca4df2f137c2721227d3cba67121fbf4397b69cee7f83714b0a58a \ + --hash=sha256:54d47b8ae8b579633a3902ca5b4ad6c1e132a5626d64447b2e22a66394e79987 \ + --hash=sha256:5af0dcda57d471440a7b5b623e70e0a61003518443d9098f211a96ecfbbc25be \ + --hash=sha256:5e1f8922ba31959cf6a9dfb51be64b7f7bc582801a3957dc0c2f3afcd3537adf \ + --hash=sha256:5e510088c27a89d53580a752f959146893563e63c330e161d159b0fee652af6f \ + --hash=sha256:6771b0cd7838c6a857a7209814158c0ad09bfef878db3033dd82d70ad101f191 \ + --hash=sha256:67e4c3cd578c65ebd81bdc09a1b6592ceafee6dfafe116dc85dfcb647b5bbb18 \ + --hash=sha256:68408341f2312836fbbdf6b3c78047f65b2d8752f5fd221c3e72d348f5b34f8b \ + --hash=sha256:69647db5746941c793d6e445a4cd349323ffb87d9cc958c2ad84a659b4832d30 \ + --hash=sha256:6be943cb68bc6660ead58c55b3aa6366cba2ef7feb06460fbcce32360376f19f \ + --hash=sha256:7389b77ed2ab0552f46d9a90b81b7b8e6dfcdc42adc36c37a0865799843e0e3e \ + --hash=sha256:7f33a781e12b1e53b278deb2f5373c2e55ec4f10727be3440c0cfb5cda9f944f \ + --hash=sha256:83235693abde86e5e0129998f80ee39fc7f58e6d56a88fafb28a9278833e9d5f \ + --hash=sha256:88a2a27dd9691ae448dfae4b26f59036be90c3c28757edd3553a29559d00859f \ + --hash=sha256:89b193b255f4f6f7948dbcee3691f4f341ab05d9a8874a67b45ddb4182922eda \ + --hash=sha256:8b14eb22961fe865efb0e4ff167e333e428908b00115a8d800ccb65ee108e481 \ + --hash=sha256:9601a1e90be21e4884c53b4f3dc3ee0544654946f9975258d691f1c2e2f119c6 \ + --hash=sha256:96f4bdeea33a8d15a071dbfe6d119451b1d719c733ac666d65357082901a9099 \ + --hash=sha256:9a076f4fc5cdc43fdf510f5981418d25c2db4973418d9f22d8bb3dc8045ada78 \ + --hash=sha256:9fdf1c818ab05d0e74002091ddaf414478a3a449ec9d51c8976d45be7e3a01e2 \ + --hash=sha256:ac104be2768ffdd8655db9e71b768cbb45f2b9aa7b450cf1595e8f65d3822319 \ + --hash=sha256:ae30c6109848ac0f9fa36c5d6270938487614c47ba31860bd5361266dabc5685 \ + --hash=sha256:aee55e9041211bf84302ab55ec3965df18dd90ae19f8b58332a7feaf208bfe83 \ + --hash=sha256:b0a19dcf73406d3746d25a5ed42d713604c9a3e024d129b102852b0d941cb9f3 \ + --hash=sha256:b4c78ceb2f11bcac7389d305cda17aeb1f4586a857854ab5780bd3dd8dbfc407 \ + --hash=sha256:b7cf158e7add54a8d51ac9b5a84abd6d4e13ed4951b4f25f1c5139f41c2addb2 \ + --hash=sha256:b937b9dba5f5f6c1e31c47abe2186c865c0914fd18f2ce0dfc39c9adcef5951d \ + --hash=sha256:ba8f811b8ddfac493734d6af0b2dff96919d0c28ca0d641858dab4262777c6ea \ + --hash=sha256:c52f7ad20ef476806ed212380b1d54d20310c8b86bdc2c9a68b51f0024a44472 \ + --hash=sha256:c90be0b73568da4f662afac580956a76e308437e641b4a45aa08925eeb67d95f \ + --hash=sha256:d2ace7273b9a5061a3b420918a16fae1f2dc5dfee1abcc13aba71b5d94b1820c \ + --hash=sha256:dadfe80797174e2984aae3be0b77594a3c72d2c0a40fbd4a0de48d2728caf3ae \ + --hash=sha256:e15ef41507f3d525f46154ac9e3ae785dacde9f20e593a25de8986267892ef74 \ + --hash=sha256:e4b9ac2f1f607ecda2af90a5232beee2af7582fce1cc30c4b6a1b012dc21ee99 \ + --hash=sha256:f2912f647f3fbe1ccf085f91e213936f9101bead81a5e670565b1f1b3712f4fb +matplotlib-inline==0.2.2 \ + --hash=sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6 \ + --hash=sha256:72f3fe8fce36b70d4a5b612f899090cd0401deddc4ea90e1572b9f4bfb058c79 +mpmath==1.3.0 \ + --hash=sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f \ + --hash=sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c +msgpack==1.2.1 \ + --hash=sha256:01e2dd6c9b19d333a00282330cc8a73d38d8dabc306dc5b42cd668c3ac82e833 \ + --hash=sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a \ + --hash=sha256:04c721c2c7448767e9e3f2520a475663d8ee0f09c31890f6d2bd70fd636a9647 \ + --hash=sha256:05f340e47e7e47d2da8db9b53e1bb1d294369e9ef45a747441309f6650b8351d \ + --hash=sha256:0a70e3cf2804a300d921bb0940426e35f4e489a23adfb77a808892241db0a064 \ + --hash=sha256:0adcf06ffde0777c0e1a9b771a2b1c4226ba1bbf748c8efcc02fcdeca3299107 \ + --hash=sha256:0c0d9802354507bcba62af19c17918e3eb437cc25e6f50657d511b5856a77aac \ + --hash=sha256:0e2bf9280bceb5efca998435904b5d3e9fdbcc11d90dc9df30aec7973252b720 \ + --hash=sha256:1233ee2dd0cefba127583de50ea654677277047d238303521db35def3d7b2e7c \ + --hash=sha256:146ee4e9ce80b365c6d4c47073da9da7bcec473e58194ceee5dd7620ace77e06 \ + --hash=sha256:1548006a91aa93c5da81f3bdcebc1a0d10cea2d25969754fbe848da622b2b895 \ + --hash=sha256:196300e7e5d6e74d50f1607ab9c06c4a1484c383cd22defd727902591f7e8dde \ + --hash=sha256:1dabedcd0f23559f3596428c6589c1cd8c6eaed3a0d720795b07b0225d769203 \ + --hash=sha256:20466cca18c49c7292a8984bc15d65857b171e7264bdcb5f96baf8be238791fc \ + --hash=sha256:298872ecf9e61950f1c6af4ca969b859ee91783bb920ef6e6172697d0c8aad74 \ + --hash=sha256:29a3f6e9667868429d8240dfd063ea5ffdc1321c13d783aa23827a38de0dcb22 \ + --hash=sha256:2eda0b7ebb1283a98d3e4492ac933c8af6aff59fd3df1c3ed024f536af4b1dc8 \ + --hash=sha256:2ef59c659f289eddf8aa6623823f19fa2f40a4029266889eac7a2505dd210c35 \ + --hash=sha256:2ff164c1b0bcb740b073b99e945234d0212852fa378e44a208c425379140dbeb \ + --hash=sha256:33f14fba63278b714efe6ad07e50ea5f03d91537aa6a1c5f1ceca4cf44013ca9 \ + --hash=sha256:350cb813d0af6e65d2f7ef0d729f7ff5be5a8bce03665892f43e5883d4ecc1b8 \ + --hash=sha256:4202c74688ca06591f78cb18988228bd4cca2cc75d57b60008372892d2f1e6e6 \ + --hash=sha256:4227224aaec8f7fbcbfbd4272319347b2bb4030366502600f8c45588c5187b07 \ + --hash=sha256:491cc39455ca765fad51fb451bf2915eb2cf41192ab5801ce8d67c1d614fe056 \ + --hash=sha256:575957e79cd51903a4e8495a242442949641e08f1efd5197b43bebd3ea7682b4 \ + --hash=sha256:5ad5467fc3f68b5468e06c5f788d712e9f8ffc8b0cd1bcb160c105c1ee92dae7 \ + --hash=sha256:5bb9c386f0a329c035ddbab4b72d1028bf9627add8dda41070288563d57ed1b1 \ + --hash=sha256:5c24aa15d5963051e1a5c62b12c50cd705992502b5ec1f3bece6046f33c9fc24 \ + --hash=sha256:5f6277e5f783c36786a145e0247fc189a03f35f84b251646e53592d2bc12b355 \ + --hash=sha256:60926b75d00c8e816ef98f3034f484a8bc64242d66839cef4cf7e503142316a0 \ + --hash=sha256:633727297ed063441fd1cda2288865487f33ad14eeb8831afb5f0c396a62cfce \ + --hash=sha256:67f6dd22fa72a93752643f07889796d62739a13415ee630169a8ce764f86cf9f \ + --hash=sha256:6d09badf350af2be9d189184e04e64cf54ad93569ab3d96fca58bd3e84aad707 \ + --hash=sha256:6ee967f7c7e1df2890c671ff2ee51a28ded0efc95da3e507176dee881ce36c66 \ + --hash=sha256:74847557e28ce71bd3c438a447ca90e4b507e997ddbdef8a12a7b283b86c156b \ + --hash=sha256:779197a6513bab3c3632265e3d0f7cb3227e62510841a6f34f1eaa37efbb345e \ + --hash=sha256:787c9bebb5833e8f6fc8abca3c0597683d8d87f56a8842b6b89c75a5f3176e2d \ + --hash=sha256:7d31c0ac0c640f877804c67cb2bc9f4e23dc2db97e96c2e67fa27d38283b41f8 \ + --hash=sha256:810b916696c86ef0deb3b74588480224df4c1b071136c34183e4a2a4284d7ac7 \ + --hash=sha256:83efa1c898e0fc5380fc0cabbf75164c52e3b5cbb45973710d75821928380c73 \ + --hash=sha256:85f57e960d877f2977f6430896191b04a21f8901b3b4baf2e4604329f4db5402 \ + --hash=sha256:8b267ce94efb76fbd1b3373511420074ee3187f0f7811bf394531de13294735a \ + --hash=sha256:8c2ed1e48cc0f460bf3c7780e7137ff21a4e18433451916f2442c1b21036cd7d \ + --hash=sha256:8c7b398c56ff125feae96c2737abfec5595f1fa0aa186df60c56040b8accb95c \ + --hash=sha256:8d00f177ca88a77c1cf848d204a38f249751650b601cb6532acc68805d8a8273 \ + --hash=sha256:8ff92d7feeaf5bc26c51495b69e2f99ed97ab79346fb6555f44be7dd2ac6503b \ + --hash=sha256:91054a783328e0ea7954b8771095705c8d2243b814743fbaadf14552c9c52c5d \ + --hash=sha256:98b58bdb89c46190e4609bb36abe17c6d4105ad13f9c5f8f6f64d320f8ced3fb \ + --hash=sha256:a28d076ca7c82b9c8728ad90b7147489449557038bed50e4241eb832395169b4 \ + --hash=sha256:aa6c4be5d1c02a42b066ca6ddb71adf36432868fdcdb6ee87e634e86e0674190 \ + --hash=sha256:aded5bdf32609dc7987a49bbbd15a8ef096193f96dd8bbeb791de729e650acf5 \ + --hash=sha256:afc5febcd4c99effbc02b528e49d6fd0760b2b7d48c05239e345a5fa6e743d9a \ + --hash=sha256:b50b727bd652bdc37d950336c848ef20ec54a4cafc38dce19b1cd86ad625d0f7 \ + --hash=sha256:c1c79a604a2969a868a78b6ebd27a887e00c624f14f66b3038e0590cb23332d1 \ + --hash=sha256:ca0dacff965c47afdc3749a8469d7302a8f801d6a28758d55120d75e66ce6889 \ + --hash=sha256:d3567748a5107cb40cdf66a275430c2f87c07777698f4bfd25c35f44d533258c \ + --hash=sha256:dc871b997a9370d855b7394465f2f350e847a5b806dd38dcc9c989e7d87da155 \ + --hash=sha256:dd3bfe82d53edfe4b7fc9a7ec9761e23a7a5b1dac22264505af428253c29ed24 \ + --hash=sha256:e3dc2feb0876209d9c38aa56cb1de169bd6c4348f1aa48271f241226590993e6 \ + --hash=sha256:e4f1d0f8f98ade9634e01fb704a408f9336c0a8f1117b369f5db83dc7551d8b1 \ + --hash=sha256:ec0e675d59150a6269ddc9139087c722292664a37d071a849c05c473350f1f2d \ + --hash=sha256:ee1d9ed27d0497b848923746cf762ed2e7db24f4be7eec8e5cbe8c766aa707b7 \ + --hash=sha256:f02cf17a6ca1abe29b5f980644f7551f94d71f2011509b26d8625ce038f0df64 \ + --hash=sha256:f12038a35fabd52e56a3547bab42401af49a45caa6dd00b34c44de235bc93ee2 \ + --hash=sha256:f310233ef7fb9c14e201c93639fe5f5260b005f56f0b29048e999c30935596cc \ + --hash=sha256:f9389552ecf4784886345ead0647e4edc96bee37cbab05b75540f542f766c48c +msgpack-numpy==0.4.8 \ + --hash=sha256:773c19d4dfbae1b3c7b791083e2caf66983bb19b40901646f61d8731554ae3da \ + --hash=sha256:c667d3180513422f9c7545be5eec5d296dcbb357e06f72ed39cc683797556e69 +narwhals==2.24.0 \ + --hash=sha256:42fdedf44e5b2ca7505630d45b4ac3058f38d8485cba9fe1652ca23152df7489 \ + --hash=sha256:b5c0f684ccd9d7475b564111e319a4964abcf2baf79d3cf6b1003d06ac9b828d +networkx==3.6.1 \ + --hash=sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509 \ + --hash=sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762 +numpy==1.26.4 \ + --hash=sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b \ + --hash=sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818 \ + --hash=sha256:1af303d6b2210eb850fcf03064d364652b7120803a0b872f5211f5234b399f20 \ + --hash=sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0 \ + --hash=sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010 \ + --hash=sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a \ + --hash=sha256:3373d5d70a5fe74a2c1bb6d2cfd9609ecf686d47a2d7b1d37a8f3b6bf6003aea \ + --hash=sha256:47711010ad8555514b434df65f7d7b076bb8261df1ca9bb78f53d3b2db02e95c \ + --hash=sha256:4c66707fabe114439db9068ee468c26bbdf909cac0fb58686a42a24de1760c71 \ + --hash=sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110 \ + --hash=sha256:52b8b60467cd7dd1e9ed082188b4e6bb35aa5cdd01777621a1658910745b90be \ + --hash=sha256:60dedbb91afcbfdc9bc0b1f3f402804070deed7392c23eb7a7f07fa857868e8a \ + --hash=sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a \ + --hash=sha256:666dbfb6ec68962c033a450943ded891bed2d54e6755e35e5835d63f4f6931d5 \ + --hash=sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed \ + --hash=sha256:679b0076f67ecc0138fd2ede3a8fd196dddc2ad3254069bcb9faf9a79b1cebcd \ + --hash=sha256:7349ab0fa0c429c82442a27a9673fc802ffdb7c7775fad780226cb234965e53c \ + --hash=sha256:7ab55401287bfec946ced39700c053796e7cc0e3acbef09993a9ad2adba6ca6e \ + --hash=sha256:7e50d0a0cc3189f9cb0aeb3a6a6af18c16f59f004b866cd2be1c14b36134a4a0 \ + --hash=sha256:95a7476c59002f2f6c590b9b7b998306fba6a5aa646b1e22ddfeaf8f78c3a29c \ + --hash=sha256:96ff0b2ad353d8f990b63294c8986f1ec3cb19d749234014f4e7eb0112ceba5a \ + --hash=sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b \ + --hash=sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0 \ + --hash=sha256:a354325ee03388678242a4d7ebcd08b5c727033fcff3b2f536aea978e15ee9e6 \ + --hash=sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2 \ + --hash=sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a \ + --hash=sha256:afedb719a9dcfc7eaf2287b839d8198e06dcd4cb5d276a3df279231138e83d30 \ + --hash=sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218 \ + --hash=sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5 \ + --hash=sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07 \ + --hash=sha256:cd25bcecc4974d09257ffcd1f098ee778f7834c3ad767fe5db785be9a4aa9cb2 \ + --hash=sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4 \ + --hash=sha256:d5241e0a80d808d70546c697135da2c613f30e28251ff8307eb72ba696945764 \ + --hash=sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef \ + --hash=sha256:f870204a840a60da0b12273ef34f7051e98c3b5961b61b0c2c1be6dfd64fbcd3 \ + --hash=sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f +nvidia-cublas==13.1.1.3 \ + --hash=sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436 \ + --hash=sha256:b6cdce694e47ff6aadf0a69df1cab6628d696f5ff56e8d16af50309d855fa20f \ + --hash=sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5 +nvidia-cuda-cupti==13.0.85 \ + --hash=sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8 \ + --hash=sha256:683f58d301548deeefcb8f6fac1b8d907691b9d8b18eccab417f51e362102f00 \ + --hash=sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151 +nvidia-cuda-nvrtc==13.0.88 \ + --hash=sha256:6bcd4e7f8e205cbe644f5a98f2f799bef9556fefc89dd786e79a16312ce49872 \ + --hash=sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575 \ + --hash=sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b +nvidia-cuda-runtime==13.0.96 \ + --hash=sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548 \ + --hash=sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55 \ + --hash=sha256:f79298c8a098cec150a597c8eba58ecdab96e3bdc4b9bc4f9983635031740492 +nvidia-cudnn-cu13==9.20.0.48 \ + --hash=sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304 \ + --hash=sha256:af8139732b99c0118be65ea5aac97f0d46018f8c552889e49d2fb0c6261a4a24 \ + --hash=sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1 +nvidia-cufft==12.0.0.61 \ + --hash=sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5 \ + --hash=sha256:2abce5b39d2f5ae12730fb7e5db6696533e36c26e2d3e8fd1750bdd2853364eb \ + --hash=sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3 +nvidia-cufile==1.15.1.6 \ + --hash=sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44 \ + --hash=sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1 +nvidia-curand==10.4.0.35 \ + --hash=sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a \ + --hash=sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc \ + --hash=sha256:65b1710aa6961d326b411e314b374290904c5ddf41dc3f766ebc3f1d7d4ca69f +nvidia-cusolver==12.0.4.66 \ + --hash=sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2 \ + --hash=sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112 \ + --hash=sha256:16515bd33a8e76bb54d024cfa068fa68d30e80fc34b9e1090813ea9362e0cb65 +nvidia-cusparse==12.6.3.3 \ + --hash=sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b \ + --hash=sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c \ + --hash=sha256:cbcf42feb737bd7ec15b4c0a63e62351886bd3f975027b8815d7f720a2b5ea79 +nvidia-cusparselt-cu13==0.8.1 \ + --hash=sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f \ + --hash=sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0 \ + --hash=sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215 +nvidia-nccl-cu13==2.29.7 \ + --hash=sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5 \ + --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d +nvidia-nvjitlink==13.3.33 \ + --hash=sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5 \ + --hash=sha256:4297ee49639b4f2e07255a1d69b3acc7ab2d011bb892b403e91ac98368962e3b \ + --hash=sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e +nvidia-nvshmem-cu13==3.4.5 \ + --hash=sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80 \ + --hash=sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9 +nvidia-nvtx==13.0.85 \ + --hash=sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4 \ + --hash=sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6 \ + --hash=sha256:d66ea44254dd3c6eacc300047af6e1288d2269dd072b417e0adffbf479e18519 +packaging==26.2 \ + --hash=sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e \ + --hash=sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661 +pandas==3.0.3 \ + --hash=sha256:0383c72c75cdcca61a9e116e611143902dbfd08bff356829c2f6d1cf40a9ca8c \ + --hash=sha256:05f1f1752b8533ea03f7f39a9c15b1a058d067bb48f4748948e7a8691e0510f2 \ + --hash=sha256:08d789b41f87e0905880e293cedf6197ce71fe67cc081358b1e148a491b9bd13 \ + --hash=sha256:0d589105b3c14645af1738ff279b2995102d8f7a03b0a66dc8d95550eb513e04 \ + --hash=sha256:13fc1e853d9e04743d11ba75a985ccbc2a317fe07d8af61e445a6fd24dacd6a6 \ + --hash=sha256:14da8316da4d0c5a77618425996bfb1248ca87fc2c1486e6fde4652bd18b5824 \ + --hash=sha256:1928e07221f82db493cd4af1e23c1bfca524a19a4699887975bff68f49a72bfb \ + --hash=sha256:261e308dfb22448384b7580cf719d2f998fe2966c92893c3e77d14008af1f066 \ + --hash=sha256:275c14e0fce14a2ec20eee474aecd305478ea3c1e6f6a9d8fe219a165542717e \ + --hash=sha256:335f62418ed562cfc3c49e9e196375c28b729dcef8543abf4f9438e381bf3c76 \ + --hash=sha256:3650109c0f22879df8bd6179ab9ee3d7f1d1d4e7e0094a3f0032d9f51e2e64ac \ + --hash=sha256:39436b377d56d2a2e52d0395bdbee171f01068e99af5250509aceeb929f765c7 \ + --hash=sha256:3c20a521bbb85902f79f7270c80a59e1b5452d96d170c034f207181870f97ac5 \ + --hash=sha256:3e91cec1879ada0624fc3dc9953c5cbd60208e59c0db28f540c5d6d47502422f \ + --hash=sha256:455f6f8139d4282188f526868dbc3c828470e88a3d9d59a891bd46a455f21b98 \ + --hash=sha256:46997386d528eb40376ecd6b033cf4a8a1e5282580f68f43de875b78cba2199d \ + --hash=sha256:4db8c527972a821cf5286b40ccc57642a39bc62e62022b42f99f8a67fca8c3a1 \ + --hash=sha256:4e15135e2ee5df1063313e2425ceef8ac0f4ae775893815b0923651b806a5639 \ + --hash=sha256:51b1fe551acb77dac643c6fda86084d8d446c10fe64b06a9cc29c4cc8540e7f2 \ + --hash=sha256:557409bc4178e70ee8d9ddb494798e51ebf6ea59330f6be22c51bab2a7db6c49 \ + --hash=sha256:5cc09a68b3120e0f54870dede8287a7bb1fa463907e4fcec1ea77cab6179bf7a \ + --hash=sha256:60ae316d3fd75d1858d450d0db0103ea2be3e7d4a95ec2f064f7e2ae63f7b028 \ + --hash=sha256:6674ab18ad8c57802867264b00e15e7bb904700cdd9046e3b2fa1fce237439ea \ + --hash=sha256:67b3b64c11910cfa29f4e94a14d3bff9ee693b6fc76055e7cad549cee0aec5fa \ + --hash=sha256:696a4a00a2a2a35d4e5deb3fc946641b96c944f02230e4f76137fe35d806c4fc \ + --hash=sha256:6dc0b3fd2169c9157deed50b4d519553a3655c8c6a96027136d654592be973a9 \ + --hash=sha256:7e65d5407dc0b394f509699650e4a2ec01c0514f21850f453fa60f3be79a5dbf \ + --hash=sha256:819959dab7bbd0049c15623fbac4e29a191b9528160a61fb1032242d8ced2d9c \ + --hash=sha256:8a1e45c80cceb3b4a21bc5939d52e8cbd8d9b7305309219d59e9754d9ce09e27 \ + --hash=sha256:9c39be2d709d01fa972a0cabc522389fceca4f3969332ba25a7d6c5802cf976a \ + --hash=sha256:9d71c63ae4ebdbf70209742096f1fc46a83a0613c99d4b23766cced9ff8cd62a \ + --hash=sha256:a2d2dff8a04f3917b55ab3910c32990f8ddf7eceba114947838cefa976a68977 \ + --hash=sha256:a4eeb6830daf35a71cc09649bd823e2b542dac246cdee9614c6e4bd65028cd6a \ + --hash=sha256:a55066a0505dae0ba2b50a46637db34b46f9094c65c5d4800794ef6335010938 \ + --hash=sha256:a82d532a3351d435432cd913edbccaf8b8e01d4dd0e5ced5a8d2e8ecd94c7e44 \ + --hash=sha256:b168fc218fd80a6cbdbdbc1a97ddc7889ed057d7eb45f50d866ceab5f39904c4 \ + --hash=sha256:b2c95f8bfc1ee412bf482605d7bfd30c12d1d26bd59fdd91efeef1d4718decb1 \ + --hash=sha256:ba7e08b9ac1d54569cd1e256e3668975ed624d6826f7b68df0342b012007bddb \ + --hash=sha256:bab900348131a7db1f69a7309ef141fd5680f1487094193bcbbb61791573bf8f \ + --hash=sha256:bd3a518890b400d32f9023722dc9a9a5c969f00b415419a3c06c043f09bb5d7d \ + --hash=sha256:c7be265b62cef88e253a941e4698604973736dcfe242fdb5198f0f7bc473cdcc \ + --hash=sha256:d26cbe1fcfc12e8fd900e2454163e466b2d3af84f7c75481df7683ffc073d870 \ + --hash=sha256:d4be06d68f9ddcfc645b87534911da79a8fbffc7573c80e0edcf42a5020624d8 \ + --hash=sha256:d72828c20c6d6e83e1e22a6a3b47b326b71664112fa9705dcbccfd7a39b62085 \ + --hash=sha256:dd1a5d1def6a46002e964510bdc67c368aa0951df5d1d9f8365336f5a1f490cd \ + --hash=sha256:e3a2ec42c98ffa2565a67e08e218d06d72576d758d90facb7c00805194d8f360 \ + --hash=sha256:f8894dc474d648fe7b6ff0ca9b0bd73950d19952bc1a6534540762c5d79d305c \ + --hash=sha256:fed2ff7fd9779120e388e285fc029bd5cf9490cdd2e4166a9ee22c0e49a9ab09 +parso==0.8.7 \ + --hash=sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c \ + --hash=sha256:eaaac4c9fdd5e9e8852dc778d2d7405897ec510f2a298071453e5e3a07914bb1 +pexpect==4.9.0 \ + --hash=sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523 \ + --hash=sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f +pillow==12.3.0 \ + --hash=sha256:00808c5e14ef63ac5161091d242999076604ff74b883423a11e5d7bbb38bf756 \ + --hash=sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a \ + --hash=sha256:06ff022112bc9cbf83b60f8e028d94ad87b60621706487e65f673de61610ab59 \ + --hash=sha256:0740a512dc522224c77d9aa5a8d70d8b7d73fb91f2c21125d8d025d3b8990e45 \ + --hash=sha256:0847a763afefb695bc912d7c131e7e0632d4edc1d8698f58ddabec8e46b8b6d3 \ + --hash=sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df \ + --hash=sha256:0feb2e9d6ad6c9e3c06effe9d00f3f1e618a6643273576b016f591e9315a7139 \ + --hash=sha256:10e41f0fbf1eec8cfd234b8fe17a4caac7c9d0db4c204d3c173a8f9f6ef3232b \ + --hash=sha256:1182d52bc2d5e5d7d0949503aa7e36d12f42205dc287e4883f407b1988820d39 \ + --hash=sha256:164b31cd1a0490ab6efae01aa5df49da7061be0af1b30e035b6e9a1bfe34ee6e \ + --hash=sha256:1657923d2d45afb66526e5b933e5b3052e6bdea196c90d3abb2424e18c77dae8 \ + --hash=sha256:186941b6aef820ad110fb01fb06eb925374dc3a21b17e37ec9a53b250c6fe2d1 \ + --hash=sha256:1cca606cd25738df4ed873d5ad46bbdb3d83b5cbca291f6b4ff13a4df6b0bbe8 \ + --hash=sha256:21900ce7ba264168cd50defae43cd75d25c833ad4ad6e73ffc5596d12e25ac89 \ + --hash=sha256:236ff70b9312fb68943c703aa842ca6a758abfa45ac187a5e7c1452e96ef72b5 \ + --hash=sha256:23aceaa007d6172b02c277f0cd359c79492bbb14f7072b4ede9fbcaf20648130 \ + --hash=sha256:23d27a3e0307ec2244cc51e7287b919aa68d097504ebe19df4e76a98a3eea5bd \ + --hash=sha256:24870b09b224f7ae3c39ed07d10e819d06f8720bc551847b1d623832b5b0e28d \ + --hash=sha256:251bf95b67017e27b13d82f5b326234ca62d70f9cf4c2b9032de2358a3b12c7b \ + --hash=sha256:25b9b82bb22e6e2b3cd07b39c68b7b862001226cb3dff7130d1cb914121b39ed \ + --hash=sha256:28ce87c5ab450a9dd970b52e5aca5fe63ed432d18a2eaddd1979a00a1ba24ace \ + --hash=sha256:300557495eb45ebb8aec96c2da9c4be642fbf7cd937278b4013ba894ea8eb0eb \ + --hash=sha256:30f2aa603c41533cc25c05acd0da21636e84a315768feb631c937177db558931 \ + --hash=sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510 \ + --hash=sha256:37d6d0a00072fd2948eb22bce7e1475f34569d90c87c59f7a2ec59541b77f7a6 \ + --hash=sha256:37dc8f7bbb66efe481bb60defacef820c950c24713fb44962ed6aa2a50966de1 \ + --hash=sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce \ + --hash=sha256:3edce1d53195db527e0191f84b71d02022de0540bf43a16ed734ed7537b07385 \ + --hash=sha256:446c34dcc4324b084a53b705127dc15717b22c5e140ae0a3c38349d4efec071e \ + --hash=sha256:4998562bf62a445225f22e07c896bb04b35b1b1f2eb6d760584c9c51d7a5f78c \ + --hash=sha256:4b0a7fe987b14c31ebda6083f74f22b561fd3739bc0ac51e019622e3d72668c7 \ + --hash=sha256:4e8c2a84d977f50b9daed6eeaf3baef67d00d5d74d932288f02cb94518ee3ace \ + --hash=sha256:4f883547d4b7f0495ebe7056b0cc2aea76094e7a4abc8e933540f3271df27d9c \ + --hash=sha256:514435a37670e3e5e08f3945b68718b6ed329bb84367777e16f9f4dfe1e61a0f \ + --hash=sha256:53aa02d20d10c3d814d536aa4e5ac9b84ca0ff5a88377963b085ad6822f93e64 \ + --hash=sha256:5594fc43d548a7ed94949d139aa1341b270f1863f11cfd37f5a6c8b778a6b67f \ + --hash=sha256:571b9fcb07b97ef3a492028fb3d2dc0993ca23a06138b0315286566d29ef718a \ + --hash=sha256:57b3d78c95ba9059768b10e28b813002261d3f3dfc55cc48b0c988f625175827 \ + --hash=sha256:5afb51d599ea772b8365ae807ae557f18bccfe46ab261fd1c2a9ed700fc6eb17 \ + --hash=sha256:6b02afb9b97f65fbca5f31db6a2a3ba21aa93030225f150fa3f249717e938fb4 \ + --hash=sha256:6c0016e7b354317c4e9e525b937ac8596c38d2d232b419529b9cd7a1cd46e39a \ + --hash=sha256:71d6097b330eea8fd15097780c8e89cb1a8ce7838669f48c5bacd6f663dd4701 \ + --hash=sha256:756c768d0c9c2955feb7a56c37ea24aea2e369f8d36a88da270b6a9f19e62b5e \ + --hash=sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91 \ + --hash=sha256:7a743ff716f746fc19a9557f60dab1600d4613255f8a7aeb3cdde4db7eb15a66 \ + --hash=sha256:85f998ea1848bc6757289e739cfbdda3a04adfd58b02fc018ce54d754a5ce468 \ + --hash=sha256:8728f216dcdb6e6d555cf971cb34076139ad74b31fc2c14da4fafc741c5f6217 \ + --hash=sha256:877c3f311ff35410f690861c4409e7ccbf0cd2f878e50628a28e5a0bb689e658 \ + --hash=sha256:8cd2f7bdda092d99c9fc2fb7391354f306d01443d22785d0cbfafa2e2c8bb418 \ + --hash=sha256:8e95e1385e4998ae9694eeaa4730ba5457ff61185b3a55e2e7bea0880aef452a \ + --hash=sha256:962864dc93511324d51ddbb5b9f8731bf71675b93ca612a07441896f4688fb8c \ + --hash=sha256:9cf95fe4d0f84c82d282745d9bb08ad9f926efa00be4697e767b814ce40d4330 \ + --hash=sha256:9e881fca225083806662a5c43d627d215f258ff43c890f831966c7d7ba9c7402 \ + --hash=sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09 \ + --hash=sha256:a45650e8ce7fafffd731db8550230db6b0d306d181a90b67d3e6bca2f1990930 \ + --hash=sha256:a876864214e136f0eb367788dbd7df045f4806801518e2cfe9e13229cfe06d8f \ + --hash=sha256:ae26d61dfa7a47befdc7572b521024e8745f3d809bd95ca9505a7bba9ef849ec \ + --hash=sha256:af8d94b0db561cf68b88a267c5c44b49e134f525d0dc2cb7ed413a66bc23559a \ + --hash=sha256:b343699e8308bdc51978310e1c959c584e7869cc8c40780058c87da7781a1e94 \ + --hash=sha256:b3c777e849237620b022f7f297dd67705f9f5cf1685f09f02e46f93e92725468 \ + --hash=sha256:b629de27fda84b42cde7edef0d85f13b958b47f6e9bbcbba9b673c562a89bd8b \ + --hash=sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965 \ + --hash=sha256:ba54cfebe86920a559a7c4d6b9050791c20513650a1952ebe3368c7dc70306f8 \ + --hash=sha256:bcb46e2f9feff8d06323983bd83ed00c201fdcab3d74973e7072a889b3979fcd \ + --hash=sha256:bcc33feacfaefce60c12fd500a277533bdc02b10a19f7f6d348763d8140bbba7 \ + --hash=sha256:bf16ba1b4d0b6b7c8e534936632270cf70eb00dbe09005bc345b2677b726855c \ + --hash=sha256:cf1845d02ad822a369a49f2bb9345b1614744267682e7a03527dc3bf6eea1777 \ + --hash=sha256:d69141514cc30b774ceea5e3ed3a6635c8d8a96edf664689b890f4089111fb35 \ + --hash=sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9 \ + --hash=sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f \ + --hash=sha256:dc624f6bc473dacdf7ef7eb8678d0d08edf15cd94fad6ae5c7d6cc67a4e4902f \ + --hash=sha256:e158cb00350dc278f3b91551101aa7d12415a66ebf2c91d8d5ac14e56ddd3ad0 \ + --hash=sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c \ + --hash=sha256:e795b7eb908249c4e43c7c99fac7c2c75dab0c43566e37db472a355f63693d71 \ + --hash=sha256:e7e480451b9fa137494bccd3a7d69adbe8ac65a87d97be61e11f1b1050a5bac3 \ + --hash=sha256:e91206ee562682b51b98ef4b26a6ef48fd84e15fd4c4bc5ec768eb641d206838 \ + --hash=sha256:e9871b1ffbfa9656b60aeee92ed5136a5742696006fa322b29ea3d8da0ecc9cf \ + --hash=sha256:e9aeb04d6aef139de265b29683e119b638208f88cf73cdd1658aa07221165321 \ + --hash=sha256:ebaea975e03d3141d9d3a507df75c9b3ec90fa9d2ffd07567b3a978d9d790b26 \ + --hash=sha256:f0606c8bf2cdefea14a43530f7657cbbb7ecf1c4222512492ef4a4434a9501ec \ + --hash=sha256:f13c32a3abd6079a66d9526e18dad9b6d280384d49d7c54040cd57b6424041d9 \ + --hash=sha256:f7401aebd7f581d7f83a439d87d474999317ee099218e5ad25d125290990ba65 \ + --hash=sha256:fa4ecea169a355be7a3ade2c783e2ed12f0e40d2c5621cda8b3297faf7fbb9f5 \ + --hash=sha256:fbd139c8447d25dd750ab79ee274cc5e1fe80fc56340ab10b18a195e1b6eca3e \ + --hash=sha256:fdafc9cce40277e0f7a0feabce0ee50dd2fa1800f3b38015e51296b5e814048d \ + --hash=sha256:fe3cca2e4e8a592be0f269a1ca4835c25199d9f3ce815c8491048f785b0a0198 \ + --hash=sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7 +pluggy==1.6.0 \ + --hash=sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3 \ + --hash=sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746 +prompt-toolkit==3.0.52 \ + --hash=sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855 \ + --hash=sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955 +psutil==7.2.2 \ + --hash=sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372 \ + --hash=sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9 \ + --hash=sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841 \ + --hash=sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63 \ + --hash=sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979 \ + --hash=sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a \ + --hash=sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b \ + --hash=sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9 \ + --hash=sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee \ + --hash=sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312 \ + --hash=sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b \ + --hash=sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9 \ + --hash=sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e \ + --hash=sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc \ + --hash=sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1 \ + --hash=sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf \ + --hash=sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea \ + --hash=sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988 \ + --hash=sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486 \ + --hash=sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00 \ + --hash=sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8 +ptyprocess==0.7.0 \ + --hash=sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35 \ + --hash=sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220 +pure-eval==0.2.3 \ + --hash=sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0 \ + --hash=sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42 +py3dmol==2.5.5 \ + --hash=sha256:9717ea9a899ec641b458f53de538e5bae758281f5053be78bc2db0706ae60bcb \ + --hash=sha256:cb102cc3370800a9ca7679774759f008e34a0f3a283778d85b0410849d370298 +pydssp==0.9.1 \ + --hash=sha256:74fb8129c07c1625bb687b80f7e94ae7ebf1277725258d7fc75fc1f3d12a67dc +pygments==2.20.0 \ + --hash=sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f \ + --hash=sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176 +pygtrie==2.5.0 \ + --hash=sha256:203514ad826eb403dab1d2e2ddd034e0d1534bbe4dbe0213bb0593f66beba4e2 \ + --hash=sha256:8795cda8105493d5ae159a5bef313ff13156c5d4d72feddefacaad59f8c8ce16 +pyparsing==3.3.2 \ + --hash=sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d \ + --hash=sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc +pytest==9.0.2 \ + --hash=sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b \ + --hash=sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11 +python-dateutil==2.9.0.post0 \ + --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \ + --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 +pyyaml==6.0.3 \ + --hash=sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c \ + --hash=sha256:0150219816b6a1fa26fb4699fb7daa9caf09eb1999f3b70fb6e786805e80375a \ + --hash=sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3 \ + --hash=sha256:02ea2dfa234451bbb8772601d7b8e426c2bfa197136796224e50e35a78777956 \ + --hash=sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6 \ + --hash=sha256:10892704fc220243f5305762e276552a0395f7beb4dbf9b14ec8fd43b57f126c \ + --hash=sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65 \ + --hash=sha256:1d37d57ad971609cf3c53ba6a7e365e40660e3be0e5175fa9f2365a379d6095a \ + --hash=sha256:1ebe39cb5fc479422b83de611d14e2c0d3bb2a18bbcb01f229ab3cfbd8fee7a0 \ + --hash=sha256:214ed4befebe12df36bcc8bc2b64b396ca31be9304b8f59e25c11cf94a4c033b \ + --hash=sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1 \ + --hash=sha256:22ba7cfcad58ef3ecddc7ed1db3409af68d023b7f940da23c6c2a1890976eda6 \ + --hash=sha256:27c0abcb4a5dac13684a37f76e701e054692a9b2d3064b70f5e4eb54810553d7 \ + --hash=sha256:28c8d926f98f432f88adc23edf2e6d4921ac26fb084b028c733d01868d19007e \ + --hash=sha256:2e71d11abed7344e42a8849600193d15b6def118602c4c176f748e4583246007 \ + --hash=sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310 \ + --hash=sha256:37503bfbfc9d2c40b344d06b2199cf0e96e97957ab1c1b546fd4f87e53e5d3e4 \ + --hash=sha256:3c5677e12444c15717b902a5798264fa7909e41153cdf9ef7ad571b704a63dd9 \ + --hash=sha256:3ff07ec89bae51176c0549bc4c63aa6202991da2d9a6129d7aef7f1407d3f295 \ + --hash=sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea \ + --hash=sha256:418cf3f2111bc80e0933b2cd8cd04f286338bb88bdc7bc8e6dd775ebde60b5e0 \ + --hash=sha256:44edc647873928551a01e7a563d7452ccdebee747728c1080d881d68af7b997e \ + --hash=sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac \ + --hash=sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9 \ + --hash=sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7 \ + --hash=sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35 \ + --hash=sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb \ + --hash=sha256:5cf4e27da7e3fbed4d6c3d8e797387aaad68102272f8f9752883bc32d61cb87b \ + --hash=sha256:5e0b74767e5f8c593e8c9b5912019159ed0533c70051e9cce3e8b6aa699fcd69 \ + --hash=sha256:5ed875a24292240029e4483f9d4a4b8a1ae08843b9c54f43fcc11e404532a8a5 \ + --hash=sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b \ + --hash=sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c \ + --hash=sha256:6344df0d5755a2c9a276d4473ae6b90647e216ab4757f8426893b5dd2ac3f369 \ + --hash=sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd \ + --hash=sha256:652cb6edd41e718550aad172851962662ff2681490a8a711af6a4d288dd96824 \ + --hash=sha256:66291b10affd76d76f54fad28e22e51719ef9ba22b29e1d7d03d6777a9174198 \ + --hash=sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065 \ + --hash=sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c \ + --hash=sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c \ + --hash=sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764 \ + --hash=sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196 \ + --hash=sha256:8098f252adfa6c80ab48096053f512f2321f0b998f98150cea9bd23d83e1467b \ + --hash=sha256:850774a7879607d3a6f50d36d04f00ee69e7fc816450e5f7e58d7f17f1ae5c00 \ + --hash=sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac \ + --hash=sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8 \ + --hash=sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e \ + --hash=sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28 \ + --hash=sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3 \ + --hash=sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5 \ + --hash=sha256:9c57bb8c96f6d1808c030b1687b9b5fb476abaa47f0db9c0101f5e9f394e97f4 \ + --hash=sha256:9c7708761fccb9397fe64bbc0395abcae8c4bf7b0eac081e12b809bf47700d0b \ + --hash=sha256:9f3bfb4965eb874431221a3ff3fdcddc7e74e3b07799e0e84ca4a0f867d449bf \ + --hash=sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5 \ + --hash=sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702 \ + --hash=sha256:b30236e45cf30d2b8e7b3e85881719e98507abed1011bf463a8fa23e9c3e98a8 \ + --hash=sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788 \ + --hash=sha256:b865addae83924361678b652338317d1bd7e79b1f4596f96b96c77a5a34b34da \ + --hash=sha256:b8bb0864c5a28024fac8a632c443c87c5aa6f215c0b126c449ae1a150412f31d \ + --hash=sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc \ + --hash=sha256:bdb2c67c6c1390b63c6ff89f210c8fd09d9a1217a465701eac7316313c915e4c \ + --hash=sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba \ + --hash=sha256:c2514fceb77bc5e7a2f7adfaa1feb2fb311607c9cb518dbc378688ec73d8292f \ + --hash=sha256:c3355370a2c156cffb25e876646f149d5d68f5e0a3ce86a5084dd0b64a994917 \ + --hash=sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5 \ + --hash=sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26 \ + --hash=sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f \ + --hash=sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b \ + --hash=sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be \ + --hash=sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c \ + --hash=sha256:efd7b85f94a6f21e4932043973a7ba2613b059c4a000551892ac9f1d11f5baf3 \ + --hash=sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6 \ + --hash=sha256:fa160448684b4e94d80416c0fa4aac48967a969efe22931448d853ada8baf926 \ + --hash=sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0 +rdkit==2026.3.4 \ + --hash=sha256:1218f1d3d1e348ced44f815b2acc8d8a941da1ad67e3cf64e9c20fa090522882 \ + --hash=sha256:19634f932126c300d0fb0e02287d9eea601f43262a8980645df16c7dd55ba1a7 \ + --hash=sha256:1d6f75e5b853eba73f7be1832a4b6f4caff4ff22fa4ba078c53352c9483c3f48 \ + --hash=sha256:1f6fd3e9ee43211b9458268a0a21a5c3db7d27d2aa2eef9cb837dc4302d4426d \ + --hash=sha256:5c0b86cd5243713c07203e602899b623f169752cbe5a20b311e5cbc3ae12dca0 \ + --hash=sha256:5cda49e4e338789090ababbb1ebf8577b67aee7386a1ea9df4cee7d3fee67cec \ + --hash=sha256:65fdda1efb852a22faafe2ff060691b23c60b3c74a8f2ca41198483862fcf352 \ + --hash=sha256:68f27f12063091115df76d0adbd3d1ec7d7ae8a0ddf103cfb88a1fee4670f2cd \ + --hash=sha256:71cc4404614c05f37a2cdb2af028d2e86f0a314af07fa2d175b01fe064c76e29 \ + --hash=sha256:7d5eeb9212ff64410fd37c4b561603da52193f9cdbe40a787a5159e1076c0981 \ + --hash=sha256:974830cdcdb95f825fc1038824af0031e7ac537e526485a5f30e58260d03bc39 \ + --hash=sha256:98cb809bc8e3f7168c7909670a418096830c95d5b386e4b124d1381f9730e5c1 \ + --hash=sha256:a41dde42ecb24e7d93d62b89e8e5d267b02bbac764f965f3968296c99234c685 \ + --hash=sha256:a9a71fe0f14e9a08818530aed77e89bae8506f479bb99b0e362ebb7fe156899a \ + --hash=sha256:bec605934ad781e91aa10d434ac986ee18e191ae9dfa65fcabb1784aadf33edb \ + --hash=sha256:c15431fc33f456998411c00091f1403e484b8a26e77742591e2b318c4ca9a857 \ + --hash=sha256:c6d25f79fe7c22ab0b6172f33417dba83518e4ad098bcf5c84f7d22ca08b1f18 \ + --hash=sha256:cee9a591c4f315b9e10138d9e7a69033abc4b229a960b45fef5cd179f2f168fc \ + --hash=sha256:d63ac1384a7a6af9742cafa23139f9b1ce7639e48d2a6b1b793600065442a133 \ + --hash=sha256:df1ec162c226c177e37a8325aa217e89fb0e4d4db68f919e53b9fce144c13dc9 +regex==2026.7.19 \ + --hash=sha256:062f8cb7a9739c4835d22bd96f370c59aba89f257adcfa53be3cc209e08d3ae0 \ + --hash=sha256:064f1760a5a4ade65c5419be23e782f29147528e8a66e0c42dd4cedb8d4e9fc6 \ + --hash=sha256:09523a592938aa9f587fb74467c63ff0cf88fc3df14c82ab0f0517dcf76aaa62 \ + --hash=sha256:09d3007fc76249a83cdd33de160d50e6cb77f54e09d8fa9e7148e10607ce24af \ + --hash=sha256:09f3e5287f94f17b709dc9a9e70865855feee835c861613be144218ce4ca82cc \ + --hash=sha256:0c41c63992bf1874cebb6e7f56fd7d3c007924659a604ae3d90e427d40d4fd13 \ + --hash=sha256:0e9554c8785eac5cffe6300f69a91f58ba72bc88a5f8d661235ad7c6aa5b8ccd \ + --hash=sha256:1123ef4211d763ee771d47916a1596e2f4915794f7aabdc1adcb20e4249a6951 \ + --hash=sha256:15b364b9b98d6d2fe1a85034c23a3180ff913f46caddc3895f6fd65186255ccc \ + --hash=sha256:1649eb39fcc9ea80c4d2f110fde2b8ab2aef3877b98f02ab9b14e961f418c511 \ + --hash=sha256:17ed5692f6acc4183e98331101a5f9e4f64d72fe58b753da4d444a2c77d05b12 \ + --hash=sha256:199535629f25caf89698039af3d1ad5fcae7f933e2112c73f1cdf49165c99518 \ + --hash=sha256:1c398716054621aa300b3d411f467dda903806c5da0df6945ab73982b8d115db \ + --hash=sha256:1d3372064506b94dd2c67c845f2db8062e9e9ba84d04e33cb96d7d33c11fe1ae \ + --hash=sha256:1d58561843f0ff7dc78b4c28b5e2dc388f3eff94ebc8a232a3adba961fc00009 \ + --hash=sha256:1d793a7988e04fcb1e2e135567443d82173225d657419ec09414a9b5a145b986 \ + --hash=sha256:1ebac3474b8589fce2f9b225b650afd61448f7c73a5d0255a10cc6366471aed1 \ + --hash=sha256:20568e182eb82d39a6bf7cff3fd58566f14c75c6f74b2c8c96537eecf9010e3a \ + --hash=sha256:22a992de9a0d91bda927bf02b94351d737a0302905432c88a53de7c4b9ce62e2 \ + --hash=sha256:2955907b7157a6660f27079edf7e0229e9c9c5325c77a2ef6a890cba91efa6f0 \ + --hash=sha256:2c4e61e2e1be56f63ec3cc618aa9e0de81ef6f43d177205451840022e24f5b78 \ + --hash=sha256:2cc3460cedf7579948486eab03bc9ad7089df4d7281c0f47f4afe03e8d13f02d \ + --hash=sha256:2ce9e679f776649746729b6c86382da519ef649c8e34cc41df0d2e5e0f6c36d4 \ + --hash=sha256:2ef7eeb108c47ce7bcc9513e51bcb1bf57e8f483d52fce68a8642e3527141ae0 \ + --hash=sha256:3080a7fd38ef049bd489e01c970c97dd84ff446a885b0f1f6b26d9b1ad13ce11 \ + --hash=sha256:343a4504e3fb688c47cad451221ca5d4814f42b1e16c0065bde9cbf7f473bd52 \ + --hash=sha256:36aacfb15faaff3ced55afbf35ec72f50d4aee22082c4f7fe0573a33e2fca92e \ + --hash=sha256:3d3143f159261b1ce5b24c261c590e5913370c3200c5e9ebbb92b5aa5e111902 \ + --hash=sha256:40b34dd88658e4fedd2fddbf0275ac970d00614b731357f425722a3ed1983d11 \ + --hash=sha256:4458124d71339f505bf1fb94f69fd1bb8fa9d2481eebfef27c10ef4f2b9e12f6 \ + --hash=sha256:4896db1f4ce0576765b8272aa922df324e0f5b9bb2c3d03044ff32a7234a9aba \ + --hash=sha256:4a0530bb1b8c1c985e7e2122e2b4d3aedd8a3c21c6bfddae6767c4405668b56e \ + --hash=sha256:4aa5435cdb3eb6f55fe98a171b05e3fbcd95fadaa4aa32acf62afd9b0cfdbcac \ + --hash=sha256:4c3501bfa814ab07b5580741f9bf78dfdfe146a04057f82df9e2402d2a975939 \ + --hash=sha256:4e5413bd5f13d3a4e3539ca98f70f75e7fca92518dd7f117f030ebedd10b60cb \ + --hash=sha256:4e6883a021db30511d9fb8cfb0f222ce1f2c369f7d4d8b0448f449a93ba0bdfc \ + --hash=sha256:52579c60a6078be70a0e49c81d6e56d677f34cd439af281a0083b8c7bc75c095 \ + --hash=sha256:555497390743af1a65045fa4527782d10ff5b88970359412baa4a1e628fe393b \ + --hash=sha256:56ad4d9f77df871a99e25c37091052a02528ec0eb059de928ee33956b854b45b \ + --hash=sha256:571fde9741eb0ccde23dd4e0c1d50fbae910e901fa7e629faf39b2dda740d220 \ + --hash=sha256:572fc57b0009c735ee56c175ea021b637a15551a312f56734277f923d6fd0f6c \ + --hash=sha256:59787bd5f8c70aa339084e961d2996b53fbdeab4d5393bba5c1fe1fc32e02bae \ + --hash=sha256:5a2721c8720e2cb3c209925dfb9200199b4b07361c9e01d321719404b21458b3 \ + --hash=sha256:5cc26a66e212fa5d6c6170c3a40d99d888db3020c6fdab1523250d4341382e44 \ + --hash=sha256:5ebee1ee89c39c953baac6924fcde08c5bb427c4057510862f9d7c7bdb3d8665 \ + --hash=sha256:60be8693a1dadc210bbcbc0db3e26da5f7d01d1d5a3da594e99b4fa42df404f5 \ + --hash=sha256:618a0aed532be87294c4477b0481f3aa0f1520f4014a4374dd4cf789b4cd2c97 \ + --hash=sha256:61bb1bd45520aacd56dd80943bd34991fb5350afdd1f36f2282230fd5154a218 \ + --hash=sha256:6383cd2ed53a646c659ba1fe65727db76437fdaa069e697a0b44a51d5843d864 \ + --hash=sha256:64729333167c2dcaaa56a331d40ee097bd9c5617ffd51dabb09eaddafb1b532e \ + --hash=sha256:64b6ca7391a1395c2638dd5c7456d67bea44fc6c5e8e92c5dc8aa6a8f23292b4 \ + --hash=sha256:65dcd28d3eba2ab7c2fd906485cc301392b47cc2234790d27d4e4814e02cdfda \ + --hash=sha256:65fa6cb38ed5e9c3637e68e544f598b39c3b86b808ed0627a67b68320384b459 \ + --hash=sha256:66bd62c59a5427746e8c44becae1d9b99d22fb13f30f492083dfb9ad7c45cc18 \ + --hash=sha256:6e44c0e7c5664be20aee92085153150c0a7967310a73a43c0f832b7cd35d0dd3 \ + --hash=sha256:6f8c6e7a1cfa3dc9d0ee2de0e65e834537fa29992cc3976ffec914afc35c5dd5 \ + --hash=sha256:7322ec6cc9fba9d49ab888bb82d67ac5625627aa168f0165139b17018df3fb8a \ + --hash=sha256:73b133a9e6fb512858e7f065e96f1180aa46646bc74a83aea62f1d314f3dd035 \ + --hash=sha256:73f272fba87b8ccfe70a137d02a54af386f6d27aa509fbffdd978f5947aae1aa \ + --hash=sha256:7e77b324909c1617cbb4c668677e2c6ae13f44d7c1de0d4f15f2e3c10f3315b5 \ + --hash=sha256:80115dd39481fd3a4b4080220799dbcacb921a844de4b827264ececacbe17c78 \ + --hash=sha256:87ccab0db8d5f4fbb0272642113c1adb2ffc698c16d3a0944580222331fa7a20 \ + --hash=sha256:89dfee3319f5ae3f75ebd5c2445a809bb320252ba5529ffdafea4ef25d79cf1a \ + --hash=sha256:8ac59a0900474a52b7c04af8196affc22bd9842acb0950df12f7b813e983609a \ + --hash=sha256:8cae6fd77a5b72dae505084b1a2ee0360139faf72fedbab667cd7cc65aae7a6a \ + --hash=sha256:8d3469c91dd92ee41b7c95280edbd975ef1ba9195086686623a1c6e8935ce965 \ + --hash=sha256:90c633e7e8d6bf4e992b8b36ce69e018f834b641dd6de8cea6d78c06ffa119c5 \ + --hash=sha256:93db40c8de0815baab96a06e08a984bac71f989d13bab789e382158c5d426797 \ + --hash=sha256:9724e6cb5e478cd7d8cabf027826178739cb18cf0e117d0e32814d479fa02276 \ + --hash=sha256:98c6ac18480fcdb33f35439183f1d2e79760ab41930309c6d951cb1f8e46694c \ + --hash=sha256:9a15e785f244f3e07847b984ce8773fc3da10a9f3c131cc49a4c5b4d672b4547 \ + --hash=sha256:9b60d7814174f059e5de4ab98271cc5ba9259cfea55273a81544dceea32dc8d9 \ + --hash=sha256:9be2a6647740dd3cca6acb24e87f03d7632cd280dbce9bbe40c26353a215a45d \ + --hash=sha256:9c7472192ebfad53a6be7c4a8bfb2d64b81c0e93a1fc8c57e1dd0b638297b5d1 \ + --hash=sha256:9dce8ec9695f531a1b8a6f314fd4b393adcccf2ea861db480cdf97a301d01a68 \ + --hash=sha256:9e50d748a32da622f256e8d505867f5d3c43a837c6a9f0efb149655fadd1042a \ + --hash=sha256:a81758ed242b861b72e778ba34d41366441a2e10b16b472784c88da2dea7e2dd \ + --hash=sha256:ac777001cdfc28b72477d93c8564bb7583081ea8fb45cdca3d568e0a4f87183c \ + --hash=sha256:b2b506b1788df5fecd270a10d5e70a95fe77b87ea2b370a318043f6f5f817ee6 \ + --hash=sha256:b2ea4a3e8357be8849e833beeae757ac3c7a6b3fc055c03c808a53c91ad30d82 \ + --hash=sha256:bf1516fe58fc104f39b2d1dbe2d5e27d0cd45c4be2e42ba6ee0cc763701ec3c7 \ + --hash=sha256:c0d702548d89d572b2929879bc883bb7a4c4709efafe4512cadee56c55c9bd15 \ + --hash=sha256:c10b82c2634df08dfb13b1f04e38fe310d086ee092f4f69c0c8da234251e556e \ + --hash=sha256:c42572142ed0b9d5d261ba727157c426510da78e20828b66bbb855098b8a4e38 \ + --hash=sha256:c4585c3e64b4f9e583b4d2683f18f5d5d872b3d71dcf24594b74ecc23602fa96 \ + --hash=sha256:c639ea314df70a7b2811e8020448c75af8c9445f5a60f8a4ced81c306a9380c2 \ + --hash=sha256:c670fe7be5b6020b76bc6e8d2196074657e1327595bca93a389e1a76ab130ad8 \ + --hash=sha256:cc1b2440423a851fad781309dd87843868f4f66a6bcd1ddb9225cf4ec2c84732 \ + --hash=sha256:cd3584591ea4429026cdb931b054342c2bcf189b44ff367f8d5c15bc092a2966 \ + --hash=sha256:d15df07081d91b76ff20d43f94592ee110330152d617b730fdbe5ef9fb680053 \ + --hash=sha256:d19662dbedbe783d323196312d38f5ba53cf56296378252171985da6899887d3 \ + --hash=sha256:d24ecb4f5e009ea0bd275ee37ad9953b32005e2e5e60f8bbae16da0dbbf0d3a0 \ + --hash=sha256:d446c6ac40bb6e05025ccee55b84d80fe9bf8e93010ffc4bb9484f13d498835f \ + --hash=sha256:d51ffd3427640fa2da6ade574ceba932f210ad095f65fcc450a2b0a0d454868e \ + --hash=sha256:d6ce43a0269d68cee79a7d1ade7def53c20f8f2a047b92d7b5d5bcc73ae88327 \ + --hash=sha256:d721e53758b2cca74990185eb0671dd466d7a388a1a45d0c6f4c13cef41a68ac \ + --hash=sha256:d7da47a0f248977f08e2cb659ff3c17ddc13a4d39b3a7baa0a81bf5b415430f6 \ + --hash=sha256:db47b561c9afd884baa1f96f797c9ca369872c4b65912bc691cfa99e68340af2 \ + --hash=sha256:dbe6493fbd27321b1d1f2dd4f5c7e5bd4d8b1d7cab7f32fd67db3d0b2ed8248a \ + --hash=sha256:dbece16025afda5e3031af0c4059207e61dcf73ef13af844964f57f387d1c435 \ + --hash=sha256:ddd67571c10869f65a5d7dde536d1e066e306cc90de57d7de4d5f34802428bb5 \ + --hash=sha256:de9208bb427130c82a5dbfd104f92c8876fc9559278c880b3002755bbbe9c83d \ + --hash=sha256:e30d40268a28d54ce0437031750497004c22602b8e3ab891f759b795a003b312 \ + --hash=sha256:e8b0abe7d870f53ca5143895fef7d1041a0c831a140d3dc2c760dd7ba25d4a8b \ + --hash=sha256:f035d9dc1d25eff9d361456572231c7d27b5ccd473ca7dc0adfce732bd006d40 \ + --hash=sha256:f04b9f56b0e0614c0126be12c2c2d9f8850c1e57af302bd0a63bed379d4af974 \ + --hash=sha256:f0fa4fa9c3632d708742baf2282f2055c11d888a790362670a403cbf48a2c404 \ + --hash=sha256:f2e7f8e2ab6c2922be02c7ec45185aa5bd771e2e57b95455ee343a44d8130dff \ + --hash=sha256:f8f6fa298bb4f7f58a33334406218ba74716e68feddf5e4e54cd5d8082705abf \ + --hash=sha256:fbf300e2070bb35038660b3be1be4b91b0024edb41517e6996320b49b92b4175 \ + --hash=sha256:fce7760bf283405b2c7999cab3da4e72f7deca6396013115e3f7a955db9760da \ + --hash=sha256:fcee38cd8e5089d6d4f048ba1233b3ad76e5954f545382180889112ff5cb712d \ + --hash=sha256:fe31f28c94402043161876a258a9c6f757cb485905c7614ce8d6cd40e6b7bdc1 \ + --hash=sha256:ffd8893ccc1c2fce6e0d6ca402d716fe1b29db70c7132609a05955e31b2aa8f2 +requests==2.34.2 \ + --hash=sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0 \ + --hash=sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed +s3transfer==0.19.1 \ + --hash=sha256:d3d6371dc3f1e5c5427b2b457bcf13bcf87bec334c95aed18642eae61f6926f3 \ + --hash=sha256:d5fd7005ee39307455ad5f310b5ea67f4b1960d7fed5b3671ee50c249de675de +safetensors==0.5.3 \ + --hash=sha256:1077f3e94182d72618357b04b5ced540ceb71c8a813d3319f1aba448e68a770d \ + --hash=sha256:11bce6164887cd491ca75c2326a113ba934be596e22b28b1742ce27b1d076467 \ + --hash=sha256:21d01c14ff6c415c485616b8b0bf961c46b3b343ca59110d38d744e577f9cce7 \ + --hash=sha256:32c3ef2d7af8b9f52ff685ed0bc43913cdcde135089ae322ee576de93eae5135 \ + --hash=sha256:37f1521be045e56fc2b54c606d4455573e717b2d887c579ee1dbba5f868ece04 \ + --hash=sha256:391ac8cab7c829452175f871fcaf414aa1e292b5448bd02620f675a7f3e7abb9 \ + --hash=sha256:4a243be3590bc3301c821da7a18d87224ef35cbd3e5f5727e4e0728b8172411e \ + --hash=sha256:799021e78287bac619c7b3f3606730a22da4cda27759ddf55d37c8db7511c74b \ + --hash=sha256:836cbbc320b47e80acd40e44c8682db0e8ad7123209f69b093def21ec7cafd11 \ + --hash=sha256:8bd84b12b1670a6f8e50f01e28156422a2bc07fb16fc4e98bded13039d688a0d \ + --hash=sha256:b6b0d6ecacec39a4fdd99cc19f4576f5219ce858e6fd8dbe7609df0b8dc56965 \ + --hash=sha256:bd20eb133db8ed15b40110b7c00c6df51655a2998132193de2f75f72d99c7073 \ + --hash=sha256:cead1fa41fc54b1e61089fa57452e8834f798cb1dc7a09ba3524f1eb08e0317a \ + --hash=sha256:cfc0ec0846dcf6763b0ed3d1846ff36008c6e7290683b61616c4b040f6a54ace \ + --hash=sha256:df26da01aaac504334644e1b7642fa000bfec820e7cef83aeac4e355e03195ff +scikit-learn==1.9.0 \ + --hash=sha256:051075bda8b7aab87b1906ab3d4740a1e1224a19d7b3781a576736edc94e76aa \ + --hash=sha256:056c92bb67ad4c28463c2f2653d9701449201e7e7a9e94e321be0f71c4fef2b8 \ + --hash=sha256:147e9329ef0e39f75d4cffa02b2aa48d827832684926cd5210d9a2cb5c57246b \ + --hash=sha256:1b944b6db288f6b926e3650026ddafb988929de95d11fc2cc5fa117773c9ba42 \ + --hash=sha256:1fea2cc5677ab49d6f5bade978c866da44957b712d92e9635e8b4f723013c3cb \ + --hash=sha256:24360002ae845e7866522b0a5bbf690802e7bc388cac8663502e78aa98598aa2 \ + --hash=sha256:26e22435f63bcdcf396b574273f29f13dd531f5ea035801f5be10ba1540a4e60 \ + --hash=sha256:2bd41b0d201bc81575531b96b713d3eb5e5f50fb0b82101ff0f92294fdc236ac \ + --hash=sha256:366652351f092b219c248f1e72821e841960a63d8f358f1dcfd54dc1cbdbbc28 \ + --hash=sha256:38c3dcb9a1ffb85505ec53d54c7b4aea0cff70050425a7760c2af661ac85df05 \ + --hash=sha256:4306775fad04cc4b472a1b15af1ae9cede1540fbfcc17fbce3767cd8dc7ae283 \ + --hash=sha256:4ccacf04ca5f4b492158a5f28afe0ace43f81b2571e4b9a66d34848b46128949 \ + --hash=sha256:5162ad10a418c8a282dde04c9aa06965de3e9a65f33c1440c0ae69bb1a09d913 \ + --hash=sha256:5808d98f15c6bf6d9d96d2348c1997392a5888ce7097e664105f930c4bca1277 \ + --hash=sha256:5b934c45c252844a91d69fda3a34cff5e7307e1db10d77cb10a3980312c74713 \ + --hash=sha256:5bad8f8b9950321b54c965fdcbac6c6c55e79e16646b49977bcf3668d3870a1a \ + --hash=sha256:5be45aa4a42a68a533913a6ed736cf309de2226411c79ef8d609a5456f1939b1 \ + --hash=sha256:5dc1818c77575d149e25fce9ef82dd7b7263ae372f03494158668ad632a69759 \ + --hash=sha256:5e50ed4da51974e86e940690e9a3d82e729b62b5a49f7c9bac534d515d39d86f \ + --hash=sha256:64fa347efc1c839c487433e40c5144d38c336e8a2b59c81aa8660373945c2673 \ + --hash=sha256:78fc56eafd4edb9575d2d8950d1dd152061abb573341a1cb7e099fc40f6c6666 \ + --hash=sha256:80746d63bd4b6eaca54d36fe5feaf4d28bb38dc6f9470f81c7cad7c40155f119 \ + --hash=sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557 \ + --hash=sha256:9656acd4e93f74e0b66c8a36c88830a99252dfa900044d36bc2212ae89a47162 \ + --hash=sha256:9db6f4d34e68c8899e4cab27fdf8eafe6ed21f2ba52ceb25ea250cd237f8e47b \ + --hash=sha256:d77f54c017633791bc0225a43e2f8d03745fdcfe4880268fcc4df15f505dec2e \ + --hash=sha256:da76d09304a4706db7cc1e3ebaa3b6b98a67365cc11d2996c4f1e58ba47df714 \ + --hash=sha256:ee1a8db2c18c08e34c7412d4b10be1cac214cd4ea7dc9715a6a327eb49a37c96 \ + --hash=sha256:f401448645a3e7bc115aa3c094097865155b34bff1cba8101857d9104e99074c \ + --hash=sha256:f7e254636164090da847715a27f8e5478feb98c40a9e0ee90cbd277de9e5ceb8 \ + --hash=sha256:fd3a8ef0c758555a3b23c03adaa858af32f7736785ded50ad5991f59c4ed03fa +scipy==1.17.1 \ + --hash=sha256:010f4333c96c9bb1a4516269e33cb5917b08ef2166d5556ca2fd9f082a9e6ea0 \ + --hash=sha256:02ae3b274fde71c5e92ac4d54bc06c42d80e399fec704383dcd99b301df37458 \ + --hash=sha256:08b900519463543aa604a06bec02461558a6e1cef8fdbb8098f77a48a83c8118 \ + --hash=sha256:131f5aaea57602008f9822e2115029b55d4b5f7c070287699fe45c661d051e39 \ + --hash=sha256:158dd96d2207e21c966063e1635b1063cd7787b627b6f07305315dd73d9c679e \ + --hash=sha256:1cc682cea2ae55524432f3cdff9e9a3be743d52a7443d0cba9017c23c87ae2f6 \ + --hash=sha256:1f95b894f13729334fb990162e911c9e5dc1ab390c58aa6cbecb389c5b5e28ec \ + --hash=sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21 \ + --hash=sha256:2040ad4d1795a0ae89bfc7e8429677f365d45aa9fd5e4587cf1ea737f927b4a1 \ + --hash=sha256:2b64ca7d4aee0102a97f3ba22124052b4bd2152522355073580bf4845e2550b6 \ + --hash=sha256:2ceb2d3e01c5f1d83c4189737a42d9cb2fc38a6eeed225e7515eef71ad301dce \ + --hash=sha256:35c3a56d2ef83efc372eaec584314bd0ef2e2f0d2adb21c55e6ad5b344c0dcb8 \ + --hash=sha256:37425bc9175607b0268f493d79a292c39f9d001a357bebb6b88fdfaff13f6448 \ + --hash=sha256:3877ac408e14da24a6196de0ddcace62092bfc12a83823e92e49e40747e52c19 \ + --hash=sha256:3fd1fcdab3ea951b610dc4cef356d416d5802991e7e32b5254828d342f7b7e0b \ + --hash=sha256:41b71f4a3a4cab9d366cd9065b288efc4d4f3c0b37a91a8e0947fb5bd7f31d87 \ + --hash=sha256:43af8d1f3bea642559019edfe64e9b11192a8978efbd1539d7bc2aaa23d92de4 \ + --hash=sha256:45abad819184f07240d8a696117a7aacd39787af9e0b719d00285549ed19a1e9 \ + --hash=sha256:4b400bdc6f79fa02a4d86640310dde87a21fba0c979efff5248908c6f15fad1b \ + --hash=sha256:4eb6c25dd62ee8d5edf68a8e1c171dd71c292fdae95d8aeb3dd7d7de4c364082 \ + --hash=sha256:581b2264fc0aa555f3f435a5944da7504ea3a065d7029ad60e7c3d1ae09c5464 \ + --hash=sha256:5cf36e801231b6a2059bf354720274b7558746f3b1a4efb43fcf557ccd484a87 \ + --hash=sha256:5e3c5c011904115f88a39308379c17f91546f77c1667cea98739fe0fccea804c \ + --hash=sha256:6609bc224e9568f65064cfa72edc0f24ee6655b47575954ec6339534b2798369 \ + --hash=sha256:6e3dcd57ab780c741fde8dc68619de988b966db759a3c3152e8e9142c26295ad \ + --hash=sha256:6fac755ca3d2c3edcb22f479fceaa241704111414831ddd3bc6056e18516892f \ + --hash=sha256:744b2bf3640d907b79f3fd7874efe432d1cf171ee721243e350f55234b4cec4c \ + --hash=sha256:74cbb80d93260fe2ffa334efa24cb8f2f0f622a9b9febf8b483c0b865bfb3475 \ + --hash=sha256:766e0dc5a616d026a3a1cffa379af959671729083882f50307e18175797b3dfd \ + --hash=sha256:7bdf2da170b67fdf10bca777614b1c7d96ae3ca5794fd9587dce41eb2966e866 \ + --hash=sha256:7ff200bf9d24f2e4d5dc6ee8c3ac64d739d3a89e2326ba68aaf6c4a2b838fd7d \ + --hash=sha256:844e165636711ef41f80b4103ed234181646b98a53c8f05da12ca5ca289134f6 \ + --hash=sha256:8a604bae87c6195d8b1045eddece0514d041604b14f2727bbc2b3020172045eb \ + --hash=sha256:94055a11dfebe37c656e70317e1996dc197e1a15bbcc351bcdd4610e128fe1ca \ + --hash=sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0 \ + --hash=sha256:9cdc1a2fcfd5c52cfb3045feb399f7b3ce822abdde3a193a6b9a60b3cb5854ca \ + --hash=sha256:9ecb4efb1cd6e8c4afea0daa91a87fbddbce1b99d2895d151596716c0b2e859d \ + --hash=sha256:a3472cfbca0a54177d0faa68f697d8ba4c80bbdc19908c3465556d9f7efce9ee \ + --hash=sha256:a4328d245944d09fd639771de275701ccadf5f781ba0ff092ad141e017eccda4 \ + --hash=sha256:a48a72c77a310327f6a3a920092fa2b8fd03d7deaa60f093038f22d98e096717 \ + --hash=sha256:a720477885a9d2411f94a93d16f9d89bad0f28ca23c3f8daa521e2dcc3f44d49 \ + --hash=sha256:a77cbd07b940d326d39a1d1b37817e2ee4d79cb30e7338f3d0cddffae70fcaa2 \ + --hash=sha256:a9956e4d4f4a301ebf6cde39850333a6b6110799d470dbbb1e25326ac447f52a \ + --hash=sha256:adb2642e060a6549c343603a3851ba76ef0b74cc8c079a9a58121c7ec9fe2350 \ + --hash=sha256:beeda3d4ae615106d7094f7e7cef6218392e4465cc95d25f900bebabfded0950 \ + --hash=sha256:c80be5ede8f3f8eded4eff73cc99a25c388ce98e555b17d31da05287015ffa5b \ + --hash=sha256:cc90d2e9c7e5c7f1a482c9875007c095c3194b1cfedca3c2f3291cdc2bc7c086 \ + --hash=sha256:cd96a1898c0a47be4520327e01f874acfd61fb48a9420f8aa9f6483412ffa444 \ + --hash=sha256:d2650c1fb97e184d12d8ba010493ee7b322864f7d3d00d3f9bb97d9c21de4068 \ + --hash=sha256:d30e57c72013c2a4fe441c2fcb8e77b14e152ad48b5464858e07e2ad9fbfceff \ + --hash=sha256:d59c30000a16d8edc7e64152e30220bfbd724c9bbb08368c054e24c651314f0a \ + --hash=sha256:dbc12c9f3d185f5c737d801da555fb74b3dcfa1a50b66a1a93e09190f41fab50 \ + --hash=sha256:e18f12c6b0bc5a592ed23d3f7b891f68fd7f8241d69b7883769eb5d5dfb52696 \ + --hash=sha256:e19ebea31758fac5893a2ac360fedd00116cbb7628e650842a6691ba7ca28a21 \ + --hash=sha256:e30bdeaa5deed6bc27b4cc490823cd0347d7dae09119b8803ae576ea0ce52e4c \ + --hash=sha256:eb092099205ef62cd1782b006658db09e2fed75bffcae7cc0d44052d8aa0f484 \ + --hash=sha256:eee2cfda04c00a857206a4330f0c5e3e56535494e30ca445eb19ec624ae75118 \ + --hash=sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3 \ + --hash=sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea \ + --hash=sha256:f8885db0bc2bffa59d5c1b72fad7a6a92d3e80e7257f967dd81abb553a90d293 \ + --hash=sha256:fcb310ddb270a06114bb64bbe53c94926b943f5b7f0842194d585c65eb4edd76 +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 +six==1.17.0 \ + --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \ + --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81 +stack-data==0.6.3 \ + --hash=sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9 \ + --hash=sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695 +sympy==1.14.0 \ + --hash=sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517 \ + --hash=sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5 +tenacity==9.1.4 \ + --hash=sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55 \ + --hash=sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a +threadpoolctl==3.6.0 \ + --hash=sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb \ + --hash=sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e +tokenizers==0.22.2 \ + --hash=sha256:143b999bdc46d10febb15cbffb4207ddd1f410e2c755857b5a0797961bbdc113 \ + --hash=sha256:1a62ba2c5faa2dd175aaeed7b15abf18d20266189fb3406c5d0550dd34dd5f37 \ + --hash=sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e \ + --hash=sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001 \ + --hash=sha256:1e50f8554d504f617d9e9d6e4c2c2884a12b388a97c5c77f0bc6cf4cd032feee \ + --hash=sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7 \ + --hash=sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd \ + --hash=sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4 \ + --hash=sha256:319f659ee992222f04e58f84cbf407cfa66a65fe3a8de44e8ad2bc53e7d99012 \ + --hash=sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67 \ + --hash=sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a \ + --hash=sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5 \ + --hash=sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917 \ + --hash=sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c \ + --hash=sha256:64d94e84f6660764e64e7e0b22baa72f6cd942279fdbb21d46abd70d179f0195 \ + --hash=sha256:753d47ebd4542742ef9261d9da92cd545b2cacbb48349a1225466745bb866ec4 \ + --hash=sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a \ + --hash=sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc \ + --hash=sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92 \ + --hash=sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5 \ + --hash=sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48 \ + --hash=sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b \ + --hash=sha256:e10bf9113d209be7cd046d40fbabbaf3278ff6d18eb4da4c500443185dc1896c \ + --hash=sha256:f01a9c019878532f98927d2bacb79bbb404b43d3437455522a00a30718cdedb5 +torch==2.13.0+cu130 \ + --hash=sha256:2efab1e83604ca628c6d85b9e188c153690980498d1297081a9dad704919303c \ + --hash=sha256:45e97bd9bc0416f4f4190b5098c55119a389fa5a7c8bbf2639f08f1d04e0a0dc \ + --hash=sha256:5ebd552c887e707c8e64927aceb8377ca2e81588c4e7494bcd23cb8ac0aca14d \ + --hash=sha256:794664c05a3470a5e738447b54495cc6bbf1efca48bf2794270a897d9f4356c4 \ + --hash=sha256:83767298109148f10de124287b392bc58d74f9938e9430f26f9fd669efb8af9c \ + --hash=sha256:8db7338e6895c3d4bd89a02ff4209507d1f0cf2ffeb3b898538b5a07d1ea8c1e \ + --hash=sha256:b8a6b58c0176dd532254f6622fb1dffef6b38ad7cb67fcd0beb673066c8c710c \ + --hash=sha256:cd11e8876fce3cbc941c576c9f87730f0ff2ca96a9d3e1aa6913bf48304c8bcb \ + --hash=sha256:e85d18b0c51744b25fab85dbf590a4f00644da432af0e9d03c62acbd2f96ea94 +tqdm==4.69.0 \ + --hash=sha256:700c5e85dcd5f009dd6222588a29180a193a748247a5d855b4d67db93d79a53b \ + --hash=sha256:9979978912be667a6ef21fd5d8abf54e324e63d82f7f43c360792ebc2bc4e622 +traitlets==5.15.1 \ + --hash=sha256:770a53705f84b81ac107e83a1b3328ff2dae16094d8fc3cfc004e4b22dfd8e92 \ + --hash=sha256:7b1c07854fe25acb39e009bae49f11b79ff6cbb2f27999104e9110e7a6b53722 +triton==3.7.1 \ + --hash=sha256:018d1a55f4fc01673cffe86d5bc35de7f111f87d579609a9af2e5ad1827826a6 \ + --hash=sha256:1f046a8d2615e389b1b8e8a4d100a03f3cb1df4bc1a107a6b678b186443bb2c9 \ + --hash=sha256:225910e79149807de74a0ca63160aa285956a8d64a2951d558997c57e2584ae1 \ + --hash=sha256:34894d51aff1abf7b017fbd0c5e9fe7211305c3599157fca251e3f42bc8f00cf \ + --hash=sha256:3d62b1df358f49fcc229d5b77e7a346d66a93e93c19f6872936eadba9aac5f01 \ + --hash=sha256:410d3786b2dcf1ef792046670787dc742ba8f800cf2c6dddf56e0656f2260db6 \ + --hash=sha256:498302da4866b62d7184825bb3ced2348fe57b5caeaef882db23fc3f746e5cc1 \ + --hash=sha256:59688b9b924f887316dcd0fae9e8cbe697ee1d1f6ab386723982cc9ecdebee01 \ + --hash=sha256:6aadad4b7b001f2bc025667af14b0b32fb803c69534381abbae24971891c1859 \ + --hash=sha256:757802f248035abef085ffe28e9f36416f303c769aace1d346a9909eb2e90fd0 \ + --hash=sha256:7f02d894c176e9b80323c4679049cb12219890027550194550548a04c70b5c2e \ + --hash=sha256:996b460a85a134bbf810f87dea9ad42fc434d9ae1cbb948313ae83dd27f432f3 \ + --hash=sha256:9e291fca77d74a2c1ce0d3a5eded017a2dd4c688c8986c36e2b0a5f33d956b4f \ + --hash=sha256:c56cb810349d3699020b5b7695429b9b997574dab1b9f46b2b63fa23d7954894 \ + --hash=sha256:e04b4252ddbcbc6d4f518bb3eb3d586b2e77d40ed32801d9a75116d8d6c1169b \ + --hash=sha256:f091fe850657971c7191f7aa5bc9037669ad045da9d37eabc751926adbb456de +typing-extensions==4.16.0 \ + --hash=sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8 \ + --hash=sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5 +urllib3==2.7.0 \ + --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \ + --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897 +wcwidth==0.8.2 \ + --hash=sha256:91fbef97204b96a3d4d421609b80340b760cf33e26da123ff243d76b1fda8dda \ + --hash=sha256:d63947694a0539a1d51e01eda7caf800c291020e6cdd7e28ad7b14dd33ad4f85 +wheel==0.47.0 \ + --hash=sha256:212281cab4dff978f6cedd499cd893e1f620791ca6ff7107cf270781e587eced \ + --hash=sha256:cc72bd1009ba0cf63922e28f94d9d83b920aa2bb28f798a31d0691b02fa3c9b3 +widgetsnbextension==4.0.15 \ + --hash=sha256:8156704e4346a571d9ce73b84bee86a29906c9abfd7223b7228a28899ccf3366 \ + --hash=sha256:de8610639996f1567952d763a5a41af8af37f2575a41f9852a38f947eb82a3b9 +zstd==1.5.6.1 \ + --hash=sha256:0eb553d2979645b373730bc1b46fdaf30d668000785b5943de1f5bbd8206b7c7 \ + --hash=sha256:1239f4ad597c8796cfceddcf1d9062125e0b4ea75dbd369e18f1f13bf687ce63 \ + --hash=sha256:136b6ead97209b0309e50751da171a056a95e9b742b4cbfd11d8bab68d379533 \ + --hash=sha256:1573926cdd0ad6620df2bf61f8a8e242205e0fc73efab7ff69fc425ab20aa906 \ + --hash=sha256:159b79922e90e387b42986074339c670b224c43d15feb9b98463ba1abd212947 \ + --hash=sha256:1692fd5e120f1b86aeda81492c1eec54c42631bcd73943c68f29c5a13fbc96f1 \ + --hash=sha256:16cfe4ce5987587dcd6a79358dcdedc091a5c1f8ab9defac59024400de34ce77 \ + --hash=sha256:177f19b76efd6e2e7d91c0557638dc49de6c8794c109ccdf0ea4684f63d2fe24 \ + --hash=sha256:17917107e3a2766629646b21dde257e085408086b95a4a1fd62153e77784fcc2 \ + --hash=sha256:1e3a20d23b2e8207d6786bdeff01a3c8e9d3da34e46e50a7ddcac71bec5915b9 \ + --hash=sha256:1ebd16e902398161c0de066c65a7ee36eb0a3b91087eafb75e48a120a955de61 \ + --hash=sha256:2128ea4bb903471c1e763e50a1f301bd37e70735eee9979dbef60244605fa5c7 \ + --hash=sha256:22e6829cbf1994999e281ffe749f110d4544861566e3f0a5e3c93d63f52cbab0 \ + --hash=sha256:23aef9526a4ea90cf47030c7a57e4cedd398cf4ba1ec2798b51cafe47aff75fd \ + --hash=sha256:2498c1e02398b724249d648e11de196ceca112d1e37270af4ddf359bc29ef10a \ + --hash=sha256:251dd039bc11da139e26007b04b3a43d8c7e9ea0f399367dbe6634213f96c562 \ + --hash=sha256:255a957f94c618f0af47c3949f27934c953084eac8941a7df455743c9fc2904d \ + --hash=sha256:2644cbd2b0cfe4a7492a71bb89c780c27f0a461984e443855624e6fe3d1c29e6 \ + --hash=sha256:29c5629d5d6f18ce5d52b54b82ce7abc5077e7557279947c050430537db251d2 \ + --hash=sha256:2dc857557e19afef00dd90e4d4cd7451b3add21826003153302e90767ec854d4 \ + --hash=sha256:343540559f75e4cbe1ab7145f03f276bbd6b29d98e0efa490e05b8e0f9fb5b80 \ + --hash=sha256:35cefbf0a5349f2d772ae41957f5f2a88fe012d374eab801683f644db4edc5a6 \ + --hash=sha256:3c9efe68de0ed648848324510ea7b100e702bdd866f7105e499b209ccde2ceb0 \ + --hash=sha256:4396d9317ee2e9032466bf006b6254d17ab5a71c6d904fa6b2bc2ab377903cb1 \ + --hash=sha256:44a731b65b7b13dff89a1067c0deab75c12bb1d07b95b89e4c18e91a59a6e290 \ + --hash=sha256:472a4d546afaf58bf0e35f388ac330b60b6aaf94713d3b7cf1c6753ff1b4d10c \ + --hash=sha256:4d61e5145c382aeeaa70dfa19b42e9e8049ab3a24bb2064e9f79cda03bb7ff22 \ + --hash=sha256:4fa797232937c7d9140fc19e83267691fa169756e9a376cd9d1520a76919cc1f \ + --hash=sha256:51cede4f54f27f62e303328f790a55e2f8f3e14060605d19fb874bb6318a0742 \ + --hash=sha256:532df36d563611d44e1d715fb68a082e0d7fbc9f1d83237e99e27725940d03c5 \ + --hash=sha256:54eb3f22357c0dc69ef5267c6e5df3e496b3beaf5d1857e50caeeea551528bbc \ + --hash=sha256:55aaff3a5b3ec1fa9fa000b8bb7926e1c129aa0597096d19be9635998eac0a33 \ + --hash=sha256:563d36a666107b9f388e766c832630c7b0eb101ead536d0cd5f6582e75cc334a \ + --hash=sha256:58c7e6441005b3c05e32ea62da034ef427ede4f74d59593a98a8d3d067abbb30 \ + --hash=sha256:5a672013a0159255f37e56ac799b8f75b3c5e5a4ebeb2c37df7fea818ce0adf8 \ + --hash=sha256:63ebd8a3efc876afa6b1a930f15853d3e691e5652441487c26cea5ed49e9b239 \ + --hash=sha256:64a01e79d8d9592cd35f9de2ebc0376e0f94dc8150d6e3ae891a55f190d3490e \ + --hash=sha256:68869460a1f6dc6dfc4de5c8da6101e06f21b1ca7b6d4bbabfd87ec1ba3bc0ca \ + --hash=sha256:6b06146e5b69568dea5506dd870071e2938e260c78db06f40c6184ae411438a3 \ + --hash=sha256:6f76f4eba38cb30d89bad0b3591f4ef3f5f1de3bd526fe33895f1d4187277634 \ + --hash=sha256:70e9892f8902bc48b596f393d4d2a3ff2ca7755783f3137069c85e0517b8d8da \ + --hash=sha256:7956e2efc29e84af71243cd9ece05d580a339fac9525ee531f3c3bdf83c3f073 \ + --hash=sha256:7a9e914b58edad15492eb9dff5f6941cbe8f431d82490f1b4b270a8e51dddbad \ + --hash=sha256:7b1cb1e2dc3cb17a879a7299dcbaf405238256bd9012a35722fa9c1677f6774e \ + --hash=sha256:8f3abdb53c2996d74dfc1a38fc58e55e794ec9557610185519280b6aa9cb90c8 \ + --hash=sha256:8fb485163456835005312499b0210c969b8e680b8ed97387a06e3ac8c573ad81 \ + --hash=sha256:9011b55afcbc2e93d8e2b1dd15314185f9aa6a32e7f7654b1ccf3273bb9643e4 \ + --hash=sha256:972eb2cf749a17f2f489374b8bb69af17b49a62306b35dccaf24767cac3c0194 \ + --hash=sha256:a2cb819f34bb13456c8468125c290dd2bb227a2b784c4e6ed8bca4a07f50400b \ + --hash=sha256:a9b98a129e74593e11b25d4b42b8f5cfcc33db33bf5b8d4dc17c434d0cf1a9bd \ + --hash=sha256:acb96a40da4974d5633c2a1577a26928ff98ff3905db3f0c21e9fc409922c4a3 \ + --hash=sha256:ad6927e917637b679cb9808815800f91c474c3b6c819e7edee898d06d1998615 \ + --hash=sha256:ad9e48f2ef5dc2db69610a5567acd34adbc2195902a0971021490e925023fd52 \ + --hash=sha256:aef69dc30beac397517a7219bdf7873620e8a0de74301e40ad51356c49f724b9 \ + --hash=sha256:af31fc67fe76a1acb1e350b9968a7476a9686a10951c24b3f52bf6173573e722 \ + --hash=sha256:b20e203ef2a4b05271420c038e6d7e004920124d33b0aa1f6c6d966b1f9676fa \ + --hash=sha256:b33c4effbdff51a9e1c32e84b6f3b28a7ea2c7278a482f856540584e7ca8016f \ + --hash=sha256:b3b2b4ae40a8404137ba76139b46210ab73faa3f9429e1512ea444d588661b70 \ + --hash=sha256:b75fe9d8af12dd9235f6d932d8601195fd82b818a0c5780ce28cd068cdad885b \ + --hash=sha256:b8dbd7e78a1d9f9a751ea93984febdb68ded6c88f9b36448f1bb6ec5f6621aed \ + --hash=sha256:bd1a162a40ca090f72766ab01ad7748b16f9b14f6bd36c00fb9b1a288db2a985 \ + --hash=sha256:c327cf199979cd559c713a9cef29a8938bd9d5a17a582b024078b4af1b70fcf2 \ + --hash=sha256:c533ecfd64285c4f228cd61909e35eb9888da4adc2a7ea213b2b96a8084b96be \ + --hash=sha256:ca1f2eb3fac59c01bdabcb767427f32d485ead735d599d12bba48b3137d0ce36 \ + --hash=sha256:d36c1ba71ba155a653f21d67d038fd1a887b343ed65f1660b3911ece5b934d33 \ + --hash=sha256:d5f3e3881058262d29ee05a981e11545cc13b722eaa40ce58b2108a14796bd93 \ + --hash=sha256:d6409b2816e8c578946d3fa5197b2fc92feb4cc6de28adcf5569f521801f927b \ + --hash=sha256:d7b9fc00a98b16b360d95d1e5f34a5a20f4334ab467ca2f172f09d814580c3cf \ + --hash=sha256:dbb0ecc30e0eee5e17a6a83028c0003f6270bcd2488f2927d79b19982cbc198c \ + --hash=sha256:de1f2b7326ca6d3c998c0ea9efc8d92fc716e11a3d02dd977c20be8694ab0cc5 \ + --hash=sha256:ea5039f112334b43f894e74333ec6b6b9ebcf4926db75136ecfd17f938d4a0d8 \ + --hash=sha256:ea9a8dfe02f26b1c68c72f063ec2106edb2d71a5279ea5b1a05e952dc44eb668 \ + --hash=sha256:f2523683758345238d544fbda97b0ecc025d440214c4424bd7bd20c0f8904479 \ + --hash=sha256:f7203e02d33932a6f37622a94ea72f50310e91e2f21d30e8ecc2ae9e67329b53 \ + --hash=sha256:fc5981846390755256eaa9e910fabd6460fac380703aa1030fba943cb3a7fe73 \ + --hash=sha256:fcebc4b59406fe04c18349c496a3d5ce5971f8cc6e30a8fb6ff86c8871f4ab3c \ + --hash=sha256:fe034eadceba325671decc0be79c0825cc25b8f8c4c11f34bcc09a793d790d18 \ + --hash=sha256:ff4d8a4c1fe37af3c4eddd632303f8fbf891ec422c5e117a70024e3757e2edff diff --git a/docker/constraints/biohub-transformers-source.json b/docker/constraints/biohub-transformers-source.json new file mode 100644 index 0000000..299afd7 --- /dev/null +++ b/docker/constraints/biohub-transformers-source.json @@ -0,0 +1,8 @@ +{ + "import_name": "transformers", + "import_root": "src/transformers", + "package_version": "4.57.6", + "schema_version": 1, + "source_revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "tree_sha256": "28b910cc18b821870db2fb6d1c50376c2d14287ae18485080699e03fa4ba4f43" +} diff --git a/docker/constraints/biohub.txt b/docker/constraints/biohub.txt new file mode 100644 index 0000000..69acf9e --- /dev/null +++ b/docker/constraints/biohub.txt @@ -0,0 +1,6 @@ +torch==2.13.0 +transformers==4.57.6 +numpy==1.26.4 +accelerate==1.13.0 +huggingface-hub==0.36.2 +safetensors==0.5.3 diff --git a/docker/constraints/dplm.txt b/docker/constraints/dplm.txt new file mode 100644 index 0000000..ea42fd0 --- /dev/null +++ b/docker/constraints/dplm.txt @@ -0,0 +1,7 @@ +torch==2.2.0 +torchtext==0.17.0 +torch-scatter==2.1.2+pt22cu121 +transformers==4.39.2 +datasets==2.20.0 +biopython==1.79 +numpy==1.26.4 diff --git a/docker/constraints/e1.txt b/docker/constraints/e1.txt new file mode 100644 index 0000000..b2407d9 --- /dev/null +++ b/docker/constraints/e1.txt @@ -0,0 +1,3 @@ +torch==2.8.0 +transformers==4.56.2 +tokenizers==0.22.1 diff --git a/docker/constraints/esmfold.txt b/docker/constraints/esmfold.txt new file mode 100644 index 0000000..ec4c850 --- /dev/null +++ b/docker/constraints/esmfold.txt @@ -0,0 +1,18 @@ +numpy==1.26.4 +biopython==1.84 +dm-tree==0.1.8 +einops==0.8.0 +ml-collections==0.1.1 +omegaconf==2.3.0 +scipy==1.13.1 +ninja==1.11.1.4 +safetensors==0.6.2 +pytest==8.4.2 +huggingface-hub==0.36.0 +wheel==0.46.3 +# The pinned OpenFold package imports these modules eagerly. They belong only +# to the isolated native ESMFold reference environment, never FastPLMs runtime. +pytorch-lightning==1.9.5 +torchmetrics==0.11.4 +lightning-utilities==0.15.2 +dllogger @ git+https://github.com/NVIDIA/dllogger.git@0478734ff7be75adde8d160e04872664d1c62e5f diff --git a/docker/constraints/openfold-sm90.patch b/docker/constraints/openfold-sm90.patch new file mode 100644 index 0000000..86aba5a --- /dev/null +++ b/docker/constraints/openfold-sm90.patch @@ -0,0 +1,29 @@ +diff --git a/setup.py b/setup.py +--- a/setup.py ++++ b/setup.py +@@ -30,3 +30,3 @@ + extra_cuda_flags = [ +- '-std=c++14', ++ '-std=c++17', + '-maxrregcount=50', +@@ -49,16 +49,5 @@ +-compute_capabilities = set([ +- (3, 7), # K80, e.g. +- (5, 2), # Titan X +- (6, 1), # GeForce 1000-series +-]) +- +-compute_capabilities.add((7, 0)) +-_, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME) +-if int(bare_metal_major) >= 11: +- compute_capabilities.add((8, 0)) +- +-compute_capability, _ = get_nvidia_cc() +-if compute_capability is not None: +- compute_capabilities = set([compute_capability]) +- ++# FastPLMs builds this copied reference source for the declared H100 ++# compliance workstation. The pinned submodule itself remains unchanged. ++compute_capabilities = {(9, 0)} ++ + cc_flag = [] diff --git a/docker/docker-bake.hcl b/docker/docker-bake.hcl new file mode 100644 index 0000000..176ddc7 --- /dev/null +++ b/docker/docker-bake.hcl @@ -0,0 +1,179 @@ +variable "REGISTRY" { + default = "local" +} + +variable "TAG" { + default = "dev" +} + +group "default" { + targets = ["runtime", "candidate"] +} + +group "check" { + targets = ["candidate", "candidate-structure", "candidate-fp8", "candidate-artifact"] +} + +group "references" { + targets = [ + "biohub-biotraj-wheel", + "reference-ankh", + "reference-biohub-esm", + "reference-boltz2", + "reference-dplm", + "reference-e1", + "reference-esm2", + "reference-esmfold", + "reference-esmfold2", + "reference-protein-ttt", + ] +} + +target "common" { + context = "." + dockerfile = "docker/Dockerfile" +} + +target "biohub-biotraj-wheel" { + context = "." + dockerfile = "docker/biohub-reference-lock.Dockerfile" + target = "biotraj-wheel-artifact" + tags = ["${REGISTRY}/fastplms-biohub-biotraj-wheel:${TAG}"] +} + +target "runtime" { + inherits = ["common"] + target = "runtime" + tags = ["${REGISTRY}/fastplms-runtime:${TAG}"] + args = { + FASTPLMS_RUNTIME_PROFILE = "core" + } +} + +target "runtime-fp8" { + inherits = ["common"] + target = "runtime" + tags = ["${REGISTRY}/fastplms-runtime-fp8:${TAG}"] + args = { + FASTPLMS_RUNTIME_PROFILE = "esmfold2-fp8" + } +} + +target "candidate" { + inherits = ["common"] + target = "candidate" + tags = ["${REGISTRY}/fastplms-candidate:${TAG}"] +} + +target "candidate-structure" { + inherits = ["common"] + target = "candidate-structure" + tags = ["${REGISTRY}/fastplms-structure:${TAG}"] +} + +target "candidate-fp8" { + inherits = ["common"] + target = "candidate-fp8" + tags = ["${REGISTRY}/fastplms-fp8:${TAG}"] +} + +target "candidate-artifact" { + inherits = ["common"] + target = "candidate-artifact" + tags = ["${REGISTRY}/fastplms-artifact:${TAG}"] +} + +target "reference-ankh" { + inherits = ["common"] + target = "reference-ankh" + tags = ["${REGISTRY}/fastplms-reference-ankh:${TAG}"] + contexts = { + upstream_ankh = "vendor/upstream/ankh" + } +} + +target "reference-biohub-esm" { + inherits = ["common"] + target = "reference-biohub-esm" + tags = ["${REGISTRY}/fastplms-reference-biohub-esm:${TAG}"] + contexts = { + biohub_biotraj_wheel = "target:biohub-biotraj-wheel" + upstream_biohub_esm = "vendor/upstream/biohub-esm" + upstream_biohub_transformers = "vendor/upstream/biohub-transformers" + } +} + +target "reference-boltz2" { + inherits = ["common"] + target = "reference-boltz2" + tags = ["${REGISTRY}/fastplms-reference-boltz2:${TAG}"] + contexts = { + upstream_boltz = "vendor/upstream/boltz" + } +} + +target "reference-boltz2-same-runtime" { + inherits = ["common"] + target = "reference-boltz2-same-runtime" + tags = ["${REGISTRY}/fastplms-reference-boltz2-same-runtime:${TAG}"] + contexts = { + upstream_boltz = "vendor/upstream/boltz" + } +} + +target "reference-dplm" { + inherits = ["common"] + target = "reference-dplm" + tags = ["${REGISTRY}/fastplms-reference-dplm:${TAG}"] + contexts = { + upstream_dplm = "vendor/upstream/dplm" + } +} + +target "reference-e1" { + inherits = ["common"] + target = "reference-e1" + tags = ["${REGISTRY}/fastplms-reference-e1:${TAG}"] + contexts = { + upstream_e1 = "vendor/upstream/e1" + } +} + +target "reference-esm2" { + inherits = ["common"] + target = "reference-esm2" + tags = ["${REGISTRY}/fastplms-reference-esm2:${TAG}"] + contexts = { + upstream_fair_esm = "vendor/upstream/fair-esm" + } +} + +target "reference-esmfold" { + inherits = ["common"] + target = "reference-esmfold" + tags = ["${REGISTRY}/fastplms-reference-esmfold:${TAG}"] + contexts = { + upstream_fair_esm = "vendor/upstream/fair-esm" + upstream_openfold = "vendor/upstream/openfold" + } +} + +target "reference-esmfold2" { + inherits = ["common"] + target = "reference-esmfold2" + tags = ["${REGISTRY}/fastplms-reference-esmfold2:${TAG}"] + contexts = { + biohub_biotraj_wheel = "target:biohub-biotraj-wheel" + upstream_biohub_esm = "vendor/upstream/biohub-esm" + upstream_biohub_transformers = "vendor/upstream/biohub-transformers" + } +} + +target "reference-protein-ttt" { + inherits = ["common"] + target = "reference-protein-ttt" + tags = ["${REGISTRY}/fastplms-reference-protein-ttt:${TAG}"] + contexts = { + upstream_protein_ttt = "vendor/upstream/protein-ttt" + } +} diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..7cf7816 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,53 @@ +# FastPLMs documentation + +This directory explains the public API, model-family contracts, release +evidence, and contributor workflows for FastPLMs 1.0. The model manifest at +`src/fastplms/models.toml` remains authoritative when prose and generated data +disagree. + +## Start here + +| Goal | Read | +| --- | --- | +| Understand repository boundaries and loading flow | [Architecture](architecture.md) | +| Find a supported checkpoint or AutoClass | [Models](models.md) and the [generated support matrix](generated/support.md) | +| Trace every capability to docs, examples, and tests | [Capability-to-evidence manifest](generated/capability_evidence.md) | +| Embed sequences or FASTA datasets | [Embedding API](embedding_api.md) | +| Select SDPA, Flex Attention, or a pinned FlashAttention kernel | [Attention backends](attention_backends.md) | +| Build and validate an offline Hub artifact | [Artifacts](artifacts.md) | +| Run parity, structure, or release tests | [Testing](testing.md) | +| Measure throughput or memory | [Benchmarking](benchmarking.md) | + +## Model and research workflows + +- [ESMFold2](esmfold2.md): folding, learned representations, the distinct full + and Fast MSA contracts, BF16, and experimental FP8. +- [Test-time training](ttt.md): opt-in low-rank adaptation and its evidence + boundary. +- [Binder design](binder_design.md): differentiable ESMFold2 and ESM++ research + example. +- [Fine-tuning](finetuning.md): Trainer, PEFT, data splits, and reproducibility. +- [Vector benchmark embeddings](vector_embeddings/README.md): reusable + embedding artifacts for Protify evaluation. +- [Runnable examples](../examples/README.md): local-only, offline-safe commands + for embeddings, attention, generation, RAG, TTT, and structure preparation. + +## Maintenance + +- [Contributing](contributing.md): adding models, tests, docs, and examples. +- [Licensing](licensing.md): project, source, and checkpoint terms. +- [Migration to 1.0](migration.md): intentional API and repository-layout + changes. + +Generated files carry a marker stating that they come from +`src/fastplms/models.toml`. Edit the typed manifest or renderer, then run: + +```bash +PYTHONPATH=src python -m tools.artifacts.generate_docs +PYTHONPATH=src python -m tools.artifacts.generate_docs --check +``` + +Documentation examples distinguish model output from experimental evidence. +Structure confidence, language-model likelihood, and generated sequences are +prioritization signals unless an independent experiment establishes the +corresponding biological claim. diff --git a/docs/architecture.md b/docs/architecture.md index 757c72f..2ee69e9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,191 +1,143 @@ -# Architecture Overview - -FastPLMs provides optimized, HuggingFace-compatible implementations of protein language models (PLMs) with pluggable attention backends. - -## Repository Layout - -``` -FastPLMs/ - fastplms/ # Main package - ankh/ # ANKH (Elnaggar Lab) - boltz/ # Boltz2 (structure prediction) - dplm/ # DPLM (ByteDance) - dplm2/ # DPLM2 (ByteDance) - e1/ # E1 (Profluent Bio) - esm2/ # ESM2 (Meta AI) - esm3/ # ESM3 (Biohub) - esm_plusplus/ # ESM++ / ESMC (Biohub) - esmfold/ # ESMFold (structure prediction) - esmfold2/ # ESMFold2 / ESMFold2-Fast (structure prediction) - attention.py # Shared attention backend code - embedding_mixin.py # Shared pooling & embedding utilities - weight_parity_utils.py # Weight comparison utilities - fine_tuning_example.py # LoRA fine-tuning example - official/ # Official reference repos (git submodules) - boltz/ # Official Boltz - e1/ # Official E1 - dplm/ # Official DPLM - esm/ # Official Biohub ESM (sys.path-injected, not pip-installed) - entrypoint_setup.py # PyTorch runtime config - cookbook/ # FastPLMs examples and tutorials - tutorials/ - binder_design_fastplms.py - binder_design_fastplms.ipynb - testing/ # Test suite + benchmarks - official/ # Official model loaders for compliance / parity - test_parity.py # Rigorous per-family parity suite - Dockerfile # Monolithic image (legacy) - Dockerfile.base # Shared base image (per-family layout) - Dockerfile. # One per family: esm2, esm_plusplus, esm3, esmfold2, e1, dplm, dplm2, ankh - build_images.sh # Builds base + selected family images - update_HF.py # Pushes composite modeling files + weights to HF Hub - docs/ # Documentation -``` - -Each model family lives in its own package directory containing: - -| File | Purpose | -|------|---------| -| `modeling_*.py` | HuggingFace-compatible `PreTrainedModel` + `PretrainedConfig` subclasses | -| `get_weights.py` | Script to convert official checkpoints to FastPLM format | -| `README.md` | Per-model HuggingFace card README | -| `LICENSE` | Per-model license file | -| `__init__.py` | Package init (often minimal; models load via `trust_remote_code`) | - -## How Model Loading Works - -All FastPLMs models are distributed on the [HuggingFace Hub](https://huggingface.co/Synthyra) and loaded with `trust_remote_code=True`: - -```python -from transformers import AutoModelForMaskedLM - -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-150M", - trust_remote_code=True, -) -``` - -When `trust_remote_code=True` is passed, HuggingFace downloads the `modeling_*.py` file from the Hub repo and executes it locally. The Hub copy is kept in sync with the canonical copy in this repository via `update_HF.py`. - -The model's `config.json` on the Hub contains an `auto_map` entry that tells `AutoModel` which class to instantiate: - -```json -{ - "auto_map": { - "AutoConfig": "modeling_fastesm.FastEsmConfig", - "AutoModelForMaskedLM": "modeling_fastesm.FastEsmForMaskedLM" - } -} +# Architecture + +FastPLMs separates runtime code, release evidence, and official reference code. +The separation is enforced by tests and container build contexts, not only by +convention. + +## Repository boundaries + +```text +src/fastplms/ runtime source copied into Hugging Face artifacts +tests/ unit, integration, parity, structure, and release tests +benchmarks/ standalone, exact-device Hopper/SM90 performance harness +docker/ candidate, runtime, and reference container definitions +examples/ runnable training and protein-design workflows +tools/ artifact, conversion, remote, and debugging commands +vendor/upstream/ pinned official Git submodules +LICENSES/ distributable third-party legal texts +model_cards/ generated checkpoint cards ``` -## EmbeddingMixin - -Every sequence model (ESM2, ESM++, E1, DPLM, DPLM2) inherits from `EmbeddingMixin` (`fastplms/embedding_mixin.py`), which provides: - -- `embed_dataset()`: Batch embedding pipeline with pooling, SQLite/pth storage, FASTA parsing, and deduplication -- `_embed()`: Abstract method implemented by each model to return last hidden states -- `load_embeddings_from_pth()` / `load_embeddings_from_db()`: Load previously saved embeddings - -The mixin supports two modes: - -1. **Tokenizer mode** (ESM2, ESM++, DPLM, DPLM2): The caller provides a tokenizer; `_embed(input_ids, attention_mask)` is called -2. **Sequence mode** (E1): The caller passes `tokenizer=None`; `_embed(sequences, return_attention_mask=True)` is called, which returns `(embeddings, mask)` - -See [Embedding & Pooling API](embedding_api.md) for full details. - -## Cross-Model Workflows - -FastPLMs also includes cookbook workflows that combine multiple model families. -The binder design tutorial is the main example: - -- ESMFold2 experimental models provide differentiable structure losses through - `res_type_soft`. -- ESM++ provides the masked-LM pseudoperplexity regularizer. -- ESMFold2 hero critics provide final pTM, iPTM, pLDDT, PDB/CIF structures, and - selection metrics. - -The workflow supports local CUDA Docker runs and Modal execution through -[`cookbook/tutorials/binder_design_fastplms.py`](../cookbook/tutorials/binder_design_fastplms.py). -See [Binder Design Example](binder_design.md) for the EGFR 128 amino acid -minibinder result, CLI, artifacts, and scoring details. - -## Attention Backend System - -Most models share a common attention backend abstraction controlled by `config.attn_backend`. Four backends are available: - -| Backend | Key | Numerics | Speed | -|---------|-----|----------|-------| -| PyTorch SDPA | `"sdpa"` | Exact | Fast | -| Flash Attention | `"kernels_flash"` | Approximate | Fastest | -| Flex Attention | `"flex"` | Near-exact | Very fast | -| Auto | `"auto"` | Varies | Best available | - -Each model's attention layer stores an `AttentionBackend` enum and dispatches accordingly. See [Attention Backends](attention_backends.md) for implementation details. - -**Backend setting is uniform across families:** - -- At load time, every family accepts `config.attn_backend = "..."` before `from_pretrained`. -- At runtime, every family exposes a mutable `model.attn_backend` property whose setter propagates to every attention submodule. Use this to benchmark backends on the same weights without reloading. -- Exception: ANKH silently resolves `kernels_flash` to `flex` (or `sdpa`), because T5 relative position bias is incompatible with the flash kernels. -- Exception: ESM3 supports `sdpa` and `flex` in the FastPLMs wrapper. - -## Entrypoint Setup - -`entrypoint_setup.py` configures PyTorch runtime defaults for optimal GPU performance: - -- TensorFloat32 matmul precision (`torch.set_float32_matmul_precision('high')`) -- TF32 enabled for matmul and cuDNN -- cuDNN autotuner (`benchmark=True`) -- Deterministic mode off for speed -- Inductor max autotune GEMM backends (ATEN, CUTLASS, FBGEMM) -- Dynamo scalar output capture and recompile limit - -This module is imported at the top of standalone scripts (`throughput.py`, `compliance.py`) but is not imported by the model files themselves. - -## Docker Layout - -There are two coexisting layouts. - -### Per-family layout (recommended) - -A shared base image plus one image per model family. This isolates conflicting native deps (notably Biohub `esm` vs `fair-esm`, and DPLM's torchtext pin) so each family can be tested against its own native reference without breaking the others. - -- `Dockerfile.base` produces `fastplms-base`: CUDA 12.8, Python 3.12, PyTorch 2.11.0, transformers, FastPLMs source at `/app`. No native reference packages. -- `Dockerfile.` (esm2, esm_plusplus, esm3, esmfold2, e1, dplm, dplm2, ankh) layers on top of `fastplms-base` and installs only that family's native reference deps. -- `build_images.sh` is a convenience script that builds the base then any subset of family images. - -`testing/official/.py` provides the `load_official_model(...)` wrapper that the parity tests call. For ESM++ and ESM3, the Biohub `esm` package itself is **not** pip-installed (it depends on a Biohub `transformers` fork); instead `testing/official/__init__.py` injects the in-tree `official/esm` submodule onto `sys.path` at import time. - -### Monolithic layout (legacy) - -The original top-level `Dockerfile` (image tag `fastplms`) bundles every dep that can coexist into a single image. Used by the broad test suites and throughput benchmarks where per-family isolation isn't needed. - -### Common environment - -Both layouts use: - -- **Base image**: `nvidia/cuda:12.8.0-cudnn-runtime-ubuntu24.04` with Python 3.12 -- **Source code**: Copied to `/app` (`PYTHONPATH=/app`) -- **Runtime workdir**: `/workspace` for outputs, caches, and volume mounts -- **Caches**: `HF_HOME=/workspace/.cache/huggingface`, `TORCH_HOME=/workspace/.cache/torch` - -## Weight Conversion - -Each model family has a `get_weights.py` script that: - -1. Loads the official checkpoint (from HuggingFace or a local file) -2. Remaps parameter names and shapes to match the FastPLM architecture -3. Exports `config.json`, `pytorch_model.bin`, and the modeling source files -4. The exported directory can be pushed to HuggingFace via `update_HF.py` - -## Parity & Compliance Testing - -Each family has a corresponding module in `testing/official/` (e.g., `testing/official/esm2.py`) that wraps the original model in a standardized interface returning `(model, tokenizer)`. The parity / compliance suites load both implementations side-by-side and compare: - -- **Tokenizer parity** (`test_parity.py`): vocab, every token id, every special token id -- **Weight parity**: bit-exact equality of every parameter, with family-specific allowlisted extras (e.g. ANKH's `lm_head.weight`, since native is a T5 encoder without a head) -- **Forward parity (fp32 + bf16)**: per-layer relative-std AND relative-maxabs hidden-state diff (two complementary metrics so a localized regression can't hide inside a collapsed scalar), `last_hidden_state` absolute + relative maxabs, logits MSE, padding-isolation across SDPA and Flex, end-to-end `embed_dataset` pipeline for every family -- **Backend consistency**: every family's supported backends (typically SDPA vs Flex vs `kernels_flash`) agree on the fast side to per-backend tolerance; ANKH compares SDPA vs Flex only because kernels_flash silently falls back -- **Backend setter propagation**: `model.attn_backend = X` actually updates every attention submodule. If this regresses, every backend-parametrized test becomes a no-op. - -The parity suite runs per family in its own Docker image (per the per-family layout above). See [Testing & Benchmarking](testing.md) for full details. +Production modules live only under `src/fastplms`. Their direct Python +dependencies are declared under `requirements/`, but they may not import code +from `vendor`, alter `sys.path` to reach an official checkout, or download code +at import time. Importing runtime source must not create a tokenizer, initialize +a model, compile a kernel, log, change global Torch settings, or access the +network. + +Official repositories live under `vendor/upstream` as real Git submodules. +Reference adapters call their public APIs and normalize outputs for comparison. +An adapter may not import FastPLMs, patch an upstream class, use a FastPLMs +loader, or reconstruct an official forward pass. + +## Manifest-driven release data + +`src/fastplms/models.toml` is the sole source of truth for supported models. Its +typed loader in `fastplms.registry` validates: + +- immutable checkpoint and upstream revisions; +- file identities and explicitly unresolved release blockers; +- AutoClass mappings and tokenizer modes; +- state transformations and conversion records; +- attention, dtype, precision, dependency, VRAM, and test contracts; +- code and checkpoint licenses; +- reference containers and documentation state. + +Tests derive model cases from this registry. Container validation checks that +every declared reference target exists. Documentation support tables and model +cards are generated from it. The artifact builder selects the declared source, +verifies every pinned input, applies the named transformation, and records the +same source record in the output. + +Adding a model only in Python code is therefore insufficient. A release-visible +model must be represented completely in the manifest and pass all generated +consistency checks. + +## Loading and artifact flow + +The same model contract is used from source conversion through downstream +inference: + +1. the manifest selects an immutable checkpoint and pinned official source; +2. the artifact builder verifies checkpoint, tokenizer, source, and legal file + identities; +3. the named state transformation produces canonical FastPLMs weights; +4. the builder writes a self-contained runtime bundle and generated model card; +5. Transformers loads the artifact through an advertised AutoClass with + `trust_remote_code=True`; +6. model-specific preparation identifies biological residues or structure + entities before shared APIs transform the output; +7. parity and artifact suites compare the same declared behavior against the + isolated official reference. + +This flow keeps user loading simple without making the official checkout a +runtime dependency. A local artifact under `dist/hub/` and a published +copy use the same Transformers interface. + +## Runtime source + +`fastplms.attention` owns backend names, mask construction, and explicit +dispatch. Models use Transformers' `attn_implementation` and +`set_attn_implementation()` contract. Mask builders produce the 4D masks used +by eager and SDPA, the packed 2D token masks used by declared precompiled +Hugging Face kernels, and Flex `BlockMask` objects. Flex functions and masks are +cached only after explicit use, keyed by device, dtype, execution shape, and +mask semantics rather than the exact row-length tuple. FastPLMs exposes bounded +cache cleanup without clearing process-global Torch compiler state. Original +padding masks must have exact `(batch, sequence)` shape before any backend +branch. + +`fastplms.embeddings` owns ordered records, biological-residue masks, pooling, +persistence, and resume. Model-specific adapters only prepare the representation +and residue mask. E1 keeps its tokenizer-free raw-sequence adapter. ESMFold2 +produces its learned width-256 representation through a dedicated mixin. + +`fastplms.models` contains model-family implementations. Parameter names remain +compatible with existing checkpoints where possible. If a schema must change, +`models.toml` names a deterministic converter and the release suite compares the +converted key set, shape, dtype, and values exactly. + +`fastplms.runtime` reports source and runtime capabilities without mutating +global state. Optional dependencies are imported only when their feature is +requested. + +## Checkpoint and artifact boundary + +Hub checkpoint files remain external assets pinned by immutable revision and +hash. `tools/artifacts/build.py` consumes an already downloaded snapshot and +never logs in, downloads, creates a repository, or uploads. It writes a local +artifact under `dist/hub//` with unchanged runtime source modules, +AutoClass metadata, tokenizer assets, legal files, source records, and deterministic +safetensors shards. + +An artifact is valid only if it loads in a fresh offline environment with +FastPLMs absent from `sys.path`, `HF_HUB_OFFLINE=1`, `local_files_only=True`, and +`trust_remote_code=True`. Every advertised AutoClass must load, run, save, +reload, and match the repository-source implementation. + +Runtime bundling uses tracked, clean regular files selected by path, +extension, and size allowlists. It rejects untracked inputs, symlinks, +credentials, unknown binaries, and path escapes. Release records separate weight +and runtime revisions, records source-tree and embedded-bundle digests plus +generator/schema version, and provides distinct complete-artifact and +runtime-only attestations. Publication rehashes validated bytes at preflight. + +## Container boundary + +`docker/Dockerfile` is a digest-pinned multi-stage build. Candidate stages use +the release validation stack. Reference stages install each upstream's native +environment and receive only the corresponding submodule and required legal +files. Runtime stages receive neither submodules nor checkpoint weights. + +`docker/docker-bake.hcl` names build targets. `docker/compose.yaml` centralizes +GPU access, `ipc: host`, caches, source mounts, and output mounts. +`tools/remote/run.py` creates an isolated source archive, sends it to a host +specified at invocation time, runs Docker there, and returns JUnit, JSON, and +benchmark outputs. Hostnames, identities, and secrets are never tracked. + +## Design rule + +The preferred implementation is the shortest clear implementation that meets +the exact behavioral contract. More complex code is retained only when a +repeatable benchmark establishes a useful speed or memory benefit and the strict +compliance suite still passes. diff --git a/docs/artifacts.md b/docs/artifacts.md new file mode 100644 index 0000000..bf361de --- /dev/null +++ b/docs/artifacts.md @@ -0,0 +1,292 @@ +# Local Hub artifacts + +`tools/artifacts/build.py` creates deterministic, offline-loadable Hugging Face +artifacts under `dist/hub//`. It operates only on an already downloaded, +manifest-pinned checkpoint snapshot. It never authenticates, downloads, creates +a Hub repository, uploads, deletes, commits, pushes, or opens a pull request. + +## Dependencies + +Artifact tooling uses Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13. +From a normal checkout, without official parity submodules, install the +artifact profile before building: + +```bash +uv venv +uv pip install \ + -r requirements/profiles/artifact.in \ + -c requirements/constraints/validation.txt +``` + +Building an artifact is offline and does not require a GPU. Live compliance is +a separate workflow and is the only stage that requires the official reference +submodules. + +## Build + +```bash +PYTHONPATH=src python -m tools.artifacts.build \ + esm2_8m \ + /cache/hub/models--Synthyra--ESM2-8M/snapshots/ \ + --tokenizer-dir \ + /cache/hub/models--facebook--esm2_t6_8M_UR50D/snapshots/ \ + --output-root dist/hub +``` + +Tokenizer-mode artifacts built from a FastPLMs checkpoint require +`--tokenizer-dir` pointing to the manifest-pinned official snapshot. The +builder copies and records only the official tokenizer files declared by the +manifest. Artifacts whose selected checkpoint is already the official snapshot +may omit the option. + +Before writing output, the builder validates: + +- the model and every required file identity are resolved in `models.toml`; +- the local snapshot matches the declared immutable revision and file hashes; +- every required upstream submodule is initialized at its pinned commit; +- canonical and distributable license files match their declared hashes; +- the family has a complete mechanism-first conversion record. +- every scoped runtime source is a tracked, clean regular file selected by an + extension, path, and size allowlist; +- no scoped input is an untracked file, symlink, credential-shaped path, + unknown binary, or Windows path escape. + +An unresolved file or hash mismatch stops the build. The builder does not infer +an identity from a similarly named checkpoint. + +## Output + +Each artifact contains: + +```text +config.json +model-00001-of-000NN.safetensors +model.safetensors.index.json +tokenizer assets, when applicable +modeling_fastplms.py +fastplms_bundle.py +fastplms/... +README.md +provenance.json +artifact-manifest.json +LICENSES/... +THIRD_PARTY_NOTICES.md +``` + +Normal runtime source modules are copied unchanged under `fastplms/`. The +builder also writes those exact bytes into a deterministic compressed archive +inside `fastplms_bundle.py`. This flat bundle is required because Transformers' +remote-module loader follows flat relative Python imports; it does not import +the copied package tree directly. + +`modeling_fastplms.py` imports the flat bundle and verifies its SHA-256 identity +and canonical embedded file inventory. It rejects unsafe or repeated paths, +non-regular archive entries, bytecode, encryption, unexpected compression, and +non-canonical modes before extraction. The bridge extracts with exclusive file +creation into a loader-owned private `TemporaryDirectory`. It then +re-hashes the exact extracted inventory. It +rejects symlinks, bytecode, non-file entries, or any missing, added, or changed +file. Imports run with bytecode writing disabled. +Nothing is extracted into or trusted from the Transformers module cache; only +an ordinary writable temporary directory is required. This step performs no +network access or compilation. + +Release validation runs each artifact in a fresh interpreter. Complementary +family bundles from the same FastPLMs release can extend the runtime loaded by +an earlier artifact in one interpreter. Runtime files shared by two bundles +must have identical hashes; an overlapping source conflict fails explicitly +and leaves the first runtime intact. Isolate incompatible releases in separate +Python processes. Production code never imports a submodule checkout. + +The release artifact tier selects the checkpoint source declared by the +manifest. This matters for ANKH, where the artifact uses the official +sequence-to-sequence checkpoint so the official decoder and LM head are +present. The named state transform is deterministic. +Weights are written as explicit safetensors shards no larger than 5 GiB with a +sorted index. Trusted legacy `.bin` input is loaded with `weights_only=True` and +is never copied into the output. + +`provenance.json` separates `weights_revision` from `runtime_revision` and +records source-tree and runtime-bundle SHA-256 digests, generator/schema +version, scope-specific complete and runtime-only attestations, both checkpoint +identities, conversion record, BF16 execution policy, upstream revisions, +runtime assets, legal files, `weights_license_status`, and `redistributable`. +The runtime attestation repeats the two license fields. Every currently +publishable family, including DPLM1 and DPLM2, uses `"resolved"` and `true`. +Synthetic unresolved-license tests retain the fail-closed `"unresolved"` and +`false` path. The model card states the same policy; artifact validation rejects +a disagreement. `artifact-manifest.json` contains a SHA-256 identity for every +artifact file. Upload preflight rehashes selected bytes so a post-validation +mutation cannot cross the time-of-check/time-of-use boundary. + +The generated `config.json` also records packaging-only FastPLMs model ID, +selected checkpoint repository and immutable revision, and a deterministic hash +of the checkpoint file identities. Offline embedding metadata uses these fields +when no Hub commit is available. Semantic configuration parity explicitly +removes them. + +## Validate + +```bash +PYTHONPATH=src python -m tools.artifacts.build \ + esm2_8m /cache/fast-snapshot --tokenizer-dir /cache/official-snapshot \ + --output-root dist/hub --replace +``` + +The command validates the completed content manifest. The release artifact tier +then creates a fresh environment where FastPLMs is absent from `sys.path` and +sets: + +```text +HF_HUB_OFFLINE=1 +local_files_only=True +trust_remote_code=True +``` + +Load the built artifact through the same Transformers API used after +publication: + +```python +from transformers import AutoModel + +artifact_path = "dist/hub/ESM2-8M" +model = AutoModel.from_pretrained( + artifact_path, + local_files_only=True, + trust_remote_code=True, +) +``` + +After publication, replace `artifact_path` with the Hub repository ID and pass +the immutable revision of that published FastPLMs 1.0 artifact. The source +checkpoint revision in `models.toml` identifies the input weights; it must not +be reused as the revision of newly generated remote code. + +It loads every advertised AutoClass, performs inference, saves, reloads, and +compares configuration, state, and output against repository source. Network access, +an undeclared import, a missing legal text, or a missing conversion record fails +the tier. + +## Publish files without weights + +Use the separate publisher to update runtime code, configuration, tokenizer or +processor assets, model cards, licenses, and notices from an existing local +artifact: + +```bash +PYTHONPATH=src python -m tools.artifacts.publish \ + --files-only \ + --artifact-root dist/hub \ + --dry-run \ + esm2_8m +``` + +`--artifact-root` points to the local parent directory that contains each +manifest-built model artifact, such as `dist/hub/ESM2-8M`. It selects the +validated local files to publish. It is not a Hub repository path or a request +to upload the directory wholesale. + +Review the complete file plan, then repeat without `--dry-run`: + +```bash +PYTHONPATH=src python -m tools.artifacts.publish \ + --files-only \ + --artifact-root dist/hub \ + esm2_8m +``` + +Files-only publication rejects every ANKH model because the 1.0 ANKH migration +requires complete weight replacement. Since an implicit all-model selection +includes ANKH, callers must pass one or more explicit non-ANKH model IDs. +Authentication comes from `HF_TOKEN` or the cached Hugging Face login; tokens +are not accepted as command-line arguments. + +Before the first commit, the publisher: + +1. Checks the local artifact and model identities. +2. Verifies every selected non-weight file against `artifact-manifest.json`. +3. Verifies the remote checkpoint weight identities against `models.toml`. +4. Records the current remote commit as `parent_commit`. +5. Preflights every selected repository. + +Each Hub update is one add-only commit made from explicit +`CommitOperationAdd` entries. The command never creates a repository, adds a +delete operation, uploads a weight-shaped path, or changes repository settings. +The parent commit makes the update fail if another process changes the target +branch between preflight and commit. + +`artifact-manifest.json` and `provenance.json` are intentionally withheld. They +describe the complete local artifact, including its canonical weight shard +layout, which may differ from the unchanged remote layout. All safetensors, +PyTorch checkpoint formats, and weight index files are also withheld. + +Files-only publishing does not construct a fresh artifact. Build and validate +the artifact first whenever runtime sources, generated model cards, legal +inventory, configuration, or tokenizer assets have changed. The publisher does +not read or hash local checkpoint shards, so the upload step itself performs no +large weight I/O. + +The completed command prints each new Hub commit. Before declaring those +commits as a new FastPLMs release baseline, update the corresponding +`fast_revision` values and the digests of any changed manifest-pinned +non-weight `fast_files`, such as `config.json` or tokenizer assets. + +## Publish a complete checkpoint atomically + +The Synthyra ANKH repositories now contain the complete 1.0 encoder-decoder +state. Use complete mode for a future checkpoint replacement that changes +weights. It publishes the full state together with runtime code, configuration, +tokenizer, card, legal files, source records, and scoped attestations: + +```bash +PYTHONPATH=src python -m tools.artifacts.publish \ + --complete \ + --artifact-root dist/hub \ + --dry-run \ + ankh_base +``` + +`--complete` requires explicit model IDs and never accepts implicit selection or +`--all`. After reviewing every path, remove `--dry-run`. The publisher makes one +parent-guarded atomic commit per selected repository, preserving atomicity and +failing if remote head changes after preflight. Complete mode may delete only +an obsolete path that is pinned in the current registry `spec.fast.files`, is +absent from the validated new inventory, and still matches the preflight remote +digest and parent. This permits a monolithic ANKH weight file to be replaced by +indexed shards without granting general remote deletion authority. It +validates that `AutoModel` can load the encoder/shared view without decoder allocation and that +`AutoModelForSeq2SeqLM` consumes the complete encoder, decoder, cross-attention, +and LM-head state from the same artifact. + +DPLM1 and DPLM2 complete publication is permitted under Apache-2.0 after every +normal preflight passes. At the pinned ByteDance revision, the +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) +is Apache-2.0 and the [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +defines the repository release as including pretrained DPLM1 and DPLM2 weights. +Built artifacts therefore carry Hub license `apache-2.0`, +`weights_license_status="resolved"`, and `redistributable=true`, plus the +verbatim license and `LICENSES/dplm/PROVENANCE.md`. Synthetic unresolved-license +fixtures still prove that complete publication fails before any Hub API call or +mutation when either license field is unresolved. + +## ESMFold2 runtime asset + +The ESMFold2 source record includes `ccd.pkl` from +`biohub/ESMFold2`: 417,306,584 bytes, +SHA-256 `9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5`, +MIT license, and trust kind `hash_pinned_pickle`. Deserialization occurs only +from a private loader-owned temporary snapshot after verifying that snapshot's +size and SHA-256. User-supplied asset and `cache_dir` symlinks are rejected. The +exact manifest repository/revision Hugging Face snapshot link is the sole +exception and must resolve within that repository's contained blob directory. +This prevents path replacement and in-place source mutation across the trust +boundary. Offline execution requires the exact verified cache object and does +not fetch a substitute. + +## Generated cards and support data + +Run `PYTHONPATH=src python -m tools.artifacts.generate_docs` to render model cards and the +support matrix from `models.toml`. Run the same command with `--check` in CI to +reject stale generated files. Generated files state the validation boundary and +do not turn manifest declarations into unverified performance or biological +claims. diff --git a/docs/attention_backends.md b/docs/attention_backends.md index 83a78c4..c0aba35 100644 --- a/docs/attention_backends.md +++ b/docs/attention_backends.md @@ -1,142 +1,282 @@ -# Attention Backends +# Attention backends -All FastPLMs sequence models share a common attention backend system controlled by `config.attn_backend`. This document covers how each backend works, when to use it, and how to configure it. +FastPLMs uses the Transformers attention interface. Callers select a backend at +load time with `attn_implementation` or after loading with +`set_attn_implementation()`. -## Overview +## Dependencies and platform requirements -| Backend | Key | Numerical Equivalence | Speed | Availability | -|---------|-----|----------------------|-------|-------------| -| PyTorch SDPA | `"sdpa"` | Exact | Fast | Any PyTorch >= 2.0 | -| Flash Attention | `"kernels_flash"` | Approximate | Fastest | `pip install kernels` | -| Flex Attention | `"flex"` | Near-exact | Very fast | PyTorch >= 2.11 in FastPLMs Docker images | -| Auto | `"auto"` | Varies | Best available | Always | - -## SDPA (Default) - -PyTorch's `scaled_dot_product_attention` dispatches to a fused CUDA kernel (cuDNN or memory-efficient attention) that is faster and more memory-efficient than naive attention while being mathematically identical. - -**When to use:** Reproducibility, numerical sensitivity, general-purpose inference. - -**Implementation:** Each attention layer calls `F.scaled_dot_product_attention(query, key, value, attn_mask)` with a 4D mask of shape `(batch, 1, 1, seq_len)`. - -**Attention weights:** SDPA does not natively return attention weights. When `output_attentions=True` is requested, all backends (including SDPA) compute attention weights via a separate naive matrix multiplication: `scores = Q @ K^T`, softmax, then `context = scores @ V`. This separate computation negates the memory savings of fused attention, so `output_attentions=True` should only be used for inspection or contact prediction, not during high-throughput inference. - -## Flash Attention (`kernels_flash`) - -Flash Attention 2/3 tiles the attention computation into blocks that fit in SRAM and applies an online softmax algorithm. This avoids materializing the full `(seq_len, seq_len)` attention matrix in HBM, achieving O(n) memory and typically 2-4x faster throughput than SDPA on Ampere (A100) and Hopper (H100) GPUs at long sequence lengths. - -**When to use:** Maximum throughput on A100/H100, long sequences, large batch sizes. - -**Numerical properties:** The online softmax and tiling introduce floating-point rounding differences compared to standard attention. These are typically small but not guaranteed to be inconsequential. They can compound across layers and interact with low-precision dtypes (bf16/fp16). If exact reproducibility matters, use `"sdpa"`. - -**Installation:** FastPLMs uses the HuggingFace `kernels` package for pre-built Flash Attention binaries: +FastPLMs supports Python 3.11 through 3.14. The release validation environment +uses PyTorch 2.13 and Transformers 5.13. Install the core dependencies directly; +Transformers loads FastPLMs runtime source from the Hugging Face model: ```bash -pip install kernels +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" ``` -No C++ compiler or CUDA toolkit version pinning required. The `kernels` package fetches a pre-compiled binary matched to your GPU architecture (SM80 for Ampere, SM90 for Hopper). If no compatible binary exists, the model gracefully falls back to `"flex"` or `"sdpa"`. - -**Implementation details:** - -1. Q, K, V are transposed from `(batch, heads, seq, dim)` to `(batch, seq, heads, dim)` for the kernels layout -2. For variable-length batches, padding tokens are removed via `_unpad_input()` which computes cumulative sequence lengths -3. The kernels flash function is called with the unpadded tensors -4. `pad_input()` reconstructs the full padded layout -5. Flash Attention 3 is tried first (Hopper GPUs), falling back to Flash Attention 2 - -## Flex Attention (`flex`) - -PyTorch's `flex_attention` (PyTorch >= 2.11 in FastPLMs Docker images) generates a fused Triton kernel customized to the mask pattern. The primary advantage is **block masks** that skip padding tokens entirely at the CUDA block level, providing meaningful speedups on variable-length batches. - -**When to use:** Variable-length batches with significant padding, best sustained throughput with `torch.compile`. - -**Numerical properties:** Near-exact to SDPA. Differences are typically within floating-point rounding of naive computation. - -**First-call compilation:** The first forward pass triggers JIT compilation via Triton, which takes 30-120 seconds. All subsequent calls with the same mask shape are fast. When combined with `torch.compile`, this yields the best sustained throughput. - -**Implementation:** - -1. A block mask is created from the 2D attention mask via `create_block_mask(mask_mod, batch, 1, seq_len, seq_len)` -2. The mask mod function returns True for positions that should attend to each other -3. `flex_attention(query, key, value, block_mask=block_mask)` generates and runs the fused kernel -4. E1 uses a block-causal variant where within-sequence attention is bidirectional but cross-sequence attention is causal - -## Auto (`auto`) - -Selects the best available backend in priority order: `kernels_flash` -> `flex` -> `sdpa`. Useful when you want maximum speed without manual configuration. The resolved backend may differ across machines depending on installed packages and GPU architecture. - -## Per-Family Caveats - -- **ANKH** supports only `sdpa` and `flex`. The flash-attention kernels can't accept the additive T5 relative position bias, so requesting `kernels_flash` (or `auto` resolving to it) silently falls back to `flex` (or `sdpa` if flex is unavailable). T5 attention is also unscaled (no `1/sqrt(d_kv)` factor) - the learned position bias absorbs the temperature. -- **ESM3** supports `sdpa` and `flex` in the FastPLMs wrapper. -- **E1** uses a block-causal flex variant: bidirectional within a sequence, causal across sequences in a packed multi-sequence batch. -- **DPLM2** packs amino-acid and structure tokens in the same sequence; the attention mask logic accounts for the multimodal layout but the backend choice is otherwise unchanged. +FlashAttention 2 and 3 additionally require Hugging Face `kernels`, a +compatible Linux CUDA device, BF16 execution, and the manifest-pinned kernel +already in cache for offline use: -## Setting the Backend - -Every FastPLMs sequence model (ESM2, ESM++, ESM3, E1, DPLM, DPLM2, ANKH) supports **both** load-time and post-load backend switching. Pick whichever fits your workflow. - -### At Load Time +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" \ + "kernels>=0.15,<0.16" +``` -Set `config.attn_backend` before calling `from_pretrained`: +The Hub quick start below needs network access for the first model download. +Build a manifest-pinned local artifact and pass `local_files_only=True` for an +air-gapped run. ```python -from transformers import AutoConfig, AutoModelForMaskedLM +from transformers import AutoModel -config = AutoConfig.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -config.attn_backend = "flex" -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-150M", config=config, trust_remote_code=True +model = AutoModel.from_pretrained( + "Synthyra/ESM2-150M", + trust_remote_code=True, + attn_implementation="flex_attention", ) +model.set_attn_implementation("sdpa") ``` -### After Load Time - -Every family's `PreTrainedModel` subclass exposes a mutable `attn_backend` property whose setter propagates the change to every attention submodule in-place: +Published Hub repositories are the normal user-facing model IDs. Pin +`revision` when the exact published model-code snapshot matters. Contributors +can use a manifest-built artifact under `dist/hub/` for local, +offline validation before publishing an update. + +If the caller does not choose a backend, FastPLMs leaves the value unspecified +and Transformers normally selects SDPA. FastPLMs does not implement an `auto` +backend. An unavailable requested implementation raises. The only per-call +substitution is the explicit, warning-emitting eager path required by +`output_attentions=True`; it does not mutate the configured backend. + +## Implementations + +| Name | Transformation | Mask | Main limitation | +| --- | --- | --- | --- | +| `eager` | Explicit score, softmax, and value products | Additive 4D mask | Materializes attention scores | +| `sdpa` | `scaled_dot_product_attention` | Boolean or additive 4D mask | Kernel dispatch is selected by Torch | +| `flex_attention` | Compiled Flex Attention score function | `BlockMask` | First shape and semantics require compilation | +| `flash_attention_2` | Precompiled `kernels-community/flash-attn2` handler at revision `db6b51744f0c` | Packed 2D mask | ESM2 and ESM++ only | +| `flash_attention_3` | Precompiled `kernels-community/flash-attn3` handler at revision `43f0bd269777` | Packed 2D mask | ESM2, ESM++, and DPLM only | + +The manifest lists the subset supported by each family. A requested name that +is not listed for that family raises. A listed optional implementation that +cannot be imported also raises, because dependency absence is a configuration +error rather than evidence that another kernel was tested. + +## Choosing a backend + +| Need | Start with | Why | +| --- | --- | --- | +| Official-parity or general inference | `sdpa` | It is the stable declared path for every sequence family | +| Attention maps or `parti` pooling | `eager` | It materializes the attention graph required by those outputs | +| Variable-length batches with compiled masks | `flex_attention` | Its `BlockMask` can avoid padded attention work | +| Precompiled BF16 CUDA kernel on a declared family | `flash_attention_2` or `flash_attention_3` | The immutable binary and compatible runtime are validated before import | +| Reproducible benchmark comparison | Set the exact backend explicitly | Leaving it unspecified delegates selection to Transformers and Torch | + +Start from the family row in the +[generated support matrix](generated/support.md). Do not request a backend +because another model family exposes it. ESMC-6B, DPLM2, and ANKH have +family-specific numerical boundaries described below. + +`output_attentions=True` requires the full materialized attention-probability +matrix. PyTorch SDPA and Flex Attention do not return that matrix, and the +pinned FlashAttention kernels do not expose it through the FastPLMs contract. +FastPLMs therefore uses eager attention for that forward call and emits a +single `RuntimeWarning` naming the configured backend, effective `eager` +backend, and full-attention-matrix reason. The eager 4-D mask is derived from +the original padding and causal semantics. The configured backend is retained +for later calls. + +## FlashAttention compatibility policy + +The Flash dependency is Hugging Face `kernels`, not the `flash-attn` Python +distribution. The adapters follow the +[Transformers kernel-loading contract](https://huggingface.co/docs/transformers/v5.13.0/kernel_doc/loading_kernels) +and resolve only the snapshot-pinned `kernels-community` repositories recorded +in the manifest. The immutable snapshot revisions are +`db6b51744f0cd7061386442c09df890fc6d9f47e` for FlashAttention 2 and +`43f0bd269777115d94ff826e0d113ce9c1c9087b` for FlashAttention 3. The tracked +`kernels.lock` records the exact hash of every published binary variant. The +loader asks `kernels` to download and hash-validate the compatible variant +before importing it. It never falls back to a branch, compiles source, imports +the `flash_attn` package, or substitutes one FlashAttention version for another. + +After installing `kernels`, `kernels download .` may be used during +image build or cache preparation to fetch both locked binaries. This command +downloads precompiled artifacts only. It is not required when the runtime can +populate its Hugging Face cache on first use. + +An explicit kernel-load failure reports the manifest-pinned repository and +revision together with the underlying cause. The exception is not replaced by +a generic dependency error, and no alternate backend is selected. + +Both pinned FlashAttention kernels are BF16-only. The Q, K, and V tensors must +share one dtype and one CUDA device. CPU tensors and mixed-device inputs raise +before binary download or import. Direct FP32 and FP16 calls raise before +kernel loading. An +FP32-resident model may use an advertised FlashAttention backend +only inside CUDA BF16 autocast, where the operation resolves to BF16 while the +stored parameters remain FP32. Parity, artifact, embedding, and benchmark +paths derive their backend and dtype combinations from this manifest contract; +they do not probe or fall back to an undeclared precision. + +FlashAttention validation must cover dense and mixed-padding forward and +backward passes, including LoRA gradients, on the frozen release environment. +Both offline variables, `HF_HUB_OFFLINE=1` and `TRANSFORMERS_OFFLINE=1`, block +kernel downloads. A required uncached binary therefore raises with its pinned +repository, revision, and original loader error. + +All five ESM++/ESMC backends are selectable. SDPA is the default and is +recommended for highest numerical fidelity. Flex Attention and FlashAttention +3 remain supported and non-experimental even though their BF16 arithmetic is +known to diverge from SDPA. Their accuracy metrics are diagnostics and warnings, +not strict parity release gates. Dispatch integrity, finite values, exact mask +semantics, output shape, and catastrophic biological disagreement remain hard +failures. + +| ESMC backend | Status | Relative L2 | Q99.9 | Residue cosine | Pooled cosine | Top-1 | Jensen-Shannon | +| --- | --- | --- | --- | --- | --- | --- | --- | +| `sdpa` | Recommended fidelity path | Exact-head report required | Exact-head report required | Exact-head report required | Exact-head report required | Exact-head report required | Exact-head report required | +| `eager` | Supported fallback semantics | Pending measured frozen-head GH200/aarch64 set | Pending | Pending | Pending | Pending | Pending | +| `flash_attention_2` | Supported; unavailable on current lock | Structured unavailable record required; prior execution evidence is historical and separate | Not measured | Not measured | Not measured | Not measured | Not measured | +| `flex_attention` | Supported, numerically divergent | Pending measured frozen-head GH200/aarch64 set | Pending | Pending | Pending | Pending | Pending | +| `flash_attention_3` | Supported, non-experimental; unavailable on current lock | Structured unavailable record required | Not measured | Not measured | Not measured | Not measured | Not measured | + +The release candidate must replace eager, SDPA, and Flex pending cells with +distributions measured for each checkpoint, dtype, hardware, and locked +sequence panel. FlashAttention 2 and 3 instead require structured unavailable +records with no numerical fields. A threshold is not a measurement, and a +result from another head or revision is not carried forward. The current release-confirmation target is the exact +GH200/aarch64 workstation and repository container build. H100 and H200 remain +Hopper-class deployment examples, but they are not current release evidence. + +Diagnostic jobs write immutable JSON reports under +`artifacts/diagnostics/esmc/`. Published accuracy bands produce warnings. The +separate corruption/catastrophe guardrails are relative L2 at most `0.25`, +relative Q99.9 at most `0.50`, first-percentile residue cosine at least `0.90`, +pooled cosine at least `0.95`, confident-position top-1 at least `0.80`, and +Jensen-Shannon divergence at most `0.05`. These broad limits catch broken +dispatch, masking, or output semantics; they are not parity or quality claims. + +Default documentation generation remains pending even when that directory or +`FASTPLMS_DIAGNOSTIC_REPORTS` exists. On a frozen release head, explicitly +select the complete 30-record schema-v3 set and then check the generated cards: -```python -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -model.attn_backend = "flex" # every attention layer now uses flex - -model.attn_backend = "kernels_flash" # flip to flash without reloading +```bash +PYTHONPATH=src python -m tools.artifacts.generate_docs \ + --source-root . \ + --esmc-report-root artifacts/diagnostics/esmc + +PYTHONPATH=src python -m tools.artifacts.generate_docs \ + --source-root . \ + --esmc-report-root artifacts/diagnostics/esmc \ + --check ``` -This is useful for benchmarking multiple backends on the same weights, or for falling back at runtime if a backend turns out to be unavailable on the current GPU. The setter validates that the requested backend is installed and raises `AssertionError` otherwise. ANKH's `kernels_flash` request silently falls back to `flex` or `sdpa` because the flash kernels cannot accept ANKH's additive relative position bias. - -## Backend Resolution - -Each model has a `resolve_attention_backend()` function that: - -1. Validates the requested backend string -2. For `"auto"`, probes available backends in order: kernels_flash -> flex -> sdpa -3. Prints the resolved backend once (globally, to avoid log spam) -4. Returns an `AttentionBackend` enum value +Use `--require-esmc-release-evidence` to select +`FASTPLMS_DIAGNOSTIC_REPORTS` or the default report directory without silently +falling back to pending output. Either evidence option fails closed on a +missing, extra, malformed, stale, self-digest-invalid, wrong-device, or +cross-device report, or on a missing/stale dependency lock, installed inventory, +container build/image identity, or official-reference source attestation. The +generated capability manifest and applicable model cards record the exact +candidate and official-source records, context, aggregate metric ranges, and +per-case minimum/median/maximum distributions. + +The current locked GH200/aarch64 release image has no validated FlashAttention +2 kernel. Prior real execution was captured in separate workstation JUnit, but +the immutable report and environment attestation are not bundled in this +repository. It is not copied into the current ESMC release distribution or +used for a numerical claim. The manifest-pinned FlashAttention 3 revision contains x86-64 +variants but no locked PyTorch 2.13, CUDA 13 aarch64 artifact. Older ARM +artifacts target different PyTorch/CUDA combinations and are not substituted. +Both backends remain supported and non-experimental, but current-platform +requests raise before dispatch and their schema-v3 records explicitly attest +that unavailability. + +DPLM advertises eager, SDPA, Flex Attention, and FlashAttention 3. Its pinned +official BF16 contract keeps parameter storage in FP32 and uses CUDA BF16 +autocast. Historical, non-release H100 diagnostics recorded eager and Flex +worst hidden-state relative L2 errors of `0.009212` and `0.006768`, respectively. +Static BF16 parameter storage is not the official DPLM precision path and is +not used to justify backend support, and those values are not current GH200 +release evidence. + +DPLM2 advertises SDPA only. Its pinned BF16 contract also keeps parameters in +FP32 and evaluates them under CUDA BF16 autocast; static BF16 parameter storage +raises before inference. A historical, non-release H100 diagnostic recorded +worst hidden-state relative L2 errors of `0.011772` for eager, `0.011231` for Flex +Attention, `0.013495` for FlashAttention 2, and `0.012656` for FlashAttention 3. +Each exceeds the fixed `0.01` engineering target. Explicit requests for any of +those backends therefore raise, and their dead kernel paths are not retained in +the DPLM2 implementation. These values are not current GH200 release evidence. + +ANKH advertises eager attention and SDPA only. Selection is local to the model +instance and does not mutate process-global CUDA SDPA reduction policy. This +family support applies to the optimized ANKH encoder. The full +sequence-to-sequence checkpoint retains the decoder's declared implementation +boundary. + +## Mask semantics + +`fastplms.attention` centralizes mask conversion. The same biological validity +mask is normalized into: + +- a packed 2D token mask for FlashAttention; +- a 4D mask for eager attention and SDPA; +- a Flex `BlockMask` for padding, causal, block-causal, or declared custom + semantics. + +The original attention mask must have exact shape `(batch, sequence)` before +backend dispatch. FlashAttention calls with a packed 2D padding mask always use the varlen kernel, +including causal self-attention. The causal flag is passed to the varlen kernel, +and padded query rows are restored as exact zeros after repadding. Masked calls +reject shapes or devices that do not match Q, K, and V before loading a kernel. + +E1's block-causal pattern is a distinct semantic key. It is never represented +as ordinary padding attention. Mixed-length and skewed-padding parity cases +exercise every required representation. + +Flex functions and masks are cached only after explicit execution. Compilation +is keyed by execution shape, device, dtype, and attention semantics rather than +the exact row-length tuple, so compatible batches reuse compiled work. Mask +content remains correct for each call. `clear_flex_attention_caches()` provides +bounded cleanup of FastPLMs compiled-function and `BlockMask` caches without +clearing process-global Torch compiler state. Importing FastPLMs does not +compile Flex or modify Dynamo or Inductor settings. + +## Attention outputs and `parti` + +The `parti` embedding pooler constructs an attention graph, so it requires: -The resolved enum is stored on each attention layer as `self.attn_backend` and on the encoder as `self.attention_backend`. - -## Mask Transformations - -`fastplms/attention.py::get_attention_mask()` builds a shared set of padding masks once per forward, and every family consumes the same output: - -| Backend | Mask produced by `get_attention_mask` | Shape | -|---------|---------------------------------------|-------| -| SDPA | Boolean 4D mask (True = valid) | `(batch, 1, 1, seq_len)` | -| Flash | Boolean 2D mask (True = valid) | `(batch, seq_len)` | -| Flex | `BlockMask` via `create_block_mask` | Opaque block mask object | - -Families that need an **additive** (float, `0.0`/`-inf`) mask -- ANKH is currently the only one, because T5 relative position bias is added directly to attention scores -- convert the shared bool mask with the `bool_to_additive_mask(bool_mask, dtype)` helper in `fastplms/attention.py`. Use the helper, don't hand-roll it: - -> Never call `.masked_fill(bool_mask, float("-inf"))` on a bool tensor. `bool(float("-inf"))` is `True`, so the result is a bool tensor and the mask is silently dropped. `bool_to_additive_mask` allocates the float tensor correctly and is the only sanctioned way to produce an additive mask inside the codebase. - -## Interaction with `torch.compile` - -- **SDPA**: Works well with `torch.compile` out of the box -- **Flex**: Best performance when the entire model is compiled; the Triton kernel generation integrates with the compiler -- **Flash**: `torch.compile` wraps the kernels call; dynamic warmup detects when compilation has stabilized - -The throughput benchmark (`testing/throughput.py`) applies `torch.compile` to all backends and uses dynamic warmup stabilization to ensure measurements reflect compiled performance. - -## s_max Tracking +```python +attn_implementation="eager" +``` -When `output_s_max=True` is passed (ESM2, E1), each attention layer computes the per-head maximum attention score bound: `max(||Q|| * ||K||)` per head. This is useful for numerical stability analysis and debugging but adds overhead and should not be enabled during production inference. +It rejects sequences longer than 2,048 biological residues. Other backends do +not materialize the complete attention graph as a side effect. Models that do +not expose meaningful sequence attention, including ESMFold2, reject `parti`. + +## Validation + +Backend validation uses the same valid biological positions as official parity. +It measures relative L2 error, relative 99.9th-percentile error, first-percentile +residue cosine, per-sequence pooled cosine, confident-position top-1 agreement, +and Jensen-Shannon divergence for probability tensors. ESMC SDPA remains the +exact, recommended path; eager validates fallback and mask semantics; Flex is +the measured supported diagnostic backend. A published Flex-band miss warns +and records all six metric distributions, while dispatch, finiteness, +mask/shape integrity, and the separate corruption limits remain hard failures. +FlashAttention 2 and 3 remain supported, non-experimental interfaces, but the +current locked GH200/aarch64 image records them as unavailable and fails closed +before dispatch. Historical FlashAttention 2 execution evidence remains +separate from current release acceptance. + +Performance is measured separately from correctness. See +[benchmarking](benchmarking.md) for compile-time, steady-state, padding, memory, +and regression methodology. diff --git a/docs/benchmarking.md b/docs/benchmarking.md new file mode 100644 index 0000000..e8a0382 --- /dev/null +++ b/docs/benchmarking.md @@ -0,0 +1,250 @@ +# Benchmarking + +Performance measurements run outside pytest on the current validated NVIDIA +GH200 workstation in the exact containerized Linux aarch64 environment. H100 +and H200 remain supported Hopper-class devices, but are not interchangeable +with or accepted as the current release benchmark evidence. +Pytest contains only a short CUDA-event smoke test for the harness. Correctness, +parity, and structure compliance remain separate release gates. + +Every report records the exact Python, PyTorch, CUDA, Transformers, FastPLMs, +Hugging Face `kernels`, Transformer Engine, driver, and accelerator environment. +Performance claims apply only when the current report exactly matches its +baseline's accelerator name, compute capability, total GPU memory, driver, and +software environment. Results are not transferable between GH200, H100, and +H200, or between aarch64 and x86-64 environments. Remote orchestration resolves +Bake to native `linux/arm64` on the GH200 and verifies every loaded image's +architecture and content digest; Docker does not erase ABI differences. + +## Fixed release matrix + +The manifest selects one deep representative for each sequence architecture: +ESM2, ESMC or ESM++, ESM3, E1, DPLM, DPLM2, and ANKH. For every attention +backend declared by that family, the suite measures: + +- latency at `b=1`, `l=512`; +- throughput at `b=8`, `l=1024`; +- padding efficiency for lengths `(1024, 512, 256, 128, 64, 64, 32, 32)`. + +The steady-state operation receives pre-tokenized, preallocated GPU tensors. +Loading, first forward, compilation, steady-state forward, and complete embedding +are separate records. Compilation and first-forward costs are never amortized +into steady-state throughput. + +Within one checkpoint and precision, the suite loads model weights once. It +changes attention only through `set_attn_implementation()`. The first record +contains `load_ms`; later records set `model_reused=true` and leave `load_ms` +empty. This removes repeated checkpoint deserialization without moving +tokenization or host transfer into the measured forward. + +Each case also records the manifest field `bf16_execution`. Families declaring +`static_parameters` load BF16 parameters. Families declaring +`fp32_parameters_autocast` retain FP32 parameters and time model computation +inside CUDA BF16 autocast. This includes structure-family entry points such as +folding and diffusion, not only token-model forwards. Results from the two +storage and execution policies are not mixed under one cache key. The +single-case CLI derives this field for registered checkpoint IDs. An +unregistered local path must provide `--bf16-execution` explicitly, and an +override that conflicts with a registered manifest entry raises. + +ESMFold and Boltz2 are represented by explicit `structure_startup` records in +this harness. Their model loading is measurable, but a generic token forward is +not a folding-throughput contract. End-to-end folds, feature preparation, +sampling, and structure outputs run in the dedicated structure suite and are not +reported as tokens per second. + +Run the fixed matrix with: + +```bash +python -m benchmarks.suite \ + --backends eager sdpa flex_attention \ + --junit-output artifacts/junit/benchmark.xml \ + --output artifacts/benchmarks/h100.json +``` + +The `h100.json` filename is retained as a legacy automation identifier. The +report's hardware fingerprint, not its filename, determines the device on which +it is valid. + +### Pre-publication local artifacts + +Capture a baseline from the final locally built Hub artifacts before publishing +them by building the exact benchmark subset in the same clean, frozen source +checkout and passing its root to the suite: + +```bash +python -m tools.artifacts.build_all \ + --benchmark-suite \ + --source-root . \ + --output-root dist/hub +python -m benchmarks.suite \ + --artifact-root dist/hub \ + --local-files-only \ + --backends eager sdpa flex_attention \ + --junit-output artifacts/junit/benchmark-capture.xml \ + --output artifacts/benchmarks/h100-baseline-candidate.json +``` + +The build step resolves only manifest-pinned snapshots. The benchmark step sets +both Hugging Face offline variables, validates every selected artifact before +CUDA or model work, and loads models and tokenizers from the local directories. +ESMFold2 also requires and validates its separately packaged ESMC-6B backbone. +Missing, linked, corrupt, swapped, unresolved, or stale artifacts fail closed. + +Local paths are load-only implementation details. Case keys remain the +registry-owned Synthyra repository and immutable revision, including for ANKH +and DPLM2 artifacts constructed from official-source weights. Each case and the +top-level report instead record path-free manifest, selected checkpoint, +weights, canonical state, runtime revision, runtime bundle, and source-tree +digests. The regression gate rejects a missing or different artifact inventory. +Published-artifact comparisons should materialize the same immutable snapshot +under an artifact root so they retain that inventory without turning a +workstation path into benchmark identity. + +## FlashAttention source policy + +FastPLMs accepts only immutable, precompiled Flash artifacts loaded by the +Hugging Face `kernels` package. The current GH200/aarch64 release benchmark +explicitly measures eager, SDPA, and Flex only. It never downloads, builds, or +executes FA2/FA3, installs the source `flash-attn` distribution, or substitutes +another implementation. The report records FA2 as prior revision-pinned focused +evidence and FA3 as unavailable in the current linux/arm64 lock. A request to +include either Flash backend fails the remote capability preflight. + +## ESMFold2 representation modes + +ESMFold2 keeps three representation measurements distinct: + +- `projection` receives precomputed BF16 hidden states H with shape + `(b, l, 81, 2560)` and measures only the learned map to Z with shape + `(b, l, 256)`. It is labeled BF16 only. An FP8 label would be misleading + because this operation does not run ESMC. +- `esmc_projection` receives preallocated residue tensors, runs ESMC with all 81 + hidden states, and applies the learned projection. This is the end-to-end + representation path measured separately in BF16 and FP8 across every ESMFold2 + attention backend and each fixed shape. `esmc_reload_ms` records construction + of runtime precision modules from canonical BF16 weights. +- `esmfold2_embed` measures one complete call through the shared embedding API, + including residue encoding, ESMC inference, learned projection, residue-only + pooling, and result construction. BF16 and FP8 are separate records. + +None of these modes runs the folding trunk or diffusion sampler. Full ESMFold2 +folding remains in the structure suite, where geometry and confidence metrics +are meaningful. + +Run a single ESMC-plus-projection case with: + +```bash +python -m benchmarks \ + --model dist/hub/ESMFold2 \ + --auto-class AutoModel \ + --backend sdpa \ + --precision bf16 \ + --mode esmc_projection \ + --batch-size 1 \ + --sequence-length 512 \ + --local-files-only \ + --output artifacts/benchmarks/esmfold2-esmc-projection.json +``` + +Build and validate `dist/hub/ESMFold2` from the manifest-pinned checkpoint +before running this command. For a published artifact, pass its exact +`Synthyra/ESMFold2` revision instead of benchmarking a mutable upstream +identifier. + +Explicit FP8 fails when Transformer Engine or compatible hardware is +unavailable. It never falls back to BF16 under an FP8 label. + +## Descriptive exhaustive sweep + +The exhaustive entry point covers every manifest checkpoint, every declared +sequence backend, batch sizes `(1, 2, 4, 8)`, and lengths +`(128, 256, 512, 1024)`. It also covers both ESMFold2 representation operations. +Structure-only families retain startup records because folding belongs to the +structure suite. + +```bash +python -m benchmarks.suite \ + --exhaustive \ + --output artifacts/benchmarks/h100-exhaustive.json +``` + +This legacy output name does not imply H100 execution. Current release reports +must identify the exact GH200/aarch64 target, and no H100 or H200 report is +GH200-equivalent. + +Exhaustive records use `matrix_kind="exhaustive"`, +`claim_scope="descriptive_only"`, and `claim_eligible=false`. The command rejects +a regression baseline. Its output can diagnose scaling behavior, but it cannot +establish a release gate or speed claim. + +## Timing protocol + +GPU work is timed with CUDA events. Warmup continues until the medians of two +consecutive ten-sample windows differ by less than 2 percent. A case that does +not stabilize fails so clocks, thermals, and competing workloads can be +investigated. + +The measurement phase collects seven blocks. Each block lasts at least 250 ms +and contains at least five forwards. Reports retain every raw event sample, +logical and padded tokens per second, median and P95 latency, peak allocated and +reserved memory, compile time, first-forward time, load time, GPU temperature +and clocks before and after, and the complete environment fingerprint. + +Logical throughput counts valid biological tokens. Padded throughput counts +allocated token positions. Both are reported so padding savings remain visible +without disguising the tensor shape executed by the backend. + +## Regression gate + +Let scalar throughput ratio `r` be: + +```text +r = throughput_current / throughput_base +``` + +The gate compares matched measurement blocks with a deterministic paired +bootstrap interval: + +- fail when the one-sided 95 percent upper confidence bound for `r` is below + `0.95`; +- fail unconditionally when median `r` is below `0.90`; +- fail memory growth above the larger of 5 percent or 256 MiB; +- fail unconditionally when memory growth exceeds 10 percent; +- support a speed claim only when the one-sided lower confidence bound is at + least `1.05`. + +```bash +python -m benchmarks.regression \ + artifacts/benchmarks/current.json \ + benchmarks/baselines/h100.json \ + --output artifacts/benchmarks/gate.json +``` + +The command never updates a baseline. A baseline change is a separate, +reviewable file change supported by raw results and a matching environment. +The `benchmarks/baselines/h100.json` path is retained for compatibility, but the +current release baseline must identify the exact GH200 device and Linux aarch64 +environment that produced it. +The regression gate rejects missing or different machine architecture, GPU +name, compute capability, total memory, NVIDIA driver, Python/Torch/CUDA/cuDNN, +Transformers, optional runtime versions, and artifact identities. A capture +report includes a mechanical promotion contract plus separate `compile_ms`, +`first_forward_ms`, warmup samples, and steady-state blocks; compilation is +never amortized into warm throughput. +The suite writes an initial incomplete JUnit sentinel before model work and +atomically replaces it with the capture or regression result only after the +report is complete. Remote phase timeouts use TERM followed by a bounded +kill-after interval, and the remote run report retains the failing phase if a +timeout or cancellation interrupts the benchmark. +This repository does not yet contain the required release baseline. Create it +only from the frozen exact release head and attach the raw immutable report; +until then, regression and speed claims are blocked. Never fabricate or copy a +baseline from a different runtime, checkpoint revision, or accelerator model. + +## Interpreting results + +Each dense, throughput, and mixed-padding case is evaluated independently. A +backend can improve padded batches while regressing dense batches. A throughput +improvement also does not relax parity: every advertised backend must pass its +correctness contract before its performance result can support a claim. diff --git a/docs/binder_design.md b/docs/binder_design.md index 903c847..d05988a 100644 --- a/docs/binder_design.md +++ b/docs/binder_design.md @@ -1,187 +1,124 @@ -# FastPLMs Binder Design Example +# Binder design example -This guide documents the FastPLMs-only binder design workflow in -[`cookbook/tutorials/binder_design_fastplms.py`](../cookbook/tutorials/binder_design_fastplms.py) -and [`cookbook/tutorials/binder_design_fastplms.ipynb`](../cookbook/tutorials/binder_design_fastplms.ipynb). -It mirrors the Biohub ESM binder design tutorial while using only FastPLMs model -repos and FastPLMs loading paths. +`examples/binder_design_fastplms.py` is a research workflow that +optimizes a soft binder sequence against ESMFold2 structural objectives and an +ESM++ sequence prior. It is a source-level example, not a published model +service or a claim that a designed sequence binds experimentally. -![FastPLMs EGFR minibinder design](assets/egfr_fastplms_binder_design.png) - -The rendered example above is the verified EGFR domain III target, shown in teal, -with a 128 amino acid de novo minibinder, shown in orange. - -## Model Roles - -The optimizer uses three model roles: - -| Role | FastPLMs checkpoints | Used for | -| :--- | :--- | :--- | -| Inversion models | `Synthyra/ESMFold2-Experimental-Fast`, `Synthyra/ESMFold2-Experimental-Fast-Cutoff2025` | Differentiable folding losses during sequence optimization | -| LM regularizer | `Synthyra/ESMplusplus_6B` | ESMC-style pseudoperplexity loss on mutable binder residues | -| Hero critics | `Synthyra/ESMFold2-Experimental-Fast`, `Synthyra/ESMFold2-Experimental-Fast-Cutoff2025`, `Synthyra/ESMFold2-Experimental`, `Synthyra/ESMFold2-Experimental-Cutoff2025` | Final confidence-head pTM, iPTM, pLDDT, structures, and consensus gate | - -Optional scaling critics can be enabled with `--use-scaling-critics`. They follow -the official ranking strategy and contribute `distogram_iptm_proxy` scores, but -they are not part of the confidence-head all-hero-iPTM gate because those -checkpoints are used as distogram proxy critics. - -## Strategy - -The FastPLMs script follows the official binder design workflow: - -1. Build a target plus binder prompt. Fixed residues are held fixed and `#` - residues are optimized. -2. Initialize a differentiable amino acid distribution for each mutable binder - residue. Cysteine logits are set very low and cysteine gradients are masked. -3. Anneal soft amino acid logits toward a discrete sequence over the optimization - trajectory. -4. Fold the target and current binder with the inversion ESMFold2 models and - backpropagate through `res_type_soft`. -5. Optimize three structure losses from the ESMFold2 distogram: binder - intra-contact confidence, target-binder inter-contact confidence, and binder - globularity. -6. Add an ESM++ masked-LM pseudoperplexity loss on mutable binder positions. -7. During the final low-temperature steps, run confidence scoring and keep the - argmax sequence state with the best iPTM. -8. Fold the selected sequence with all hero critics and write structures, - confidence metrics, logits, trajectory, and selection tables. -9. Rank with the official selection rule: minibinders with pI >= 6 are filtered, - hero critics contribute mean iPTM, optional scaling critics contribute mean - distogram iPTM proxy, and the selection score is `0.5 * mean_iPTM + 0.5 * - mean_proxy`. - -The script also reports an extra `all_hero_critics_pass` field for stricter -internal screening. It is true only when the minimum hero-critic iPTM is greater -than the configured consensus threshold, currently `0.9`. - -## Local Docker Run - -Run on a Linux CUDA workstation with the ESMFold2 Docker image: - -```bash -cd /home/ubuntu/FastPLMs - -sudo -n docker run --gpus all --rm \ - -v /home/ubuntu/FastPLMs:/app \ - -v /home/ubuntu/FastPLMs:/workspace \ - -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \ - -w /workspace fastplms-esmfold2 \ - python /app/cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli -``` - -`--binder-sequence` is 128 `#` characters, so the binder is generated from -scratch. Use `--binder-name minibinder` to sample a minibinder length from the -official 60 to 200 amino acid range, or provide a scaffolded prompt with fixed -residues and mutable `#` positions. - -## Modal Run +## Input, transformation, and output -The same script can be deployed to Modal: - -```bash -modal deploy cookbook/tutorials/binder_design_fastplms.py -``` - -Then run the CLI against the deployed app: - -```bash -python cookbook/tutorials/binder_design_fastplms.py \ - --backend modal \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir binder_design_egfr_len128_seed103 -``` +The input is one target protein chain and either a mutable minibinder prompt or +an antibody framework with mutable CDR positions. A fixed random seed creates +initial sequence logits. -Modal jobs return the same result rows and write the same local -`results.parquet` and `selection.parquet` tables after the remote call returns. +Each optimization step: -## Output Files +1. maps binder logits to residue probabilities; +2. constructs the target and binder folding input; +3. evaluates differentiable intra-chain and inter-chain distogram objectives; +4. adds an ESM++ masked-language-model regularizer; +5. updates only mutable binder logits; +6. retains the lowest-loss discrete candidate. -Each run writes: +The two supported Cutoff2025 experimental ESMFold2 variants then act as +critics. Candidates are ranked by mean iPTM across those critics. The workflow +writes sequences, loss trajectories, structures, confidence fields, and a +selection table under the requested output directory. -| File | Contents | -| :--- | :--- | -| `best_sequences.fasta` | Target and selected binder sequence for each batch item | -| `trajectory.jsonl` | Per-step structure, LM, and total losses | -| `results.parquet` | One row per final critic with iPTM, pTM, pLDDT, distogram proxy, PDB text, CIF text, and logits path | -| `selection.parquet` | Official-style post-filtered ranking with `selection_score`, `iptm_score`, `iptm_proxy_score`, pI, and `all_hero_critics_pass` | -| `batch*_*.pdb` and `batch*_*.cif` | Final structures from each critic | -| `batch*_*_logits.pt` | Final binder logits saved for reproducibility and inspection | +Prepared atom tensors are padded to the largest observed atom table in the +batch, rounded upward for kernel alignment. They are never sized from the first +sequence or rounded downward, so dense binder batches cannot truncate atoms. -## Verified EGFR Result +![FastPLMs EGFR minibinder design](assets/egfr_fastplms_binder_design.png) -The following result was generated on the workstation with the local Docker -command above. +## Run -| Field | Value | -| :--- | :--- | -| Target | EGFR domain III crop from `TARGET_SEQUENCES["egfr"]` | -| Binder type | 128 amino acid de novo minibinder | -| Seed | `103` | -| Steps | `150` | -| Batch size | `1` | -| Output directory | `/home/ubuntu/FastPLMs/campaign_egfr_len128_b1_s150_seed103_consensus_cli` | -| Official selection score | `0.456935` | -| Hero mean iPTM | `0.913870` | -| Hero min iPTM | `0.904600` | -| All hero critics above 0.9 | `True` | +Run from a source checkout with the `binder` dependency profile. The published +workflow requires Python 3.11-3.14, +PyTorch 2.13, Transformers 5.13, verified ESMFold2 runtime assets, and CUDA. The +current release evidence target is the exact containerized Linux aarch64 +environment on the NVIDIA GH200 workstation. CPU-only, x86-64, Windows, macOS, +H100, and H200 binder runs do not substitute for that evidence. -Binder sequence: +The script intentionally has no standalone PEP 723 dependency block. +`requirements/profiles/binder.in` composes the core, structure, and bounded +binder-design dependencies. Its binder feature pins AbNumber 0.4.4 and ANARCII +2.0.8, plus pandas and PyArrow: -```text -SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ +```bash +uv pip install \ + -r requirements/profiles/binder.in \ + -c requirements/constraints/validation.txt +PYTHONPATH=src python examples/binder_design_fastplms.py \ + --target-name pd-l1 \ + --binder-name minibinder \ + --batch-size 4 \ + --steps 150 \ + --output-dir artifacts/binder-design ``` -Per-critic metrics: - -| Critic | iPTM | pTM | Mean pLDDT | Distogram iPTM proxy | -| :--- | ---: | ---: | ---: | ---: | -| `ESMFold2-Experimental-Fast` | `0.910996` | `0.940850` | `0.919280` | `0.869852` | -| `ESMFold2-Experimental-Fast-Cutoff2025` | `0.906549` | `0.933330` | `0.903867` | `0.851969` | -| `ESMFold2-Experimental` | `0.904600` | `0.935770` | `0.910045` | `0.827132` | -| `ESMFold2-Experimental-Cutoff2025` | `0.933336` | `0.953066` | `0.933806` | `0.888729` | - -Nearby cheaper step counts were tested with the same seed and 128-residue prompt: - -| Steps | Hero min iPTM | Hero mean iPTM | All hero critics above 0.9 | Notes | -| ---: | ---: | ---: | :---: | :--- | -| `140` | `0.876993` | `0.897606` | `False` | Passed pI filter, failed consensus | -| `145` | `0.863137` | `0.895119` | `False` | Filtered by pI | -| `148` | `0.882141` | `0.908492` | `False` | Passed pI filter, failed consensus | -| `149` | `0.894900` | `0.903115` | `False` | Filtered by pI | -| `150` | `0.904600` | `0.913870` | `True` | Cheapest verified passing run in this bracket | - -This is an ESMFold2-critic result, not experimental validation. Other structure -predictors or docking pipelines can disagree, so high FastPLMs ESMFold2 iPTM -should be treated as a screening signal that still needs orthogonal validation. - -## Test Commands - -The focused binder tests run in the ESMFold2 Docker image: +Pass `--target-sequence` instead of `--target-name` for a custom target. Pass +`--binder-sequence` with `#` at mutable positions instead of a named binder +prompt. + +The output directory must not already exist, including as an empty directory. +The CLI checks this before loading models, and the design call creates the path +exclusively before optimization. A concurrent, interrupted, or stale run is +therefore rejected instead of having its files mixed with a new campaign. +`run_manifest.json` is written atomically and last; if it is absent, treat the +directory as an incomplete run and preserve or move it for diagnosis before +choosing a new output path. + +The default inversion, critic, and ESM++ repositories are loaded at the +immutable FastPLMs commits declared in `src/fastplms/models.toml`; the example +never follows a mutable Hub branch. For a fully cached, network-free run, add +`--local-files-only`. That option passes `local_files_only=True` to every +top-level model load and sets both `HF_HUB_OFFLINE=1` and +`TRANSFORMERS_OFFLINE=1` before loading nested runtime assets. Missing cached +files fail the run instead of downloading them. + +Custom repositories require an explicit immutable commit for every model +(replace the example 40-character values below): ```bash -docker run --rm -v /home/ubuntu/FastPLMs:/app -w /app fastplms-esmfold2 \ - python -m pytest /app/testing/test_binder_design_fastplms.py -m "not gpu" -v - -docker run --gpus all --rm -v /home/ubuntu/FastPLMs:/app -w /app fastplms-esmfold2 \ - python -m pytest /app/testing/test_binder_design_fastplms.py \ - -k tiny_design_dry_run_writes_outputs -v +PYTHONPATH=src python examples/binder_design_fastplms.py \ + --inversion-model lab/esmfold2-inversion \ + --critic-model lab/esmfold2-critic \ + --lm-model lab/esmplusplus \ + --model-revision lab/esmfold2-inversion=1111111111111111111111111111111111111111 \ + --model-revision lab/esmfold2-critic=2222222222222222222222222222222222222222 \ + --model-revision lab/esmplusplus=3333333333333333333333333333333333333333 \ + --local-files-only ``` -The verified run used the current focused test set: - -- `11 passed, 2 deselected` for non-GPU binder tests. -- `1 passed, 12 deselected` for the CUDA tiny design dry run. +Repeat `--inversion-model`, `--critic-model`, and `--model-revision` when a +campaign uses multiple checkpoints. + +The example writes `trajectory.jsonl`, `best_sequences.fasta`, +`results.parquet`, `selection.parquet`, and critic-specific structure and +confidence files plus `run_manifest.json`. The example records the complete +command and normalized configuration; exact optimizer; ESMFold2 critic and +ESM++ weight and runtime revisions; tokenizer identity; backend; parameter and +compute dtype; Torch, Transformers, CUDA runtime, Python, and package +environment; all random seeds; and target, prompt, and input-file hashes. Each +model record separates the requested Hub commit, resolved Hub commit, +`fastplms_weights_revision`, and `fastplms_runtime_revision`. The tokenizer +record carries the ESM++ snapshot and runtime identity alongside its vocabulary +hash. +Antibody CDR positions are obtained through AbNumber's public +`Chain.multiple_domains` API with ANARCII-backed Chothia numbering; the workflow +does not depend on AbNumber's private modules. +Retain the full output directory when comparing campaigns. CUDA driver identity +and ranked-output-table hashes are useful promotion evidence, but the current +example does not emit them and this manifest must not be cited as if it did. + +## Validation boundary + +Feature tests use short seeded runs to verify prompt construction, mutable masks, +loss finiteness, gradient scope, critic output schema, deterministic ranking, +and structure serialization. They do not validate affinity, specificity, +developability, expression, immunogenicity, toxicity, or therapeutic utility. + +Candidates require independent structural review, orthogonal computational +checks, synthesis, and experimental binding and functional validation. Confidence +scores are model outputs, not measurements of biochemical activity. diff --git a/docs/contributing.md b/docs/contributing.md index d6c55fe..4c108cd 100644 --- a/docs/contributing.md +++ b/docs/contributing.md @@ -1,182 +1,99 @@ # Contributing -## Getting Started +Changes should preserve scientific behavior, repository boundaries, and a clear +audit trail. Do not commit or upload model weights as part of a source change. -1. Fork the repository and clone your fork -2. Use Python 3.12 (`.python-version` pins 3.12.3 for pyenv/uv-style tools) -3. Pin packaging tools: `python -m pip install --upgrade pip==26.1.1 setuptools==70.2.0` -4. Install cu128 PyTorch: `python -m pip install torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128` -5. Install pinned direct dependencies: `python -m pip install -r requirements.txt` -6. Tests run in Docker only (see below) +## Setup -## Code Style - -- **Python**: PEP 8, type hints in function signatures -- **No unnecessary comments**: Only where logic is not self-evident -- **Hard asserts**: No silent recovery or defensive error handling -- **No `.get()`, `getattr()`, `hasattr()`**: Let missing keys throw `KeyError`/`AttributeError` - -### Import Ordering - -1. Standard library: `import xyz` -2. Third-party: `from xyz import q` -3. Local/repo: `from models.foo import Bar` - -Within each group, `import x` lines come before `from x import y` lines. - -### Import Grouping - -Never import from the same package on separate lines: - -```python -# Bad -from torch import nn -from torch import optim - -# Good -from torch import nn, optim - -# For many names, use parenthesized form -from transformers import ( - AutoConfig, - AutoModel, - AutoModelForMaskedLM, - PreTrainedModel, -) -``` - -## Adding a New Model - -### 1. Create the Package - -``` -fastplms/new_model/ - __init__.py - modeling_new_model.py # PreTrainedModel + PretrainedConfig - get_weights.py # Weight conversion from official checkpoint - README.md # HuggingFace model card README - LICENSE # Model license - -testing/official/new_model.py # Load official model for compliance testing -``` - -### 2. Implement the Model - -Your `modeling_*.py` should: - -- Subclass `PreTrainedModel` and `EmbeddingMixin` -- Define a `PretrainedConfig` subclass with `attn_backend` attribute -- Implement the `AttentionBackend` enum and backend resolution -- Implement `_embed(input_ids, attention_mask)` returning last hidden states -- Register in `config.json` via `auto_map`: - -```json -{ - "auto_map": { - "AutoConfig": "modeling_new_model.NewModelConfig", - "AutoModelForMaskedLM": "modeling_new_model.NewModelForMaskedLM" - } -} -``` - -### 3. Add Weight Conversion - -`get_weights.py` should: -1. Load the official checkpoint -2. Remap parameter names to match your architecture -3. Export `config.json`, `pytorch_model.bin`, and modeling source files -4. The output directory can be pushed to HuggingFace - -### 4. Add Compliance Testing - -`testing/official/new_model.py` should expose: - -```python -def load_official_model(reference_repo_id: str, device: torch.device, dtype: torch.dtype): - # Load and wrap the official model - # Return (wrapped_model, tokenizer) where wrapped_model has .logits and .hidden_states outputs - ... +```bash +uv venv +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu +python -m pytest tests/cpu -m cpu_contract -n auto --dist=loadscope ``` -### 5. Register in Test Configuration - -Add your model to `testing/conftest.py`: +The CPU validation profile contains only the direct dependencies needed by that +lane. `--torch-backend cpu` selects the CPU PyTorch index. CUDA-only +cuEquivariance and FP8 profiles belong in separate environments. + +Official submodules are not required for routine CPU work. Initialize them only +for a live compliance run with `git submodule update --init --recursive`. Run +release GPU verification on the configured Linux aarch64 GH200 host through +`tools/remote/run.py`. H100 and H200 are supported Hopper-class devices, but do +not substitute for the current exact-device release evidence. The runner binds +Bake to native `linux/arm64` and records the GPU UUID; never use emulated images +for CUDA evidence. +Candidate PyTorch containers must use `ipc: host`. + +## Code rules + +- Keep production code under `src/fastplms` and examples under `examples`. +- Do not import `vendor/upstream` from production code. +- Do not download, compile, construct tokenizers, log, or mutate global Torch + settings at import time. +- Use optional imports only inside the feature that requires them. +- Prefer the shortest clear implementation. Retain complexity only with a + measured benefit and strict parity coverage. +- Preserve checkpoint keys and aliases. If that is impossible, add a named + deterministic transform and an exact conversion test. +- Use type annotations for public interfaces and explain non-obvious numerical + choices near the implementation. + +Python identifiers use PEP 8 snake case. In prose and mathematical comments, +scalar quantities and dimensions are lowercase, tensors and matrices use an +uppercase alias, and shapes use parentheses: ```python -# In MODEL_REGISTRY (for fast CI, pick the smallest checkpoint) -"new_model": { - "fast_path": "Synthyra/NewModel-150M", - "official_path": "org/official-model", - "load_official": "testing.official.new_model", - "model_type": "NewModel", - "uses_tokenizer": True, -}, - -# In FULL_MODEL_REGISTRY (all checkpoints with size_category) -"new_model_150m": { - "fast_path": "Synthyra/NewModel-150M", - "official_path": "org/official-model-150m", - "load_official": "testing.official.new_model", - "model_type": "NewModel", - "uses_tokenizer": True, - "size_category": "small", -}, +# H is the hidden-state tensor with shape (b, l, n, d). +hidden_states = model_output.hidden_states ``` -### 6. Add HuggingFace README - -Create `fastplms/new_model/README.md` with the HuggingFace model card content and `fastplms/new_model/LICENSE` with the model license. +Do not write square-bracket shape signatures or uppercase dimension symbols in +shapes. Run the notation checker before review. -### 7. Add a Per-Family Dockerfile +## Adding or changing a model -Create `Dockerfile.` that layers on top of `fastplms-base` and installs your model's native reference deps. Add the family to `build_images.sh` so `./build_images.sh` picks it up. +First freeze the official configuration, tokenizer assets and behavior, state +schema and aliases, representative outputs, source revision, environment, and +licenses. Then: -If your model's native package conflicts with another (e.g. transformers version pin, torchtext pin), prefer either: -- Loading the native package from a `sys.path`-injected submodule (see `testing/official/__init__.py` for the ESM++ pattern), or -- Using a HuggingFace `transformers` reference class instead (DPLM uses `EsmForMaskedLM` for this reason). +1. update `src/fastplms/models.toml` with immutable identities and a complete + conversion record; +2. implement or change runtime source without importing the official checkout; +3. update a public-API reference adapter in + `tests/parity/support/reference_adapters`; +4. add exact configuration, tokenizer, state, alias, FP32, BF16, feature, and + backend cases; +5. build and validate its offline local artifact; +6. regenerate support data and model cards; +7. verify the generated capability-to-evidence row points to a guide, runnable + offline/local example, and every required test tier; +8. run the required remote tiers. -### 8. Add Parity Tolerances +Never create a family-specific tolerance to make a failing comparison pass. Fix +the implementation or remove the unsupported capability from the manifest and +documentation. -Add a `ParityTolerances(...)` entry in `FAMILY_TOLERANCES` at the top of `testing/test_parity.py`. Start with the default, then tighten as you investigate failures. +## Documentation -### 9. Update `update_HF.py` +State the input, transformation, output, validation evidence, and limitation. +Avoid unsupported equivalence, performance, or biological claims. Keep +first-party model cards under `model_cards/`, legal texts under `LICENSES/`, +and runnable scripts directly under `examples/`. Generated cards and support +tables must be changed through `src/fastplms/models.toml` or their renderer. -Add entries for pushing your model's files to the Hub. - -## Running Tests - -All tests must run in Docker. Never run tests natively on Windows (missing Triton, flash-attention, CUDA kernels). Always pass `--ipc=host`. +Execute code snippets, validate internal links, and run: ```bash -# Build base + your family image -./build_images.sh new_model - -# Run your model's parity suite (its own image) -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-new_model \ - python -m pytest /workspace/testing/test_parity.py -k new_model -v - -# Broader smoke tests in the monolithic image -docker build -t fastplms . -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -k new_model -v +PYTHONPATH=src python -m tools.artifacts.generate_docs --check +python -m tools.debug.check_notation +python -m pytest tests/release/test_documentation.py \ + tests/release/test_model_card_licenses.py -v ``` -## Required Passing Tests - -Before submitting a PR for a new model, ensure inside the family's Docker image: - -1. `test_parity.py::test_tokenizer_parity[]` -2. `test_parity.py::test_weight_parity_fp32[]` -3. `test_parity.py::test_forward_parity_fp32[-{single,uniform,skewed}]` (all three padding scenarios) -4. `test_parity.py::test_forward_parity_bf16[-{single,uniform,skewed}]` -5. `test_parity.py::test_padding_does_not_pollute_valid_positions_fp32[]` (tokenizer-mode families) -6. `test_parity.py::test_backend_consistency_fp32[]` - -And in the monolithic image: - -7. `test_automodel_loads` and `test_automodel_forward_pass` -8. `test_nan_stability` -9. `test_batch_single_match` (tokenizer-mode models) - -## Reporting Issues +## Review scope -Found a bug or have a feature request? Open a [GitHub Issue](https://github.com/Synthyra/FastPLMs/issues). +Keep unrelated user changes intact. Do not commit, push, upload, delete a live +Hub repository, or open a pull request unless the maintainer explicitly asks. diff --git a/docs/embedding_api.md b/docs/embedding_api.md index 67edf29..0ee2ed1 100644 --- a/docs/embedding_api.md +++ b/docs/embedding_api.md @@ -1,277 +1,433 @@ -# Embedding & Pooling API +# Embedding API -The `EmbeddingMixin` class (`fastplms/embedding_mixin.py`) provides a standardized interface for extracting protein representations from tokenizer-based FastPLMs sequence models and E1. ESM3 exposes the same `embed_dataset()` user API through its wrapper. +## Dependencies and platform requirements -## Pooler +The shared sequence embedding API requires Python 3.11-3.14, PyTorch 2.13, and +Transformers 5.13. Install those dependencies directly. Transformers loads the +runtime source from the pinned Hugging Face model repository: -The `Pooler` class aggregates per-residue representations `(batch, seq_len, hidden_size)` into fixed-size vectors `(batch, hidden_size)`. - -### Construction - -```python -from fastplms.embedding_mixin import Pooler - -pooler = Pooler(pooling_types=["mean", "max"]) +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" ``` -### Strategies +Core tokenizer-mode embeddings run on CPU or CUDA. E1 uses its raw-sequence +adapter rather than a tokenizer. Structure models and optional FlashAttention +backends require the additional dependencies and CUDA platforms declared in +the support matrix. -| Strategy | Key | Description | -|----------|-----|-------------| -| Mean | `"mean"` | Mask-aware average over all residues | -| Max | `"max"` | Element-wise maximum (masked positions zeroed) | -| CLS | `"cls"` | First token's representation | -| L2 Norm | `"norm"` | L2 norm over the sequence dimension | -| Median | `"median"` | Element-wise median (masked positions zeroed) | -| Variance | `"var"` | Variance over non-masked positions, computed correctly via mean-centered squared diffs | -| Std Dev | `"std"` | Square root of variance pooling | -| PageRank | `"parti"` | Experimental: uses `networkx.pagerank` over attention matrices to weight token importance | +## Quick start -### Calling Convention +Published models expose `model.embed_dataset(...)`. A source checkout also +exposes the same implementation as `fastplms.embed_dataset(model, ...)` when +run with `PYTHONPATH=src`. This minimal Hugging Face example returns one +mean-pooled vector per sequence: ```python -# emb: (batch, seq_len, hidden_size) -# attention_mask: (batch, seq_len) - 1 for real tokens, 0 for padding -# attentions: (batch, num_layers, seq_len, seq_len) - required only for "parti" -pooled = pooler(emb, attention_mask=attention_mask, attentions=attentions) -# pooled: (batch, num_pooling_types * hidden_size) +from transformers import AutoModel + +model = AutoModel.from_pretrained( + "Synthyra/ESM2-150M", + trust_remote_code=True, +).eval() +result = model.embed_dataset( + [ + ("protein-a", "MSTNPKPQRKTKRNT"), + ("protein-b", "MKTIIALSYIFCLVFA"), + ], + batch_size=2, + pooling=("mean",), +) +print(result[0].id, result[0].tensor.shape) ``` -When multiple strategies are specified, their outputs are concatenated along the last dimension. - -### PageRank Pooling (`parti`) - -The `parti` strategy: -1. Max-pools attention matrices across all layers to get `(batch, seq_len, seq_len)` -2. Converts each attention matrix to a directed graph via `networkx` -3. Runs PageRank (alpha=0.85, tol=1e-6, max_iter=100) to get per-token importance scores -4. Computes a weighted average of embeddings using importance scores as weights - -This requires `output_attentions=True` when calling the model. - ---- - -## EmbeddingMixin - -### `embed_dataset()` - -The primary entry point for batch embedding. +The operation accepts sequences, `(id, sequence)` pairs, `EmbeddingInput` +values, an insertion-ordered `{id: sequence}` mapping, or a FASTA path. It +preserves input order, mapping and FASTA identifiers, and duplicate records. + +## Argument reference + +| Argument | Meaning | +| --- | --- | +| `inputs` | Sequence iterable, `(id, sequence)` pairs, `EmbeddingInput` values, `{id: sequence}` mapping, or FASTA path | +| `batch_size` | Number of records prepared together; must be positive | +| `pooling` | One pooler or an ordered pooler sequence; `None` selects mean unless `full_embeddings=True` | +| `full_embeddings` | Return residue-level tensors instead of pooled vectors | +| `output` | Safetensors directory or SQLite file; omit for in-memory results | +| `format` | `safetensors` or `sqlite` when `output` is set | +| `resume` | Reuse an exact compatible ordered prefix when persistent output exists | +| `tokenizer` | Explicit tokenizer override for a compatible tokenizer-mode family | +| `max_length` | Optional maximum number of biological residues, excluding tokenizer-added special tokens | +| `truncate` | Truncate biological residues to `max_length`; when false, an over-length record raises | +| `batch_window_size` | Bounded number of records eligible for stable length bucketing; defaults to `16 * batch_size` | +| `max_tokens_per_batch` | Optional padded biological-residue budget for one inference batch | +| `dtype` | Output tensor dtype; `None` retains the model output dtype | +| `shard_size` | Target safetensors shard size in bytes | +| `model_state_fingerprint` | Caller-supplied state identity for offloaded or externally managed models | +| `**model_kwargs` | Family-specific embedding controls such as hidden-state selection | + +`store_all_hidden_states=True` is a model keyword and requires +`full_embeddings=True`. `full_embeddings=True` cannot be combined with an +explicit pooler. The output format and every invalid argument combination are +validated before input hashing, model inference, or output creation. + +## Bounded streaming and length policy + +FASTA input is read line by line into an immutable, incrementally fingerprinted +spool. The runner never reads the complete FASTA file into memory. It keeps only +one bounded `batch_window_size` group eligible for length bucketing, applies +`max_tokens_per_batch` to the padded biological-residue count, and then restores +exact source order. Omitting the window size resolves it to sixteen times the +batch size; an explicit value always wins. The resolved window is included in +the result metadata and resume fingerprint. Result descriptors are stored once; +tensor payloads remain lazy for persistent outputs. + +SQLite prefixes commit at completed batch-window boundaries. Safetensors packs +windows into bounded shards and publishes a resumable prefix whenever a shard +flushes; an interruption replays the unflushed in-memory shard. Set +`batch_window_size=batch_size` when per-batch inference boundaries matter more +than the default padding-efficiency lookahead. + +`result.metadata["batching"]["resume_commit_granularity"]` records +`"batch-window"` for persistent SQLite, `"shard-flush"` for persistent +safetensors, and `"not-applicable"` for an in-memory result. A new or replacement +SQLite run remains staged or deferred and does not replace the default readable +run until its first batch window commits. + +`max_length` always counts amino-acid residues. Tokenizer-mode families add +their required BOS, EOS, or modality boundary width when constructing the model +token budget. `truncate=False` does not silently exceed the model contract: an +input longer than `max_length` raises with the record position and identifier. + +## Result types + +The following source-level imports are for contributor workflows run with +`PYTHONPATH=src`; Hugging Face users can pass `(id, sequence)` pairs directly +to the model method. ```python -embeddings = model.embed_dataset( - sequences=["MALWMRLLPLLALL", "MKTLLILAVVAAALA"], - batch_size=32, - pooling_types=["mean"], - save=True, - save_path="embeddings.pth", -) -``` +from fastplms import EmbeddingInput -#### Parameters +inputs = [ + EmbeddingInput("a", "MSTNPKPQRKTKRNT"), + EmbeddingInput("a", "MKTIIALSYIFCLVFA"), +] +result = model.embed_dataset(inputs, batch_size=2) -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `sequences` | `List[str]` | `None` | Protein sequences to embed | -| `fasta_path` | `str` | `None` | Path to a FASTA file; sequences are parsed and combined with `sequences` | -| `tokenizer` | `PreTrainedTokenizerBase` | `None` | Tokenizer for tokenizer-mode models. Defaults to `model.tokenizer` when available. Pass `None` for E1 sequence mode | -| `batch_size` | `int` | `2` | Batch size for inference | -| `max_len` | `int` | `512` | Maximum sequence length (longer sequences are truncated if `truncate=True`) | -| `truncate` | `bool` | `True` | Whether to truncate sequences exceeding `max_len` | -| `full_embeddings` | `bool` | `False` | If True, return per-residue embeddings instead of pooled vectors | -| `embed_dtype` | `torch.dtype` | `torch.float32` | Dtype for stored embeddings | -| `pooling_types` | `List[str]` | `["mean"]` | Pooling strategies to apply (concatenated) | -| `num_workers` | `int` | `0` | DataLoader workers (tokenizer mode only) | -| `sql` | `bool` | `False` | Use SQLite storage instead of in-memory dict | -| `sql_db_path` | `str` | `"embeddings.db"` | Path to SQLite database | -| `save` | `bool` | `True` | Save embeddings to `.pth` file | -| `save_path` | `str` | `"embeddings.pth"` | Path to `.pth` output file | -| `hidden_state_index` | `int` | `-1` | Hidden-state tuple index to embed from. `-1` preserves current last-hidden-state behavior | -| `store_all_hidden_states` | `bool` | `False` | Store every hidden state as layer-first token-wise embeddings. Requires `full_embeddings=True` | +for record in result.records: + print(record.id, record.sequence, record.tensor) +``` -At least one of `sequences` or `fasta_path` must be provided. If both are given, the two sources are merged. +`EmbeddingRecord(id, sequence, tensor)` is ordered and retains the original +sequence. `EmbeddingResult(records, metadata)` is sequence-like. Persisted +records may hold a `LazyTensorReference`; call `record.load_tensor()` to load +that tensor. -#### Return Value +`result.as_dict(key="id")` raises when keys repeat. Callers must explicitly +choose a duplicate policy if they want `first` or `last`. This +prevents silent loss of repeated FASTA identifiers. -- **In-memory mode** (`sql=False`): Returns `Dict[str, torch.Tensor]` mapping each sequence to its embedding -- **SQLite mode** (`sql=True`): Returns `None`; embeddings are written to the database +## Biological-residue policy -#### Deduplication & Resumability +Models return a representation `X` and a biological residue mask `M`: -- Sequences are deduplicated before embedding -- Sorted by length (longest first) for efficient padding -- If `save_path` already exists, previously embedded sequences are loaded and only new sequences are processed -- SQLite mode similarly checks which sequences are already in the database +```text +X: (b, l, d) +M: (b, l) +``` -### Two Modes +Pooling includes positions where `M` is true. BOS, EOS, padding, chain +delimiters, and non-protein structure tokens are excluded. E1 derives `M` from +its native raw-sequence preparation because it has no tokenizer. DPLM2 accepts +raw amino-acid sequences through a model adapter that adds its modality-specific +boundaries and invokes the exact tokenizer with `add_special_tokens=False`. +Each persisted run records the token policy and tokenizer metadata. + +## Pooling + +The supported operations are: + +| Name | Transformation | Limitation | +| --- | --- | --- | +| `mean` | Arithmetic mean over valid residues | None | +| `max` | Elementwise maximum over valid residues | None | +| `norm` | Elementwise L2 norm across valid residues | None | +| `median` | Elementwise median over valid residues | More expensive than mean | +| `std` | Elementwise population standard deviation | Requires at least one residue | +| `var` | Elementwise population variance | Requires at least one residue | +| `cls` | Model-defined classification position | Rejected without meaningful CLS semantics | +| `parti` | Attention-graph weighted residue summary | Eager only and at most 2,048 residues | + +Multiple poolers are concatenated in request order. Metadata records the output +slice for each operation. `parti` uses Torch power iteration with damping 0.85, +tolerance `1e-6`, and at most 100 iterations. It requires an explicit +`attn_implementation="eager"` because it materializes the attention graph. -**Tokenizer mode** (ESM2, ESM++, DPLM, DPLM2): ```python -# The model tokenizer is used by default -embeddings = model.embed_dataset( - sequences=sequences, - batch_size=32, +result = model.embed_dataset( + inputs, + batch_size=8, + pooling=("mean", "max", "std"), ) +print(result.metadata["pool_slices"]) ``` -The mixin builds a `DataLoader` with `build_collator(tokenizer)` and calls `_embed(input_ids, attention_mask)`. Pass `tokenizer=...` only when you need a custom tokenizer wrapper. +Choose poolers based on the downstream object. `mean` is a stable sequence +summary, `max` highlights large per-feature responses, and `std` or `var` +captures within-sequence dispersion. Concatenating poolers increases output +width and should be treated as a feature-design decision rather than a free +accuracy improvement. -**Sequence mode** (E1): -```python -# Pass tokenizer=None -embeddings = model.embed_dataset( - sequences=sequences, - tokenizer=None, - batch_size=32, -) -``` +## Full residue embeddings -The mixin iterates over chunks and calls `_embed(sequences, return_attention_mask=True)`, which returns `(embeddings, attention_mask)`. +`full_embeddings=True` returns one ragged residue tensor per input and cannot be +combined with pooling: -**ESM3 wrapper mode**: ```python -embeddings = model.embed_dataset( - sequences=sequences, +result = model.embed_dataset( + inputs, batch_size=4, - pooling_types=["mean", "cls"], + full_embeddings=True, ) ``` -ESM3 supports pooled `mean`, `cls`, and `max` embeddings, plus residue-wise embeddings with `full_embeddings=True`. SQLite streaming is not enabled for ESM3. +Each tensor has shape `(l_i, d)`, where `l_i` is the number of retained +biological residues for record `i`. Padding is never persisted as a residue +embedding. ---- +Passing `store_all_hidden_states=True` requires `full_embeddings=True` and +returns one tensor with shape `(n, l_i, d)` per input, where `n` follows the +model's hidden-state output order. The biological residue mask is applied only +to the token axis. Safetensors and SQLite preserve this rank without flattening +the state axis. -## Hidden-State Selection +ESMFold2 returns the learned projection with shape `(l_i, 256)`. Its dataset +path accepts only single-chain sequences and FASTA records and supports the +residue-statistic poolers. It rejects `cls` and `parti`. -By default, `embed_dataset()` embeds from the same final representation as before. Pass `hidden_state_index` when you want pooling from a specific hidden-state tuple entry: +### ANKH encoder and decoder layers -```python -embeddings = model.embed_dataset( - sequences=sequences, - batch_size=32, - pooling_types=["mean"], - hidden_state_index=12, -) -``` +The Synthyra ANKH repositories contain the complete encoder-decoder +checkpoints. `AutoModel` exposes the encoder view and +`AutoModelForSeq2SeqLM` exposes the full sequence-to-sequence view. -To save token-wise embeddings for every hidden state, use `store_all_hidden_states=True` with `full_embeddings=True`: +ANKH defaults to the encoder final state: ```python -model.embed_dataset( - sequences=sequences, - batch_size=4, +encoder = model.embed_dataset( + inputs, + hidden_state_source="encoder", + hidden_state_index=-1, full_embeddings=True, - store_all_hidden_states=True, - sql=True, - sql_db_path="all_layers.db", ) ``` -Saved all-layer tensors are layer-first per sequence: - -```python -embeddings["MALWMRLL..."].shape == (num_hidden_states, num_real_tokens, hidden_size) -``` +`hidden_state_index` is applied to the selected stack, and +`store_all_hidden_states=True` stores every state from that stack. Decoder +extraction requires the full `AutoModelForSeq2SeqLM` view and exactly one +explicit aligned `decoder_inputs` sequence or `decoder_input_ids` tensor: -You can later pool a selected layer without rerunning the model: +Use raw protein strings such as `MSTNPK`, not space-separated residues. +Decoder sentinels must be adjacent to their residues, as in +`M`. FastPLMs applies this normalization consistently to the +model-owned tokenizer and an explicitly supplied tokenizer object. ```python -pooled = model.load_pooled_embeddings_from_db( - "all_layers.db", - pooling_types=["mean", "cls"], - hidden_state_index=12, -) - -pooled_pth = model.load_pooled_embeddings_from_pth( - "all_layers.pth", - pooling_types=["mean"], +decoder = seq2seq.embed_dataset( + inputs, + hidden_state_source="decoder", + decoder_inputs=["M" for _ in inputs], hidden_state_index=-1, + full_embeddings=True, ) ``` -`pool_embeddings()` provides the same conversion for an in-memory dictionary returned by `load_embeddings_from_pth()` or `load_embeddings_from_db()`. - ---- +There is no implicit shifted-source decoder input. Official ANKH tasks use +task-dependent prompts, sentinels, or generated tokens. A +`decoder_attention_mask` is valid only with `decoder_input_ids`. Decoder pooling +uses the decoder biological mask and excludes start, EOS, padding, sentinel, +and other tokenizer-special positions. Metadata records stack, layer, decoder +input and mask fingerprints, input-position alignment, and mask policy. -## Storage Formats +### E1 MSA-aware embeddings -### `.pth` (PyTorch) +E1 keeps its native raw-sequence and retrieval preparation, but returns the +same ordered, duplicate-preserving `EmbeddingResult` as the shared embedding +API. Record IDs are the zero-based input positions, so repeated query sequences +remain independently addressable as `"0"`, `"1"`, and so on. -A dictionary serialized via `torch.save`: ```python -{ - "MALWMRLLPLLALL": tensor(...), # (hidden_size,), (seq_len, hidden_size), or (num_hidden_states, seq_len, hidden_size) - "MKTLLILAVVAAALA": tensor(...), -} +result = model.embed_dataset_with_msa( + [query, query], + msa_lookup={query: "/data/query.a3m"}, + batch_size=2, + max_len=len(query), + pooling_types=["mean"], + seed=7, + batch_window_size=2, + max_tokens_per_batch=2 * len(query), + output="e1-msa.sqlite", + format="sqlite", + resume=True, +) +assert [record.id for record in result] == ["0", "1"] ``` -Load with: -```python -embeddings = model.load_embeddings_from_pth("embeddings.pth") -``` +`max_len` is measured in biological residues. `matrix_embed=True` selects full +residue output. `output`, `format`, `resume`, `shard_size`, and +`model_state_fingerprint` have the same persistence and compatibility meaning +as ordinary dataset embedding. Local A3M input is offline; homology search and +Hub MSA acquisition are separate, explicit networked workflows. -### SQLite +## Safetensors storage -Schema: -```sql -CREATE TABLE embeddings ( - sequence TEXT PRIMARY KEY, - embedding BLOB NOT NULL -); +With `format="safetensors"`, `output` names an output directory. FastPLMs writes +generation-scoped shards and then transactionally publishes: + +```text +output/ + embeddings-run--00001.safetensors + embeddings-records-run--00001.jsonl + embeddings-index-run--00001.json + index.json + run.json ``` -- `embedding`: Compact tensor blob containing dtype, rank, shape, and raw bytes +The default maximum shard size is 2 GiB. Tensors are packed across inference +batches and written one shard at a time, so the complete tensor dataset is +never materialized in host memory. Each flushed shard publishes an incomplete +ordered prefix that a matching `resume=True` call can continue. An interrupted, +unflushed shard is recomputed. Generation descriptors preserve record position, +identifier, sequence, shape, dtype, tensor hash, and shard key. Loading the +result creates lazy references rather than reading every shard into memory. +`run.json` is the transactional commit marker. It points to one immutable +generation index by filename and SHA-256 digest and is atomically replaced only +after that index, its descriptor shards, and every tensor shard are durable. +`index.json` is a non-authoritative convenience pointer; reopening follows +`run.json` even when the convenience pointer is missing or interrupted. + +Successful overwrites retain earlier immutable generation indexes, descriptors, +and tensor shards. This is required for correctness: an `EmbeddingResult` opened +before the overwrite still resolves its lazy tensors through the earlier paths. +FastPLMs never guesses when those readers have been released. Preview and then +explicitly collect stale generations only after guaranteeing that no reader or +writer for the output remains active: -Load with: ```python -# All embeddings -embeddings = model.load_embeddings_from_db("embeddings.db") +from fastplms.embeddings import garbage_collect_safetensors_generations -# Specific sequences -embeddings = model.load_embeddings_from_db("embeddings.db", sequences=["MALWMRLLPLLALL"]) +stale = garbage_collect_safetensors_generations("output") # dry run +garbage_collect_safetensors_generations( + "output", + dry_run=False, + confirm_no_active_readers_or_writers=True, +) ``` -SQLite mode writes asynchronously and commits when the writer queue drains. - ---- +Destructive collection invalidates any older `EmbeddingResult`, +`EmbeddingRecord`, or `LazyTensorReference` that still names a collected shard. +It also removes abandoned generation files from interrupted writers. Never run +it concurrently with embedding, overwrite, resume, or result retrieval. -## FASTA Parsing +## SQLite streaming, retrieval, and resume -The `parse_fasta()` utility reads a FASTA file and returns a list of sequences: +Use `format="sqlite"` when a long run should commit each batch: ```python -from fastplms.embedding_mixin import parse_fasta - -sequences = parse_fasta("proteins.fasta") +result = model.embed_dataset( + inputs, + batch_size=16, + output="embeddings.sqlite", + format="sqlite", + resume=True, +) ``` -Multi-line sequences are concatenated. Header lines (starting with `>`) are discarded. Empty lines are skipped. +Tensor payloads store raw bytes and an explicit dtype, so BF16 is lossless. +Each completed batch window is committed transactionally. Resume is allowed +only when the full run fingerprint matches and existing records form the exact +ordered prefix of the request. + +SQLite keeps runs under their full fingerprint. With `resume=False`, a new or +restarted run becomes the default result as soon as its first batch commits; +other fingerprints remain available through `run_id`. An interrupted overwrite +therefore exposes a resumable incomplete prefix while retaining the previous +complete run. This is batch-transactional behavior, not the full-run atomic +replacement provided by safetensors generations. -You can pass a FASTA file directly to `embed_dataset`: +Reopening uses SQLite read-only mode. Filtered retrieval accepts exactly one +ordered selector and preserves request order and duplicates: ```python -model.embed_dataset( - fasta_path="proteins.fasta", - batch_size=64, - pooling_types=["mean"], - sql=True, - sql_db_path="proteins.db", +from fastplms.embeddings import load_sqlite_result + +selected = load_sqlite_result( + "embeddings.sqlite", + record_ids=["protein-b", "protein-a", "protein-b"], ) +print([record.id for record in selected]) ``` ---- - -## Full Embeddings (Per-Residue) +Selectors are `positions`, `record_ids`, or `sequences`; `run_id` may select a +specific compatible run. A writable connection is never opened by the result +reader. -When `full_embeddings=True`, the pooler is bypassed and per-residue embeddings are returned. Padding tokens are stripped using the attention mask: +Convert an older FastPLMs SQLite database once, then use the current read-only +reader: ```python -embeddings = model.embed_dataset( - sequences=sequences, - batch_size=32, - full_embeddings=True, - save=False, -) -# embeddings["MALWMRLL..."].shape == (num_real_tokens, hidden_size) +from fastplms.embeddings import convert_legacy_sqlite + +convert_legacy_sqlite("legacy.sqlite", "embeddings-v1.sqlite") ``` -Each sequence's embedding has shape `(num_real_tokens, hidden_size)` for one hidden state, or `(num_hidden_states, num_real_tokens, hidden_size)` when `store_all_hidden_states=True`. `num_real_tokens` excludes padding according to the model attention mask. +Compact and weights-only tensor blobs convert without pickle. An unsupported +pickle payload is rejected unless `allow_unsafe_pickle=True` is explicitly set +for a trusted source. + +## Run metadata + +Persisted results include: + +- model ID, immutable revision, checkpoint hash, and package versions; +- Torch and Transformers versions, backend/device policy, checkpoint identity, + and adapter identity; +- tensor dtype and resolved attention backend; +- selected layer or projection; +- tokenizer and biological-residue policy; +- pooling names and output slices; +- truncation settings; +- input and complete-run fingerprints; +- fingerprint schema version and exact model-state fingerprint; +- generation-indexed output tensor shapes and SHA-256 hashes. + +When a model is loaded from `dist/hub/`, Transformers does not assign a +Hub commit to `config._commit_hash`. The artifact therefore carries +packaging-only model ID, checkpoint repository, immutable revision, and +checkpoint-identity hash fields. Embedding metadata and resume fingerprints use +those fields as the fallback, so local offline runs retain complete traceability. +The packaging fields are excluded from semantic configuration parity. + +Run-fingerprint schema v3 binds the exact current bytes, names, dtypes, and +shapes of every model parameter and persistent buffer. State tensors are copied +to CPU in bounded chunks rather than duplicating the complete model. The digest +is recomputed from authoritative bytes for every persisted run; object identity, +autograd version counters, and cached state digests are never trusted. Mutations +through `Parameter.data` or another storage alias therefore change both the +model-state digest and resume identity. Changing any material input, model +state, or setting prevents resume into an incompatible output. Results written +by older fingerprint schemas cannot be resumed. + +Models with meta-device tensors, custom offloading, or an externally managed +state identity may pass the keyword-only `model_state_fingerprint` override. +The caller is responsible for changing this value whenever the effective model +state changes; metadata records whether the identity was computed or supplied +by the caller. + +## Legacy `.pth` files + +FastPLMs never writes pickle-based `.pth` embeddings. A read-only importer is +available for existing files only when the caller explicitly enables unsafe +pickle loading. Treat such files as executable input and use the opt-in only for +trusted data. Convert imported records to safetensors or SQLite immediately. diff --git a/docs/esmfold2.md b/docs/esmfold2.md new file mode 100644 index 0000000..aaf7a08 --- /dev/null +++ b/docs/esmfold2.md @@ -0,0 +1,339 @@ +# ESMFold2 + +FastPLMs supports exactly four Biohub ESMFold2 variants: + +| Official checkpoint | FastPLMs mirror | Folding blocks | MSA conditioning | +| --- | --- | ---: | --- | +| `biohub/ESMFold2` | `Synthyra/ESMFold2` | 48 | Optional; single-sequence and MSA-conditioned inference are supported | +| `biohub/ESMFold2-Fast` | `Synthyra/ESMFold2-Fast` | 24 | None; inference-optimized single-sequence conditioning | +| `biohub/ESMFold2-Experimental-Cutoff2025` | `Synthyra/ESMFold2-Experimental-Cutoff2025` | 48 | Optional; experimental full-checkpoint contract | +| `biohub/ESMFold2-Experimental-Fast-Cutoff2025` | `Synthyra/ESMFold2-Experimental-Fast-Cutoff2025` | 24 | None; experimental Fast single-sequence-conditioning contract | + +The Fast variants are inference-optimized for single-sequence conditioning and +are not MSA-conditioned. They still support the checkpoint's typed +multichain and multimolecule inputs, but every protein chain must use +single-sequence mode with `msa=None`. Use the corresponding full ESMFold2 +variant when any protein input carries an MSA. This distinction follows the +official Biohub architecture description: Appendix A.2.1 reports 24 folding +blocks for Fast versus 48 for full ESMFold2 and describes Fast as operating +without MSA conditioning for single-sequence inference +([Biohub preprint](https://biohub.ai/papers/esm_protein.pdf)). + +Other snapshots are not advertised in code, artifacts, tests, or documentation. +Local artifact building does not modify the Hub. Files-only publication is a +separate, add-only workflow described in [Hub artifacts](artifacts.md). + +## Dependencies and platform requirements + +ESMFold2 requires the structure dependencies, Python 3.11-3.14, PyTorch +2.13, Transformers 5.13, and a CUDA device for its published execution +contract. The current validated release target is the exact containerized Linux +aarch64 environment on the NVIDIA GH200 workstation. CPU-only, x86-64, +Windows, macOS, H100, and H200 structure runs do not substitute for that release +evidence: + +```bash +uv pip install \ + -r requirements/core.in \ + -r requirements/features/structure.in \ + -c requirements/constraints/validation.txt +``` + +These files are dependency declarations, not a FastPLMs installation. The +loading example below obtains runtime source from the pinned Hugging Face model +with `trust_remote_code=True`. + +The structure dependency file retains Accelerate specifically for the documented +`device_map`-based, memory-safe loading of the 6B ESMC backbone. It retains +OmegaConf for the explicit trusted-deserialization boundary in +`Boltz2Model.from_boltz_checkpoint`, where official Lightning checkpoints may +contain OmegaConf objects. Plotting, reporting, table, and antibody-numbering +packages are not structure runtime dependencies and remain in the reporting or +binder dependency files. + +The reference folding path is included in the structure dependencies. The named +`cuequivariance` kernel backend is a separate opt-in because it adds NVIDIA's +CUDA-specific binary runtime: + +```bash +uv pip install \ + -r requirements/core.in \ + -r requirements/features/structure.in \ + -r requirements/features/cueq.in \ + -c requirements/constraints/validation.txt +``` + +The cuEquivariance dependency file pins the version-aligned frontend and CUDA +kernels used by the release contract: `cuequivariance==0.10.0`, +`cuequivariance-torch==0.10.0`, and +`cuequivariance-ops-torch-cu13==0.10.0`. It selects NVIDIA's +CUDA 13 build because FastPLMs validates PyTorch 2.13 on CUDA 13.0. Do not +install the CUDA 12 and CUDA 13 kernel packages into the same environment. +FastPLMs requires both the frontend and the CUDA ops package before accepting +`model.set_kernel_backend("cuequivariance")`; a frontend-only installation is +not treated as backend availability. + +This backend is available only on Linux with an NVIDIA GPU, a compatible +CUDA 13 driver, and CPython 3.11-3.14. NVIDIA publishes both x86-64 and ARM64 +manylinux wheels for those interpreters, so the Linux aarch64 GH200 validation +workstation can resolve this exact package set. H100 and H200 remain supported +Hopper-class execution devices, but only the exact GH200/aarch64 environment is +the current release evidence target. Results must identify the exact device and +architecture, and performance baselines from different accelerator models are +not interchangeable. Windows, macOS, CPU-only hosts, and the FastPLMs CUDA 12 +legacy reference images are not supported execution paths. +The cuEquivariance Python frontend is Apache-2.0, while the CUDA ops wheels are +distributed under the NVIDIA Software License Agreement and are described by +NVIDIA as beta software. Installing the cuEquivariance dependencies means +accepting those separate NVIDIA terms; FastPLMs does not redistribute the +wheels. + +## Loading + +```python +from transformers import AutoModel + +model = AutoModel.from_pretrained( + "Synthyra/ESMFold2-Fast", + trust_remote_code=True, + attn_implementation="flex_attention", + device_map={"": "cuda:0"}, + esmc_precision="auto", +).eval() +``` + +This Fast quick start intentionally demonstrates the no-MSA path. It does not +turn MSA conditioning on implicitly. For MSA-conditioned inference, load +`Synthyra/ESMFold2` or the experimental full Cutoff2025 checkpoint instead. + +The folding model and its ESMC backbone use the same explicitly resolved +attention implementation. Unsupported names raise. The ESMC checkpoint is +loaded directly on the requested CUDA device when a CUDA device is used. For a +declared ESMC-6B Hub identifier, FastPLMs forwards the immutable revision from +`models.toml` to both configuration and weight loading. Other remote ESMC +checkpoints are rejected because they do not satisfy the learned 81-state, +2560-width projection contract. A local checkpoint directory remains a +supported explicit source and has no Hub revision. + +The folding checkpoint itself loads with FP32 parameters. Learned projection, +folding trunk, and diffusion computation run under CUDA BF16 autocast. ESMC has +an independent precision policy: its canonical BF16 weights may remain BF16 or +be used to reconstruct the transient FP8 inference path without changing the +folding checkpoint's FP32 storage. + +## Hash-pinned CCD asset + +Structure preparation requires `ccd.pkl` from the immutable snapshot +of `biohub/ESMFold2`. The manifest records its 417,306,584-byte size, MIT +license, and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5`. + +This file is a pickle and therefore an explicit trusted-deserialization +boundary. FastPLMs rejects user-supplied asset and `cache_dir` symlinks. The +only allowed link is the Hugging Face snapshot entry for the exact manifest +repository and revision, and it must resolve inside that repository's contained +blob directory. The loader creates a private, loader-owned temporary snapshot, +verifies that snapshot's size and SHA-256, and unpickles only the verified +snapshot. This closes path-replacement and in-place source-write races between +validation and deserialization. `HF_HUB_OFFLINE=1` and +`TRANSFORMERS_OFFLINE=1` require the exact cache object to exist; an offline +call never downloads or substitutes an asset. The release record stores the identity +and cache policy for every artifact. + +## Inputs and outputs + +For a single protein, pass the amino-acid sequence to `fold_protein`: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(result.ptm, result.plddt.mean().item()) +``` + +For complexes, use the typed input schema exposed by the loaded model: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(result.ptm, result.plddt.mean().item()) +``` + +The shared schema also supports RNA, modifications, covalent bonds, and +distogram conditioning. Full ESMFold2 checkpoints additionally accept protein +MSAs. Fast and experimental Fast checkpoints reject a non-null +`ProteinInput.msa`; this does not prevent no-MSA multichain or multimolecule +inference. `PocketConditioning` is recognized by the schema, but +the pinned official runtime drops it and hard-codes a zero pocket feature. +FastPLMs rejects a non-null pocket request rather than silently discarding it; +pocket conditioning is not supported in 1.0. No known target structure is +required. Prepared feature tensors include `ref_pos`, but this is component +reference geometry created during featurization, not the target coordinates. +Atomic coordinates and confidence fields are model outputs. +The offline [`structure_preparation.py`](../examples/structure_preparation.py) +example constructs the supported MSA, protein-complex, RNA, DNA, ligand, +modification, covalent-bond, and distogram inputs and executes the pocket +rejection contract. Its MSA path is for the full variants. Fast variants may +use the other typed modalities only when every protein input has `msa=None`. + +## Learned sequence representation + +Biohub ESMC-6B provides the embedding state followed by 80 transformer-layer +states. FastPLMs validates that exact 81-state ordering and width before applying +the folding checkpoint's learned projection: + +```text +H: (b, l, 81, 2560) +``` + +For each state, `base_z_linear` applies layer normalization and a bias-free +linear map to width 256. The softmax of `base_z_combine` gives 81 scalar weights. +The weighted states are combined with the same matrix multiplication order as +Biohub. The low-level operation maps `H` and its residue mask `M` to +`Z = model.project_esmc_hidden_states(H, residue_mask=M)`. + +The result is: + +```text +Z: (b, l, 256) +``` + +This is the learned sequence summary returned before `base_z_mlp` expands it +into pair features. The refactor retains the checkpoint names +`base_z_linear`, `base_z_combine`, and `base_z_mlp`. + +Projection compliance compares identical 81-state inputs. FP32 output must be +exact. The BF16 engineering target for relative L2 error is `5e-4`, with a hard +limit of `1e-3`. + +## Dataset embeddings + +```python +result = model.embed_dataset( + "proteins.fasta", + batch_size=2, + full_embeddings=True, +) +``` + +Each output record contains a residue tensor with shape `(l, 256)`. Pooling uses +only real residues. Supported poolers are the residue statistics `mean`, `max`, +`norm`, `median`, `std`, and `var`. ESMFold2 rejects `cls` because the learned +representation has no classification token semantics and rejects `parti` +because it is not an attention-graph representation. + +Run metadata records the folding checkpoint identity and the resolved ESMC +repository, immutable revision, and manifest file hashes. These fields are part +of the resume fingerprint, so an embedding run cannot resume after either +checkpoint identity changes. + +The dataset path accepts a single protein chain or FASTA records of single +chains. It rejects gaps, chain separators, structured complexes, ligands, MSAs, +and non-protein tokens. Those remain inputs to the folding preparation path, +not the embedding utility. + +## ESMC precision policy + +The accepted precision values are `auto`, `bf16`, `fp32`, and `fp8`. The +manifest marks `fp8` as experimental: + +```python +model.reload_esmc(precision="auto", device="cuda") +status = model.esmc_precision_status +print(status.as_dict()) +``` + +The status contains requested and resolved precision, reason, device, and the +installed Transformer Engine version. Resolution is fail-closed: + +- `auto` always selects BF16, including on FP8-capable GPUs; +- experimental explicit `fp8` raises when the device or Transformer Engine path is + unavailable; +- explicit `bf16` and `fp32` remain supported. + +FP8 must therefore be requested deliberately: + +```python +model.reload_esmc(precision="fp8", device="cuda") +assert model.esmc_precision_status.resolved == "fp8" +``` + +Converting every ESMC linear compounds quantization error across 80 layers. +The experimental path instead converts exactly each layer's attention output +projection, for 80 Transformer Engine linears in total. Their canonical +parameters remain BF16. `Float8CurrentScaling` quantizes the GEMMs during the +inference context, and sequence inputs are padded to a multiple of 16 before +ESMC execution. + +Three fresh BF16-to-FP8 reload cycles on the historical locked H100 environment +produced identical metrics in each cycle: projection relative L2 `0.0375936`, +first-percentile residue cosine `0.999091`, and minimum per-sequence pooled +cosine `0.999754`. These satisfy the engineering targets of `0.04`, `0.995`, +and `0.999`, respectively. The runtime loads canonical BF16 weights and rebuilds +Transformer Engine modules on every startup or reload. Transformer Engine +workspaces and quantized caches are transient and are excluded from folding +checkpoints. + +That historical locked H100 image uses Transformer Engine `2.12.0` with its +CUDA 13 core. These H100 measurements do not satisfy the current GH200/aarch64 +release gate. +The `uv` override excludes the package's default CUDA 12 core, so the FP8 image +contains one Transformer Engine runtime rather than two. Transformer Engine +`2.13` through `2.16` are not advertised on this validation stack because their +precompiled CUDA 13 cores require a newer cuBLAS ABI than the pinned CUDA 13.0 +toolchain provides. + +FastPLMs does not replace CUDA libraries, compile code at import time, or +serialize runtime-quantized state. The FP8 image builds Transformer Engine's +small PyTorch binding once against the locked Torch/CUDA ABI; its CUDA core is +precompiled. Core FastPLMs imports remain independent of Transformer Engine +because the optional runtime is loaded only when the precision policy needs its +capability probe or execution context. + +## Historical FP8 diagnostic + +A non-release H100 diagnostic over multiple real proteins and seeds found exact +official-versus-candidate BF16 parity but model- and sequence-dependent FP8 +folding deviations. FP8 passed the historical hard structure limits in 48 of +60 cases. The panel and its fixtures are not part of the release suite; current +coverage is one explicit FP8 smoke per variant plus three reload cycles on the +standard variant. This evidence motivates the BF16 `auto` policy and precludes +an FP8 numerical-equivalence claim. + +## Gradient-enabled paths + +Test-time training and other gradient-enabled ESMC execution use canonical BF16 +weights. A future FP8 implementation must retain this boundary. + +## Folding compliance + +Structure tests hash prepared features and sampled diffusion noise before +comparing implementations. They require exact discrete features and masks, +finite outputs, valid geometry, and no NaNs. + +For official versus local BF16, engineering targets are C-alpha RMSD at most +0.10 angstrom, lDDT-C-alpha at least 0.995, pLDDT MAE at most 0.001, PAE MAE at +most 0.10 angstrom, and pTM or ipTM error at most 0.002. Corresponding hard +limits are 0.25 angstrom, 0.99, 0.005, 0.50 angstrom, and 0.005. + +For FP8 versus BF16, engineering targets are C-alpha RMSD at most 0.75 angstrom, +lDDT-C-alpha at least 0.97, pLDDT MAE at most 0.01, PAE MAE at most 0.5 +angstrom, pTM or ipTM error at most 0.01, and mean probability Jensen-Shannon +divergence at most 0.002. Hard limits are 1.5 angstrom, 0.95, 0.02, 1.0 +angstrom, 0.02, and 0.005. diff --git a/docs/finetuning.md b/docs/finetuning.md index 4d8dc30..2974796 100644 --- a/docs/finetuning.md +++ b/docs/finetuning.md @@ -1,173 +1,226 @@ -# Fine-Tuning Guide - -FastPLMs models can be fine-tuned for downstream tasks using LoRA (recommended) or full fine-tuning via the HuggingFace `Trainer` API. This guide is based on `fastplms/fine_tuning_example.py`. - -## Quick Start - -```python -from transformers import AutoModelForSequenceClassification -from peft import LoraConfig, get_peft_model - -# Load with a classification head -model = AutoModelForSequenceClassification.from_pretrained( - "Synthyra/ESM2-150M", - trust_remote_code=True, - num_labels=1, # regression -) - -# Apply LoRA -lora_config = LoraConfig( - r=16, - lora_alpha=32, - target_modules="all-linear", - lora_dropout=0.1, - bias="none", - task_type="SEQ_CLS", -) -model = get_peft_model(model, lora_config) -model.print_trainable_parameters() +# Fine-tuning + +FastPLMs models follow Transformers `PreTrainedModel` conventions, so +compatible family and task-head combinations can use Trainer, Accelerate, +distributed, and adapter workflows. Core training dependencies are isolated in +`requirements/features/train.in`. FastPLMs 1.0 supports Python 3.11-3.14 with PyTorch 2.13 and +Transformers 5.13. CPU is suitable for tiny contract runs; practical checkpoint +fine-tuning normally requires a CUDA accelerator whose memory fits the selected +model, optimizer state, and batch. + +```bash +uv pip install \ + -r requirements/core.in \ + -r requirements/features/train.in \ + -c requirements/constraints/validation.txt ``` -## Dataset Classes - -### Single-Sequence Classification/Regression - -`SequenceDatasetHF` wraps a HuggingFace `datasets` split: - -```python -from datasets import load_dataset - -dataset = load_dataset("your/dataset") -train_data = SequenceDatasetHF( - dataset["train"], - col_name="sequence", - label_col="label", - max_length=2048, -) +The runnable classification and regression example currently targets ESM2, +whose artifacts advertise `AutoModelForSequenceClassification`. Plotting is +disabled by default, so a normal run needs only the training dependencies. To +request plots, install the reporting profile and pass `--plot-results`: + +```bash +uv pip install \ + -r requirements/profiles/reporting.in \ + -c requirements/constraints/validation.txt +PYTHONPATH=src python examples/fine_tuning.py \ + --task classification \ + --model_path Synthyra/ESM2-8M \ + --model-revision 185ecbd45665d050a8dae326d91886d330c5f9d0 \ + --classification-dataset-source GleghornLab/DL2_reg \ + --classification-dataset-revision 7e18f1b98859b0a3e3da283f63d0a153b774cf1f \ + --attn-backend sdpa \ + --output-dir artifacts/fine-tuning \ + --batch_size 8 \ + --epochs 2 \ + --seed 7 \ + --full-determinism \ + --plot-results ``` -### Protein-Protein Interaction (Pair) - -`PairDatasetHF` handles two-sequence inputs (e.g., PPI binding affinity): - -```python -train_data = PairDatasetHF( - dataset["train"], - col_a="protein_a", - col_b="protein_b", - label_col="pkd", - max_length=2048, -) +Install only `requirements/core.in` and `requirements/features/train.in`, then +omit `--plot-results`, for the default plot-free run. +Generated plots are saved at 300 dpi as +`/classification_results.png` or +`/regression_results.png`. Existing plot files are never replaced. + +Regression uses three independently pinned dataset snapshots: + +```bash +PYTHONPATH=src python examples/fine_tuning.py \ + --task regression \ + --model_path Synthyra/ESM2-8M \ + --model-revision 185ecbd45665d050a8dae326d91886d330c5f9d0 \ + --regression-train-dataset-source Synthyra/ProteinProteinAffinity \ + --regression-train-dataset-revision f4a51e5e9f2c2a0185693f9fbcffc02d9dae08db \ + --regression-validation-dataset-source Synthyra/AffinityBenchmarkv5.5 \ + --regression-validation-dataset-revision 826ccfb1488d52b7b361802fbde161373247d084 \ + --regression-test-dataset-source Synthyra/haddock_benchmark \ + --regression-test-dataset-revision 4e22f014745728fca2d9c10f2f2cfd5a29a4981c \ + --attn-backend sdpa \ + --output-dir artifacts/fine-tuning \ + --seed 7 \ + --full-determinism ``` -For pair tasks, a custom collate function tokenizes both sequences and concatenates them with the tokenizer's separator handling. - -## Trainer Configuration - -The example uses a shared `BASE_TRAINER_KWARGS` dict: - -```python -BASE_TRAINER_KWARGS = { - "warmup_steps": 500, - "weight_decay": 0.01, - "logging_steps": 100, - "eval_strategy": "steps", - "eval_steps": 500, - "save_strategy": "steps", - "save_steps": 500, - "load_best_model_at_end": True, - "metric_for_best_model": "eval_loss", - "greater_is_better": False, - "report_to": "none", - "label_names": ["labels"], -} - -training_args = TrainingArguments( - output_dir="./results", - num_train_epochs=10, - per_device_train_batch_size=8, - per_device_eval_batch_size=16, - learning_rate=1e-4, - **BASE_TRAINER_KWARGS, -) +The exact shipped sources above use those pinned revisions automatically. +Custom remote sources reject omitted revisions, branches, and tags. For a fully +local run, pass existing local model and dataset directory paths through +`--model_path` and the corresponding `--*-dataset-source` options. Revisions +may be omitted for local directories because the example records an immutable +SHA-256 over the complete tree and rejects any tree that changes between model +initialization and final persistence. Pre-populate every pinned model and +dataset in the cache before starting a network-isolated Hub-backed run. +Local dataset directories must be layouts accepted by `datasets.load_dataset`; +the example does not currently accept arbitrary `Dataset.save_to_disk()` trees. + +The classification source must resolve to a `DatasetDict` with the required +`train`, `valid`, and `test` splits. Every split must be non-empty and +contain a `seqs` column of non-empty protein strings plus an integer `labels` +column. Training labels must be the contiguous zero-based set `0..K-1`; +validation and test labels must be subsets of the training labels. A local +file-backed layout can therefore use `train.csv`, `valid.csv`, and `test.csv` +in one directory, each with this header: + +```text +seqs,labels +MKT...,0 +GAV...,1 ``` -## Metrics - -### Regression (Spearman Correlation) - -```python -from scipy.stats import spearmanr +Regression uses three separate sources. Each source must expose a non-empty +`train` split with non-empty string columns `SeqA` and `SeqB` and a `labels` +column containing only finite real numbers. A local directory for each source +can contain `train.csv` with this header: -def compute_regression_metrics(eval_pred: EvalPrediction): - predictions = eval_pred.predictions.flatten() - labels = eval_pred.label_ids.flatten() - rho, p_value = spearmanr(predictions, labels) - return {"spearman_rho": rho, "spearman_p": p_value} +```text +SeqA,SeqB,labels +MKT...,GAV...,-8.42 ``` -### Classification (Confusion Matrix) +All split, column, and label contracts are checked before model initialization. +Malformed data cannot allocate a model or begin training. -```python -from sklearn.metrics import confusion_matrix - -def compute_classification_metrics(eval_pred: EvalPrediction): - predictions = eval_pred.predictions.argmax(axis=-1) - labels = eval_pred.label_ids - cm = confusion_matrix(labels, predictions) - accuracy = (predictions == labels).mean() - return {"accuracy": accuracy} -``` +## Sequence tasks -## Training +Tokenize proteins with the tokenizer belonging to the exact checkpoint revision. +Construct labels only for the task being trained and mask padding or ignored +positions explicitly. For residue-level tasks, align labels to biological +residues rather than assuming tokenizer position equals residue position. +For paired proteins, filter using the exact tokenizer encoding including +special tokens. The collator applies the same `max_length` with longest-first +truncation at the boundary. Throughout the CLI and manifest, `max_length` is an +encoded token budget, not a biological-residue count. It includes all +tokenizer-added BOS, EOS, separator, and other special tokens for a single +sequence or complete pair. ```python -from transformers import Trainer, EarlyStoppingCallback - -trainer = Trainer( - model=model, - args=training_args, - train_dataset=train_data, - eval_dataset=val_data, - compute_metrics=compute_regression_metrics, - callbacks=[EarlyStoppingCallback(early_stopping_patience=3)], -) - -trainer.train() -``` - -## LoRA Configuration - -The recommended defaults from `fastplms/fine_tuning_example.py`: +from transformers import AutoModelForMaskedLM, AutoTokenizer -| Parameter | Value | Notes | -|-----------|-------|-------| -| `r` | 16 | Rank of LoRA matrices | -| `lora_alpha` | 32 | Scaling factor | -| `target_modules` | `"all-linear"` | Applies LoRA to all linear layers (valid in peft >= 0.7) | -| `lora_dropout` | 0.1 | Dropout on LoRA paths | -| `bias` | `"none"` | Do not train bias terms | -| `task_type` | `"SEQ_CLS"` | Sequence classification task | - -## Model-Specific Notes - -- **ESM2, ESM++, DPLM, DPLM2**: Standard tokenizer-based fine-tuning via `AutoModelForSequenceClassification` -- **E1**: Requires custom collation because it uses sequence mode (no standard tokenizer). You need to override the Trainer's data collation to call `model.model.prep_tokens.get_batch_kwargs()` instead of the standard tokenizer -- **ESMC (ESM++)**: Ensure `sequence_id` is included in the forward pass inputs when batching - -## Saving and Loading - -```python -# Save LoRA adapter -model.save_pretrained("./lora_adapter") - -# Load later -from peft import PeftModel - -base_model = AutoModelForSequenceClassification.from_pretrained( - "Synthyra/ESM2-150M", +model_id = "Synthyra/ESM2-150M" +model_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + revision=model_revision, trust_remote_code=True, - num_labels=1, ) -model = PeftModel.from_pretrained(base_model, "./lora_adapter") +model = AutoModelForMaskedLM.from_pretrained( + model_id, + revision=model_revision, + trust_remote_code=True, + attn_implementation="sdpa", +) ``` + +E1 is the exception: it has no tokenizer and must use its native raw-sequence +preparation adapter. + +## Precision and attention + +Select a declared attention backend explicitly for a reproducible training run. +Do not rely on a silent fallback. Record the resolved backend, Torch and +Transformers versions, CUDA environment, checkpoint revision, tokenizer hashes, +dtype, optimizer, seed, and data fingerprint. + +The runnable example exposes `eager`, `sdpa`, and `flex_attention` through +`--attn-backend` and defaults to `sdpa`. It intentionally rejects +FlashAttention because the compact CLI does not expose the explicit BF16 CUDA +model-loading and placement policy that Flash training requires. The current +GH200/aarch64 validation environment is not expected to provide compatible +bundled Flash kernels, and this workflow does not recommend source-building +them. Prefer SDPA for the default training path and use Flex only on supported +Torch platforms. Any Flash gradient evidence remains a separate, +device-specific compliance result. `--output-dir` names a parent directory; +regression and classification are written into separate task- and +LoRA-specific children so Trainer state, final artifacts, and manifests do not +collide. + +Every selected task child must be absent before the command starts. The CLI +preflights all selected children, then each training function atomically +creates its own child before loading a model or dataset. An existing file, +directory, or broken symlink is rejected rather than reused. If an in-process +run fails, its newly reserved child is removed; an interrupted process leaves a +reservation marker and partial output that a later run refuses to mix. Inspect +and explicitly relocate or remove such a failed run before retrying. + +ESMFold2 FP8 is inference-only. Gradient-enabled ESMC execution reloads BF16, +and canonical training checkpoints contain only BF16 or FP32 weights. Runtime +Transformer Engine quantization state is never serialized. + +## Parameter-efficient adaptation + +Low-rank adapters are appropriate when full-model optimization is unnecessary. +Record exact target module names and verify that only intended parameters have +`requires_grad=True`. A separately trained task head, such as `classifier`, must +also be listed in PEFT's `modules_to_save`; setting `requires_grad=True` alone +does not include that head in the adapter checkpoint. Save adapter configuration +and base checkpoint revision together. The runnable example enforces this by +requiring the adapter configuration and safetensors payload, reloading into the +same immutable base, and comparing every persisted adapter and task-head tensor +hash. Full fine-tuning verifies the complete model state, including buffers. It +separately compares logits from the prepared Trainer and the +independently reloaded model on the first one or two held-out rows. + +LoRA is enabled by default in the runnable example. Pass `--no-use-lora` to +fine-tune the full model instead. + +## Reproducibility and evaluation + +Split homologous proteins at an identity threshold appropriate to the task to +reduce sequence leakage. Report class balance, length distribution, ambiguous +residue handling, truncation, and any structure-derived labels. Use task-specific +metrics and retain per-sequence predictions so aggregate improvements can be +audited. + +The CLI rejects moving Hub references before loading a model or dataset. The +example refuses a pre-existing task output before loading either source and +writes `run_manifest.json` with the requested immutable model source; +resolved base-weight and FastPLMs runtime revisions; tokenizer identity and +vocabulary hash; attention backend; parameter and compute dtypes; Trainer +device, optimizer, and scheduler; Torch, Transformers, PEFT, Datasets, CUDA, +Python, and package environment; command line and normalized configuration; +seeds and deterministic settings; adapter configuration and trainable target +modules; and each dataset's immutable source, split, consumed columns, row +count, and `ordered_rows_sha256` over ordered post-filter values. + +Final persistence is part of the run contract. The example saves into a +temporary sibling directory, requires the expected PEFT configuration and +safetensors weights, reloads that staged artifact against the same immutable +base, and refuses to continue if any persisted training-state hash changes. A +deterministic reload check then compares the +prepared Trainer's logits with independently reloaded-model logits for the +first `min(2, len(test_dataset))` rows, using the same collator, device, and +autocast policy with dtype-specific tolerances. Only after both checks pass is +the staging directory atomically renamed to `final_model`, without overwriting +an existing artifact. The manifest records the final +tree SHA-256, verified parameter hashes, `reload_verified: true`, and the +`held_out_inference` comparison metrics. The subsequent full +`trainer.predict(test_dataset)` evaluation deliberately remains on the original +prepared Trainer; it is not a full-dataset pass through the independently +reloaded instance. Retain the manifest, final artifact, and per-sequence +evaluation outputs together when making biological claims. + +The runnable minimal pattern is in [`examples/fine_tuning.py`](../examples/fine_tuning.py). +Model-family compliance establishes that the starting checkpoint is represented +correctly; it does not validate a fine-tuned model's biological claims. diff --git a/docs/generated/capability_evidence.md b/docs/generated/capability_evidence.md new file mode 100644 index 0000000..8972941 --- /dev/null +++ b/docs/generated/capability_evidence.md @@ -0,0 +1,201 @@ + + +# Capability-to-evidence manifest + +This manifest maps every advertised FastPLMs 1.0 capability to its user +documentation, runnable example, and required validation tier. It is a +coverage contract, not a statement that an unreported run passed. The exact +checkpoint list and family declarations come from `src/fastplms/models.toml`. + +The Example column links a curated CLI when that interface exposes the whole +capability. Programmatic-only forms instead link their runnable CPU contract so +the manifest does not imply broader CLI coverage than the example provides. + +## Frozen ESMC release evidence + +**Status: pending.** Default documentation generation never discovers or +trusts reports implicitly. Release rendering requires an explicitly selected, +complete schema-v3 set of exactly 30 records on one exact GH200/aarch64 +target: 18 measured eager, SDPA, and Flex records plus 12 structured +FlashAttention 2/3 unavailable records across three checkpoints and two +immutable sequence panels. +The set must also carry the final candidate/reference image identities, +dependency lock, installed inventory, and official source attestations. +A partial, stale, malformed, self-digest-invalid, or cross-device set fails +closed and cannot replace this status. + +### Locked oracle package compatibility exception + +The frozen oracle lock permits exactly one nonzero `pip check` diagnostic: +`nvidia-cusparselt-cu13 0.8.1 is not supported on this platform`. It applies only to +`nvidia-cusparselt-cu13==0.8.1` on +`NVIDIA GH200 480GB` / `linux` / +`aarch64`. The vendor filename tag is +`py3-none-manylinux2014_aarch64`, while the wheel metadata declares +`py3-none-manylinux2014_sbsa`. The exact wheel is +`nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl` with SHA-256 `4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f`. +FastPLMs accepts this vendor metadata mismatch only after the lock, installed +inventory, wheel bytes, metadata tag, and target identity all match. The wheel +is not rewritten (`validated-vendor-metadata-exception-no-wheel-rewrite`). Any additional diagnostic or +identity drift fails closed. + +## Executable evidence selectors + +Only the selectors below are claimed. Their scopes are intentionally narrower +than a whole family or validation tier. A tier appearing on another row does not +automatically apply to the capability in this row. + +| Selector | Tier/job | Executable target | Scope | +| --- | --- | --- | --- | +| `cpu:autoclass-runtime` | `cpu_contract` | `tests/cpu/test_autoclass_evidence_matrix.py::test_autoclass_runtime_evidence_matrix_exactly_matches_all_37_entries`
`tests/cpu/test_autoclass_evidence_matrix.py::test_autoclass_runtime_evidence_targets_are_collected_cpu_tests` | Every family-level AutoClass entry and its explicit tiny runtime contracts. | +| `artifact:checkpoint-autoclasses` | `artifact` | `tests/release/test_published_automodel.py::test_local_artifact_offline_autoclass_parity` | Every advertised AutoClass for every built checkpoint, grouped by checkpoint. | +| `compliance:sequence-primary-head` | `compliance` | `tests/parity/test_native_results.py::test_native_exact_checkpoint_contract`
`tests/parity/test_native_results.py::test_native_every_checkpoint_bf16_inference` | The official-parity head only: AutoModel for ANKH or a family without MaskedLM; otherwise AutoModelForMaskedLM. | +| `compliance:ankh-seq2seq` | `compliance` | `tests/parity/test_native_results.py::test_native_ankh_explicit_decoder_prompt_generation` | ANKH AutoModelForSeq2SeqLM explicit-prompt generation only. | +| `compliance:structure-automodel` | `compliance` | `tests/structure/test_esmfold_folding_compliance.py`
`tests/structure/test_esmfold2_folding_compliance.py` | ESMFold and ESMFold2 AutoModel folding paths only. | +| `benchmark:claim-eligible-primary-head` | `benchmark` | `benchmarks/suite.py::benchmark_cases[claim_eligible=True]` | The benchmark-selected head for representative sequence checkpoints and ESMFold2 projection cases; startup and embedding cases are excluded. | +| `cpu:attention-contracts` | `cpu_contract` | `tests/cpu/test_attention_contracts.py` | Portable dispatch, masks, fallback, fake FA2/FA3, ESMC Flex/FA3, and eager/SDPA gradient contracts. | +| `nightly:sequence-backends` | `nightly` | `tests/integration/test_backend_consistency.py` | Current GH200 eager, SDPA, and Flex forward/backward paths. Flash kernels are not downloaded, built, or executed in the current locked environment. | +| `historical:fa2-focused` | `historical` | `tools/remote/run.py::_kernel_capability_preflight` | Policy records prior real FlashAttention 2 focused execution, but the immutable execution report is not bundled in this repository and no current GH200 numerical claim is inferred from it. | +| `compliance:flash-unavailable-gh200` | `compliance` | `tests/parity/test_native_results.py::test_esmc_bf16_calibration_and_biological_holdout` | Complete report-bound FA2/FA3 unavailability records and fail-closed dispatch on the frozen release environment. | +| `compliance:deep-backends` | `compliance` | `tests/parity/test_native_results.py::test_native_representatives_all_backends` | Every advertised backend on the pinned deep sequence representative per family. | +| `benchmark:claim-eligible-backends` | `benchmark` | `benchmarks/suite.py::benchmark_cases[claim_eligible=True]` | Backends emitted by claim-eligible sequence and ESMFold2 benchmark cases. | +| `cpu:embedding-contracts` | `cpu_contract` | `tests/cpu/test_embedding_contracts.py` | Ordered inputs, biological masking, pooling, streaming, and persistence. | +| `cpu:e1-embeddings` | `cpu_contract` | `tests/cpu/test_e1_contracts.py` | E1 raw-sequence and MSA embedding persistence. | +| `feature:e1-rag` | `feature` | `tests/integration/test_e1_rag.py` | E1 retrieval, MSA preparation, cache, scoring, and embedding flows. | +| `cpu:ankh-contracts` | `cpu_contract` | `tests/cpu/test_ankh_contracts.py` | ANKH encoder and explicit-decoder embeddings, layers, masks, and T5 views. | +| `cpu:generation-contracts` | `cpu_contract` | `tests/cpu/test_generation_contracts.py` | Tiny deterministic DPLM, DPLM2, and ESM3 generation contracts. | +| `feature:generation` | `feature` | `tests/integration/test_dplm_generation.py`
`tests/integration/test_esm3.py` | DPLM, DPLM2, and ESM3 generation behavior in the feature suite. | +| `cpu:peft` | `cpu_contract` | `tests/cpu/test_peft_contracts.py` | Real initializer, collators, one optimizer step, and adapter/classifier reload. | +| `nightly:peft` | `nightly` | `tests/unit/test_fine_tuning_example.py` | Fine-tuning example contracts in the nightly feature job. | +| `cpu:ttt` | `cpu_contract` | `tests/cpu/test_ttt_contracts.py` | Seeded TTT initialization, update, reset, save, reload, and family isolation. | +| `feature:ttt` | `feature` | `tests/integration/test_ttt.py` | TTT integration behavior in the feature suite. | +| `cpu:structure-contracts` | `cpu_contract` | `tests/cpu/test_structure_contracts.py` | Tiny injected structure cores, public outputs, save/reload, and binder batching. | +| `structure:public-contracts` | `structure` | `tests/structure/test_structure_public_helpers.py` | Seeded Boltz helper, linker masking, real features, losses, and binder gradients. | +| `structure:full-suite` | `structure` | `tests/structure` | The declared GPU structure suite for folding and preparation behavior. | +| `feature:binder` | `feature` | `tests/integration/test_binder_design.py` | Seeded binder workflow, atom padding, critic ranking, and traceability. | +| `cpu:artifact-example` | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_artifact_loading_example_executes_local_only_autoconfig` | The offline local-artifact example with AutoConfig. | +| `cpu:task-head-example` | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_task_head_example_executes_all_advertised_heads_offline` | Offline ESM2 masked-LM scoring, contacts, sequence classification, and token classification through the documented example. | + +## Curated offline example execution + +Every curated example is routed to the exact collected CPU test nodes below. +These tests run under the required offline `cpu_contract` gate. + +| Example | Tier | Exact executable CPU node | +| --- | --- | --- | +| [`embedding_and_retrieval.py`](../../examples/embedding_and_retrieval.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_embedding_and_retrieval_example_executes_with_ordered_sqlite` | +| [`attention_switching.py`](../../examples/attention_switching.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_attention_switching_main_executes_optimized_and_masked_fallback` | +| [`ankh_embeddings.py`](../../examples/ankh_embeddings.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_ankh_embedding_example_executes_encoder_and_decoder_layers` | +| [`generation.py`](../../examples/generation.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_generation_example_executes_seeded_dplm_branch_offline` | +| [`generation.py`](../../examples/generation.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_generation_example_executes_seeded_dplm2_branch_offline` | +| [`generation.py`](../../examples/generation.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_generation_example_executes_seeded_esm3_trace` | +| [`e1_rag.py`](../../examples/e1_rag.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_e1_rag_example_executes_local_msa_and_shared_persistence` | +| [`ttt.py`](../../examples/ttt.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_ttt_example_executes_seeded_adapt_save_and_reset` | +| [`structure_preparation.py`](../../examples/structure_preparation.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_structure_preparation_example_executes_each_public_branch` | +| [`artifact_loading.py`](../../examples/artifact_loading.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_artifact_loading_example_executes_local_only_autoconfig` | +| [`task_heads.py`](../../examples/task_heads.py) | `cpu_contract` | `tests/cpu/test_documentation_contracts.py::test_task_head_example_executes_all_advertised_heads_offline` | +| [`fine_tuning.py`](../../examples/fine_tuning.py) | `cpu_contract` | `tests/cpu/test_peft_contracts.py::test_fine_tuning_main_wires_both_tasks_without_external_io` | +| [`fine_tuning.py`](../../examples/fine_tuning.py) | `cpu_contract` | `tests/cpu/test_peft_contracts.py::test_shipped_collators_create_tokenizer_aware_sequence_and_pair_batches` | +| [`fine_tuning.py`](../../examples/fine_tuning.py) | `cpu_contract` | `tests/cpu/test_peft_contracts.py::test_shipped_initializer_drives_one_peft_step_and_atomic_final_reload` | +| [`binder_design_fastplms.py`](../../examples/binder_design_fastplms.py) | `cpu_contract` | `tests/cpu/test_structure_contracts.py::test_public_binder_workflow_pads_heterogeneous_prepared_atoms_without_truncation` | +| [`binder_design_fastplms.py`](../../examples/binder_design_fastplms.py) | `cpu_contract` | `tests/cpu/test_structure_contracts.py::test_binder_example_main_wires_explicit_offline_cli_arguments` | +| [`binder_design_fastplms.py`](../../examples/binder_design_fastplms.py) | `cpu_contract` | `tests/cpu/test_structure_contracts.py::test_binder_structure_loss_is_finite_and_differentiable` | + +## Families and AutoClasses + +| Family | Tokenizer mode | AutoClass | Weight status | Guide | Family workflow and runnable entry-point contract | Required evidence | +| --- | --- | --- | --- | --- | --- | --- | +| `esm2` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#esm2) | [family workflow](../../examples/embedding_and_retrieval.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esm2` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#esm2) | [family workflow](../../examples/embedding_and_retrieval.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esm2` | `tokenizer` | `AutoModelForMaskedLM` | `pretrained` | [guide](../models.md#esm2) | [family workflow](../../examples/task_heads.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `cpu:task-head-example`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `esm2` | `tokenizer` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | [guide](../models.md#esm2) | [family workflow](../../examples/task_heads.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `cpu:task-head-example` | +| `esm2` | `tokenizer` | `AutoModelForTokenClassification` | `base weights + untrained task head` | [guide](../models.md#esm2) | [family workflow](../../examples/task_heads.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `cpu:task-head-example` | +| `esm_plusplus` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#esm-and-esmc) | [family workflow](../../examples/attention_switching.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esm_plusplus` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#esm-and-esmc) | [family workflow](../../examples/attention_switching.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esm_plusplus` | `tokenizer` | `AutoModelForMaskedLM` | `pretrained` | [guide](../models.md#esm-and-esmc) | [family workflow](../../examples/attention_switching.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `esm3` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#esm3) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esm3` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#esm3) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `e1` | `sequence` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#e1) | [family workflow](../../examples/e1_rag.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `e1` | `sequence` | `AutoModel` | `pretrained` | [guide](../models.md#e1) | [family workflow](../../examples/e1_rag.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `e1` | `sequence` | `AutoModelForMaskedLM` | `pretrained` | [guide](../models.md#e1) | [family workflow](../../examples/e1_rag.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `e1` | `sequence` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | [guide](../models.md#e1) | [family workflow](../../examples/e1_rag.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `e1` | `sequence` | `AutoModelForTokenClassification` | `base weights + untrained task head` | [guide](../models.md#e1) | [family workflow](../../examples/e1_rag.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#dplm) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#dplm) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm` | `tokenizer` | `AutoModelForMaskedLM` | `pretrained` | [guide](../models.md#dplm) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `dplm` | `tokenizer` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | [guide](../models.md#dplm) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm` | `tokenizer` | `AutoModelForTokenClassification` | `base weights + untrained task head` | [guide](../models.md#dplm) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm2` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#dplm2) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm2` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#dplm2) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm2` | `tokenizer` | `AutoModelForMaskedLM` | `pretrained` | [guide](../models.md#dplm2) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `dplm2` | `tokenizer` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | [guide](../models.md#dplm2) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `dplm2` | `tokenizer` | `AutoModelForTokenClassification` | `base weights + untrained task head` | [guide](../models.md#dplm2) | [family workflow](../../examples/generation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `ankh` | `tokenizer` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `ankh` | `tokenizer` | `AutoModel` | `pretrained` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:sequence-primary-head`, `benchmark:claim-eligible-primary-head` | +| `ankh` | `tokenizer` | `AutoModelForMaskedLM` | `FastPLMs extension` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `ankh` | `tokenizer` | `AutoModelForSeq2SeqLM` | `pretrained` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:ankh-seq2seq` | +| `ankh` | `tokenizer` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `ankh` | `tokenizer` | `AutoModelForTokenClassification` | `base weights + untrained task head` | [guide](../models.md#ankh) | [family workflow](../../examples/ankh_embeddings.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `boltz2` | `structure` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#boltz2) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `boltz2` | `structure` | `AutoModel` | `pretrained` | [guide](../models.md#boltz2) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esmfold` | `structure` | `AutoConfig` | `FastPLMs extension` | [guide](../models.md#esmfold) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esmfold` | `structure` | `AutoModel` | `pretrained` | [guide](../models.md#esmfold) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:structure-automodel` | +| `esmfold2` | `structure` | `AutoConfig` | `FastPLMs extension` | [guide](../esmfold2.md) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses` | +| `esmfold2` | `structure` | `AutoModel` | `pretrained` | [guide](../esmfold2.md) | [family workflow](../../examples/structure_preparation.py); [runnable AutoClass contract](../../tests/cpu/test_autoclass_evidence_matrix.py) | `cpu:autoclass-runtime`, `artifact:checkpoint-autoclasses`, `compliance:structure-automodel`, `benchmark:claim-eligible-primary-head` | + +## Attention backends + +| Backend | Advertising families | Guide | Example | Required evidence | +| --- | --- | --- | --- | --- | +| `eager` | `ankh`, `boltz2`, `dplm`, `esm2`, `esm3`, `esm_plusplus`, `esmfold`, `esmfold2` | [guide](../attention_backends.md) | [example](../../examples/attention_switching.py) | `cpu:attention-contracts`, `nightly:sequence-backends`, `compliance:deep-backends`, `benchmark:claim-eligible-backends` | +| `flash_attention_2` | `esm2`, `esm_plusplus` | [guide](../attention_backends.md) | [example](../../examples/attention_switching.py) | `cpu:attention-contracts`, `historical:fa2-focused`, `compliance:flash-unavailable-gh200` | +| `flash_attention_3` | `dplm`, `esm2`, `esm_plusplus` | [guide](../attention_backends.md) | [example](../../examples/attention_switching.py) | `cpu:attention-contracts`, `compliance:flash-unavailable-gh200` | +| `flex_attention` | `dplm`, `e1`, `esm2`, `esm3`, `esm_plusplus`, `esmfold`, `esmfold2` | [guide](../attention_backends.md) | [example](../../examples/attention_switching.py) | `cpu:attention-contracts`, `nightly:sequence-backends`, `compliance:deep-backends`, `benchmark:claim-eligible-backends` | +| `sdpa` | `ankh`, `dplm`, `dplm2`, `e1`, `esm2`, `esm3`, `esm_plusplus`, `esmfold`, `esmfold2` | [guide](../attention_backends.md) | [example](../../examples/attention_switching.py) | `cpu:attention-contracts`, `nightly:sequence-backends`, `compliance:deep-backends`, `benchmark:claim-eligible-backends` | + +## Input, embedding, and storage contracts + +| Capability | Guide | Example | Required evidence | +| --- | --- | --- | --- | +| Sequence list or streaming FASTA | [embedding API](../embedding_api.md) | [embedding and retrieval](../../examples/embedding_and_retrieval.py) | `cpu:embedding-contracts` | +| Ordered mapping or one-shot generator | [embedding API](../embedding_api.md) | [runnable API contracts](../../tests/cpu/test_embedding_contracts.py) | `cpu:embedding-contracts` | +| Biological-residue `max_length`, bounded token windows, and stable order | [embedding API](../embedding_api.md#bounded-streaming-and-length-policy) | [runnable API contracts](../../tests/cpu/test_embedding_contracts.py) | `cpu:embedding-contracts` | +| Mean and standard-deviation pooling | [embedding API](../embedding_api.md#pooling) | [embedding and retrieval](../../examples/embedding_and_retrieval.py) | `cpu:embedding-contracts` | +| Max/norm/median/variance/CLS/PARTI pooling | [embedding API](../embedding_api.md#pooling) | [runnable pooler contract](../../tests/cpu/test_embedding_contracts.py) | `cpu:embedding-contracts` | +| Full-residue and all-selected-layer output | [embedding API](../embedding_api.md#full-residue-embeddings) | [ANKH layers](../../examples/ankh_embeddings.py) | `cpu:embedding-contracts`, `cpu:ankh-contracts` | +| Transactional sharded safetensors and exact resume | [embedding API](../embedding_api.md#safetensors-storage) | [embedding and retrieval](../../examples/embedding_and_retrieval.py) | `cpu:embedding-contracts` | +| Read-only SQLite and ordered duplicate-preserving filters | [embedding API](../embedding_api.md#sqlite-streaming-retrieval-and-resume) | [embedding and retrieval](../../examples/embedding_and_retrieval.py) | `cpu:embedding-contracts` | +| Legacy SQLite conversion without pickle deserialization | [embedding API](../embedding_api.md#sqlite-streaming-retrieval-and-resume) | [runnable converter contract](../../tests/cpu/test_embedding_contracts.py) | `cpu:embedding-contracts` | +| E1 raw-sequence and MSA-aware ordered embeddings | [E1 guide](../models.md#e1) | [E1 RAG](../../examples/e1_rag.py) | `cpu:e1-embeddings`, `feature:e1-rag` | +| ANKH encoder/explicit-decoder hidden-state selection | [ANKH guide](../models.md#ankh) | [ANKH layers](../../examples/ankh_embeddings.py) | `cpu:ankh-contracts` | + +## Generation and adaptation contracts + +| Capability | Guide | Example | Required evidence | +| --- | --- | --- | --- | +| ESM2 pretrained masked-LM scoring and contact prediction | [ESM2](../models.md#esm2) | [task heads](../../examples/task_heads.py) | `cpu:task-head-example`, `cpu:autoclass-runtime` | +| ESM2 sequence/token classification with explicitly untrained task heads | [ESM2](../models.md#esm2) | [task heads](../../examples/task_heads.py) | `cpu:task-head-example`, `cpu:autoclass-runtime` | +| DPLM amino-acid diffusion generation | [DPLM](../models.md#dplm) | [generation](../../examples/generation.py) | `cpu:generation-contracts`, `feature:generation` | +| DPLM2 modality-aware sequence/structure co-generation | [DPLM2](../models.md#dplm2) | [generation](../../examples/generation.py) | `cpu:generation-contracts`, `feature:generation` | +| ESM3 multimodal-conditioned generation | [ESM3](../models.md#esm3) | [generation](../../examples/generation.py) | `cpu:generation-contracts`, `feature:generation` | +| ANKH task-prompted sequence-to-sequence generation | [ANKH](../models.md#ankh) | [ANKH embeddings and generation](../../examples/ankh_embeddings.py) | `cpu:ankh-contracts`, `compliance:ankh-seq2seq` | +| Trainer/PEFT LoRA with immutable inputs and verified save/reload | [fine-tuning](../finetuning.md) | [fine-tuning](../../examples/fine_tuning.py) | `cpu:peft`, `nightly:peft` | +| Seeded TTT adapter initialize/update/reset/save/reload | [TTT](../ttt.md) | [TTT](../../examples/ttt.py) | `cpu:ttt`, `feature:ttt` | + +## Structure contracts + +| Capability | Guide | Example | Required evidence | +| --- | --- | --- | --- | +| ESMFold single-chain folding and multimer-linker confidence masking | [models](../models.md#esmfold) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:public-contracts`, `structure:full-suite`, `compliance:structure-automodel` | +| Seed-scoped Boltz2 protein helper and BF16 execution policy | [Boltz2](../models.md#boltz2) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:public-contracts`, `structure:full-suite` | +| Atom-dense binder optimization and critic reporting | [binder design](../binder_design.md) | [binder design](../../examples/binder_design_fastplms.py) | `cpu:structure-contracts`, `structure:public-contracts`, `feature:binder` | +| Offline local artifact AutoClass loading | [artifacts](../artifacts.md) | [artifact loading](../../examples/artifact_loading.py) | `cpu:artifact-example`, `artifact:checkpoint-autoclasses` | +| `esmfold2` 48-block full ESMFold2: single-sequence or optional MSA-conditioned protein inputs, typed complexes, ligands, nucleic acids, modifications, bonds, and distograms; pocket requests fail closed | [ESMFold2](../esmfold2.md) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:full-suite`, `compliance:structure-automodel` | +| `esmfold2_fast` 24-block Fast ESMFold2: inference-optimized single-sequence conditioning with typed multichain and multimolecule inputs; every protein must have `msa=None` and MSA inputs fail closed | [ESMFold2](../esmfold2.md) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:full-suite`, `compliance:structure-automodel` | +| `esmfold2_experimental_cutoff2025` 48-block full ESMFold2: single-sequence or optional MSA-conditioned protein inputs, typed complexes, ligands, nucleic acids, modifications, bonds, and distograms; pocket requests fail closed | [ESMFold2](../esmfold2.md) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:full-suite`, `compliance:structure-automodel` | +| `esmfold2_experimental_fast_cutoff2025` 24-block Fast ESMFold2: inference-optimized single-sequence conditioning with typed multichain and multimolecule inputs; every protein must have `msa=None` and MSA inputs fail closed | [ESMFold2](../esmfold2.md) | [structure preparation](../../examples/structure_preparation.py) | `cpu:structure-contracts`, `structure:full-suite`, `compliance:structure-automodel` | + +Release evidence must name the exact head, checkpoint and runtime revisions, +tokenizer identity, backend, dtype, hardware, sequence or structure panel, +seed, environment, and input hash. Missing evidence remains visibly pending; +it must not be replaced by a synthetic benchmark number or an inferred claim. diff --git a/docs/generated/support.md b/docs/generated/support.md new file mode 100644 index 0000000..6ac8888 --- /dev/null +++ b/docs/generated/support.md @@ -0,0 +1,138 @@ + + +# Model support + +This file is generated from `src/fastplms/models.toml`. A listed capability is +selectable. Strict-parity exceptions are documented in the checkpoint cards. + +## Family interfaces + +| Family | Architecture | Checkpoints | Public input | AutoClasses | Tokenizer class | +| --- | --- | ---: | --- | --- | --- | +| `esm2` | ESM2 | 5 | Amino-acid sequences tokenized to residue IDs | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` | `n/a` | +| `esm_plusplus` | ESMC | 3 | Amino-acid sequences tokenized to residue IDs | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM` | `n/a` | +| `esm3` | ESM3 | 1 | Sequence, structure, and function tracks prepared through the multimodal helpers | `AutoConfig`, `AutoModel` | `n/a` | +| `e1` | E1 | 3 | Raw amino-acid sequences prepared by the native E1 adapter | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` | `n/a` | +| `dplm` | DPLM | 3 | Amino-acid sequences tokenized to masked or partially masked residue IDs | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` | `n/a` | +| `dplm2` | DPLM2 | 3 | Tokenized amino-acid and structure tracks with explicit modality boundaries | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` | `fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer` | +| `ankh` | ANKH | 5 | Amino-acid sequences tokenized for encoder or sequence-to-sequence use | `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` | `n/a` | +| `boltz2` | Boltz2 | 1 | Raw amino-acid sequences through the convenience API, or prepared model features | `AutoConfig`, `AutoModel` | `n/a` | +| `esmfold` | ESMFold | 1 | Raw amino-acid sequences through folding helpers, or prepared residue tensors | `AutoConfig`, `AutoModel` | `n/a` | +| `esmfold2` | ESMFold2 | 4 | Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors | `AutoConfig`, `AutoModel` | `n/a` | + +## AutoClass weight status + +`pretrained` means the advertised head is present in the checkpoint. `base weights + untrained task head` means the task head must be trained before use. `FastPLMs extension` is an integration or head that is not an official pretrained ANKH capability. + +| Family | AutoClass | Weight status | +| --- | --- | --- | +| `esm2` | `AutoConfig` | `FastPLMs extension` | +| `esm2` | `AutoModel` | `pretrained` | +| `esm2` | `AutoModelForMaskedLM` | `pretrained` | +| `esm2` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | +| `esm2` | `AutoModelForTokenClassification` | `base weights + untrained task head` | +| `esm_plusplus` | `AutoConfig` | `FastPLMs extension` | +| `esm_plusplus` | `AutoModel` | `pretrained` | +| `esm_plusplus` | `AutoModelForMaskedLM` | `pretrained` | +| `esm3` | `AutoConfig` | `FastPLMs extension` | +| `esm3` | `AutoModel` | `pretrained` | +| `e1` | `AutoConfig` | `FastPLMs extension` | +| `e1` | `AutoModel` | `pretrained` | +| `e1` | `AutoModelForMaskedLM` | `pretrained` | +| `e1` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | +| `e1` | `AutoModelForTokenClassification` | `base weights + untrained task head` | +| `dplm` | `AutoConfig` | `FastPLMs extension` | +| `dplm` | `AutoModel` | `pretrained` | +| `dplm` | `AutoModelForMaskedLM` | `pretrained` | +| `dplm` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | +| `dplm` | `AutoModelForTokenClassification` | `base weights + untrained task head` | +| `dplm2` | `AutoConfig` | `FastPLMs extension` | +| `dplm2` | `AutoModel` | `pretrained` | +| `dplm2` | `AutoModelForMaskedLM` | `pretrained` | +| `dplm2` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | +| `dplm2` | `AutoModelForTokenClassification` | `base weights + untrained task head` | +| `ankh` | `AutoConfig` | `FastPLMs extension` | +| `ankh` | `AutoModel` | `pretrained` | +| `ankh` | `AutoModelForMaskedLM` | `FastPLMs extension` | +| `ankh` | `AutoModelForSeq2SeqLM` | `pretrained` | +| `ankh` | `AutoModelForSequenceClassification` | `base weights + untrained task head` | +| `ankh` | `AutoModelForTokenClassification` | `base weights + untrained task head` | +| `boltz2` | `AutoConfig` | `FastPLMs extension` | +| `boltz2` | `AutoModel` | `pretrained` | +| `esmfold` | `AutoConfig` | `FastPLMs extension` | +| `esmfold` | `AutoModel` | `pretrained` | +| `esmfold2` | `AutoConfig` | `FastPLMs extension` | +| `esmfold2` | `AutoModel` | `pretrained` | + +## Family execution + +| Family | Attention | Precision | BF16 execution | Extra | Reference | +| --- | --- | --- | --- | --- | --- | +| `esm2` | `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | `default` | `fp32_parameters_autocast` | `core` | `reference-esm2` | +| `esm_plusplus` | `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | `default` | `static_parameters` | `core` | `reference-biohub-esm` | +| `esm3` | `eager`, `sdpa`, `flex_attention` | `default` | `fp32_parameters_autocast` | `core` | `reference-biohub-esm` | +| `e1` | `sdpa`, `flex_attention` | `default` | `static_parameters` | `core` | `reference-e1` | +| `dplm` | `eager`, `sdpa`, `flex_attention`, `flash_attention_3` | `default` | `fp32_parameters_autocast` | `core` | `reference-dplm` | +| `dplm2` | `sdpa` | `default` | `fp32_parameters_autocast` | `core` | `reference-dplm` | +| `ankh` | `eager`, `sdpa` | `default` | `static_parameters` | `core` | `reference-ankh` | +| `boltz2` | `eager` | `default` | `fp32_parameters_autocast` | `structure` | `reference-boltz2` | +| `esmfold` | `eager`, `sdpa`, `flex_attention` | `default` | `fp32_parameters_autocast` | `structure` | `reference-esmfold` | +| `esmfold2` | `eager`, `sdpa`, `flex_attention` | `auto`, `fp32`, `bf16`, `fp8` (experimental) | `fp32_parameters_autocast` | `structure` | `reference-esmfold2` | + +## Family release contracts + +| Family | Checkpoint terms | Hub license | Weight publication | Tiers | +| --- | --- | --- | --- | --- | +| `esm2` | MIT | `mit` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `esm_plusplus` | MIT | `mit` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `esm3` | MIT | `mit` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `e1` | Profluent-E1-Agreement | `other` ([Profluent-E1 Clickthrough License Agreement](https://github.com/Profluent-AI/E1/blob/main/LICENSE)) | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `dplm` | Apache-2.0 | `apache-2.0` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `dplm2` | Apache-2.0 | `apache-2.0` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `ankh` | CC-BY-NC-SA-4.0 | `cc-by-nc-sa-4.0` | manifest policy | `check`, `compliance`, `feature`, `artifact`, `benchmark` | +| `boltz2` | MIT | `mit` | manifest policy | `structure`, `artifact`, `benchmark` | +| `esmfold` | MIT | `mit` | manifest policy | `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` | +| `esmfold2` | MIT | `mit` | manifest policy | `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` | + +## Runtime assets + +| ID | Family | Repository | Path | SHA-256 | Size | License | Trust boundary | Offline behavior | +| --- | --- | --- | --- | --- | ---: | --- | --- | --- | +| `esmfold2_ccd` | `esmfold2` | `biohub/ESMFold2` | `ccd.pkl` | `9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` | 417306584 | `MIT` | `hash_pinned_pickle` | `requires_cached_verified_file` | + +## Checkpoints + +| ID | Family | Size | FastPLMs checkpoint | Official checkpoint | Artifact source | State transform | Generation contract | MSA conditioning | Unresolved files | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | ---: | +| `esm2_8m` | `esm2` | `small` | [Synthyra/ESM2-8M](https://huggingface.co/Synthyra/ESM2-8M) | [facebook/esm2_t6_8M_UR50D](https://huggingface.co/facebook/esm2_t6_8M_UR50D) | `fast` | `esm2_hf_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esm2_35m` | `esm2` | `small` | [Synthyra/ESM2-35M](https://huggingface.co/Synthyra/ESM2-35M) | [facebook/esm2_t12_35M_UR50D](https://huggingface.co/facebook/esm2_t12_35M_UR50D) | `fast` | `esm2_hf_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esm2_150m` | `esm2` | `medium` | [Synthyra/ESM2-150M](https://huggingface.co/Synthyra/ESM2-150M) | [facebook/esm2_t30_150M_UR50D](https://huggingface.co/facebook/esm2_t30_150M_UR50D) | `fast` | `esm2_hf_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esm2_650m` | `esm2` | `large` | [Synthyra/ESM2-650M](https://huggingface.co/Synthyra/ESM2-650M) | [facebook/esm2_t33_650M_UR50D](https://huggingface.co/facebook/esm2_t33_650M_UR50D) | `fast` | `esm2_hf_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esm2_3b` | `esm2` | `xlarge` | [Synthyra/ESM2-3B](https://huggingface.co/Synthyra/ESM2-3B) | [facebook/esm2_t36_3B_UR50D](https://huggingface.co/facebook/esm2_t36_3B_UR50D) | `fast` | `esm2_hf_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esmc_small` | `esm_plusplus` | `medium` | [Synthyra/ESMplusplus_small](https://huggingface.co/Synthyra/ESMplusplus_small) | [biohub/ESMC-300M](https://huggingface.co/biohub/ESMC-300M) | `fast` | `esmc_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esmc_large` | `esm_plusplus` | `large` | [Synthyra/ESMplusplus_large](https://huggingface.co/Synthyra/ESMplusplus_large) | [biohub/ESMC-600M](https://huggingface.co/biohub/ESMC-600M) | `fast` | `esmc_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esmc_6b` | `esm_plusplus` | `xlarge` | [Synthyra/ESMplusplus_6B](https://huggingface.co/Synthyra/ESMplusplus_6B) | [biohub/ESMC-6B](https://huggingface.co/biohub/ESMC-6B) | `fast` | `esmc_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esm3_small` | `esm3` | `large` | [Synthyra/ESM3_small](https://huggingface.co/Synthyra/ESM3_small) | [biohub/esm3-sm-open-v1](https://huggingface.co/biohub/esm3-sm-open-v1) | `fast` | `esm3_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `e1_150m` | `e1` | `small` | [Synthyra/Profluent-E1-150M](https://huggingface.co/Synthyra/Profluent-E1-150M) | [Profluent-Bio/E1-150m](https://huggingface.co/Profluent-Bio/E1-150m) | `fast` | `e1_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `e1_300m` | `e1` | `medium` | [Synthyra/Profluent-E1-300M](https://huggingface.co/Synthyra/Profluent-E1-300M) | [Profluent-Bio/E1-300m](https://huggingface.co/Profluent-Bio/E1-300m) | `fast` | `e1_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `e1_600m` | `e1` | `large` | [Synthyra/Profluent-E1-600M](https://huggingface.co/Synthyra/Profluent-E1-600M) | [Profluent-Bio/E1-600m](https://huggingface.co/Profluent-Bio/E1-600m) | `fast` | `e1_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `dplm_150m` | `dplm` | `small` | [Synthyra/DPLM-150M](https://huggingface.co/Synthyra/DPLM-150M) | [airkingbd/dplm_150m](https://huggingface.co/airkingbd/dplm_150m) | `fast` | `dplm_to_fastplms_v1` | `required` | not applicable | 0 | +| `dplm_650m` | `dplm` | `large` | [Synthyra/DPLM-650M](https://huggingface.co/Synthyra/DPLM-650M) | [airkingbd/dplm_650m](https://huggingface.co/airkingbd/dplm_650m) | `fast` | `dplm_to_fastplms_v1` | `required` | not applicable | 0 | +| `dplm_3b` | `dplm` | `xlarge` | [Synthyra/DPLM-3B](https://huggingface.co/Synthyra/DPLM-3B) | [airkingbd/dplm_3b](https://huggingface.co/airkingbd/dplm_3b) | `fast` | `dplm_to_fastplms_v1` | `required` | not applicable | 0 | +| `dplm2_150m` | `dplm2` | `small` | [Synthyra/DPLM2-150M](https://huggingface.co/Synthyra/DPLM2-150M) | [airkingbd/dplm2_150m](https://huggingface.co/airkingbd/dplm2_150m) | `official` | `dplm2_to_fastplms_v1` | `required` | not applicable | 0 | +| `dplm2_650m` | `dplm2` | `large` | [Synthyra/DPLM2-650M](https://huggingface.co/Synthyra/DPLM2-650M) | [airkingbd/dplm2_650m](https://huggingface.co/airkingbd/dplm2_650m) | `official` | `dplm2_to_fastplms_v1` | `required` | not applicable | 0 | +| `dplm2_3b` | `dplm2` | `xlarge` | [Synthyra/DPLM2-3B](https://huggingface.co/Synthyra/DPLM2-3B) | [airkingbd/dplm2_3b](https://huggingface.co/airkingbd/dplm2_3b) | `official` | `dplm2_to_fastplms_v1` | `official_unavailable` | not applicable | 0 | +| `ankh_base` | `ankh` | `medium` | [Synthyra/ANKH_base](https://huggingface.co/Synthyra/ANKH_base) | [ElnaggarLab/ankh-base](https://huggingface.co/ElnaggarLab/ankh-base) | `official` | `ankh_t5_to_fastplms_v1` | `required` | not applicable | 0 | +| `ankh_large` | `ankh` | `large` | [Synthyra/ANKH_large](https://huggingface.co/Synthyra/ANKH_large) | [ElnaggarLab/ankh-large](https://huggingface.co/ElnaggarLab/ankh-large) | `official` | `ankh_t5_to_fastplms_v1` | `required` | not applicable | 0 | +| `ankh2_large` | `ankh` | `large` | [Synthyra/ANKH2_large](https://huggingface.co/Synthyra/ANKH2_large) | [ElnaggarLab/ankh2-ext2](https://huggingface.co/ElnaggarLab/ankh2-ext2) | `official` | `ankh_t5_to_fastplms_v1` | `required` | not applicable | 0 | +| `ankh3_large` | `ankh` | `large` | [Synthyra/ANKH3_large](https://huggingface.co/Synthyra/ANKH3_large) | [ElnaggarLab/ankh3-large](https://huggingface.co/ElnaggarLab/ankh3-large) | `official` | `ankh_t5_to_fastplms_v1` | `required` | not applicable | 0 | +| `ankh3_xl` | `ankh` | `xlarge` | [Synthyra/ANKH3_xl](https://huggingface.co/Synthyra/ANKH3_xl) | [ElnaggarLab/ankh3-xl](https://huggingface.co/ElnaggarLab/ankh3-xl) | `official` | `ankh_t5_to_fastplms_v1` | `required` | not applicable | 0 | +| `boltz2` | `boltz2` | `structure` | [Synthyra/Boltz2](https://huggingface.co/Synthyra/Boltz2) | [boltz-community/boltz-2](https://huggingface.co/boltz-community/boltz-2) | `fast` | `boltz2_inference_core_v1` | `not_applicable` | not applicable | 0 | +| `esmfold` | `esmfold` | `structure` | [Synthyra/FastESMFold](https://huggingface.co/Synthyra/FastESMFold) | [facebook/esmfold_v1](https://huggingface.co/facebook/esmfold_v1) | `fast` | `esmfold_meta_to_fastplms_v1` | `not_applicable` | not applicable | 0 | +| `esmfold2` | `esmfold2` | `structure` | [Synthyra/ESMFold2](https://huggingface.co/Synthyra/ESMFold2) | [biohub/ESMFold2](https://huggingface.co/biohub/ESMFold2) | `fast` | `identity` | `not_applicable` | `optional` (full checkpoint) | 0 | +| `esmfold2_fast` | `esmfold2` | `structure` | [Synthyra/ESMFold2-Fast](https://huggingface.co/Synthyra/ESMFold2-Fast) | [biohub/ESMFold2-Fast](https://huggingface.co/biohub/ESMFold2-Fast) | `fast` | `identity` | `not_applicable` | `none` (Fast; MSA inputs rejected) | 0 | +| `esmfold2_experimental_cutoff2025` | `esmfold2` | `structure` | [Synthyra/ESMFold2-Experimental-Cutoff2025](https://huggingface.co/Synthyra/ESMFold2-Experimental-Cutoff2025) | [biohub/ESMFold2-Experimental-Cutoff2025](https://huggingface.co/biohub/ESMFold2-Experimental-Cutoff2025) | `fast` | `identity` | `not_applicable` | `optional` (full checkpoint) | 0 | +| `esmfold2_experimental_fast_cutoff2025` | `esmfold2` | `structure` | [Synthyra/ESMFold2-Experimental-Fast-Cutoff2025](https://huggingface.co/Synthyra/ESMFold2-Experimental-Fast-Cutoff2025) | [biohub/ESMFold2-Experimental-Fast-Cutoff2025](https://huggingface.co/biohub/ESMFold2-Experimental-Fast-Cutoff2025) | `fast` | `identity` | `not_applicable` | `none` (Fast; MSA inputs rejected) | 0 | + +A nonzero unresolved-file count blocks release. It is not permission to +omit that file from checkpoint, tokenizer, artifact, or compliance checks. diff --git a/docs/licensing.md b/docs/licensing.md new file mode 100644 index 0000000..6b87a1a --- /dev/null +++ b/docs/licensing.md @@ -0,0 +1,72 @@ +# Licensing and attribution + +FastPLMs source is distributed under the Apache License 2.0. Model checkpoints, +official source repositories, and copied third-party components retain their own +terms. The model manifest records code and checkpoint licenses separately. + +Each family also records Hugging Face model-card metadata separately from its +human-readable checkpoint terms. Standard checkpoints use the Hub identifiers +`mit`, `apache-2.0`, or `cc-by-nc-sa-4.0`. E1 uses `other` with the name and +source link for its clickthrough agreement. DPLM1 and DPLM2 use +`apache-2.0`. Missing or mismatched identifiers block generation and artifact +validation. + +`LICENSES/` contains distributable copies of required legal texts. +`THIRD_PARTY_NOTICES.md` maps each model family and component to its source, +revision, terms, modifications, and attribution. Release validation compares +these files with the canonical pinned upstream files by hash. + +## ANKH + +ANKH implementations and mirrored weights are retained under CC BY-NC-SA 4.0. +Artifacts and model cards display those terms prominently. FastPLMs does not +implement a runtime restriction or decide whether a particular use satisfies +the license. + +## E1 + +E1 retains the upstream agreement, `ATTRIBUTION`, `NOTICE`, Apache and BSD +texts, modified-file notices, and documentation attribution. Relevant launches +display `Profluent-E1`. Redistribution and use remain subject to the upstream +agreement; review `LICENSES/e1/` before use. + +## DPLM + +The ByteDance DPLM repository carries an +[Apache-2.0 license](https://github.com/bytedance/dplm/blob/main/LICENSE) +and its [official README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +defines the repository release as including pretrained weights for both DPLM1 +and DPLM2. FastPLMs therefore records both checkpoint families as Apache-2.0, +with Hub metadata `license: apache-2.0`, +`weights_license_status="resolved"`, and `redistributable=true`. + +The verbatim license and immutable evidence record are distributed from +`LICENSES/dplm/`. Complete publication remains fail-closed unless the artifact, +legal inventory, state-parity evidence, and atomic Hub preflight all pass. + +## Biohub, Boltz, Meta, OpenFold, and ProteinTTT + +Biohub MIT and Apache notices, including `THIRD_PARTY_NOTICE`, are retained for +ESM++, ESM3, and ESMFold2. Boltz MIT terms, Meta ESM and ESMFold notices, +OpenFold notices, and ProteinTTT source records are included where their code or +derived behavior is distributed. + +ESMFold2 additionally uses a 417,306,584-byte `ccd.pkl` runtime asset under MIT +terms. The manifest pins repository, immutable revision, path, size, and +SHA-256. Because it is pickle, validated deserialization is an explicit trust +boundary. The loader rejects user and `cache_dir` symlinks, except for the exact +manifest snapshot link resolving into its repository's contained blob +directory. It copies into a private loader-owned temporary snapshot, verifies +that snapshot's size and hash, and unpickles only those verified bytes. Offline +runs require the exact cache object and never fetch a substitute. + +## Artifact and container rules + +An artifact or reference container fails validation when a required source +revision, license, notice, attribution, modified-file record, or conversion +record is absent. Reference images receive only the legal files relevant to +their corresponding upstream. Runtime images contain no official submodules or +checkpoint weights. + +This guide summarizes repository policy and is not legal advice. The complete +texts in `LICENSES/` and the original upstream sources control. diff --git a/docs/migration.md b/docs/migration.md new file mode 100644 index 0000000..fbfba5f --- /dev/null +++ b/docs/migration.md @@ -0,0 +1,309 @@ +# Migration to FastPLMs 1.0 + +FastPLMs 1.0 intentionally changes the repository layout, embedding return +types and storage, attention selection, artifact publication, and several model +contracts. There are no compatibility imports or silent keyword aliases. Run +the snippets in this guide in the offline CPU documentation job whenever this +contract changes. + +## Dependencies and source layout + +FastPLMs 1.0 requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13. +The project is no longer installed as a Python distribution. Published models +carry their runtime source in the Hugging Face repository and load it with +`trust_remote_code=True`. Install dependencies directly: + +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" +``` + +Repository tools and source-level APIs run with `PYTHONPATH=src`. Their +dependencies are composed under `requirements/` rather than distribution +metadata: + +```bash +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu +PYTHONPATH=src python -m pytest tests/cpu -m cpu_contract +``` + +Source remains under `src`: + +| Pre-1.0 path/import | FastPLMs 1.0 | +| --- | --- | +| `fastplms/esm2/...` | `src/fastplms/models/esm2/...` | +| family-local attention helpers | `fastplms.attention` | +| `fastplms.embedding_mixin` | `fastplms.embeddings` and `fastplms.embed_dataset` | +| family-local TTT helpers | `fastplms.models.ttt` | +| `testing/...` operational scripts | `tools.remote`, `tools.artifacts`, `tools.goldens`, or `benchmarks` | + +Model implementations remain loadable through the AutoClasses declared in the +[generated support matrix](generated/support.md). Runtime source must not import +`vendor/upstream`; those repositories are isolated compliance oracles. + +## Attention selection and outputs + +Replace `config.attn_backend`, `model.attn_backend`, `flex`, +`kernels_flash`, and `auto` with the Transformers attention interface: + +```python +from transformers import AutoModel + +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="flex_attention", +) +model.set_attn_implementation("sdpa") +``` + +The 1.0 names are `eager`, `sdpa`, `flex_attention`, +`flash_attention_2`, and `flash_attention_3`, restricted by the manifest. +There is no `auto` backend, source compilation, or unavailable-kernel fallback. + +`output_attentions=True` is a documented exception: when an optimized backend +cannot return the full matrix, FastPLMs emits one warning containing the +configured backend, effective `eager` backend, and reason. It derives the eager +4-D mask from the original padding and causal semantics for that call only. +Configuration and later calls are unchanged. + +ESMC Flex Attention and FlashAttention 3 remain supported and +non-experimental. SDPA is recommended for highest numerical fidelity. Their +published deviations are diagnostic warnings, while dispatch, masks, finite +outputs, shapes, and catastrophic biological disagreement remain hard gates. +Do not convert a threshold into a measured number. + +## Embedding arguments + +The historical method deduplicated and length-sorted sequences, returned a +sequence-keyed dictionary, and defaulted to writing pickle. The 1.0 operation +preserves order and duplicates and returns `EmbeddingResult`. + +| Pre-1.0 argument | FastPLMs 1.0 replacement | +| --- | --- | +| `sequences=values` | positional `inputs`, accepting sequences, pairs, records, mapping, generator, or FASTA path | +| `fasta_path=path` | positional `inputs=path`; combine sources explicitly with a generator when needed | +| `max_len=n` | `max_length=n`, always measured in biological residues | +| `pooling_types=[...]` | `pooling=(...)` | +| `embed_dtype=dtype` | `dtype=dtype` | +| `save=True, save_path=...` | `output=..., format="safetensors"` | +| `sql=True, sql_db_path=...` | `output=..., format="sqlite"` | +| `padding="max_length"` or `"longest"` | bounded `batch_window_size` and optional `max_tokens_per_batch` | +| `num_workers` | removed; FASTA parsing and immutable spooling are bounded and deterministic | +| `hidden_state_index`, `store_all_hidden_states` | retained, applied to the model-selected hidden-state stack | + +```python +result = model.embed_dataset( + inputs, + batch_size=8, + batch_window_size=64, + max_tokens_per_batch=8192, + max_length=1024, + pooling=("mean",), + output="embeddings.sqlite", + format="sqlite", + resume=True, +) +``` + +`pooling=None` selects mean pooling unless `full_embeddings=True`. An explicit +pooler with `full_embeddings=True` raises. Full embeddings contain only +biological residues; BOS, EOS, padding, chain delimiters, and structure-only +tokens are removed. + +E1's pre-1.0 `embed_dataset_with_msa` dictionary return is also replaced by an +ordered `EmbeddingResult`. Duplicate queries are preserved, and record IDs are +their zero-based input positions. Its native names `max_len`, `pooling_types`, +and `matrix_embed` remain because E1 has no tokenizer, while `output`, `format`, +`resume`, `shard_size`, `model_state_fingerprint`, `batch_window_size`, and +`max_tokens_per_batch` use the shared persistence and bounded-batching +contracts. See [`examples/e1_rag.py`](../examples/e1_rag.py). + +## Return values and storage readers + +Replace dictionary indexing with ordered records: + +```python +for record in result: + tensor = record.load_tensor() + print(record.id, record.sequence, tensor.shape) +``` + +`result.as_dict(key="id")` raises for duplicate keys unless an explicit +duplicate policy is chosen. Sharded safetensors and SQLite are the writable +formats. The run manifest is the authoritative safetensors commit and can +recover a valid committed run when the standalone index is interrupted. + +Replace old readers as follows: + +```python +from fastplms.embeddings import ( + convert_legacy_sqlite, + load_legacy_pth, + load_result, + load_sqlite_result, +) + +current = load_result("embeddings") +selected = load_sqlite_result( + "embeddings.sqlite", + record_ids=["b", "a", "b"], +) +convert_legacy_sqlite("legacy.db", "embeddings-v1.sqlite") +trusted_pickle = load_legacy_pth( + "legacy.pth", + allow_unsafe_pickle=True, +) +``` + +SQLite retrieval opens read-only and preserves selector order and duplicates. +Legacy pickle remains executable input and therefore requires explicit opt-in. + +## ANKH full-checkpoint replacement + +FastPLMs 1.0 replaced each Synthyra encoder-only ANKH repository with the full +official-compatible T5 state. This increases the default checkpoint size while +preserving encoder output parity. + +`AutoModel` loads the encoder and shared state without decoder allocation. +`AutoModelForSeq2SeqLM` loads encoder, decoder, cross-attention, and LM head from +the same repository. Set `ankh_id` to a Synthyra ANKH repository or a validated +local artifact: + +```python +from transformers import AutoModel, AutoModelForSeq2SeqLM + +encoder = AutoModel.from_pretrained( + ankh_id, + revision=ankh_revision, + trust_remote_code=True, +) +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + ankh_id, + revision=ankh_revision, + trust_remote_code=True, +) +``` + +ANKH embeddings default to `hidden_state_source="encoder"` and +`hidden_state_index=-1`. Decoder extraction requires exactly one explicit +aligned decoder text list or ID tensor: + +Remove the pre-1.0 residue-spacing workaround. Pass raw sequences such as +`MSTNPK`, not `M S T N P K`, and write sentinel prompts as +`M`, not `M `. FastPLMs now applies the same safe +pre-tokenizer and normalization to model-owned and explicitly supplied ANKH +tokenizers. + +```python +encoder_layers = encoder.embed_dataset( + inputs, + hidden_state_source="encoder", + store_all_hidden_states=True, + full_embeddings=True, +) +decoder_layer = seq2seq.embed_dataset( + inputs, + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M" for _ in inputs], + full_embeddings=True, +) +``` + +No shifted-source decoder input is invented. The official ANKH workflows use +task prompts, sentinels, or generated tokens. Decoder pooling excludes special +tokens and persisted metadata fingerprints the decoder input and alignment. + +Files-only publication is forbidden for this migration. Every weight shard, +weight index, tokenizer asset, configuration, runtime source, model card, and +release record must land in one immutable Hub commit. Both AutoClass views +must pass artifact and live parity from that same commit. + +Complete publication may remove a superseded monolithic weight path in that +same commit only when the path is pinned in the current registry, absent from +the validated sharded inventory, and its remote digest and parent still match +preflight. Files-only publication remains strictly add-only. + +The separately named `FastAnkhForMaskedLMExtension` remains a FastPLMs +extension, not an official ANKH head. Sequence and token classification views +load pretrained base weights with newly initialized task heads. + +## AutoClass weight meaning + +Every family and entry point is classified in the +[capability-to-evidence manifest](generated/capability_evidence.md): + +- `pretrained`: the advertised base or head exists in checkpoint state; +- `base weights + untrained task head`: train the classification head before + interpreting logits; +- `FastPLMs extension`: integration code or a head that is not an official + pretrained capability. + +All advertised classes must honor `return_dict`, output flags, tuple order, +embedding resize and setters, initialization, forward/loss/backward, and +save/reload. + +ESM2, ESMC, DPLM, and DPLM2 task outputs now keep the Transformers task +prefix: `loss` when labels are present, then `logits`, `hidden_states` when +requested, and `attentions` when requested. FastPLMs diagnostics such as +`s_max` follow those standard fields. Masked-LM outputs that expose +`last_hidden_state` place it after the diagnostic extension. Tuple output is +exactly `output.to_tuple()`, so disabled or unavailable fields, including an +unconfigured pooler, are omitted rather than represented by `None`. These +forwards also reject misspelled or unsupported keyword arguments instead of +silently discarding them. + +ESM3 `AutoModel` tuple output now begins with the standard base-model fields in +this order: `last_hidden_state`, `hidden_states` when requested, and +`attentions` when requested. Sequence, structure, function, residue, and other +multimodal outputs follow that prefix. Prefer named fields when consuming the +additional tracks. This is a deliberate v1 tuple-order correction for callers +that previously treated element zero as sequence logits. + +## ESMFold2 and structure dependencies + +Only the standard, fast, experimental cutoff 2025, and experimental fast +cutoff 2025 ESMFold2 variants remain supported. Dataset embeddings accept a +single protein chain and expose the learned width-256 representation. + +The full `ESMFold2` and `ESMFold2-Experimental-Cutoff2025` checkpoints have 48 +folding blocks and retain optional MSA conditioning. The two Fast checkpoints +have 24 folding blocks and were trained without MSA conditioning, so they +reject MSA-derived inputs rather than silently ignoring them. This includes +`ProteinInput.msa` and low-level MSA-derived features. Fast still supports the +declared multichain and multimolecule inputs when every protein chain uses +`msa=None`. This distinction follows the official model description in +[Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). + +`esmc_precision="auto"` resolves to BF16. FP8 is an explicit, experimental, +inference-only request and raises when the validated Transformer Engine path is +unavailable. General structure dependencies live in +`requirements/features/structure.in`; reporting and binder dependencies remain +separate. + +## Commands and validation tiers + +Initialize official references only when running compliance: + +```bash +git submodule update --init --recursive +python -m tools.remote --host user@gpu-host --identity /path/to/key --suite compliance +``` + +Routine pull requests use the offline CPU contract and static/source checks: + +```bash +python -m pytest tests/cpu -m cpu_contract -n auto --dist=loadscope \ + --durations=25 --junitxml=artifacts/junit/cpu-contract.xml +PYTHONPATH=src python -m tools.artifacts.generate_docs --check +``` + +The routine `check` tier consumes immutable goldens and does not build live +official references. Exact-device Hopper/SM90 golden smoke, nightly +kernel/throughput work, and the +frozen-head compliance release candidate remain separate cost tiers. diff --git a/docs/models.md b/docs/models.md index ac9ff17..80fb409 100644 --- a/docs/models.md +++ b/docs/models.md @@ -1,538 +1,397 @@ -# Per-Model Guides +# Models -This document covers each model family supported by FastPLMs: loading, configuration, special handling, and available checkpoints. +The generated [support matrix](generated/support.md) is the current list of +families, checkpoints, AutoClasses, backends, precisions, licenses, and release +tiers. It is produced from `src/fastplms/models.toml`; edit the manifest, not the +table. -Most sequence models (ESM2, ESM++, E1, DPLM, DPLM2, ANKH) share the same embedding pipeline via `EmbeddingMixin`. ESM3 exposes its own compatible `embed_dataset()` method for sequence embeddings. They support most attention backends, with these exceptions: ANKH supports only `sdpa` and `flex`, and ESM3 supports `sdpa` and `flex`. Structure prediction models (Boltz2, ESMFold, ESMFold2, and ESMFold2-Fast) have their own APIs. +## Dependencies and platform requirements -Experimental test-time training (TTT) is available for ESM2, ESM++, ESM3, E1, -DPLM, DPLM2, ANKH, FastESMFold, and ESMFold2. It is disabled by default and is -only activated by explicit calls such as `model.ttt(...)`, -`fold_protein(..., ttt=True)`, or `fold_protein_ttt(...)`. TTT trains small -LoRA adapters with a masked language modeling objective on the test protein. It -can improve difficult or low-confidence cases, but it adds test-time compute and -can degrade already strong predictions. Treat it as experimental. +FastPLMs 1.0 sequence models require Python 3.11-3.14, PyTorch 2.13, and +Transformers 5.13. Install those dependencies directly. The runtime source is +loaded from the pinned Hugging Face model repository: ---- - -## ESM2 - -**Organization:** Meta AI -**Architecture:** Transformer encoder with rotary position embeddings (RoPE) -**Checkpoints:** 8M, 35M, 150M, 650M, 3B - -### Loading - -```python -from transformers import AutoModelForMaskedLM, AutoConfig - -# Default (SDPA backend) -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) - -# With a specific backend -config = AutoConfig.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -config.attn_backend = "flex" -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ESM2-150M", config=config, trust_remote_code=True) +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" ``` -### Key Details - -- Uses the standard ESM tokenizer (`EsmTokenizer` from `transformers`) -- Tokenizer accessible via `model.tokenizer` -- Backend can be set on the config before `from_pretrained` OR via the mutable `model.attn_backend` property after load (same mechanism as every other family). -- Pre-LayerNorm architecture with a final `emb_layer_norm_after` -- Supports `output_attentions=True` and `output_hidden_states=True` -- Experimental TTT is opt-in via `model.ttt(seq=...)`; no LoRA adapters are injected during normal inference. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| ESM2-8M | `Synthyra/ESM2-8M` | `facebook/esm2_t6_8M_UR50D` | -| ESM2-35M | `Synthyra/ESM2-35M` | `facebook/esm2_t12_35M_UR50D` | -| ESM2-150M | `Synthyra/ESM2-150M` | `facebook/esm2_t30_150M_UR50D` | -| ESM2-650M | `Synthyra/ESM2-650M` | `facebook/esm2_t33_650M_UR50D` | -| ESM2-3B | `Synthyra/ESM2-3B` | `facebook/esm2_t36_3B_UR50D` | +Eager and SDPA are portable CPU or CUDA paths. Optimized backends and structure +families have additional dependency, dtype, CUDA, and platform requirements; +check the generated support matrix and the relevant family guide before +selecting one. ---- +## Shared loading contract -## ESM++ (ESMC) - -**Organization:** Biohub -**Architecture:** Transformer encoder with configurable rotary embeddings (scaling, interleaving) -**Checkpoints:** Small (300M), Large (600M), 6B - -### Loading +Tokenizer-based sequence models load through the normal Transformers auto +classes: ```python from transformers import AutoModelForMaskedLM -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ESMplusplus_small", trust_remote_code=True) -``` - -### Key Details - -- Uses the ESM tokenizer (same as ESM2) -- **Requires `sequence_id`** parameter for batched inference: `sequence_id = attention_mask.to(dtype=torch.bool)` -- Uses `einops` for tensor reshaping operations -- Rotary embeddings support `interleaved` mode and `scale_base`/`scaling_factor` for dynamic scaling -- Backend can be set on the config before `from_pretrained` OR via the mutable `model.attn_backend` property after load. -- Experimental TTT is opt-in via `model.ttt(seq=...)`; it adapts the PLM backbone only. - -### Batched Forward Pass - -```python -tokenizer = model.tokenizer -tokenized = tokenizer(sequences, return_tensors="pt", padding=True) -tokenized = {k: v.to(device) for k, v in tokenized.items()} -tokenized["sequence_id"] = tokenized["attention_mask"].to(dtype=torch.bool) - -with torch.inference_mode(): - output = model(**tokenized) -``` - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| ESM++ Small (300M) | `Synthyra/ESMplusplus_small` | `biohub/ESMC-300M` | -| ESM++ Large (600M) | `Synthyra/ESMplusplus_large` | `biohub/ESMC-600M` | -| ESM++ 6B | `Synthyra/ESMplusplus_6B` | `biohub/ESMC-6B` | - -### Binder Design Role - -The FastPLMs binder design tutorial uses ESM++ as the masked-LM -pseudoperplexity regularizer on mutable binder residues. The verified EGFR run -uses `Synthyra/ESMplusplus_6B` by default, paired with FastPLMs ESMFold2 -experimental checkpoints for differentiable folding and final criticism. - -See [Binder Design Example](binder_design.md) for the full local and Modal -workflow, official selection metrics, and the verified EGFR 128 amino acid -minibinder result. - ---- - -## ESM3 - -**Organization:** Biohub -**Architecture:** Multimodal protein model over sequence, structure, and function tracks -**Checkpoints:** Open Small - -### Loading - -```python -import torch -from transformers import AutoModel - -model = AutoModel.from_pretrained( - "Synthyra/ESM3_small", +model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-150M", trust_remote_code=True, - dtype=torch.bfloat16, - device_map="cuda", + attn_implementation="sdpa", ) ``` -`AutoModelForMaskedLM` also resolves to the same ESM3 wrapper class, which returns sequence logits plus ESM3 track logits. - -### Key Details - -- Supports sequence-only inference by default via `input_ids` and `attention_mask`. -- Additional ESM3 tracks can be passed as tensors: `structure_tokens`, `ss8_tokens`, `sasa_tokens`, `function_tokens`, `residue_annotation_tokens`, `average_plddt`, `per_res_plddt`, `structure_coords`, `chain_id`, and `sequence_id`. -- Exposes `forward_sequence()` and `tokenize_sequences()` helpers for ergonomic sequence inference. -- Supports `embed_dataset()` with pooled `mean`, `cls`, and `max` embeddings, plus residue-wise embeddings through `full_embeddings=True`. -- Supports `sdpa` and `flex` attention backends. -- Includes the Biohub ESM MIT license in the Hub `LICENSE` file and model card metadata. -- Experimental TTT is opt-in via `model.ttt(seq=...)` and uses `sequence_logits` only. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| ESM3 Small | `Synthyra/ESM3_small` | `biohub/esm3-sm-open-v1` | - ---- - -## E1 - -**Organization:** Profluent Bio -**Architecture:** Transformer with within-sequence and global (block-causal) attention layers -**Checkpoints:** 150M, 300M, 600M - -### Loading - -```python -from transformers import AutoModelForMaskedLM - -model = AutoModelForMaskedLM.from_pretrained("Synthyra/Profluent-E1-150M", trust_remote_code=True) -``` - -### Key Details - -- **Sequence mode**: E1 does not use a standard HuggingFace tokenizer. Tokenization is built into the model via `E1BatchPreparer` -- Uses RMSNorm (not LayerNorm) -- Grouped-query attention (num_key_value_heads < num_heads) -- Two attention layer types that alternate: - - `WITHIN_SEQ`: Attention within individual sequences only - - `GLOBAL`: Cross-sequence block-causal attention -- Separate RoPE configurations for within-sequence and global attention (different `rope_theta`) -- KV caching via `DynamicCache` for efficient generation -- Backend can be set on the config before `from_pretrained` OR via the mutable `model.attn_backend` property after load. -- Experimental TTT is opt-in via `model.ttt(seq=...)`; the E1 path carries `input_ids`, `within_seq_position_ids`, `global_position_ids`, and `sequence_ids`. - -### Tokenization (Sequence Mode) - -E1's tokenization happens via `model.model.prep_tokens`: - -```python -batch_kwargs = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) -# Returns dict with: -# input_ids, within_seq_position_ids, global_position_ids, -# sequence_ids, labels, context, context_len -``` - -For the `embed_dataset()` pipeline, pass `tokenizer=None` and the mixin handles E1's sequence mode automatically. - -### Embedding Extraction - -```python -embeddings = model.embed_dataset( - sequences=["MKTLLILAVVAAALA", "MALWMRLLPLLALL"], - batch_size=2, - tokenizer=None, # E1 uses sequence mode - pooling_types=["mean"], - save=False, -) -``` - -### MSA Context And PPLL - -FastPLMs exposes E1 MSA context utilities directly on the model object: - -```python -a3m_path = model.search_homologues( - sequence="MALWMRLLPLLALLALWGPDPAAA", - output_dir="msas", - provider="colabfold", -) - -scores = model.score_ppll( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - a3m_path=a3m_path, - ensemble=True, -) - -embeddings = model.embed_with_msa( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - a3m_path=a3m_path, - pooling_types=["mean"], -) -``` - -MSA parsing and context sampling match Profluent's official E1 `msa_sampling` behavior. `score_ppll()` intentionally differs from the official `E1Scorer`: FastPLMs reports mean correct-token probability for each scored sequence and optionally averages across sampled contexts, rather than computing mutant deltas against a parent sequence. This is much cheaper and is the preferred scoring path here. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| E1-150M | `Synthyra/Profluent-E1-150M` | `Profluent-Bio/E1-150m` | -| E1-300M | `Synthyra/Profluent-E1-300M` | `Profluent-Bio/E1-300m` | -| E1-600M | `Synthyra/Profluent-E1-600M` | `Profluent-Bio/E1-600m` | - -### Compliance Dependencies - -E1 compliance tests require the official E1 package: - -```bash -pip install E1 @ git+https://github.com/Profluent-AI/E1.git -``` - -This is pre-installed in the Docker image. - ---- - -## DPLM - -**Organization:** ByteDance -**Architecture:** Diffusion-optimized transformer based on ESM2 architecture -**Checkpoints:** 150M, 650M, 3B - -### Loading - -```python -from transformers import AutoModelForMaskedLM - -model = AutoModelForMaskedLM.from_pretrained("Synthyra/DPLM-150M", trust_remote_code=True) -``` - -### Key Details - -- Uses the ESM tokenizer (same as ESM2) -- Backend can be set on the config before `from_pretrained` or via the mutable `model.attn_backend` property after load. -- Architecture extends `EsmConfig` and `EsmPreTrainedModel` from HuggingFace -- Supports cross-attention and KV caching for generation -- `ModifiedEsmSelfAttention` extends the official `EsmSelfAttention` with multi-backend support -- Experimental TTT is opt-in via `model.ttt(seq=...)`; normal DPLM inference and diffusion APIs are unchanged. - -### Post-Load Backend Switching +Each manifest entry declares its valid AutoClasses. Unsupported class or +attention combinations fail explicitly. The default attention implementation +is left unspecified so Transformers can select its standard SDPA path. + +## Sequence model families + +### ESM2 + +FastPLMs preserves ESM2 tokenization, encoder outputs, masked-language-model +head, contact head, tied weights, and checkpoint keys. Eager, SDPA, Flex +Attention, and both revision-pinned Hugging Face FlashAttention kernels are +tested against the pinned checkpoint and Transformers references. +Plain `AutoModel` omits the optional ESM pooler because the published +masked-language-model checkpoints contain no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training +that head. + +### ESM++ and ESMC + +The ESM++ family provides the ESMC sequence encoders with Hugging Face auto +loading, shared embeddings, and residue-aware pooling. It is also the language +model used by ESMFold2. The model records the resolved attention implementation +and rejects an unavailable requested kernel. + +ESMC follows the pinned Biohub mask precedence. When `sequence_id` is supplied, +it is authoritative: non-negative values identify chains and `-1` identifies +padding, while `attention_mask` is ignored. Without `sequence_id`, +`attention_mask` is the ordinary padding mask and defaults from the tokenizer +padding ID. Callers that need both chain isolation and padding must encode both +in `sequence_id`; the two public masks are not intersected. + +Exact semantic configuration, tokenizer, state, alias, and SDPA contracts are +validated against the pinned Biohub implementation. SDPA is the recommended +highest-fidelity path. Flex Attention and FlashAttention 3 are supported, +non-experimental backends with diagnostic numerical-deviation warnings rather +than strict parity gates. Every checkpoint card exposes the required relative +L2, Q99.9, residue-cosine, pooled-cosine, top-1, and Jensen-Shannon table. Cells +remain explicitly pending until a frozen-head report from the exact +GH200/aarch64 validation target for the backend, dtype, software stack, and +sequence panel is attached. H100 and H200 remain supported Hopper-class +devices, but measurements from them are not interchangeable with or accepted +as the current GH200 release evidence. + +### ESM3 + +ESM3 retains sequence, structure, and function tracks and generation helpers +without importing the official checkout at runtime. Feature tests cover encoding, +multimodal inputs, forward outputs, and seeded generation. Dataset embedding +uses the sequence representation and excludes non-protein track tokens. +With `return_dict=False`, the standard base-model prefix is +`(last_hidden_state, hidden_states, attentions)` with disabled optional fields +omitted; multimodal logits and extensions follow. Named output fields are the +recommended interface for individual tracks. + +### DPLM + +DPLM generation starts from a tokenized sequence whose biological positions +define the requested length. The sampler replaces those positions with the mask +token, predicts all positions, retains the most confident predictions, and +repeats until no masks remain. The optional `partial_masks` argument is a +boolean tensor with the same shape and device as `input_tokens`; `True` marks +positions that must remain fixed. ```python -model = AutoModelForMaskedLM.from_pretrained("Synthyra/DPLM-150M", trust_remote_code=True) -model.attn_backend = "flex" # Updates every attention layer in-place -``` - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| DPLM-150M | `Synthyra/DPLM-150M` | `airkingbd/dplm_150m` | -| DPLM-650M | `Synthyra/DPLM-650M` | `airkingbd/dplm_650m` | -| DPLM-3B | `Synthyra/DPLM-3B` | `airkingbd/dplm_3b` | - ---- - -## DPLM2 - -**Organization:** ByteDance -**Architecture:** Multimodal diffusion transformer handling both sequence and structure tokens -**Checkpoints:** 150M, 650M, 3B - -### Loading +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer -```python -from transformers import AutoModelForMaskedLM +checkpoint = "Synthyra/DPLM-150M" +tokenizer = AutoTokenizer.from_pretrained(checkpoint) +model = AutoModelForMaskedLM.from_pretrained( + checkpoint, + trust_remote_code=True, +).cuda().eval() -model = AutoModelForMaskedLM.from_pretrained("Synthyra/DPLM2-150M", trust_remote_code=True) +X = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() +with torch.inference_mode(): + generated_tokens = model.generate(X, max_iter=100) +sequence = tokenizer.decode(generated_tokens[0], skip_special_tokens=True).replace(" ", "") ``` -### Key Details - -- Uses the ESM tokenizer -- **Multimodal input**: Handles both amino acid tokens and structure tokens in packed sequences -- Mutable `model.attn_backend` property (same as DPLM) -- Special token normalization: `_normalize_dplm2_input_ids()` maps tokens above vocab_size back into range -- Packed multimodal layout detection: `_has_packed_multimodal_layout()` checks if input_ids contain interleaved AA and structure tokens -- The official DPLM2 has an extra `contact_head` not present in the FastPLM version, so weight compliance testing is skipped for this family -- Experimental TTT is opt-in via `model.ttt(seq=...)`; it adapts only LoRA weights on the PLM backbone. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| DPLM2-150M | `Synthyra/DPLM2-150M` | `airkingbd/dplm2_150m` | -| DPLM2-650M | `Synthyra/DPLM2-650M` | `airkingbd/dplm2_650m` | -| DPLM2-3B | `Synthyra/DPLM2-3B` | `airkingbd/dplm2_3b` | +`sampling_strategy` accepts `gumbel_argmax`, `argmax`, or `vanilla`. The +official 500-step schedule is used when `max_iter` is omitted. A shorter +schedule reduces latency but changes the sampling process. ---- - -## ANKH - -**Organization:** Elnaggar Lab -**Architecture:** T5-style encoder with bidirectional gated GELU FFN and learned relative position bias (bucketed) -**Checkpoints:** Base, Large, ANKH2-Large, ANKH3-Large, ANKH3-XL - -### Loading +DPLM advertises eager, SDPA, Flex Attention, and the precompiled Hugging Face +kernels implementation of FlashAttention 3. The official BF16 inference path +loads FP32-resident parameters and enters an explicit CUDA BF16 autocast +context: ```python -from transformers import AutoModelForMaskedLM, AutoConfig - -# Default (SDPA) -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True) - -# Flex backend (block-mask aware) -config = AutoConfig.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True) -config.attn_backend = "flex" -model = AutoModelForMaskedLM.from_pretrained("Synthyra/ANKH_base", config=config, trust_remote_code=True) -``` - -### Key Details - -- Uses the checkpoint-matched ANKH T5 tokenizer exposed through each Synthyra checkpoint -- Tokenizer accessible via `model.tokenizer` -- Backend can be set on the config before `from_pretrained` OR via the mutable `model.attn_backend` property after load (same mechanism as every other family). -- **Attention is unscaled** (no `1/sqrt(d_kv)` factor). T5 trains without scaling; the learned relative position bias absorbs the temperature. -- Only `sdpa` and `flex` are supported. Requesting `kernels_flash` silently falls back to `flex` (or `sdpa` if flex is unavailable) because flash kernels can't accept additive position bias. -- Layer 0 owns the relative-position-bias `nn.Embedding`; subsequent layers receive the precomputed bias through the forward call. -- The native ANKH checkpoint is a T5 encoder-decoder; FastPLMs uses the encoder only and bolts on a separate `lm_head` for the `ForMaskedLM` variant. For weight-parity comparisons against `transformers.T5EncoderModel`, the FastPLMs `lm_head.weight` is allowlisted as an expected extra parameter. -- ANKH3 checkpoints use 256-token tokenizers, while ANKH v1/v2 checkpoints use 144-token tokenizers. Use the checkpoint tokenizer through `model.tokenizer` or `AutoTokenizer.from_pretrained()`. -- Experimental TTT is opt-in via `model.ttt(seq=...)`; the ANKH MaskedLM head is not pretrained for standard MLM, so results should be treated especially cautiously. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Official Reference | -|------------|----------------|-------------------| -| ANKH-base | `Synthyra/ANKH_base` | `ElnaggarLab/ankh-base` | -| ANKH-large | `Synthyra/ANKH_large` | `ElnaggarLab/ankh-large` | -| ANKH2-large | `Synthyra/ANKH2_large` | `ElnaggarLab/ankh2-ext2` | -| ANKH3-large | `Synthyra/ANKH3_large` | `ElnaggarLab/ankh3-large` | -| ANKH3-XL | `Synthyra/ANKH3_xl` | `ElnaggarLab/ankh3-xl` | - ---- - -## Boltz2 - -**Organization:** MIT / Various -**Architecture:** Diffusion-based structure prediction model -**Checkpoints:** Standard - -### Loading - -```python -from transformers import AutoModel - -model = AutoModel.from_pretrained("Synthyra/Boltz2", trust_remote_code=True, dtype=torch.float32) -``` - -### Key Details - -- **Structure prediction model**, not a sequence encoder. Does not inherit from `EmbeddingMixin` -- Uses `AutoModel` (not `AutoModelForMaskedLM`) -- Custom featurization pipeline via `minimal_featurizer.build_boltz2_features()` -- Outputs atomic coordinates, pLDDT, pTM, ipTM confidence scores -- Can export predictions as CIF files via `model.save_as_cif(output, "pred.cif")` -- TTT is not supported for Boltz2 in FastPLMs. - -### Predict Structure +model = AutoModelForMaskedLM.from_pretrained( + checkpoint, + trust_remote_code=True, + attn_implementation="sdpa", + dtype=torch.float32, +).cuda().eval() -```python -output = model.predict_structure( - amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG", - recycling_steps=3, - num_sampling_steps=200, - diffusion_samples=1, -) -print(output.sample_atom_coords.shape) # (N_atoms, 3) -print(output.plddt.shape) # (N_residues,) +with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + output = model(X) ``` -### Compliance Testing - -Boltz2 compliance is tested via a standalone script (`testing/run_boltz2_compliance.py`) that compares coordinates, pairwise distances, and TM-scores against the official implementation. - ---- - -## ESMFold2 - -**Organization:** Biohub -**Architecture:** ESMC-backed diffusion structure predictor -**Checkpoints:** Full, Fast, Experimental, Experimental-Fast, Cutoff2025 variants - -### Loading +Static BF16 parameter storage is not an advertised DPLM precision path. The +official autocast contract is exact at every hidden state with SDPA and meets +the fixed engineering target with eager, Flex Attention, and FlashAttention 3. +Released DPLM checkpoints configure attention-probability dropout as zero. +For custom fine-tuning configurations with a nonzero value, FastPLMs applies +that dropout in eager and SDPA training calls. Flex Attention and +FlashAttention 3 fail closed for nonzero training dropout because those paths +do not implement the declared stochastic contract. Evaluation always uses +zero attention dropout. Decoder cross-attention supports eager and SDPA and +fails closed when another backend is requested. +Plain DPLM and DPLM2 `AutoModel` loads likewise omit the optional untrained ESM +pooler by default; it remains available through explicit +`add_pooling_layer=True`. + +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The pinned ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) +and [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +provide the immutable license basis; the latter explicitly scopes the official +repository release to the pretrained weights for both model families. + +### DPLM2 + +DPLM2 applies the same confidence-based unmasking separately to amino-acid and +structure tokens. Co-generation input X contains two equal-length tracks, +including their boundary tokens, in the official order: structure first and +amino acids second. The output mapping contains `output_tokens` with the same +shape as X. ```python import torch -from transformers import AutoModel +from transformers import AutoModelForMaskedLM, AutoTokenizer -model = AutoModel.from_pretrained( - "Synthyra/ESMFold2-Fast", +checkpoint = "Synthyra/DPLM2-150M" +tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) +model = AutoModelForMaskedLM.from_pretrained( + checkpoint, trust_remote_code=True, - dtype=torch.float32, -).eval().cuda() -``` - -### Key Details - -- Uses `AutoModel`, not `AutoModelForMaskedLM`. -- Can fold single chains and multi-chain complexes. -- Loads the embedded PLM through `Synthyra/ESMplusplus_6B` by default, using - FastPLMs ESM++ with `config.esmc_attn_backend = "flex"`. -- Exposes `fold_protein()`, `fold()`, `prepare_structure_input()`, - `result_to_cif()`, and `result_to_pdb()`. -- Experimental checkpoints support `res_type_soft` for differentiable binder - sequence optimization. -- Final scoring can report pTM, iPTM, pLDDT, structures, and distogram logits. -- `set_kernel_backend()`, `set_chunk_size()`, and `apply_torch_compile()` are - available on the ESMFold2 wrappers. -- Experimental TTT is opt-in via `fold_protein(..., ttt=True)` or - `fold_protein_ttt(...)`; it trains LoRA adapters only on `_esmc`, not the - folding trunk, confidence head, or diffusion head. - -### Available Checkpoints - -| Checkpoint | HuggingFace ID | Use | -|------------|----------------|-----| -| ESMFold2 | `Synthyra/ESMFold2` | Full released structure predictor | -| ESMFold2-Fast | `Synthyra/ESMFold2-Fast` | Faster released structure predictor | -| ESMFold2-Experimental-Fast | `Synthyra/ESMFold2-Experimental-Fast` | Binder inversion and hero critic | -| ESMFold2-Experimental-Fast-Cutoff2025 | `Synthyra/ESMFold2-Experimental-Fast-Cutoff2025` | Binder inversion and hero critic | -| ESMFold2-Experimental | `Synthyra/ESMFold2-Experimental` | Final hero critic | -| ESMFold2-Experimental-Cutoff2025 | `Synthyra/ESMFold2-Experimental-Cutoff2025` | Final hero critic | - -### Binder Design Example - -The FastPLMs binder workflow lives at -`cookbook/tutorials/binder_design_fastplms.py` and supports local CUDA or Modal -execution. The optimizer follows the official ESM strategy: continuous amino acid -logits, cysteine suppression, ESMFold2 distogram structure losses, ESM++ -pseudoperplexity, late-trajectory iPTM selection, and final critic ranking. - -```bash -python cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir binder_design_egfr_len128_seed103 -``` - -The verified EGFR result had hero mean iPTM `0.913870`, hero min iPTM -`0.904600`, and all four hero critics above `0.9`. - -Binder sequence: - -```text -SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ -``` - -See [Binder Design Example](binder_design.md) for output files, per-critic -metrics, pI filtering, optional scaling critics, and caveats. - ---- - -## ESMFold - -**Organization:** Meta AI (wrapped by Synthyra) -**Architecture:** ESM2 backbone + structure module with optional experimental Test-Time Training (TTT) -**Checkpoints:** Standard - -### Loading +).cuda().eval() +vocab = tokenizer.get_vocab() +l = 64 +structure = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +amino_acids = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +X = torch.tensor([structure + amino_acids], device="cuda") -```python -from transformers import AutoModel - -model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True, dtype=torch.float32) +with torch.inference_mode(): + generated_tokens = model.generate(X, max_iter=100)["output_tokens"] ``` -### Key Details - -- Inherits from `transformers.EsmForProteinFolding` with the ESM2 backbone replaced by `FastEsmBackbone` -- Supports all attention backends via `config.attn_backend` -- TTT is disabled by default and must be requested with `ttt=True` or `fold_protein_ttt(...)` -- Optional experimental TTT adapts the ESM2 backbone via LoRA + masked LM before folding -- TTT can improve low-confidence sequences, but it adds compute and may degrade already strong predictions - -### Structure Prediction +The DPLM2 tokenizer preserves separate amino-acid and structure boundary, +unknown, and mask tokens. Generic `cls_token`, `eos_token`, `mask_token`, and +`unk_token` aliases are intentionally unset, so multimodal tracks must add the +corresponding ``/`` or ``/`` tokens +explicitly and tokenize them with `add_special_tokens=False`. +`model.embed_dataset(...)` and amino-acid TTT still accept raw sequences: the +model adapter adds `` and ``, invokes this exact tokenizer with +`add_special_tokens=False`, and excludes those boundaries from residue outputs. + +The default sampler uses `annealing@2.0:0.1` token sampling and +`stochastic1.0` confidence selection. `argmax` with `deterministic` unmasking +provides a deterministic compliance path. The checkpoint sets +`tie_word_embeddings=False`, so input and output embeddings are intentionally +distinct. The trained contact head remains under `esm.contact_head`. + +DPLM2 advertises SDPA only. Its BF16 inference contract keeps checkpoint +parameters in FP32 and evaluates the forward pass under CUDA BF16 autocast. +Loading static BF16 parameters for evaluation raises before inference. Eager, +Flex Attention, FlashAttention 2, and FlashAttention 3 requests fail explicitly +because their representative deep hidden-state comparisons miss the release +engineering target. + +### E1 + +E1 has no tokenizer dependency. Its dedicated adapter accepts raw protein +sequences and preserves official boundary-token, context, and retrieval-augmented +generation preparation. Launches display `Profluent-E1` as required by the +upstream agreement. E1 legal files and modified-file notices are distributed +with relevant artifacts and containers. E1 advertises SDPA and Flex Attention; +its eager path is not advertised because it misses the pinned output contract. + +MSA-aware embedding returns the same ordered, duplicate-preserving +`EmbeddingResult` and uses the same safetensors or SQLite persistence as +ordinary dataset embedding. Record IDs are zero-based input positions, and +`max_len` is measured in biological residues. `matrix_embed=True` selects full +residue output. Local A3M input is deterministic and offline. Homology search +and Hub MSA download are separate, networked acquisition steps and are never +triggered by an offline embedding call. + +Local MMseqs2 search defaults to the official multi-architecture CPU image +`ghcr.io/soedinglab/mmseqs2:18-8cc5c` pinned to manifest digest +`sha256:41b12b0d5f41432fa1b9976123da6e2e06e7fab49a34964f3b54ec038e5845d9`. +It never pulls implicitly. The container runs with `--network none`, each phase +has a bounded timeout, and every local image inspection must match the requested +repository digest, Linux platform, host architecture, and a valid image ID. ```python -# Without TTT -with torch.no_grad(): - output = model.infer("MKTLLILAVVAAALA") -pdb_string = model.output_to_pdb(output) - -# With experimental TTT (default: 10 optimizer steps) -result = model.fold_protein("MKTLLILAVVAAALA", ttt=True) -# result = {"plddt": float, "ptm": float, "pdb_string": str, ...} +from fastplms.models.e1.retrieval import HomologueSearcher + +searcher = HomologueSearcher( + target_db="databases/uniref30", + use_gpu=False, + allow_pull=False, + allow_network=False, + phase_timeout=1800, + target_db_identity="uniref30-release-2025-02", +) +a3m_path = searcher.search("MSTNPKPQRKTKRNT", "msa-results", seq_id="query-1") ``` -### TTT Defaults - -TTT is disabled by default. Standard `fold_protein(...)` is a baseline fold and -returns `best_step=0`. You can also call `model.infer(...)` directly for raw -ESMFold outputs. Use `model._ttt_cfg` to tune optimizer steps, LoRA rank, and -masking parameters before calling `fold_protein(..., ttt=True)`. - -```python -model._ttt_cfg.steps = 3 -result = model.fold_protein_ttt("MKTLLILAVVAAALA") -``` +The database and output must resolve beneath the current working directory; +symlink escapes are rejected before Docker runs. Successful searches write +`search-provenance.json` beside the A3M with the image version, manifest digest, +local image ID, platform, database identity, parameters, and sequence hash. +The sidecar also records the A3M size and SHA-256; cached output is reused only +when both source-record and result integrity checks match. `allow_pull=True` is an explicit +network acquisition opt-in. `allow_network=True` separately permits network +access inside the search container, which a local database search does not +require. + +GPU MMseqs2 is not selected automatically. The stable official CUDA image is +AMD64-only, so it is incompatible with GH200/ARM64. `use_gpu=True` requires a +caller-supplied, digest-pinned image that is compatible with the current host; +the CPU default fails closed instead of silently attempting GPU execution. + +### ANKH + +The ANKH 1.0 migration replaced every Synthyra ANKH repository with full +official-compatible T5 state. +`FastAnkhModel` and `AutoModel` load the encoder and shared embeddings cleanly +without decoder allocation, while `FastAnkhForConditionalGeneration` and +`AutoModelForSeq2SeqLM` load the encoder, decoder, cross-attention, and LM head +from the same repository. The larger full checkpoint changes the +default Hub contents while preserving encoder output parity. + +Encoder embeddings are the default and select the final state unless +`hidden_state_index` or `store_all_hidden_states` requests another view. +Decoder embeddings require exactly one explicit aligned `decoder_inputs` list +or `decoder_input_ids` tensor. FastPLMs does not shift the source implicitly, +because official ANKH tasks use prompts, sentinels, or generated tokens. The +decoder biological mask excludes start, EOS, padding, sentinel, and other +special tokens; persistence fingerprints stack, layer, decoder input, mask, and +alignment. + +The encoder is the representative throughput architecture and supports the +manifest-declared eager and SDPA attention implementations. Exact encoder and +sequence-to-sequence weights, aliases, seeded inference, and save/reload must be +validated from the same artifact and new Hub revision before that revision is +advertised. Files-only publication is forbidden for the migration. +The previous synthesized masked-language-model head remains available only as +the separately named `FastAnkhForMaskedLMExtension`; it is a FastPLMs extension, +not an official equivalent. + +ANKH code and mirrored weights retain CC BY-NC-SA 4.0 terms. FastPLMs displays +those terms but does not enforce a runtime usage policy. + +## Structure model families + +### ESMFold + +ESMFold retains the official ESM2 language-model trunk, folding trunk, output +heads, and structure export. Structure compliance hashes prepared features, +uses seeded stochastic inputs, checks geometry and finite values, and compares +coordinates and confidence outputs with the pinned reference. + +The structure-only ESM2 backbone omits the masked-LM and contact-regression +heads because folding consumes hidden states only. Both independently +implemented checkpoint transforms declare and test those omissions. Reported +pLDDT remains on the conventional `(0, 100)` scale; compliance normalizes it to +`(0, 1)` before computing mean absolute error. +For multimer inputs, summary mean pLDDT excludes synthetic linker residues and +includes only biological residues from the requested chains. + +The folding checkpoint remains in FP32 parameter storage. CUDA BF16 inference +enters autocast around the folding operation; loading the checkpoint itself as +static BF16 is not the declared compliance path. + +### ESMFold2 + +Supported variants are restricted to: + +| Official checkpoint | Folding blocks | MSA conditioning | Intended path | +| --- | ---: | --- | --- | +| `biohub/ESMFold2` | 48 | Optional | Full sequence or complex inference, including MSA-conditioned requests | +| `biohub/ESMFold2-Fast` | 24 | None; MSA-derived inputs are rejected | Inference-optimized single-sequence use | +| `biohub/ESMFold2-Experimental-Cutoff2025` | 48 | Optional | Experimental-cutoff full inference, including MSA-conditioned requests | +| `biohub/ESMFold2-Experimental-Fast-Cutoff2025` | 24 | None; MSA-derived inputs are rejected | Experimental-cutoff, inference-optimized single-sequence use | + +The Fast distinction is architectural, not merely a speed label. Biohub's +[Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf) describes Fast as a +model with 24 folding blocks trained without MSA conditioning for +single-sequence inference, compared with 48 folding blocks in the full model. +The Fast variants are not necessarily single-chain-only: supported multichain +and multimolecule requests remain available, but each protein chain uses +single-sequence mode. Fast variants reject MSA-derived inputs. Use a full variant +whenever optional MSA conditioning is part of the request. + +All four expose the learned ESMC projection and the `auto`, `bf16`, `fp32`, and +`fp8` ESMC precision policy. The manifest marks `fp8` as experimental; it is an +explicit inference-only opt-in rather than a release numerical-parity claim. +See [ESMFold2](esmfold2.md) for the exact embedding, reload, and folding contracts. + +The ESMFold2 folding checkpoint remains FP32 and folding computation uses CUDA +BF16 autocast. Its ESMC backbone is governed independently by the requested +ESMC precision, so selecting BF16 or FP8 ESMC does not change folding-parameter +storage. + +### Boltz2 + +Boltz2 accepts a raw amino-acid sequence through its protein helper or prepared +model features through its lower-level interface. It preserves trunk, +diffusion, confidence, and export behavior. Its larger scientific dependency +set is isolated in `requirements/features/structure.in` and the structure +candidate image. Chemistry and plotting dependencies are not part of the core +runtime requirements. + +Boltz2 retains FP32 parameters and runs supported CUDA BF16 structure inference +inside autocast. Static BF16 parameter loading is not its declared compliance +or artifact-validation path. + +`predict_structure(..., seed=...)` owns a scoped Python, NumPy, CPU Torch, and +CUDA RNG context and restores caller state on return. Prepared features and +parameters remain FP32; the supported CUDA compute path enters BF16 autocast +inside that scope. `seed` accepts a Python `int` or `None`; booleans, floats, +strings, NumPy integer scalars, and other coercible values are rejected before +any RNG state is read or changed. + +Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared +inference-core state, feature preparation, and seeded execution remain covered, +but native-environment BF16 end-to-end inference currently exceeds the fixed +numerical-equivalence limits. FastPLMs does not yet claim official inference +equivalence for Boltz2. Its ongoing structure tests remain available without +blocking the ESM++ and ESMFold2 release gates. + +## Test-time training + +ProteinTTT-derived adaptation is opt-in and covered as a feature, not a default +loading behavior. It never runs during model construction or ordinary inference. +ESMFold2 reloads ESMC in BF16 before a gradient-enabled path. See +[test-time training](ttt.md). + +## Adding a checkpoint + +Before code changes, capture the official configuration, tokenizer files and +behavior, state-key and alias schema, outputs, source revision, environment, and +licenses. Add the immutable checkpoint identities and conversion record to the +manifest. Then add official-generated goldens, a live reference case, artifact +loading, and all feature tests declared for the family. A selectable backend +must have an explicit implementation, failure behavior, and documented +numerical boundary. Only backends that meet the applicable release thresholds +may be described as parity paths. diff --git a/docs/testing.md b/docs/testing.md index ef2e03b..686dd06 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -1,287 +1,526 @@ -# Testing & Benchmarking - -FastPLMs uses pytest with Docker for all GPU testing. Tests cover model loading, attention backend consistency, weight/forward parity against official implementations, embedding stability, and throughput benchmarking. - -**Requires PyTorch 2.11+**. Flex attention uses Flash Attention 4 (FA4) as its backend on Hopper/Blackwell GPUs. The Dockerfiles pin PyTorch 2.11.0 with CUDA 12.8. - -## Docker Layout - -Two Docker layouts are supported. - -### Per-family images (recommended for parity / compliance work) - -A shared base image plus one image per model family. Each family image installs only that family's native reference deps, so we can run e.g. ESM++ tests against Biohub's `esm` package without breaking ESM2 tests that depend on `fair-esm` / `transformers.EsmForMaskedLM`. - -| Image tag | Native reference package | -|-----------|---------------------------| -| `fastplms-base` | none (torch 2.11.0, transformers 4.57.6, FastPLMs source, shared deps) | -| `fastplms-esm2` | uses `transformers.EsmForMaskedLM` | -| `fastplms-esm_plusplus` | Biohub `esm` runtime deps + `official/esm` submodule on `sys.path`. The `esm` package itself is **not** pip-installed (it depends on a Biohub `transformers` fork). | -| `fastplms-esm3` | Biohub ESM3 runtime deps + `official/esm` submodule on `sys.path`; requires gated source-model access for official ESM3 parity/compliance. | -| `fastplms-e1` | `pip install -e /app/official/e1` | -| `fastplms-dplm` | uses `transformers.EsmForMaskedLM` (DPLM's native package conflicts with our torchtext pin) | -| `fastplms-dplm2` | none beyond base | -| `fastplms-ankh` | uses `transformers.T5EncoderModel` | -| `fastplms-esmfold2` | Biohub `transformers` fork, ESMFold2 runtime deps, and structure export deps | - -Build: +# Testing and compliance + +FastPLMs treats official equivalence as a release property. Routine goldens make +development faster, but a release requires live comparison with the pinned +official implementation in its native reference container. + +GPU release and benchmark suites run through repository Docker images on the +current exact NVIDIA GH200/aarch64 workstation. H100 and H200 remain +Hopper-class deployment examples, but they are not current release-confirmation +evidence. Portable unit and documentation checks can also run locally. Every +Dockerized PyTorch run uses `--ipc=host` directly or receives `ipc: host` from +Compose. Remote preflight records the GH200 GPU UUID and native OCI platform, +then rejects platform or device drift between preflight and the loaded images. + +## Run tiers + +| Tier | Purpose | +| --- | --- | +| `cpu_contract` | Required offline CPU confidence gate with tiny models, no checkpoints, network, Docker, skips, or xfails | +| `check` | Candidate-only units, imports, local integration, release checks, and immutable checkpoint goldens; no artifacts, live references, or kernel downloads | +| `gpu-golden-smoke` | Conditional exact-device comparison with checked-in sequence and structure goldens; no live reference build | +| `compliance` | Every checkpoint whose manifest declares the release compliance tier against its live pinned official implementation | +| `structure` | ESMFold, four ESMFold2 variants, provisional Boltz2 diagnostics, feature preparation, export, and seeded stochastic output | +| `feature` | DPLM generation, DPLM2 generation, ESM3 multimodal generation, TTT, E1 sequence and RAG adapters, binder flow, pooling, and conversion | +| `artifact` | Fresh offline remote-code loading and save-reload for every local artifact | +| `benchmark` | Separate GH200/aarch64 latency, throughput, padding, memory, and exact-device regression suite | +| `python-matrix` | Isolated repository-source smokes with runtime dependencies on Python 3.11-3.14 | + +Routine `check` consumes goldens and never builds an official reference image. +Live references are reserved for the frozen exact-head `compliance` release +candidate. Missing expected +dependencies, checkpoints, reference containers, or backends are failures, not +skips. + +Boltz2 is intentionally outside the FastPLMs 1.0 `check` and `compliance` +claims while its native-environment BF16 numerical gap remains under +investigation. Its exact state/configuration, feature, seeded-execution, +artifact, and benchmark diagnostics remain in the focused tiers. This is an +explicit provisional boundary, not a relaxed tolerance or silent skip. + +## Required offline CPU gate + +Run this positive allowlist on the validation workstation before merge: ```bash -git submodule update --init --recursive - -# Build base + every family image -./build_images.sh - -# Build a specific subset -./build_images.sh esm2 esm_plusplus +python -m pytest tests/cpu -m cpu_contract -n auto --dist=loadscope \ + --durations=25 --junitxml=artifacts/junit/cpu-contract.xml ``` -`build_images.sh` always builds `fastplms-base` first and then layers each requested family on top, with `--cache-from` chained so dep changes don't invalidate the base. - -### Monolithic image (legacy) - -The original `Dockerfile` (image tag `fastplms`) bundles everything compatible into a single image. Used by the broad test suites that don't need per-family isolation. +The command uses Python 3.12, CPU-only Torch 2.13, Transformers 5.13, four CPU +workers, fixed seeds and thread counts, hidden CUDA, and empty temporary caches. +It sets `HF_HUB_OFFLINE=1` and `TRANSFORMERS_OFFLINE=1`, guards sockets and Hub +download functions, and fails on a skip or xfail. Tiny models use one layer, +hidden width 8-16, two heads, and a mixed-padding batch of two. The suite budget +is under five minutes and 4 GiB, with an approximately ten-second per-test cap. +The controller samples its complete live Linux process tree from `/proc` every +50 ms and deduplicates PIDs before summing concurrent memory. Complete +`smaps_rollup` proportional-set-size evidence enforces the 4 GiB gate; if any +live process lacks PSS, the gate fails closed to the larger concurrent RSS sum +and records the reason. Per-worker `RUSAGE_SELF` and `RUSAGE_CHILDREN` maxima +remain in telemetry as a temporal upper-bound diagnostic, not as the budget +metric, because those maxima can occur at different times. +Prepare the exact CPU validation environment before running the gate: ```bash -git submodule update --init --recursive -docker build -t fastplms . +uv venv --python 3.12 +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu ``` -## Running Tests +This consumes the named CPU validation profile and constrains Torch and +Transformers to the release versions in +`requirements/constraints/validation.txt`. `--torch-backend cpu` routes Torch +to the CPU index. CUDA-only cuEquivariance and FP8 dependencies belong in +separate environments. -**Always pass `--ipc=host`** with PyTorch, otherwise multi-worker DataLoader and CUDA can deadlock. +The gate statically covers all 29 checkpoints and executes every advertised +AutoClass once per family, including forward/loss/backward, resize, tuple and +dictionary output, and save/reload. It also covers ANKH stack selection and +state views; backend masks, fallback warnings, fake Flash dispatch, E1 cache and +ANKH concurrency; ESMC diagnostics; sequence-family embeddings and persistence; +bounded disk-spooled generator and FASTA streaming; generation and TTT; PEFT; +injected-core structure and binder flows; publication security; and curated +offline documentation examples. -### Per-family parity / compliance +Run the focused quality and source checks from the same environment: ```bash -# ESM2 -- model_key in conftest.py is "esm2" -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm2 \ - python -m pytest /workspace/testing/test_parity.py -k esm2 -v - -# ESM++ -- model_key is "esmc" -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm_plusplus \ - python -m pytest /workspace/testing/test_parity.py -k esmc -v - -# ESM3 -- requires accepted access to biohub/esm3-sm-open-v1 for official parity -docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-esm3 \ - python -m pytest /workspace/testing/test_parity.py -k esm3 -v - -# Everything else -for fam in e1 dplm dplm2 ankh; do - docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-$fam \ - python -m pytest /workspace/testing/test_parity.py -k $fam -v -done +python -m ruff check src tests tools examples benchmarks +mapfile -t MYPY_TARGETS < tools/typing-critical-files.txt +python -m mypy --python-version 3.12 --ignore-missing-imports \ + --explicit-package-bases --follow-imports=silent "${MYPY_TARGETS[@]}" +PYTHONPATH=src python -m tools.artifacts.generate_docs --check +python -m pytest \ + tests/release/test_binder_example_contracts.py \ + tests/release/test_documentation.py \ + tests/release/test_dependency_contracts.py \ + tests/release/test_flash_source_policy.py \ + tests/release/test_manifest_readiness.py \ + tests/release/test_model_card_licenses.py \ + tests/release/test_optimized_mode_contracts.py \ + tests/release/test_production_source_boundary.py \ + tests/release/test_python_support_matrix.py \ + tests/release/test_static_typing_scope.py \ + tests/release/test_typing_no_regression_gate.py \ + tests/release/test_workflow_declared_inputs.py \ + -q +python tools/remote/runtime_import_closure.py \ + --source-root src/fastplms --requirements-root requirements ``` -### Broader suites in the monolithic image +The repository has no GitHub Actions workflows. Run lint, bounded strict +typing, generated-document checks, release contracts, repository-source smoke, +and runtime import closure directly on the workstation. Cross-version source +smokes, live official implementations, and GPU suites remain release-time +responsibilities through the explicit `python-matrix`, `check`, `compliance`, +and release suites. + +## Cost-controlled schedule + +- Before merge: offline CPU contracts and static/source checks. +- Conditional GH200/aarch64 smoke: when a relevant sequence or structure path + changes, run `gpu-golden-smoke` against the exact candidate head. Candidate + output is compared with checked-in goldens; no reference image is built. +- Extended workstation tier: sharded real-checkpoint goldens, + eager/SDPA/Flex execution, generation, PEFT, structure, artifacts, FP8, and + throughput by family. The GH200 job does not download, build, or execute + FA2/FA3 kernels; FA2 retains separate prior focused evidence and FA3 is + explicitly unavailable in the current linux/arm64 lock. +- Release candidate: every live pinned reference, checkpoint/state/tokenizer + contract, published artifact, structure panel, and benchmark on one frozen + exact head. + +GPU phases enforce explicit timeouts and cancellation, group all AutoClasses +for one checkpoint in one isolated process, build independent Buildx targets in +parallel, and publish JUnit, duration, cache, environment, and immutable report +telemetry. GH200/aarch64 benchmarking records cold compilation separately from warm +throughput; a missing baseline remains a release blocker rather than a synthetic +placeholder. All GH200 tiers are invoked directly with `python -m tools.remote`; +there is no hosted CI or scheduled accelerator workflow. + +Every remote report contains a structured kernel-capability record. It names +the measured eager, SDPA, and Flex backends, identifies the pinned FA2 revision +as prior-focused-evidence-only, identifies FA3 as unavailable on linux/arm64, +and records that network downloads and source builds were disabled. Asking the +GH200 runner to execute either Flash backend fails before source archiving. + +## Frozen ESMC release evidence + +The ESMC compliance run writes one schema-v3 JSON record for every combination +of three checkpoints, five BF16 attention backends, and two immutable sequence +panels. A release set therefore contains exactly 30 records under one explicit +directory: 18 measured eager, SDPA, and Flex records, plus 12 structured +FlashAttention 2 and 3 locked-platform unavailable records. All 30 records must +come from the current exact GH200/aarch64 +accelerator, repository container images, dependency lock, installed inventory, +and official source attestations. H100 and H200 remain supported Hopper-class +examples, but their measurements are not current release evidence and are not +combined with GH200 results. Candidate and official-reference measurements must +carry the same preflight hardware identity. + +Default documentation generation does not inspect the environment or discover +reports. It deliberately renders ESMC measurements as pending. On the frozen +release head, select the completed report directory explicitly: ```bash -# Fast tests (small models, no compliance, no structure) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "gpu and not slow and not large and not structure" -v - -# All sequence model tests except 3B -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "not large and not structure" -v - -# Full suite including 3B models (requires 40+ GB VRAM) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "not structure" -v - -# Structure models only (Boltz2, ESMFold, ESMFold2, ESMFold2-Fast) -docker run --gpus all --ipc=host fastplms python -m pytest /app/testing/ -m "structure" -v - -# Throughput benchmark (saves JSON/CSV/PNG) -docker run --gpus all --ipc=host -v ${PWD}:/workspace fastplms python -m pytest /app/testing/test_throughput.py -v -s - -# Throughput benchmark, standalone, more configurable -docker run --gpus all --ipc=host -v ${PWD}:/workspace fastplms python -m testing.throughput \ - --model_paths Synthyra/ESM2-8M Synthyra/ESMplusplus_small \ - --backends sdpa flex kernels_flash \ - --batch_sizes 2 4 8 \ - --sequence_lengths 64 128 256 512 1024 2048 - -# Interactive shell -docker run --gpus all --ipc=host -v ${PWD}:/workspace -it fastplms bash +PYTHONPATH=src python -m tools.artifacts.generate_docs \ + --source-root . \ + --esmc-report-root artifacts/diagnostics/esmc + +PYTHONPATH=src python -m tools.artifacts.generate_docs \ + --source-root . \ + --esmc-report-root artifacts/diagnostics/esmc \ + --check ``` -On Windows, replace `${PWD}` with `$(pwd)`. - -### Binder design workflow tests - -The FastPLMs binder design tutorial is tested in the ESMFold2 family image -because it combines FastPLMs ESMFold2 experimental folding with the ESM++ -masked-LM regularizer. +The release-evidence form may instead read +`FASTPLMS_DIAGNOSTIC_REPORTS`, falling back to the same repository-relative +location: ```bash -# Unit tests for prompt reproducibility, input validation, pI filtering, -# official-style selection scoring, and differentiable LM regularization. -docker run --rm -v /home/ubuntu/FastPLMs:/app -w /app fastplms-esmfold2 \ - python -m pytest /app/testing/test_binder_design_fastplms.py -m "not gpu" -v - -# CUDA dry run that writes trajectory, FASTA, results, selection, and structures -# with fake folding functions so it is fast but still checks artifact plumbing. -docker run --gpus all --rm -v /home/ubuntu/FastPLMs:/app -w /app fastplms-esmfold2 \ - python -m pytest /app/testing/test_binder_design_fastplms.py \ - -k tiny_design_dry_run_writes_outputs -v +FASTPLMS_DIAGNOSTIC_REPORTS=/validated/esmc-reports \ +PYTHONPATH=src python -m tools.artifacts.generate_docs \ + --source-root . \ + --require-esmc-release-evidence \ + --check ``` -The verified EGFR example in [Binder Design Example](binder_design.md) used this -focused test set: `11 passed, 2 deselected` for the non-GPU tests and `1 passed, -12 deselected` for the CUDA dry run. - -## Pytest Markers - -| Marker | Description | VRAM | -|--------|-------------|------| -| `gpu` | Requires CUDA GPU | Varies | -| `slow` | Loads two models simultaneously (compliance tests) | 2x model size | -| `large` | 3B parameter models | 24+ GB | -| `structure` | Structure prediction models (Boltz2, ESMFold, ESMFold2, ESMFold2-Fast) | 8+ GB | - -Use `-m` to filter and `-k` to select by name: +Supplying either evidence option is fail closed. Missing or extra files, +duplicate JSON keys, non-finite values, malformed schema, a stale manifest, +checkpoint, weights, source-tree, runtime, kernel, dtype, backend, or panel +identity, either missing or stale official-reference source attestation, an invalid +canonical self-digest, a failed measured-record catastrophe gate, a false Flash +execution claim, and any cross-device record set all stop generation. Successful ingestion records the exact device, +runtime and source identities, official import revision/tree/attestation, +aggregate ranges, and per-case minimum/median/maximum distributions in the +generated capability manifest and applicable checkpoint cards. It never +fabricates a value from a threshold or copies a measurement from another +checkpoint. Flash records contain no numerical metrics. They attest that the +current GH200/aarch64 lock fails closed without dispatch. Prior real +FlashAttention 2 execution was captured in separate workstation JUnit, but its +immutable report and environment attestation are not bundled in this +repository. It is not a current release distribution or numerical claim. + +## Portable remote execution + +The runner accepts connection data at invocation time, synchronizes an isolated +workspace, preserves external model and container caches, and retrieves test and +benchmark outputs: ```bash -# Only compliance tests -python -m pytest /workspace/testing/ -m slow -v - -# Exclude large models -python -m pytest /workspace/testing/ -m "not large" -v - -# Only a specific model family -python -m pytest /workspace/testing/ -k esm2 -v +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/private-key \ + --suite check ``` -## Test File Map - -| File | What it tests | Markers | -|------|---------------|---------| -| `test_parity.py` | **Rigorous parity** vs official implementations: tokenizer parity, bit-exact weight parity, per-layer relative-std hidden-state diff (fp32 + bf16) across `single`/`uniform`/`skewed` padding scenarios, padding-isolation (`[short]` alone vs `[short, long_]` padded), backend consistency, end-to-end `embed_dataset` pipeline parity. Runs per family in its own Docker image. | `gpu` | -| `test_automodel.py` | Model loading + forward pass validity (no NaN/Inf) | `gpu` | -| `test_backend_consistency.py` | SDPA, Flex, Flash backends produce equivalent predictions (>= 95% agreement) | `gpu` | -| `test_compliance.py` | Original (looser, bf16-only) weight/forward compliance against official implementations. Kept as a smoke layer; `test_parity.py` is the source of truth. | `slow`, `gpu` | -| `test_embedding_mixin.py` | NaN stability, batch-vs-single match, FASTA parsing, DPLM2 utilities | `gpu` | -| `test_binder_design_fastplms.py` | FastPLMs-only binder prompt reproducibility, input validation, ESM++ pseudoperplexity gradient, official selection metrics, pI filtering, and dry-run artifact writing | `gpu` for the dry run | -| `test_throughput.py` | Throughput benchmark across models/backends/batch sizes; saves JSON/CSV/PNG | `slow`, `gpu` | -| `test_structure_models.py` | Boltz2 and ESMFold loading + forward pass | `structure`, `slow`, `gpu` | - -Each test file has both **default registry** tests (one small model per family for fast CI) and **full registry** tests (all 21+ checkpoints with size-based markers). +Workstation names, identity paths, cache credentials, and secrets are not stored +in tracked files. The runner archive excludes Git metadata, known credential +names, private-key extensions, and ignored files. Before synchronization, every +initialized upstream must have a clean tracked tree and a `HEAD` equal to its +parent Git link. The runner records those revisions, the exact tracked-file +inventory, and a tree digest in `.fastplms-source-provenance.json`. Artifact +validation uses this record only in an extracted tree with no Git metadata and +rejects missing, added, or modified upstream files. -## Model Registries +Python 3.12 is the canonical GPU validation interpreter. Source compatibility +for Python 3.11, 3.13, and 3.14 is checked separately on the same workstation: -### Default Registry (`MODEL_REGISTRY`) - -Used by the base parametrized tests. One small model per family: - -| Key | Model | Family | -|-----|-------|--------| -| `esm2` | ESM2-8M | ESM2 | -| `esmc` | ESMplusplus_small | ESM++ | -| `esm3` | ESM3_small | ESM3 | -| `e1` | Profluent-E1-150M | E1 | -| `dplm` | DPLM-150M | DPLM | -| `dplm2` | DPLM2-150M | DPLM2 | -| `ankh` | ANKH_base | ANKH | - -### Full Registry (`FULL_MODEL_REGISTRY`) - -Used by the `test_full_*` parametrized tests. All checkpoints with `size_category`: - -| Category | Models | Marker | -|----------|--------|--------| -| `small` | ESM2-8M, ESM2-35M, E1-150M, DPLM-150M, DPLM2-150M | (none) | -| `medium` | ESM2-150M, ESMC-small, E1-300M, ANKH-base | `slow` | -| `large` | ESM2-650M, ESMC-large, ESM3-small, E1-600M, DPLM-650M, DPLM2-650M, ANKH-large, ANKH2-large, ANKH3-large | `slow` | -| `xlarge` | ESM2-3B, ESMC-6B, DPLM-3B, DPLM2-3B, ANKH3-xl | `large` | - -### Structure Registry (`STRUCTURE_MODEL_REGISTRY`) - -| Key | Model | -|-----|-------| -| `boltz2` | Synthyra/Boltz2 | -| `esmfold` | Synthyra/FastESMFold | -| `esmfold2` | Synthyra/ESMFold2 | -| `esmfold2_fast` | Synthyra/ESMFold2-Fast | - -## Parity Testing (`test_parity.py`) - -The parity suite is the source of truth for "FastPLMs matches the official implementation." It is intentionally strict: when it passes, FastPLMs and the native model agree at every layer to documented numerical tolerance, including under variable-length padded batches. - -### What each test checks - -| Test | Checks | -|------|--------| -| `test_tokenizer_parity` | Vocab size, every token id, every special token id (`pad`, `cls`, `eos`, `mask`, `unk`) match exactly. | -| `test_weight_parity_fp32` | Per-parameter bit-exact equality. Expected extras (e.g. ANKH's `lm_head.weight`) are allowlisted. | -| `test_forward_parity_fp32` | Per-layer relative-std-of-diff (`std(fast - native) / std(native)`), `last_hidden_state` MSE/maxabs, logits MSE. Parametrized over `single`/`uniform`/`skewed` padding scenarios. | -| `test_forward_parity_bf16` | Same as fp32 with documented per-family tolerances. | -| `test_padding_does_not_pollute_valid_positions_fp32` | Runs `[short]` alone and `[short, long_]` padded; asserts the short sequence's valid-position output matches across both. Catches mask-handling bugs that uniform-length tests miss. | -| `test_backend_consistency_fp32` | SDPA vs `kernels_flash` vs `flex` on FastPLMs side, against SDPA as ground truth. | -| `test_embed_dataset_pipeline_parity` | End-to-end `embed_dataset()` vs manual native forward + mean-pool. This is what downstream users actually call. | - -### Tolerances - -Per-family tolerances live in `FAMILY_TOLERANCES` at the top of `test_parity.py`. fp32 tolerances are tight (machine precision); bf16 tolerances are looser per-family because accumulated rounding scales with depth (ESMC has 30 layers, ANKH-base has 48, ESM2-8M has 6). - -### Adding a new family - -1. Add a registry entry in `testing/conftest.py` (`MODEL_REGISTRY` and `FULL_MODEL_REGISTRY`). -2. Implement `testing/official/.py` exporting `load_official_model(reference_repo_id, device, dtype)` that returns `(wrapped_model, tokenizer)` with `.forward()` returning `.logits`, `.hidden_states`. -3. Add a `Dockerfile.` that installs the family's native deps on top of `fastplms-base`, and add it to `build_images.sh`. -4. Add a `ParityTolerances(...)` entry in `FAMILY_TOLERANCES` with reasonable starting values, then tighten as you investigate failures. +```bash +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/private-key \ + --suite python-matrix +``` -## Compliance Testing (`test_compliance.py`) +The matrix creates an isolated uv-managed environment for each version, +installs `requirements/profiles/runtime.in`, enables offline Hub behavior, +disables CUDA visibility, compiles the repository source, loads `models.toml`, +and runs a small ESM2 CPU forward from an explicit source root. Results are +recorded in JSON and JUnit. Python 3.12 remains the only environment used for +the pinned CUDA 13.0, PyTorch 2.13.0, and Transformers 5.13.0 GPU release gates. -Older, looser test layer kept for backward compatibility. Compares FastPLM and official outputs in bf16 with MSE < 0.05 and prediction accuracy > 0.90. Use `test_parity.py` instead for new work. +For direct execution in an already synchronized checkout: -DPLM2 is excluded from weight compliance because the official model has an extra `contact_head` not present in the FastPLM version. +```bash +sudo docker buildx bake -f docker/docker-bake.hcl candidate-structure --load +sudo docker compose -f docker/compose.yaml run --rm structure \ + python -m pytest tests/unit tests/integration tests/release \ + -m "not gpu and not slow and not structure and not artifact" -v +``` -## Throughput Benchmarking +## Manifest-generated cases + +`tests/conftest.py` loads `src/fastplms/models.toml`. Each checkpoint contributes +its AutoClasses, tokenizer mode, source revisions, state transform, reference +container, dependencies, backends, dtypes, precision paths, and declared test +tiers. Release tests fail when the manifest, Docker targets, artifact metadata, +support tables, or model cards diverge. + +Any manifest `unresolved_files` entry is an explicit release blocker. A test may +report it precisely, but must not guess a hash or silently omit the asset. + +## Reference isolation + +Each official adapter runs in a reference stage built from its pinned submodule +and native dependency set. The adapter may call official public APIs and +normalize output names and layout. It may not: + +- import `fastplms`; +- add the FastPLMs source tree to `sys.path`; +- patch official classes; +- reuse a FastPLMs tokenizer or checkpoint loader; +- duplicate an official forward pass in test code. + +Candidate code runs in a separate container. Comparisons exchange serialized +inputs and outputs, not live Python objects. + +The Biohub ESM oracle does not resolve the upstream package's mutable +`transformers @ ... @main` dependency. Its reference image derives only the +remaining dependency set, installs ESM without dependency resolution, and +force-installs the manifest-pinned Biohub Transformers checkout last. Both the +image build and every ESMC, ESM3, or ESMFold2 adapter or standalone loader run +verify the exact source revision, complete tracked-tree inventory and digest, +pre-import origin, and package version before the official Transformers API is +imported, then rehash after import. The locally built, non-editable wheel owns +distribution metadata and dependency validation; imports execute from the +separately attested source tree. Native `metadata.json` and ESMFold2 reference +bundles preserve the exact versioned `reference_sources` mapping for both +Biohub ESM and the pinned Biohub Transformers checkout, including each +attestation-file hash and source-relative import identity. + +The GH200/SM90 ESMFold reference image makes one documented build-only modification: +`docker/constraints/openfold-sm90.patch` restricts the copied OpenFold +`setup.py` extension architecture list to `sm90`. BuildKit cannot discover the +host GPU, and the upstream fallback list contains architectures rejected by +CUDA 12.1. The patch does not change the pinned submodule, extension source, +model classes, checkpoint data, or the public ESMFold API. Its modified-file +notice is `LICENSES/openfold/MODIFICATIONS.md`. + +The same native stage pins PyTorch Lightning `1.9.5`, TorchMetrics `0.11.4`, +Lightning Utilities `0.15.2`, and NVIDIA DLLogger revision +`0478734ff7be75adde8d160e04872664d1c62e5f`. Pinned OpenFold imports those +packages eagerly; they are reference-container dependencies and are excluded +from FastPLMs runtime images and direct dependency profiles. + +## Exact contracts + +Every release-gated checkpoint declaring `compliance` must establish: + +- exact semantic configuration equality, excluding only declared packaging + fields; +- exact tokenizer assets, vocabulary, special IDs, normalization, and behavior + for canonical, ambiguous, lowercase, whitespace, empty, truncated, and padded + inputs; +- exact state-key sets, tensor shapes, dtypes, and `torch.equal` values after the + declared transform; +- exact tied-weight and parameter-alias contracts; +- one live BF16 mixed-length inference; +- representative deep FP32 and BF16 comparisons across all layers, public + outputs, embeddings, skewed padding, and required attention backends. + +DPLM2 specifically asserts that input and output embeddings are not aliased and +that the trained `esm.contact_head` keys are present. ANKH covers the official +encoder and sequence-to-sequence heads separately; the named masked-LM extension +is tested as an extension, not as official parity. + +## Numerical metrics + +Metrics include only valid biological positions. Padding and special positions +cannot improve a score. The suite reports relative L2 error, relative +99.9th-percentile error, first-percentile residue cosine, per-sequence pooled +cosine, confident-position top-1 agreement, and Jensen-Shannon divergence for +probability tensors. + +The tables below record the fixed engineering targets and hard limits enforced +by the release tests. Repeated runs and official baselines are diagnostic +evidence, but they do not silently widen those limits. A new implementation +must meet the engineering target, not only the hard limit. + +| Contract | Engineering target | Hard limit | +| --- | ---: | ---: | +| FP32 official relative L2 | `2e-6` | `2e-5` | +| FP32 relative Q99.9 error | `1e-5` | `1e-4` | +| BF16 official or backend relative L2, except scoped rows below | `1e-2` | `3e-2` | +| ESM2 optimized-backend BF16 relative L2 | `2e-2` | `3e-2` | +| ESMC eager BF16 relative L2 | `2.9e-2` | `3e-2` | +| BF16 relative Q99.9 error | `2.5e-2` | `5e-2` | +| ESMC eager relative Q99.9 error | `4.9e-2` | `5e-2` | +| BF16 residue cosine, first percentile | `>=0.999` | `>=0.995` | +| ESMC eager residue cosine, first percentile | `>=0.997` | `>=0.995` | +| BF16 pooled cosine, every sequence | `>=0.9995` | `>=0.995` | +| BF16 confident top-1 agreement | `>=99.5%` | `>=99.0%` | +| BF16 Jensen-Shannon divergence | `1e-4` | `1e-3` | +| ESMC eager Jensen-Shannon divergence | `4e-4` | `1e-3` | + +The ESM2 row applies to SDPA, Flex Attention, FlashAttention 2, and +FlashAttention 3; eager retains the global contract. The current ESMC row +applies to eager. SDPA remains bit-for-bit exact; Q99.9, pooled-cosine, and +top-1 thresholds remain global. ESMC Flex Attention is supported, +non-experimental, and measured as a diagnostic path. Crossing a published Flex +band emits a warning and writes JSON under `artifacts/diagnostics/esmc/`; it is +not an xfail or release failure. The old FlashAttention 2 threshold is retained +only with its separate historical focused evidence and is not a current GH200 +acceptance claim. + +Measured Flex still fails on broken dispatch, non-finite output, invalid mask +or shape, and these corruption/catastrophe limits: + +| ESMC diagnostic | Catastrophe limit | +| --- | ---: | +| Relative L2 | `<=0.25` | +| Relative Q99.9 | `<=0.50` | +| Residue cosine, first percentile | `>=0.90` | +| Pooled cosine | `>=0.95` | +| Confident-position top-1 | `>=0.80` | +| Jensen-Shannon divergence | `<=0.05` | + +These broad limits detect corruption and do not claim parity or biological +quality. Each measured release record must publish full distributions for the +exact checkpoint, backend, dtype, hardware, and locked sequence panel. The +FlashAttention 2 and 3 records contain structured current-platform +unavailability instead. Both remain supported, non-experimental interfaces; +the current locked GH200/aarch64 image raises before dispatch. Model-card cells +remain pending until frozen-head, exact-device GH200/aarch64 evidence exists. + +The pinned ESM2-3B SDPA BF16 path has a checkpoint-specific calibration: +relative L2 target/hard limit `0.06`/`0.07`, relative Q99.9 `0.15`/`0.18`, +first-percentile residue cosine `0.994`/`0.992`, and pooled cosine +`0.998`/`0.997`. Its confident top-1 and Jensen-Shannon thresholds remain +global. Exact state identity and perfect confident-token agreement still gate +this checkpoint. + +ESMFold2 FP8 is experimental and is not a release numerical-parity gate. Its +smoke coverage on the locked, exact-device GH200/aarch64 stack verifies explicit opt-in, finite +outputs, exactly 80 converted ESMC attention output projections, transient +runtime state, and strict failure when unavailable. `auto` always resolves to +BF16 so model behavior does not change with hardware or optional dependencies. + +## ESMFold2 projection and structure + +Projection from identical ordered ESMC states is exact in FP32. The BF16 +relative L2 target is `5e-4`, with a hard limit of `1e-3`. Experimental FP8 +smoke runs once on each of the four variants and performs three fresh +BF16-to-FP8 reload cycles only on the standard variant. + +Folding tests hash prepared features and sampled diffusion noise. They require +exact discrete features and masks, valid geometry, and no NaNs. Coordinate and +confidence thresholds are documented in [ESMFold2](esmfold2.md) and encoded once +in the strict metric module. The pinned five-protein, three-seed, four-variant +panel found exact official-versus-candidate BF16 parity in all 60 cases. A +prior FP8 diagnostic passed its historical structure limits in 48 of 60 cases; +that result is retained as evidence, not as a release gate or equivalence claim. + +## Goldens + +Official-generated goldens use safetensors. Each includes the official source +revision, checkpoint revision and hashes, environment fingerprint, deterministic +generation command, input fingerprint, tensor names and shapes, dtypes, and +output hashes. Goldens are read-only fixtures. They accelerate `check`, but they +never replace live `compliance`. + +The manifest declares a required golden only through an `official_golden` +record on a model entry. Both files are SHA-256 pinned and use fixed paths: + +```toml +official_golden = { metadata = "tests/goldens/.json=sha256:", tensors = "tests/goldens/.safetensors=sha256:" } +``` -### Pytest Test (`test_throughput.py`) +Absence of this record means that the checkpoint golden is not complete. It +must be reported as release work rather than represented by a placeholder or a +synthetic fixture. Presence makes a missing, modified, or +source-record-inconsistent bundle a `check` failure. Other tiers do not infer a +golden requirement. -Benchmarks multiple model families across all 3 backends, batch sizes, and sequence lengths. Saves structured results: +The manifest-driven converter consumes only normalized output from an isolated +native reference container. It does not load a model, import an upstream +package, or download a checkpoint: -- `throughput_results.json`: machine-readable -- `throughput_results.csv`: spreadsheet-friendly -- `throughput_comparison.png`: visualization plot +```bash +python -m tools.goldens \ + --native-root artifacts/reference \ + --output-root tests/goldens \ + --model esm2_8m + +python -m tools.goldens \ + --status-only \ + --report-matrix \ + --native-root artifacts/reference \ + --output-root tests/goldens +``` -The pytest test uses fewer timed batches (25 vs 100) for faster execution. +The matrix is the manifest-wide source of truth for all `check` checkpoints. It +reports each native request, reference container, normalized native result, +compact bundle, and declaration state. A targeted native sequence run accepts +one or more explicit model IDs: -### Standalone Script (`testing/throughput.py`) +```bash +python -m tests.parity.support.native_reference \ + --request-dir artifacts/reference/requests/reference-esm2 \ + --output-dir artifacts/reference/results \ + --model esm2_8m +``` -More configurable, with CLI arguments: +An official public-generation limitation is not generation parity. Only a +checkpoint whose manifest capability is `official_unavailable` may carry a +normalized limitation record, and that record must match the public method, +exception type, and semantic reason exactly. All `required` DPLM and DPLM2 +checkpoints fail when generation output is absent. Feature tests retain viable +family representatives even when one checkpoint's pinned official sampler is +unusable. + +For sequence models, the input is `metadata.json` plus `bf16.safetensors`. The +converter retains token inputs, the biological-residue mask, the final hidden +state, and logits when the official head returns them. For structure models, +the input is the official `metadata.json` plus `bundle.safetensors`. In both +cases it validates the model ID, checkpoint revision and file identities, +reference environment, normalized tensor contract, and source-result hashes. +It rejects candidate-produced structure bundles and legacy native results that +do not carry an environment record. + +The output is a compact safetensors file and JSON sidecar. The sidecar records +upstream revisions, official checkpoint file identities, the native environment +fingerprint, canonical generation command, deterministic input fingerprint, +source-result file hashes, tensor hashes, and the output tensor-file hash. The +converter prints a TOML declaration only when output is written to the canonical +`tests/goldens` directory. It never edits `models.toml`; a reviewer adds the +printed declaration only after validating both generated files. The read-only +validator then verifies every recorded identity, shape, dtype, and hash. + +The sequence regression resolves the current repository-source class from the manifest +`auto_map` and loads only the pinned checkpoint weights. Generated remote-code +artifacts have their own offline suite. This separation prevents stale Hub code +from substituting for the repository implementation under test. + +The pinned Biohub ESMC loader has a standalone reproducer so a native-loader +failure cannot be mistaken for a FastPLMs inference failure: ```bash -python -m testing.throughput \ - --model_paths Synthyra/ESM2-8M Synthyra/ESMplusplus_small \ - --backends sdpa flex kernels_flash \ - --batch_sizes 2 4 8 \ - --sequence_lengths 64 128 256 512 1024 2048 \ - --warmup_batches 10 \ - --timed_batches 100 \ - --output_path /workspace/throughput_comparison.png +python -m tests.parity.support.reference_adapters.biohub_loader_reproducer \ + --repo-id biohub/ESMC-300M \ + --revision a59b831785f907e96e6a246b1d142bfb76df31ee ``` -Both pytest and standalone output JSON and CSV in addition to the plot. +Run it inside `reference-biohub-esm`. It prints the native package versions and +then invokes the pinned public `ESMC.from_pretrained` method unchanged. The +declared ESMC goldens were generated through that native path; the suite does +not patch the loader or replace it with a FastPLMs path. -### How throughput is measured +## Artifacts -1. Model is compiled via `torch.compile()` -2. Dynamic warmup: 10-100 batches until latency stabilizes (relative change <= 10% over a 3-batch window) -3. Timed phase: N batches with `torch.cuda.synchronize()` around the timing loop -4. Reports non-padding tokens per second +Artifact tests build from a pinned local checkpoint snapshot. They create a +fresh environment with FastPLMs absent from `sys.path`, set +`HF_HUB_OFFLINE=1`, pass `local_files_only=True` and `trust_remote_code=True`, +load every advertised AutoClass, run inference, save, reload, and compare with +the repository-source implementation. Network access during this tier is a test +failure. -### Boltz2 Compliance +## Test markers -Boltz2 has its own compliance script (`testing/run_boltz2_compliance.py`) that compares: -- Coordinate MAE/RMSE (raw and Kabsch-aligned) -- Pairwise distance MAE -- TM-score comparison - -```bash -python -m testing.run_boltz2_compliance \ - --device cuda \ - --dtype float32 \ - --seed 42 \ - --num-sequences 3 \ - --recycling-steps 3 \ - --num-sampling-steps 200 -``` +The suite uses `gpu`, `slow`, `large`, and `structure` markers to describe +resource needs. Markers do not authorize skipping a required release case. The +runner chooses the appropriate image and tier, then treats an unmet declared +requirement as a failure. diff --git a/docs/ttt.md b/docs/ttt.md new file mode 100644 index 0000000..11ac838 --- /dev/null +++ b/docs/ttt.md @@ -0,0 +1,103 @@ +# Test-time training + +FastPLMs includes an opt-in ProteinTTT-derived adaptation path for supported +sequence models and the ESMC language-model backbone used by ESMFold2. Ordinary +construction, inference, embedding, folding, and `state_dict()` do not perform +adaptation. + +## Mechanism + +Given one protein sequence, TTT samples masked views, computes the model's +masked-language-model loss, and updates only injected low-rank adapter +parameters. Base checkpoint parameters remain frozen. The returned metrics +record per-step loss and any explicitly enabled evaluation values. + +Tokenizer-based paths use the tokenizer assets and immutable revision attached +to the loaded checkpoint. + +```python +metrics = model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={ + "steps": 3, + "ags": 1, + "batch_size": 1, + "seed": 7, + }, +) +model.ttt_reset() +``` + +`ttt_reset()` restores the initial adapter state. With +`initial_state_reset=True`, each `ttt()` call begins from that initial state. +Seeded tests compare mask sampling, losses, updated parameter scope, reset, and +checkpoint state. + +Adapter initialization uses the TTT seed without advancing the caller's Python, +NumPy, CPU Torch, or CUDA RNG streams. Random BERT-style replacements are drawn +only from the 20 canonical biological amino acids supported by the family, not +from arbitrary vocabulary entries, boundary tokens, or structure modalities. +Uneven final sample batches remain finite and use their actual valid count. + +## Main controls + +`TTTConfig` controls learning rate, optimization steps, gradient accumulation, +sample batch size, mask ratio, crop size, BERT-style leave and replacement +probabilities, optimizer, seed, low-rank width and scale, target modules, reset, +optional step evaluation, and gradient clipping. + +`lora_alpha` is a direct multiplier on the low-rank adapter output. It is not +divided by `lora_rank`. This intentionally matches the pinned ProteinTTT call +`inject_trainable_lora(..., scale=lora_alpha)` and differs from the common PEFT +LoRA `alpha / rank` convention. The direct scale is serialized with the TTT +configuration, so changing this interpretation would alter reloaded adapters. + +Changing low-rank width, scale, or target modules after adapter initialization +is rejected because it would change the parameter schema. + +## Save and reload + +Initialized adapter tensors, their reset baseline, and normalized TTT +configuration are part of `save_pretrained`: + +```python +from transformers import AutoModelForMaskedLM + +model.ttt(seq="MSTNPKPQRKTKRNT", ttt_config={"steps": 3, "seed": 7}) +model.save_pretrained("adapted", safe_serialization=True) +reloaded = AutoModelForMaskedLM.from_pretrained( + "adapted", + trust_remote_code=True, + local_files_only=True, +) +reloaded.ttt_reset() +``` + +Reload preserves both the adapted state and the deterministic reset state. +Models whose adapters attach to transient modules excluded from checkpoint +state fail closed and require a model-specific export rather than silently +dropping adaptation. + +## Folding + +ESMFold2 exposes a family-specific opt-in folding helper. Adaptation affects +only its language-model backbone. ESMFold2 uses canonical BF16 ESMC weights +before a gradient-enabled path. If serving selected FP8, entering TTT reloads +canonical BF16 weights while preserving the requested serving policy in +configuration and status metadata. + +Meta ESMFold does not expose TTT. Its pinned checkpoint contains the folding +language model but no trained masked-language-model head for the ProteinTTT +objective. `ttt()`, `ttt_reset()`, `fold_protein(ttt=True)`, and +`fold_protein_ttt()` therefore raise explicitly. FastPLMs does not construct or +serialize an untrained replacement head. + +## Limitations + +TTT increases latency and GPU memory and can worsen a prediction. It is not a +calibration method and does not establish biological function. Compare the +unadapted output, retain complete seeds and configuration, and validate on an +independent task-specific set before drawing a scientific conclusion. + +Boltz2 remains inference-only in FastPLMs. The manifest and feature suite define +the model families that advertise TTT. diff --git a/docs/vector_embeddings/README.md b/docs/vector_embeddings/README.md index 1328320..25db881 100644 --- a/docs/vector_embeddings/README.md +++ b/docs/vector_embeddings/README.md @@ -15,7 +15,15 @@ tags: Precomputed pooled protein embeddings for the Protify vector benchmark. -This dataset stores ready-to-use `.pth.gz` embedding artifacts for a broad panel of protein language models and controls. It is intended for fast downstream benchmarking in Protify without repeatedly embedding the same benchmark sequences on local GPUs. +> [!CAUTION] +> This page documents a historical external dataset, not the FastPLMs 1.0 +> embedding format. FastPLMs writes sharded safetensors or transactional SQLite +> and never writes new `.pth` files. Prefer the current +> [embedding API](../embedding_api.md) for new runs. + +This dataset stores ready-to-use `.pth.gz` embedding artifacts for a broad +panel of protein language models and controls. They support downstream Protify +benchmarks without repeatedly embedding the same sequences on local GPUs. ## What Is Included @@ -34,7 +42,8 @@ The filename convention mirrors the Protify embedding cache settings: | `mean_var` | Mean and variance pooling used for the vector benchmark | | `.pth.gz` | PyTorch serialization compressed with gzip | -Use these files when you want fixed-size protein vectors for classical ML, vector search, nearest-neighbor analysis, low-shot benchmarking, or model comparison. +Use these files for fixed-size protein vectors in classical ML, vector search, +nearest-neighbor analysis, low-shot benchmarking, or model comparison. ## Quick Start @@ -65,7 +74,7 @@ import gzip import torch with gzip.open(path, "rb") as handle: - embeddings = torch.load(handle, map_location="cpu") + embeddings = torch.load(handle, map_location="cpu", weights_only=True) first_key = next(iter(embeddings)) first_vector = embeddings[first_key] @@ -74,6 +83,11 @@ print(first_key) print(first_vector.shape, first_vector.dtype) ``` +Only load a legacy pickle from a source whose bytes and expected hash you +trust. If an old payload cannot be read with `weights_only=True`, use the +FastPLMs read-only legacy importer and its explicit unsafe-pickle opt-in. Never +enable unsafe pickle loading for an untrusted download. + Download with the Hugging Face CLI: ```bash @@ -92,11 +106,15 @@ hf download Synthyra/vector_embeddings \ --local-dir vector_embeddings ``` -The full `embeddings/` directory is about 189.22 GiB. +The full directory is large. Inspect the current Hub repository size before a +bulk download. ## Protify Usage -Protify can use these precomputed embeddings as model-ready vector caches for benchmark runs. This avoids recomputing embeddings for every model and dataset combination, which is especially useful for large PLMs and laptop-scale downstream analysis. +Protify can use these precomputed embeddings as vector caches for benchmark +runs. This avoids recomputing every model and dataset combination, which is +especially useful for large protein language models or laptop-scale downstream +analysis. Typical workflow: @@ -107,6 +125,9 @@ Typical workflow: ## Model Inventory +This inventory is a historical snapshot of the external dataset. It is not the +FastPLMs support matrix and is not generated from `src/fastplms/models.toml`. + | Model | File | Size | | --- | --- | ---: | | AMPLIFY-120 | `AMPLIFY-120_False_mean_var.pth.gz` | 3.63 GiB | @@ -142,12 +163,18 @@ Typical workflow: ## Notes -These files are large. Prefer single-file downloads with `hf_hub_download()` or `hf download --include` unless you explicitly need every model. +These files are large. Prefer single-file downloads with `hf_hub_download()` or +`hf download --include` unless you need every model. -This repository is an artifact store, not a tabular Hugging Face Datasets dataset. Use `huggingface_hub` or the `hf` CLI rather than `datasets.load_dataset()`. +This repository is an artifact store, not a tabular Hugging Face Datasets +dataset. Use `huggingface_hub` or the `hf` CLI rather than +`datasets.load_dataset()`. -Embedding artifacts are intended to support reproducible Protify benchmarking. For redistribution or derived work, check the licenses and usage terms of the original benchmark datasets and the upstream model checkpoints used to generate each embedding set. +The artifacts support reproducible Protify benchmarking. For redistribution or +derived work, check the licenses and usage terms of the original benchmark +datasets and upstream model checkpoints used to generate each embedding set. ## Citation -If these embeddings help your work, please cite the relevant upstream model papers or model cards, Protify, and Synthyra resources used in your analysis. +If these embeddings help your work, cite the relevant upstream model papers or +model cards, Protify, and the Synthyra resources used in your analysis. diff --git a/entrypoint_setup.py b/entrypoint_setup.py deleted file mode 100644 index 4fdcacf..0000000 --- a/entrypoint_setup.py +++ /dev/null @@ -1,22 +0,0 @@ -import torch -import torch._inductor.config as inductor_config -import torch._dynamo as dynamo - -# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs) -# Provides significant speedup with minimal precision loss -torch.set_float32_matmul_precision('high') - -# Enable TF32 for matrix multiplications and cuDNN operations -torch.backends.cuda.matmul.allow_tf32 = True -torch.backends.cudnn.allow_tf32 = True - -# Enable cuDNN autotuner - finds fastest algorithms for your hardware -# Best when input sizes are consistent; may slow down first iterations -torch.backends.cudnn.benchmark = True - -# Deterministic operations off for speed (set True if reproducibility needed) -torch.backends.cudnn.deterministic = False -inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM" - -dynamo.config.capture_scalar_outputs = True -torch._dynamo.config.recompile_limit = 16 diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..bd105df --- /dev/null +++ b/examples/README.md @@ -0,0 +1,153 @@ +# FastPLMs runnable examples + +The examples are executable entry points, not performance or biological-validity +claims. The curated offline examples consume local, manifest-built Hugging Face +artifacts, set both Hub offline variables, pass `local_files_only=True`, and do +not download models, tokenizers, kernels, or runtime assets. + +## Dependencies and platform requirements + +FastPLMs 1.0 requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13. +Published models carry their runtime source in the Hugging Face repository. +Install the core dependencies directly, then load the model with +`trust_remote_code=True`: + +```bash +python -m pip install \ + "torch>=2.13,<2.14" \ + "transformers>=5.13,<5.14" +``` + +These examples run from a source checkout. Install the profile needed by the +workflow, then invoke them with `PYTHONPATH=src`. The CPU validation profile +covers the portable examples: + +```bash +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu +``` + +Core sequence examples accept `--device cpu|cuda[:index]` and +`--dtype float32|bfloat16`; the defaults are the portable CPU/FP32 path. +Structure examples require `requirements/features/structure.in`, verified +runtime assets, and CUDA for the published execution contract. Binder design +uses `requirements/profiles/binder.in`. FlashAttention requires +`requirements/features/flash.in`, compatible CUDA hardware, a pre-populated +pinned kernel cache, and BF16. Fine-tuning uses +`requirements/features/train.in`; add `requirements/features/reporting.in` +only for plots and statistical reports. + +For ESMFold2, choose the artifact by conditioning contract. The full +`ESMFold2` and `ESMFold2-Experimental-Cutoff2025` checkpoints have 48 folding +blocks and support optional MSA conditioning. The Fast and experimental Fast +checkpoints have 24 folding blocks, are optimized for single-sequence +inference, and reject MSA-derived inputs. The distinction follows Biohub +[Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). +Fast is not necessarily single-chain-only: supported multichain and +multimolecule requests remain available, but each protein chain uses +single-sequence mode. + +## Prepare an offline artifact + +Build and validate the local artifact before disconnecting the network: + +```bash +PYTHONPATH=src python -m tools.artifacts.build \ + esm2_8m /cache/fast-snapshot \ + --tokenizer-dir /cache/official-tokenizer-snapshot \ + --output-root dist/hub +``` + +Start with `--help` for any entry point. Representative portable commands are: + +```bash +PYTHONPATH=src python examples/artifact_loading.py dist/hub/ESM2-8M --auto-class AutoModel +PYTHONPATH=src python examples/embedding_and_retrieval.py dist/hub/ESM2-8M \ + --sequence MSTNPKPQRKTKRNT --device cpu --dtype float32 +PYTHONPATH=src python examples/attention_switching.py dist/hub/ESM2-8M \ + --backend sdpa --device cpu --dtype float32 +PYTHONPATH=src python examples/task_heads.py dist/hub/ESM2-8M \ + --attn-backend eager --device cpu --dtype float32 +``` + +The structure-preparation example deliberately constructs an MSA-conditioned +request, so point it at a full ESMFold2 artifact: + +```bash +PYTHONPATH=src python examples/structure_preparation.py \ + esmfold2 dist/hub/ESMFold2 --device cuda:0 +``` + +Use Flex when a compiled path is wanted on the current GH200/aarch64 validation +target: + +```bash +PYTHONPATH=src python examples/attention_switching.py dist/hub/ESM2-8M \ + --backend flex_attention --device cuda:0 --dtype bfloat16 +``` + +The CLI retains explicit FlashAttention 2 and 3 choices for supported +family/platform combinations with a pre-populated pinned kernel cache. The +current locked GH200/aarch64 environment has no expected Flash kernels, so use +SDPA or Flex there. Do not build an unpinned Flash kernel from source. Prior +FlashAttention 2 results remain historical exact-environment evidence; +FlashAttention 3 is supported but unavailable on the current locked target. + +## Example inventory and evidence boundary + +| Workflow | Example | Demonstrated contract | Boundary | +| --- | --- | --- | --- | +| Offline AutoClass loading | [`artifact_loading.py`](artifact_loading.py) | Load any advertised AutoClass from a local artifact | Loading only; forward/loss/save-reload are CPU contract tests | +| MLM, contacts, and task heads | [`task_heads.py`](task_heads.py) | ESM2 masked-residue scoring, trained contact head, sequence and token classification loss | Sequence and token classifiers use base weights + untrained task head unless a separately fine-tuned head is supplied | +| Ordered embeddings and retrieval | [`embedding_and_retrieval.py`](embedding_and_retrieval.py) | Repeated sequences or FASTA, mean/std pooling, safetensors or SQLite, duplicate-preserving SQLite retrieval | Full-residue, all-layer, mapping, generator, and other poolers remain shared-API examples/tests | +| Attention switching | [`attention_switching.py`](attention_switching.py) | Eager, SDPA, Flex, explicit Flash requirements, warning-emitting masked eager fallback without configuration mutation | Not a parity or throughput benchmark; the current GH200/aarch64 lock has no expected Flash kernels | +| ANKH stack selection | [`ankh_embeddings.py`](ankh_embeddings.py) | Encoder final/all layers, decoder layer with explicit prompt, deterministic seq2seq generation | The offline example accepts a validated local artifact and loads both views, so budget device memory accordingly | +| Diffusion and multimodal generation | [`generation.py`](generation.py) | Seeded DPLM, DPLM2, and conditioned ESM3 generation | One representative deterministic strategy per family | +| E1 RAG | [`e1_rag.py`](e1_rag.py) | Local A3M retrieval, ordered duplicate records, shared persistence | No remote MSA search or network fallback | +| Test-time training | [`ttt.py`](ttt.py) | Seeded update, atomic save, reset, local reload | Output must be absent and outside the source artifact | +| Structure preparation | [`structure_preparation.py`](structure_preparation.py) | Typed ESMFold2 multimolecule/MSA/modification/bond/distogram input, pocket rejection, seeded ESMFold/Boltz helpers | The MSA branch requires a full 48-block ESMFold2 variant; Fast variants reject MSA-derived inputs; tiny preparation and helper contracts are not full folding parity | +| Fine-tuning | [`fine_tuning.py`](fine_tuning.py) | ESM2 classification/regression, LoRA or full tuning, eager/SDPA/Flex selection, immutable inputs, atomic verified final artifact | LoRA is the demonstrated PEFT method; Flash training requires a separate explicit BF16 CUDA policy; other PEFT methods are not claimed by this example | +| Binder design | [`binder_design_fastplms.py`](binder_design_fastplms.py) | Differentiable ESMFold2/ESM++ optimization and critic consensus | Research prioritization only; no experimental binding claim | + +The generated [capability-to-evidence manifest](../docs/generated/capability_evidence.md) +maps each curated example to its required CPU, feature, structure, nightly, or +compliance evidence. A capability absent from the table above is not implied by +an example merely because its model class exists. + +## Embedding coverage matrix + +| Surface | Runnable CLI coverage | Where the remaining contract is shown | +| --- | --- | --- | +| Repeated sequence list | `--sequence` may be repeated | `embedding_and_retrieval.py` | +| FASTA streaming | `--fasta` | `embedding_and_retrieval.py` | +| Insertion-ordered mapping | Not a CLI encoding | [Embedding API](../docs/embedding_api.md) and CPU contracts | +| One-shot generator | Not a CLI encoding | [Embedding API](../docs/embedding_api.md) and CPU contracts | +| In-memory output | Omit `--output` | `embedding_and_retrieval.py` | +| Safetensors write/reopen | `--format safetensors` | Runnable example plus persistence CPU contracts | +| SQLite write/read-only filtered retrieval | `--output PATH --format sqlite --select-id ID` | Runnable example and duplicate-order CPU contracts; other `--select-id` combinations fail before loading | +| Mean and standard-deviation pooling | Always demonstrated together | `embedding_and_retrieval.py` | +| Full-residue and all-layer tensors | Not exposed by this compact CLI | [Embedding API](../docs/embedding_api.md) and ANKH example | +| Other declared poolers | Not exposed by this compact CLI | [Embedding API](../docs/embedding_api.md) and CPU contracts | + +## Network and output policy + +`fine_tuning.py` and `binder_design_fastplms.py` are checkpoint workflows, not +members of the fully offline example gate. Their shipped remote defaults are +pinned automatically. Custom remote model or dataset sources reject omitted, +branch, and tag revisions; pre-populate every snapshot before a network-isolated +run. Local fine-tuning dataset directories must be layouts accepted by +`datasets.load_dataset`; arbitrary `Dataset.save_to_disk()` trees are not +currently accepted. + +Fine-tuning writes separate task-specific children beneath `--output-dir` and +records requested and effective attention backends. Binder design refuses any +pre-existing output directory. It writes `run_manifest.json` atomically last; +its absence identifies an incomplete run. Retain the complete directory for +reproducibility. + +The CPU gate executes CLI wiring and dependency-free preparation with tiny +local artifacts. Full checkpoints, real optimized kernels, GPU parity, +structure prediction, and throughput remain in the feature, nightly, +compliance, structure, and benchmark tiers. diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..c560711 --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1 @@ +"""Runnable FastPLMs research and training examples.""" diff --git a/examples/_runtime.py b/examples/_runtime.py new file mode 100644 index 0000000..dc8582e --- /dev/null +++ b/examples/_runtime.py @@ -0,0 +1,33 @@ +"""Shared fail-closed execution arguments for runnable examples.""" + +from __future__ import annotations + +import argparse +from typing import Any + + +def add_execution_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--device", default="cpu", help="cpu or cuda[:index]") + parser.add_argument( + "--dtype", + choices=("float32", "bfloat16"), + default="float32", + help="Model parameter and compute dtype", + ) + + +def resolve_execution(device_name: str, dtype_name: str) -> tuple[Any, Any]: + """Resolve a CPU/CUDA device and dtype before loading a checkpoint.""" + + import torch + + try: + device = torch.device(device_name) + except (RuntimeError, TypeError) as error: + raise ValueError(f"Invalid execution device {device_name!r}") from error + if device.type not in {"cpu", "cuda"}: + raise ValueError(f"Only CPU and CUDA devices are supported, got {device.type!r}") + if device.type == "cuda" and not torch.cuda.is_available(): + raise ValueError(f"CUDA device {device} was requested but CUDA is unavailable") + dtype = torch.float32 if dtype_name == "float32" else torch.bfloat16 + return device, dtype diff --git a/examples/ankh_embeddings.py b/examples/ankh_embeddings.py new file mode 100644 index 0000000..1fdf10f --- /dev/null +++ b/examples/ankh_embeddings.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Extract ANKH hidden states and run task-prompted seq2seq generation.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +if __package__: + from ._runtime import add_execution_arguments, resolve_execution +else: + from _runtime import add_execution_arguments, resolve_execution + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def extract_ankh_layers( + encoder: Any, + seq2seq: Any, + tokenizer: Any, + sequences: list[str], + decoder_prompts: list[str], +) -> tuple[Any, Any, Any]: + encoder_final = encoder.embed_dataset( + sequences, + tokenizer=tokenizer, + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, + ) + encoder_all = encoder.embed_dataset( + sequences, + tokenizer=tokenizer, + hidden_state_source="encoder", + store_all_hidden_states=True, + full_embeddings=True, + ) + decoder_final = seq2seq.embed_dataset( + sequences, + tokenizer=tokenizer, + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=decoder_prompts, + full_embeddings=True, + ) + return encoder_final, encoder_all, decoder_final + + +def generate_ankh_task( + model: Any, + tokenizer: Any, + sequence: str, + decoder_prompt: str, + *, + max_new_tokens: int, +) -> Any: + """Generate after an explicit task prompt without shifting the source.""" + import torch + + from fastplms.models.ankh.modeling_ankh import ( + tokenize_ankh_decoder_prompts, + tokenize_ankh_sequences, + ) + + encoded = tokenize_ankh_sequences( + tokenizer, + sequence, + return_tensors="pt", + ) # input_ids/attention_mask: (1, l_s) + prompt = tokenize_ankh_decoder_prompts( + tokenizer, + decoder_prompt, + return_tensors="pt", + add_special_tokens=False, + ) # input_ids: (1, l_p) + prompt_ids = prompt["input_ids"] # (1, l_p) + decoder_start_token_id = getattr(model.config, "decoder_start_token_id", None) + if not isinstance(decoder_start_token_id, int): + raise RuntimeError("The ANKH artifact does not declare decoder_start_token_id.") + decoder_input_ids = torch.cat( + ( + prompt_ids.new_full((prompt_ids.shape[0], 1), decoder_start_token_id), + prompt_ids, + ), + dim=1, + ) # (1, l_p + 1) + device = model.device + generation_inputs = { + "input_ids": encoded["input_ids"].to(device), # (1, l_s) + "decoder_input_ids": decoder_input_ids.to(device), # (1, l_p + 1) + "decoder_attention_mask": torch.ones_like( + decoder_input_ids, device=device + ), # (1, l_p + 1) + } + if "attention_mask" in encoded: + generation_inputs["attention_mask"] = encoded["attention_mask"].to( + device + ) # (1, l_s) + with torch.inference_mode(): + generated = model.generate( + **generation_inputs, + do_sample=False, + num_beams=1, + use_cache=True, + max_new_tokens=max_new_tokens, + ) # (1, l_o) + return generated # (1, l_o) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path, help="Full local ANKH 1.0 artifact") + parser.add_argument("--sequence", default="MSTNPKPQRKTKRNT") + parser.add_argument("--decoder-prompt", default="M") + parser.add_argument("--max-new-tokens", type=int, default=4) + add_execution_arguments(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + from transformers import AutoModel, AutoModelForSeq2SeqLM + + common = { + "trust_remote_code": True, + "local_files_only": True, + "dtype": dtype, + } + encoder = AutoModel.from_pretrained(artifact, **common).to(device).eval() + seq2seq = AutoModelForSeq2SeqLM.from_pretrained(artifact, **common).to(device).eval() + tokenizer = seq2seq.tokenizer + results = extract_ankh_layers( + encoder, + seq2seq, + tokenizer, + [arguments.sequence], + [arguments.decoder_prompt], + ) + print("encoder-final", tuple(results[0][0].load_tensor().shape)) + print("encoder-all", tuple(results[1][0].load_tensor().shape)) + print("decoder-final", tuple(results[2][0].load_tensor().shape)) + generated = generate_ankh_task( + seq2seq, + tokenizer, + arguments.sequence, + arguments.decoder_prompt, + max_new_tokens=arguments.max_new_tokens, + ) # (1, l_o) + print("generated", tokenizer.decode(generated[0], skip_special_tokens=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/artifact_loading.py b/examples/artifact_loading.py new file mode 100644 index 0000000..a11cba3 --- /dev/null +++ b/examples/artifact_loading.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Load one manifest-built Hugging Face artifact without network access.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +AUTO_CLASS_NAMES = ( + "AutoConfig", + "AutoModel", + "AutoModelForMaskedLM", + "AutoModelForSeq2SeqLM", + "AutoModelForSequenceClassification", + "AutoModelForTokenClassification", +) + + +def require_local_artifact(value: str) -> Path: + artifact = Path(value).expanduser().resolve() + if not artifact.is_dir() or not (artifact / "config.json").is_file(): + raise argparse.ArgumentTypeError( + f"Expected a local artifact directory containing config.json: {artifact}" + ) + return artifact + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def load_local_artifact(artifact: Path, auto_class_name: str) -> Any: + configure_offline() + import transformers + + auto_class = getattr(transformers, auto_class_name) + loaded = auto_class.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + ) + if hasattr(loaded, "eval"): + loaded = loaded.eval() + return loaded + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=require_local_artifact) + parser.add_argument("--auto-class", choices=AUTO_CLASS_NAMES, default="AutoModel") + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + loaded = load_local_artifact(arguments.artifact, arguments.auto_class) + print(type(loaded).__name__) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/attention_switching.py b/examples/attention_switching.py new file mode 100644 index 0000000..20c8011 --- /dev/null +++ b/examples/attention_switching.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""Run an optimized attention backend and inspect its per-call eager fallback.""" + +from __future__ import annotations + +import argparse +import os +import warnings +from pathlib import Path +from typing import Any + + +FLASH_BACKENDS = frozenset({"flash_attention_2", "flash_attention_3"}) +DTYPE_NAMES = ("float32", "bfloat16") + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def attention_configuration_snapshot(model: Any) -> tuple[tuple[str, str, Any], ...]: + """Capture model/config backend state so a per-call fallback cannot mutate it.""" + + attributes = ( + "attn_backend", + "attn_implementation", + "_attn_implementation", + "_attn_implementation_internal", + ) + state: list[tuple[str, str, Any]] = [] + for owner, value in (("model", model), ("config", getattr(model, "config", None))): + if value is None: + continue + for attribute in attributes: + if hasattr(value, attribute): + state.append((owner, attribute, getattr(value, attribute))) + if not state: + raise RuntimeError("The loaded model does not expose its configured attention backend") + return tuple(state) + + +def run_optimized_attention_example( + model: Any, + tokenizer: Any, + sequences: list[str], +) -> Any: + """Execute the configured backend without requesting attention tensors.""" + + import torch + + batch = tokenizer( + sequences, padding=True, return_tensors="pt" + ) # each tensor: (b, l) + batch = { + name: tensor.to(model.device) for name, tensor in batch.items() + } # each tensor: (b, l) + with torch.inference_mode(): + output = model( + **batch, output_attentions=False + ) # last_hidden_state: (b, l, d) + return output + + +def run_attention_example( + model: Any, + tokenizer: Any, + sequences: list[str], +) -> tuple[Any, list[str]]: + import torch + + batch = tokenizer( + sequences, padding=True, return_tensors="pt" + ) # each tensor: (b, l) + batch = { + name: tensor.to(model.device) for name, tensor in batch.items() + } # each tensor: (b, l) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", RuntimeWarning) + with torch.inference_mode(): + output = model( + **batch, output_attentions=True + ) # last_hidden_state: (b, l, d); attentions: layer-wise (b, h, l, l) + return output, [str(item.message) for item in caught] + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path) + parser.add_argument( + "--backend", + choices=( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ), + default="sdpa", + ) + parser.add_argument( + "--device", + default="cpu", + help="Execution device. Use cpu for the portable path or cuda[:index] for CUDA.", + ) + parser.add_argument( + "--dtype", + choices=DTYPE_NAMES, + default="float32", + help="Model compute dtype. FlashAttention 2 and 3 require bfloat16.", + ) + return parser + + +def resolve_execution(backend: str, device_name: str, dtype_name: str) -> tuple[Any, Any]: + """Resolve and validate the documented backend, device, and dtype contract.""" + + import torch + + try: + device = torch.device(device_name) + except (RuntimeError, TypeError) as error: + raise ValueError(f"Invalid execution device {device_name!r}") from error + if device.type not in {"cpu", "cuda"}: + raise ValueError( + f"The attention example supports only CPU or CUDA devices, got {device.type!r}" + ) + if dtype_name not in DTYPE_NAMES: + raise ValueError(f"Unsupported compute dtype {dtype_name!r}") + dtype = torch.float32 if dtype_name == "float32" else torch.bfloat16 + + if backend in FLASH_BACKENDS: + if device.type != "cuda": + raise ValueError(f"{backend} requires a CUDA device; pass --device cuda[:index]") + if dtype is not torch.bfloat16: + raise ValueError(f"{backend} requires --dtype bfloat16") + if device.type == "cuda": + if not torch.cuda.is_available(): + raise ValueError(f"CUDA device {device} was requested but CUDA is unavailable") + if backend in FLASH_BACKENDS and not torch.cuda.is_bf16_supported(): + raise ValueError(f"{backend} requires CUDA hardware with BF16 support") + return device, dtype + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + + try: + device, dtype = resolve_execution( + arguments.backend, + arguments.device, + arguments.dtype, + ) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + from transformers import AutoModel, AutoTokenizer + + from fastplms.attention import clear_flex_attention_caches + + model = ( + AutoModel.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + attn_implementation=arguments.backend, + dtype=dtype, + ) + .to(device) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + ) + model.set_attn_implementation(arguments.backend) + configured_state = attention_configuration_snapshot(model) + optimized_output = run_optimized_attention_example( + model, + tokenizer, + ["MSTNPKPQRKTKRNT", "MKTII"], + ) # last_hidden_state: (2, l, d) + if attention_configuration_snapshot(model) != configured_state: + raise RuntimeError("The optimized attention call mutated the configured backend") + fallback_output, warning_messages = run_attention_example( + model, + tokenizer, + ["MSTNPKPQRKTKRNT", "MKTII"], + ) # last_hidden_state: (2, l, d) + if attention_configuration_snapshot(model) != configured_state: + raise RuntimeError("The output_attentions eager fallback mutated the configured backend") + print("optimized", tuple(optimized_output.last_hidden_state.shape)) + print("fallback", tuple(fallback_output.last_hidden_state.shape)) + print( + "execution", + f"backend={arguments.backend}", + f"device={device}", + f"dtype={arguments.dtype}", + ) + for message in warning_messages: + print("warning", message) + + clear_flex_attention_caches() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/binder_design_fastplms.py b/examples/binder_design_fastplms.py new file mode 100644 index 0000000..01ca61b --- /dev/null +++ b/examples/binder_design_fastplms.py @@ -0,0 +1,1951 @@ +"""Local FastPLMs binder-design research example. + +This is a FastPLMs-only variant of the Biohub ESMFold2 binder design workflow. +It uses FastPLMs ESMFold2 experimental checkpoints for folding and FastPLMs +ESM++ checkpoints for the masked-LM regularizer. +""" + +from __future__ import annotations + +import argparse +import hashlib +import inspect +import json +import logging +import math +import os +import platform +import random +import re +import secrets +import sys +import torch +import torch.nn.functional as F +import torch.optim as optim +from collections.abc import Mapping +from contextlib import suppress +from dataclasses import dataclass +from functools import cache +from importlib import metadata +from pathlib import Path +from typing import Any +from tqdm.auto import tqdm +from transformers import AutoModel, AutoModelForMaskedLM + +from fastplms.models.esm_plusplus.modeling_esm_plusplus import EsmSequenceTokenizer +from fastplms.models.esmfold2 import seed_context +from fastplms.models.esmfold2.esmfold2_constants import ( + ELEMENT_NUMBER_TO_SYMBOL, + PROTEIN_1TO3, + PROTEIN_3TO1, + RES_TYPE_TO_CCD, +) + + +logger = logging.getLogger(__name__) + + +TOKENS = ["", "-"] + [RES_TYPE_TO_CCD[i] for i in range(2, 33)] +ELEMENTS = ["X"] * (max(ELEMENT_NUMBER_TO_SYMBOL) + 1) +ELEMENTS[0] = "" +for _atomic_num, _symbol in ELEMENT_NUMBER_TO_SYMBOL.items(): + ELEMENTS[_atomic_num] = _symbol[:1] + _symbol[1:].lower() +TOKEN_IDS = {token: idx for idx, token in enumerate(TOKENS)} +AA_DIMS = 20 +CYS_IDX = TOKEN_IDS[PROTEIN_1TO3["C"]] - 2 +MUTABLE_TOKEN = "#" +BinderPromptStr = str + +LOSS_WEIGHTS = {"intra_contact": 0.5, "inter_contact": 0.5, "glob": 0.2} +DEFAULT_STEPS = 150 +DEFAULT_LOG_INTERVAL = 5 +DEFAULT_LEARNING_RATE = 0.1 +DEFAULT_TEMPERATURE_MIN = 1e-2 +DEFAULT_ESMC_MASK_FRACTION = 0.15 +DEFAULT_SELECTION_TOP_K = 84 +MINIBINDER_PI_CUTOFF = 6.0 +DEFAULT_CONSENSUS_IPTM_THRESHOLD = 0.9 + + +@dataclass(frozen=True) +class PromptFactory: + name: str + template: str + length_ranges: dict[str, tuple[int, int]] + is_antibody: bool + + def sample(self, seed: int) -> BinderPromptStr: + rng = random.Random(seed) + sampled_lengths = { + key: MUTABLE_TOKEN * rng.randint(low, high) + for key, (low, high) in self.length_ranges.items() + } + return self.template.format(**sampled_lengths) + + +BINDER_PROMPT_FACTORIES = { + "minibinder": PromptFactory( + name="minibinder", + template="{seq}", + length_ranges={"seq": (60, 200)}, + is_antibody=False, + ), + "trastuzumab_framework_vhvl": PromptFactory( + name="trastuzumab_framework_vhvl", + template=( + "EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}YIHWVRQAPGKGLEWVARI{hcdr2}" + "TRYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCSR{hcdr3}WGQGTLVTVSS" + "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" + "{lcdr2}GVPSRFSGSRSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" + ), + length_ranges={ + "hcdr1": (7, 9), + "hcdr2": (5, 6), + "hcdr3": (9, 15), + "lcdr1": (11, 16), + "lcdr2": (7, 7), + "lcdr3": (9, 9), + }, + is_antibody=True, + ), + "atezolizumab_framework_vhvl": PromptFactory( + name="atezolizumab_framework_vhvl", + template=( + "EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}WIHWVRQAPGKGLEWVAWI{hcdr2}" + "TYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCAR{hcdr3}WGQGTLVTVSS" + "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" + "{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" + ), + length_ranges={ + "hcdr1": (7, 9), + "hcdr2": (5, 6), + "hcdr3": (9, 15), + "lcdr1": (11, 16), + "lcdr2": (7, 7), + "lcdr3": (9, 9), + }, + is_antibody=True, + ), + "ocankitug_framework_vhvl": PromptFactory( + name="ocankitug_framework_vhvl", + template=( + "QVQLVQSGAEVKKPGSSVKVSCKAS{hcdr1}WMHWVRQAPGQGLEWMGII{hcdr2}" + "TSLNQKFQGRVTITADTSTSTAYMELSSLRSEDTAVYYCAR{hcdr3}WGQGTLVTVSS" + "GGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY" + "{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK" + ), + length_ranges={ + "hcdr1": (7, 9), + "hcdr2": (5, 6), + "hcdr3": (8, 14), + "lcdr1": (11, 16), + "lcdr2": (7, 7), + "lcdr3": (9, 9), + }, + is_antibody=True, + ), +} + +TARGET_SEQUENCES = { + "cd45": ( + "GSPGEPQIIFCRSEAAHQGVITWNPPQRSFHNFTLCYIKETEKDCLNLDKNLIKYDLQNLKPYT" + "KYVLSLHAYIIAKVQRNGSAAMCHFTTKSAPPSQVWNMTVSMTSDNSMHVKCRPPRDRNGPHE" + "RYHLEVEAGNTLVRNESHKNCDFRVKDLQYSTDYTFKAYFHNGDYPGEPFILHHSTSY" + ), + "ctla4": ( + "MHVAQPAVVLASSRGIASFVCEYASPGKATEVRVTVLRQADSQVTEVCAATYMMGNELTFLDDSI" + "CTGTSSGNQVNLTIQGLRAMDTGLYICKVELMYPPPYYLGIGNGTQIYVIDPE" + ), + "egfr": ( + "RKVCNGIGIGEFKDSLSINATNIKHFKNCTSISGDLHILPVAFRGDSFTHTPPLDPQELDILKTV" + "KEITGFLLIQAWPENRTDLHAFENLEIIRGRTKQHGQFSLAVVSLNITSLGLRSLKEISDGDV" + "IISGNKNLCYANTINWKKLFGTSGQKTKIISNRGENSCKATGQVCHALCSPEGCWGPEPRDCV" + ), + "pd-l1": ( + "AFTVTVPKDLYVVEYGSNMTIECKFPVEKQLDLAALIVYWEMEDKNIIQFVHGEEDLKVQHSSYR" + "QRARLLKDQLSLGNAALQITDVKLQDAGVYRCMISYGGADYKRITVKVNA" + ), + "pdgfr": ( + "GFLPNDAEELFIFLTEITEITIPCRVTDPQLVVTLHEKKGDVALPVPYDHQRGFSGIFEDRSYIC" + "KTTIGDREVDSDAYYVYRLQVSSINVSVNAVQTVVRQGENITLMCIVIGNEVVNFEWTYPRKES" + "GRLVEPVTDFLLDMPYHIRSILHIPSAELEDSGTYTCNVTESVNDHQDEKAINITVVE" + ), +} + + +def _repo_name(name: str) -> str: + if "/" in name: + return name + return f"Synthyra/{name}" + + +_IMMUTABLE_HUB_REVISION = re.compile(r"[0-9a-fA-F]{40}\Z") + + +@cache +def _registered_fast_revisions() -> dict[str, str]: + from fastplms.registry import get_model_registry + + return {model.fast.repo_id: model.fast.revision for model in get_model_registry().values()} + + +def _normalize_model_revisions( + revisions: Mapping[str, str] | None, +) -> dict[str, str]: + normalized: dict[str, str] = {} + for raw_name, raw_revision in (revisions or {}).items(): + repo_id = _repo_name(str(raw_name).strip()) + revision = str(raw_revision).strip().lower() + namespace, separator, repository = repo_id.partition("/") + if ( + not separator + or not namespace + or not repository + or "/" in repository + or _IMMUTABLE_HUB_REVISION.fullmatch(revision) is None + ): + raise ValueError( + "Model revisions must map a repository to an immutable 40-character " + f"Git commit; got {raw_name!r}={raw_revision!r}." + ) + previous = normalized.get(repo_id) + if previous is not None and previous != revision: + raise ValueError( + f"Conflicting revisions were supplied for {repo_id!r}: " + f"{previous!r} and {revision!r}." + ) + normalized[repo_id] = revision + return normalized + + +def _parse_model_revision_args(values: list[str] | None) -> dict[str, str]: + parsed: dict[str, str] = {} + for value in values or (): + repo_id, separator, revision = value.partition("=") + if not separator or not repo_id.strip() or not revision.strip(): + raise ValueError( + f"--model-revision must use REPO=40_CHARACTER_COMMIT syntax; got {value!r}." + ) + normalized = _normalize_model_revisions({repo_id: revision}) + normalized_repo, normalized_revision = next(iter(normalized.items())) + previous = parsed.get(normalized_repo) + if previous is not None and previous != normalized_revision: + raise ValueError( + f"Conflicting revisions were supplied for {normalized_repo!r}: " + f"{previous!r} and {normalized_revision!r}." + ) + parsed[normalized_repo] = normalized_revision + return parsed + + +def _resolve_model_source( + model_name: str, + revisions: Mapping[str, str], +) -> tuple[str, str]: + repo_id = _repo_name(model_name) + revision = revisions.get(repo_id) or _registered_fast_revisions().get(repo_id) + if revision is None: + raise ValueError( + f"Custom model repository {repo_id!r} requires an immutable revision. " + f"Pass --model-revision {repo_id}=<40-character-commit>." + ) + if _IMMUTABLE_HUB_REVISION.fullmatch(revision) is None: + raise ValueError( + f"Resolved revision for {repo_id!r} is not an immutable Git commit: {revision!r}." + ) + return repo_id, revision.lower() + + +def _configure_offline_mode(local_files_only: bool) -> None: + if local_files_only: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def _record_model_load_identity( + model: Any, + *, + repo_id: str, + revision: str, + local_files_only: bool, +) -> None: + model.__dict__["_fastplms_binder_load_identity"] = { + "repo_id": repo_id, + "requested_revision": revision, + "local_files_only": local_files_only, + } + + +def build_initial_soft_sequence_logits(sequence: str, batch_size: int) -> torch.Tensor: + if all(aa == MUTABLE_TOKEN for aa in sequence): + logits = 0.01 * torch.randn( + [batch_size, len(sequence), AA_DIMS] + ) # (b, l, a) + logits[:, :, CYS_IDX] = -1e6 # (b, l, a) + else: + logits = torch.zeros([batch_size, len(sequence), AA_DIMS]) # (b, l, a) + for i, aa in enumerate(sequence): + if aa == MUTABLE_TOKEN: + logits[:, i, :] = 0.01 * torch.randn(batch_size, AA_DIMS) # (b, a) + logits[:, i, CYS_IDX] = -1e6 # (b,) + else: + if aa not in PROTEIN_1TO3: + raise ValueError( + f"Unsupported fixed binder residue {aa!r} at position {i}; " + "use an uppercase canonical amino acid or '#'." + ) + token_id = TOKEN_IDS[PROTEIN_1TO3[aa]] + logits[:, i, token_id - 2] = 10.0 # (b,) + return logits.requires_grad_(True) # (b, l, a) + + +def build_gradient_mask(sequence: str, batch_size: int) -> torch.Tensor: + mask = torch.ones([batch_size, len(sequence), AA_DIMS]) # (b, l, a) + fixed_positions = [i for i, aa in enumerate(sequence) if aa != MUTABLE_TOKEN] + mask[:, fixed_positions, :] = 0.0 # (b, l, a) + mask[:, :, CYS_IDX] = 0.0 # (b, l, a) + return mask # (b, l, a) + + +def sequence_to_one_hot(sequence: str, device: torch.device | str = "cuda") -> torch.Tensor: + target_index = [TOKEN_IDS[PROTEIN_1TO3[letter]] for letter in sequence] + one_hot = F.one_hot( + torch.tensor(target_index), num_classes=len(TOKENS) + ) # (l, v_f) + return one_hot.to(device).unsqueeze(0).float() # (1, l, v_f) + + +def get_mid_points() -> torch.Tensor: + boundaries = torch.linspace(2, 52.0, 127) # (127,) + lower = torch.tensor([1.0]) # (1,) + upper = torch.tensor([57.0]) # (1,) + exp_boundaries = torch.cat((lower, boundaries, upper)) # (129,) + return (exp_boundaries[:-1] + exp_boundaries[1:]) / 2 # (z=128,) + + +def binned_entropy(dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float) -> torch.Tensor: + # dgram: (..., z); bin_distance: (z,) + bin_mask = ~(bin_distance < cutoff) # (z,) + masked_dgram = dgram - (1e7 * bin_mask) # (..., z) + px = torch.softmax(masked_dgram, dim=-1) # (..., z) + log_px = torch.log_softmax(dgram, dim=-1) # (..., z) + return -(px * log_px).sum(-1) # (...) + + +def masked_min_k(x: torch.Tensor, mask: torch.Tensor, k: int) -> torch.Tensor: + # x/mask: (..., n) + mask = mask.bool() # (..., n) + y = torch.sort(torch.where(mask, x, float("nan")))[0] # (..., n) + k_mask = (torch.arange(y.shape[-1]).to(y.device) < k) & ( + ~torch.isnan(y) + ) # (..., n) + return torch.where(k_mask, y, 0).sum(-1) / (k_mask.sum(-1) + 1e-8) # (...) + + +def masked_average(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + # x/mask: (..., n) + mask = mask.bool() # (..., n) + return torch.where(mask, x, 0).sum(-1) / ( + torch.where(mask, 1, 0).sum(-1) + 1e-8 + ) # (...) + + +def compute_contact_loss( + distogram_logits: torch.Tensor, + bin_distance: torch.Tensor, + num_contacts: int, + min_sep: int, + cutoff: float, + chain_mask: torch.Tensor, + binder_mask: torch.Tensor, +) -> torch.Tensor: + # distogram_logits: (b, l, l, z); bin_distance: (z,) + # chain_mask/binder_mask: (l,) or broadcast-compatible with (b, l, l) + con_loss = binned_entropy(distogram_logits, bin_distance, cutoff) # (b, l, l) + position = torch.arange(distogram_logits.shape[1]) # (l,) + p_dist = position[:, None] - position[None, :] # (l, l) + if min_sep > 0: + separation_mask = (torch.abs(p_dist) >= min_sep).to( + distogram_logits.device + ) # (l, l) + binder_mask = torch.logical_and( + separation_mask, binder_mask + ) # (l, l) + per_residue = masked_min_k(con_loss, mask=binder_mask, k=num_contacts).to( + distogram_logits.device + ) # (b, l) + return masked_average(per_residue, mask=chain_mask).to( + distogram_logits.device + ) # (b,) + + +def compute_intra_contact_loss( + distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor +) -> torch.Tensor: + full_len = distogram_logits.shape[1] + is_binder = torch.ones(full_len, device=distogram_logits.device) # (L,) + is_binder[:-binder_length] *= 0.0 # (L,) + return compute_contact_loss( + distogram_logits, + bin_distance, + num_contacts=2, + min_sep=9, + cutoff=14.0, + chain_mask=is_binder, + binder_mask=is_binder, + ) + + +def compute_inter_contact_loss( + distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor +) -> torch.Tensor: + full_len = distogram_logits.shape[1] + is_binder = torch.ones(full_len, device=distogram_logits.device) # (L,) + is_binder[:-binder_length] *= 0.0 # (L,) + return compute_contact_loss( + distogram_logits, + bin_distance, + num_contacts=1, + min_sep=0, + cutoff=22.0, + chain_mask=1 - is_binder, + binder_mask=is_binder, + ) + + +def compute_globularity_loss( + distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor +) -> torch.Tensor: + # distogram_logits: (b, l, l, z) + binder_disto = distogram_logits[ + :, -binder_length:, -binder_length:, : + ] # (b, l, l, z) + n = binder_disto.shape[1] + disto_probs = torch.softmax(binder_disto, dim=-1) # (b, l, l, z) + bin_distance = bin_distance.clamp(max=27) # (z,) + e_sq_dist = torch.sum( + disto_probs * torch.square(bin_distance), dim=-1 + ) # (b, l, l) + sum_sq_dist = torch.sum( + torch.tril(e_sq_dist, diagonal=-1), dim=(1, 2) + ) # (b,) + rg_term = torch.sqrt(sum_sq_dist / (n * n)) # (b,) + rg_th = 2.38 * (n**0.365) + return F.elu(rg_term - rg_th) # (b,) + + +def compute_structure_losses( + distogram_logits: torch.Tensor, binder_length: int +) -> dict[str, torch.Tensor]: + # distogram_logits: (b, l, l, z) + bin_distance = get_mid_points().to(distogram_logits.device) # (z,) + losses: dict[str, torch.Tensor] = {} + losses["intra_contact_loss"] = compute_intra_contact_loss( + distogram_logits, binder_length, bin_distance + ) # (b,) + losses["inter_contact_loss"] = compute_inter_contact_loss( + distogram_logits, binder_length, bin_distance + ) # (b,) + losses["glob_loss"] = compute_globularity_loss( + distogram_logits, binder_length, bin_distance + ) # (b,) + batch = distogram_logits.size(0) + total = torch.tensor([0.0] * batch, device=distogram_logits.device) # (b,) + total = total + LOSS_WEIGHTS["intra_contact"] * losses["intra_contact_loss"] # (b,) + total = total + LOSS_WEIGHTS["inter_contact"] * losses["inter_contact_loss"] # (b,) + total = total + LOSS_WEIGHTS["glob"] * losses["glob_loss"] # (b,) + losses["total_loss"] = total # (b,) + return losses + + +def _binding_confidence_entropy( + dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float +) -> torch.Tensor: + # dgram: (..., z); bin_distance: (z,) + probs = torch.softmax(dgram, dim=-1) # (..., z) + cutoff_mask = bin_distance < cutoff # (z,) + p_cut = probs[..., cutoff_mask] # (..., z_c) + p_cut = p_cut / (p_cut.sum(-1, keepdim=True) + 1e-8) # (..., z_c) + return -(p_cut * torch.log(p_cut + 1e-10)).sum(-1) # (...) + + +def _entropy_to_confidence(mean_entropy: float) -> float: + return float(max(0.0, min(1.0, 1.0 - mean_entropy / math.log(51)))) + + +def _cdr_indices(binder_sequence: str) -> list[int]: + from abnumber import Chain + + chains = list( + Chain.multiple_domains( + binder_sequence, + scheme="chothia", + allowed_species=None, + use_anarcii=True, + ) + ) + if not chains: + raise ValueError("AbNumber did not identify an antibody domain in the binder sequence.") + + indices: list[int] = [] + search_start = 0 + for chain in chains: + domain_sequence = str(chain.seq) + domain_start = binder_sequence.find(domain_sequence, search_start) + if domain_start < 0: + raise RuntimeError( + "AbNumber returned an antibody domain that cannot be aligned back to " + "the supplied binder sequence." + ) + indices.extend( + domain_start + offset + for offset, (position, _residue) in enumerate(chain) + if position.is_in_cdr() + ) + search_start = domain_start + len(domain_sequence) + if not indices: + raise ValueError("AbNumber identified antibody domains but no Chothia CDR residues.") + return indices + + +def compute_distogram_iptm_proxy( + distogram_logits: torch.Tensor, + target_length: int, + binder_sequence: str, + is_antibody: bool, + cdr_indices: list[int] | None = None, +) -> dict[str, float]: + # distogram_logits: (b, l, l, z) or (l, l, z) + if distogram_logits.ndim == 4: + distogram_logits = distogram_logits[0] # (l, l, z) + binder_length = len(binder_sequence) + expected_length = target_length + binder_length + if distogram_logits.ndim != 3 or distogram_logits.shape[:2] != ( + expected_length, + expected_length, + ): + raise ValueError( + "Distogram logits must have shape " + f"({expected_length}, {expected_length}, bins); got " + f"{tuple(distogram_logits.shape)}." + ) + + bin_distance = get_mid_points().to(distogram_logits.device) # (z,) + binder_start = target_length + + def _mean_lowest_k(entropies: torch.Tensor, k: int) -> float: + # entropies: (...) + sorted_entropies, _ = torch.sort(entropies.reshape(-1)) # (n,) + k = min(k, sorted_entropies.numel()) + return float(sorted_entropies[:k].mean()) + + binder_to_target_entropy = _binding_confidence_entropy( + distogram_logits[binder_start:, :target_length, :], bin_distance, cutoff=22.0 + ) # (l_b, l_t) + distogram_iptm_proxy = _entropy_to_confidence( + _mean_lowest_k(binder_to_target_entropy, k=binder_length) + ) + + if not is_antibody: + cdr_distogram_iptm_proxy = float("nan") + else: + if cdr_indices is None: + cdr_indices = _cdr_indices(binder_sequence) + cdr_rows = [binder_start + i for i in cdr_indices] + cdr_to_target_entropy = _binding_confidence_entropy( + distogram_logits[cdr_rows, :target_length, :], bin_distance, cutoff=22.0 + ) # (l_cdr, l_t) + cdr_distogram_iptm_proxy = _entropy_to_confidence( + _mean_lowest_k(cdr_to_target_entropy, k=len(cdr_indices)) + ) + return { + "distogram_iptm_proxy": distogram_iptm_proxy, + "cdr_distogram_iptm_proxy": cdr_distogram_iptm_proxy, + } + + +_ATOM_FEATURE_DIMS = { + "ref_pos": 1, + "ref_element": 1, + "ref_charge": 1, + "ref_atom_name_chars": 1, + "ref_space_uid": 1, + "atom_attention_mask": 1, + "atom_to_token": 1, + "is_resolved": 1, + "gt_coords": 2, +} + + +def _resize_tensor(tensor: torch.Tensor, *, dim: int, size: int) -> torch.Tensor: + # tensor: (..., n_dim, ...) + current = tensor.shape[dim] + if current > size: + raise ValueError( + f"Refusing to truncate atom features from {current} to {size}; " + "batch padding must preserve every prepared atom." + ) + if current == size: + return tensor + pad_shape = list(tensor.shape) + pad_shape[dim] = size - current + pad = torch.zeros( + pad_shape, dtype=tensor.dtype, device=tensor.device + ) # (..., size - n_dim, ...) + return torch.cat((tensor, pad), dim=dim) # (..., size, ...) + + +def _prepared_atom_count(features: dict[str, torch.Tensor]) -> int: + sizes = {features[key].shape[dim] for key, dim in _ATOM_FEATURE_DIMS.items() if key in features} + if not sizes: + raise ValueError("Prepared ESMFold2 features contain no atom-axis tensors.") + if len(sizes) != 1: + raise ValueError(f"Prepared ESMFold2 atom axes disagree: {sorted(sizes)}") + return sizes.pop() + + +def _pad_prepared_atom_features( + prepared: list[tuple[dict[str, torch.Tensor], list[Any]]], +) -> list[tuple[dict[str, torch.Tensor], list[Any]]]: + """Pad a prepared batch to its largest atom table without truncation.""" + + if not prepared: + raise ValueError("At least one prepared ESMFold2 input is required.") + largest = max(_prepared_atom_count(features) for features, _ in prepared) + max_atoms = ((largest + 31) // 32) * 32 + padded: list[tuple[dict[str, torch.Tensor], list[Any]]] = [] + for features, chain_infos in prepared: + resized = dict(features) # tensor values retain model-defined feature shapes + for key, dim in _ATOM_FEATURE_DIMS.items(): + if key in resized: + resized[key] = _resize_tensor( + resized[key], dim=dim, size=max_atoms + ) # atom axis -> max_atoms + padded.append((resized, chain_infos)) + return padded + + +def prepare_esmfold2_tensors( + model: Any, + input_data: Any, + max_atoms: int | None = None, + seed: int | None = None, +) -> tuple[dict[str, torch.Tensor], list[Any]]: + features, chain_infos = model.prepare_structure_input( + input_data, seed=seed + ) # tensor values: model-defined feature shapes + if max_atoms is not None: + for key, dim in _ATOM_FEATURE_DIMS.items(): + if key in features: + features[key] = _resize_tensor( + features[key], dim=dim, size=max_atoms + ) # atom axis -> max_atoms + return features, chain_infos + + +def _filter_model_forward_kwargs( + model: Any, kwargs: dict[str, torch.Tensor | int | bool | None] +) -> dict[str, torch.Tensor | int | bool | None]: + signature = inspect.signature(model.forward) + parameters = signature.parameters + accepts_kwargs = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ) + if accepts_kwargs: + return kwargs + return {key: value for key, value in kwargs.items() if key in parameters} + + +def fold_and_get_distogram( + model: Any, + target_seq: str, + target_one_hot: torch.Tensor, + design: torch.Tensor, + num_loops: int = 0, + num_sampling_steps: int = 1, + calculate_confidence: bool = False, + seed: int | None = None, +) -> dict[str, Any]: + # target_one_hot: (1, l_t, v_f); design: (b, l_b, a) + padding = (2, 11) + padded_design = F.pad( + design, padding, mode="constant", value=0 + ) # (b, l_b, v_f) + + token_lists = torch.argmax(padded_design, dim=-1) # (b, l_b) + designed_seq = [ + [PROTEIN_3TO1[TOKENS[int(tkn.item())]] for tkn in token_list] for token_list in token_lists + ] + seq_list = [target_seq + "|" + "".join(seq) for seq in designed_seq] + prepared_inputs: list[tuple[dict[str, torch.Tensor], list[Any]]] = [] + for seq in seq_list: + target, binder = seq.split("|") + input_types = model.input_types + inputs_raw = input_types.StructurePredictionInput( + sequences=[ + input_types.ProteinInput(id="A", sequence=target, msa=None), + input_types.ProteinInput(id="B", sequence=binder, msa=None), + ] + ) + prepared_inputs.append(prepare_esmfold2_tensors(model, inputs_raw, seed=seed)) + + prepared_inputs = _pad_prepared_atom_features( + prepared_inputs + ) # tensor values share a padded atom axis + inputs_list = [features for features, _ in prepared_inputs] + chain_info_list = [chain_infos for _, chain_infos in prepared_inputs] + + inputs = { + key: torch.cat([inp[key] for inp in inputs_list], dim=0).to(design.device) + for key in inputs_list[0] + } # tensor values: (b, ...) with model-defined trailing feature shapes + inputs["res_type_soft"] = torch.cat( + (target_one_hot.repeat(design.size(0), 1, 1), padded_design), dim=1 + ) # (b, l_t + l_b, v_f) + + forward_kwargs: dict[str, torch.Tensor | int | bool | None] = dict(inputs) + forward_kwargs.update( + { + "num_diffusion_samples": 1, + "num_sampling_steps": num_sampling_steps, + "num_loops": num_loops, + "calculate_confidence": calculate_confidence, + "seed": seed, + } + ) + + with seed_context(seed): + output = model( + **_filter_model_forward_kwargs(model, forward_kwargs) + ) # distogram_logits: (b, l_t + l_b, l_t + l_b, z) + + result: dict[str, Any] = { + "distogram_logits": output["distogram_logits"], + "inputs": inputs, + "chain_info_list": chain_info_list, + "output": output, + "seq_list": seq_list, + } + if calculate_confidence: + for key in ("ptm", "iptm", "plddt"): + if key in output: + result[key] = output[key] + return result + + +@cache +def _folding_trunk_to_lm_aa_vocab_matrix(device: torch.device) -> torch.Tensor: + three_to_one_map = {v: k for k, v in PROTEIN_1TO3.items()} + ft_aas = [three_to_one_map[tok_3letter] for tok_3letter in TOKENS[2:22]] + tokenizer = EsmSequenceTokenizer() + lm_vocab = sorted(tokenizer.vocab.items(), key=lambda x: x[1]) + lm_aas = [lm_vocab[i][0] for i in range(4, 24)] + ft_to_lm_aa_matrix = torch.zeros(20, 20) # (a, a) + for ft_idx, ft_aa in enumerate(ft_aas): + lm_idx = lm_aas.index(ft_aa) + ft_to_lm_aa_matrix[ft_idx, lm_idx] = 1 # scalar assignment; (a, a) + return ft_to_lm_aa_matrix.to(device=device) # (a, a) + + +def _one_hot_from_probs(probs: torch.Tensor) -> torch.Tensor: + # probs: (..., a) + return F.one_hot(torch.argmax(probs, dim=-1), num_classes=probs.size(-1)).to( + probs.dtype + ) # (..., a) + + +def _straight_through(discrete: torch.Tensor, continuous: torch.Tensor) -> torch.Tensor: + # discrete/continuous: (..., a) + return continuous + (discrete - continuous).detach() # (..., a) + + +def compute_fastplms_pseudoperplexity_nll( + lm_model: Any, + binder_design: torch.Tensor, + score_mask: torch.Tensor, + batch_size: int = 4, + n_passes: int = 4, + mask_fraction: float = DEFAULT_ESMC_MASK_FRACTION, +) -> torch.Tensor: + # binder_design: (b, l, a); score_mask: (b, l) or (l,) + device = binder_design.device + lm_vocab_size = lm_model.config.vocab_size + model_dtype = lm_model.embed.weight.dtype + + target_esm = binder_design @ _folding_trunk_to_lm_aa_vocab_matrix( + device + ) # (b, l, a) + input_esm = _straight_through( + _one_hot_from_probs(target_esm), target_esm + ) # (b, l, a) + input_ids = torch.zeros( + (binder_design.size(0), binder_design.size(1) + 2, lm_vocab_size), + dtype=model_dtype, + device=device, + ) # (b, l + 2, v) + tokenizer = lm_model.tokenizer + input_ids[:, 0, tokenizer.cls_token_id] = 1 # (b, l + 2, v) + input_ids[:, -1, tokenizer.eos_token_id] = 1 # (b, l + 2, v) + input_ids[:, 1:-1, 4:24] = input_esm.to(model_dtype) # (b, l, a) + + if score_mask.ndim == 1: + score_mask = score_mask.unsqueeze(0).expand( + binder_design.size(0), -1 + ) # (b, l) + elif score_mask.shape != binder_design.shape[:2]: + raise ValueError( + f"Expected score_mask with shape {(binder_design.size(0), binder_design.size(1))}, " + f"got {tuple(score_mask.shape)}" + ) + score_mask = score_mask.to(device=device, dtype=torch.bool) # (b, l) + + mask_token = torch.zeros(lm_vocab_size, dtype=model_dtype, device=device) # (v,) + mask_token[tokenizer.mask_token_id] = 1 # (v,) + losses = [] + for batch_idx in range(binder_design.size(0)): + position_indices = ( + score_mask[batch_idx].nonzero(as_tuple=False).flatten() + ) # (n,) + num_positions = int(position_indices.numel()) + if num_positions == 0: + raise ValueError("Pseudoperplexity score mask selected zero positions.") + + num_masked = max(1, math.ceil(mask_fraction * num_positions)) + random_scores = torch.rand((n_passes, num_positions), device=device) # (p, n) + masked_offsets = random_scores.topk( + num_masked, dim=-1, largest=False + ).indices # (p, m) + pass_masks = torch.zeros( + (n_passes, binder_design.size(1)), dtype=torch.bool, device=device + ) # (p, l) + pass_masks[ + torch.arange(n_passes, device=device)[:, None], + position_indices[masked_offsets], + ] = True # (p, m) selected within (p, l) + + masked_sequences = input_ids[batch_idx : batch_idx + 1].repeat( + n_passes, 1, 1 + ) # (p, l + 2, v) + mask_rows, mask_cols = pass_masks.nonzero( + as_tuple=True + ) # each: (p * m,) + masked_sequences[mask_rows, mask_cols + 1] = mask_token # (p * m, v) + + target_weights = target_esm[batch_idx] # (l, a) + masked_nlls = [] + for start in range(0, n_passes, batch_size): + stop = min(start + batch_size, n_passes) + chunk = masked_sequences[start:stop] # (p_c, l + 2, v) + with torch.autocast( + device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda" + ): + hidden = lm_model.transformer( + x=chunk @ lm_model.embed.weight.to(chunk.dtype), + attention_mask=None, + output_hidden_states=False, + output_attentions=False, + ).last_hidden_state # (p_c, l + 2, d) + logits = lm_model.sequence_head(hidden) # (p_c, l + 2, v) + log_probs = logits.log_softmax(dim=-1)[:, 1:-1, 4:24] # (p_c, l, a) + nlls = -( + log_probs * target_weights.to(log_probs.dtype).unsqueeze(0) + ).sum(dim=-1) # (p_c, l) + masked_nlls.append(nlls[pass_masks[start:stop]]) # (p_c * m,) + losses.append(torch.cat(masked_nlls, dim=0).mean()) # () + return torch.stack(losses, dim=0) # (b,) + + +def normalized_gradient_tensor(grad: torch.Tensor, gradient_mask: torch.Tensor) -> torch.Tensor: + # grad/gradient_mask: (b, l, a) + masked_grad = grad * gradient_mask # (b, l, a) + index_has_nonzero_grad = torch.square(masked_grad).sum(-1) > 0 # (b, l) + eff_l = index_has_nonzero_grad.sum(-1) # (b,) + grad_norm = torch.linalg.norm(masked_grad, axis=(-1, -2)) # (b,) + normalized_grad = (masked_grad / (grad_norm[:, None, None] + 1e-7)) * torch.sqrt( + eff_l[:, None, None] + ) # (b, l, a) + return normalized_grad * gradient_mask # (b, l, a) + + +def _tensor_mean_float(tensor: torch.Tensor) -> float: + # tensor: (...) + return float(tensor.detach().float().mean().cpu().item()) + + +def _metric_float(output: dict[str, Any], key: str) -> float | None: + if key not in output: + return None + value = output[key] + if value is None: + return None + if isinstance(value, torch.Tensor): + # value: (...) + return float(value.detach().float().mean().cpu().item()) + return float(value) + + +def _require_fresh_output_directory(output_dir: str | Path | None) -> Path | None: + if output_dir is None: + return None + result_dir = Path(output_dir) + if result_dir.exists() or result_dir.is_symlink(): + raise FileExistsError( + f"Binder output directory {result_dir} already exists. Choose a new path; " + "existing, partial, and empty run directories are never reused." + ) + return result_dir + + +def _reserve_output_directory(output_dir: str | Path | None) -> Path | None: + result_dir = _require_fresh_output_directory(output_dir) + if result_dir is None: + return None + result_dir.parent.mkdir(parents=True, exist_ok=True) + try: + result_dir.mkdir() + except FileExistsError as error: + raise FileExistsError( + f"Binder output directory {result_dir} was created by another run. Choose a new path." + ) from error + return result_dir + + +def _validate_design_sequence(name: str, sequence: str, *, allow_mutable: bool) -> None: + if not sequence: + raise ValueError(f"{name} must not be empty.") + allowed = set(PROTEIN_1TO3) + if allow_mutable: + allowed.add(MUTABLE_TOKEN) + invalid = [(index, residue) for index, residue in enumerate(sequence) if residue not in allowed] + if invalid: + index, residue = invalid[0] + mutable_note = " or '#'" if allow_mutable else "" + raise ValueError( + f"{name} contains unsupported residue {residue!r} at position {index}; " + f"use uppercase canonical amino acids{mutable_note}." + ) + + +def design_binder( + inversion_models: dict[str, Any], + critic_models: dict[str, Any], + lm_model: Any, + target_name: str | None, + target_sequence: str | None, + binder_name: str | None, + binder_sequence: str | None, + is_antibody: bool | None, + seed: int, + batch_size: int = 1, + steps: int = DEFAULT_STEPS, + log_interval: int = DEFAULT_LOG_INTERVAL, + learning_rate: float = DEFAULT_LEARNING_RATE, + temperature_min: float = DEFAULT_TEMPERATURE_MIN, + output_dir: str | Path | None = None, + device: torch.device | str = "cuda", +) -> tuple[list[str], dict[int, dict[str, torch.Tensor]], list[dict[str, Any]]]: + if (target_name is None) == (target_sequence is None): + raise ValueError("Provide exactly one of target_name or target_sequence.") + if (binder_name is None) == (binder_sequence is None): + raise ValueError("Provide exactly one of binder_name or binder_sequence.") + if not inversion_models: + raise ValueError("At least one inversion model is required.") + if not critic_models: + raise ValueError("At least one critic model is required.") + if batch_size <= 0: + raise ValueError(f"batch_size must be positive; got {batch_size}.") + if steps <= 0: + raise ValueError(f"steps must be positive; got {steps}.") + if log_interval <= 0: + raise ValueError(f"log_interval must be positive; got {log_interval}.") + if not math.isfinite(learning_rate) or learning_rate <= 0: + raise ValueError(f"learning_rate must be finite and positive; got {learning_rate}.") + if not math.isfinite(temperature_min) or not 0 < temperature_min <= 1: + raise ValueError( + f"temperature_min must be finite and in the interval (0, 1]; got {temperature_min}." + ) + + device = torch.device(device) + if target_name is not None: + if target_name not in TARGET_SEQUENCES: + raise ValueError( + f"Unknown target_name {target_name!r}; choose one of " + f"{sorted(TARGET_SEQUENCES)} or pass target_sequence." + ) + target_sequence = TARGET_SEQUENCES[target_name] + if target_sequence is None: + raise RuntimeError("Target sequence resolution failed.") + + if binder_name is None: + if binder_sequence is None: + raise RuntimeError("Binder sequence resolution failed.") + if is_antibody is None: + is_antibody = False + else: + if binder_name not in BINDER_PROMPT_FACTORIES: + raise ValueError( + f"Unknown binder_name {binder_name!r}; choose one of " + f"{sorted(BINDER_PROMPT_FACTORIES)} or pass binder_sequence." + ) + binder_prompt_factory = BINDER_PROMPT_FACTORIES[binder_name] + if is_antibody is not None and binder_prompt_factory.is_antibody != is_antibody: + raise ValueError( + f"Binder prompt {binder_name!r} has is_antibody=" + f"{binder_prompt_factory.is_antibody}, not {is_antibody}." + ) + is_antibody = binder_prompt_factory.is_antibody + binder_sequence = binder_prompt_factory.sample(seed=seed) + if binder_sequence is None or is_antibody is None: + raise RuntimeError("Binder prompt resolution failed.") + _validate_design_sequence("target_sequence", target_sequence, allow_mutable=False) + _validate_design_sequence("binder_sequence", binder_sequence, allow_mutable=True) + mutable_binder_indices = [i for i, aa in enumerate(binder_sequence) if aa == MUTABLE_TOKEN] + binder_length = len(binder_sequence) + result_dir = _reserve_output_directory(output_dir) + target_one_hot = sequence_to_one_hot( + target_sequence, device=device + ) # (1, l_t, v_f) + + with seed_context(seed), torch.device(device): + logits = build_initial_soft_sequence_logits( + binder_sequence, batch_size=batch_size + ) # (b, l_b, a) + gradient_mask = build_gradient_mask( + binder_sequence, batch_size=batch_size + ) # (b, l_b, a) + logits = logits.to(device) # (b, l_b, a) + gradient_mask = gradient_mask.to(device) # (b, l_b, a) + + trajectory: dict[int, dict[str, torch.Tensor]] = {} + optimizer = optim.SGD([logits], lr=learning_rate) + best_iptm: list[float] = [-1.0] * batch_size + best_loss: list[float] = [float("inf")] * batch_size + best_sequences: list[str] = [""] * batch_size + best_logits: list[torch.Tensor | None] = [None] * batch_size + best_steps: list[int | None] = [None] * batch_size + model_names = list(inversion_models) + + progress = tqdm(range(steps), desc="design", dynamic_ncols=True) + for step in progress: + optimizer.zero_grad() + t = (step + 1) / steps + remaining = 0.5 * (1 + math.cos(math.pi * t)) + temperature = temperature_min + (1 - temperature_min) * remaining + + replicate_choice = random.Random(seed + step).randint(0, len(model_names) - 1) + inversion_model = inversion_models[model_names[replicate_choice]] + design = F.softmax(logits / temperature, dim=-1) # (b, l_b, a) + calculate_confidence = temperature < 0.05 + + fold_result = fold_and_get_distogram( + inversion_model, + target_sequence, + target_one_hot, + design, + num_loops=1, + num_sampling_steps=50 if calculate_confidence else 1, + calculate_confidence=calculate_confidence, + seed=seed + step, + ) + sequences: list[str] = fold_result["seq_list"] + losses = compute_structure_losses( + fold_result["distogram_logits"], binder_length + ) # each tensor: (b,) + structure_loss = losses["total_loss"] # (b,) + structure_grad = torch.autograd.grad( + structure_loss.mean(), logits + )[0] # (b, l_b, a) + + design = F.softmax(logits / temperature, dim=-1) # (b, l_b, a) + score_mask = gradient_mask.sum(dim=-1) > 0 # (b, l_b) + with seed_context(seed + step): + plm_loss = compute_fastplms_pseudoperplexity_nll( + lm_model=lm_model, + binder_design=design, + score_mask=score_mask, + batch_size=4, + n_passes=4, + ) # (b,) + plm_grad = torch.autograd.grad(plm_loss.mean(), logits)[0] # (b, l_b, a) + candidate_logits = logits.detach().clone() # (b, l_b, a) + + logits.grad = normalized_gradient_tensor(structure_grad, gradient_mask) + ( + 0.05 if is_antibody else 0.15 + ) * normalized_gradient_tensor(plm_grad, gradient_mask) # (b, l_b, a) + for group in optimizer.param_groups: + group["lr"] = learning_rate * temperature + optimizer.step() + + step_losses = { + key: value.detach().cpu() for key, value in losses.items() + } # each tensor: (b,) + step_losses["plm_loss"] = plm_loss.detach().cpu() # (b,) + step_losses["total_loss"] = ( + structure_loss + plm_loss + ).detach().cpu() # (b,) + trajectory[step] = step_losses + + iptm = fold_result.get("iptm") # (b,) or None + for batch_idx in range(batch_size): + current_loss = float(step_losses["total_loss"][batch_idx].item()) + if iptm is not None and iptm[batch_idx] is not None: + current_iptm = float(iptm[batch_idx].item()) + if current_iptm > best_iptm[batch_idx]: + best_iptm[batch_idx] = current_iptm + best_sequences[batch_idx] = sequences[batch_idx] + best_loss[batch_idx] = current_loss + best_logits[batch_idx] = candidate_logits[ + batch_idx + ].cpu() # (l_b, a) + best_steps[batch_idx] = step + elif current_loss < best_loss[batch_idx]: + best_sequences[batch_idx] = sequences[batch_idx] + best_loss[batch_idx] = current_loss + best_logits[batch_idx] = candidate_logits[ + batch_idx + ].cpu() # (l_b, a) + best_steps[batch_idx] = step + + if step % log_interval == 0: + loss_str = " ".join( + f"{key}={_tensor_mean_float(value):.4f}" for key, value in step_losses.items() + ) + logger.info("step %3d | %s T=%.4f", step, loss_str, temperature) + progress.set_postfix( + loss=f"{_tensor_mean_float(step_losses['total_loss']):.3f}", + temp=f"{temperature:.3f}", + ) + + if any(not sequence for sequence in best_sequences): + raise RuntimeError("Optimization completed without selecting every binder sequence.") + if any(value is None for value in best_logits): + raise RuntimeError("Optimization completed without retaining every selected logit tensor.") + if any(value is None for value in best_steps): + raise RuntimeError("Optimization completed without retaining every selected step.") + if result_dir is not None: + _write_trajectory(result_dir / "trajectory.jsonl", trajectory) + _write_fasta(result_dir / "best_sequences.fasta", best_sequences) + + critic_results: list[dict[str, Any]] = [] + target_length = len(target_sequence.replace("|", "")) + for batch_idx, best_seq in enumerate(best_sequences): + binder_seq = best_seq.split("|")[-1] + binder_design = sequence_to_one_hot( + binder_seq, device=device + )[..., 2:22] # (1, l_b, a) + for critic_name, critic_model in critic_models.items(): + final_fold = fold_and_get_distogram( + critic_model, + target_sequence, + target_one_hot, + binder_design, + num_loops=3, + num_sampling_steps=200, + calculate_confidence=True, + seed=seed, + ) + final_output = final_fold["output"] + final_inputs = final_fold[ + "inputs" + ] # tensor values: (1, ...) with model-defined feature shapes + chain_infos = final_fold["chain_info_list"][0] + complex_result = critic_model.input_builder.decode( + final_output, + final_inputs, + chain_infos, + num_diffusion_samples=1, + complex_id=f"{critic_name}-{batch_idx}", + ) + cif_text = critic_model.result_to_cif(complex_result) + pdb_text = critic_model.result_to_pdb(complex_result) + iptm_proxy_scores = compute_distogram_iptm_proxy( + final_fold["distogram_logits"], + target_length, + binder_seq, + is_antibody, + cdr_indices=mutable_binder_indices if is_antibody else None, + ) + iptm_value = None + if "iptm" in final_fold: + iptm_value = float(final_fold["iptm"][0].item()) + ptm_value = _metric_float(final_fold, "ptm") + mean_plddt = _metric_float(final_fold, "plddt") + + structure_stem = f"batch{batch_idx}_{critic_name.replace('/', '_')}" + logits_path = None + if result_dir is not None: + cif_path = result_dir / f"{structure_stem}.cif" + pdb_path = result_dir / f"{structure_stem}.pdb" + logits_path_obj = result_dir / f"{structure_stem}_logits.pt" + cif_path.write_text(cif_text, encoding="utf-8") + pdb_path.write_text(pdb_text, encoding="utf-8") + torch.save(best_logits[batch_idx], logits_path_obj) + logits_path = str(logits_path_obj) + + row = { + "is_antibody": is_antibody, + "critic_name": critic_name, + "batch_idx": batch_idx, + "designed_sequence": best_seq, + "binder_sequence": binder_seq, + "target_length": target_length, + "binder_length": len(binder_seq), + "final_loss": best_loss[batch_idx], + "selected_step": best_steps[batch_idx], + "ptm": ptm_value, + "iptm": iptm_value, + "mean_plddt": mean_plddt, + "pdb": pdb_text, + "cif": cif_text, + "logits_path": logits_path, + } + row.update(iptm_proxy_scores) + critic_results.append(row) + + if result_dir is not None: + _write_results_table(result_dir / "results.parquet", critic_results) + _write_official_selection_table( + result_dir / "selection.parquet", + critic_results, + required_hero_critics=tuple(critic_models), + ) + _write_run_manifest( + result_dir / "run_manifest.json", + seed=seed, + batch_size=batch_size, + steps=steps, + log_interval=log_interval, + learning_rate=learning_rate, + temperature_min=temperature_min, + target_sequence=target_sequence, + binder_sequence=binder_sequence, + is_antibody=is_antibody, + inversion_models=inversion_models, + critic_models=critic_models, + lm_model=lm_model, + device=device, + ) + return best_sequences, trajectory, critic_results + + +def _write_trajectory(path: Path, trajectory: dict[int, dict[str, torch.Tensor]]) -> None: + with path.open("w", encoding="utf-8") as handle: + for step, losses in trajectory.items(): + row = {"step": step} + for key, value in losses.items(): + row[key] = [ + float(x) for x in value.reshape(-1).tolist() + ] # (b,) -> b scalars + handle.write(json.dumps(row) + "\n") + + +def _package_version(package: str) -> str | None: + try: + return metadata.version(package) + except metadata.PackageNotFoundError: + return None + + +def _configuration_digest(config: Any) -> str | None: + if config is None or not hasattr(config, "to_dict"): + return None + payload = json.dumps(config.to_dict(), sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _model_load_identity(model: Any) -> dict[str, Any]: + identity = getattr(model, "__dict__", {}).get( + "_fastplms_binder_load_identity", + {}, + ) + return dict(identity) if isinstance(identity, dict) else {} + + +def _tokenizer_identity( + tokenizer: Any, + *, + model: Any | None = None, +) -> dict[str, Any] | None: + if tokenizer is None: + return None + config = getattr(model, "config", None) + load_identity = _model_load_identity(model) + get_vocab = getattr(tokenizer, "get_vocab", None) + vocab = get_vocab() if callable(get_vocab) else getattr(tokenizer, "vocab", None) + vocab_digest = None + if isinstance(vocab, dict): + payload = json.dumps( + sorted((str(token), int(index)) for token, index in vocab.items()), + separators=(",", ":"), + ) + vocab_digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() + init_kwargs = getattr(tokenizer, "init_kwargs", {}) + tokenizer_revision = ( + init_kwargs.get("revision") or getattr(tokenizer, "_commit_hash", None) + if isinstance(init_kwargs, dict) + else getattr(tokenizer, "_commit_hash", None) + ) + return { + "class": type(tokenizer).__name__, + "name_or_path": getattr(tokenizer, "name_or_path", None), + "repo_id": load_identity.get("repo_id") or getattr(config, "_name_or_path", None), + "requested_revision": ( + load_identity.get("requested_revision") or getattr(config, "_commit_hash", None) + ), + "hub_revision": getattr(config, "_commit_hash", None), + "weights_revision": getattr(config, "fastplms_weights_revision", None), + "runtime_revision": getattr(config, "fastplms_runtime_revision", None), + "revision": tokenizer_revision or load_identity.get("requested_revision"), + "local_files_only": load_identity.get("local_files_only"), + "vocab_size": len(vocab) if isinstance(vocab, dict) else None, + "vocab_sha256": vocab_digest, + "special_token_ids": { + name: getattr(tokenizer, f"{name}_token_id", None) + for name in ("bos", "cls", "eos", "mask", "pad", "sep", "unk") + }, + } + + +def _parameter_dtype_identity(model: Any) -> tuple[str | None, list[str], dict[str, int]]: + named_parameters = getattr(model, "named_parameters", None) + if callable(named_parameters): + parameters = (parameter for _name, parameter in named_parameters()) + else: + raw_parameters = getattr(model, "parameters", None) + parameters = iter(raw_parameters()) if callable(raw_parameters) else iter(()) + + dtype_numel: dict[str, int] = {} + for parameter in parameters: + dtype = str(parameter.dtype) + dtype_numel[dtype] = dtype_numel.get(dtype, 0) + int(parameter.numel()) + dtypes = sorted(dtype_numel) + if not dtypes: + summary = None + elif len(dtypes) == 1: + summary = dtypes[0] + else: + summary = f"mixed[{','.join(dtypes)}]" + return summary, dtypes, {dtype: dtype_numel[dtype] for dtype in dtypes} + + +def _effective_precision_identity(model: Any) -> dict[str, Any] | None: + status = getattr(model, "esmc_precision_status", None) + as_dict = getattr(status, "as_dict", None) + if callable(as_dict): + return dict(as_dict()) + return None + + +def _model_identity(name: str, model: Any) -> dict[str, Any]: + config = getattr(model, "config", None) + load_identity = _model_load_identity(model) + parameter_dtype, parameter_dtypes, parameter_dtype_numel = _parameter_dtype_identity(model) + attention_backend = next( + ( + getattr(config, field) + for field in ( + "esmc_attn_backend", + "attn_backend", + "attention_backend", + "_attn_implementation", + ) + if config is not None and getattr(config, field, None) is not None + ), + None, + ) + return { + "name": name, + "requested": name, + "resolved": getattr(config, "_name_or_path", None), + "repo_id": load_identity.get("repo_id") or getattr(config, "_name_or_path", None), + "requested_revision": ( + load_identity.get("requested_revision") or getattr(config, "_commit_hash", None) + ), + "hub_revision": getattr(config, "_commit_hash", None), + "weights_revision": getattr(config, "fastplms_weights_revision", None), + "runtime_revision": getattr(config, "fastplms_runtime_revision", None), + "revision": ( + getattr(config, "_commit_hash", None) or load_identity.get("requested_revision") + ), + "local_files_only": load_identity.get("local_files_only"), + "attention_backend": attention_backend, + "kernel_backend": getattr(model, "_kernel_backend", None), + "parameter_dtype": parameter_dtype, + "parameter_dtypes": parameter_dtypes, + "parameter_dtype_numel": parameter_dtype_numel, + "effective_precision": _effective_precision_identity(model), + "configuration_sha256": _configuration_digest(config), + } + + +def _write_run_manifest( + path: Path, + *, + seed: int, + batch_size: int, + steps: int, + log_interval: int, + learning_rate: float, + temperature_min: float, + target_sequence: str, + binder_sequence: str, + is_antibody: bool, + inversion_models: dict[str, Any], + critic_models: dict[str, Any], + lm_model: Any, + device: torch.device, +) -> None: + def sequence_hash(sequence: str) -> str: + return hashlib.sha256(sequence.encode("ascii")).hexdigest() + + manifest = { + "schema_version": 2, + "seed": seed, + "batch_size": batch_size, + "steps": steps, + "learning_rate": learning_rate, + "temperature_min": temperature_min, + "target_sequence_sha256": sequence_hash(target_sequence), + "binder_prompt_sha256": sequence_hash(binder_sequence), + "device": str(device), + "configuration": { + "batch_size": batch_size, + "steps": steps, + "log_interval": log_interval, + "optimizer": { + "class": "torch.optim.SGD", + "learning_rate": learning_rate, + }, + "temperature_min": temperature_min, + "is_antibody": is_antibody, + "loss_weights": LOSS_WEIGHTS, + "plm_batch_size": 4, + "plm_passes": 4, + "compute_dtype": ("torch.bfloat16" if device.type == "cuda" else "torch.float32"), + }, + "command": list(sys.argv), + "environment": { + "python": platform.python_version(), + "torch": torch.__version__, + "transformers": _package_version("transformers"), + "fastplms": _package_version("fastplms"), + "cuda": torch.version.cuda, + "hf_hub_offline": os.environ.get("HF_HUB_OFFLINE") == "1", + "transformers_offline": os.environ.get("TRANSFORMERS_OFFLINE") == "1", + }, + "models": { + "inversion": [_model_identity(name, model) for name, model in inversion_models.items()], + "critics": [_model_identity(name, model) for name, model in critic_models.items()], + "language_model": _model_identity( + getattr(getattr(lm_model, "config", None), "_name_or_path", "language_model"), + lm_model, + ), + }, + "tokenizer": _tokenizer_identity( + getattr(lm_model, "tokenizer", None), + model=lm_model, + ), + } + payload = json.dumps(manifest, indent=2, sort_keys=True) + "\n" + temporary = path.with_name(f".{path.name}.{os.getpid()}.{secrets.token_hex(8)}.tmp") + try: + temporary.write_text(payload, encoding="utf-8") + os.replace(temporary, path) + finally: + with suppress(FileNotFoundError): + temporary.unlink() + + +def _write_fasta(path: Path, sequences: list[str]) -> None: + with path.open("w", encoding="utf-8") as handle: + for idx, sequence in enumerate(sequences): + handle.write(f">design_{idx}\n{sequence}\n") + + +def _write_results_table(path: Path, rows: list[dict[str, Any]]) -> None: + import pandas as pd + + pd.DataFrame(rows).to_parquet(path, index=False) + + +def _binder_sequence_from_designed_sequence(designed_sequence: str) -> str: + parts = designed_sequence.split("|") + if len(parts) != 2 or not all(parts): + raise ValueError( + "designed_sequence must contain nonempty target and binder chains " + f"separated by one '|'; got {designed_sequence!r}." + ) + return parts[1] + + +def _compute_isoelectric_points(sequences: list[str]) -> list[float]: + from Bio.SeqUtils.ProtParam import ProteinAnalysis + + return [float(ProteinAnalysis(sequence).isoelectric_point()) for sequence in sequences] + + +def annotate_official_selection_scores(result_df: Any) -> Any: + """Add the official binder-design selection components to critic rows. + + Mirrors the paper Appendix A.3.1.2 and the official notebook selection cell: + minibinders with pI >= 6 are filtered and the approved critics contribute + mean iPTM. + """ + import pandas as pd + + df = result_df.copy() if isinstance(result_df, pd.DataFrame) else pd.DataFrame(result_df) + required_columns = [ + "critic_name", + "designed_sequence", + "is_antibody", + "iptm", + "distogram_iptm_proxy", + ] + missing_columns = [column for column in required_columns if column not in df.columns] + if missing_columns: + raise ValueError(f"Missing selection columns: {missing_columns}") + if df["distogram_iptm_proxy"].isna().any(): + raise ValueError("distogram_iptm_proxy must be present for every selection row") + + binder_sequences = [ + _binder_sequence_from_designed_sequence(sequence) + for sequence in df["designed_sequence"].tolist() + ] + is_antibody = df["is_antibody"].astype(bool) + df["binder_sequence"] = binder_sequences + df["isoelectric_point"] = _compute_isoelectric_points(binder_sequences) + df["passes_official_pi_filter"] = is_antibody | df["isoelectric_point"].lt(MINIBINDER_PI_CUTOFF) + df["official_iptm_score_component"] = df["iptm"] + return df + + +def select_official_designs( + result_df: Any, + top_k: int = DEFAULT_SELECTION_TOP_K, + consensus_iptm_threshold: float = DEFAULT_CONSENSUS_IPTM_THRESHOLD, + group_columns: tuple[str, ...] = ("target_name", "binder_name"), + required_hero_critics: tuple[str, ...] | None = None, +) -> Any: + """Rank candidates using the official ESM binder-design selection strategy.""" + df = annotate_official_selection_scores(result_df) + available_group_columns = [column for column in group_columns if column in df.columns] + selection_columns = [ + *available_group_columns, + "designed_sequence", + "iptm_score", + "iptm_proxy_score", + "hero_iptm_min", + "hero_iptm_median", + "hero_iptm_max", + "critic_count", + "hero_critic_count", + "required_hero_critic_count", + "batch_idx", + "binder_sequence", + "is_antibody", + "isoelectric_point", + "selection_score", + "all_hero_critics_pass", + "consensus_iptm_threshold", + ] + df = df[df["passes_official_pi_filter"]].copy() + if df.empty: + import pandas as pd + + return pd.DataFrame(columns=selection_columns) + + if required_hero_critics is None: + required_hero_critics = tuple(dict.fromkeys(df["critic_name"].dropna().tolist())) + else: + required_hero_critics = tuple(dict.fromkeys(required_hero_critics)) + if not required_hero_critics: + raise ValueError("At least one required hero critic must be specified") + required_hero_critic_count = len(required_hero_critics) + is_required_hero_critic = df["critic_name"].isin(required_hero_critics) + df["hero_iptm_score_component"] = df["official_iptm_score_component"].where( + is_required_hero_critic + ) + df["scored_hero_critic_name"] = df["critic_name"].where( + is_required_hero_critic & df["hero_iptm_score_component"].notna() + ) + + key_columns = [*available_group_columns, "designed_sequence"] + summary_columns = [ + "batch_idx", + "binder_sequence", + "is_antibody", + "isoelectric_point", + ] + summary_aggregations = { + column: (column, "first") for column in summary_columns if column in df.columns + } + scores = df.groupby(key_columns, as_index=False).agg( + iptm_score=("hero_iptm_score_component", "mean"), + iptm_proxy_score=("distogram_iptm_proxy", "mean"), + hero_iptm_min=("hero_iptm_score_component", "min"), + hero_iptm_median=("hero_iptm_score_component", "median"), + hero_iptm_max=("hero_iptm_score_component", "max"), + critic_count=("critic_name", "nunique"), + hero_critic_count=("scored_hero_critic_name", "nunique"), + **summary_aggregations, + ) + scores["required_hero_critic_count"] = required_hero_critic_count + scores["selection_score"] = scores["iptm_score"].fillna(0.0) + scores["all_hero_critics_pass"] = scores["hero_critic_count"].eq( + required_hero_critic_count + ) & scores["hero_iptm_min"].gt(consensus_iptm_threshold) + scores["consensus_iptm_threshold"] = consensus_iptm_threshold + + if available_group_columns: + sort_columns = [*available_group_columns, "selection_score"] + ascending = [*[True] * len(available_group_columns), False] + scores = scores.sort_values(sort_columns, ascending=ascending) + return ( + scores.groupby(available_group_columns, group_keys=False, sort=False) + .head(top_k) + .reset_index(drop=True) + ) + return scores.nlargest(min(len(scores), top_k), "selection_score").reset_index(drop=True) + + +def _write_official_selection_table( + path: Path, + rows: list[dict[str, Any]], + required_hero_critics: tuple[str, ...] | None = None, +) -> None: + selection_df = select_official_designs( + rows, + required_hero_critics=required_hero_critics, + ) + selection_df.to_parquet(path, index=False) + + +def _log_official_selection_summary(rows: list[dict[str, Any]]) -> None: + selection_df = select_official_designs(rows) + if selection_df.empty: + logger.info("Official selection table is empty after pI filtering") + return + top = selection_df.iloc[0] + logger.info( + "Top official selection | score=%.4f hero_mean=%.4f proxy_mean=%.4f " + "hero_min=%.4f all_hero_pass=%s binder=%s", + float(top["selection_score"]), + float(top["iptm_score"]), + float(top["iptm_proxy_score"]) if not math.isnan(top["iptm_proxy_score"]) else 0.0, + float(top["hero_iptm_min"]), + bool(top["all_hero_critics_pass"]), + top["binder_sequence"], + ) + + +_ESMC_CACHE: Any | None = None +_ESMC_CACHE_KEY: tuple[str, str] | None = None +_ESMC_CACHE_CONTEXT: dict[str, Any] = {} + + +_ESMC_CONTEXT_FIELDS = ( + "_esmc_fp8", + "_esmc_fp8_module_paths", + "_esmc_source", + "_esmc_source_revision", + "_esmc_source_files", + "_esmc_local_files_only", + "_esmc_precision_policy", + "_esmc_precision_status", +) + + +def _load_fold_model( + model_name: str, + revision: str, + lm_dropout: float, + cache_esmc: bool, + device: torch.device | str, + kernel_backend: str | None, + compile_model: bool, + local_files_only: bool, +) -> Any: + global _ESMC_CACHE, _ESMC_CACHE_CONTEXT, _ESMC_CACHE_KEY + repo_id = _repo_name(model_name) + model = AutoModel.from_pretrained( + repo_id, + revision=revision, + local_files_only=local_files_only, + trust_remote_code=True, + load_esmc=not cache_esmc, + dtype=torch.float32, + ) + _record_model_load_identity( + model, + repo_id=repo_id, + revision=revision, + local_files_only=local_files_only, + ) + model = model.to(device=device) + if cache_esmc: + esmc_cache_key = (str(model.config.esmc_id), str(torch.device(device))) + if _ESMC_CACHE is None or esmc_cache_key != _ESMC_CACHE_KEY: + model.load_esmc( + model.config.esmc_id, + device=device, + local_files_only=local_files_only, + ) + _ESMC_CACHE = model._esmc + _ESMC_CACHE_KEY = esmc_cache_key + _ESMC_CACHE_CONTEXT = {field: getattr(model, field) for field in _ESMC_CONTEXT_FIELDS} + else: + model._esmc = _ESMC_CACHE + for field, value in _ESMC_CACHE_CONTEXT.items(): + setattr(model, field, value.copy() if isinstance(value, dict) else value) + model.configure_lm_dropout(lm_dropout, force_lm_dropout_during_inference=True) + if kernel_backend is not None: + model.set_kernel_backend(kernel_backend) + if compile_model: + model.apply_torch_compile() + return model.eval().requires_grad_(False) + + +class FastPLMsBinderDesign: + lm_name: str = "Synthyra/ESMplusplus_6B" + inversion_model_names: tuple[str, ...] = ("ESMFold2-Experimental-Fast-Cutoff2025",) + hero_critic_model_names: tuple[str, ...] = ( + "ESMFold2-Experimental-Fast-Cutoff2025", + "ESMFold2-Experimental-Cutoff2025", + ) + + def load( + self, + device: str = "cuda", + kernel_backend: str | None = None, + compile_model: bool = False, + inversion_model_names: tuple[str, ...] | None = None, + critic_model_names: tuple[str, ...] | None = None, + lm_name: str | None = None, + model_revisions: Mapping[str, str] | None = None, + local_files_only: bool = False, + ) -> None: + _configure_offline_mode(local_files_only) + revisions = _normalize_model_revisions(model_revisions) + selected_inversion_models = inversion_model_names or self.inversion_model_names + selected_critic_models = critic_model_names or self.hero_critic_model_names + selected_lm = lm_name or self.lm_name + if len(set(selected_inversion_models)) != len(selected_inversion_models): + raise ValueError("Inversion model names must be unique.") + if len(set(selected_critic_models)) != len(selected_critic_models): + raise ValueError("Critic model names must be unique.") + inversion_sources = { + model_name: _resolve_model_source(model_name, revisions) + for model_name in selected_inversion_models + } + critic_sources = { + model_name: _resolve_model_source(model_name, revisions) + for model_name in selected_critic_models + } + lm_repo_id, lm_revision = _resolve_model_source(selected_lm, revisions) + selected_repositories = { + repo_id + for repo_id, _revision in ( + *inversion_sources.values(), + *critic_sources.values(), + (lm_repo_id, lm_revision), + ) + } + unused_revisions = sorted(set(revisions) - selected_repositories) + if unused_revisions: + raise ValueError( + "Model revisions were supplied for repositories that are not loaded: " + f"{unused_revisions}." + ) + + self.device = torch.device(device) + self.inversion_models = { + model_name: _load_fold_model( + model_name, + revision=inversion_sources[model_name][1], + lm_dropout=0.5, + cache_esmc=True, + device=device, + kernel_backend=kernel_backend, + compile_model=compile_model, + local_files_only=local_files_only, + ) + for model_name in selected_inversion_models + } + self.critic_models = { + model_name: _load_fold_model( + model_name, + revision=critic_sources[model_name][1], + lm_dropout=0.25, + cache_esmc=True, + device=device, + kernel_backend=kernel_backend, + compile_model=compile_model, + local_files_only=local_files_only, + ) + for model_name in selected_critic_models + } + self.lm_model = ( + AutoModelForMaskedLM.from_pretrained( + lm_repo_id, + revision=lm_revision, + local_files_only=local_files_only, + trust_remote_code=True, + dtype=torch.float32, + ) + .to(device=device) + .eval() + .requires_grad_(False) + ) + _record_model_load_identity( + self.lm_model, + repo_id=lm_repo_id, + revision=lm_revision, + local_files_only=local_files_only, + ) + self.inversion_model_names = tuple(selected_inversion_models) + self.hero_critic_model_names = tuple(selected_critic_models) + self.lm_name = selected_lm + + def design( + self, + target_name: str | None = None, + target_sequence: str | None = None, + binder_name: str | None = None, + binder_sequence: str | None = None, + is_antibody: bool | None = None, + seed: int = 0, + batch_size: int = 1, + steps: int = DEFAULT_STEPS, + output_dir: str | None = None, + ) -> tuple[list[str], dict[int, dict[str, torch.Tensor]], list[dict[str, Any]]]: + return design_binder( + self.inversion_models, + self.critic_models, + self.lm_model, + target_name=target_name, + target_sequence=target_sequence, + binder_name=binder_name, + binder_sequence=binder_sequence, + is_antibody=is_antibody, + seed=seed, + batch_size=batch_size, + steps=steps, + output_dir=output_dir, + device=self.device, + ) + + +def _design_kwargs_from_args(args: argparse.Namespace) -> dict[str, Any]: + return { + "target_name": args.target_name, + "target_sequence": args.target_sequence, + "binder_name": args.binder_name, + "binder_sequence": args.binder_sequence, + "is_antibody": args.is_antibody, + "seed": args.seed, + "batch_size": args.batch_size, + "steps": args.steps, + "output_dir": args.output_dir, + } + + +def run_local(args: argparse.Namespace) -> None: + _require_fresh_output_directory(args.output_dir) + runner = FastPLMsBinderDesign() + runner.load( + kernel_backend=args.kernel_backend, + compile_model=args.compile_model, + inversion_model_names=( + tuple(args.inversion_model_names) if args.inversion_model_names is not None else None + ), + critic_model_names=( + tuple(args.critic_model_names) if args.critic_model_names is not None else None + ), + lm_name=args.lm_model, + model_revisions=args.model_revisions, + local_files_only=args.local_files_only, + ) + best_sequences, _, results = runner.design(**_design_kwargs_from_args(args)) + logger.info("Designed sequences: %s", best_sequences) + logger.info("Returned %d critic rows", len(results)) + _log_official_selection_summary(results) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--target-name", default="pd-l1") + parser.add_argument("--target-sequence", default=None) + parser.add_argument("--binder-name", default="minibinder") + parser.add_argument("--binder-sequence", default=None) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--steps", type=int, default=DEFAULT_STEPS) + parser.add_argument("--output-dir", default="binder_design_out") + parser.add_argument( + "--inversion-model", + dest="inversion_model_names", + action="append", + default=None, + help="FastPLMs inversion checkpoint; repeat to use multiple checkpoints.", + ) + parser.add_argument( + "--critic-model", + dest="critic_model_names", + action="append", + default=None, + help="FastPLMs critic checkpoint; repeat to use multiple checkpoints.", + ) + parser.add_argument( + "--lm-model", + default=None, + help="FastPLMs masked-language-model checkpoint.", + ) + parser.add_argument( + "--model-revision", + dest="model_revisions", + action="append", + default=[], + metavar="REPO=COMMIT", + help=( + "Immutable 40-character Hub commit for a custom model repository; " + "repeat once per custom repository. Registered defaults use models.toml." + ), + ) + parser.add_argument( + "--local-files-only", + action="store_true", + help=( + "Use cached snapshots only and set HF_HUB_OFFLINE=1 plus " + "TRANSFORMERS_OFFLINE=1 before model loading." + ), + ) + parser.add_argument("--kernel-backend", default=None) + parser.add_argument("--compile-model", action="store_true") + parser.add_argument("--is-antibody", dest="is_antibody", action="store_true") + parser.add_argument("--not-antibody", dest="is_antibody", action="store_false") + parser.set_defaults(is_antibody=None) + args = parser.parse_args(argv) + try: + args.model_revisions = _parse_model_revision_args(args.model_revisions) + except ValueError as error: + parser.error(str(error)) + if args.target_sequence is not None: + args.target_name = None + if args.binder_sequence is not None: + args.binder_name = None + return args + + +def main(argv: list[str] | None = None) -> int: + """Run the local binder-design workflow from explicit CLI arguments.""" + + run_local(parse_args(argv)) + return 0 + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s: %(message)s", + ) + raise SystemExit(main()) diff --git a/examples/e1_rag.py b/examples/e1_rag.py new file mode 100644 index 0000000..e0769b6 --- /dev/null +++ b/examples/e1_rag.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Embed duplicate E1 queries with a local A3M and shared persistence.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +if __package__: + from ._runtime import add_execution_arguments, resolve_execution +else: + from _runtime import add_execution_arguments, resolve_execution + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def embed_local_msa( + model: Any, + sequence: str, + a3m_path: Path, + *, + output: Path | None, + output_format: str, + seed: int, +) -> Any: + return model.embed_dataset_with_msa( + [sequence, sequence], + msa_lookup={sequence: str(a3m_path)}, + batch_size=2, + max_len=len(sequence), + pooling_types=["mean"], + seed=seed, + progress=False, + batch_window_size=2, + max_tokens_per_batch=2 * len(sequence), + output=output, + format=output_format, + resume=output is not None, + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path, help="Local E1 artifact") + parser.add_argument("a3m", type=Path, help="Local A3M whose query matches --sequence") + parser.add_argument("--sequence", default="MSTNPKPQRKTKRNT") + parser.add_argument("--output", type=Path) + parser.add_argument( + "--format", + choices=("safetensors", "sqlite"), + default="safetensors", + dest="output_format", + ) + parser.add_argument("--seed", type=int, default=7) + add_execution_arguments(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + a3m_path = arguments.a3m.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + if not a3m_path.is_file(): + raise SystemExit(f"Not a local A3M: {a3m_path}") + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + from transformers import AutoModelForMaskedLM + + model = AutoModelForMaskedLM.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ).to(device).eval() + result = embed_local_msa( + model, + arguments.sequence, + a3m_path, + output=arguments.output, + output_format=arguments.output_format, + seed=arguments.seed, + ) + print([(record.id, record.sequence) for record in result]) + + if arguments.output is not None and arguments.output_format == "sqlite": + from fastplms.embeddings import load_sqlite_result + + selected = load_sqlite_result(arguments.output, positions=[1, 0, 1]) + print("selected", [record.id for record in selected]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/embedding_and_retrieval.py b/examples/embedding_and_retrieval.py new file mode 100644 index 0000000..7df8661 --- /dev/null +++ b/examples/embedding_and_retrieval.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Stream ordered embeddings and reopen an optional SQLite result read-only.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +if __package__: + from ._runtime import add_execution_arguments, resolve_execution +else: + from _runtime import add_execution_arguments, resolve_execution + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def run_embeddings( + model: Any, + tokenizer: Any, + inputs: Any, + *, + output: Path | None, + output_format: str, + max_length: int, +) -> Any: + return model.embed_dataset( + inputs, + tokenizer=tokenizer, + batch_size=8, + batch_window_size=64, + max_tokens_per_batch=4096, + max_length=max_length, + pooling=("mean", "std"), + output=output, + format=output_format, + resume=output is not None, + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path, help="Local manifest-built artifact") + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--fasta", type=Path) + source.add_argument("--sequence", action="append", dest="sequences") + parser.add_argument("--output", type=Path) + parser.add_argument( + "--format", + choices=("safetensors", "sqlite"), + default="safetensors", + dest="output_format", + ) + parser.add_argument("--max-length", type=int, default=1024) + parser.add_argument("--select-id", action="append", default=[]) + add_execution_arguments(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + if arguments.select_id and (arguments.output is None or arguments.output_format != "sqlite"): + raise SystemExit("--select-id requires both --output and --format sqlite") + inputs: Any = arguments.fasta if arguments.fasta is not None else arguments.sequences + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + from transformers import AutoModel, AutoTokenizer + + model = ( + AutoModel.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ) + .to(device) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + ) + result = run_embeddings( + model, + tokenizer, + inputs, + output=arguments.output, + output_format=arguments.output_format, + max_length=arguments.max_length, + ) + for record in result: + print(record.id, record.sequence, tuple(record.load_tensor().shape)) + + if arguments.output is not None and arguments.output_format == "sqlite" and arguments.select_id: + from fastplms.embeddings import load_sqlite_result + + selected = load_sqlite_result( + arguments.output, + record_ids=arguments.select_id, + ) + print("selected", [record.id for record in selected]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/fine_tuning.py b/examples/fine_tuning.py new file mode 100644 index 0000000..3c694a0 --- /dev/null +++ b/examples/fine_tuning.py @@ -0,0 +1,1976 @@ +#! /usr/bin/env python3 +""" +This script shows how to fine-tune a Synthyra FastPLM model for protein +sequence regression or classification. +For regression we look at the binding affinity of two proteins (pkd) +For classification we look at the solubility of a protein (membrane bound or not) +""" + +import argparse +import contextlib +import hashlib +import inspect +import json +import math +import os +import platform +import re +import shutil +import sys +import tempfile +import uuid +import numpy as np +import torch +from collections.abc import Callable, Iterator, Mapping +from functools import wraps +from importlib import metadata +from numbers import Integral, Real +from pathlib import Path +from typing import Any, ParamSpec, TypeVar, cast +from datasets import load_dataset +from peft import LoraConfig, PeftModel, get_peft_model +from torch.utils.data import Dataset as TorchDataset +from transformers import ( + AutoModelForSequenceClassification, + EarlyStoppingCallback, + EvalPrediction, + Trainer, + TrainingArguments, + set_seed, +) + + +DEFAULT_MODEL = "Synthyra/ESM2-8M" +DEFAULT_MODEL_REVISION = "185ecbd45665d050a8dae326d91886d330c5f9d0" +DEFAULT_CLASSIFICATION_DATASET = "GleghornLab/DL2_reg" +DEFAULT_CLASSIFICATION_DATASET_REVISION = "7e18f1b98859b0a3e3da283f63d0a153b774cf1f" +DEFAULT_REGRESSION_TRAIN_DATASET = "Synthyra/ProteinProteinAffinity" +DEFAULT_REGRESSION_TRAIN_DATASET_REVISION = "f4a51e5e9f2c2a0185693f9fbcffc02d9dae08db" +DEFAULT_REGRESSION_VALIDATION_DATASET = "Synthyra/AffinityBenchmarkv5.5" +DEFAULT_REGRESSION_VALIDATION_DATASET_REVISION = "826ccfb1488d52b7b361802fbde161373247d084" +DEFAULT_REGRESSION_TEST_DATASET = "Synthyra/haddock_benchmark" +DEFAULT_REGRESSION_TEST_DATASET_REVISION = "4e22f014745728fca2d9c10f2f2cfd5a29a4981c" +CLASSIFIER_MODULE_NAME = "classifier" +EXAMPLE_ATTENTION_BACKENDS = ("eager", "sdpa", "flex_attention") +_OUTPUT_RESERVATION_FILE = ".fastplms-output-reservation.json" +_IMMUTABLE_REVISION = re.compile(r"^[0-9a-f]{40}$", re.IGNORECASE) +_P = ParamSpec("_P") +_R = TypeVar("_R") +_PINNED_DEFAULT_REVISIONS = { + ("model", DEFAULT_MODEL): DEFAULT_MODEL_REVISION, + ("dataset", DEFAULT_CLASSIFICATION_DATASET): DEFAULT_CLASSIFICATION_DATASET_REVISION, + ("dataset", DEFAULT_REGRESSION_TRAIN_DATASET): (DEFAULT_REGRESSION_TRAIN_DATASET_REVISION), + ("dataset", DEFAULT_REGRESSION_VALIDATION_DATASET): ( + DEFAULT_REGRESSION_VALIDATION_DATASET_REVISION + ), + ("dataset", DEFAULT_REGRESSION_TEST_DATASET): DEFAULT_REGRESSION_TEST_DATASET_REVISION, +} + +BASE_TRAINER_KWARGS = { + "warmup_steps": 500, + "weight_decay": 0.01, + "logging_steps": 100, + "eval_strategy": "steps", + "eval_steps": 500, + "save_strategy": "steps", + "save_steps": 500, + "load_best_model_at_end": True, + "metric_for_best_model": "eval_loss", + "greater_is_better": False, + "report_to": "none", + "label_names": ["labels"], +} + + +def _output_path_exists(path: Path) -> bool: + """Return true for files, directories, and broken symlinks.""" + + return os.path.lexists(path) + + +@contextlib.contextmanager +def _reserved_output_directory(output_dir: str | Path) -> Iterator[Path]: + """Atomically reserve a new run directory and clean it after a failed run.""" + + destination = Path(output_dir).expanduser().absolute() + destination.parent.mkdir(parents=True, exist_ok=True) + try: + destination.mkdir() + except FileExistsError as error: + raise FileExistsError( + "Refusing to start because the task output path already exists and could " + f"mix state from another run: {destination}" + ) from error + + reservation_token = uuid.uuid4().hex + reservation_path = destination / _OUTPUT_RESERVATION_FILE + try: + with reservation_path.open("x", encoding="utf-8") as handle: + json.dump( + { + "schema_version": 1, + "reservation_token": reservation_token, + }, + handle, + sort_keys=True, + ) + handle.write("\n") + except BaseException: + destination.rmdir() + raise + + try: + yield destination + except BaseException as error: + try: + reservation = json.loads(reservation_path.read_text(encoding="utf-8")) + if reservation.get("reservation_token") != reservation_token: + raise RuntimeError( + "The output reservation identity changed while the run was active; " + f"preserving {destination} for manual inspection." + ) + shutil.rmtree(destination) + except BaseException as cleanup_error: + error.add_note( + "FastPLMs could not clean the failed run's reserved output directory: " + f"{cleanup_error}" + ) + raise + else: + reservation_path.unlink() + + +def _guard_training_output( + *, + lora_default: str, + full_default: str, +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + """Reserve a task output before the decorated function performs any work.""" + + def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]: + function_signature = inspect.signature(function) + + @wraps(function) + def guarded(*args: _P.args, **kwargs: _P.kwargs) -> _R: + bound = function_signature.bind(*args, **kwargs) + bound.apply_defaults() + output_dir = bound.arguments["output_dir"] + if output_dir is None: + output_dir = lora_default if bound.arguments["use_lora"] else full_default + with _reserved_output_directory(output_dir) as reserved: + bound.arguments["output_dir"] = reserved + return function(*bound.args, **bound.kwargs) + + return guarded + + return decorate + + +def _ensure_output_paths_available(paths: list[Path]) -> None: + """Preflight all requested task outputs before a multi-task CLI starts.""" + + collisions = [str(path) for path in paths if _output_path_exists(path)] + if collisions: + raise FileExistsError( + "Refusing to start because task output paths already exist and could mix " + f"prior state: {collisions}" + ) + + +def _ensure_classifier_persistence(lora_config: Any) -> Any: + """Ensure PEFT saves the independently trained classification head.""" + modules_to_save = list(lora_config.modules_to_save or ()) + if CLASSIFIER_MODULE_NAME not in modules_to_save: + modules_to_save.append(CLASSIFIER_MODULE_NAME) + lora_config.modules_to_save = modules_to_save + return lora_config + + +class PairDatasetHF(TorchDataset): + """ + Dataset class for protein pair data (e.g., protein-protein interactions). + + Args: + data: The dataset containing protein sequences and labels + col_a: Column name for the first protein sequence + col_b: Column name for the second protein sequence + label_col: Column name for the labels + max_length: Encoded token budget for the complete pair, including + tokenizer-added separator and special tokens + """ + + def __init__( + self, dataset: Any, col_a: str, col_b: str, label_col: str, max_length: int = 2048 + ): + self.seqs_a = dataset[col_a] + self.seqs_b = dataset[col_b] + self.labels = dataset[label_col] + self.max_length = max_length + + def __len__(self) -> int: + return len(self.seqs_a) + + def __getitem__(self, idx: int) -> tuple[str, str, float | int]: + # Token-budget filtering and any defensive truncation belong to the + # tokenizer-aware collator. Slicing both strings to max_length could + # still create a pair longer than the model context and ignores special + # tokens inserted between the proteins. + seq_a = self.seqs_a[idx] + seq_b = self.seqs_b[idx] + label = self.labels[idx] + return seq_a, seq_b, label + + +class SequenceDatasetHF(TorchDataset): + """ + Dataset class for single protein sequence data. + + Args: + dataset: The dataset containing protein sequences and labels + col_name: Column name for the protein sequences + label_col: Column name for the labels + max_length: Encoded token budget including tokenizer-added special tokens + """ + + def __init__( + self, + dataset: Any, + col_name: str = "seqs", + label_col: str = "labels", + max_length: int = 2048, + ): + self.seqs = dataset[col_name] + self.labels = dataset[label_col] + self.max_length = max_length + + def __len__(self) -> int: + return len(self.seqs) + + def __getitem__(self, idx: int) -> tuple[str, float | int]: + seq = self.seqs[idx] + label = self.labels[idx] + return seq, label + + +def _tokenization_kwargs(max_length: int | None, *, pair: bool) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "padding": "longest", + "return_tensors": "pt", + } + if max_length is not None: + if max_length <= 0: + raise ValueError("max_length must be positive") + kwargs.update( + { + "truncation": "longest_first" if pair else True, + "max_length": max_length, + } + ) + if max_length % 8 == 0: + kwargs["pad_to_multiple_of"] = 8 + else: + kwargs["pad_to_multiple_of"] = 8 + return kwargs + + +def _encoded_length(tokenizer: Any, sequence: str, pair: str | None = None) -> int: + encoded = tokenizer( + sequence, + pair, + add_special_tokens=True, + truncation=False, + return_attention_mask=False, + ) + input_ids = encoded["input_ids"] # (l,) or (1, l) when returned as a tensor + if isinstance(input_ids, torch.Tensor): + return int(input_ids.shape[-1]) + if input_ids and isinstance(input_ids[0], list): + input_ids = input_ids[0] + return len(input_ids) + + +def _fits_token_budget( + tokenizer: Any, + sequence: str, + pair: str | None, + max_length: int, +) -> bool: + """Include tokenizer-added control tokens in the context-length gate.""" + + return _encoded_length(tokenizer, sequence, pair) <= max_length + + +class PairCollator: + """ + Collator for protein pair data that handles tokenization and tensor conversion. + + Args: + tokenizer: The tokenizer to use for encoding sequences + regression: Whether this is a regression task (True) or classification (False) + max_length: Encoded token budget for the complete pair, including + tokenizer-added separator and special tokens + """ + + def __init__( + self, + tokenizer: Any, + regression: bool = False, + max_length: int | None = None, + ): + self.tokenizer = tokenizer + self.regression = regression + self.max_length = max_length + + def __call__(self, batch: list[tuple[str, str, float | int]]) -> dict[str, torch.Tensor]: + seqs_a, seqs_b, labels = zip(*batch, strict=True) + labels = torch.tensor(labels) # (b,) + labels = labels.float() if self.regression else labels.long() # (b,) + tokenized = self.tokenizer( + seqs_a, + seqs_b, + **_tokenization_kwargs(self.max_length, pair=True), + ) # input_ids/attention_mask: (b, l) + return { + "input_ids": tokenized["input_ids"], # (b, l) + "attention_mask": tokenized["attention_mask"], # (b, l) + "labels": labels, # (b,) + } + + +class SequenceCollator: + """ + Collator for single protein sequence data that handles tokenization and tensor conversion. + + Args: + tokenizer: The tokenizer to use for encoding sequences + regression: Whether this is a regression task (True) or classification (False) + max_length: Encoded token budget including tokenizer-added special tokens + """ + + def __init__( + self, + tokenizer: Any, + regression: bool = False, + max_length: int | None = None, + ): + self.tokenizer = tokenizer + self.regression = regression + self.max_length = max_length + + def __call__(self, batch: list[tuple[str, float | int]]) -> dict[str, torch.Tensor]: + seqs, labels = zip(*batch, strict=True) + labels = torch.tensor(labels) # (b,) + labels = labels.float() if self.regression else labels.long() # (b,) + tokenized = self.tokenizer( + seqs, + **_tokenization_kwargs(self.max_length, pair=False), + ) # input_ids/attention_mask: (b, l) + return { + "input_ids": tokenized["input_ids"], # (b, l) + "attention_mask": tokenized["attention_mask"], # (b, l) + "labels": labels, # (b,) + } + + +def _tree_sha256(path: Path) -> str: + """Hash every regular byte in a local model or dataset tree.""" + + root = path.resolve(strict=True) + if not root.is_dir(): + raise ValueError(f"Local reproducibility source must be a directory: {root}") + digest = hashlib.sha256() + files = sorted(root.rglob("*"), key=lambda item: item.relative_to(root).as_posix()) + for file_path in files: + if file_path.is_symlink(): + raise ValueError(f"Local reproducibility sources may not contain symlinks: {file_path}") + if not file_path.is_file(): + continue + relative = file_path.relative_to(root).as_posix().encode("utf-8") + digest.update(len(relative).to_bytes(8, "big")) + digest.update(relative) + digest.update(file_path.stat().st_size.to_bytes(8, "big")) + with file_path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _immutable_source_identity( + source: str, + revision: str | None, + *, + source_kind: str, +) -> dict[str, Any]: + """Resolve a Hub commit or immutable local tree, rejecting moving refs.""" + + local_path = Path(source).expanduser() + if local_path.exists(): + return { + "kind": "local_tree", + "source_kind": source_kind, + "path": str(local_path.resolve(strict=True)), + "tree_sha256": _tree_sha256(local_path), + } + if local_path.is_absolute() or source.startswith((".", "~")): + raise FileNotFoundError(f"Local {source_kind} source does not exist: {source}") + if revision is None: + revision = _PINNED_DEFAULT_REVISIONS.get((source_kind, source)) + if revision is None or _IMMUTABLE_REVISION.fullmatch(revision) is None: + raise ValueError( + f"Remote {source_kind} source {source!r} requires an immutable 40-character " + "Hub commit revision; branches, tags, and omitted revisions are rejected." + ) + return { + "kind": "hub", + "source_kind": source_kind, + "repo_id": source, + "revision": revision.lower(), + } + + +def _load_dataset_immutable( + source: str, + revision: str | None, + *, + split: str | None = None, +) -> tuple[Any, dict[str, Any]]: + identity = _immutable_source_identity( + source, + revision, + source_kind="dataset", + ) + kwargs: dict[str, Any] = {} + if identity["kind"] == "hub": + kwargs["revision"] = identity["revision"] + if split is not None: + kwargs["split"] = split + return load_dataset(source, **kwargs), identity + + +def _available_columns(dataset: Any) -> set[str]: + column_names = getattr(dataset, "column_names", None) + if column_names is not None and not isinstance(column_names, Mapping): + return {str(column) for column in column_names} + if isinstance(dataset, Mapping): + return {str(column) for column in dataset} + features = getattr(dataset, "features", None) + if isinstance(features, Mapping): + return {str(column) for column in features} + raise TypeError("Dataset must expose column_names, features, or mapping keys.") + + +def _require_dataset_columns( + dataset: Any, + *, + split: str, + required: tuple[str, ...], +) -> None: + available = _available_columns(dataset) + missing = sorted(set(required).difference(available)) + if missing: + raise ValueError( + f"Dataset split {split!r} is missing required columns {missing}; " + f"available columns are {sorted(available)}." + ) + if len(dataset) == 0: + raise ValueError(f"Dataset split {split!r} must contain at least one row.") + + +def _validate_sequence_column(dataset: Any, *, split: str, column: str) -> None: + for row_index, sequence in enumerate(dataset[column]): + if not isinstance(sequence, str) or not sequence.strip(): + raise ValueError( + f"Dataset split {split!r} column {column!r} row {row_index} must " + "contain a non-empty protein sequence string." + ) + + +def _classification_label_set(dataset: Any, *, split: str) -> set[int]: + labels: set[int] = set() + for row_index, label in enumerate(dataset["labels"]): + if isinstance(label, bool) or not isinstance(label, Integral): + raise ValueError( + f"Dataset split {split!r} label at row {row_index} must be an " + f"integer, got {label!r}." + ) + labels.add(int(label)) + return labels + + +def _validate_classification_dataset_dict(data: Any) -> int: + """Validate the complete classification schema before model initialization.""" + + if not isinstance(data, Mapping): + raise TypeError( + "Classification data must be a DatasetDict-style mapping with train, " + "valid, and test splits." + ) + required_splits = ("train", "valid", "test") + missing_splits = [split for split in required_splits if split not in data] + if missing_splits: + raise ValueError( + "Classification data is missing required splits: " + f"{missing_splits}; expected train, valid, and test." + ) + + label_sets: dict[str, set[int]] = {} + for split in required_splits: + dataset = data[split] + _require_dataset_columns( + dataset, + split=split, + required=("seqs", "labels"), + ) + _validate_sequence_column(dataset, split=split, column="seqs") + label_sets[split] = _classification_label_set(dataset, split=split) + + train_labels = label_sets["train"] + expected_train_labels = set(range(len(train_labels))) + if train_labels != expected_train_labels: + raise ValueError( + "Classification training labels must be contiguous zero-based integers; " + f"observed {sorted(train_labels)}, expected {sorted(expected_train_labels)}." + ) + for split in ("valid", "test"): + unseen = sorted(label_sets[split].difference(train_labels)) + if unseen: + raise ValueError( + f"Classification split {split!r} contains labels absent from train: {unseen}." + ) + return len(train_labels) + + +def _validate_regression_dataset(dataset: Any, *, split: str) -> None: + """Validate a protein-pair regression split before model initialization.""" + + _require_dataset_columns( + dataset, + split=split, + required=("SeqA", "SeqB", "labels"), + ) + _validate_sequence_column(dataset, split=split, column="SeqA") + _validate_sequence_column(dataset, split=split, column="SeqB") + for row_index, label in enumerate(dataset["labels"]): + if ( + isinstance(label, bool) + or not isinstance(label, Real) + or not math.isfinite(float(label)) + ): + raise ValueError( + f"Regression split {split!r} label at row {row_index} must be a " + f"finite real number, got {label!r}." + ) + + +def _require_non_empty_filtered_split(dataset: Any, *, split: str) -> None: + if len(dataset) == 0: + raise ValueError(f"Dataset split {split!r} has no rows within the encoded token budget.") + + +def _verify_training_source_unchanged( + model: Any, + model_name: str, + model_revision: str | None, +) -> dict[str, Any]: + """Fail if the exact local tree loaded at initialization has since drifted.""" + + expected = getattr(model, "_fastplms_training_source_identity", None) + if not isinstance(expected, Mapping): + raise RuntimeError( + "The model is missing its initialization-time source identity; initialize it " + "with initialize_model() before saving a reproducible training artifact." + ) + observed = _immutable_source_identity( + model_name, + model_revision, + source_kind="model", + ) + if dict(expected) != observed: + raise RuntimeError( + "The training model source changed after initialization. Refusing to save or " + f"record a stale artifact identity: expected={dict(expected)}, observed={observed}." + ) + return observed + + +def _effective_attention_backend(model: Any) -> str | None: + config = model.config + for field in ("attn_backend", "attention_backend", "_attn_implementation"): + value = getattr(config, field, None) + if value is not None: + return str(getattr(value, "value", value)) + return None + + +def initialize_model( + model_name: str, + num_labels: int, + use_lora: bool = True, + lora_config: Any = None, + model_revision: str | None = None, + attn_backend: str = "sdpa", +) -> tuple[Any, Any]: + """ + Initialize a model with optional LoRA support + + Args: + model_name: Name or path of the pretrained model + num_labels: Number of labels for the task (1 for regression) + use_lora: Whether to use LoRA for fine-tuning + lora_config: Custom LoRA configuration (optional) + model_revision: Immutable Hub commit for a remote model + attn_backend: Explicit eager, SDPA, or Flex implementation + + Returns: + model: The initialized model + tokenizer: The model's tokenizer + """ + if attn_backend not in EXAMPLE_ATTENTION_BACKENDS: + raise ValueError( + f"The fine-tuning example supports {EXAMPLE_ATTENTION_BACKENDS}, got " + f"{attn_backend!r}. FlashAttention training requires an explicit BF16 CUDA " + "loading and placement policy that this compact CLI does not expose." + ) + + print(f"Loading model {model_name} with {num_labels} labels...") + + source_identity = _immutable_source_identity( + model_name, + model_revision, + source_kind="model", + ) + revision_kwargs = ( + {"revision": source_identity["revision"]} if source_identity["kind"] == "hub" else {} + ) + + model = AutoModelForSequenceClassification.from_pretrained( + model_name, + trust_remote_code=True, + num_labels=num_labels, + attn_implementation=attn_backend, + **revision_kwargs, + ) + effective_backend = _effective_attention_backend(model) + if effective_backend != attn_backend: + raise RuntimeError( + f"Requested attention backend {attn_backend!r}, but the loaded model " + f"configured {effective_backend!r}. Refusing to train with ambiguous dispatch." + ) + tokenizer = model.tokenizer + + if use_lora: + if lora_config is None: + # Target modules for the ESM2 sequence-classification artifacts. + target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"] + + lora_config = LoraConfig( + r=8, + lora_alpha=16, + lora_dropout=0.01, + bias="none", + target_modules=target_modules, + modules_to_save=[CLASSIFIER_MODULE_NAME], + ) + + # A custom configuration must preserve the task head too. Marking the + # head trainable without modules_to_save would omit it from an adapter + # checkpoint and restore an untrained head after reload. + lora_config = _ensure_classifier_persistence(lora_config) + + model = get_peft_model(model, lora_config) + + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + non_trainable_params = total_params - trainable_params + print(f"Total parameters: {total_params}") + print(f"Trainable parameters: {trainable_params}") + print(f"Non-trainable parameters: {non_trainable_params}") + print( + f"Percentage of parameters being trained: {100 * trainable_params / total_params:.2f}%" + ) + + model._fastplms_training_source_identity = source_identity + return model, tokenizer + + +def _package_version(package: str) -> str | None: + try: + return metadata.version(package) + except metadata.PackageNotFoundError: + return None + + +def _json_safe(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _json_safe(nested) for key, nested in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted((_json_safe(item) for item in value), key=str) + if value is None or isinstance(value, (bool, float, int, str)): + return value + return str(value) + + +def _sha256_json(value: Any) -> str: + payload = json.dumps( + _json_safe(value), + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _tokenizer_identity(tokenizer: Any) -> dict[str, Any]: + get_vocab = getattr(tokenizer, "get_vocab", None) + vocab = get_vocab() if callable(get_vocab) else getattr(tokenizer, "vocab", None) + init_kwargs = getattr(tokenizer, "init_kwargs", {}) + return { + "class": type(tokenizer).__name__, + "name_or_path": getattr(tokenizer, "name_or_path", None), + "revision": ( + init_kwargs.get("revision") or getattr(tokenizer, "_commit_hash", None) + if isinstance(init_kwargs, dict) + else getattr(tokenizer, "_commit_hash", None) + ), + "vocab_size": len(vocab) if isinstance(vocab, dict) else None, + "vocab_sha256": _sha256_json(vocab) if isinstance(vocab, dict) else None, + "special_token_ids": { + name: getattr(tokenizer, f"{name}_token_id", None) + for name in ("bos", "cls", "eos", "mask", "pad", "sep", "unk") + }, + } + + +def _ordered_rows_sha256(dataset: Any, columns: tuple[str, ...]) -> str: + """Hash ordered post-filter rows and only the columns consumed by training.""" + + digest = hashlib.sha256() + digest.update(_sha256_json(list(columns)).encode("ascii")) + row_count = 0 + for row_count, row in enumerate(dataset, start=1): + if not isinstance(row, Mapping): + raise TypeError("Dataset iteration must yield row mappings.") + missing = [column for column in columns if column not in row] + if missing: + raise KeyError(f"Dataset row is missing required columns: {missing}") + payload = json.dumps( + _json_safe({column: row[column] for column in columns}), + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + digest.update(row_count.to_bytes(8, "big")) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + if row_count != len(dataset): + raise RuntimeError( + "Dataset length changed while hashing ordered training rows: " + f"expected {len(dataset)}, observed {row_count}." + ) + return digest.hexdigest() + + +def _dataset_identity( + dataset: Any, + *, + columns: tuple[str, ...], + source: Mapping[str, Any], + split: str, +) -> dict[str, Any]: + fingerprint = getattr(dataset, "_fingerprint", None) + info = getattr(dataset, "info", None) + return { + "source": _json_safe(source), + "split": split, + "columns": list(columns), + "ordered_rows_sha256": _ordered_rows_sha256(dataset, columns), + "library_fingerprint_advisory": fingerprint, + "rows": len(dataset), + "builder_name": getattr(info, "builder_name", None), + "config_name": getattr(info, "config_name", None), + "version": ( + str(info.version) + if info is not None and getattr(info, "version", None) is not None + else None + ), + } + + +def _write_training_manifest( + output_dir: str, + *, + task: str, + model: Any, + tokenizer: Any, + model_name: str, + model_revision: str | None, + seed: int, + max_length: int, + use_lora: bool, + batch_size: int, + gradient_accumulation_steps: int, + learning_rate: float, + num_epochs: float, + full_determinism: bool, + datasets: dict[str, Any], + dataset_contracts: dict[str, dict[str, Any]], + training_arguments: TrainingArguments, + patience: int, + final_artifact: Mapping[str, Any], + requested_attention_backend: str = "sdpa", +) -> None: + """Persist the execution identity needed to reproduce a training run.""" + + config = model.config + model_configuration = _json_safe(config.to_dict()) + parameter = next(iter(model.parameters()), None) # model-defined parameter shape + attention_backend = _effective_attention_backend(model) + adapter_configuration = { + str(name): _json_safe(adapter_config.to_dict()) + for name, adapter_config in (getattr(model, "peft_config", None) or {}).items() + } + if training_arguments.bf16: + compute_dtype = "torch.bfloat16" + elif training_arguments.fp16: + compute_dtype = "torch.float16" + else: + compute_dtype = "torch.float32" + manifest = { + "schema_version": 1, + "task": task, + "command": list(sys.argv), + "model": { + "requested": model_name, + "requested_source": _verify_training_source_unchanged( + model, + model_name, + model_revision, + ), + "resolved": getattr(config, "_name_or_path", None), + "revision": getattr(config, "_commit_hash", None), + "weights_revision": getattr(config, "fastplms_weights_revision", None), + "runtime_revision": getattr(config, "fastplms_runtime_revision", None), + "attention_backend": attention_backend, + "requested_attention_backend": requested_attention_backend, + "effective_attention_backend": attention_backend, + "parameter_dtype": str(parameter.dtype) if parameter is not None else None, + "configuration": model_configuration, + "configuration_sha256": _sha256_json(model_configuration), + "adapters": adapter_configuration, + "adapter_configuration_sha256": _sha256_json(adapter_configuration), + }, + "tokenizer": _tokenizer_identity(tokenizer), + "training": { + "seed": seed, + "max_length": max_length, + "max_length_semantics": ( + "encoded token budget including tokenizer-added special and pair tokens" + ), + "use_lora": use_lora, + "batch_size": batch_size, + "gradient_accumulation_steps": gradient_accumulation_steps, + "effective_single_process_batch_size": (batch_size * gradient_accumulation_steps), + "learning_rate": learning_rate, + "num_epochs": num_epochs, + "full_determinism": full_determinism, + "device": str(training_arguments.device), + "compute_dtype": compute_dtype, + "optimizer": getattr( + training_arguments.optim, + "value", + str(training_arguments.optim), + ), + "scheduler": getattr( + training_arguments.lr_scheduler_type, + "value", + str(training_arguments.lr_scheduler_type), + ), + "warmup_steps": training_arguments.warmup_steps, + "weight_decay": training_arguments.weight_decay, + "early_stopping_patience": patience, + "evaluation_strategy": str(training_arguments.eval_strategy), + "eval_steps": training_arguments.eval_steps, + "save_strategy": str(training_arguments.save_strategy), + "save_steps": training_arguments.save_steps, + "logging_strategy": str(training_arguments.logging_strategy), + "logging_steps": training_arguments.logging_steps, + "load_best_model_at_end": training_arguments.load_best_model_at_end, + "metric_for_best_model": training_arguments.metric_for_best_model, + "greater_is_better": training_arguments.greater_is_better, + "label_names": list(training_arguments.label_names or ()), + "report_to": list(training_arguments.report_to or ()), + "filtering": "tokenizer encoding including added special tokens", + "truncation": "longest_first", + }, + "datasets": { + split_name: _dataset_identity( + dataset, + columns=tuple(dataset_contracts[split_name]["columns"]), + source=dataset_contracts[split_name]["source"], + split=str(dataset_contracts[split_name]["split"]), + ) + for split_name, dataset in datasets.items() + }, + "final_artifact": _json_safe(final_artifact), + "environment": { + "python": platform.python_version(), + "platform": platform.platform(), + "torch": torch.__version__, + "transformers": _package_version("transformers"), + "peft": _package_version("peft"), + "datasets": _package_version("datasets"), + "fastplms": _package_version("fastplms"), + "accelerate": _package_version("accelerate"), + "cuda": torch.version.cuda, + "cuda_device": ( + torch.cuda.get_device_name(training_arguments.device) + if training_arguments.device.type == "cuda" + else None + ), + "cuda_capability": ( + list(torch.cuda.get_device_capability(training_arguments.device)) + if training_arguments.device.type == "cuda" + else None + ), + "cuda_driver": ( + getattr(torch._C, "_cuda_getDriverVersion", lambda: None)() + if training_arguments.device.type == "cuda" + else None + ), + }, + } + destination = Path(output_dir) + destination.mkdir(parents=True, exist_ok=True) + (destination / "run_manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + # tensor: (...) + value = tensor.detach().cpu().contiguous() # (...) + return hashlib.sha256(value.view(torch.uint8).numpy().tobytes()).hexdigest() + + +def _persisted_parameter_hashes(model: Any, *, use_lora: bool) -> dict[str, str]: + if not use_lora: + hashes = { + name: _tensor_sha256(value) + for name, value in model.state_dict().items() + if isinstance(value, torch.Tensor) + } + else: + hashes = { + name: _tensor_sha256(parameter) + for name, parameter in model.named_parameters() + if "lora_" in name or "modules_to_save" in name + } + if not hashes: + raise RuntimeError("No persisted model state was found for final-artifact verification.") + return hashes + + +def _primary_prediction_tensor(predictions: Any) -> torch.Tensor: + """Normalize Trainer and model prediction containers to a CPU tensor.""" + + value = predictions[0] if isinstance(predictions, tuple) else predictions # (...) + if not isinstance(value, (np.ndarray, torch.Tensor)): + raise TypeError( + "Held-out verification expected logits as a NumPy array or Torch tensor, " + f"got {type(value).__name__}." + ) + return torch.as_tensor(value).detach().cpu() # (...) + + +def _held_out_reload_verification( + trainer: Trainer, + reloaded_model: Any, + *, + verification_dataset: Any, + data_collator: Any, +) -> dict[str, Any]: + """Compare prepared-Trainer and independently reloaded held-out logits.""" + + row_count = len(verification_dataset) + if row_count < 1: + raise ValueError("Final-artifact verification requires a non-empty held-out dataset.") + held_out_rows = [verification_dataset[index] for index in range(min(2, row_count))] + prediction_output = trainer.predict(held_out_rows) # type: ignore[arg-type] + prediction_values = getattr(prediction_output, "predictions", None) + if prediction_values is None: + prediction_values = prediction_output[0] + expected = _primary_prediction_tensor(prediction_values) # (b_v, c) + + batch = data_collator(held_out_rows) + if not isinstance(batch, Mapping): + raise TypeError("The verification data collator must return a mapping.") + device = torch.device(getattr(trainer.args, "device", "cpu")) + prepared_batch = { + key: value.to(device) if isinstance(value, torch.Tensor) else value + for key, value in batch.items() + } # input_ids/attention_mask: (b_v, l); labels: (b_v,) + reloaded_model = reloaded_model.to(device).eval() + use_bf16 = bool(getattr(trainer.args, "bf16", False)) + use_fp16 = bool(getattr(trainer.args, "fp16", False)) + autocast_dtype: torch.dtype | None = None + if use_bf16 and device.type in {"cpu", "cuda"}: + autocast_dtype = torch.bfloat16 + elif use_fp16 and device.type == "cuda": + autocast_dtype = torch.float16 + with torch.inference_mode(): + if autocast_dtype is None: + output = reloaded_model(**prepared_batch) # logits: (b_v, c) + else: + with torch.autocast(device_type=device.type, dtype=autocast_dtype): + output = reloaded_model(**prepared_batch) # logits: (b_v, c) + observed = _primary_prediction_tensor( + output.logits if hasattr(output, "logits") else output[0] + ) # (b_v, c) + if observed.shape != expected.shape: + raise RuntimeError( + "Final model reload changed held-out prediction shape: " + f"expected {tuple(expected.shape)}, observed {tuple(observed.shape)}." + ) + if use_bf16: + rtol, atol = 5e-3, 5e-3 + elif use_fp16: + rtol, atol = 1e-3, 1e-3 + else: + rtol, atol = 1e-5, 1e-6 + expected_float = expected.float() # (b_v, c) + observed_float = observed.float() # (b_v, c) + try: + torch.testing.assert_close( + observed_float, + expected_float, + rtol=rtol, + atol=atol, + ) + except AssertionError as error: + raise RuntimeError( + "Final model reload changed held-out logits beyond the configured " + f"dtype tolerance (rtol={rtol}, atol={atol})." + ) from error + absolute_error = (observed_float - expected_float).abs() # (b_v, c) + relative_error = absolute_error / expected_float.abs().clamp_min(atol) # (b_v, c) + return { + "rows": len(held_out_rows), + "shape": list(expected.shape), + "device": str(device), + "autocast_dtype": str(autocast_dtype) if autocast_dtype is not None else None, + "rtol": rtol, + "atol": atol, + "max_absolute_error": float(absolute_error.max().item()), + "max_relative_error": float(relative_error.max().item()), + } + + +def _reload_final_model( + artifact_dir: Path, + *, + model_name: str, + model_revision: str | None, + num_labels: int, + use_lora: bool, + attn_backend: str = "sdpa", +) -> Any: + source = _immutable_source_identity( + model_name, + model_revision, + source_kind="model", + ) + revision_kwargs = {"revision": source["revision"]} if source["kind"] == "hub" else {} + if use_lora: + base_model = AutoModelForSequenceClassification.from_pretrained( + model_name, + trust_remote_code=True, + num_labels=num_labels, + attn_implementation=attn_backend, + **revision_kwargs, + ) + return PeftModel.from_pretrained( + base_model, + artifact_dir, + local_files_only=True, + ) + return AutoModelForSequenceClassification.from_pretrained( + artifact_dir, + trust_remote_code=True, + local_files_only=True, + attn_implementation=attn_backend, + ) + + +def _save_reload_verify_final_artifact( + trainer: Trainer, + tokenizer: Any, + *, + output_dir: str, + model_name: str, + model_revision: str | None, + num_labels: int, + use_lora: bool, + verification_dataset: Any, + data_collator: Any, + attn_backend: str = "sdpa", +) -> dict[str, Any]: + """Atomically save, locally reload, and verify trained adapter/head state.""" + + output_root = Path(output_dir).resolve() + output_root.mkdir(parents=True, exist_ok=True) + final_dir = output_root / "final_model" + if final_dir.exists(): + raise FileExistsError(f"Refusing to overwrite an existing final artifact: {final_dir}") + base_source = _verify_training_source_unchanged( + trainer.model, + model_name, + model_revision, + ) + expected_hashes = _persisted_parameter_hashes(trainer.model, use_lora=use_lora) + staging = Path(tempfile.mkdtemp(prefix=".final-model-", dir=output_root)) + try: + trainer.save_model(os.fspath(staging)) + save_tokenizer = getattr(tokenizer, "save_pretrained", None) + if not callable(save_tokenizer): + raise TypeError("The training tokenizer must implement save_pretrained().") + save_tokenizer(staging) + if use_lora and not (staging / "adapter_config.json").is_file(): + raise RuntimeError("Trainer did not save the required PEFT adapter configuration.") + if not any(staging.glob("*.safetensors")): + raise RuntimeError("Final artifact does not contain safetensors weights.") + reloaded = _reload_final_model( + staging, + model_name=model_name, + model_revision=model_revision, + num_labels=num_labels, + use_lora=use_lora, + attn_backend=attn_backend, + ).eval() + observed_hashes = _persisted_parameter_hashes(reloaded, use_lora=use_lora) + if observed_hashes != expected_hashes: + missing = sorted(set(expected_hashes).difference(observed_hashes)) + unexpected = sorted(set(observed_hashes).difference(expected_hashes)) + changed = sorted( + name + for name in set(expected_hashes).intersection(observed_hashes) + if expected_hashes[name] != observed_hashes[name] + ) + raise RuntimeError( + "Final model reload changed persisted training state: " + f"missing={missing}, unexpected={unexpected}, changed={changed}." + ) + inference_verification = _held_out_reload_verification( + trainer, + reloaded, + verification_dataset=verification_dataset, + data_collator=data_collator, + ) + metadata_payload = { + "schema_version": 1, + "base_source": base_source, + "use_lora": use_lora, + "verified_parameter_sha256": expected_hashes, + "held_out_inference": inference_verification, + } + (staging / "artifact_metadata.json").write_text( + json.dumps(metadata_payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + os.replace(staging, final_dir) + except Exception: + if staging.exists(): + shutil.rmtree(staging) + raise + + return { + "path": str(final_dir), + "tree_sha256": _tree_sha256(final_dir), + "verified_parameter_sha256": expected_hashes, + "reload_verified": True, + "held_out_inference": inference_verification, + } + + +def _rankdata(values: np.ndarray) -> np.ndarray: + """Return average ranks for ties without requiring the reporting extra.""" + + # values: (...) + flattened = np.asarray(values).reshape(-1) # (n,) + order = np.argsort(flattened, kind="mergesort") # (n,) + sorted_values = flattened[order] # (n,) + ranks = np.empty(flattened.size, dtype=np.float64) # (n,) + start = 0 + while start < flattened.size: + stop = start + 1 + while stop < flattened.size and sorted_values[stop] == sorted_values[start]: + stop += 1 + ranks[order[start:stop]] = (start + stop - 1) / 2 + 1 # (stop - start,) + start = stop + return ranks # (n,) + + +def _spearman_correlation(predictions: np.ndarray, labels: np.ndarray) -> float: + # predictions/labels: same-size arrays of arbitrary rank + prediction_ranks = _rankdata(predictions) # (n,) + label_ranks = _rankdata(labels) # (n,) + if prediction_ranks.size < 2: + return float("nan") + correlation = np.corrcoef(prediction_ranks, label_ranks)[0, 1] # () + return float(correlation) + + +def compute_metrics_regression(p: EvalPrediction) -> dict[str, float]: + """Compute Spearman correlation for regression tasks.""" + predictions, labels = p.predictions, p.label_ids # predictions: (...); labels: (...) + predictions = ( + predictions[0] if isinstance(predictions, tuple) else predictions + ) # (...) + return { + "spearman_correlation": _spearman_correlation( + predictions, + cast(np.ndarray, labels), + ) + } + + +def compute_metrics_classification(p: EvalPrediction) -> dict[str, float]: + """Compute accuracy for classification tasks""" + predictions, labels = p.predictions, p.label_ids # predictions: (n, c); labels: (n,) + predictions = ( + predictions[0] if isinstance(predictions, tuple) else predictions + ) # (n, c) + predictions = np.argmax(predictions, axis=-1) # (n,) + + accuracy = (predictions.flatten() == labels.flatten()).mean() # () + + return {"accuracy": accuracy} + + +def _save_figure_exclusive(figure: Any, output_path: str | Path) -> Path: + """Save one PNG without replacing an existing report artifact.""" + + destination = Path(output_path) + if not destination.parent.is_dir(): + raise FileNotFoundError(f"Plot output directory does not exist: {destination.parent}") + try: + handle = destination.open("xb") + except FileExistsError as error: + raise FileExistsError( + f"Refusing to overwrite an existing result plot: {destination}" + ) from error + try: + with handle: + figure.savefig(handle, format="png", dpi=300) + except BaseException: + destination.unlink(missing_ok=True) + raise + return destination + + +def plot_regression_results( + preds: np.ndarray, + labels: np.ndarray, + output_path: str | Path, + task_name: str = "Regression", +) -> float: + """ + Plot regression results with Spearman correlation + + Args: + preds: Predicted values + labels: True values + output_path: New PNG path inside this task's reserved output directory + task_name: Name of the task for the plot title + + Returns: + correlation: Spearman correlation coefficient + """ + import matplotlib.pyplot as plt + import seaborn as sns + from scipy.stats import spearmanr + + # preds/labels: (n,) + correlation, p_value = spearmanr(preds, labels) + + figure, axis = plt.subplots(figsize=(10, 8)) + sns.scatterplot(x=labels, y=preds, alpha=0.6, ax=axis) + + sns.regplot(x=labels, y=preds, scatter=False, color="red", ax=axis) + + axis.set_title(f"{task_name} - Spearman Correlation: {correlation:.3f} (p={p_value:.3e})") + axis.set_xlabel("True Values") + axis.set_ylabel("Predicted Values") + + axis.annotate( + f"rho = {correlation:.3f}", + xy=(0.05, 0.95), + xycoords="axes fraction", + fontsize=12, + bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8), + ) + + figure.tight_layout() + try: + _save_figure_exclusive(figure, output_path) + finally: + plt.close(figure) + return correlation + + +def plot_classification_results( + trainer: Trainer, + test_dataset: Any, + output_path: str | Path, + task_name: str = "Classification", +) -> float: + """ + Plot classification results with confusion matrix + + Args: + trainer: The trained model trainer + test_dataset: Dataset to evaluate on + output_path: New PNG path inside this task's reserved output directory + task_name: Name of the task for the plot title + + Returns: + accuracy: Classification accuracy + """ + import matplotlib.pyplot as plt + from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix + + predictions, labels, _ = trainer.predict( + test_dataset + ) # predictions: (n, c); labels: (n,) + preds = predictions[0] if isinstance(predictions, tuple) else predictions # (n, c) + pred_values = np.argmax(preds, axis=1) # (n,) + + accuracy = (pred_values == labels).mean() # () + + cm = confusion_matrix(labels, pred_values) # (c, c) + + figure, axis = plt.subplots(figsize=(10, 8)) + disp = ConfusionMatrixDisplay(confusion_matrix=cm) + disp.plot(cmap=plt.cm.Blues, ax=axis) + + axis.set_title(f"{task_name} - Accuracy: {accuracy:.3f}") + figure.tight_layout() + try: + _save_figure_exclusive(figure, output_path) + finally: + plt.close(figure) + + return accuracy + + +@_guard_training_output( + lora_default="./results_regression_lora", + full_default="./results_regression", +) +def train_regression_model( + model_name: str = DEFAULT_MODEL, + model_revision: str | None = None, + train_dataset_source: str = DEFAULT_REGRESSION_TRAIN_DATASET, + train_dataset_revision: str | None = None, + validation_dataset_source: str = DEFAULT_REGRESSION_VALIDATION_DATASET, + validation_dataset_revision: str | None = None, + test_dataset_source: str = DEFAULT_REGRESSION_TEST_DATASET, + test_dataset_revision: str | None = None, + use_lora: bool = True, + custom_lora_config: Any = None, + batch_size: int = 8, + learning_rate: float = 5e-5, + num_epochs: int = 10, + max_length: int = 1024, + gradient_accumulation_steps: int = 1, + patience: int = 3, + seed: int = 42, + full_determinism: bool = False, + plot_results: bool = False, + attn_backend: str = "sdpa", + output_dir: str | Path | None = None, +) -> tuple[Trainer, Any]: + """ + Train a regression model for protein-protein affinity prediction + + Args: + model_name: Name or path of the pretrained model + use_lora: Whether to use LoRA for fine-tuning + custom_lora_config: Custom LoRA configuration (optional) + batch_size: Batch size for training + learning_rate: Learning rate for training + num_epochs: Number of epochs for training + max_length: Encoded token budget for the complete protein pair, + including tokenizer-added separator and special tokens + gradient_accumulation_steps: Number of gradient accumulation steps + patience: Number of evaluation calls without improvement before + training stops + seed: Shared model, data-loader, and training seed + full_determinism: Request Transformers deterministic algorithms + plot_results: Generate a reporting-extra scatter plot after training + attn_backend: Explicit eager, SDPA, or Flex implementation + output_dir: New task-specific output directory; existing paths are rejected + + Returns: + trainer: The trained model trainer + test_dataset: The test dataset used for evaluation + """ + print("Loading datasets for regression task...") + if max_length <= 0: + raise ValueError("max_length must be a positive encoded token budget.") + set_seed(seed) + + # Validate every source contract before allocating or initializing a model. + train_data, train_source = _load_dataset_immutable( + train_dataset_source, + train_dataset_revision, + split="train", + ) + valid_data, validation_source = _load_dataset_immutable( + validation_dataset_source, + validation_dataset_revision, + split="train", + ) + test_data, test_source = _load_dataset_immutable( + test_dataset_source, + test_dataset_revision, + split="train", + ) + _validate_regression_dataset(train_data, split="train") + _validate_regression_dataset(valid_data, split="validation") + _validate_regression_dataset(test_data, split="test") + + # Resolve the tokenizer only after source validation. Raw residue counts + # omit the special separator/EOS tokens inserted for a protein pair. + model, tokenizer = initialize_model( + model_name=model_name, + model_revision=model_revision, + num_labels=1, + use_lora=use_lora, + lora_config=custom_lora_config, + attn_backend=attn_backend, + ) + + def _filter_pair_by_length(example: Any) -> bool: + return _fits_token_budget( + tokenizer, + example["SeqA"], + example["SeqB"], + max_length, + ) + + train_data = train_data.filter(_filter_pair_by_length) + valid_data = valid_data.filter(_filter_pair_by_length) + test_data = test_data.filter(_filter_pair_by_length) + _require_non_empty_filtered_split(train_data, split="train") + _require_non_empty_filtered_split(valid_data, split="validation") + _require_non_empty_filtered_split(test_data, split="test") + dataset_contracts = { + "train": { + "source": train_source, + "split": "train", + "columns": ("SeqA", "SeqB", "labels"), + }, + "validation": { + "source": validation_source, + "split": "train", + "columns": ("SeqA", "SeqB", "labels"), + }, + "test": { + "source": test_source, + "split": "train", + "columns": ("SeqA", "SeqB", "labels"), + }, + } + + train_dataset = PairDatasetHF(train_data, "SeqA", "SeqB", "labels", max_length=max_length) + valid_dataset = PairDatasetHF(valid_data, "SeqA", "SeqB", "labels", max_length=max_length) + test_dataset = PairDatasetHF(test_data, "SeqA", "SeqB", "labels", max_length=max_length) + + data_collator = PairCollator(tokenizer, regression=True, max_length=max_length) + + if output_dir is None: + raise RuntimeError("The output reservation guard did not supply a directory.") + output_dir = str(Path(output_dir)) + logging_dir = str(Path(output_dir) / "logs") + + training_args = TrainingArguments( + output_dir=output_dir, + num_train_epochs=num_epochs, + gradient_accumulation_steps=gradient_accumulation_steps, + per_device_train_batch_size=batch_size, + per_device_eval_batch_size=batch_size, + logging_dir=logging_dir, + learning_rate=learning_rate, + seed=seed, + data_seed=seed, + full_determinism=full_determinism, + **BASE_TRAINER_KWARGS, + ) + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=valid_dataset, + data_collator=data_collator, + compute_metrics=compute_metrics_regression, + callbacks=[EarlyStoppingCallback(early_stopping_patience=patience)], + ) + + metrics = trainer.evaluate(test_dataset) + print(f"Initial metrics: {metrics}") + print("Training regression model...") + trainer.train() + + final_artifact = _save_reload_verify_final_artifact( + trainer, + tokenizer, + output_dir=output_dir, + model_name=model_name, + model_revision=model_revision, + num_labels=1, + use_lora=use_lora, + verification_dataset=test_dataset, + data_collator=data_collator, + attn_backend=attn_backend, + ) + _write_training_manifest( + output_dir, + task="protein_pair_regression", + model=trainer.model, + tokenizer=tokenizer, + model_name=model_name, + model_revision=model_revision, + seed=seed, + max_length=max_length, + use_lora=use_lora, + batch_size=batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + learning_rate=learning_rate, + num_epochs=num_epochs, + full_determinism=full_determinism, + datasets={"train": train_data, "validation": valid_data, "test": test_data}, + dataset_contracts=dataset_contracts, + training_arguments=training_args, + patience=patience, + final_artifact=final_artifact, + requested_attention_backend=attn_backend, + ) + + print("Evaluating and visualizing results...") + predictions, labels, _prediction_metrics = trainer.predict( + test_dataset + ) # predictions: (n, 1); labels: (n,) + preds = predictions[0] if isinstance(predictions, tuple) else predictions # (n, 1) + label_values = cast(np.ndarray, labels) # (n,) + if plot_results: + correlation = plot_regression_results( + preds.flatten(), + label_values.flatten(), + Path(output_dir) / "regression_results.png", + "Protein-Protein Affinity", + ) + else: + correlation = _spearman_correlation(preds, label_values) + print(f"Final Spearman correlation on test set: {correlation:.3f}") + return trainer, test_dataset + + +@_guard_training_output( + lora_default="./results_classification_lora", + full_default="./results_classification", +) +def train_classification_model( + model_name: str = DEFAULT_MODEL, + model_revision: str | None = None, + dataset_source: str = DEFAULT_CLASSIFICATION_DATASET, + dataset_revision: str | None = None, + use_lora: bool = True, + custom_lora_config: Any = None, + batch_size: int = 8, + learning_rate: float = 5e-5, + num_epochs: int = 10, + max_length: int = 512, + gradient_accumulation_steps: int = 1, + patience: int = 3, + seed: int = 42, + full_determinism: bool = False, + plot_results: bool = False, + attn_backend: str = "sdpa", + output_dir: str | Path | None = None, +) -> Trainer: + """ + Train a classification model for protein solubility prediction + + Args: + model_name: Name or path of the pretrained model + use_lora: Whether to use LoRA for fine-tuning + custom_lora_config: Custom LoRA configuration (optional) + batch_size: Batch size for training + learning_rate: Learning rate for training + num_epochs: Number of epochs for training + max_length: Encoded token budget including tokenizer-added special tokens + gradient_accumulation_steps: Number of gradient accumulation steps + patience: Number of evaluation calls without improvement before + training stops + seed: Shared model, data-loader, and training seed + full_determinism: Request Transformers deterministic algorithms + plot_results: Generate a reporting-extra confusion matrix after training + attn_backend: Explicit eager, SDPA, or Flex implementation + output_dir: New task-specific output directory; existing paths are rejected + + Returns: + trainer: The trained model trainer + """ + print("Loading datasets for classification task...") + if max_length <= 0: + raise ValueError("max_length must be a positive encoded token budget.") + set_seed(seed) + + data, dataset_source_identity = _load_dataset_immutable( + dataset_source, + dataset_revision, + ) + num_labels = _validate_classification_dataset_dict(data) + model, tokenizer = initialize_model( + model_name=model_name, + model_revision=model_revision, + num_labels=num_labels, + use_lora=use_lora, + lora_config=custom_lora_config, + attn_backend=attn_backend, + ) + + def _filter_by_length(example: Any) -> bool: + return _fits_token_budget(tokenizer, example["seqs"], None, max_length) + + train_data = data["train"].filter(_filter_by_length) + valid_data = data["valid"].filter(_filter_by_length) + test_data = data["test"].filter(_filter_by_length) + _require_non_empty_filtered_split(train_data, split="train") + _require_non_empty_filtered_split(valid_data, split="valid") + _require_non_empty_filtered_split(test_data, split="test") + filtered_num_labels = _validate_classification_dataset_dict( + {"train": train_data, "valid": valid_data, "test": test_data} + ) + if filtered_num_labels != num_labels: + raise ValueError( + "Encoded-token filtering removed every training row for at least one " + "declared class; increase max_length or repair the split before training." + ) + dataset_contracts = { + split_name: { + "source": dataset_source_identity, + "split": source_split, + "columns": ("seqs", "labels"), + } + for split_name, source_split in ( + ("train", "train"), + ("validation", "valid"), + ("test", "test"), + ) + } + + train_dataset = SequenceDatasetHF(train_data, "seqs", "labels", max_length=max_length) + valid_dataset = SequenceDatasetHF(valid_data, "seqs", "labels", max_length=max_length) + test_dataset = SequenceDatasetHF(test_data, "seqs", "labels", max_length=max_length) + + data_collator = SequenceCollator(tokenizer, regression=False, max_length=max_length) + + if output_dir is None: + raise RuntimeError("The output reservation guard did not supply a directory.") + output_dir = str(Path(output_dir)) + logging_dir = str(Path(output_dir) / "logs") + + training_args = TrainingArguments( + output_dir=output_dir, + num_train_epochs=num_epochs, + gradient_accumulation_steps=gradient_accumulation_steps, + per_device_train_batch_size=batch_size, + per_device_eval_batch_size=batch_size, + logging_dir=logging_dir, + learning_rate=learning_rate, + seed=seed, + data_seed=seed, + full_determinism=full_determinism, + **BASE_TRAINER_KWARGS, + ) + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=valid_dataset, + data_collator=data_collator, + compute_metrics=compute_metrics_classification, + callbacks=[EarlyStoppingCallback(early_stopping_patience=patience)], + ) + + metrics = trainer.evaluate(test_dataset) + print(f"Initial metrics: {metrics}") + print("Training classification model...") + trainer.train() + + final_artifact = _save_reload_verify_final_artifact( + trainer, + tokenizer, + output_dir=output_dir, + model_name=model_name, + model_revision=model_revision, + num_labels=num_labels, + use_lora=use_lora, + verification_dataset=test_dataset, + data_collator=data_collator, + attn_backend=attn_backend, + ) + _write_training_manifest( + output_dir, + task="protein_sequence_classification", + model=trainer.model, + tokenizer=tokenizer, + model_name=model_name, + model_revision=model_revision, + seed=seed, + max_length=max_length, + use_lora=use_lora, + batch_size=batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + learning_rate=learning_rate, + num_epochs=num_epochs, + full_determinism=full_determinism, + datasets={"train": train_data, "validation": valid_data, "test": test_data}, + dataset_contracts=dataset_contracts, + training_arguments=training_args, + patience=patience, + final_artifact=final_artifact, + requested_attention_backend=attn_backend, + ) + + print("Evaluating and visualizing results...") + if plot_results: + accuracy = plot_classification_results( + trainer, + test_dataset, + Path(output_dir) / "classification_results.png", + "Protein Solubility", + ) + else: + predictions, labels, _ = trainer.predict( + test_dataset + ) # predictions: (n, c); labels: (n,) + preds = predictions[0] if isinstance(predictions, tuple) else predictions # (n, c) + label_values = cast(np.ndarray, labels) # (n,) + accuracy = float( + ( + np.argmax(preds, axis=-1).reshape(-1) + == label_values.reshape(-1) + ).mean() + ) # () + print(f"Final accuracy on test set: {accuracy:.3f}") + + return trainer + + +MODEL_LIST = [ + "Synthyra/ESM2-8M", + "Synthyra/ESM2-35M", + "Synthyra/ESM2-150M", + "Synthyra/ESM2-650M", +] + + +def build_parser() -> argparse.ArgumentParser: + """Build the public fine-tuning command-line interface.""" + + parser = argparse.ArgumentParser(description="Train models for protein tasks") + parser.add_argument( + "--task", + type=str, + choices=["regression", "classification", "both"], + default="both", + help="Task to train model for", + ) + parser.add_argument( + "--model_path", type=str, default=DEFAULT_MODEL, help="Path to the model to train" + ) + parser.add_argument( + "--model-revision", + help=( + "Immutable 40-character Hub commit for the model; the shipped default model " + "is pinned automatically, and custom remote models require this argument" + ), + ) + parser.add_argument( + "--classification-dataset-source", + default=DEFAULT_CLASSIFICATION_DATASET, + ) + parser.add_argument( + "--classification-dataset-revision", + help="Immutable 40-character Hub commit for the classification dataset", + ) + parser.add_argument( + "--regression-train-dataset-source", + default=DEFAULT_REGRESSION_TRAIN_DATASET, + ) + parser.add_argument("--regression-train-dataset-revision") + parser.add_argument( + "--regression-validation-dataset-source", + default=DEFAULT_REGRESSION_VALIDATION_DATASET, + ) + parser.add_argument("--regression-validation-dataset-revision") + parser.add_argument( + "--regression-test-dataset-source", + default=DEFAULT_REGRESSION_TEST_DATASET, + ) + parser.add_argument("--regression-test-dataset-revision") + parser.add_argument( + "--use-lora", + "--use_lora", + action=argparse.BooleanOptionalAction, + default=True, + help="Use LoRA (pass --no-use-lora to fine-tune the full model)", + ) + parser.add_argument("--batch_size", type=int, default=2, help="Batch size for training") + parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate for training") + parser.add_argument("--epochs", type=float, default=1.0, help="Number of epochs for training") + parser.add_argument( + "--max_length", + type=int, + default=512, + help=( + "Maximum encoded token count, including tokenizer-added special and pair " + "separator tokens" + ), + ) + parser.add_argument( + "--attn-backend", + choices=EXAMPLE_ATTENTION_BACKENDS, + default="sdpa", + help=( + "Explicit eager, SDPA, or Flex implementation recorded in the run manifest; " + "FlashAttention training needs a separate explicit BF16 CUDA loading policy" + ), + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("artifacts/fine-tuning"), + help=("Parent for isolated task runs; each selected task path must not already exist"), + ) + parser.add_argument( + "--grad_accum", type=int, default=1, help="Number of gradient accumulation steps" + ) + parser.add_argument( + "--patience", + type=int, + default=3, + help=("Number of evaluation calls without improvement before early stopping"), + ) + parser.add_argument("--seed", type=int, default=42, help="Training and data-loader seed") + parser.add_argument( + "--full-determinism", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable Transformers deterministic algorithms (potentially slower)", + ) + parser.add_argument( + "--plot-results", + action=argparse.BooleanOptionalAction, + default=False, + help="Generate 300 dpi result plots using the reporting extra", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + """Run sequence classification, protein-pair regression, or both.""" + + args = build_parser().parse_args(argv) + planned_output_dirs: list[Path] = [] + if args.task in ("regression", "both"): + planned_output_dirs.append( + args.output_dir / ("regression_lora" if args.use_lora else "regression") + ) + if args.task in ("classification", "both"): + planned_output_dirs.append( + args.output_dir / ("classification_lora" if args.use_lora else "classification") + ) + _ensure_output_paths_available(planned_output_dirs) + + print("\n" + "=" * 50) + print("TRAINING CONFIGURATION") + print("=" * 50) + print(f"Task: {args.task}") + print(f"Model revision: {args.model_revision}") + print(f"Using LoRA: {args.use_lora}") + print(f"Batch size: {args.batch_size}") + print(f"Learning rate: {args.lr}") + print(f"Number of epochs: {args.epochs}") + print(f"Max encoded token budget: {args.max_length}") + print(f"Attention backend: {args.attn_backend}") + print(f"Output root: {args.output_dir}") + print(f"Gradient Accumulation Steps: {args.grad_accum}") + print(f"Early stopping patience: {args.patience}") + print(f"Seed: {args.seed}") + print(f"Full determinism: {args.full_determinism}") + print("=" * 50 + "\n") + + if args.task in ["regression", "both"]: + print("\n" + "=" * 50) + print("TRAINING REGRESSION MODEL") + print("=" * 50) + train_regression_model( + model_name=args.model_path, + model_revision=args.model_revision, + train_dataset_source=args.regression_train_dataset_source, + train_dataset_revision=args.regression_train_dataset_revision, + validation_dataset_source=args.regression_validation_dataset_source, + validation_dataset_revision=args.regression_validation_dataset_revision, + test_dataset_source=args.regression_test_dataset_source, + test_dataset_revision=args.regression_test_dataset_revision, + use_lora=args.use_lora, + batch_size=args.batch_size, + learning_rate=args.lr, + num_epochs=args.epochs, + max_length=args.max_length, + gradient_accumulation_steps=args.grad_accum, + patience=args.patience, + seed=args.seed, + full_determinism=args.full_determinism, + plot_results=args.plot_results, + attn_backend=args.attn_backend, + output_dir=(args.output_dir / ("regression_lora" if args.use_lora else "regression")), + ) + + if args.task in ["classification", "both"]: + print("\n" + "=" * 50) + print("TRAINING CLASSIFICATION MODEL") + print("=" * 50) + train_classification_model( + model_name=args.model_path, + model_revision=args.model_revision, + dataset_source=args.classification_dataset_source, + dataset_revision=args.classification_dataset_revision, + use_lora=args.use_lora, + batch_size=args.batch_size, + learning_rate=args.lr, + num_epochs=args.epochs, + max_length=args.max_length, + gradient_accumulation_steps=args.grad_accum, + patience=args.patience, + seed=args.seed, + full_determinism=args.full_determinism, + plot_results=args.plot_results, + attn_backend=args.attn_backend, + output_dir=( + args.output_dir / ("classification_lora" if args.use_lora else "classification") + ), + ) + + print("\nTraining completed!") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/generation.py b/examples/generation.py new file mode 100644 index 0000000..b57fa4b --- /dev/null +++ b/examples/generation.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""Run deterministic DPLM, DPLM2, or conditioned ESM3 generation offline.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +if __package__: + from ._runtime import add_execution_arguments, resolve_execution +else: + from _runtime import add_execution_arguments, resolve_execution + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def generate_dplm(model: Any, tokenizer: Any, length: int, steps: int, seed: int) -> Any: + import torch + + input_ids = tokenizer("A" * length, return_tensors="pt")["input_ids"].to( + model.device + ) # (1, l_t) + with torch.random.fork_rng(), torch.inference_mode(): + torch.manual_seed(seed) + output_tokens = model.generate( + input_ids, + max_iter=steps, + sampling_strategy="argmax", + disable_resample=True, + ) # (1, l_t) + return output_tokens # (1, l_t) + + +def generate_dplm2(model: Any, tokenizer: Any, length: int, steps: int, seed: int) -> Any: + import torch + + vocab = tokenizer.get_vocab() + structure = [ + vocab[""], + *([vocab[""]] * length), + vocab[""], + ] + amino_acids = [ + vocab[""], + *([vocab[""]] * length), + vocab[""], + ] + input_ids = torch.tensor( + [structure + amino_acids], device=model.device + ) # (1, 2 * (l + 2)) + with torch.random.fork_rng(), torch.inference_mode(): + torch.manual_seed(seed) + output_tokens = model.generate( + input_ids, + max_iter=steps, + sampling_strategy="argmax", + unmasking_strategy="deterministic", + )["output_tokens"] # (1, 2 * (l + 2)) + return output_tokens # (1, 2 * (l + 2)) + + +def generate_esm3(model: Any, request: str | dict[str, Any], steps: int, seed: int) -> Any: + from fastplms.models.esm3.modeling_esm3 import FastESM3GenerationConfig + + return model.generate( + request, + FastESM3GenerationConfig(num_steps=steps, temperature=1.0, seed=seed), + ) + + +def build_esm3_multimodal_request(model: Any, prompt: str) -> dict[str, Any]: + """Build a synthetic request carrying every supported conditioning track. + + Replace these valid placeholder tracks with model-prepared biological + conditioning in a scientific workflow. + """ + import torch + + encoded = model.encode( + prompt, device=model.device + ) # input_ids/attention_mask: (b, l) + sequence_tokens = encoded["input_ids"] # (b, l) + shape = sequence_tokens.shape + device = sequence_tokens.device + return { + "sequence_tokens": sequence_tokens, # (b, l) + "attention_mask": encoded["attention_mask"], # (b, l) + "structure_tokens": torch.zeros( + shape, dtype=torch.long, device=device + ), # (b, l) + "ss8_tokens": torch.zeros(shape, dtype=torch.long, device=device), # (b, l) + "sasa_tokens": torch.zeros(shape, dtype=torch.long, device=device), # (b, l) + "function_tokens": torch.zeros( + (*shape, 8), dtype=torch.long, device=device + ), # (b, l, 8) + "residue_annotation_tokens": torch.zeros( + (*shape, 16), dtype=torch.long, device=device + ), # (b, l, 16) + "average_plddt": torch.ones(shape, device=device), # (b, l) + "per_res_plddt": torch.zeros(shape, device=device), # (b, l) + "structure_coords": torch.full( + (*shape, 3, 3), float("nan"), device=device + ), # (b, l, 3, 3) + "chain_id": torch.zeros(shape, dtype=torch.long, device=device), # (b, l) + "sequence_id": torch.ones(shape, dtype=torch.bool, device=device), # (b, l) + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("family", choices=("dplm", "dplm2", "esm3")) + parser.add_argument("artifact", type=Path) + parser.add_argument("--length", type=int, default=32) + parser.add_argument("--steps", type=int, default=8) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--esm3-prompt", default="MK____A") + add_execution_arguments(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + if arguments.family == "esm3": + from transformers import AutoModel + + model = AutoModel.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ).to(device).eval() + request = build_esm3_multimodal_request(model, arguments.esm3_prompt) + output = generate_esm3(model, request, arguments.steps, arguments.seed) + else: + from transformers import AutoModelForMaskedLM, AutoTokenizer + + model = AutoModelForMaskedLM.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ).to(device).eval() + tokenizer = AutoTokenizer.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + ) + if arguments.family == "dplm": + output = generate_dplm( + model, + tokenizer, + arguments.length, + arguments.steps, + arguments.seed, + ) + else: + output = generate_dplm2( + model, + tokenizer, + arguments.length, + arguments.steps, + arguments.seed, + ) + print(output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/structure_preparation.py b/examples/structure_preparation.py new file mode 100644 index 0000000..9b2e1c6 --- /dev/null +++ b/examples/structure_preparation.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Prepare ESMFold2 complexes or run seeded local structure helpers. + +The ESMFold2 branch deliberately includes an MSA and therefore requires a full +48-block checkpoint, not an inference-optimized Fast checkpoint. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +from typing import Any + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def build_esmfold2_conditioned_complex(types: Any) -> Any: + """Construct the full-variant multimolecule, MSA, bond, and distogram schema.""" + import numpy as np + + protein = "MSTNPKPQRKTKRNT" + msa = types.MSA.from_sequences([protein, "MSTNPKPQRKTKRNS"]) + distances = np.full( + (len(protein), len(protein)), 8.0, dtype=np.float32 + ) # (l, l) + np.fill_diagonal(distances, 0.0) # (l, l) + return types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence=protein, msa=msa), + types.ProteinInput( + id="B", + sequence="MKTIIALSYIFCLVFA", + modifications=[types.Modification(position=0, ccd="MSE")], + ), + types.RNAInput(id="R", sequence="AUGC"), + types.DNAInput(id="D", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ], + distogram_conditioning=[ + types.DistogramConditioning(chain_id="A", distogram=distances) + ], + covalent_bonds=[ + types.CovalentBond( + chain_id1="B", + res_idx1=0, + atom_idx1=0, + chain_id2="L", + res_idx2=0, + atom_idx2=0, + ) + ], + ) + + +def prepare_esmfold2_complex(model: Any, seed: int) -> tuple[Any, Any]: + types = model.input_types + request = build_esmfold2_conditioned_complex(types) + return model.prepare_structure_input(request, seed=seed) + + +def verify_esmfold2_pocket_rejection(model: Any, seed: int) -> str: + """Show the explicit boundary inherited from the published feature pipeline.""" + types = model.input_types + request = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="target", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="binder", sequence="MKTIIALSYIFCLVFA"), + ], + pocket=types.PocketConditioning( + binder_chain_id="binder", + contacts=[("target", 0)], + ), + ) + try: + model.prepare_structure_input(request, seed=seed) + except NotImplementedError as error: + return str(error) + raise RuntimeError("ESMFold2 unexpectedly accepted unsupported pocket conditioning.") + + +def run_structure_helper(model: Any, family: str, sequence: str, seed: int) -> Any: + if family == "boltz2": + return model.predict_structure( + amino_acid_sequence=sequence, + recycling_steps=1, + num_sampling_steps=8, + diffusion_samples=1, + seed=seed, + ) + if family == "esmfold": + return model.fold_protein(sequence) + raise ValueError(f"No direct structure helper for {family!r}.") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("family", choices=("esmfold2", "esmfold", "boltz2")) + parser.add_argument( + "artifact", + type=Path, + help=( + "Local artifact; the ESMFold2 MSA branch requires a full 48-block " + "checkpoint and rejects Fast variants" + ), + ) + parser.add_argument( + "--sequence", + default="MSTNPKPQRKTKRNT", + help="Protein sequence; ESMFold also accepts colon-delimited multimers", + ) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--device", default="cuda:0") + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + + configure_offline() + from transformers import AutoModel + + model = AutoModel.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + device_map={"": arguments.device}, + ).eval() + if arguments.family == "esmfold2": + features, chain_info = prepare_esmfold2_complex(model, arguments.seed) + print("feature-keys", sorted(features)) + print("chains", len(chain_info)) + pocket_contract = verify_esmfold2_pocket_rejection(model, arguments.seed) + print("pocket-contract", pocket_contract) + else: + result = run_structure_helper( + model, + arguments.family, + arguments.sequence, + arguments.seed, + ) + print(type(result).__name__) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/task_heads.py b/examples/task_heads.py new file mode 100644 index 0000000..58a0af2 --- /dev/null +++ b/examples/task_heads.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +"""Run ESM2 masked-LM, contact, sequence, and token task heads offline.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path +from typing import Any + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def resolve_execution(device_name: str, dtype_name: str) -> tuple[Any, Any]: + """Validate the portable CPU/CUDA execution requested by the user.""" + + import torch + + try: + device = torch.device(device_name) + except (RuntimeError, TypeError) as error: + raise ValueError(f"Invalid execution device {device_name!r}") from error + if device.type not in {"cpu", "cuda"}: + raise ValueError(f"Only CPU and CUDA devices are supported, got {device.type!r}") + if device.type == "cuda" and not torch.cuda.is_available(): + raise ValueError(f"CUDA device {device} was requested but CUDA is unavailable") + dtype = torch.float32 if dtype_name == "float32" else torch.bfloat16 + return device, dtype + + +def _biological_mask(tokenizer: Any, batch: dict[str, Any]) -> Any: + # batch input_ids/attention_mask: (b, l) + mask = batch["attention_mask"].bool() # (b, l) + for token_id in getattr(tokenizer, "all_special_ids", ()): + mask &= batch["input_ids"].ne(int(token_id)) # (b, l) + return mask # (b, l) + + +def _loading_key(value: Any) -> str: + if isinstance(value, str): + return value + if isinstance(value, (tuple, list)) and value: + return str(value[0]) + return str(value) + + +def _require_checkpoint_heads( + loading_info: dict[str, Any], + prefixes: tuple[str, ...], +) -> None: + """Reject a checkpoint load that silently initialized an advertised trained head.""" + + problems: dict[str, list[str]] = {} + for field in ("missing_keys", "mismatched_keys"): + matching = [ + key + for item in loading_info.get(field, ()) + if (key := _loading_key(item)).startswith(prefixes) + ] + if matching: + problems[field] = matching + error_messages = [str(value) for value in loading_info.get("error_msgs", ())] + if error_messages: + problems["error_msgs"] = error_messages + if problems: + raise RuntimeError( + "The local artifact does not contain the complete checkpoint-provided " + f"masked-LM/contact head state: {problems}" + ) + + +def _require_finite_tensor(name: str, value: Any) -> None: + import torch + + # value: (...) + if not bool(torch.isfinite(value).all().item()): + raise RuntimeError(f"{name} contained non-finite values") + + +def run_task_heads( + artifact: Path, + sequences: list[str], + *, + device: Any, + dtype: Any, + attn_backend: str, + num_labels: int, +) -> dict[str, Any]: + """Run the trained ESM2 heads and smoke separately initialized task heads.""" + + import torch + from transformers import ( + AutoModelForMaskedLM, + AutoModelForSequenceClassification, + AutoModelForTokenClassification, + AutoTokenizer, + ) + + common = { + "trust_remote_code": True, + "local_files_only": True, + "attn_implementation": attn_backend, + "dtype": dtype, + } + tokenizer = AutoTokenizer.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + ) + batch = tokenizer( + sequences, padding=True, return_tensors="pt" + ) # each tensor: (b, l) + batch = { + name: tensor.to(device) for name, tensor in batch.items() + } # each tensor: (b, l) + biological_mask = _biological_mask(tokenizer, batch) # (b, l) + if not biological_mask.any(dim=1).all(): # (b,) -> () + raise ValueError("Every input sequence must contain at least one biological residue") + + masked_lm, loading_info = AutoModelForMaskedLM.from_pretrained( + artifact, + output_loading_info=True, + **common, + ) + _require_checkpoint_heads(loading_info, ("lm_head.", "esm.contact_head.")) + masked_lm = masked_lm.to(device).eval() + masked_ids = batch["input_ids"].clone() # (b, l) + labels = torch.full_like(masked_ids, -100) # (b, l) + mask_token_id = getattr(tokenizer, "mask_token_id", None) + if mask_token_id is None: + raise ValueError("Masked-LM scoring requires a tokenizer mask token") + scored_positions: list[int] = [] + for row in range(masked_ids.shape[0]): + position = int(torch.nonzero(biological_mask[row], as_tuple=False)[0, 0]) + scored_positions.append(position) + labels[row, position] = masked_ids[row, position] # scalar assignment; (b, l) + masked_ids[row, position] = int(mask_token_id) # scalar assignment; (b, l) + + with torch.inference_mode(): + mlm_output = masked_lm( + input_ids=masked_ids, + attention_mask=batch["attention_mask"], + labels=labels, + ) # loss: (); logits: (b, l, v) + contacts = masked_lm.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) # (b, r, r) + probabilities = mlm_output.logits.float().softmax(dim=-1) # (b, l, v) + residue_probabilities = [ + float(probabilities[row, position, labels[row, position]].item()) + for row, position in enumerate(scored_positions) + ] + _require_finite_tensor("Masked-LM probabilities", probabilities) + _require_finite_tensor("Contact predictions", contacts) + + sequence_model = ( + AutoModelForSequenceClassification.from_pretrained( + artifact, + num_labels=num_labels, + **common, + ) + .to(device) + .eval() + ) + sequence_labels = torch.zeros( + len(sequences), dtype=torch.long, device=device + ) # (b,) + with torch.inference_mode(): + sequence_output = sequence_model( + **batch, labels=sequence_labels + ) # loss: (); logits: (b, c) + + token_model = ( + AutoModelForTokenClassification.from_pretrained( + artifact, + num_labels=num_labels, + **common, + ) + .to(device) + .eval() + ) + token_labels = torch.full_like(batch["input_ids"], -100) # (b, l) + token_labels[biological_mask] = 0 # (b, l) + with torch.inference_mode(): + token_output = token_model( + **batch, labels=token_labels + ) # loss: (); logits: (b, l, c) + + losses = { + "masked_lm": float(mlm_output.loss.item()), + "sequence_classification": float(sequence_output.loss.item()), + "token_classification": float(token_output.loss.item()), + } + if not all(math.isfinite(value) for value in losses.values()): + raise RuntimeError(f"A task-head loss was non-finite: {losses}") + return { + "sequences": len(sequences), + "device": str(device), + "dtype": str(dtype), + "attention_backend": attn_backend, + "masked_lm": { + "status": "checkpoint-provided pretrained head", + "loss": losses["masked_lm"], + "scored_positions": scored_positions, + "residue_probabilities": residue_probabilities, + }, + "contacts": { + "status": "checkpoint-provided pretrained head", + "shape": list(contacts.shape), + "finite": True, + }, + "sequence_classification": { + "status": "base weights + untrained task head", + "loss": losses["sequence_classification"], + "logits_shape": list(sequence_output.logits.shape), + }, + "token_classification": { + "status": "base weights + untrained task head", + "loss": losses["token_classification"], + "logits_shape": list(token_output.logits.shape), + }, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path, help="Local manifest-built ESM2 artifact") + parser.add_argument("--sequence", action="append", dest="sequences") + parser.add_argument("--device", default="cpu", help="cpu or cuda[:index]") + parser.add_argument("--dtype", choices=("float32", "bfloat16"), default="float32") + parser.add_argument( + "--attn-backend", + choices=("eager", "sdpa", "flex_attention"), + default="sdpa", + help="Portable backend used by every loaded head", + ) + parser.add_argument("--num-labels", type=int, default=2) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local ESM2 artifact: {artifact}") + if arguments.num_labels < 2: + raise SystemExit("--num-labels must be at least 2") + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + configure_offline() + summary = run_task_heads( + artifact, + arguments.sequences or ["MSTNPKPQRKTKRNT", "MKTII"], + device=device, + dtype=dtype, + attn_backend=arguments.attn_backend, + num_labels=arguments.num_labels, + ) + print(json.dumps(summary, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/ttt.py b/examples/ttt.py new file mode 100644 index 0000000..386c744 --- /dev/null +++ b/examples/ttt.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Run seeded test-time training, persist it, reset, and reload offline.""" + +from __future__ import annotations + +import argparse +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any + + +if __package__: + from ._runtime import add_execution_arguments, resolve_execution +else: + from _runtime import add_execution_arguments, resolve_execution + + +def configure_offline() -> None: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + + +def adapt_and_save(model: Any, sequence: str, output: Path, seed: int) -> Any: + if output.exists(): + raise FileExistsError(f"Refusing to overwrite an existing TTT artifact: {output}") + output.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{output.name}-", dir=output.parent)) + adaptation_started = False + try: + adaptation_started = True + metrics = model.ttt( + seq=sequence, + ttt_config={ + "steps": 3, + "ags": 1, + "batch_size": 1, + "seed": seed, + "initial_state_reset": True, + }, + ) + model.save_pretrained(staging, safe_serialization=True) + if not (staging / "config.json").is_file() or not any(staging.glob("*.safetensors")): + raise RuntimeError("TTT staging output is missing config or safetensors weights") + if output.exists(): + raise FileExistsError(f"TTT output appeared during staging: {output}") + staging.rename(output) + return metrics + finally: + try: + if adaptation_started: + model.ttt_reset() + finally: + if staging.exists(): + shutil.rmtree(staging) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--sequence", default="MSTNPKPQRKTKRNT") + parser.add_argument("--seed", type=int, default=7) + add_execution_arguments(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + artifact = arguments.artifact.expanduser().resolve() + output = arguments.output.expanduser().resolve() + if not (artifact / "config.json").is_file(): + raise SystemExit(f"Not a local artifact: {artifact}") + if output == artifact or artifact in output.parents: + raise SystemExit("TTT output must not be the source artifact or a directory inside it") + if output.exists(): + raise SystemExit(f"Refusing to overwrite an existing TTT artifact: {output}") + try: + device, dtype = resolve_execution(arguments.device, arguments.dtype) + except ValueError as error: + raise SystemExit(str(error)) from error + + configure_offline() + from transformers import AutoModelForMaskedLM + + model = ( + AutoModelForMaskedLM.from_pretrained( + artifact, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ) + .to(device) + .eval() + ) + metrics = adapt_and_save(model, arguments.sequence, output, arguments.seed) + reloaded = ( + AutoModelForMaskedLM.from_pretrained( + output, + trust_remote_code=True, + local_files_only=True, + dtype=dtype, + ) + .to(device) + .eval() + ) + reloaded.ttt_reset() + print(metrics) + print("reloaded", type(reloaded).__name__) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/fastplms/ankh/README.md b/fastplms/ankh/README.md deleted file mode 100644 index 68f8a13..0000000 --- a/fastplms/ankh/README.md +++ /dev/null @@ -1,153 +0,0 @@ ---- -library_name: transformers -tags: -- protein language model -- biology ---- - -# FastANKH - -Fast, optimized implementations of ANKH protein language models (T5-based) with multi-backend attention support. - -**Requires PyTorch 2.11+** for Flash Attention 4 (FA4) backend support via flex attention. - -## Models - -| Model | Params | Layers | Hidden | Heads | Activation | Source | -|-------|--------|--------|--------|-------|------------|--------| -| ANKH_base | 453.3M | 48 | 768 | 12 | gelu_new | ElnaggarLab/ankh-base | -| ANKH_large | 1.15B | 48 | 1536 | 16 | gelu_new | ElnaggarLab/ankh-large | -| ANKH2_large | 1.15B | 24 | 1536 | 16 | silu | ElnaggarLab/ankh2-ext2 | -| ANKH3_large | 1.15B | 48 | 1536 | 16 | silu | ElnaggarLab/ankh3-large | -| ANKH3_xl | 3.49B | 48 | 2560 | 32 | silu | ElnaggarLab/ankh3-xl | - -## Usage - -```python -from transformers import AutoModel, AutoTokenizer - -model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True) -tokenizer = AutoTokenizer.from_pretrained("Synthyra/ANKH_base") - -# Set attention backend before loading for best performance -# Options: "sdpa" (default, exact), "flex" (FA4 on Hopper/Blackwell) -model.config.attn_backend = "flex" - -sequences = ["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"] -inputs = tokenizer(sequences, return_tensors="pt", padding=True) -outputs = model(**inputs) - -# Per-residue embeddings -embeddings = outputs.last_hidden_state # (batch, seq_len, hidden_dim) -``` - -## Experimental Test-Time Training - -TTT is disabled by default. Normal ANKH inference, embeddings, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -with masked language modeling on the test protein. ANKH's FastPLMs MaskedLM -head is encoder-only and not pretrained for standard MLM, so treat TTT results -with extra caution. - -```python -from transformers import AutoModelForMaskedLM - -mlm = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ANKH_base", - trust_remote_code=True, -).cuda().eval() - -metrics = mlm.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -mlm.ttt_reset() -print(metrics["losses"]) -``` - -## Batch Embedding - -```python -model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True).to("cuda") -embeddings = model.embed_dataset( - sequences=["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"], - batch_size=8, - max_len=512, - full_embeddings=True, -) -``` - -## Attention Backends - -| Backend | Key | Notes | -|---------|-----|-------| -| SDPA | `"sdpa"` | Default. Exact attention with position bias as additive mask. | -| Flex | `"flex"` | Uses FA4 on Hopper/Blackwell GPUs (PyTorch 2.11+). Position bias computed via `score_mod`. Triton fallback on older hardware. | -| Flash | `"kernels_flash"` | Not supported for ANKH (no arbitrary bias support). Falls back to flex/sdpa. | - -## Architecture - -ANKH models are T5 encoder-only architectures: -- **No absolute position embeddings**: Uses T5-style relative position bias (log-bucketed, bidirectional) -- **RMS LayerNorm**: No mean subtraction, no bias term -- **Gated FFN**: `activation(wi_0(x)) * wi_1(x) -> wo(x)` with gelu_new (v1) or silu (v2/v3) -- **Pre-layer normalization**: Norm before attention and FFN, residual after -- **No bias in projections**: All q/k/v/o and FFN linear layers are bias=False - -The relative position bias is computed once (materialized as a full tensor) and shared across all encoder layers. For the flex backend, the bias is passed as a `score_mod` closure for optimal performance. - -## Notes - -- The `FastAnkhForMaskedLM` variant includes an LM head initialized from the shared embedding weights. The original ANKH models were trained with T5's span corruption objective using an encoder-decoder architecture. This encoder-only MaskedLM head is **not pre-trained for standard MLM** and requires additional fine-tuning. -- `model.tokenizer` is loaded from the checkpoint hub id so each ANKH model uses its matching tokenizer. -- ANKH3 models use a vocabulary of 256 tokens (vs 144 for v1/v2) and were trained with dual objectives ([NLU] for embeddings, [S2S] for generation). - -## Citations - -```bibtex -@article{elnaggar2023ankh, - title={Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling}, - author={Elnaggar, Ahmed and Essam, Hazem and Salah-Eldin, Wafaa and Moustafa, Walid and Elkerdawy, Mohamed and Rochereau, Charlotte and Rost, Burkhard}, - journal={arXiv preprint arXiv:2301.06568}, - year={2023} -} -``` - -```bibtex -@article{alsamkary2025ankh3, - title={Ankh3: Multi-Task Pretraining with Sequence Denoising and Completion Enhances Protein Representations}, - author={Alsamkary, Hazem and Elshaffei, Mohamed and Elkerdawy, Mohamed and Elnaggar, Ahmed}, - journal={arXiv preprint arXiv:2505.20052}, - year={2025} -} -``` - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/ankh/ankh_license.txt b/fastplms/ankh/ankh_license.txt deleted file mode 100644 index 6c9215c..0000000 --- a/fastplms/ankh/ankh_license.txt +++ /dev/null @@ -1,114 +0,0 @@ -License for ANKH models, from their repo https://github.com/agemagician/Ankh?tab=License-1-ov-file -Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International - -Creative Commons Corporation ("Creative Commons") is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an "as-is" basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible. - -Using Creative Commons Public Licenses - -Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses. - -Considerations for licensors: Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. More considerations for licensors : wiki.creativecommons.org/Considerations_for_licensors - -Considerations for the public: By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor's permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. More considerations for the public : wiki.creativecommons.org/Considerations_for_licensees - -Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License - -By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions. - -Section 1 – Definitions. - -a. Adapted Material means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image. -b. Adapter's License means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License. -c. BY-NC-SA Compatible License means a license listed at creativecommons.org/compatiblelicenses, approved by Creative Commons as essentially the equivalent of this Public License. -d. Copyright and Similar Rights means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights. -e. Effective Technological Measures means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements. -f. Exceptions and Limitations means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material. -g. License Elements means the license attributes listed in the name of a Creative Commons Public License. The License Elements of this Public License are Attribution, NonCommercial, and ShareAlike. -h. Licensed Material means the artistic or literary work, database, or other material to which the Licensor applied this Public License. -i. Licensed Rights means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license. -j. Licensor means the individual(s) or entity(ies) granting rights under this Public License. -k. NonCommercial means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange. -l. Share means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them. -m. Sui Generis Database Rights means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world. -n. You means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning. -Section 2 – Scope. - -a. License grant. -Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to: -A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and -B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only. -Exceptions and Limitations. For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions. -Term. The term of this Public License is specified in Section 6(a). -Media and formats; technical modifications allowed. The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material. -Downstream recipients. -A. Offer from the Licensor – Licensed Material. Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License. -B. Additional offer from the Licensor – Adapted Material. Every recipient of Adapted Material from You automatically receives an offer from the Licensor to exercise the Licensed Rights in the Adapted Material under the conditions of the Adapter's License You apply. -C. No downstream restrictions. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material. -No endorsement. Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i). -b. Other rights. -Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise. -Patent and trademark rights are not licensed under this Public License. -To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes. -Section 3 – License Conditions. - -Your exercise of the Licensed Rights is expressly made subject to the following conditions. - -a. Attribution. -If You Share the Licensed Material (including in modified form), You must: -A. retain the following if it is supplied by the Licensor with the Licensed Material: - -i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated); -ii. a copyright notice; -iii. a notice that refers to this Public License; -iv. a notice that refers to the disclaimer of warranties; -v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable; -B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and - -C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License. - -You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information. -If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable. -b. ShareAlike.In addition to the conditions in Section 3(a), if You Share Adapted Material You produce, the following conditions also apply. -The Adapter's License You apply must be a Creative Commons license with the same License Elements, this version or later, or a BY-NC-SA Compatible License. -You must include the text of, or the URI or hyperlink to, the Adapter's License You apply. You may satisfy this condition in any reasonable manner based on the medium, means, and context in which You Share Adapted Material. -You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, Adapted Material that restrict exercise of the rights granted under the Adapter's License You apply. -Section 4 – Sui Generis Database Rights. - -Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material: - -a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only; -b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material, including for purposes of Section 3(b); and -c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database. -For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights. -Section 5 – Disclaimer of Warranties and Limitation of Liability. - -a. Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You. -b. To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You. -c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability. -Section 6 – Term and Termination. - -a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically. - -b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates: - -automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or -upon express reinstatement by the Licensor. -For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License. - -c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License. - -d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License. - -Section 7 – Other Terms and Conditions. - -a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed. -b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License. -Section 8 – Interpretation. - -a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License. -b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions. -c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor. -d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority. -Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the "Licensor." The text of the Creative Commons public licenses is dedicated to the public domain under the CC0 Public Domain Dedication. Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at creativecommons.org/policies, Creative Commons does not authorize the use of the trademark "Creative Commons" or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses. - -Creative Commons may be contacted at creativecommons.org. \ No newline at end of file diff --git a/fastplms/ankh/get_weights.py b/fastplms/ankh/get_weights.py deleted file mode 100644 index e57bd4a..0000000 --- a/fastplms/ankh/get_weights.py +++ /dev/null @@ -1,223 +0,0 @@ -import copy -import os - -import torch -from typing import Dict, List, Optional, Tuple - -from huggingface_hub import HfApi, login -from transformers import T5ForConditionalGeneration, T5Config, AutoTokenizer, AutoModel - -from fastplms.ankh.modeling_ankh import FastAnkhConfig, FastAnkhForMaskedLM -from fastplms.weight_parity_utils import assert_model_parameters_fp32 - - -MODEL_DICT = { - "Synthyra/ANKH_base": "ElnaggarLab/ankh-base", - "Synthyra/ANKH_large": "ElnaggarLab/ankh-large", - "Synthyra/ANKH2_large": "ElnaggarLab/ankh2-ext2", - "Synthyra/ANKH3_large": "ElnaggarLab/ankh3-large", - "Synthyra/ANKH3_xl": "ElnaggarLab/ankh3-xl", -} -SHARDED_REPO_IDS = {"Synthyra/ANKH3_xl"} -SHARD_SIZE = "5GB" - - -def _map_encoder_state_dict(official_sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """Map T5ForConditionalGeneration state dict to FastAnkh format (encoder only). - - FastAnkh mirrors T5 key naming, so encoder.* keys pass through directly. - We add shared.weight (= encoder.embed_tokens.weight) and lm_head.weight. - """ - new_sd = {} - for key, value in official_sd.items(): - if key.startswith("encoder.") or key == "shared.weight": - new_sd[key] = value.clone() - - assert "encoder.embed_tokens.weight" in new_sd, "Missing embed_tokens in mapped state dict" - new_sd["lm_head.weight"] = new_sd["encoder.embed_tokens.weight"].clone() - - return new_sd - - -def _build_config(source_repo: str) -> FastAnkhConfig: - """Build FastAnkhConfig from official T5 config.""" - t5_config = T5Config.from_pretrained(source_repo) - - # Determine activation function - act_info = t5_config.feed_forward_proj.split("-") - if t5_config.feed_forward_proj == "gated-gelu": - dense_act_fn = "gelu_new" - elif len(act_info) > 1: - dense_act_fn = act_info[-1] - else: - dense_act_fn = act_info[0] - - config = FastAnkhConfig( - vocab_size=t5_config.vocab_size, - d_model=t5_config.d_model, - d_kv=t5_config.d_kv, - d_ff=t5_config.d_ff, - num_heads=t5_config.num_heads, - num_layers=t5_config.num_layers, - relative_attention_num_buckets=t5_config.relative_attention_num_buckets, - relative_attention_max_distance=t5_config.relative_attention_max_distance, - dense_act_fn=dense_act_fn, - layer_norm_epsilon=t5_config.layer_norm_epsilon, - initializer_factor=t5_config.initializer_factor, - pad_token_id=t5_config.pad_token_id or 0, - eos_token_id=t5_config.eos_token_id or 1, - ) - config.auto_map = { - "AutoConfig": "modeling_ankh.FastAnkhConfig", - "AutoModel": "modeling_ankh.FastAnkhModel", - "AutoModelForMaskedLM": "modeling_ankh.FastAnkhForMaskedLM", - "AutoModelForSequenceClassification": "modeling_ankh.FastAnkhForSequenceClassification", - "AutoModelForTokenClassification": "modeling_ankh.FastAnkhForTokenClassification", - } - config.tie_word_embeddings = False - return config - - -def _delete_legacy_unsharded_weights_if_present(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - if "model.safetensors" in repo_files: - print(f"Deleting legacy unified model.safetensors from {repo_id}") - api.delete_file(path_in_repo="model.safetensors", repo_id=repo_id, repo_type="model") - - -def _assert_repo_has_sharded_weights(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - assert "model.safetensors.index.json" in repo_files, f"{repo_id} missing index file." - has_shards = any(f.startswith("model-") and f.endswith(".safetensors") for f in repo_files) - assert has_shards, f"{repo_id} has no shard files." - assert "model.safetensors" not in repo_files, f"{repo_id} still has unified weights." - - -def _push_model_with_expected_format(model: FastAnkhForMaskedLM, api: HfApi, repo_id: str) -> None: - if repo_id in SHARDED_REPO_IDS: - print(f"Pushing sharded weights for {repo_id} with max_shard_size={SHARD_SIZE}") - model.push_to_hub(repo_id, max_shard_size=SHARD_SIZE) - _delete_legacy_unsharded_weights_if_present(api, repo_id) - _assert_repo_has_sharded_weights(api, repo_id) - return - model.push_to_hub(repo_id) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return list(MODEL_DICT.items()) - selected = [] - for repo_id in repo_ids: - assert repo_id in MODEL_DICT, ( - f"Unknown repo {repo_id}. Valid: {sorted(MODEL_DICT.keys())}" - ) - selected.append((repo_id, MODEL_DICT[repo_id])) - return selected - - -if __name__ == "__main__": - # py -m fastplms.ankh.get_weights - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--dry_run", action="store_true") - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - api = HfApi() - - if args.hf_token is not None: - assert len(args.hf_token) > 0 - login(token=args.hf_token) - - script_root = os.path.dirname(os.path.abspath(__file__)) - - for repo_id, source_repo in _resolve_repo_items(args.repo_ids): - print(f"\n{'='*60}") - print(f"Processing {repo_id} <- {source_repo}") - print(f"{'='*60}") - - config = _build_config(source_repo) - - if args.skip_weights: - if args.dry_run: - print(f"[skip-weights][dry-run] validated config for {repo_id}") - continue - tokenizer = AutoTokenizer.from_pretrained(source_repo) - config.push_to_hub(repo_id) - tokenizer.push_to_hub(repo_id) - print(f"[skip-weights] uploaded config+tokenizer for {repo_id}") - continue - - # Load official T5 encoder-decoder - print(f"Loading official T5ForConditionalGeneration from {source_repo}...") - official_model = T5ForConditionalGeneration.from_pretrained( - source_repo, dtype=torch.float32, device_map="cpu", - ) - - # Map encoder weights to FastAnkh format - mapped_sd = _map_encoder_state_dict(official_model.state_dict()) - print(f"Mapped {len(mapped_sd)} parameters from encoder") - - # Create FastAnkh model and load mapped weights - model = FastAnkhForMaskedLM(config) - model = model.to(dtype=torch.float32) - - # Verify all expected keys are present - model_keys = set(model.state_dict().keys()) - mapped_keys = set(mapped_sd.keys()) - missing = model_keys - mapped_keys - unexpected = mapped_keys - model_keys - assert not missing, f"Missing keys in mapped state dict:\n{missing}" - assert not unexpected, f"Unexpected keys in mapped state dict:\n{unexpected}" - - model.load_state_dict(mapped_sd, strict=True) - - assert_model_parameters_fp32(model=model, model_name=f"FastAnkh ({repo_id})") - - # Verify encoder weight parity (compare mapped values directly) - print("Verifying encoder weight parity...") - for key, mapped_val in mapped_sd.items(): - if key == "lm_head.weight": - continue # LM head is a copy, not from official encoder - model_val = model.state_dict()[key] - mse = (model_val.float() - mapped_val.float()).pow(2).mean().item() - assert mse == 0.0, f"Weight mismatch at {key}: MSE={mse}" - print("Weight parity verified (MSE=0.0 for all encoder parameters)") - - if args.dry_run: - print(f"[dry_run] validated ANKH parity for {repo_id} <- {source_repo}") - continue - - # Push tokenizer - tokenizer = AutoTokenizer.from_pretrained(source_repo) - tokenizer.push_to_hub(repo_id) - - # Push model - _push_model_with_expected_format(model, api, repo_id) - - # Upload modeling file - api.upload_file( - path_or_fileobj=os.path.join(script_root, "modeling_ankh.py"), - path_in_repo="modeling_ankh.py", - repo_id=repo_id, - repo_type="model", - ) - - # Verify download - print(f"Verifying download from {repo_id}...") - downloaded_model = AutoModel.from_pretrained( - repo_id, dtype=torch.float32, device_map="cpu", - force_download=True, trust_remote_code=True, - ) - for key in downloaded_model.encoder.state_dict(): - orig = model.encoder.state_dict()[key] - dl = downloaded_model.encoder.state_dict()[key] - mse = (orig.float() - dl.float()).pow(2).mean().item() - assert mse == 0.0, f"Post-download mismatch at {key}: MSE={mse}" - print(f"Download verification passed for {repo_id}") diff --git a/fastplms/ankh/modeling_ankh.py b/fastplms/ankh/modeling_ankh.py deleted file mode 100644 index 0e290b1..0000000 --- a/fastplms/ankh/modeling_ankh.py +++ /dev/null @@ -1,936 +0,0 @@ -from __future__ import annotations - -import math - -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Optional, Tuple, Dict, Any -from dataclasses import dataclass -from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer -from transformers.modeling_outputs import ModelOutput - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, bool_to_additive_mask, - _get_flex_attention_fn, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -# --------------------------------------------------------------------------- -# Output dataclasses -# --------------------------------------------------------------------------- - -@dataclass -class AnkhEncoderOutput(ModelOutput): - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - - -@dataclass -class AnkhMaskedLMOutput(ModelOutput): - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- - -class FastAnkhConfig(PretrainedConfig): - model_type = "fast_ankh" - attribute_map = {"hidden_size": "d_model"} - - def __init__( - self, - vocab_size: int = 144, - d_model: int = 768, - d_kv: int = 64, - d_ff: int = 3072, - num_heads: int = 12, - num_layers: int = 48, - relative_attention_num_buckets: int = 64, - relative_attention_max_distance: int = 128, - dense_act_fn: str = "gelu_new", - layer_norm_epsilon: float = 1e-6, - initializer_factor: float = 1.0, - pad_token_id: int = 0, - eos_token_id: int = 1, - attn_backend: str = "sdpa", - **kwargs, - ): - super().__init__( - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - **kwargs, - ) - self.vocab_size = vocab_size - self.d_model = d_model - self.d_kv = d_kv - self.d_ff = d_ff - self.num_heads = num_heads - self.num_layers = num_layers - self.relative_attention_num_buckets = relative_attention_num_buckets - self.relative_attention_max_distance = relative_attention_max_distance - self.dense_act_fn = dense_act_fn - self.layer_norm_epsilon = layer_norm_epsilon - self.initializer_factor = initializer_factor - self.tie_word_embeddings = False - self.attn_backend = attn_backend - - def to_dict(self) -> Dict[str, Any]: - output = super().to_dict() - return output - - -def _load_ankh_tokenizer(config: FastAnkhConfig): - """Load the checkpoint-matched tokenizer, falling back only for bare configs.""" - name_or_path = config._name_or_path - if isinstance(name_or_path, str) and len(name_or_path) > 0: - return AutoTokenizer.from_pretrained(name_or_path) - return AutoTokenizer.from_pretrained("ElnaggarLab/ankh-base") - - -# --------------------------------------------------------------------------- -# Submodules -# --------------------------------------------------------------------------- - -class AnkhRMSNorm(nn.Module): - """T5-style RMS layer norm: scales without mean subtraction or bias.""" - - def __init__(self, hidden_size: int, eps: float = 1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(self.weight.dtype) - - -def _gelu_new(x: torch.Tensor) -> torch.Tensor: - return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) - - -class AnkhGatedFFN(nn.Module): - """T5-style gated feed-forward: activation(wi_0(x)) * wi_1(x) -> wo.""" - - def __init__(self, config: FastAnkhConfig): - super().__init__() - self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.act = F.silu if config.dense_act_fn == "silu" else _gelu_new - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.wo(self.act(self.wi_0(hidden_states)) * self.wi_1(hidden_states)) - - -# --------------------------------------------------------------------------- -# Attention -# --------------------------------------------------------------------------- - -class AnkhSelfAttention(nn.Module): - """T5-style self-attention with relative position bias and multi-backend dispatch. - - Only layer 0 has ``has_relative_attention_bias=True`` and owns the - ``nn.Embedding`` that produces the position bias. All other layers - receive the precomputed bias through the forward call. - """ - - def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False): - super().__init__() - self.num_heads = config.num_heads - self.d_kv = config.d_kv - self.inner_dim = self.num_heads * self.d_kv - self.has_relative_attention_bias = has_relative_attention_bias - self.relative_attention_num_buckets = config.relative_attention_num_buckets - self.relative_attention_max_distance = config.relative_attention_max_distance - - self.q = nn.Linear(config.d_model, self.inner_dim, bias=False) - self.k = nn.Linear(config.d_model, self.inner_dim, bias=False) - self.v = nn.Linear(config.d_model, self.inner_dim, bias=False) - self.o = nn.Linear(self.inner_dim, config.d_model, bias=False) - # T5/ANKH attention is unscaled: scores = Q K^T (no 1/sqrt(d_kv)). - # The learned relative position bias absorbs any temperature. - self.scale = 1.0 - - if self.has_relative_attention_bias: - self.relative_attention_bias = nn.Embedding( - config.relative_attention_num_buckets, config.num_heads - ) - - self.attn_backend: AttentionBackend = AttentionBackend.SDPA # set by encoder - - # ---- T5 relative position bucketing ---- - - @staticmethod - def _relative_position_bucket( - relative_position: torch.Tensor, - num_buckets: int = 32, - max_distance: int = 128, - ) -> torch.Tensor: - """Bidirectional log-bucketed relative position mapping (T5 style).""" - # Bidirectional: half buckets for negative, half for positive - num_buckets //= 2 - relative_buckets = (relative_position > 0).to(torch.long) * num_buckets - relative_position = torch.abs(relative_position) - - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - relative_position_if_large = max_exact + ( - torch.log(relative_position.float() / max_exact) - / math.log(max_distance / max_exact) - * (num_buckets - max_exact) - ).to(torch.long) - relative_position_if_large = torch.clamp(relative_position_if_large, max=num_buckets - 1) - - relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) - return relative_buckets - - def compute_bias(self, query_length: int, key_length: int, device: torch.device) -> torch.Tensor: - """Compute (1, H, Q, K) position bias tensor for SDPA / manual paths.""" - context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] - relative_position = memory_position - context_position - buckets = self._relative_position_bucket( - relative_position, - num_buckets=self.relative_attention_num_buckets, - max_distance=self.relative_attention_max_distance, - ) - values = self.relative_attention_bias(buckets) # (Q, K, H) - return values.permute(2, 0, 1).unsqueeze(0) # (1, H, Q, K) - - # ---- Forward ---- - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - position_bias: Optional[torch.Tensor] = None, - flex_score_mod=None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: - """Returns (attn_output, attn_weights_or_none, position_bias).""" - batch_size, seq_length = hidden_states.shape[:2] - hidden_shape = (batch_size, seq_length, self.num_heads, self.d_kv) - - query_BHLD = self.q(hidden_states).view(hidden_shape).transpose(1, 2) - key_BHLD = self.k(hidden_states).view(hidden_shape).transpose(1, 2) - value_BHLD = self.v(hidden_states).view(hidden_shape).transpose(1, 2) - - # Compute position bias on first layer (SDPA/manual only; flex uses score_mod) - if position_bias is None and self.has_relative_attention_bias and self.attn_backend != AttentionBackend.FLEX: - position_bias = self.compute_bias(seq_length, seq_length, hidden_states.device) - # Fold padding mask into position bias so layers don't need separate mask. - if attention_mask_4d is not None: - position_bias = position_bias + bool_to_additive_mask(attention_mask_4d, position_bias.dtype) - - if output_attentions: - attn_output, attn_weights = self._manual_attn(query_BHLD, key_BHLD, value_BHLD, position_bias) - return self.o(attn_output), attn_weights, position_bias - - if self.attn_backend == AttentionBackend.FLEX: - attn_output = self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask, flex_score_mod) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output = self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, position_bias) - else: - raise AssertionError(f"Unsupported backend for ANKH: {self.attn_backend}") - - return self.o(attn_output), None, position_bias - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - position_bias: Optional[torch.Tensor], - ) -> torch.Tensor: - # SDPA: position_bias is (1, H, Q, K) additive bias (includes padding mask) - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, - attn_mask=position_bias, - scale=self.scale, - ) - return context_BHLD.transpose(1, 2).contiguous().view( - query_BHLD.shape[0], -1, self.inner_dim - ) - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask], - flex_score_mod, - ) -> torch.Tensor: - assert flex_attention is not None, "Flex attention is not available." - fn = _get_flex_attention_fn() - context_BHLD = fn( - query_BHLD, key_BHLD, value_BHLD, - score_mod=flex_score_mod, - block_mask=flex_block_mask, - scale=self.scale, - ) - return context_BHLD.transpose(1, 2).contiguous().view( - query_BHLD.shape[0], -1, self.inner_dim - ) - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - position_bias: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) * self.scale - if position_bias is not None: - attn_weights = attn_weights + position_bias - attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = context_BHLD.transpose(1, 2).contiguous().view( - query_BHLD.shape[0], -1, self.inner_dim - ) - return attn_output, attn_weights - - -# --------------------------------------------------------------------------- -# Encoder block & stack (T5-compatible key naming) -# --------------------------------------------------------------------------- - -class AnkhSelfAttentionLayer(nn.Module): - """Wraps AnkhSelfAttention + layer_norm to match T5Block.layer[0] key naming.""" - - def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False): - super().__init__() - self.SelfAttention = AnkhSelfAttention(config, has_relative_attention_bias) - self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - position_bias: Optional[torch.Tensor] = None, - flex_score_mod=None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: - normed = self.layer_norm(hidden_states) - attn_output, attn_weights, position_bias = self.SelfAttention( - normed, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - position_bias=position_bias, - flex_score_mod=flex_score_mod, - output_attentions=output_attentions, - ) - hidden_states = hidden_states + attn_output - return hidden_states, attn_weights, position_bias - - -class AnkhFFLayer(nn.Module): - """Wraps AnkhGatedFFN + layer_norm to match T5Block.layer[1] key naming.""" - - def __init__(self, config: FastAnkhConfig): - super().__init__() - self.DenseReluDense = AnkhGatedFFN(config) - self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - normed = self.layer_norm(hidden_states) - hidden_states = hidden_states + self.DenseReluDense(normed) - return hidden_states - - -class AnkhBlock(nn.Module): - """Single transformer block with T5-compatible .layer ModuleList naming.""" - - def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False): - super().__init__() - self.layer = nn.ModuleList([ - AnkhSelfAttentionLayer(config, has_relative_attention_bias), - AnkhFFLayer(config), - ]) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - position_bias: Optional[torch.Tensor] = None, - flex_score_mod=None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: - hidden_states, attn_weights, position_bias = self.layer[0]( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - position_bias=position_bias, - flex_score_mod=flex_score_mod, - output_attentions=output_attentions, - ) - hidden_states = self.layer[1](hidden_states) - return hidden_states, attn_weights, position_bias - - -# --------------------------------------------------------------------------- -# PreTrainedModel base -# --------------------------------------------------------------------------- - -class AnkhPreTrainedModel(PreTrainedModel): - config_class = FastAnkhConfig - base_model_prefix = "encoder" - supports_gradient_checkpointing = True - _no_split_modules = ["AnkhBlock"] - - @classmethod - def is_remote_code(cls) -> bool: - return True - - @torch.no_grad() - def _init_weights(self, module: nn.Module) -> None: - factor = self.config.initializer_factor - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=factor * (self.config.d_model ** -0.5)) - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=factor * 1.0) - elif isinstance(module, AnkhRMSNorm): - module.weight.data.fill_(1.0) - - def post_init(self) -> None: - super().post_init() - - def get_output_embeddings(self): - return None - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, ( - f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - ) - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - if resolved == AttentionBackend.KERNELS_FLASH: - print("ANKH: kernels_flash -> flex/sdpa fallback") - resolved = AttentionBackend.FLEX if flex_attention is not None else AttentionBackend.SDPA - for module in self.modules(): - if isinstance(module, FAST_ANKH_ENCODER): - module.attention_backend = resolved - elif isinstance(module, AnkhSelfAttention): - module.attn_backend = resolved - - -# --------------------------------------------------------------------------- -# FAST_ANKH_ENCODER (mirrors T5Stack key naming) -# --------------------------------------------------------------------------- - -class FAST_ANKH_ENCODER(AnkhPreTrainedModel, EmbeddingMixin): - """Inner encoder that mirrors T5Stack attribute naming for weight compliance. - - State dict keys: embed_tokens.*, block.{i}.layer.0.SelfAttention.*, - block.{i}.layer.1.DenseReluDense.*, final_layer_norm.*. - """ - - def __init__(self, config: FastAnkhConfig, **kwargs): - AnkhPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - - resolved = resolve_attention_backend(config.attn_backend) - if resolved == AttentionBackend.KERNELS_FLASH: - print("ANKH: kernels_flash not supported (relative position bias); falling back to flex/sdpa") - resolved = AttentionBackend.FLEX if flex_attention is not None else AttentionBackend.SDPA - self.attention_backend = resolved - - self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model) - self.block = nn.ModuleList([ - AnkhBlock(config, has_relative_attention_bias=(i == 0)) - for i in range(config.num_layers) - ]) - for blk in self.block: - blk.layer[0].SelfAttention.attn_backend = self.attention_backend - - self.final_layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) - self.gradient_checkpointing = False - self.tokenizer = _load_ankh_tokenizer(config) - self.post_init() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, value): - self.embed_tokens = value - - @torch.compiler.disable - def _compute_materialized_bias(self, seq_len: int, device: torch.device) -> torch.Tensor: - """Precompute full (Q, K, H) bias tensor for flex score_mod lookup.""" - bias_embedding = self.block[0].layer[0].SelfAttention.relative_attention_bias - context_position = torch.arange(seq_len, dtype=torch.long, device=device)[:, None] - memory_position = torch.arange(seq_len, dtype=torch.long, device=device)[None, :] - relative_position = memory_position - context_position - buckets = AnkhSelfAttention._relative_position_bucket( - relative_position, - num_buckets=self.config.relative_attention_num_buckets, - max_distance=self.config.relative_attention_max_distance, - ) - return bias_embedding(buckets) # (Q, K, H) - - def _build_flex_score_mod(self, seq_len: int, device: torch.device): - """Build score_mod closure that reads from materialized bias tensor.""" - bias = self._compute_materialized_bias(seq_len, device) - - def score_mod(score, b, h, q_idx, kv_idx): - return score + bias[q_idx, kv_idx, h] - - return score_mod - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - hidden_states = self.embed_tokens(input_ids) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - encoder_output = self._run_encoder( - hidden_states, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - ) - return select_hidden_state_embeddings( - encoder_output.last_hidden_state, - encoder_output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def _run_encoder( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: bool = False, - output_attentions: bool = False, - ) -> AnkhEncoderOutput: - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - batch_size, seq_len = hidden_states.shape[:2] - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=batch_size, - seq_len=seq_len, - device=hidden_states.device, - attention_mask=attention_mask, - ) - - flex_score_mod = None - position_bias = None - if self.attention_backend == AttentionBackend.FLEX: - flex_score_mod = self._build_flex_score_mod(seq_len, hidden_states.device) - - for layer_module in self.block: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - hidden_states, attn_weights, position_bias = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask_2d, - attention_mask_4d, - flex_block_mask, - position_bias, - flex_score_mod, - output_attentions, - ) - else: - hidden_states, attn_weights, position_bias = layer_module( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - position_bias=position_bias, - flex_score_mod=flex_score_mod, - output_attentions=output_attentions, - ) - - if all_attentions is not None: - all_attentions = all_attentions + (attn_weights,) - - hidden_states = self.final_layer_norm(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - return AnkhEncoderOutput( - last_hidden_state=hidden_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - **kwargs, - ) -> AnkhEncoderOutput: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - hidden_states = self.embed_tokens(input_ids) - elif inputs_embeds is not None: - hidden_states = inputs_embeds - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - return self._run_encoder( - hidden_states, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states or False, - output_attentions=output_attentions or False, - ) - - -# --------------------------------------------------------------------------- -# Model classes -# --------------------------------------------------------------------------- - -class FastAnkhModel(AnkhPreTrainedModel, EmbeddingMixin): - """ANKH encoder model for embedding extraction.""" - - def __init__(self, config: FastAnkhConfig, **kwargs): - AnkhPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.shared = nn.Embedding(config.vocab_size, config.d_model) - self.encoder = FAST_ANKH_ENCODER(config) - self.post_init() - - @property - def tokenizer(self): - return self.encoder.tokenizer - - def get_input_embeddings(self): - return self.encoder.embed_tokens - - def set_input_embeddings(self, value): - self.encoder.embed_tokens = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.encoder._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - **kwargs, - ) -> AnkhEncoderOutput: - return self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - ) - - -class FastAnkhForMaskedLM(FastPLMTestTimeTrainingMixin, AnkhPreTrainedModel, EmbeddingMixin): - """ANKH encoder with LM head for masked language modeling. - - NOTE: The LM head is initialized from the shared embedding weights but is NOT - tied. The original ANKH models were trained with T5's span corruption objective - using an encoder-decoder architecture. This encoder-only MaskedLM variant is - not pre-trained for standard MLM and requires additional fine-tuning. - """ - - def __init__(self, config: FastAnkhConfig, **kwargs): - AnkhPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.shared = nn.Embedding(config.vocab_size, config.d_model) - self.encoder = FAST_ANKH_ENCODER(config) - self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - self.init_ttt({"lora_target_replace_module": "AnkhSelfAttention"}) - - @property - def tokenizer(self): - return self.encoder.tokenizer - - def get_input_embeddings(self): - return self.encoder.embed_tokens - - def set_input_embeddings(self, value): - self.encoder.embed_tokens = value - - def get_output_embeddings(self): - return self.lm_head - - def set_output_embeddings(self, new_embeddings): - self.lm_head = new_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.encoder._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.encoder] - - def _ttt_tokenize( - self, - seq: str | list[str] | None = None, - input_ids: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - del kwargs - if input_ids is not None: - return input_ids - assert seq is not None, "Pass either seq or input_ids for ANKH TTT." - sequences = [seq] if isinstance(seq, str) else seq - spaced_sequences = [" ".join(sequence) for sequence in sequences] - tokenized = self.tokenizer(spaced_sequences, return_tensors="pt", padding=True) - return tokenized["input_ids"] - - def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: - amino_acids = "ACDEFGHIKLMNPQRSTVWY" - ids = [self.tokenizer.convert_tokens_to_ids(aa) for aa in amino_acids] - return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - **kwargs, - ) -> AnkhMaskedLMOutput: - outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - ) - sequence_output = outputs.last_hidden_state - logits = self.lm_head(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) - - return AnkhMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - -class FastAnkhForSequenceClassification(AnkhPreTrainedModel, EmbeddingMixin): - def __init__(self, config: FastAnkhConfig, **kwargs): - AnkhPreTrainedModel.__init__(self, config, **kwargs) - self.num_labels = config.num_labels - self.config = config - self.shared = nn.Embedding(config.vocab_size, config.d_model) - self.encoder = FAST_ANKH_ENCODER(config) - self.classifier = nn.Linear(config.d_model, config.num_labels) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - self.post_init() - - @property - def tokenizer(self): - return self.encoder.tokenizer - - def get_input_embeddings(self): - return self.encoder.embed_tokens - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.encoder._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - **kwargs, - ) -> AnkhMaskedLMOutput: - outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - ) - # Pool: mean over non-padding tokens - sequence_output = outputs.last_hidden_state - if attention_mask is not None: - mask = attention_mask.unsqueeze(-1).to(sequence_output.dtype) - pooled = (sequence_output * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1) - else: - pooled = sequence_output.mean(dim=1) - logits = self.classifier(pooled) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss = self.mse(logits.squeeze(), labels.squeeze()) if self.num_labels == 1 else self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return AnkhMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - -class FastAnkhForTokenClassification(AnkhPreTrainedModel, EmbeddingMixin): - def __init__(self, config: FastAnkhConfig, **kwargs): - AnkhPreTrainedModel.__init__(self, config, **kwargs) - self.num_labels = config.num_labels - self.shared = nn.Embedding(config.vocab_size, config.d_model) - self.encoder = FAST_ANKH_ENCODER(config) - self.classifier = nn.Linear(config.d_model, config.num_labels) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - - @property - def tokenizer(self): - return self.encoder.tokenizer - - def get_input_embeddings(self): - return self.encoder.embed_tokens - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.encoder._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - **kwargs, - ) -> AnkhMaskedLMOutput: - outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - ) - sequence_output = outputs.last_hidden_state - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return AnkhMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) diff --git a/fastplms/attention.py b/fastplms/attention.py deleted file mode 100644 index 92c59e6..0000000 --- a/fastplms/attention.py +++ /dev/null @@ -1,377 +0,0 @@ -"""Shared attention infrastructure for all FastPLMs models. - -Contains: AttentionBackend enum, backend resolution, mask creation, -flex attention helpers, flash kernel detection/dispatch, and pad/unpad utilities. -""" -from __future__ import annotations - -from enum import Enum -from typing import Dict, List, Optional, Tuple - -import torch -import torch.nn as nn -from torch.nn import functional as F -from einops import rearrange - -try: - from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask -except ImportError: - create_block_mask = None - flex_attention = None - BlockMask = None - -_compiled_flex_attention = None - - -def _get_flex_attention_fn(): - """Return flex_attention callable: compiled (fused kernel) by default, or eager when debug flag is set.""" - global _compiled_flex_attention - if flex_attention is None: - return None - flex_mod = torch.nn.attention.flex_attention - if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False): - return flex_attention - if _compiled_flex_attention is None: - _compiled_flex_attention = torch.compile( - flex_attention, - dynamic=False, - ) - return _compiled_flex_attention - - -# HuggingFace `kernels` exposes slightly different APIs for Flash Attention 2 -# and 3. Detect the loaded variant once so every caller uses the same dispatch. -def _infer_kernels_flash_variant(kernel) -> Optional[str]: - if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"): - return "flash_attn2" - if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"): - return "flash_attn3" - return None - - -def _try_get_kernels_flash(): - try: - from kernels import get_kernel - except ImportError: - return None, None - - flash_kernel = None - flash_kernel_variant = None - try: - flash_kernel = get_kernel("kernels-community/flash-attn3") - flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel) - assert flash_kernel_variant is not None, "Loaded flash-attn3 kernel does not expose a supported API." - except Exception: - try: - flash_kernel = get_kernel("kernels-community/flash-attn2") - flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel) - assert flash_kernel_variant is not None, "Loaded flash-attn2 kernel does not expose a supported API." - except Exception: - flash_kernel = None - flash_kernel_variant = None - return flash_kernel, flash_kernel_variant - - -_FLASH_KERNELS_LOADED = False -FLASH_KERNEL = None -FLASH_KERNEL_VARIANT = None - - -def _ensure_flash_kernels_loaded(): - global _FLASH_KERNELS_LOADED, FLASH_KERNEL, FLASH_KERNEL_VARIANT - if _FLASH_KERNELS_LOADED: - return - _FLASH_KERNELS_LOADED = True - FLASH_KERNEL, FLASH_KERNEL_VARIANT = _try_get_kernels_flash() - - -def _kernels_flash_forward( - query_states: torch.Tensor, - key_states: torch.Tensor, - value_states: torch.Tensor, - causal: bool = False, - softmax_scale: Optional[float] = None, -) -> torch.Tensor: - """Flash-attention forward, optionally overriding the softmax scale. - - When `softmax_scale is None`, the flash kernel applies its default - `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already - pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold). - Failing to override when Q is pre-scaled applies the scale twice. On - DPLM-150M, that produced pooled-embedding cosine around -0.12 and argmax - agreement around 0.27 vs SDPA. - """ - assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." - if FLASH_KERNEL_VARIANT == "flash_attn2": - return FLASH_KERNEL.fwd( - q=query_states, k=key_states, v=value_states, - softmax_scale=softmax_scale, is_causal=causal, - )[0] - if FLASH_KERNEL_VARIANT == "flash_attn3": - try: - output = FLASH_KERNEL.flash_attn_func( - q=query_states, k=key_states, v=value_states, - softmax_scale=softmax_scale, causal=causal, - ) - except TypeError: - output = FLASH_KERNEL.flash_attn_func( - query_states, key_states, value_states, - 0.0, softmax_scale, causal, - ) - if isinstance(output, tuple): - return output[0] - return output - raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}") - - -def _kernels_flash_varlen_forward( - query_states: torch.Tensor, - key_states: torch.Tensor, - value_states: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_in_batch_q: int, - max_seqlen_in_batch_k: int, - causal: bool = False, - softmax_scale: Optional[float] = None, -) -> torch.Tensor: - """Varlen flash-attention forward, optionally overriding the softmax scale. - - See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be - passed when Q has been pre-scaled by the caller. - """ - assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." - if FLASH_KERNEL_VARIANT == "flash_attn2": - return FLASH_KERNEL.varlen_fwd( - q=query_states, k=key_states, v=value_states, - cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k, - softmax_scale=softmax_scale, is_causal=causal, - )[0] - if FLASH_KERNEL_VARIANT == "flash_attn3": - try: - output = FLASH_KERNEL.flash_attn_varlen_func( - q=query_states, k=key_states, v=value_states, - cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k, - softmax_scale=softmax_scale, causal=causal, - ) - except TypeError: - output = FLASH_KERNEL.flash_attn_varlen_func( - query_states, key_states, value_states, - cu_seqlens_q, cu_seqlens_k, - max_seqlen_in_batch_q, max_seqlen_in_batch_k, - 0.0, softmax_scale, causal, - ) - if isinstance(output, tuple): - return output[0] - return output - raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}") - - -# Varlen flash attention runs only on real tokens. These helpers remove padding -# before the kernel call and restore the original padded batch shape afterward. -class IndexFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, input, indices) -> torch.Tensor: - ctx.save_for_backward(indices) - assert input.ndim >= 2 - ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] - second_dim = other_shape.numel() - return torch.gather( - rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim) - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, grad_output) -> Tuple[torch.Tensor, None]: - (indices,) = ctx.saved_tensors - assert grad_output.ndim >= 2 - other_shape = grad_output.shape[1:] - grad_output = rearrange(grad_output, "b ... -> b (...)") - grad_input = torch.zeros( - [ctx.first_axis_dim, grad_output.shape[1]], device=grad_output.device, dtype=grad_output.dtype - ) - grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output) - return grad_input.reshape(ctx.first_axis_dim, *other_shape), None - - -class IndexPutFirstAxis(torch.autograd.Function): - @staticmethod - def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor: - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert values.ndim >= 2 - output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype) - output[indices] = values - return output - - @staticmethod - def backward(ctx, grad_output) -> Tuple[torch.Tensor, None, None]: - (indices,) = ctx.saved_tensors - return grad_output[indices], None, None - - -index_first_axis = IndexFirstAxis.apply -index_put_first_axis = IndexPutFirstAxis.apply - - -def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor: - output = index_put_first_axis(hidden_states, indices, batch * seqlen) - return rearrange(output, "(b s) ... -> b s ...", b=batch) - - -def _unpad_input( - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - attention_mask_2d: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]: - batch_size, seq_len, num_heads, head_dim = query_layer.shape - seqlens = attention_mask_2d.sum(dim=1).int() - cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0)) - max_seqlen = int(seqlens.max().item()) - indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten() - query_layer = index_first_axis(query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) - key_layer = index_first_axis(key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) - value_layer = index_first_axis(value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) - return query_layer, key_layer, value_layer, indices, (cu_seqlens, cu_seqlens), (max_seqlen, max_seqlen) - - -def kernels_flash_attention_func( - query_states: torch.Tensor, - key_states: torch.Tensor, - value_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, -) -> torch.Tensor: - """Public flash-attention entry point with optional padding handling. - - `softmax_scale`: - None -> kernel applies its default `1 / sqrt(head_dim)`. - float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled - by the caller). - - Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)` - before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass - `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale - again, yielding an effective `1/head_dim` scale that drifts across layers. - """ - assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." - if not causal and attention_mask_2d is not None: - batch_size, q_len = query_states.shape[:2] - ( - query_states, key_states, value_states, - indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k), - ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d) - attn_output_unpad = _kernels_flash_varlen_forward( - query_states=query_states, key_states=key_states, value_states=value_states, - cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - max_seqlen_in_batch_q=max_seqlen_q, max_seqlen_in_batch_k=max_seqlen_k, - softmax_scale=softmax_scale, - ) - return pad_input(attn_output_unpad, indices_q, batch_size, q_len) - else: - return _kernels_flash_forward( - query_states=query_states, key_states=key_states, value_states=value_states, - causal=causal, softmax_scale=softmax_scale, - ) - - -# User-facing backend strings resolve to this enum before attention dispatch. -class AttentionBackend(Enum): - AUTO = "auto" - KERNELS_FLASH = "kernels_flash" - FLEX = "flex" - SDPA = "sdpa" - - -VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend) - - -_BACKEND_CONFIRMED = False - - -def resolve_attention_backend(requested_backend: str) -> AttentionBackend: - global _BACKEND_CONFIRMED - assert requested_backend in VALID_ATTENTION_BACKENDS, ( - f"Unsupported attention backend: {requested_backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - ) - if requested_backend in (AttentionBackend.AUTO.value, AttentionBackend.KERNELS_FLASH.value): - _ensure_flash_kernels_loaded() - if requested_backend == AttentionBackend.AUTO.value: - if FLASH_KERNEL is not None: - resolved = AttentionBackend.KERNELS_FLASH - elif flex_attention is not None: - resolved = AttentionBackend.FLEX - else: - resolved = AttentionBackend.SDPA - elif requested_backend == AttentionBackend.KERNELS_FLASH.value: - assert FLASH_KERNEL is not None, "Kernels Flash Attention is not available in this environment." - resolved = AttentionBackend.KERNELS_FLASH - elif requested_backend == AttentionBackend.FLEX.value: - assert flex_attention is not None, "Flex Attention is not available in this environment." - resolved = AttentionBackend.FLEX - elif requested_backend == AttentionBackend.SDPA.value: - resolved = AttentionBackend.SDPA - else: - raise AssertionError(f"Unsupported attention backend: {requested_backend}") - if not _BACKEND_CONFIRMED: - print(f"Attention backend: config='{requested_backend}' -> resolved='{resolved.value}'") - _BACKEND_CONFIRMED = True - return resolved - - -@torch.compiler.disable -def get_attention_mask( - effective_backend: AttentionBackend, - batch_size: int, - seq_len: int, - device: torch.device, - attention_mask: Optional[torch.Tensor] = None, -) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]: - """Build padding masks once for all encoder layers. - - Returns (attention_mask_2d, attention_mask_4d, flex_block_mask). - """ - if attention_mask is None: - return None, None, None - - attention_mask_2d = attention_mask.bool() - - if effective_backend == AttentionBackend.KERNELS_FLASH: - return attention_mask_2d, None, None - - if effective_backend == AttentionBackend.FLEX: - assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable." - valid_lens = attention_mask_2d.sum(dim=-1) - - def mask_mod(batch_idx, head_idx, q_idx, kv_idx): - return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx]) - - flex_block_mask = create_block_mask(mask_mod, batch_size, 1, seq_len, seq_len, device=device) - return attention_mask_2d, None, flex_block_mask - - # SDPA/manual masks only keys. Padding queries still attend to real keys, so - # their outputs stay finite instead of softmaxing over all -inf scores. - attention_mask_4d = attention_mask_2d[:, None, None, :] - return attention_mask_2d, attention_mask_4d, None - - -def bool_to_additive_mask( - bool_mask: torch.Tensor, - dtype: torch.dtype, -) -> torch.Tensor: - """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid). - - Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))` - directly on a bool tensor returns a bool tensor because `-inf` casts to `True`. - That silently drops the mask. Always allocate a float tensor first, then fill it. - This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask. - """ - assert bool_mask.dtype == torch.bool, ( - f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}" - ) - additive = torch.zeros_like(bool_mask, dtype=dtype) - additive.masked_fill_(bool_mask.logical_not(), float("-inf")) - return additive diff --git a/fastplms/boltz/README.md b/fastplms/boltz/README.md deleted file mode 100644 index 59784bb..0000000 --- a/fastplms/boltz/README.md +++ /dev/null @@ -1,164 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements can be found [here](https://github.com/Synthyra/FastPLMs.git). - -# Boltz2 AutoModel (Inference-only) -This is a barebones Huggingface `AutoModel` compatible implementation of Boltz2 focused on fast inference workflows. - -The implementation is located in `fastplms/boltz/` and exposes: -- `Boltz2Config` -- `Boltz2Model` -- `predict_structure(amino_acid_sequence, ...)` -- `save_as_cif(structure_output, output_path, ...)` - -## Design goals -- Inference-only (no training hooks, no Lightning trainer usage). -- Lightweight runtime around `torch` + `transformers` (plus `numpy`). -- AutoModel remote-code compatibility via `trust_remote_code=True`. -- Confidence outputs included in prediction outputs (`plddt`, `ptm`, `iptm`, and derived confidence score when available). - -## Runtime note -This implementation is self-contained inside `fastplms/boltz/` and does not require -the original cloned `boltz` package at runtime. - -## Use with transformers - -### Load from an exported directory -```python -import torch -from transformers import AutoModel - -model = AutoModel.from_pretrained( - "Synthyra/Boltz2", - trust_remote_code=True, - dtype=torch.float32, -).eval() -``` - -### Predict structure from sequence -```python -out = model.predict_structure( - amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG", - recycling_steps=3, - num_sampling_steps=200, - diffusion_samples=1, -) - -print(out.sample_atom_coords.shape) -print(None if out.plddt is None else out.plddt.shape) -``` - -### Save CIF -```python -model.save_as_cif(out, "prediction.cif") -``` - -## Convert Boltz checkpoint to HF export -Use: - -```bash -py -m fastplms.boltz.get_weights --checkpoint_path fastplms/boltz/weights/boltz2_conf.ckpt --output_dir boltz2_automodel_export -``` - -The export directory contains: -- `config.json` -- `pytorch_model.bin` -- `modeling_boltz2.py` -- `minimal_featurizer.py` -- `minimal_structures.py` -- `cif_writer.py` -- `vb_*.py` (self-contained vendored Boltz2 inference modules/constants) - -## Output object fields -`predict_structure(...)` returns `Boltz2StructureOutput` with: -- `sample_atom_coords` -- `atom_pad_mask` -- `plddt` -- `complex_plddt` -- `ptm` -- `iptm` -- `confidence_score` (derived when available) -- `raw_output` - -## Limitations -- Current featurization path is protein-only and minimal. -- This implementation is meant for practical inference and export workflows, not full Boltz training parity. -- Test-time training is not supported for Boltz2 in FastPLMs. TTT is currently limited to sequence PLMs plus ESMFold and ESMFold2 PLM backbones. - -## Docker-first compliance testing - -Build the container at repo root: - -```bash -docker build -t fastplms-test -f Dockerfile . -``` - -Launch a test shell: - -```bash -docker run --rm --gpus all -it -v ${PWD}:/workspace fastplms-test bash -``` - -Inside the container, run Boltz2 compliance against pip `boltz`: - -```bash -python -m testing.run_boltz2_compliance --device cuda --dtype float32 --seed 42 --num-sequences 3 --recycling-steps 3 --num-sampling-steps 200 --diffusion-samples 1 --pass-coord-metric aligned --write-cif-artifacts -``` - -Artifacts are written to `testing/results//boltz2_compliance/` by default: -- `metrics.json` -- `metrics.csv` -- `summary.txt` -- `structures/seq_/ours_seq.cif` -- `structures/seq_/ref_seq.cif` - -Coordinate metrics now include both raw and rigid-aligned variants: -- `coord_mae`, `coord_rmse`, `coord_max_abs` (raw frame-dependent deltas) -- `coord_mae_aligned`, `coord_rmse_aligned`, `coord_max_abs_aligned` (Kabsch aligned) -- `pairwise_dist_mae` (frame-invariant pairwise-distance delta) - -Pass/fail uses `--pass-coord-metric aligned` by default. Set `--pass-coord-metric raw` to use the raw coordinate thresholds. - -## Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{passaro2025boltz2, - title={Boltz-2: Exploring the Frontiers of Biomolecular Prediction}, - author={Passaro, Saro and Corso, Gabriele and Wohlwend, Jeremy and Reveiz, Mateo and Bordes, Florian and Wicky, Basile and Dayan, Peter and Jing, Bowen}, - journal={bioRxiv}, - year={2025} -} -``` - -```bibtex -@article{wohlwend2024boltz1, - title={Boltz-1: Democratizing Biomolecular Interaction Modeling}, - author={Wohlwend, Jeremy and Corso, Gabriele and Passaro, Saro and Reveiz, Mateo and Leidal, Ken and Swanson, Wojtek and Kher, Gilmer and Lember, Tommi and Jaakkola, Tommi}, - journal={bioRxiv}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/boltz/cif_writer.py b/fastplms/boltz/cif_writer.py deleted file mode 100644 index 4ac571d..0000000 --- a/fastplms/boltz/cif_writer.py +++ /dev/null @@ -1,126 +0,0 @@ -from pathlib import Path -from typing import List, Optional - -import numpy as np -import torch - -from .minimal_structures import ProteinStructureTemplate - - -def _confidence_per_atom( - plddt: Optional[torch.Tensor], - atom_to_residue: List[int], - num_atoms: int, - sample_index: int, -) -> np.ndarray: - if plddt is None: - return np.ones((num_atoms,), dtype=np.float32) * 100.0 - - values = plddt.detach().cpu() - if values.ndim == 1: - values = values.unsqueeze(0) - assert values.ndim == 2, "Expected pLDDT with shape [samples, tokens/atoms]." - assert sample_index < values.shape[0], "sample_index out of range for pLDDT." - - selected = values[sample_index] - if selected.shape[0] == num_atoms: - return (selected.numpy() * 100.0).astype(np.float32) - - num_residues = max(atom_to_residue) + 1 - if selected.shape[0] == num_residues: - expanded = np.zeros((num_atoms,), dtype=np.float32) - selected_np = selected.numpy() - for atom_idx, residue_idx in enumerate(atom_to_residue): - expanded[atom_idx] = selected_np[residue_idx] * 100.0 - return expanded - - return np.ones((num_atoms,), dtype=np.float32) * 100.0 - - -def write_cif( - structure_template: ProteinStructureTemplate, - atom_coords: torch.Tensor, - atom_mask: torch.Tensor, - output_path: str, - plddt: Optional[torch.Tensor] = None, - sample_index: int = 0, -) -> str: - coords = atom_coords.detach().cpu() - if coords.ndim == 2: - coords = coords.unsqueeze(0) - assert coords.ndim == 3, "Expected coordinates with shape [samples, atoms, 3]." - assert sample_index < coords.shape[0], "sample_index out of range." - selected_coords_tensor = coords[sample_index] - all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor)) - assert not torch.any(all_non_finite), ( - "CIF export received non-finite coordinates. " - f"Non-finite count: {int(all_non_finite.sum().item())}" - ) - selected_coords = selected_coords_tensor.numpy() - - mask = atom_mask.detach().cpu() - if mask.ndim == 2: - mask = mask[0] - assert mask.ndim == 1, "Expected atom mask with shape [atoms]." - assert mask.shape[0] == selected_coords.shape[0], "Atom mask/coord size mismatch." - assert torch.any(mask > 0), "Atom mask has no valid atoms for CIF export." - valid_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor[mask > 0])) - assert not torch.any(valid_non_finite), ( - "CIF export has non-finite coordinates in unmasked atoms. " - f"Non-finite count: {int(valid_non_finite.sum().item())}" - ) - - b_iso = _confidence_per_atom( - plddt=plddt, - atom_to_residue=structure_template.atom_residue_index, - num_atoms=structure_template.num_atoms, - sample_index=sample_index, - ) - assert b_iso.shape[0] == structure_template.num_atoms - - lines = [ - "data_boltz2_prediction", - "#", - "loop_", - "_atom_site.group_PDB", - "_atom_site.id", - "_atom_site.type_symbol", - "_atom_site.label_atom_id", - "_atom_site.label_comp_id", - "_atom_site.label_asym_id", - "_atom_site.label_seq_id", - "_atom_site.Cartn_x", - "_atom_site.Cartn_y", - "_atom_site.Cartn_z", - "_atom_site.occupancy", - "_atom_site.B_iso_or_equiv", - "_atom_site.pdbx_PDB_model_num", - ] - - atom_id = 1 - for idx in range(structure_template.num_atoms): - if mask[idx] <= 0: - continue - - residue_idx = structure_template.atom_residue_index[idx] - residue_name = structure_template.residue_names[residue_idx] - atom_name = structure_template.atom_names[idx] - element = structure_template.atom_elements[idx] - chain_id = structure_template.atom_chain_id[idx] - x_val, y_val, z_val = selected_coords[idx].tolist() - b_factor = float(b_iso[idx]) - - line = ( - f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} " - f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1" - ) - lines.append(line) - atom_id += 1 - - lines.append("#") - text = "\n".join(lines) + "\n" - - out_path = Path(output_path) - out_path.parent.mkdir(parents=True, exist_ok=True) - out_path.write_text(text, encoding="utf-8") - return str(out_path) diff --git a/fastplms/boltz/get_weights.py b/fastplms/boltz/get_weights.py deleted file mode 100644 index b970661..0000000 --- a/fastplms/boltz/get_weights.py +++ /dev/null @@ -1,264 +0,0 @@ -import argparse -import copy -import shutil -import sys -import urllib.request -from pathlib import Path -from typing import Dict, List, Set - -import torch -from huggingface_hub import HfApi, login -from transformers import AutoConfig - -from fastplms.boltz.modeling_boltz2 import ( - Boltz2Model, - _filtered_kwargs, - _state_dict_without_wrappers, - _to_plain_python, -) -from fastplms.weight_parity_utils import assert_state_dict_equal, assert_model_parameters_fp32 - - -BOLTZ2_CKPT_URL = "https://huggingface.co/boltz-community/boltz-2/resolve/main/boltz2_conf.ckpt" - - -def _download_checkpoint_if_needed(checkpoint_path: Path) -> Path: - checkpoint_path.parent.mkdir(parents=True, exist_ok=True) - if not checkpoint_path.exists(): - urllib.request.urlretrieve(BOLTZ2_CKPT_URL, str(checkpoint_path)) # noqa: S310 - return checkpoint_path - - -def _copy_runtime_package(output_dir: Path) -> None: - source_pkg = Path(__file__).resolve().parent - project_root = source_pkg.parent - runtime_files = [ - "__init__.py", - "modeling_boltz2.py", - "minimal_featurizer.py", - "minimal_structures.py", - "cif_writer.py", - ] - for filename in runtime_files: - shutil.copyfile(source_pkg / filename, output_dir / filename) - shutil.copyfile(project_root / "entrypoint_setup.py", output_dir / "entrypoint_setup.py") - for flat_module in source_pkg.glob("vb_*.py"): - shutil.copyfile(flat_module, output_dir / flat_module.name) - - -def _ensure_local_boltz_module_on_path() -> Path: - script_root = Path(__file__).resolve().parents[1] - candidates = [script_root / "boltz" / "src"] - - cwd = Path.cwd().resolve() - for parent in [cwd, *cwd.parents]: - candidates.append(parent / "boltz" / "src") - - deduplicated_candidates: List[Path] = [] - seen: Set[str] = set() - for candidate in candidates: - candidate_resolved = candidate.resolve() - candidate_key = str(candidate_resolved) - if candidate_key not in seen: - seen.add(candidate_key) - deduplicated_candidates.append(candidate_resolved) - - for candidate in deduplicated_candidates: - package_marker = candidate / "boltz" / "__init__.py" - if package_marker.exists(): - candidate_str = str(candidate) - if candidate_str not in sys.path: - sys.path.insert(0, candidate_str) - return candidate - - raise FileNotFoundError( - "Unable to locate local boltz submodule. " - f"Checked: {', '.join([str(path) for path in deduplicated_candidates])}" - ) - - -def _load_official_boltz2_model( - checkpoint_path: Path, - use_kernels: bool, -) -> torch.nn.Module: - _ensure_local_boltz_module_on_path() - from boltz.model.models.boltz2 import Boltz2 as OfficialBoltz2 - from boltz.model.modules.diffusionv2 import AtomDiffusion - - checkpoint = torch.load( - str(checkpoint_path), - map_location="cpu", - weights_only=False, - ) - assert isinstance(checkpoint, dict), "Checkpoint must deserialize to a dictionary." - assert "hyper_parameters" in checkpoint, "Checkpoint missing 'hyper_parameters'." - assert "state_dict" in checkpoint, "Checkpoint missing 'state_dict'." - hyper_parameters = checkpoint["hyper_parameters"] - state_dict = checkpoint["state_dict"] - assert isinstance(hyper_parameters, dict), "Checkpoint hyper_parameters must be a dictionary." - assert isinstance(state_dict, dict), "Checkpoint state_dict must be a dictionary." - - init_kwargs = _filtered_kwargs( - target=OfficialBoltz2, - kwargs=_to_plain_python(copy.deepcopy(hyper_parameters)), - ) - if "use_kernels" in init_kwargs: - init_kwargs["use_kernels"] = use_kernels - assert "pairformer_args" in init_kwargs, ( - "Checkpoint hyperparameters missing pairformer_args for official Boltz2." - ) - raw_pairformer_args = init_kwargs["pairformer_args"] - assert isinstance(raw_pairformer_args, dict), "Expected pairformer_args to be a dictionary." - pairformer_args = _to_plain_python(copy.deepcopy(raw_pairformer_args)) - assert isinstance(pairformer_args, dict), "Expected normalized pairformer_args to be a dictionary." - pairformer_args["v2"] = True - init_kwargs["pairformer_args"] = pairformer_args - assert "diffusion_process_args" in init_kwargs, ( - "Checkpoint hyperparameters missing diffusion_process_args for official Boltz2." - ) - raw_diffusion_process_args = init_kwargs["diffusion_process_args"] - assert isinstance(raw_diffusion_process_args, dict), ( - "Expected diffusion_process_args to be a dictionary." - ) - filtered_diffusion_process_args = _filtered_kwargs( - target=AtomDiffusion, - kwargs=raw_diffusion_process_args, - ) - sanitized_diffusion_process_args: Dict[str, object] = {} - for key in filtered_diffusion_process_args: - if key == "score_model_args": - continue - sanitized_diffusion_process_args[key] = filtered_diffusion_process_args[key] - init_kwargs["diffusion_process_args"] = sanitized_diffusion_process_args - official_model = OfficialBoltz2(**init_kwargs) - - cleaned_state_dict = _state_dict_without_wrappers(state_dict) - target_keys = set(official_model.state_dict().keys()) - filtered_state_dict: Dict[str, torch.Tensor] = {} - for key in cleaned_state_dict: - if key in target_keys: - filtered_state_dict[key] = cleaned_state_dict[key] - missing_keys = sorted(target_keys.difference(filtered_state_dict.keys())) - assert len(missing_keys) == 0, ( - "Official Boltz2 model is missing required checkpoint keys. " - f"Missing keys (first 20): {missing_keys[:20]}" - ) - load_result = official_model.load_state_dict(filtered_state_dict, strict=False) - assert len(load_result.missing_keys) == 0, ( - "Missing keys while loading official Boltz2 checkpoint. " - f"Missing keys (first 20): {load_result.missing_keys[:20]}" - ) - assert len(load_result.unexpected_keys) == 0, ( - "Unexpected keys while loading official Boltz2 checkpoint. " - f"Unexpected keys (first 20): {load_result.unexpected_keys[:20]}" - ) - - official_model = official_model.eval().cpu().to(torch.float32) - assert_model_parameters_fp32( - model=official_model, - model_name="official Boltz2 model", - ) - return official_model - - -if __name__ == "__main__": - # py -m fastplms.boltz.get_weights - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint_path", type=str, default="fastplms/boltz/weights/boltz2_conf.ckpt") - parser.add_argument("--output_dir", type=str, default="boltz2_automodel_export") - parser.add_argument("--repo_ids", nargs="*", type=str, default=["Synthyra/Boltz2"]) - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--use_kernels", action="store_true") - parser.add_argument("--dry_run", action="store_true") - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - - # Standardization: use the first repo_id from repo_ids - repo_id = args.repo_ids[0] if args.repo_ids else "Synthyra/Boltz2" - - if args.skip_weights: - config = AutoConfig.from_pretrained(repo_id, trust_remote_code=True) - config.auto_map = { - "AutoConfig": "modeling_boltz2.Boltz2Config", - "AutoModel": "modeling_boltz2.Boltz2Model", - } - if args.dry_run: - print(f"[skip-weights][dry-run] validated Boltz2 config for {repo_id}") - raise SystemExit(0) - config.push_to_hub(repo_id) - print(f"[skip-weights] uploaded Boltz2 config to {repo_id}") - raise SystemExit(0) - - checkpoint_path = _download_checkpoint_if_needed(Path(args.checkpoint_path)) - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - official_model = _load_official_boltz2_model( - checkpoint_path=checkpoint_path, - use_kernels=args.use_kernels, - ) - model = Boltz2Model.from_boltz_checkpoint( - checkpoint_path=str(checkpoint_path), - use_kernels=args.use_kernels, - ) - model = model.eval().cpu().to(torch.float32) - assert_model_parameters_fp32( - model=model.core, - model_name="mapped Boltz2 inference core", - ) - official_state_dict = official_model.state_dict() - candidate_state_dict = model.core.state_dict() - official_keys = set(official_state_dict.keys()) - candidate_keys = set(candidate_state_dict.keys()) - missing_official_keys = sorted(candidate_keys - official_keys) - assert len(missing_official_keys) == 0, ( - "Official Boltz2 model is missing inference-core keys required by FastPLMs. " - f"Missing keys (first 20): {missing_official_keys[:20]}" - ) - excluded_official_keys = sorted(official_keys - candidate_keys) - allowed_excluded_prefixes = ( - "template_module.", - "bfactor_module.", - ) - unexpected_excluded_official_keys: List[str] = [] - for key in excluded_official_keys: - is_allowed = False - for prefix in allowed_excluded_prefixes: - if key.startswith(prefix): - is_allowed = True - break - if is_allowed is False: - unexpected_excluded_official_keys.append(key) - assert len(unexpected_excluded_official_keys) == 0, ( - "Unexpected official Boltz2 keys not present in FastPLMs inference core. " - f"Unexpected keys (first 20): {unexpected_excluded_official_keys[:20]}" - ) - filtered_official_state_dict: Dict[str, torch.Tensor] = {} - for key in candidate_state_dict: - filtered_official_state_dict[key] = official_state_dict[key] - assert_state_dict_equal( - reference_state_dict=filtered_official_state_dict, - candidate_state_dict=candidate_state_dict, - context="Boltz2 weight parity", - ) - - model.config.auto_map = { - "AutoConfig": "modeling_boltz2.Boltz2Config", - "AutoModel": "modeling_boltz2.Boltz2Model", - } - if args.dry_run: - print(f"[dry_run] validated Boltz2 parity for checkpoint {checkpoint_path}") - else: - model.save_pretrained(str(output_dir)) - _copy_runtime_package(output_dir=output_dir) - - if args.repo_id is not None and args.dry_run is False: - if args.hf_token is not None: - login(token=args.hf_token) - api = HfApi() - api.create_repo(repo_id=args.repo_id, repo_type="model", exist_ok=True) - api.upload_folder( - folder_path=str(output_dir), - repo_id=args.repo_id, - repo_type="model", - ) diff --git a/fastplms/boltz/minimal_featurizer.py b/fastplms/boltz/minimal_featurizer.py deleted file mode 100644 index cceee37..0000000 --- a/fastplms/boltz/minimal_featurizer.py +++ /dev/null @@ -1,528 +0,0 @@ -import math -from typing import Dict, List, Tuple - -import numpy as np -import torch -from torch.nn.functional import one_hot - -from .minimal_structures import ProteinStructureTemplate -from . import vb_const as const - - -_ELEMENT_TO_Z = { - "H": 1, - "C": 6, - "N": 7, - "O": 8, - "P": 15, - "S": 16, -} - - -def _normalize_sequence(sequence: str) -> str: - seq = sequence.strip().upper() - assert len(seq) > 0, "Amino acid sequence must be non-empty." - for aa in seq: - assert aa in const.prot_letter_to_token, f"Unsupported residue code '{aa}'." - return seq - - -def _atom_name_to_element(atom_name: str) -> str: - name = atom_name.strip().upper() - if len(name) == 0: - return "C" - if name[0].isdigit(): - name = name[1:] - if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"): - return name[0] - return name[0] - - -def _atom_name_to_codes(atom_name: str) -> torch.Tensor: - clipped = atom_name.strip()[:4] - vals = [ord(ch) - 32 for ch in clipped] - while len(vals) < 4: - vals.append(0) - out = torch.tensor(vals, dtype=torch.long) - assert torch.all(out >= 0) and torch.all(out < 64), ( - f"Invalid atom-name encoding for '{atom_name}'." - ) - return out - - -# Canonical RDKit conformer positions extracted from official Boltz2 mol files -# (boltz-community/boltz-2 mols.tar, first conformer, centered per residue). -# These match the geometry the model was trained on. -_RDKIT_CONFORMERS: Dict[str, Dict[str, List[float]]] = { - "ALA": {"N": [-0.944785, 0.952743, 0.876326], "CA": [-0.287002, -0.317661, 0.564429], "C": [1.098243, -0.090637, 0.019913], "O": [1.267603, 0.576468, -1.036548], "CB": [-1.13406, -1.120912, -0.42412]}, - "ARG": {"N": [3.318334, -1.792721, -0.81714], "CA": [2.256519, -0.785164, -0.782524], "C": [2.865221, 0.579152, -0.9479], "O": [2.530447, 1.305604, -1.921779], "CB": [1.446111, -0.888115, 0.531287], "CG": [0.26845, 0.099716, 0.625405], "CD": [-0.837403, -0.192693, -0.39904], "NE": [-2.021791, 0.62151, -0.123045], "CZ": [-2.953151, 0.387273, 0.951242], "NH1": [-4.01513, 1.322306, 1.16201], "NH2": [-2.857608, -0.656868, 1.721485]}, - "ASN": {"N": [-1.76737, -1.671462, 0.274097], "CA": [-0.894725, -0.51088, 0.459045], "C": [-1.437261, 0.658619, -0.316984], "O": [-1.745491, 0.527762, -1.53251], "CB": [0.535815, -0.851978, 0.002785], "CG": [1.496584, 0.259422, 0.3155], "OD1": [2.006304, 0.340489, 1.464867], "ND2": [1.806145, 1.248028, -0.666799]}, - "ASP": {"N": [-0.531828, -1.5551, 0.564551], "CA": [-0.798781, -0.116987, 0.471405], "C": [-2.220416, 0.13533, 0.042978], "O": [-2.637844, -0.292258, -1.067311], "CB": [0.18409, 0.554469, -0.501854], "CG": [1.594766, 0.436104, -0.012856], "OD1": [2.333189, -0.495262, -0.431504], "OD2": [2.076824, 1.333703, 0.934591]}, - "CYS": {"N": [0.008485, 1.680076, -0.119503], "CA": [0.001723, 0.385667, -0.808584], "C": [-1.206733, -0.441379, -0.446529], "O": [-1.665274, -0.434109, 0.728016], "CB": [1.310427, -0.380146, -0.555459], "SG": [1.551372, -0.810109, 1.202058]}, - "GLN": {"N": [-1.932297, -1.099026, -1.887635], "CA": [-1.370158, -0.7752, -0.575112], "C": [-2.304423, 0.155162, 0.147051], "O": [-2.803483, -0.184897, 1.253191], "CB": [0.038652, -0.158687, -0.716379], "CG": [0.736106, 0.008065, 0.640487], "CD": [2.117149, 0.560394, 0.455758], "OE1": [2.309056, 1.803816, 0.522783], "NE2": [3.209399, -0.309626, 0.159856]}, - "GLU": {"N": [-1.750645, -1.298566, -1.148627], "CA": [-1.677927, -0.690629, 0.182405], "C": [-2.320926, 0.670409, 0.183131], "O": [-2.310494, 1.385902, -0.855225], "CB": [-0.218639, -0.616385, 0.677413], "CG": [0.704486, 0.181646, -0.256505], "CD": [2.107661, 0.167064, 0.261748], "OE1": [2.495054, 1.060279, 1.061799], "OE2": [2.97143, -0.85972, -0.106138]}, - "GLY": {"N": [-1.416855, 0.862616, -0.1801], "CA": [-0.6149, -0.033688, 0.644675], "C": [0.809702, 0.000186, 0.194889], "O": [1.222053, -0.829114, -0.659464]}, - "HIS": {"N": [1.26543, -1.579115, 0.588736], "CA": [1.426323, -0.419409, -0.29228], "C": [2.884431, -0.088226, -0.469953], "O": [3.407859, -0.14071, -1.615113], "CB": [0.668818, 0.794698, 0.269345], "CG": [-0.807376, 0.534987, 0.323389], "ND1": [-1.680612, 0.632102, -0.797403], "CD2": [-1.481273, 0.09683, 1.377023], "CE1": [-2.846859, 0.264305, -0.380313], "NE2": [-2.836741, -0.095463, 0.996569]}, - "ILE": {"N": [-0.963969, -1.670572, 0.035929], "CA": [-1.004097, -0.255306, 0.424676], "C": [-1.714092, 0.593857, -0.603418], "O": [-1.755267, 0.246984, -1.815103], "CB": [0.416162, 0.290836, 0.741822], "CG1": [1.380111, 0.189293, -0.469113], "CG2": [0.995443, -0.424464, 1.975358], "CD1": [2.645709, 1.029371, -0.290151]}, - "LEU": {"N": [1.265833, -0.791579, -1.184426], "CA": [0.933538, 0.306792, -0.27234], "C": [2.051763, 0.509151, 0.713934], "O": [2.543725, -0.475272, 1.329405], "CB": [-0.373342, 0.011933, 0.498041], "CG": [-1.632816, -0.15124, -0.388179], "CD1": [-2.835081, -0.534398, 0.48526], "CD2": [-1.953621, 1.124614, -1.181694]}, - "LYS": {"N": [-1.908918, -1.413217, -1.088748], "CA": [-1.724515, -0.558426, 0.087009], "C": [-2.96641, 0.258056, 0.315903], "O": [-3.48391, 0.312278, 1.463711], "CB": [-0.523791, 0.388858, -0.107149], "CG": [0.821393, -0.35366, -0.108072], "CD": [1.993822, 0.633537, -0.161446], "CE": [3.338346, -0.105923, -0.176016], "NZ": [4.453983, 0.838498, -0.225193]}, - "MET": {"N": [-1.522666, -0.831762, 1.959218], "CA": [-0.996081, 0.012809, 0.885437], "C": [-2.091048, 0.292064, -0.106017], "O": [-2.445252, 1.479197, -0.337323], "CB": [0.220238, -0.653951, 0.206767], "CG": [0.915991, 0.285443, -0.785597], "SD": [2.357592, -0.540499, -1.548879], "CE": [3.561227, -0.043301, -0.273607]}, - "PHE": {"N": [3.046798, -1.689795, -0.208668], "CA": [1.829154, -0.89776, -0.391539], "C": [2.202709, 0.507016, -0.775299], "O": [1.770264, 1.003795, -1.849745], "CB": [0.962876, -0.916233, 0.885596], "CG": [-0.373751, -0.255463, 0.660531], "CD1": [-0.576773, 1.029007, 1.007455], "CD2": [-1.46686, -1.011364, -0.000788], "CE1": [-1.882158, 1.672384, 0.749291], "CE2": [-2.646851, -0.423965, -0.236002], "CZ": [-2.865409, 0.982376, 0.159168]}, - "PRO": {"N": [-0.685006, -0.370164, -0.768919], "CA": [0.404323, 0.377887, -0.13805], "C": [1.731651, -0.324892, -0.273548], "O": [1.975993, -1.024105, -1.293633], "CB": [-0.004718, 0.595698, 1.311196], "CG": [-1.517474, 0.663762, 1.260234], "CD": [-1.904769, 0.081814, -0.097281]}, - "SER": {"N": [1.015962, -1.698341, -0.119567], "CA": [0.101187, -0.56256, -0.236622], "C": [0.88759, 0.717914, -0.271742], "O": [0.652684, 1.578473, -1.161893], "CB": [-0.87741, -0.552018, 0.948557], "OG": [-1.780012, 0.516532, 0.841267]}, - "THR": {"N": [-0.05857, 1.577455, 0.452633], "CA": [0.359545, 0.187419, 0.662184], "C": [1.654837, -0.077979, -0.057809], "O": [1.861664, 0.408605, -1.20258], "CB": [-0.727051, -0.803664, 0.176961], "OG1": [-1.139906, -0.486103, -1.128492], "CG2": [-1.950518, -0.805733, 1.097103]}, - "TRP": {"N": [-3.22168, 1.119624, 0.055731], "CA": [-2.641395, 0.081435, -0.801364], "C": [-2.364426, -1.166119, -0.006875], "O": [-2.514206, -2.296758, -0.543038], "CB": [-1.384007, 0.576331, -1.552364], "CG": [-0.24222, 0.940856, -0.643481], "CD1": [-0.016615, 2.14194, -0.107979], "CD2": [0.800451, 0.051232, -0.1373], "NE1": [1.132969, 2.107427, 0.738111], "CE2": [1.575163, 0.761216, 0.659678], "CE3": [1.058378, -1.378818, -0.38353], "CZ2": [2.729299, 0.158745, 1.347069], "CZ3": [2.107441, -1.950409, 0.235009], "CH2": [2.980848, -1.146704, 1.140333]}, - "TYR": {"N": [-2.152606, 0.331726, -1.283488], "CA": [-2.201241, -0.428961, -0.031756], "C": [-3.597879, -0.441616, 0.531374], "O": [-4.172602, 0.638246, 0.83694], "CB": [-1.21174, 0.150443, 0.994648], "CG": [0.209673, 0.053123, 0.503276], "CD1": [0.911845, -1.084226, 0.654794], "CD2": [0.82966, 1.201129, -0.203917], "CE1": [2.294408, -1.180811, 0.141538], "CE2": [2.080826, 1.10993, -0.671525], "CZ": [2.85374, -0.136912, -0.490516], "OH": [4.155917, -0.21207, -0.981367]}, - "VAL": {"N": [0.631715, -1.334948, 0.719228], "CA": [0.419681, -0.467832, -0.445221], "C": [1.457942, 0.623882, -0.483986], "O": [1.935858, 1.0943, 0.583992], "CB": [-1.010701, 0.140743, -0.455789], "CG1": [-2.075373, -0.919582, -0.774314], "CG2": [-1.359123, 0.863438, 0.85609]}, - "UNK": {"N": [0.8287, -1.182096, -0.645721], "CA": [-0.174671, -0.11586, -0.67851], "C": [0.301419, 1.045573, 0.149013], "O": [0.589973, 0.885289, 1.365851], "CB": [-1.545421, -0.632906, -0.190632]}, -} - - -def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarray: - """Get the canonical RDKit conformer position for an atom. - - Uses pre-extracted positions from official Boltz2 mol files. Falls back - to a simple geometric placement for unknown residue/atom combinations. - """ - if res_name in _RDKIT_CONFORMERS and atom_name in _RDKIT_CONFORMERS[res_name]: - return np.array(_RDKIT_CONFORMERS[res_name][atom_name], dtype=np.float32) - # Fallback for unknown atoms (should not happen for canonical AAs) - angle = (atom_idx + 1) * 0.7 - radius = 1.4 + 0.03 * atom_idx - return np.array( - [ - radius * math.cos(angle), - radius * math.sin(angle), - 0.1 * ((atom_idx % 5) - 2), - ], - dtype=np.float32, - ) - - -def _build_template( - sequence: str, -) -> Tuple[ - ProteinStructureTemplate, - List[str], - List[int], - List[int], - List[int], - List[np.ndarray], - List[int], -]: - residue_names: List[str] = [] - residue_token_ids: List[int] = [] - atom_names: List[str] = [] - atom_elements: List[str] = [] - atom_residue_index: List[int] = [] - atom_chain_id: List[str] = [] - atom_positions: List[np.ndarray] = [] - residue_center_atom_idx: List[int] = [] - residue_disto_atom_idx: List[int] = [] - residue_frame_atom_idx: List[int] = [] - - global_atom_idx = 0 - for res_idx, aa in enumerate(sequence): - token_name = const.prot_letter_to_token[aa] - residue_names.append(token_name) - residue_token_ids.append(const.token_ids[token_name]) - - residue_atoms = const.ref_atoms[token_name] - assert len(residue_atoms) > 0, f"No reference atoms for residue {token_name}." - center_atom_name = const.res_to_center_atom[token_name] - disto_atom_name = const.res_to_disto_atom[token_name] - - center_idx = -1 - disto_idx = -1 - n_idx = -1 - ca_idx = -1 - c_idx = -1 - - for local_idx, atom_name in enumerate(residue_atoms): - atom_names.append(atom_name) - element = _atom_name_to_element(atom_name) - atom_elements.append(element) - atom_residue_index.append(res_idx) - atom_chain_id.append("A") - - atom_pos = _get_atom_position(token_name, atom_name, local_idx) - atom_positions.append(atom_pos) - - if atom_name == center_atom_name: - center_idx = global_atom_idx - if atom_name == disto_atom_name: - disto_idx = global_atom_idx - if atom_name == "N": - n_idx = global_atom_idx - if atom_name == "CA": - ca_idx = global_atom_idx - if atom_name == "C": - c_idx = global_atom_idx - global_atom_idx += 1 - - if center_idx == -1: - center_idx = global_atom_idx - len(residue_atoms) - if disto_idx == -1: - disto_idx = center_idx - if n_idx == -1: - n_idx = center_idx - if ca_idx == -1: - ca_idx = center_idx - if c_idx == -1: - c_idx = center_idx - - residue_center_atom_idx.append(center_idx) - residue_disto_atom_idx.append(disto_idx) - residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx]) - - template = ProteinStructureTemplate( - sequence=sequence, - residue_names=residue_names, - atom_names=atom_names, - atom_elements=atom_elements, - atom_residue_index=atom_residue_index, - atom_chain_id=atom_chain_id, - ) - - return ( - template, - residue_names, - residue_token_ids, - residue_center_atom_idx, - residue_disto_atom_idx, - atom_positions, - residue_frame_atom_idx, - ) - - -def _random_rotation_matrix() -> torch.Tensor: - """Sample a uniform random 3x3 rotation matrix (Algorithm 19 from AF2/Boltz).""" - q = torch.randn(4) - q = q / q.norm() - # Quaternion to rotation matrix - w, x, y, z = q[0], q[1], q[2], q[3] - return torch.tensor([ - [1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y)], - [2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x)], - [2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y)], - ], dtype=torch.float32) - - -def _center_and_augment_atoms_per_residue( - atom_positions: torch.Tensor, - atom_residue_index: List[int], - num_residues: int, -) -> torch.Tensor: - """Center atoms per residue and apply random rotation per residue. - - Matches the official Boltz2 featurizer which applies center_random_augmentation - to each residue's ref_pos independently (featurizerv2.py lines 1495-1500). - """ - result = atom_positions.clone() - residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long) - for residue_idx in range(num_residues): - residue_mask = residue_index_tensor == residue_idx - assert torch.any(residue_mask), f"Residue index {residue_idx} has no atoms." - residue_coords = result[residue_mask] - # Center - residue_center = residue_coords.mean(dim=0, keepdim=True) - residue_coords = residue_coords - residue_center - # Random rotation (matching official center_random_augmentation with centering=True) - R = _random_rotation_matrix() - residue_coords = residue_coords @ R.T - result[residue_mask] = residue_coords - return result - - -def build_boltz2_features( - amino_acid_sequence: str, - num_bins: int = 64, - atoms_per_window_queries: int = 32, -) -> Tuple[Dict[str, torch.Tensor], ProteinStructureTemplate]: - sequence = _normalize_sequence(amino_acid_sequence) - ( - template, - residue_names, - residue_token_ids, - residue_center_atom_idx, - residue_disto_atom_idx, - atom_positions_np, - residue_frame_atom_idx_flat, - ) = _build_template(sequence) - - num_tokens = len(residue_names) - num_atoms = len(atom_positions_np) - assert num_tokens > 0 and num_atoms > 0 - - atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32) - atom_positions = _center_and_augment_atoms_per_residue( - atom_positions=atom_positions, - atom_residue_index=template.atom_residue_index, - num_residues=num_tokens, - ) - - token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0) - residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0) - asym_id = torch.zeros((1, num_tokens), dtype=torch.long) - entity_id = torch.zeros((1, num_tokens), dtype=torch.long) - sym_id = torch.zeros((1, num_tokens), dtype=torch.long) - mol_type = torch.full( - (1, num_tokens), - fill_value=const.chain_type_ids["PROTEIN"], - dtype=torch.long, - ) - - res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long) - res_type = one_hot(res_type_ids, num_classes=const.num_tokens).float().unsqueeze(0) - - # token_bonds encodes explicit covalent cross-links from structure bonds, - # NOT backbone peptide bonds (those are implicit via residue_index + asym_id). - # For a standard single-chain protein without cross-links, this is all zeros. - # This matches the official Boltz2 featurizer (featurizerv2.py lines 696-705). - token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32) - type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long) - token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1) - type_bonds = type_bonds.unsqueeze(0) - - token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32) - token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32) - token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32) - - num_contact_classes = len(const.contact_conditioning_info) - unspecified_id = const.contact_conditioning_info["UNSPECIFIED"] - contact_ids = torch.full( - (num_tokens, num_tokens), - fill_value=unspecified_id, - dtype=torch.long, - ) - contact_conditioning = one_hot( - contact_ids, - num_classes=num_contact_classes, - ).float().unsqueeze(0) - contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32) - - assert "x-ray diffraction" in const.method_types_ids - method_feature = torch.full( - (1, num_tokens), - fill_value=const.method_types_ids["x-ray diffraction"], - dtype=torch.long, - ) - modified = torch.zeros((1, num_tokens), dtype=torch.long) - cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32) - affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32) - - ref_pos = atom_positions.unsqueeze(0) - atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32) - atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.float32) - - atom_name_codes = torch.stack( - [_atom_name_to_codes(atom_name) for atom_name in template.atom_names], - dim=0, - ) - ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).float().unsqueeze(0) - - atomic_numbers = [] - for element in template.atom_elements: - if element in _ELEMENT_TO_Z: - z_value = _ELEMENT_TO_Z[element] - else: - z_value = _ELEMENT_TO_Z["C"] - assert z_value < const.num_elements - atomic_numbers.append(z_value) - ref_element = one_hot( - torch.tensor(atomic_numbers, dtype=torch.long), - num_classes=const.num_elements, - ).float().unsqueeze(0) - - ref_charge = torch.zeros((1, num_atoms), dtype=torch.float32) - ref_chirality = torch.zeros((1, num_atoms), dtype=torch.long) - ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0) - - atom_to_token = one_hot( - torch.tensor(template.atom_residue_index, dtype=torch.long), - num_classes=num_tokens, - ).float().unsqueeze(0) - token_to_rep_atom = one_hot( - torch.tensor(residue_disto_atom_idx, dtype=torch.long), - num_classes=num_atoms, - ).float().unsqueeze(0) - token_to_center_atom = one_hot( - torch.tensor(residue_center_atom_idx, dtype=torch.long), - num_classes=num_atoms, - ).float().unsqueeze(0) - r_set_to_rep_atom = token_to_center_atom.clone() - - num_backbone_classes = ( - 1 - + len(const.protein_backbone_atom_index) - + len(const.nucleic_backbone_atom_index) - ) - backbone_ids = [] - for atom_name in template.atom_names: - if atom_name in const.protein_backbone_atom_index: - backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1) - else: - backbone_ids.append(0) - atom_backbone_feat = one_hot( - torch.tensor(backbone_ids, dtype=torch.long), - num_classes=num_backbone_classes, - ).float().unsqueeze(0) - - coords = ref_pos.unsqueeze(1).contiguous() - disto_coords = torch.stack( - [atom_positions[idx] for idx in residue_disto_atom_idx], - dim=0, - ) - disto_coords_ensemble = disto_coords.unsqueeze(0).unsqueeze(0).contiguous() - - bfactor = torch.zeros((1, num_atoms), dtype=torch.float32) - atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32) - - assert atoms_per_window_queries > 0 - pad_atoms = ( - ((num_atoms - 1) // atoms_per_window_queries + 1) * atoms_per_window_queries - - num_atoms - ) - if pad_atoms > 0: - ref_pos = torch.nn.functional.pad(ref_pos, (0, 0, 0, pad_atoms), value=0.0) - atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0) - atom_resolved_mask = torch.nn.functional.pad( - atom_resolved_mask, - (0, pad_atoms), - value=0.0, - ) - ref_atom_name_chars = torch.nn.functional.pad( - ref_atom_name_chars, - (0, 0, 0, 0, 0, pad_atoms), - value=0.0, - ) - ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0) - ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0) - ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0) - atom_backbone_feat = torch.nn.functional.pad( - atom_backbone_feat, - (0, 0, 0, pad_atoms), - value=0.0, - ) - ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0) - coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0) - atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0) - token_to_rep_atom = torch.nn.functional.pad( - token_to_rep_atom, - (0, pad_atoms), - value=0.0, - ) - token_to_center_atom = torch.nn.functional.pad( - token_to_center_atom, - (0, pad_atoms), - value=0.0, - ) - r_set_to_rep_atom = torch.nn.functional.pad( - r_set_to_rep_atom, - (0, pad_atoms), - value=0.0, - ) - bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0) - atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0) - - frames_idx = torch.tensor( - residue_frame_atom_idx_flat, - dtype=torch.long, - ).reshape(num_tokens, 3) - frames_idx = frames_idx.unsqueeze(0).unsqueeze(1) - frame_resolved_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32) - - msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0) - msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32) - deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32) - deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32) - profile = one_hot( - torch.tensor(residue_token_ids, dtype=torch.long), - num_classes=const.num_tokens, - ).float().unsqueeze(0) - - template_restype = one_hot( - torch.zeros((1, 1, num_tokens), dtype=torch.long), - num_classes=const.num_tokens, - ).float() - template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32) - template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32) - template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32) - template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32) - template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long) - visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32) - - disto_target = torch.zeros( - (1, num_tokens, num_tokens, 1, num_bins), - dtype=torch.float32, - ) - disto_center = torch.stack( - [atom_positions[idx] for idx in residue_disto_atom_idx], - dim=0, - ).unsqueeze(0) - - features: Dict[str, torch.Tensor] = { - "token_index": token_index, - "residue_index": residue_index, - "asym_id": asym_id, - "entity_id": entity_id, - "sym_id": sym_id, - "mol_type": mol_type, - "res_type": res_type, - "disto_center": disto_center, - "token_bonds": token_bonds, - "type_bonds": type_bonds, - "token_pad_mask": token_pad_mask, - "token_resolved_mask": token_resolved_mask, - "token_disto_mask": token_disto_mask, - "contact_conditioning": contact_conditioning, - "contact_threshold": contact_threshold, - "method_feature": method_feature, - "modified": modified, - "cyclic_period": cyclic_period, - "affinity_token_mask": affinity_token_mask, - "ref_pos": ref_pos, - "atom_resolved_mask": atom_resolved_mask, - "ref_atom_name_chars": ref_atom_name_chars, - "ref_element": ref_element, - "ref_charge": ref_charge, - "ref_chirality": ref_chirality, - "atom_backbone_feat": atom_backbone_feat, - "ref_space_uid": ref_space_uid, - "coords": coords, - "atom_pad_mask": atom_pad_mask, - "atom_to_token": atom_to_token, - "token_to_rep_atom": token_to_rep_atom, - "r_set_to_rep_atom": r_set_to_rep_atom, - "token_to_center_atom": token_to_center_atom, - "disto_target": disto_target, - "disto_coords_ensemble": disto_coords_ensemble, - "bfactor": bfactor, - "plddt": atom_plddt, - "frames_idx": frames_idx, - "frame_resolved_mask": frame_resolved_mask, - "msa": msa, - "msa_paired": msa_paired, - "deletion_value": deletion_value, - "has_deletion": has_deletion, - "deletion_mean": deletion_mean, - "profile": profile, - "msa_mask": msa_mask, - "template_restype": template_restype, - "template_frame_rot": template_frame_rot, - "template_frame_t": template_frame_t, - "template_cb": template_cb, - "template_ca": template_ca, - "template_mask_cb": template_mask_cb, - "template_mask_frame": template_mask_frame, - "template_mask": template_mask, - "query_to_template": query_to_template, - "visibility_ids": visibility_ids, - } - - return features, template diff --git a/fastplms/boltz/test_boltz2_featurization.py b/fastplms/boltz/test_boltz2_featurization.py deleted file mode 100644 index 7361911..0000000 --- a/fastplms/boltz/test_boltz2_featurization.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Featurization parity test: minimal_featurizer vs official Boltz2 pipeline. - -Compares feature tensors produced by both pipelines for the same input sequence. -Requires the `boltz` package to be installed for the official pipeline. - -Usage: - python -m fastplms.boltz.test_boltz2_featurization -""" - -import sys -from pathlib import Path -from typing import Dict, List - -import torch - -from .minimal_featurizer import build_boltz2_features -from . import vb_const as const - - -def _load_official_features(sequence: str) -> Dict[str, torch.Tensor]: - """Generate features using the official Boltz2 data pipeline.""" - try: - from boltz.main import check_inputs, process_inputs, BoltzProcessedInput, Manifest - from boltz.data.module.inferencev2 import BoltzInferenceDataModule - except ImportError: - print("ERROR: boltz package not installed. Install with: pip install boltz") - sys.exit(1) - - import tempfile - import yaml - - cache = Path("~/.boltz/").expanduser() - input_yaml = yaml.dump({ - "version": 1, - "sequences": [ - {"protein": {"id": "A", "sequence": sequence, "msa": "empty"}} - ], - }) - - with tempfile.TemporaryDirectory() as tmp_dir: - out_dir = Path(tmp_dir) - input_path = out_dir / "input.yaml" - input_path.write_text(input_yaml) - - data = check_inputs(input_path) - process_inputs( - data=data, - out_dir=out_dir, - ccd_path=cache / "ccd.pkl", - mol_dir=cache / "mols", - use_msa_server=False, - msa_server_url="https://api.colabfold.com", - msa_pairing_strategy="greedy", - ) - - processed_dir = out_dir / "processed" - processed = BoltzProcessedInput( - manifest=Manifest.load(processed_dir / "manifest.json"), - targets_dir=processed_dir / "structures", - msa_dir=processed_dir / "msa", - ) - - data_module = BoltzInferenceDataModule( - manifest=processed.manifest, - target_dir=processed.targets_dir, - msa_dir=processed.msa_dir, - num_workers=0, - ) - - features_dict = list(data_module.predict_dataloader())[0] - - features = {} - for k, v in features_dict.items(): - if k == "record": - continue - if torch.is_tensor(v): - features[k] = v.float() if v.is_floating_point() else v - else: - features[k] = torch.tensor(v) - return features - - -def _compare_features( - fast_feats: Dict[str, torch.Tensor], - official_feats: Dict[str, torch.Tensor], -) -> List[Dict]: - """Compare two feature dicts key-by-key. Returns list of mismatch reports.""" - all_keys = sorted(set(fast_feats.keys()) | set(official_feats.keys())) - mismatches = [] - - for key in all_keys: - if key not in fast_feats: - mismatches.append({"key": key, "issue": "MISSING in FastPLMs"}) - continue - if key not in official_feats: - mismatches.append({"key": key, "issue": "EXTRA in FastPLMs (not in official)"}) - continue - - f = fast_feats[key] - o = official_feats[key] - - if f.shape != o.shape: - mismatches.append({ - "key": key, - "issue": f"SHAPE MISMATCH: FastPLMs {tuple(f.shape)} vs official {tuple(o.shape)}", - }) - continue - - if f.dtype != o.dtype: - mismatches.append({ - "key": key, - "issue": f"DTYPE MISMATCH: FastPLMs {f.dtype} vs official {o.dtype}", - }) - - if f.is_floating_point(): - max_err = (f.float() - o.float()).abs().max().item() - mean_err = (f.float() - o.float()).abs().mean().item() - if max_err > 1e-4: - mismatches.append({ - "key": key, - "issue": f"VALUE MISMATCH: max_err={max_err:.6f}, mean_err={mean_err:.6f}", - }) - else: - n_diff = (f != o).sum().item() - if n_diff > 0: - mismatches.append({ - "key": key, - "issue": f"VALUE MISMATCH: {n_diff} differing elements out of {f.numel()}", - }) - - return mismatches - - -# Features where random augmentation makes exact comparison impossible; -# compare only shape and dtype. -_RANDOM_AUGMENTATION_KEYS = {"ref_pos", "coords", "disto_coords_ensemble"} - - -def main(): - test_sequence = "AAAAAAAAAAAAAAAAAAAAGGGGGGGGGGLLLLLLLLLLL" - print(f"Test sequence ({len(test_sequence)} residues): {test_sequence}") - print() - - print("Building FastPLMs features...") - fast_feats, _ = build_boltz2_features(test_sequence) - print(f" Keys: {len(fast_feats)}") - - print("Building official Boltz2 features...") - official_feats = _load_official_features(test_sequence) - print(f" Keys: {len(official_feats)}") - print() - - # Compare - mismatches = _compare_features(fast_feats, official_feats) - - # Filter out expected random augmentation differences - structural_mismatches = [] - augmentation_mismatches = [] - for m in mismatches: - if m["key"] in _RANDOM_AUGMENTATION_KEYS and "VALUE MISMATCH" in m["issue"]: - augmentation_mismatches.append(m) - else: - structural_mismatches.append(m) - - if augmentation_mismatches: - print(f"Expected differences (random augmentation, {len(augmentation_mismatches)} keys):") - for m in augmentation_mismatches: - print(f" {m['key']}: {m['issue']}") - print() - - if structural_mismatches: - print(f"STRUCTURAL MISMATCHES ({len(structural_mismatches)}):") - for m in structural_mismatches: - print(f" {m['key']}: {m['issue']}") - print() - print("FAIL: Feature parity not achieved.") - return 1 - else: - print("PASS: All features match (excluding expected random augmentation).") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/fastplms/boltz/vb_const.py b/fastplms/boltz/vb_const.py deleted file mode 100644 index 71fc1b6..0000000 --- a/fastplms/boltz/vb_const.py +++ /dev/null @@ -1,1184 +0,0 @@ -#################################################################################################### -# CHAINS -#################################################################################################### - -chain_types = [ - "PROTEIN", - "DNA", - "RNA", - "NONPOLYMER", -] -chain_type_ids = {chain: i for i, chain in enumerate(chain_types)} - -out_types = [ - "dna_protein", - "rna_protein", - "ligand_protein", - "dna_ligand", - "rna_ligand", - "intra_ligand", - "intra_dna", - "intra_rna", - "intra_protein", - "protein_protein", - "modified", -] - -out_types_weights_af3 = { - "dna_protein": 10.0, - "rna_protein": 10.0, - "ligand_protein": 10.0, - "dna_ligand": 5.0, - "rna_ligand": 5.0, - "intra_ligand": 20.0, - "intra_dna": 4.0, - "intra_rna": 16.0, - "intra_protein": 20.0, - "protein_protein": 20.0, - "modified": 0.0, -} - -out_types_weights = { - "dna_protein": 5.0, - "rna_protein": 5.0, - "ligand_protein": 20.0, - "dna_ligand": 2.0, - "rna_ligand": 2.0, - "intra_ligand": 20.0, - "intra_dna": 2.0, - "intra_rna": 8.0, - "intra_protein": 20.0, - "protein_protein": 20.0, - "modified": 0.0, -} - - -out_single_types = ["protein", "ligand", "dna", "rna"] - -clash_types = [ - "dna_protein", - "rna_protein", - "ligand_protein", - "protein_protein", - "dna_ligand", - "rna_ligand", - "ligand_ligand", - "rna_dna", - "dna_dna", - "rna_rna", -] - -chain_types_to_clash_type = { - frozenset(("PROTEIN", "DNA")): "dna_protein", - frozenset(("PROTEIN", "RNA")): "rna_protein", - frozenset(("PROTEIN", "NONPOLYMER")): "ligand_protein", - frozenset(("PROTEIN",)): "protein_protein", - frozenset(("NONPOLYMER", "DNA")): "dna_ligand", - frozenset(("NONPOLYMER", "RNA")): "rna_ligand", - frozenset(("NONPOLYMER",)): "ligand_ligand", - frozenset(("DNA", "RNA")): "rna_dna", - frozenset(("DNA",)): "dna_dna", - frozenset(("RNA",)): "rna_rna", -} - -chain_type_to_out_single_type = { - "PROTEIN": "protein", - "DNA": "dna", - "RNA": "rna", - "NONPOLYMER": "ligand", -} -#################################################################################################### -# RESIDUES & TOKENS -#################################################################################################### - - -canonical_tokens = [ - "ALA", - "ARG", - "ASN", - "ASP", - "CYS", - "GLN", - "GLU", - "GLY", - "HIS", - "ILE", - "LEU", - "LYS", - "MET", - "PHE", - "PRO", - "SER", - "THR", - "TRP", - "TYR", - "VAL", - "UNK", # unknown protein token -] - -tokens = [ - "", - "-", - *canonical_tokens, - "A", - "G", - "C", - "U", - "N", # unknown rna token - "DA", - "DG", - "DC", - "DT", - "DN", # unknown dna token -] - -token_ids = {token: i for i, token in enumerate(tokens)} -num_tokens = len(tokens) -unk_token = {"PROTEIN": "UNK", "DNA": "DN", "RNA": "N"} -unk_token_ids = {m: token_ids[t] for m, t in unk_token.items()} - -prot_letter_to_token = { - "A": "ALA", - "R": "ARG", - "N": "ASN", - "D": "ASP", - "C": "CYS", - "E": "GLU", - "Q": "GLN", - "G": "GLY", - "H": "HIS", - "I": "ILE", - "L": "LEU", - "K": "LYS", - "M": "MET", - "F": "PHE", - "P": "PRO", - "S": "SER", - "T": "THR", - "W": "TRP", - "Y": "TYR", - "V": "VAL", - "X": "UNK", - "J": "UNK", - "B": "UNK", - "Z": "UNK", - "O": "UNK", - "U": "UNK", - "-": "-", -} - -prot_token_to_letter = {v: k for k, v in prot_letter_to_token.items()} -prot_token_to_letter["UNK"] = "X" - -rna_letter_to_token = { - "A": "A", - "G": "G", - "C": "C", - "U": "U", - "N": "N", -} -rna_token_to_letter = {v: k for k, v in rna_letter_to_token.items()} - -dna_letter_to_token = { - "A": "DA", - "G": "DG", - "C": "DC", - "T": "DT", - "N": "DN", -} -dna_token_to_letter = {v: k for k, v in dna_letter_to_token.items()} - -#################################################################################################### -# ATOMS -#################################################################################################### - -num_elements = 128 - -chirality_types = [ - "CHI_UNSPECIFIED", - "CHI_TETRAHEDRAL_CW", - "CHI_TETRAHEDRAL_CCW", - "CHI_SQUAREPLANAR", - "CHI_OCTAHEDRAL", - "CHI_TRIGONALBIPYRAMIDAL", - "CHI_OTHER", -] -chirality_type_ids = {chirality: i for i, chirality in enumerate(chirality_types)} -unk_chirality_type = "CHI_OTHER" - -hybridization_map = [ - "S", - "SP", - "SP2", - "SP2D", - "SP3", - "SP3D", - "SP3D2", - "OTHER", - "UNSPECIFIED", -] -hybridization_type_ids = {hybrid: i for i, hybrid in enumerate(hybridization_map)} -unk_hybridization_type = "UNSPECIFIED" - -# fmt: off -ref_atoms = { - "PAD": [], - "UNK": ["N", "CA", "C", "O", "CB"], - "-": [], - "ALA": ["N", "CA", "C", "O", "CB"], - "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"], - "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"], - "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"], - "CYS": ["N", "CA", "C", "O", "CB", "SG"], - "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"], - "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"], - "GLY": ["N", "CA", "C", "O"], - "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"], - "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"], - "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"], - "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"], - "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"], - "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"], - "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"], - "SER": ["N", "CA", "C", "O", "CB", "OG"], - "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"], - "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"], # noqa: E501 - "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"], - "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"], - "A": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501 - "G": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501 - "C": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501 - "U": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"], # noqa: E501 - "N": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"], # noqa: E501 - "DA": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], # noqa: E501 - "DG": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], # noqa: E501 - "DC": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], # noqa: E501 - "DT": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"], # noqa: E501 - "DN": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'"] -} - -protein_backbone_atom_names = ["N", "CA", "C", "O"] -nucleic_backbone_atom_names = ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"] - -protein_backbone_atom_index = {name: i for i, name in enumerate(protein_backbone_atom_names)} -nucleic_backbone_atom_index = {name: i for i, name in enumerate(nucleic_backbone_atom_names)} - -ref_symmetries = { - "PAD": [], - "ALA": [], - "ARG": [], - "ASN": [], - "ASP": [[(6, 7), (7, 6)]], - "CYS": [], - "GLN": [], - "GLU": [[(7, 8), (8, 7)]], - "GLY": [], - "HIS": [], - "ILE": [], - "LEU": [], - "LYS": [], - "MET": [], - "PHE": [[(6, 7), (7, 6), (8, 9), (9, 8)]], - "PRO": [], - "SER": [], - "THR": [], - "TRP": [], - "TYR": [[(6, 7), (7, 6), (8, 9), (9, 8)]], - "VAL": [], - "A": [[(1, 2), (2, 1)]], - "G": [[(1, 2), (2, 1)]], - "C": [[(1, 2), (2, 1)]], - "U": [[(1, 2), (2, 1)]], - #"N": [[(1, 2), (2, 1)]], - "DA": [[(1, 2), (2, 1)]], - "DG": [[(1, 2), (2, 1)]], - "DC": [[(1, 2), (2, 1)]], - "DT": [[(1, 2), (2, 1)]], - #"DN": [[(1, 2), (2, 1)]] -} - - -res_to_center_atom = { - "UNK": "CA", - "ALA": "CA", - "ARG": "CA", - "ASN": "CA", - "ASP": "CA", - "CYS": "CA", - "GLN": "CA", - "GLU": "CA", - "GLY": "CA", - "HIS": "CA", - "ILE": "CA", - "LEU": "CA", - "LYS": "CA", - "MET": "CA", - "PHE": "CA", - "PRO": "CA", - "SER": "CA", - "THR": "CA", - "TRP": "CA", - "TYR": "CA", - "VAL": "CA", - "A": "C1'", - "G": "C1'", - "C": "C1'", - "U": "C1'", - "N": "C1'", - "DA": "C1'", - "DG": "C1'", - "DC": "C1'", - "DT": "C1'", - "DN": "C1'" -} - -res_to_disto_atom = { - "UNK": "CB", - "ALA": "CB", - "ARG": "CB", - "ASN": "CB", - "ASP": "CB", - "CYS": "CB", - "GLN": "CB", - "GLU": "CB", - "GLY": "CA", - "HIS": "CB", - "ILE": "CB", - "LEU": "CB", - "LYS": "CB", - "MET": "CB", - "PHE": "CB", - "PRO": "CB", - "SER": "CB", - "THR": "CB", - "TRP": "CB", - "TYR": "CB", - "VAL": "CB", - "A": "C4", - "G": "C4", - "C": "C2", - "U": "C2", - "N": "C1'", - "DA": "C4", - "DG": "C4", - "DC": "C2", - "DT": "C2", - "DN": "C1'" -} - -res_to_center_atom_id = { - res: ref_atoms[res].index(atom) - for res, atom in res_to_center_atom.items() -} - -res_to_disto_atom_id = { - res: ref_atoms[res].index(atom) - for res, atom in res_to_disto_atom.items() -} - -# fmt: on - -#################################################################################################### -# BONDS -#################################################################################################### - -atom_interface_cutoff = 5.0 -interface_cutoff = 15.0 - -bond_types = [ - "OTHER", - "SINGLE", - "DOUBLE", - "TRIPLE", - "AROMATIC", - "COVALENT", -] -bond_type_ids = {bond: i for i, bond in enumerate(bond_types)} -unk_bond_type = "OTHER" - - -#################################################################################################### -# Contacts -#################################################################################################### - - -pocket_contact_info = { - "UNSPECIFIED": 0, - "UNSELECTED": 1, - "POCKET": 2, - "BINDER": 3, -} - -contact_conditioning_info = { - "UNSPECIFIED": 0, - "UNSELECTED": 1, - "POCKET>BINDER": 2, - "BINDER>POCKET": 3, - "CONTACT": 4, -} - - -#################################################################################################### -# MSA -#################################################################################################### - -max_msa_seqs = 16384 -max_paired_seqs = 8192 - - -#################################################################################################### -# CHUNKING -#################################################################################################### - -chunk_size_threshold = 384 - -#################################################################################################### -# Method conditioning -#################################################################################################### - -# Methods -method_types_ids = { - "MD": 0, - "X-RAY DIFFRACTION": 1, - "ELECTRON MICROSCOPY": 2, - "SOLUTION NMR": 3, - "SOLID-STATE NMR": 4, - "NEUTRON DIFFRACTION": 4, - "ELECTRON CRYSTALLOGRAPHY": 4, - "FIBER DIFFRACTION": 4, - "POWDER DIFFRACTION": 4, - "INFRARED SPECTROSCOPY": 4, - "FLUORESCENCE TRANSFER": 4, - "EPR": 4, - "THEORETICAL MODEL": 4, - "SOLUTION SCATTERING": 4, - "OTHER": 4, - "AFDB": 5, - "BOLTZ-1": 6, - "FUTURE1": 7, # Placeholder for future supervision sources - "FUTURE2": 8, - "FUTURE3": 9, - "FUTURE4": 10, - "FUTURE5": 11, -} -method_types_ids = {k.lower(): v for k, v in method_types_ids.items()} -num_method_types = len(set(method_types_ids.values())) - -# Temperature -temperature_bins = [(265, 280), (280, 295), (295, 310)] -temperature_bins_ids = {temp: i for i, temp in enumerate(temperature_bins)} -temperature_bins_ids["other"] = len(temperature_bins) -num_temp_bins = len(temperature_bins_ids) - - -# pH -ph_bins = [(0, 6), (6, 8), (8, 14)] -ph_bins_ids = {ph: i for i, ph in enumerate(ph_bins)} -ph_bins_ids["other"] = len(ph_bins) -num_ph_bins = len(ph_bins_ids) - -#################################################################################################### -# VDW_RADII -#################################################################################################### - -# fmt: off -vdw_radii = [ - 1.2, 1.4, 2.2, 1.9, 1.8, 1.7, 1.6, 1.55, 1.5, 1.54, - 2.4, 2.2, 2.1, 2.1, 1.95, 1.8, 1.8, 1.88, 2.8, 2.4, - 2.3, 2.15, 2.05, 2.05, 2.05, 2.05, 2.0, 2.0, 2.0, 2.1, - 2.1, 2.1, 2.05, 1.9, 1.9, 2.02, 2.9, 2.55, 2.4, 2.3, - 2.15, 2.1, 2.05, 2.05, 2.0, 2.05, 2.1, 2.2, 2.2, 2.25, - 2.2, 2.1, 2.1, 2.16, 3.0, 2.7, 2.5, 2.48, 2.47, 2.45, - 2.43, 2.42, 2.4, 2.38, 2.37, 2.35, 2.33, 2.32, 2.3, 2.28, - 2.27, 2.25, 2.2, 2.1, 2.05, 2.0, 2.0, 2.05, 2.1, 2.05, - 2.2, 2.3, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.4, - 2.0, 2.3, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, - 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, - 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 -] -# fmt: on - -#################################################################################################### -# Excluded ligands -#################################################################################################### - -ligand_exclusion = { - "144", - "15P", - "1PE", - "2F2", - "2JC", - "3HR", - "3SY", - "7N5", - "7PE", - "9JE", - "AAE", - "ABA", - "ACE", - "ACN", - "ACT", - "ACY", - "AZI", - "BAM", - "BCN", - "BCT", - "BDN", - "BEN", - "BME", - "BO3", - "BTB", - "BTC", - "BU1", - "C8E", - "CAD", - "CAQ", - "CBM", - "CCN", - "CIT", - "CL", - "CLR", - "CM", - "CMO", - "CO3", - "CPT", - "CXS", - "D10", - "DEP", - "DIO", - "DMS", - "DN", - "DOD", - "DOX", - "EDO", - "EEE", - "EGL", - "EOH", - "EOX", - "EPE", - "ETF", - "FCY", - "FJO", - "FLC", - "FMT", - "FW5", - "GOL", - "GSH", - "GTT", - "GYF", - "HED", - "IHP", - "IHS", - "IMD", - "IOD", - "IPA", - "IPH", - "LDA", - "MB3", - "MEG", - "MES", - "MLA", - "MLI", - "MOH", - "MPD", - "MRD", - "MSE", - "MYR", - "N", - "NA", - "NH2", - "NH4", - "NHE", - "NO3", - "O4B", - "OHE", - "OLA", - "OLC", - "OMB", - "OME", - "OXA", - "P6G", - "PE3", - "PE4", - "PEG", - "PEO", - "PEP", - "PG0", - "PG4", - "PGE", - "PGR", - "PLM", - "PO4", - "POL", - "POP", - "PVO", - "SAR", - "SCN", - "SEO", - "SEP", - "SIN", - "SO4", - "SPD", - "SPM", - "SR", - "STE", - "STO", - "STU", - "TAR", - "TBU", - "TME", - "TPO", - "TRS", - "UNK", - "UNL", - "UNX", - "UPL", - "URE", -} - - -#################################################################################################### -# TEMPLATES -#################################################################################################### - -min_coverage_residues = 10 -min_coverage_fraction = 0.1 - - -#################################################################################################### -# Ambiguous atoms -#################################################################################################### - -ambiguous_atoms = { - "CA": { - "*": "C", - "OEX": "CA", - "OEC": "CA", - "543": "CA", - "OC6": "CA", - "OC1": "CA", - "OC7": "CA", - "OEY": "CA", - "OC4": "CA", - "OC3": "CA", - "ICA": "CA", - "CA": "CA", - "OC2": "CA", - "OC5": "CA", - }, - "CD": {"*": "C", "CD": "CD", "CD3": "CD", "CD5": "CD", "CD1": "CD"}, - "BR": "BR", - "CL": { - "*": "CL", - "C8P": "C", - "L3T": "C", - "TLC": "C", - "TZ0": "C", - "471": "C", - "NLK": "C", - "PGM": "C", - "PNE": "C", - "RCY": "C", - "11F": "C", - "PII": "C", - "C1Q": "C", - "4MD": "C", - "R5A": "C", - "KW2": "C", - "I7M": "C", - "R48": "C", - "FC3": "C", - "55V": "C", - "KPF": "C", - "SPZ": "C", - "0TT": "C", - "R9A": "C", - "5NA": "C", - "C55": "C", - "NIX": "C", - "5PM": "C", - "PP8": "C", - "544": "C", - "812": "C", - "NPM": "C", - "KU8": "C", - "A1AMM": "C", - "4S0": "C", - "AQC": "C", - "2JK": "C", - "WJR": "C", - "A1AAW": "C", - "85E": "C", - "MB0": "C", - "ZAB": "C", - "85K": "C", - "GBP": "C", - "A1H80": "C", - "A1AFR": "C", - "L9M": "C", - "MYK": "C", - "MB9": "C", - "38R": "C", - "EKB": "C", - "NKF": "C", - "UMQ": "C", - "T4K": "C", - "3PT": "C", - "A1A7S": "C", - "1Q9": "C", - "11R": "C", - "D2V": "C", - "SM8": "C", - "IFC": "C", - "DB5": "C", - "L2T": "C", - "GNB": "C", - "PP7": "C", - "072": "C", - "P88": "C", - "DRL": "C", - "C9W": "C", - "NTP": "C", - "4HJ": "C", - "7NA": "C", - "LPC": "C", - "T8W": "C", - "63R": "C", - "570": "C", - "R4A": "C", - "3BG": "C", - "4RB": "C", - "GSO": "C", - "BQ6": "C", - "R4P": "C", - "5CP": "C", - "TTR": "C", - "6UZ": "C", - "SPJ": "C", - "0SA": "C", - "ZL1": "C", - "BYG": "C", - "F0E": "C", - "PC0": "C", - "B2Q": "C", - "KV6": "C", - "NTO": "C", - "CLG": "C", - "R7U": "C", - "SMQ": "C", - "GM2": "C", - "Z7P": "C", - "NXF": "C", - "C6Q": "C", - "A1G": "C", - "433": "C", - "L9N": "C", - "7OX": "C", - "A1H84": "C", - "97L": "C", - "HDV": "C", - "LUO": "C", - "R6A": "C", - "1PC": "C", - "4PT": "C", - "SBZ": "C", - "EAB": "C", - "FL4": "C", - "OPS": "C", - "C2X": "C", - "SLL": "C", - "BFC": "C", - "GIP": "C", - "7CP": "C", - "CLH": "C", - "34E": "C", - "5NE": "C", - "PBF": "C", - "ABD": "C", - "ABC": "C", - "LPF": "C", - "TIZ": "C", - "4HH": "C", - "AFC": "C", - "WQH": "C", - "9JL": "C", - "CS3": "C", - "NL0": "C", - "KPY": "C", - "DNA": "C", - "B3C": "C", - "TKL": "C", - "KVS": "C", - "HO6": "C", - "NLH": "C", - "1PB": "C", - "CYF": "C", - "G4M": "C", - "R5B": "C", - "N4S": "C", - "N11": "C", - "C8F": "C", - "PIJ": "C", - "WIN": "C", - "NT1": "C", - "WJW": "C", - "HF7": "C", - "TY1": "C", - "VM1": "C", - }, - "OS": {"*": "O", "DWC": "OS", "OHX": "OS", "OS": "OS", "8WV": "OS", "OS4": "OS"}, - "PB": {"*": "P", "ZN9": "PB", "ZN7": "PB", "PBM": "PB", "PB": "PB", "CSB": "PB"}, - "CE": {"*": "C", "CE": "CE"}, - "FE": {"*": "FE", "TFR": "F", "PF5": "F", "IFC": "F", "F5C": "F"}, - "NA": {"*": "N", "CGO": "NA", "R2K": "NA", "LVQ": "NA", "NA": "NA"}, - "ND": {"*": "N", "ND": "ND"}, - "CF": {"*": "C", "CF": "CF"}, - "RU": "RU", - "BRAF": "BR", - "EU": "EU", - "CLAA": "CL", - "CLBQ": "CL", - "CM": {"*": "C", "ZCM": "CM"}, - "SN": {"*": "SN", "TAP": "S", "SND": "S", "TAD": "S", "XPT": "S"}, - "AG": "AG", - "CLN": "CL", - "CLM": "CL", - "CLA": {"*": "CL", "PII": "C", "TDL": "C", "D0J": "C", "GM2": "C", "PIJ": "C"}, - "CLB": { - "*": "CL", - "TD5": "C", - "PII": "C", - "TDL": "C", - "GM2": "C", - "TD7": "C", - "TD6": "C", - "PIJ": "C", - }, - "CR": { - "*": "C", - "BW9": "CR", - "CQ4": "CR", - "AC9": "CR", - "TIL": "CR", - "J7U": "CR", - "CR": "CR", - }, - "CLAY": "CL", - "CLBC": "CL", - "PD": { - "*": "P", - "F6Q": "PD", - "SVP": "PD", - "SXC": "PD", - "U5U": "PD", - "PD": "PD", - "PLL": "PD", - }, - "CO": { - "*": "C", - "J1S": "CO", - "OCN": "CO", - "OL3": "CO", - "OL4": "CO", - "B12": "CO", - "XCO": "CO", - "UFU": "CO", - "CON": "CO", - "OL5": "CO", - "B13": "CO", - "7KI": "CO", - "PL1": "CO", - "OCO": "CO", - "J1R": "CO", - "COH": "CO", - "SIR": "CO", - "6KI": "CO", - "NCO": "CO", - "9CO": "CO", - "PC3": "CO", - "BWU": "CO", - "B1Z": "CO", - "J83": "CO", - "CO": "CO", - "COY": "CO", - "CNC": "CO", - "3CO": "CO", - "OCL": "CO", - "R5Q": "CO", - "X5Z": "CO", - "CBY": "CO", - "OLS": "CO", - "F0X": "CO", - "I2A": "CO", - "OCM": "CO", - }, - "CU": { - "*": "C", - "8ZR": "CU", - "K7E": "CU", - "CU3": "CU", - "SI9": "CU", - "35N": "CU", - "C2O": "CU", - "SI7": "CU", - "B15": "CU", - "SI0": "CU", - "CUP": "CU", - "SQ1": "CU", - "CUK": "CU", - "CUL": "CU", - "SI8": "CU", - "IC4": "CU", - "CUM": "CU", - "MM2": "CU", - "B30": "CU", - "S32": "CU", - "V79": "CU", - "IMF": "CU", - "CUN": "CU", - "MM1": "CU", - "MP1": "CU", - "IME": "CU", - "B17": "CU", - "C2C": "CU", - "1CU": "CU", - "CU6": "CU", - "C1O": "CU", - "CU1": "CU", - "B22": "CU", - "CUS": "CU", - "RUQ": "CU", - "CUF": "CU", - "CUA": "CU", - "CU": "CU", - "CUO": "CU", - "0TE": "CU", - "SI4": "CU", - }, - "CS": {"*": "C", "CS": "CS"}, - "CLQ": "CL", - "CLR": "CL", - "CLU": "CL", - "TE": "TE", - "NI": { - "*": "N", - "USN": "NI", - "NFO": "NI", - "NI2": "NI", - "NFS": "NI", - "NFR": "NI", - "82N": "NI", - "R5N": "NI", - "NFU": "NI", - "A1ICD": "NI", - "NI3": "NI", - "M43": "NI", - "MM5": "NI", - "BF8": "NI", - "TCN": "NI", - "NIK": "NI", - "CUV": "NI", - "MM6": "NI", - "J52": "NI", - "NI": "NI", - "SNF": "NI", - "XCC": "NI", - "F0L": "NI", - "UWE": "NI", - "NFC": "NI", - "3NI": "NI", - "HNI": "NI", - "F43": "NI", - "RQM": "NI", - "NFE": "NI", - "NFB": "NI", - "B51": "NI", - "NI1": "NI", - "WCC": "NI", - "NUF": "NI", - }, - "SB": {"*": "S", "UJI": "SB", "SB": "SB", "118": "SB", "SBO": "SB", "3CG": "SB"}, - "MO": "MO", - "SEG": "SE", - "CLL": "CL", - "CLAH": "CL", - "CLC": { - "*": "CL", - "TD5": "C", - "PII": "C", - "TDL": "C", - "GM2": "C", - "TD7": "C", - "TD6": "C", - "PIJ": "C", - }, - "CLD": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"}, - "CLAD": "CL", - "CLAE": "CL", - "LA": "LA", - "RH": "RH", - "BRAC": "BR", - "BRAD": "BR", - "CLBN": "CL", - "CLAC": "CL", - "BRAB": "BR", - "BRAE": "BR", - "MG": "MG", - "IR": "IR", - "SE": { - "*": "SE", - "HII": "S", - "NT2": "S", - "R2P": "S", - "S2P": "S", - "0IU": "S", - "QMB": "S", - "81S": "S", - "0QB": "S", - "UB4": "S", - "OHS": "S", - "Q78": "S", - "0Y2": "S", - "B3M": "S", - "NT1": "S", - "81R": "S", - }, - "BRAG": "BR", - "CLF": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"}, - "CLE": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"}, - "BRAX": "BR", - "CLK": "CL", - "ZN": "ZN", - "AS": "AS", - "AU": "AU", - "PT": "PT", - "CLAS": "CL", - "MN": "MN", - "CLBE": "CL", - "CLBF": "CL", - "CLAF": "CL", - "NA'": {"*": "N", "CGO": "NA"}, - "BRAH": "BR", - "BRAI": "BR", - "BRA": "BR", - "BRB": "BR", - "BRAV": "BR", - "HG": { - "*": "HG", - "BBA": "H", - "MID": "H", - "APM": "H", - "4QQ": "H", - "0ZG": "H", - "APH": "H", - }, - "AR": "AR", - "D": "H", - "CLAN": "CL", - "SI": "SI", - "CLS": "CL", - "ZR": "ZR", - "CLAR": {"*": "CL", "ZM4": "C"}, - "HO": "HO", - "CLI": {"*": "CL", "GM2": "C"}, - "CLH": {"*": "CL", "GM2": "C"}, - "CLAP": "CL", - "CLBL": "CL", - "CLBM": "CL", - "PR": {"*": "PR", "UF0": "P", "252": "P"}, - "IN": "IN", - "CLJ": "CL", - "BRU": "BR", - "SC": {"*": "S", "SFL": "SC"}, - "CLG": {"*": "CL", "GM2": "C"}, - "BRAT": "BR", - "BRAR": "BR", - "CLAG": "CL", - "CLAB": "CL", - "CLV": "CL", - "TI": "TI", - "CLAX": "CL", - "CLAJ": "CL", - "CL'": {"*": "CL", "BNR": "C", "25A": "C", "BDA": "C"}, - "CLAW": "CL", - "BRF": "BR", - "BRE": "BR", - "RE": "RE", - "GD": "GD", - "SM": {"*": "S", "SM": "SM"}, - "CLBH": "CL", - "CLBI": "CL", - "CLAI": "CL", - "CLY": "CL", - "CLZ": "CL", - "AC": "AC", - "BR'": "BR", - "CLT": "CL", - "CLO": "CL", - "CLP": "CL", - "LU": "LU", - "BA": {"*": "B", "BA": "BA"}, - "CLAU": "CL", - "RB": "RB", - "LI": "LI", - "MOM": "MO", - "BRAQ": "BR", - "SR": {"*": "S", "SR": "SR", "OER": "SR"}, - "CLAT": "CL", - "BRAL": "BR", - "SEB": "SE", - "CLW": "CL", - "CLX": "CL", - "BE": "BE", - "BRG": "BR", - "SEA": "SE", - "BRAW": "BR", - "BRBB": "BR", - "ER": "ER", - "TH": "TH", - "BRR": "BR", - "CLBV": "CL", - "AL": "AL", - "CLAV": "CL", - "BRH": "BR", - "CLAQ": "CL", - "GA": "GA", - "X": "*", - "TL": "TL", - "CLBB": "CL", - "TB": "TB", - "CLAK": "CL", - "XE": {"*": "*", "XE": "XE"}, - "SEL": "SE", - "PU": {"*": "P", "4PU": "PU"}, - "CLAZ": "CL", - "SE'": "SE", - "CLBA": "CL", - "SEN": "SE", - "SNN": "SN", - "MOB": "MO", - "YB": "YB", - "BRC": "BR", - "BRD": "BR", - "CLAM": "CL", - "DA": "H", - "DB": "H", - "DC": "H", - "DXT": "H", - "DXU": "H", - "DXX": "H", - "DXY": "H", - "DXZ": "H", - "DY": "DY", - "TA": "TA", - "XD": "*", - "SED": "SE", - "CLAL": "CL", - "BRAJ": "BR", - "AM": "AM", - "CLAO": "CL", - "BI": "BI", - "KR": "KR", - "BRBJ": "BR", - "UNK": "*", -} diff --git a/fastplms/boltz/vb_layers_attention.py b/fastplms/boltz/vb_layers_attention.py deleted file mode 100644 index 3275667..0000000 --- a/fastplms/boltz/vb_layers_attention.py +++ /dev/null @@ -1,136 +0,0 @@ -from typing import Optional - -import torch -from einops.layers.torch import Rearrange -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -class AttentionPairBias(nn.Module): - """Attention pair bias layer.""" - - def __init__( - self, - c_s: int, - c_z: int, - num_heads: int, - inf: float = 1e6, - initial_norm: bool = True, - ) -> None: - """Initialize the attention pair bias layer. - - Parameters - ---------- - c_s : int - The input sequence dimension. - c_z : int - The input pairwise dimension. - num_heads : int - The number of heads. - inf : float, optional - The inf value, by default 1e6 - initial_norm: bool, optional - Whether to apply layer norm to the input, by default True - - """ - super().__init__() - - assert c_s % num_heads == 0 - - self.c_s = c_s - self.num_heads = num_heads - self.head_dim = c_s // num_heads - self.inf = inf - - self.initial_norm = initial_norm - if self.initial_norm: - self.norm_s = nn.LayerNorm(c_s) - - self.proj_q = nn.Linear(c_s, c_s) - self.proj_k = nn.Linear(c_s, c_s, bias=False) - self.proj_v = nn.Linear(c_s, c_s, bias=False) - self.proj_g = nn.Linear(c_s, c_s, bias=False) - - self.proj_z = nn.Sequential( - nn.LayerNorm(c_z), - nn.Linear(c_z, num_heads, bias=False), - Rearrange("b ... h -> b h ..."), - ) - - self.proj_o = nn.Linear(c_s, c_s, bias=False) - init.final_init_(self.proj_o.weight) - - def forward( - self, - s: Tensor, - z: Tensor, - mask: Tensor, - k_in: Optional[Tensor] = None, - multiplicity: int = 1, - to_keys=None, - model_cache=None, - ) -> Tensor: - """Forward pass. - - Parameters - ---------- - s : torch.Tensor - The input sequence tensor (B, S, D) - z : torch.Tensor - The input pairwise tensor (B, N, N, D) - mask : torch.Tensor - The pairwise mask tensor (B, N) - multiplicity : int, optional - The diffusion batch size, by default 1 - - Returns - ------- - torch.Tensor - The output sequence tensor. - - """ - B = s.shape[0] - - # Layer norms - if self.initial_norm: - s = self.norm_s(s) - - if to_keys is not None: - k_in = to_keys(s) - mask = to_keys(mask.unsqueeze(-1)).squeeze(-1) - else: - if k_in is None: - k_in = s - - # Compute projections - q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim) - k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim) - v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim) - - # Caching z projection during diffusion roll-out - if model_cache is None or "z" not in model_cache: - z = self.proj_z(z) - - if model_cache is not None: - model_cache["z"] = z - else: - z = model_cache["z"] - z = z.repeat_interleave(multiplicity, 0) - - g = self.proj_g(s).sigmoid() - - with torch.autocast("cuda", enabled=False): - # Compute attention weights - attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float()) - attn = attn / (self.head_dim**0.5) + z.float() - # The pairwise mask tensor (B, N) is broadcasted to (B, 1, 1, N) and (B, H, N, N) - attn = attn + (1 - mask[:, None, None].float()) * -self.inf - attn = attn.softmax(dim=-1) - - # Compute output - o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype) - o = o.reshape(B, -1, self.c_s) - o = self.proj_o(g * o) - - return o diff --git a/fastplms/boltz/vb_layers_attentionv2.py b/fastplms/boltz/vb_layers_attentionv2.py deleted file mode 100644 index ceb70d8..0000000 --- a/fastplms/boltz/vb_layers_attentionv2.py +++ /dev/null @@ -1,111 +0,0 @@ -from typing import Optional - -import torch -from einops.layers.torch import Rearrange -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -class AttentionPairBias(nn.Module): - """Attention pair bias layer.""" - - def __init__( - self, - c_s: int, - c_z: Optional[int] = None, - num_heads: Optional[int] = None, - inf: float = 1e6, - compute_pair_bias: bool = True, - ) -> None: - """Initialize the attention pair bias layer. - - Parameters - ---------- - c_s : int - The input sequence dimension. - c_z : int - The input pairwise dimension. - num_heads : int - The number of heads. - inf : float, optional - The inf value, by default 1e6 - - """ - super().__init__() - - assert c_s % num_heads == 0 - - self.c_s = c_s - self.num_heads = num_heads - self.head_dim = c_s // num_heads - self.inf = inf - - self.proj_q = nn.Linear(c_s, c_s) - self.proj_k = nn.Linear(c_s, c_s, bias=False) - self.proj_v = nn.Linear(c_s, c_s, bias=False) - self.proj_g = nn.Linear(c_s, c_s, bias=False) - - self.compute_pair_bias = compute_pair_bias - if compute_pair_bias: - self.proj_z = nn.Sequential( - nn.LayerNorm(c_z), - nn.Linear(c_z, num_heads, bias=False), - Rearrange("b ... h -> b h ..."), - ) - else: - self.proj_z = Rearrange("b ... h -> b h ...") - - self.proj_o = nn.Linear(c_s, c_s, bias=False) - init.final_init_(self.proj_o.weight) - - def forward( - self, - s: Tensor, - z: Tensor, - mask: Tensor, - k_in: Tensor, - multiplicity: int = 1, - ) -> Tensor: - """Forward pass. - - Parameters - ---------- - s : torch.Tensor - The input sequence tensor (B, S, D) - z : torch.Tensor - The input pairwise tensor or bias (B, N, N, D) - mask : torch.Tensor - The pairwise mask tensor (B, N, N) - - Returns - ------- - torch.Tensor - The output sequence tensor. - - """ - B = s.shape[0] - - # Compute projections - q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim) - k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim) - v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim) - - bias = self.proj_z(z) - bias = bias.repeat_interleave(multiplicity, 0) - - g = self.proj_g(s).sigmoid() - - with torch.autocast("cuda", enabled=False): - # Compute attention weights - attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float()) - attn = attn / (self.head_dim**0.5) + bias.float() - attn = attn + (1 - mask[:, None, None].float()) * -self.inf - attn = attn.softmax(dim=-1) - - # Compute output - o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype) - o = o.reshape(B, -1, self.c_s) - o = self.proj_o(g * o) - - return o diff --git a/fastplms/boltz/vb_layers_confidence_utils.py b/fastplms/boltz/vb_layers_confidence_utils.py deleted file mode 100644 index 5267d28..0000000 --- a/fastplms/boltz/vb_layers_confidence_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -import torch -from torch import nn - -from . import vb_const as const - - -def compute_collinear_mask(v1, v2): - norm1 = torch.norm(v1, dim=1, keepdim=True) - norm2 = torch.norm(v2, dim=1, keepdim=True) - v1 = v1 / (norm1 + 1e-6) - v2 = v2 / (norm2 + 1e-6) - mask_angle = torch.abs(torch.sum(v1 * v2, dim=1)) < 0.9063 - mask_overlap1 = norm1.reshape(-1) > 1e-2 - mask_overlap2 = norm2.reshape(-1) > 1e-2 - return mask_angle & mask_overlap1 & mask_overlap2 - - -def compute_frame_pred( - pred_atom_coords, - frames_idx_true, - feats, - multiplicity, - resolved_mask=None, - inference=False, -): - with torch.amp.autocast("cuda", enabled=False): - asym_id_token = feats["asym_id"] - asym_id_atom = torch.bmm( - feats["atom_to_token"].float(), asym_id_token.unsqueeze(-1).float() - ).squeeze(-1) - - B, N, _ = pred_atom_coords.shape - pred_atom_coords = pred_atom_coords.reshape(B // multiplicity, multiplicity, -1, 3) - frames_idx_pred = ( - frames_idx_true.clone() - .repeat_interleave(multiplicity, 0) - .reshape(B // multiplicity, multiplicity, -1, 3) - ) - - # Iterate through the batch and modify the frames for nonpolymers - for i, pred_atom_coord in enumerate(pred_atom_coords): - token_idx = 0 - atom_idx = 0 - for id in torch.unique(asym_id_token[i]): - mask_chain_token = (asym_id_token[i] == id) * feats["token_pad_mask"][i] - mask_chain_atom = (asym_id_atom[i] == id) * feats["atom_pad_mask"][i] - num_tokens = int(mask_chain_token.sum().item()) - num_atoms = int(mask_chain_atom.sum().item()) - if ( - feats["mol_type"][i, token_idx] != const.chain_type_ids["NONPOLYMER"] - or num_atoms < 3 - ): - token_idx += num_tokens - atom_idx += num_atoms - continue - dist_mat = ( - ( - pred_atom_coord[:, mask_chain_atom.bool()][:, None, :, :] - - pred_atom_coord[:, mask_chain_atom.bool()][:, :, None, :] - ) - ** 2 - ).sum(-1) ** 0.5 - if inference: - resolved_pair = 1 - ( - feats["atom_pad_mask"][i][mask_chain_atom.bool()][None, :] - * feats["atom_pad_mask"][i][mask_chain_atom.bool()][:, None] - ).to(torch.float32) - resolved_pair[resolved_pair == 1] = torch.inf - indices = torch.sort(dist_mat + resolved_pair, axis=2).indices - else: - if resolved_mask is None: - resolved_mask = feats["atom_resolved_mask"] - resolved_pair = 1 - ( - resolved_mask[i][mask_chain_atom.bool()][None, :] - * resolved_mask[i][mask_chain_atom.bool()][:, None] - ).to(torch.float32) - resolved_pair[resolved_pair == 1] = torch.inf - indices = torch.sort(dist_mat + resolved_pair, axis=2).indices - frames = ( - torch.cat( - [ - indices[:, :, 1:2], - indices[:, :, 0:1], - indices[:, :, 2:3], - ], - dim=2, - ) - + atom_idx - ) - try: - frames_idx_pred[i, :, token_idx : token_idx + num_atoms, :] = frames - except Exception as e: - print(f"Failed to process {feats['pdb_id']} due to {e}") - token_idx += num_tokens - atom_idx += num_atoms - - frames_expanded = pred_atom_coords[ - torch.arange(0, B // multiplicity, 1)[:, None, None, None].to( - frames_idx_pred.device - ), - torch.arange(0, multiplicity, 1)[None, :, None, None].to( - frames_idx_pred.device - ), - frames_idx_pred, - ].reshape(-1, 3, 3) - - # Compute masks for collinearity / overlap - mask_collinear_pred = compute_collinear_mask( - frames_expanded[:, 1] - frames_expanded[:, 0], - frames_expanded[:, 1] - frames_expanded[:, 2], - ).reshape(B // multiplicity, multiplicity, -1) - return frames_idx_pred, mask_collinear_pred * feats["token_pad_mask"][:, None, :] - - -def compute_aggregated_metric(logits, end=1.0): - # Compute aggregated metric from logits - num_bins = logits.shape[-1] - bin_width = end / num_bins - bounds = torch.arange( - start=0.5 * bin_width, end=end, step=bin_width, device=logits.device - ) - probs = nn.functional.softmax(logits, dim=-1) - plddt = torch.sum( - probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape), - dim=-1, - ) - return plddt - - -def tm_function(d, Nres): - d0 = 1.24 * (torch.clip(Nres, min=19) - 15) ** (1 / 3) - 1.8 - return 1 / (1 + (d / d0) ** 2) - - -def compute_ptms(logits, x_preds, feats, multiplicity): - # It needs to take as input the mask of the frames as they are not used to compute the PTM - _, mask_collinear_pred = compute_frame_pred( - x_preds, feats["frames_idx"], feats, multiplicity, inference=True - ) - # mask overlapping, collinear tokens and ions (invalid frames) - mask_pad = feats["token_pad_mask"].repeat_interleave(multiplicity, 0) - maski = mask_collinear_pred.reshape(-1, mask_collinear_pred.shape[-1]) - pair_mask_ptm = maski[:, :, None] * mask_pad[:, None, :] * mask_pad[:, :, None] - asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0) - pair_mask_iptm = ( - maski[:, :, None] - * (asym_id[:, None, :] != asym_id[:, :, None]) - * mask_pad[:, None, :] - * mask_pad[:, :, None] - ) - num_bins = logits.shape[-1] - bin_width = 32.0 / num_bins - end = 32.0 - pae_value = torch.arange( - start=0.5 * bin_width, end=end, step=bin_width, device=logits.device - ).unsqueeze(0) - N_res = mask_pad.sum(dim=-1, keepdim=True) - tm_value = tm_function(pae_value, N_res).unsqueeze(1).unsqueeze(2) - probs = nn.functional.softmax(logits, dim=-1) - tm_expected_value = torch.sum( - probs * tm_value, - dim=-1, - ) # shape (B, N, N) - ptm = torch.max( - torch.sum(tm_expected_value * pair_mask_ptm, dim=-1) - / (torch.sum(pair_mask_ptm, dim=-1) + 1e-5), - dim=1, - ).values - iptm = torch.max( - torch.sum(tm_expected_value * pair_mask_iptm, dim=-1) - / (torch.sum(pair_mask_iptm, dim=-1) + 1e-5), - dim=1, - ).values - - # compute ligand and protein iPTM - token_type = feats["mol_type"] - token_type = token_type.repeat_interleave(multiplicity, 0) - is_ligand_token = (token_type == const.chain_type_ids["NONPOLYMER"]).float() - is_protein_token = (token_type == const.chain_type_ids["PROTEIN"]).float() - - ligand_iptm_mask = ( - maski[:, :, None] - * (asym_id[:, None, :] != asym_id[:, :, None]) - * mask_pad[:, None, :] - * mask_pad[:, :, None] - * ( - (is_ligand_token[:, :, None] * is_protein_token[:, None, :]) - + (is_protein_token[:, :, None] * is_ligand_token[:, None, :]) - ) - ) - protein_ipmt_mask = ( - maski[:, :, None] - * (asym_id[:, None, :] != asym_id[:, :, None]) - * mask_pad[:, None, :] - * mask_pad[:, :, None] - * (is_protein_token[:, :, None] * is_protein_token[:, None, :]) - ) - - ligand_iptm = torch.max( - torch.sum(tm_expected_value * ligand_iptm_mask, dim=-1) - / (torch.sum(ligand_iptm_mask, dim=-1) + 1e-5), - dim=1, - ).values - protein_iptm = torch.max( - torch.sum(tm_expected_value * protein_ipmt_mask, dim=-1) - / (torch.sum(protein_ipmt_mask, dim=-1) + 1e-5), - dim=1, - ).values - - # Compute pair chain ipTM - chain_pair_iptm = {} - asym_ids_list = torch.unique(asym_id).tolist() - for idx1 in asym_ids_list: - chain_iptm = {} - for idx2 in asym_ids_list: - mask_pair_chain = ( - maski[:, :, None] - * (asym_id[:, None, :] == idx1) - * (asym_id[:, :, None] == idx2) - * mask_pad[:, None, :] - * mask_pad[:, :, None] - ) - - chain_iptm[idx2] = torch.max( - torch.sum(tm_expected_value * mask_pair_chain, dim=-1) - / (torch.sum(mask_pair_chain, dim=-1) + 1e-5), - dim=1, - ).values - chain_pair_iptm[idx1] = chain_iptm - - return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm diff --git a/fastplms/boltz/vb_layers_dropout.py b/fastplms/boltz/vb_layers_dropout.py deleted file mode 100644 index 588b87e..0000000 --- a/fastplms/boltz/vb_layers_dropout.py +++ /dev/null @@ -1,34 +0,0 @@ -import torch -from torch import Tensor - - -def get_dropout_mask( - dropout: float, - z: Tensor, - training: bool, - columnwise: bool = False, -) -> Tensor: - """Get the dropout mask. - - Parameters - ---------- - dropout : float - The dropout rate - z : torch.Tensor - The tensor to apply dropout to - training : bool - Whether the model is in training mode - columnwise : bool, optional - Whether to apply dropout columnwise - - Returns - ------- - torch.Tensor - The dropout mask - - """ - dropout = dropout * training - v = z[:, 0:1, :, 0:1] if columnwise else z[:, :, 0:1, 0:1] - d = torch.rand(v.shape, dtype=torch.float32, device=v.device) >= dropout - d = d * 1.0 / (1.0 - dropout) - return d diff --git a/fastplms/boltz/vb_layers_initialize.py b/fastplms/boltz/vb_layers_initialize.py deleted file mode 100644 index 9de084f..0000000 --- a/fastplms/boltz/vb_layers_initialize.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Utility functions for initializing weights and biases.""" - -# Copyright 2021 AlQuraishi Laboratory -# Copyright 2021 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math - -import torch - - -def _calculate_fan(linear_weight_shape, fan="fan_in"): - fan_out, fan_in = linear_weight_shape - - if fan == "fan_in": - f = fan_in - elif fan == "fan_out": - f = fan_out - elif fan == "fan_avg": - f = (fan_in + fan_out) / 2 - else: - raise ValueError("Invalid fan option") - - return f - - -def trunc_normal_init_(weights, scale=1.0, fan="fan_in"): - shape = weights.shape - f = _calculate_fan(shape, fan) - scale = scale / max(1, f) - std = math.sqrt(scale) - with torch.no_grad(): - torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std) - - -def lecun_normal_init_(weights): - trunc_normal_init_(weights, scale=1.0) - - -def he_normal_init_(weights): - trunc_normal_init_(weights, scale=2.0) - - -def glorot_uniform_init_(weights): - torch.nn.init.xavier_uniform_(weights, gain=1) - - -def final_init_(weights): - with torch.no_grad(): - weights.fill_(0.0) - - -def gating_init_(weights): - with torch.no_grad(): - weights.fill_(0.0) - - -def bias_init_zero_(bias): - with torch.no_grad(): - bias.fill_(0.0) - - -def bias_init_one_(bias): - with torch.no_grad(): - bias.fill_(1.0) - - -def normal_init_(weights): - torch.nn.init.kaiming_normal_(weights, nonlinearity="linear") - - -def ipa_point_weights_init_(weights): - with torch.no_grad(): - softplus_inverse_1 = 0.541324854612918 - weights.fill_(softplus_inverse_1) diff --git a/fastplms/boltz/vb_layers_outer_product_mean.py b/fastplms/boltz/vb_layers_outer_product_mean.py deleted file mode 100644 index 3445e4f..0000000 --- a/fastplms/boltz/vb_layers_outer_product_mean.py +++ /dev/null @@ -1,98 +0,0 @@ -import torch -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -class OuterProductMean(nn.Module): - """Outer product mean layer.""" - - def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None: - """Initialize the outer product mean layer. - - Parameters - ---------- - c_in : int - The input dimension. - c_hidden : int - The hidden dimension. - c_out : int - The output dimension. - - """ - super().__init__() - self.c_hidden = c_hidden - self.norm = nn.LayerNorm(c_in) - self.proj_a = nn.Linear(c_in, c_hidden, bias=False) - self.proj_b = nn.Linear(c_in, c_hidden, bias=False) - self.proj_o = nn.Linear(c_hidden * c_hidden, c_out) - init.final_init_(self.proj_o.weight) - init.final_init_(self.proj_o.bias) - - def forward(self, m: Tensor, mask: Tensor, chunk_size: int = None) -> Tensor: - """Forward pass. - - Parameters - ---------- - m : torch.Tensor - The sequence tensor (B, S, N, c_in). - mask : torch.Tensor - The mask tensor (B, S, N). - - Returns - ------- - torch.Tensor - The output tensor (B, N, N, c_out). - - """ - # Expand mask - mask = mask.unsqueeze(-1).to(m) - - # Compute projections - m = self.norm(m) - a = self.proj_a(m) * mask - b = self.proj_b(m) * mask - - # Compute outer product mean - if chunk_size is not None and not self.training: - # Compute pairwise mask - for i in range(0, mask.shape[1], 64): - if i == 0: - num_mask = ( - mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None] - ).sum(1) - else: - num_mask += ( - mask[:, i : i + 64, None, :] * mask[:, i : i + 64, :, None] - ).sum(1) - num_mask = num_mask.clamp(min=1) - - # Compute squentially in chunks - for i in range(0, self.c_hidden, chunk_size): - a_chunk = a[:, :, :, i : i + chunk_size] - sliced_weight_proj_o = self.proj_o.weight[ - :, i * self.c_hidden : (i + chunk_size) * self.c_hidden - ] - - z = torch.einsum("bsic,bsjd->bijcd", a_chunk, b) - z = z.reshape(*z.shape[:3], -1) - z = z / num_mask - - # Project to output - if i == 0: - z_out = z.to(m) @ sliced_weight_proj_o.T - else: - z_out = z_out + z.to(m) @ sliced_weight_proj_o.T - - z_out = z_out + self.proj_o.bias # add bias - return z_out - else: - mask = mask[:, :, None, :] * mask[:, :, :, None] - num_mask = mask.sum(1).clamp(min=1) - z = torch.einsum("bsic,bsjd->bijcd", a.float(), b.float()) - z = z.reshape(*z.shape[:3], -1) - z = z / num_mask - - # Project to output - z = self.proj_o(z.to(m)) - return z diff --git a/fastplms/boltz/vb_layers_pair_averaging.py b/fastplms/boltz/vb_layers_pair_averaging.py deleted file mode 100644 index a58ea69..0000000 --- a/fastplms/boltz/vb_layers_pair_averaging.py +++ /dev/null @@ -1,135 +0,0 @@ -import torch -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -class PairWeightedAveraging(nn.Module): - """Pair weighted averaging layer.""" - - def __init__( - self, - c_m: int, - c_z: int, - c_h: int, - num_heads: int, - inf: float = 1e6, - ) -> None: - """Initialize the pair weighted averaging layer. - - Parameters - ---------- - c_m: int - The dimension of the input sequence. - c_z: int - The dimension of the input pairwise tensor. - c_h: int - The dimension of the hidden. - num_heads: int - The number of heads. - inf: float - The value to use for masking, default 1e6. - - """ - super().__init__() - self.c_m = c_m - self.c_z = c_z - self.c_h = c_h - self.num_heads = num_heads - self.inf = inf - - self.norm_m = nn.LayerNorm(c_m) - self.norm_z = nn.LayerNorm(c_z) - - self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False) - self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False) - self.proj_z = nn.Linear(c_z, num_heads, bias=False) - self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False) - init.final_init_(self.proj_o.weight) - - def forward( - self, m: Tensor, z: Tensor, mask: Tensor, chunk_heads: False = bool - ) -> Tensor: - """Forward pass. - - Parameters - ---------- - m : torch.Tensor - The input sequence tensor (B, S, N, D) - z : torch.Tensor - The input pairwise tensor (B, N, N, D) - mask : torch.Tensor - The pairwise mask tensor (B, N, N) - - Returns - ------- - torch.Tensor - The output sequence tensor (B, S, N, D) - - """ - # Compute layer norms - m = self.norm_m(m) - z = self.norm_z(z) - - if chunk_heads and not self.training: - # Compute heads sequentially - o_chunks = [] - for head_idx in range(self.num_heads): - sliced_weight_proj_m = self.proj_m.weight[ - head_idx * self.c_h : (head_idx + 1) * self.c_h, : - ] - sliced_weight_proj_g = self.proj_g.weight[ - head_idx * self.c_h : (head_idx + 1) * self.c_h, : - ] - sliced_weight_proj_z = self.proj_z.weight[head_idx : (head_idx + 1), :] - sliced_weight_proj_o = self.proj_o.weight[ - :, head_idx * self.c_h : (head_idx + 1) * self.c_h - ] - - # Project input tensors - v: Tensor = m @ sliced_weight_proj_m.T - v = v.reshape(*v.shape[:3], 1, self.c_h) - v = v.permute(0, 3, 1, 2, 4) - - # Compute weights - b: Tensor = z @ sliced_weight_proj_z.T - b = b.permute(0, 3, 1, 2) - b = b + (1 - mask[:, None]) * -self.inf - w = torch.softmax(b, dim=-1) - - # Compute gating - g: Tensor = m @ sliced_weight_proj_g.T - g = g.sigmoid() - - # Compute output - o = torch.einsum("bhij,bhsjd->bhsid", w, v) - o = o.permute(0, 2, 3, 1, 4) - o = o.reshape(*o.shape[:3], 1 * self.c_h) - o_chunks = g * o - if head_idx == 0: - o_out = o_chunks @ sliced_weight_proj_o.T - else: - o_out += o_chunks @ sliced_weight_proj_o.T - return o_out - else: - # Project input tensors - v: Tensor = self.proj_m(m) - v = v.reshape(*v.shape[:3], self.num_heads, self.c_h) - v = v.permute(0, 3, 1, 2, 4) - - # Compute weights - b: Tensor = self.proj_z(z) - b = b.permute(0, 3, 1, 2) - b = b + (1 - mask[:, None]) * -self.inf - w = torch.softmax(b, dim=-1) - - # Compute gating - g: Tensor = self.proj_g(m) - g = g.sigmoid() - - # Compute output - o = torch.einsum("bhij,bhsjd->bhsid", w, v) - o = o.permute(0, 2, 3, 1, 4) - o = o.reshape(*o.shape[:3], self.num_heads * self.c_h) - o = self.proj_o(g * o) - return o diff --git a/fastplms/boltz/vb_layers_pairformer.py b/fastplms/boltz/vb_layers_pairformer.py deleted file mode 100644 index 08911e9..0000000 --- a/fastplms/boltz/vb_layers_pairformer.py +++ /dev/null @@ -1,337 +0,0 @@ -from typing import Optional, Tuple - -import torch -from torch import Tensor, nn - -from . import vb_const as const -from .vb_layers_attention import AttentionPairBias -from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2 -from .vb_layers_dropout import get_dropout_mask -from .vb_layers_transition import Transition -from .vb_tri_attn_attention import ( - TriangleAttentionEndingNode, - TriangleAttentionStartingNode, -) -from .vb_layers_triangular_mult import ( - TriangleMultiplicationIncoming, - TriangleMultiplicationOutgoing, -) - - -class PairformerLayer(nn.Module): - """Pairformer module.""" - - def __init__( - self, - token_s: int, - token_z: int, - num_heads: int = 16, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - v2: bool = False, - ) -> None: - super().__init__() - self.token_z = token_z - self.dropout = dropout - self.num_heads = num_heads - self.post_layer_norm = post_layer_norm - - self.pre_norm_s = nn.LayerNorm(token_s) - if v2: - self.attention = AttentionPairBiasV2(token_s, token_z, num_heads) - else: - self.attention = AttentionPairBias(token_s, token_z, num_heads) - - self.tri_mul_out = TriangleMultiplicationOutgoing(token_z) - self.tri_mul_in = TriangleMultiplicationIncoming(token_z) - - self.tri_att_start = TriangleAttentionStartingNode( - token_z, pairwise_head_width, pairwise_num_heads, inf=1e9 - ) - self.tri_att_end = TriangleAttentionEndingNode( - token_z, pairwise_head_width, pairwise_num_heads, inf=1e9 - ) - - self.transition_s = Transition(token_s, token_s * 4) - self.transition_z = Transition(token_z, token_z * 4) - - self.s_post_norm = ( - nn.LayerNorm(token_s) if self.post_layer_norm else nn.Identity() - ) - - def forward( - self, - s: Tensor, - z: Tensor, - mask: Tensor, - pair_mask: Tensor, - chunk_size_tri_attn: Optional[int] = None, - use_kernels: bool = False, - use_cuequiv_mul: bool = False, - use_cuequiv_attn: bool = False, - ) -> Tuple[Tensor, Tensor]: - # Compute pairwise stack - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_mul_out( - z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels - ) - - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_mul_in( - z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels - ) - - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_att_start( - z, - mask=pair_mask, - chunk_size=chunk_size_tri_attn, - use_kernels=use_cuequiv_attn or use_kernels, - ) - - dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True) - z = z + dropout * self.tri_att_end( - z, - mask=pair_mask, - chunk_size=chunk_size_tri_attn, - use_kernels=use_cuequiv_attn or use_kernels, - ) - - z = z + self.transition_z(z) - - # Compute sequence stack - with torch.autocast("cuda", enabled=False): - s_normed = self.pre_norm_s(s.float()) - s = s.float() + self.attention( - s=s_normed, z=z.float(), mask=mask.float(), k_in=s_normed - ) - s = s + self.transition_s(s) - s = self.s_post_norm(s) - - return s, z - - -class PairformerModule(nn.Module): - """Pairformer module.""" - - def __init__( - self, - token_s: int, - token_z: int, - num_blocks: int, - num_heads: int = 16, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - activation_checkpointing: bool = False, - v2: bool = False, - **kwargs, - ) -> None: - super().__init__() - self.token_z = token_z - self.num_blocks = num_blocks - self.dropout = dropout - self.num_heads = num_heads - self.post_layer_norm = post_layer_norm - self.activation_checkpointing = activation_checkpointing - - self.layers = nn.ModuleList() - for _ in range(num_blocks): - self.layers.append( - PairformerLayer( - token_s, - token_z, - num_heads, - dropout, - pairwise_head_width, - pairwise_num_heads, - post_layer_norm, - v2, - ), - ) - - def forward( - self, - s: Tensor, - z: Tensor, - mask: Tensor, - pair_mask: Tensor, - use_kernels: bool = False, - ) -> Tuple[Tensor, Tensor]: - """Perform the forward pass. - - Parameters - ---------- - s : Tensor - The sequence stack. - z : Tensor - The pairwise stack. - mask : Tensor - The mask. - pair_mask : Tensor - The pairwise mask. - use_kernels : bool - Whether to use kernels. - - """ - if not self.training: - if z.shape[1] > const.chunk_size_threshold: - chunk_size_tri_attn = 128 - else: - chunk_size_tri_attn = 512 - else: - chunk_size_tri_attn = None - - for layer in self.layers: - if self.activation_checkpointing: - s, z = torch.utils.checkpoint.checkpoint( - layer, - s, - z, - mask, - pair_mask, - chunk_size_tri_attn, - use_kernels, - use_reentrant=False, - ) - else: - s, z = layer(s, z, mask, pair_mask, chunk_size_tri_attn, use_kernels) - return s, z - - -class PairformerNoSeqLayer(nn.Module): - """Pairformer module without sequence track.""" - - def __init__( - self, - token_z: int, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - ) -> None: - super().__init__() - self.token_z = token_z - self.dropout = dropout - self.post_layer_norm = post_layer_norm - - self.tri_mul_out = TriangleMultiplicationOutgoing(token_z) - self.tri_mul_in = TriangleMultiplicationIncoming(token_z) - - self.tri_att_start = TriangleAttentionStartingNode( - token_z, pairwise_head_width, pairwise_num_heads, inf=1e9 - ) - self.tri_att_end = TriangleAttentionEndingNode( - token_z, pairwise_head_width, pairwise_num_heads, inf=1e9 - ) - - self.transition_z = Transition(token_z, token_z * 4) - - def forward( - self, - z: Tensor, - pair_mask: Tensor, - chunk_size_tri_attn: Optional[int] = None, - use_kernels: bool = False, - use_cuequiv_mul: bool = False, - use_cuequiv_attn: bool = False, - ) -> Tensor: - # Compute pairwise stack - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_mul_out( - z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels - ) - - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_mul_in( - z, mask=pair_mask, use_kernels=use_cuequiv_mul or use_kernels - ) - - dropout = get_dropout_mask(self.dropout, z, self.training) - z = z + dropout * self.tri_att_start( - z, - mask=pair_mask, - chunk_size=chunk_size_tri_attn, - use_kernels=use_cuequiv_attn or use_kernels, - ) - - dropout = get_dropout_mask(self.dropout, z, self.training, columnwise=True) - z = z + dropout * self.tri_att_end( - z, - mask=pair_mask, - chunk_size=chunk_size_tri_attn, - use_kernels=use_cuequiv_attn or use_kernels, - ) - - z = z + self.transition_z(z) - return z - - -class PairformerNoSeqModule(nn.Module): - """Pairformer module without sequence track.""" - - def __init__( - self, - token_z: int, - num_blocks: int, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - activation_checkpointing: bool = False, - **kwargs, - ) -> None: - super().__init__() - self.token_z = token_z - self.num_blocks = num_blocks - self.dropout = dropout - self.post_layer_norm = post_layer_norm - self.activation_checkpointing = activation_checkpointing - - self.layers = nn.ModuleList() - for i in range(num_blocks): - self.layers.append( - PairformerNoSeqLayer( - token_z, - dropout, - pairwise_head_width, - pairwise_num_heads, - post_layer_norm, - ), - ) - - def forward( - self, - z: Tensor, - pair_mask: Tensor, - use_kernels: bool = False, - ) -> Tensor: - if not self.training: - if z.shape[1] > const.chunk_size_threshold: - chunk_size_tri_attn = 128 - else: - chunk_size_tri_attn = 512 - else: - chunk_size_tri_attn = None - - for layer in self.layers: - if self.activation_checkpointing: - z = torch.utils.checkpoint.checkpoint( - layer, - z, - pair_mask, - chunk_size_tri_attn, - use_kernels, - use_reentrant=False, - ) - else: - z = layer( - z, - pair_mask, - chunk_size_tri_attn, - use_kernels, - ) - return z diff --git a/fastplms/boltz/vb_layers_transition.py b/fastplms/boltz/vb_layers_transition.py deleted file mode 100644 index f0bb0a2..0000000 --- a/fastplms/boltz/vb_layers_transition.py +++ /dev/null @@ -1,78 +0,0 @@ -from typing import Optional - -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -class Transition(nn.Module): - """Perform a two-layer MLP.""" - - def __init__( - self, - dim: int = 128, - hidden: int = 512, - out_dim: Optional[int] = None, - ) -> None: - """Initialize the TransitionUpdate module. - - Parameters - ---------- - dim: int - The dimension of the input, default 128 - hidden: int - The dimension of the hidden, default 512 - out_dim: Optional[int] - The dimension of the output, default None - - """ - super().__init__() - if out_dim is None: - out_dim = dim - - self.norm = nn.LayerNorm(dim, eps=1e-5) - self.fc1 = nn.Linear(dim, hidden, bias=False) - self.fc2 = nn.Linear(dim, hidden, bias=False) - self.fc3 = nn.Linear(hidden, out_dim, bias=False) - self.silu = nn.SiLU() - self.hidden = hidden - - init.bias_init_one_(self.norm.weight) - init.bias_init_zero_(self.norm.bias) - - init.lecun_normal_init_(self.fc1.weight) - init.lecun_normal_init_(self.fc2.weight) - init.final_init_(self.fc3.weight) - - def forward(self, x: Tensor, chunk_size: int = None) -> Tensor: - """Perform a forward pass. - - Parameters - ---------- - x: torch.Tensor - The input data of shape (..., D) - - Returns - ------- - x: torch.Tensor - The output data of shape (..., D) - - """ - x = self.norm(x) - - if chunk_size is None or self.training: - x = self.silu(self.fc1(x)) * self.fc2(x) - x = self.fc3(x) - return x - else: - # Compute in chunks - for i in range(0, self.hidden, chunk_size): - fc1_slice = self.fc1.weight[i : i + chunk_size, :] - fc2_slice = self.fc2.weight[i : i + chunk_size, :] - fc3_slice = self.fc3.weight[:, i : i + chunk_size] - x_chunk = self.silu((x @ fc1_slice.T)) * (x @ fc2_slice.T) - if i == 0: - x_out = x_chunk @ fc3_slice.T - else: - x_out = x_out + x_chunk @ fc3_slice.T - return x_out diff --git a/fastplms/boltz/vb_layers_triangular_mult.py b/fastplms/boltz/vb_layers_triangular_mult.py deleted file mode 100644 index 4865dcb..0000000 --- a/fastplms/boltz/vb_layers_triangular_mult.py +++ /dev/null @@ -1,215 +0,0 @@ -import importlib - -import torch -from torch import Tensor, nn - -from . import vb_layers_initialize as init - - -@torch.compiler.disable -def kernel_triangular_mult( - x, - direction, - mask, - norm_in_weight, - norm_in_bias, - p_in_weight, - g_in_weight, - norm_out_weight, - norm_out_bias, - p_out_weight, - g_out_weight, - eps, -): - triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle") - triangle_multiplicative_update = triangle_module.triangle_multiplicative_update - return triangle_multiplicative_update( - x, - direction=direction, - mask=mask, - norm_in_weight=norm_in_weight, - norm_in_bias=norm_in_bias, - p_in_weight=p_in_weight, - g_in_weight=g_in_weight, - norm_out_weight=norm_out_weight, - norm_out_bias=norm_out_bias, - p_out_weight=p_out_weight, - g_out_weight=g_out_weight, - eps=eps, - ) - - -class TriangleMultiplicationOutgoing(nn.Module): - """TriangleMultiplicationOutgoing.""" - - def __init__(self, dim: int = 128) -> None: - """Initialize the TriangularUpdate module. - - Parameters - ---------- - dim: int - The dimension of the input, default 128 - - """ - super().__init__() - - self.norm_in = nn.LayerNorm(dim, eps=1e-5) - self.p_in = nn.Linear(dim, 2 * dim, bias=False) - self.g_in = nn.Linear(dim, 2 * dim, bias=False) - - self.norm_out = nn.LayerNorm(dim) - self.p_out = nn.Linear(dim, dim, bias=False) - self.g_out = nn.Linear(dim, dim, bias=False) - - init.bias_init_one_(self.norm_in.weight) - init.bias_init_zero_(self.norm_in.bias) - - init.lecun_normal_init_(self.p_in.weight) - init.gating_init_(self.g_in.weight) - - init.bias_init_one_(self.norm_out.weight) - init.bias_init_zero_(self.norm_out.bias) - - init.final_init_(self.p_out.weight) - init.gating_init_(self.g_out.weight) - - def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor: - """Perform a forward pass. - - Parameters - ---------- - x: torch.Tensor - The input data of shape (B, N, N, D) - mask: torch.Tensor - The input mask of shape (B, N, N) - use_kernels: bool - Whether to use the kernel - - Returns - ------- - x: torch.Tensor - The output data of shape (B, N, N, D) - - """ - if use_kernels: - return kernel_triangular_mult( - x, - direction="outgoing", - mask=mask, - norm_in_weight=self.norm_in.weight, - norm_in_bias=self.norm_in.bias, - p_in_weight=self.p_in.weight, - g_in_weight=self.g_in.weight, - norm_out_weight=self.norm_out.weight, - norm_out_bias=self.norm_out.bias, - p_out_weight=self.p_out.weight, - g_out_weight=self.g_out.weight, - eps=1e-5, - ) - - # Input gating: D -> D - x = self.norm_in(x) - x_in = x - x = self.p_in(x) * self.g_in(x).sigmoid() - - # Apply mask - x = x * mask.unsqueeze(-1) - - # Split input and cast to float - a, b = torch.chunk(x.float(), 2, dim=-1) - - # Triangular projection - x = torch.einsum("bikd,bjkd->bijd", a, b) - - # Output gating - x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid() - - return x - - -class TriangleMultiplicationIncoming(nn.Module): - """TriangleMultiplicationIncoming.""" - - def __init__(self, dim: int = 128) -> None: - """Initialize the TriangularUpdate module. - - Parameters - ---------- - dim: int - The dimension of the input, default 128 - - """ - super().__init__() - - self.norm_in = nn.LayerNorm(dim, eps=1e-5) - self.p_in = nn.Linear(dim, 2 * dim, bias=False) - self.g_in = nn.Linear(dim, 2 * dim, bias=False) - - self.norm_out = nn.LayerNorm(dim) - self.p_out = nn.Linear(dim, dim, bias=False) - self.g_out = nn.Linear(dim, dim, bias=False) - - init.bias_init_one_(self.norm_in.weight) - init.bias_init_zero_(self.norm_in.bias) - - init.lecun_normal_init_(self.p_in.weight) - init.gating_init_(self.g_in.weight) - - init.bias_init_one_(self.norm_out.weight) - init.bias_init_zero_(self.norm_out.bias) - - init.final_init_(self.p_out.weight) - init.gating_init_(self.g_out.weight) - - def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor: - """Perform a forward pass. - - Parameters - ---------- - x: torch.Tensor - The input data of shape (B, N, N, D) - mask: torch.Tensor - The input mask of shape (B, N, N) - use_kernels: bool - Whether to use the kernel - - Returns - ------- - x: torch.Tensor - The output data of shape (B, N, N, D) - - """ - if use_kernels: - return kernel_triangular_mult( - x, - direction="incoming", - mask=mask, - norm_in_weight=self.norm_in.weight, - norm_in_bias=self.norm_in.bias, - p_in_weight=self.p_in.weight, - g_in_weight=self.g_in.weight, - norm_out_weight=self.norm_out.weight, - norm_out_bias=self.norm_out.bias, - p_out_weight=self.p_out.weight, - g_out_weight=self.g_out.weight, - eps=1e-5, - ) - - # Input gating: D -> D - x = self.norm_in(x) - x_in = x - x = self.p_in(x) * self.g_in(x).sigmoid() - - # Apply mask - x = x * mask.unsqueeze(-1) - - # Split input and cast to float - a, b = torch.chunk(x.float(), 2, dim=-1) - - # Triangular projection - x = torch.einsum("bkid,bkjd->bijd", a, b) - - # Output gating - x = self.p_out(self.norm_out(x)) * self.g_out(x_in).sigmoid() - - return x diff --git a/fastplms/boltz/vb_loss_diffusionv2.py b/fastplms/boltz/vb_loss_diffusionv2.py deleted file mode 100644 index 54980db..0000000 --- a/fastplms/boltz/vb_loss_diffusionv2.py +++ /dev/null @@ -1,138 +0,0 @@ -# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang - -import torch -import torch.nn.functional as F -from einops import einsum, rearrange - - -def weighted_rigid_align( - true_coords, # Float['b n 3'], # true coordinates - pred_coords, # Float['b n 3'], # predicted coordinates - weights, # Float['b n'], # weights for each atom - mask, # Bool['b n'] | None = None # mask for variable lengths -): # -> Float['b n 3']: - """Algorithm 28 : note there is a problem with the pseudocode in the paper where predicted and - GT are swapped in algorithm 28, but correct in equation (2).""" - - out_shape = torch.broadcast_shapes(true_coords.shape, pred_coords.shape) - *batch_size, num_points, dim = out_shape - weights = (mask * weights).unsqueeze(-1) - - # Compute weighted centroids - true_centroid = (true_coords * weights).sum(dim=-2, keepdim=True) / weights.sum( - dim=-2, keepdim=True - ) - pred_centroid = (pred_coords * weights).sum(dim=-2, keepdim=True) / weights.sum( - dim=-2, keepdim=True - ) - - # Center the coordinates - true_coords_centered = true_coords - true_centroid - pred_coords_centered = pred_coords - pred_centroid - - if torch.any(mask.sum(dim=-1) < (dim + 1)): - print( - "Warning: The size of one of the point clouds is <= dim+1. " - + "`WeightedRigidAlign` cannot return a unique rotation." - ) - - # Compute the weighted covariance matrix - cov_matrix = einsum( - weights * pred_coords_centered, - true_coords_centered, - "... n i, ... n j -> ... i j", - ) - - # Compute the SVD of the covariance matrix, required float32 for svd and determinant - original_dtype = cov_matrix.dtype - cov_matrix_32 = cov_matrix.to(dtype=torch.float32) - - U, S, V = torch.linalg.svd( - cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None - ) - V = V.mH - - # Catch ambiguous rotation by checking the magnitude of singular values - if (S.abs() <= 1e-15).any() and not (num_points < (dim + 1)): - print( - "Warning: Excessively low rank of " - + "cross-correlation between aligned point clouds. " - + "`WeightedRigidAlign` cannot return a unique rotation." - ) - - # Compute the rotation matrix - rot_matrix = torch.einsum("... i j, ... k j -> ... i k", U, V).to( - dtype=torch.float32 - ) - - # Ensure proper rotation matrix with determinant 1 - F = torch.eye(dim, dtype=cov_matrix_32.dtype, device=cov_matrix.device)[ - None - ].repeat(*batch_size, 1, 1) - F[..., -1, -1] = torch.det(rot_matrix) - rot_matrix = einsum(U, F, V, "... i j, ... j k, ... l k -> ... i l") - rot_matrix = rot_matrix.to(dtype=original_dtype) - - # Apply the rotation and translation - aligned_coords = ( - einsum(true_coords_centered, rot_matrix, "... n i, ... j i -> ... n j") - + pred_centroid - ) - aligned_coords.detach_() - - return aligned_coords - - -def smooth_lddt_loss( - pred_coords, # Float['b n 3'], - true_coords, # Float['b n 3'], - is_nucleotide, # Bool['b n'], - coords_mask, # Bool['b n'] | None = None, - nucleic_acid_cutoff: float = 30.0, - other_cutoff: float = 15.0, - multiplicity: int = 1, -): # -> Float['']: - """Algorithm 27 - pred_coords: predicted coordinates - true_coords: true coordinates - Note: for efficiency pred_coords is the only one with the multiplicity expanded - TODO: add weighing which overweight the smooth lddt contribution close to t=0 (not present in the paper) - """ - lddt = [] - for i in range(true_coords.shape[0]): - true_dists = torch.cdist(true_coords[i], true_coords[i]) - - is_nucleotide_i = is_nucleotide[i // multiplicity] - coords_mask_i = coords_mask[i // multiplicity] - - is_nucleotide_pair = is_nucleotide_i.unsqueeze(-1).expand( - -1, is_nucleotide_i.shape[-1] - ) - - mask = is_nucleotide_pair * (true_dists < nucleic_acid_cutoff).float() - mask += (1 - is_nucleotide_pair) * (true_dists < other_cutoff).float() - mask *= 1 - torch.eye(pred_coords.shape[1], device=pred_coords.device) - mask *= coords_mask_i.unsqueeze(-1) - mask *= coords_mask_i.unsqueeze(-2) - - valid_pairs = mask.nonzero() - true_dists_i = true_dists[valid_pairs[:, 0], valid_pairs[:, 1]] - - pred_coords_i1 = pred_coords[i, valid_pairs[:, 0]] - pred_coords_i2 = pred_coords[i, valid_pairs[:, 1]] - pred_dists_i = F.pairwise_distance(pred_coords_i1, pred_coords_i2) - - dist_diff_i = torch.abs(true_dists_i - pred_dists_i) - - eps_i = ( - F.sigmoid(0.5 - dist_diff_i) - + F.sigmoid(1.0 - dist_diff_i) - + F.sigmoid(2.0 - dist_diff_i) - + F.sigmoid(4.0 - dist_diff_i) - ) / 4.0 - - lddt_i = eps_i.sum() / (valid_pairs.shape[0] + 1e-5) - lddt.append(lddt_i) - - # average over batch & multiplicity - return 1.0 - torch.stack(lddt, dim=0).mean(dim=0) diff --git a/fastplms/boltz/vb_modules_diffusion_conditioning.py b/fastplms/boltz/vb_modules_diffusion_conditioning.py deleted file mode 100644 index 83882fd..0000000 --- a/fastplms/boltz/vb_modules_diffusion_conditioning.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations - -import torch -from torch import nn -from torch.nn import Module - -from .vb_modules_encodersv2 import ( - AtomEncoder, - PairwiseConditioning, -) - - -class DiffusionConditioning(Module): - def __init__( - self, - token_s: int, - token_z: int, - atom_s: int, - atom_z: int, - atoms_per_window_queries: int = 32, - atoms_per_window_keys: int = 128, - atom_encoder_depth: int = 3, - atom_encoder_heads: int = 4, - token_transformer_depth: int = 24, - token_transformer_heads: int = 8, - atom_decoder_depth: int = 3, - atom_decoder_heads: int = 4, - atom_feature_dim: int = 128, - conditioning_transition_layers: int = 2, - use_no_atom_char: bool = False, - use_atom_backbone_feat: bool = False, - use_residue_feats_atoms: bool = False, - ) -> None: - super().__init__() - - self.pairwise_conditioner = PairwiseConditioning( - token_z=token_z, - dim_token_rel_pos_feats=token_z, - num_transitions=conditioning_transition_layers, - ) - - self.atom_encoder = AtomEncoder( - atom_s=atom_s, - atom_z=atom_z, - token_s=token_s, - token_z=token_z, - atoms_per_window_queries=atoms_per_window_queries, - atoms_per_window_keys=atoms_per_window_keys, - atom_feature_dim=atom_feature_dim, - structure_prediction=True, - use_no_atom_char=use_no_atom_char, - use_atom_backbone_feat=use_atom_backbone_feat, - use_residue_feats_atoms=use_residue_feats_atoms, - ) - - self.atom_enc_proj_z = nn.ModuleList() - for _ in range(atom_encoder_depth): - self.atom_enc_proj_z.append( - nn.Sequential( - nn.LayerNorm(atom_z), - nn.Linear(atom_z, atom_encoder_heads, bias=False), - ) - ) - - self.atom_dec_proj_z = nn.ModuleList() - for _ in range(atom_decoder_depth): - self.atom_dec_proj_z.append( - nn.Sequential( - nn.LayerNorm(atom_z), - nn.Linear(atom_z, atom_decoder_heads, bias=False), - ) - ) - - self.token_trans_proj_z = nn.ModuleList() - for _ in range(token_transformer_depth): - self.token_trans_proj_z.append( - nn.Sequential( - nn.LayerNorm(token_z), - nn.Linear(token_z, token_transformer_heads, bias=False), - ) - ) - - def forward( - self, - s_trunk, # Float['b n ts'] - z_trunk, # Float['b n n tz'] - relative_position_encoding, # Float['b n n tz'] - feats, - ): - z = self.pairwise_conditioner( - z_trunk, - relative_position_encoding, - ) - - q, c, p, to_keys = self.atom_encoder( - feats=feats, - s_trunk=s_trunk, # Float['b n ts'], - z=z, # Float['b n n tz'], - ) - - atom_enc_bias = [] - for layer in self.atom_enc_proj_z: - atom_enc_bias.append(layer(p)) - atom_enc_bias = torch.cat(atom_enc_bias, dim=-1) - - atom_dec_bias = [] - for layer in self.atom_dec_proj_z: - atom_dec_bias.append(layer(p)) - atom_dec_bias = torch.cat(atom_dec_bias, dim=-1) - - token_trans_bias = [] - for layer in self.token_trans_proj_z: - token_trans_bias.append(layer(z)) - token_trans_bias = torch.cat(token_trans_bias, dim=-1) - - return q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias diff --git a/fastplms/boltz/vb_modules_encodersv2.py b/fastplms/boltz/vb_modules_encodersv2.py deleted file mode 100644 index e6aa39f..0000000 --- a/fastplms/boltz/vb_modules_encodersv2.py +++ /dev/null @@ -1,565 +0,0 @@ -# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang -from functools import partial -from math import pi - -import torch -from einops import rearrange -from torch import nn -from torch.nn import Linear, Module, ModuleList -from torch.nn.functional import one_hot - -from . import vb_layers_initialize as init -from .vb_layers_transition import Transition -from .vb_modules_transformersv2 import AtomTransformer -from .vb_modules_utils import LinearNoBias - - -class FourierEmbedding(Module): - """Algorithm 22.""" - - def __init__(self, dim): - super().__init__() - self.proj = nn.Linear(1, dim) - torch.nn.init.normal_(self.proj.weight, mean=0, std=1) - torch.nn.init.normal_(self.proj.bias, mean=0, std=1) - self.proj.requires_grad_(False) - - def forward( - self, - times, # Float[' b'], - ): # -> Float['b d']: - times = rearrange(times, "b -> b 1") - rand_proj = self.proj(times) - return torch.cos(2 * pi * rand_proj) - - -class RelativePositionEncoder(Module): - """Algorithm 3.""" - - def __init__( - self, token_z, r_max=32, s_max=2, fix_sym_check=False, cyclic_pos_enc=False - ): - super().__init__() - self.r_max = r_max - self.s_max = s_max - self.linear_layer = LinearNoBias(4 * (r_max + 1) + 2 * (s_max + 1) + 1, token_z) - self.fix_sym_check = fix_sym_check - self.cyclic_pos_enc = cyclic_pos_enc - - def forward(self, feats): - b_same_chain = torch.eq( - feats["asym_id"][:, :, None], feats["asym_id"][:, None, :] - ) - b_same_residue = torch.eq( - feats["residue_index"][:, :, None], feats["residue_index"][:, None, :] - ) - b_same_entity = torch.eq( - feats["entity_id"][:, :, None], feats["entity_id"][:, None, :] - ) - - d_residue = ( - feats["residue_index"][:, :, None] - feats["residue_index"][:, None, :] - ) - - if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0): - period = torch.where( - feats["cyclic_period"] > 0, - feats["cyclic_period"], - torch.zeros_like(feats["cyclic_period"]) + 10000, - ) - d_residue = (d_residue - period * torch.round(d_residue / period)).long() - - d_residue = torch.clip( - d_residue + self.r_max, - 0, - 2 * self.r_max, - ) - d_residue = torch.where( - b_same_chain, d_residue, torch.zeros_like(d_residue) + 2 * self.r_max + 1 - ) - a_rel_pos = one_hot(d_residue, 2 * self.r_max + 2) - - d_token = torch.clip( - feats["token_index"][:, :, None] - - feats["token_index"][:, None, :] - + self.r_max, - 0, - 2 * self.r_max, - ) - d_token = torch.where( - b_same_chain & b_same_residue, - d_token, - torch.zeros_like(d_token) + 2 * self.r_max + 1, - ) - a_rel_token = one_hot(d_token, 2 * self.r_max + 2) - - d_chain = torch.clip( - feats["sym_id"][:, :, None] - feats["sym_id"][:, None, :] + self.s_max, - 0, - 2 * self.s_max, - ) - d_chain = torch.where( - (~b_same_entity) if self.fix_sym_check else b_same_chain, - torch.zeros_like(d_chain) + 2 * self.s_max + 1, - d_chain, - ) - # Note: added | (~b_same_entity) based on observation of ProteinX manuscript - a_rel_chain = one_hot(d_chain, 2 * self.s_max + 2) - - p = self.linear_layer( - torch.cat( - [ - a_rel_pos.float(), - a_rel_token.float(), - b_same_entity.unsqueeze(-1).float(), - a_rel_chain.float(), - ], - dim=-1, - ) - ) - return p - - -class SingleConditioning(Module): - """Algorithm 21.""" - - def __init__( - self, - sigma_data: float, - token_s: int = 384, - dim_fourier: int = 256, - num_transitions: int = 2, - transition_expansion_factor: int = 2, - eps: float = 1e-20, - disable_times: bool = False, - ) -> None: - super().__init__() - self.eps = eps - self.sigma_data = sigma_data - self.disable_times = disable_times - - self.norm_single = nn.LayerNorm(2 * token_s) - self.single_embed = nn.Linear(2 * token_s, 2 * token_s) - if not self.disable_times: - self.fourier_embed = FourierEmbedding(dim_fourier) - self.norm_fourier = nn.LayerNorm(dim_fourier) - self.fourier_to_single = LinearNoBias(dim_fourier, 2 * token_s) - - transitions = ModuleList([]) - for _ in range(num_transitions): - transition = Transition( - dim=2 * token_s, hidden=transition_expansion_factor * 2 * token_s - ) - transitions.append(transition) - - self.transitions = transitions - - def forward( - self, - times, # Float[' b'], - s_trunk, # Float['b n ts'], - s_inputs, # Float['b n ts'], - ): # -> Float['b n 2ts']: - s = torch.cat((s_trunk, s_inputs), dim=-1) - s = self.single_embed(self.norm_single(s)) - if not self.disable_times: - fourier_embed = self.fourier_embed( - times - ) # note: sigma rescaling done in diffusion module - normed_fourier = self.norm_fourier(fourier_embed) - fourier_to_single = self.fourier_to_single(normed_fourier) - - s = rearrange(fourier_to_single, "b d -> b 1 d") + s - - for transition in self.transitions: - s = transition(s) + s - - return s, normed_fourier if not self.disable_times else None - - -class PairwiseConditioning(Module): - """Algorithm 21.""" - - def __init__( - self, - token_z, - dim_token_rel_pos_feats, - num_transitions=2, - transition_expansion_factor=2, - ): - super().__init__() - - self.dim_pairwise_init_proj = nn.Sequential( - nn.LayerNorm(token_z + dim_token_rel_pos_feats), - LinearNoBias(token_z + dim_token_rel_pos_feats, token_z), - ) - - transitions = ModuleList([]) - for _ in range(num_transitions): - transition = Transition( - dim=token_z, hidden=transition_expansion_factor * token_z - ) - transitions.append(transition) - - self.transitions = transitions - - def forward( - self, - z_trunk, # Float['b n n tz'], - token_rel_pos_feats, # Float['b n n 3'], - ): # -> Float['b n n tz']: - z = torch.cat((z_trunk, token_rel_pos_feats), dim=-1) - z = self.dim_pairwise_init_proj(z) - - for transition in self.transitions: - z = transition(z) + z - - return z - - -def get_indexing_matrix(K, W, H, device): - assert W % 2 == 0 - assert H % (W // 2) == 0 - - h = H // (W // 2) - assert h % 2 == 0 - - arange = torch.arange(2 * K, device=device) - index = ((arange.unsqueeze(0) - arange.unsqueeze(1)) + h // 2).clamp( - min=0, max=h + 1 - ) - index = index.view(K, 2, 2 * K)[:, 0, :] - onehot = one_hot(index, num_classes=h + 2)[..., 1:-1].transpose(1, 0) - return onehot.reshape(2 * K, h * K).float() - - -def single_to_keys(single, indexing_matrix, W, H): - B, N, D = single.shape - K = N // W - single = single.view(B, 2 * K, W // 2, D) - return torch.einsum("b j i d, j k -> b k i d", single, indexing_matrix).reshape( - B, K, H, D - ) # j = 2K, i = W//2, k = h * K - - -class AtomEncoder(Module): - def __init__( - self, - atom_s, - atom_z, - token_s, - token_z, - atoms_per_window_queries, - atoms_per_window_keys, - atom_feature_dim, - structure_prediction=True, - use_no_atom_char=False, - use_atom_backbone_feat=False, - use_residue_feats_atoms=False, - ): - super().__init__() - - self.embed_atom_features = Linear(atom_feature_dim, atom_s) - self.embed_atompair_ref_pos = LinearNoBias(3, atom_z) - self.embed_atompair_ref_dist = LinearNoBias(1, atom_z) - self.embed_atompair_mask = LinearNoBias(1, atom_z) - self.atoms_per_window_queries = atoms_per_window_queries - self.atoms_per_window_keys = atoms_per_window_keys - self.use_no_atom_char = use_no_atom_char - self.use_atom_backbone_feat = use_atom_backbone_feat - self.use_residue_feats_atoms = use_residue_feats_atoms - - self.structure_prediction = structure_prediction - if structure_prediction: - self.s_to_c_trans = nn.Sequential( - nn.LayerNorm(token_s), LinearNoBias(token_s, atom_s) - ) - init.final_init_(self.s_to_c_trans[1].weight) - - self.z_to_p_trans = nn.Sequential( - nn.LayerNorm(token_z), LinearNoBias(token_z, atom_z) - ) - init.final_init_(self.z_to_p_trans[1].weight) - - self.c_to_p_trans_k = nn.Sequential( - nn.ReLU(), - LinearNoBias(atom_s, atom_z), - ) - init.final_init_(self.c_to_p_trans_k[1].weight) - - self.c_to_p_trans_q = nn.Sequential( - nn.ReLU(), - LinearNoBias(atom_s, atom_z), - ) - init.final_init_(self.c_to_p_trans_q[1].weight) - - self.p_mlp = nn.Sequential( - nn.ReLU(), - LinearNoBias(atom_z, atom_z), - nn.ReLU(), - LinearNoBias(atom_z, atom_z), - nn.ReLU(), - LinearNoBias(atom_z, atom_z), - ) - init.final_init_(self.p_mlp[5].weight) - - def forward( - self, - feats, - s_trunk=None, # Float['bm n ts'], - z=None, # Float['bm n n tz'], - ): - with torch.autocast("cuda", enabled=False): - B, N, _ = feats["ref_pos"].shape - atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'], - - atom_ref_pos = feats["ref_pos"] # Float['b m 3'], - atom_uid = feats["ref_space_uid"] # Long['b m'], - - atom_feats = [ - atom_ref_pos, - feats["ref_charge"].unsqueeze(-1), - feats["ref_element"], - ] - if not self.use_no_atom_char: - atom_feats.append(feats["ref_atom_name_chars"].reshape(B, N, 4 * 64)) - if self.use_atom_backbone_feat: - atom_feats.append(feats["atom_backbone_feat"]) - if self.use_residue_feats_atoms: - res_feats = torch.cat( - [ - feats["res_type"], - feats["modified"].unsqueeze(-1), - one_hot(feats["mol_type"], num_classes=4).float(), - ], - dim=-1, - ) - atom_to_token = feats["atom_to_token"].float() - atom_res_feats = torch.bmm(atom_to_token, res_feats) - atom_feats.append(atom_res_feats) - - atom_feats = torch.cat(atom_feats, dim=-1) - - c = self.embed_atom_features(atom_feats) - - # note we are already creating the windows to make it more efficient - W, H = self.atoms_per_window_queries, self.atoms_per_window_keys - B, N = c.shape[:2] - K = N // W - keys_indexing_matrix = get_indexing_matrix(K, W, H, c.device) - to_keys = partial( - single_to_keys, indexing_matrix=keys_indexing_matrix, W=W, H=H - ) - - atom_ref_pos_queries = atom_ref_pos.view(B, K, W, 1, 3) - atom_ref_pos_keys = to_keys(atom_ref_pos).view(B, K, 1, H, 3) - - d = atom_ref_pos_keys - atom_ref_pos_queries # Float['b k w h 3'] - d_norm = torch.sum(d * d, dim=-1, keepdim=True) # Float['b k w h 1'] - d_norm = 1 / ( - 1 + d_norm - ) # AF3 feeds in the reciprocal of the distance norm - - atom_mask_queries = atom_mask.view(B, K, W, 1) - atom_mask_keys = ( - to_keys(atom_mask.unsqueeze(-1).float()).view(B, K, 1, H).bool() - ) - atom_uid_queries = atom_uid.view(B, K, W, 1) - atom_uid_keys = ( - to_keys(atom_uid.unsqueeze(-1).float()).view(B, K, 1, H).long() - ) - v = ( - ( - atom_mask_queries - & atom_mask_keys - & (atom_uid_queries == atom_uid_keys) - ) - .float() - .unsqueeze(-1) - ) # Bool['b k w h 1'] - - p = self.embed_atompair_ref_pos(d) * v - p = p + self.embed_atompair_ref_dist(d_norm) * v - p = p + self.embed_atompair_mask(v) * v - - q = c - - if self.structure_prediction: - # run only in structure model not in initial encoding - atom_to_token = feats["atom_to_token"].float() # Long['b m n'], - - s_to_c = self.s_to_c_trans(s_trunk.float()) - s_to_c = torch.bmm(atom_to_token, s_to_c) - c = c + s_to_c.to(c) - - atom_to_token_queries = atom_to_token.view( - B, K, W, atom_to_token.shape[-1] - ) - atom_to_token_keys = to_keys(atom_to_token) - z_to_p = self.z_to_p_trans(z.float()) - z_to_p = torch.einsum( - "bijd,bwki,bwlj->bwkld", - z_to_p, - atom_to_token_queries, - atom_to_token_keys, - ) - p = p + z_to_p.to(p) - - p = p + self.c_to_p_trans_q(c.view(B, K, W, 1, c.shape[-1])) - p = p + self.c_to_p_trans_k(to_keys(c).view(B, K, 1, H, c.shape[-1])) - p = p + self.p_mlp(p) - return q, c, p, to_keys - - -class AtomAttentionEncoder(Module): - def __init__( - self, - atom_s, - token_s, - atoms_per_window_queries, - atoms_per_window_keys, - atom_encoder_depth=3, - atom_encoder_heads=4, - structure_prediction=True, - activation_checkpointing=False, - transformer_post_layer_norm=False, - ): - super().__init__() - - self.structure_prediction = structure_prediction - if structure_prediction: - self.r_to_q_trans = LinearNoBias(3, atom_s) - init.final_init_(self.r_to_q_trans.weight) - - self.atom_encoder = AtomTransformer( - dim=atom_s, - dim_single_cond=atom_s, - attn_window_queries=atoms_per_window_queries, - attn_window_keys=atoms_per_window_keys, - depth=atom_encoder_depth, - heads=atom_encoder_heads, - activation_checkpointing=activation_checkpointing, - post_layer_norm=transformer_post_layer_norm, - ) - - self.atom_to_token_trans = nn.Sequential( - LinearNoBias(atom_s, 2 * token_s if structure_prediction else token_s), - nn.ReLU(), - ) - - def forward( - self, - feats, - q, - c, - atom_enc_bias, - to_keys, - r=None, # Float['bm m 3'], - multiplicity=1, - ): - B, N, _ = feats["ref_pos"].shape - atom_mask = feats["atom_pad_mask"].bool() # Bool['b m'], - - if self.structure_prediction: - # only here the multiplicity kicks in because we use the different positions r - q = q.repeat_interleave(multiplicity, 0) - r_to_q = self.r_to_q_trans(r) - q = q + r_to_q - - c = c.repeat_interleave(multiplicity, 0) - atom_mask = atom_mask.repeat_interleave(multiplicity, 0) - - q = self.atom_encoder( - q=q, - mask=atom_mask, - c=c, - bias=atom_enc_bias, - multiplicity=multiplicity, - to_keys=to_keys, - ) - - with torch.autocast("cuda", enabled=False): - q_to_a = self.atom_to_token_trans(q).float() - atom_to_token = feats["atom_to_token"].float() - atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0) - atom_to_token_mean = atom_to_token / ( - atom_to_token.sum(dim=1, keepdim=True) + 1e-6 - ) - a = torch.bmm(atom_to_token_mean.transpose(1, 2), q_to_a) - - a = a.to(q) - - return a, q, c, to_keys - - -class AtomAttentionDecoder(Module): - """Algorithm 6.""" - - def __init__( - self, - atom_s, - token_s, - attn_window_queries, - attn_window_keys, - atom_decoder_depth=3, - atom_decoder_heads=4, - activation_checkpointing=False, - transformer_post_layer_norm=False, - ): - super().__init__() - - self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s) - init.final_init_(self.a_to_q_trans.weight) - - self.atom_decoder = AtomTransformer( - dim=atom_s, - dim_single_cond=atom_s, - attn_window_queries=attn_window_queries, - attn_window_keys=attn_window_keys, - depth=atom_decoder_depth, - heads=atom_decoder_heads, - activation_checkpointing=activation_checkpointing, - post_layer_norm=transformer_post_layer_norm, - ) - - if transformer_post_layer_norm: - self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3) - init.final_init_(self.atom_feat_to_atom_pos_update.weight) - else: - self.atom_feat_to_atom_pos_update = nn.Sequential( - nn.LayerNorm(atom_s), LinearNoBias(atom_s, 3) - ) - init.final_init_(self.atom_feat_to_atom_pos_update[1].weight) - - def forward( - self, - a, # Float['bm n 2ts'], - q, # Float['bm m as'], - c, # Float['bm m as'], - atom_dec_bias, # Float['bm m m az'], - feats, - to_keys, - multiplicity=1, - ): - with torch.autocast("cuda", enabled=False): - atom_to_token = feats["atom_to_token"].float() - atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0) - - a_to_q = self.a_to_q_trans(a.float()) - a_to_q = torch.bmm(atom_to_token, a_to_q) - - q = q + a_to_q.to(q) - atom_mask = feats["atom_pad_mask"] # Bool['b m'], - atom_mask = atom_mask.repeat_interleave(multiplicity, 0) - - q = self.atom_decoder( - q=q, - mask=atom_mask, - c=c, - bias=atom_dec_bias, - multiplicity=multiplicity, - to_keys=to_keys, - ) - - r_update = self.atom_feat_to_atom_pos_update(q) - return r_update diff --git a/fastplms/boltz/vb_modules_transformersv2.py b/fastplms/boltz/vb_modules_transformersv2.py deleted file mode 100644 index 601aacc..0000000 --- a/fastplms/boltz/vb_modules_transformersv2.py +++ /dev/null @@ -1,263 +0,0 @@ -# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang - -import torch -from torch import nn, sigmoid -from torch.nn import ( - LayerNorm, - Linear, - Module, - ModuleList, - Sequential, -) - -from .vb_layers_attentionv2 import AttentionPairBias -from .vb_modules_utils import LinearNoBias, SwiGLU, default - - -class AdaLN(Module): - """Algorithm 26""" - - def __init__(self, dim, dim_single_cond): - super().__init__() - self.a_norm = LayerNorm(dim, elementwise_affine=False, bias=False) - self.s_norm = LayerNorm(dim_single_cond, bias=False) - self.s_scale = Linear(dim_single_cond, dim) - self.s_bias = LinearNoBias(dim_single_cond, dim) - - def forward(self, a, s): - a = self.a_norm(a) - s = self.s_norm(s) - a = sigmoid(self.s_scale(s)) * a + self.s_bias(s) - return a - - -class ConditionedTransitionBlock(Module): - """Algorithm 25""" - - def __init__(self, dim_single, dim_single_cond, expansion_factor=2): - super().__init__() - - self.adaln = AdaLN(dim_single, dim_single_cond) - - dim_inner = int(dim_single * expansion_factor) - self.swish_gate = Sequential( - LinearNoBias(dim_single, dim_inner * 2), - SwiGLU(), - ) - self.a_to_b = LinearNoBias(dim_single, dim_inner) - self.b_to_a = LinearNoBias(dim_inner, dim_single) - - output_projection_linear = Linear(dim_single_cond, dim_single) - nn.init.zeros_(output_projection_linear.weight) - nn.init.constant_(output_projection_linear.bias, -2.0) - - self.output_projection = nn.Sequential(output_projection_linear, nn.Sigmoid()) - - def forward( - self, - a, # Float['... d'] - s, - ): # -> Float['... d']: - a = self.adaln(a, s) - b = self.swish_gate(a) * self.a_to_b(a) - a = self.output_projection(s) * self.b_to_a(b) - - return a - - -class DiffusionTransformer(Module): - """Algorithm 23""" - - def __init__( - self, - depth, - heads, - dim=384, - dim_single_cond=None, - pair_bias_attn=True, - activation_checkpointing=False, - post_layer_norm=False, - ): - super().__init__() - self.activation_checkpointing = activation_checkpointing - dim_single_cond = default(dim_single_cond, dim) - self.pair_bias_attn = pair_bias_attn - - self.layers = ModuleList() - for _ in range(depth): - self.layers.append( - DiffusionTransformerLayer( - heads, - dim, - dim_single_cond, - post_layer_norm, - ) - ) - - def forward( - self, - a, # Float['bm n d'], - s, # Float['bm n ds'], - bias=None, # Float['b n n dp'] - mask=None, # Bool['b n'] | None = None - to_keys=None, - multiplicity=1, - ): - if self.pair_bias_attn: - B, N, M, D = bias.shape - L = len(self.layers) - bias = bias.view(B, N, M, L, D // L) - - for i, layer in enumerate(self.layers): - if self.pair_bias_attn: - bias_l = bias[:, :, :, i] - else: - bias_l = None - - if self.activation_checkpointing: - a = torch.utils.checkpoint.checkpoint( - layer, - a, - s, - bias_l, - mask, - to_keys, - multiplicity, - use_reentrant=False, - ) - - else: - a = layer( - a, # Float['bm n d'], - s, # Float['bm n ds'], - bias_l, # Float['b n n dp'] - mask, # Bool['b n'] | None = None - to_keys, - multiplicity, - ) - return a - - -class DiffusionTransformerLayer(Module): - """Algorithm 23""" - - def __init__( - self, - heads, - dim=384, - dim_single_cond=None, - post_layer_norm=False, - ): - super().__init__() - - dim_single_cond = default(dim_single_cond, dim) - - self.adaln = AdaLN(dim, dim_single_cond) - self.pair_bias_attn = AttentionPairBias( - c_s=dim, num_heads=heads, compute_pair_bias=False - ) - - self.output_projection_linear = Linear(dim_single_cond, dim) - nn.init.zeros_(self.output_projection_linear.weight) - nn.init.constant_(self.output_projection_linear.bias, -2.0) - - self.output_projection = nn.Sequential( - self.output_projection_linear, nn.Sigmoid() - ) - self.transition = ConditionedTransitionBlock( - dim_single=dim, dim_single_cond=dim_single_cond - ) - - if post_layer_norm: - self.post_lnorm = nn.LayerNorm(dim) - else: - self.post_lnorm = nn.Identity() - - def forward( - self, - a, # Float['bm n d'], - s, # Float['bm n ds'], - bias=None, # Float['b n n dp'] - mask=None, # Bool['b n'] | None = None - to_keys=None, - multiplicity=1, - ): - b = self.adaln(a, s) - - k_in = b - if to_keys is not None: - k_in = to_keys(b) - mask = to_keys(mask.unsqueeze(-1)).squeeze(-1) - - if self.pair_bias_attn: - b = self.pair_bias_attn( - s=b, - z=bias, - mask=mask, - multiplicity=multiplicity, - k_in=k_in, - ) - else: - b = self.no_pair_bias_attn(s=b, mask=mask, k_in=k_in) - - b = self.output_projection(s) * b - - a = a + b - a = a + self.transition(a, s) - - a = self.post_lnorm(a) - return a - - -class AtomTransformer(Module): - """Algorithm 7""" - - def __init__( - self, - attn_window_queries, - attn_window_keys, - **diffusion_transformer_kwargs, - ): - super().__init__() - self.attn_window_queries = attn_window_queries - self.attn_window_keys = attn_window_keys - self.diffusion_transformer = DiffusionTransformer( - **diffusion_transformer_kwargs - ) - - def forward( - self, - q, # Float['b m d'], - c, # Float['b m ds'], - bias, # Float['b m m dp'] - to_keys, - mask, # Bool['b m'] | None = None - multiplicity=1, - ): - W = self.attn_window_queries - H = self.attn_window_keys - - B, N, D = q.shape - NW = N // W - - # reshape tokens - q = q.view((B * NW, W, -1)) - c = c.view((B * NW, W, -1)) - mask = mask.view(B * NW, W) - bias = bias.repeat_interleave(multiplicity, 0) - bias = bias.view((bias.shape[0] * NW, W, H, -1)) - - to_keys_new = lambda x: to_keys(x.view(B, NW * W, -1)).view(B * NW, H, -1) - - # main transformer - q = self.diffusion_transformer( - a=q, - s=c, - bias=bias, - mask=mask.float(), - multiplicity=1, # bias term already expanded with multiplicity - to_keys=to_keys_new, - ) - - q = q.view((B, NW * W, D)) - return q diff --git a/fastplms/boltz/vb_modules_trunkv2.py b/fastplms/boltz/vb_modules_trunkv2.py deleted file mode 100644 index ad3b9f3..0000000 --- a/fastplms/boltz/vb_modules_trunkv2.py +++ /dev/null @@ -1,833 +0,0 @@ -from typing import Dict, Tuple - -import torch -from torch import Tensor, nn -from torch.nn.functional import one_hot - -from . import vb_const as const -from .vb_layers_outer_product_mean import OuterProductMean -from .vb_layers_pair_averaging import PairWeightedAveraging -from .vb_layers_pairformer import ( - PairformerNoSeqLayer, - PairformerNoSeqModule, - get_dropout_mask, -) -from .vb_layers_transition import Transition -from .vb_modules_encodersv2 import ( - AtomAttentionEncoder, - AtomEncoder, - FourierEmbedding, -) - - -class ContactConditioning(nn.Module): - def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float): - super().__init__() - - self.fourier_embedding = FourierEmbedding(token_z) - self.encoder = nn.Linear( - token_z + len(const.contact_conditioning_info) - 1, token_z - ) - self.encoding_unspecified = nn.Parameter(torch.zeros(token_z)) - self.encoding_unselected = nn.Parameter(torch.zeros(token_z)) - self.cutoff_min = cutoff_min - self.cutoff_max = cutoff_max - - def forward(self, feats): - assert const.contact_conditioning_info["UNSPECIFIED"] == 0 - assert const.contact_conditioning_info["UNSELECTED"] == 1 - contact_conditioning = feats["contact_conditioning"][:, :, :, 2:] - contact_threshold = feats["contact_threshold"] - contact_threshold_normalized = (contact_threshold - self.cutoff_min) / ( - self.cutoff_max - self.cutoff_min - ) - contact_threshold_fourier = self.fourier_embedding( - contact_threshold_normalized.flatten() - ).reshape(contact_threshold_normalized.shape + (-1,)) - - contact_conditioning = torch.cat( - [ - contact_conditioning, - contact_threshold_normalized.unsqueeze(-1), - contact_threshold_fourier, - ], - dim=-1, - ) - contact_conditioning = self.encoder(contact_conditioning) - - contact_conditioning = ( - contact_conditioning - * ( - 1 - - feats["contact_conditioning"][:, :, :, 0:2].sum(dim=-1, keepdim=True) - ) - + self.encoding_unspecified * feats["contact_conditioning"][:, :, :, 0:1] - + self.encoding_unselected * feats["contact_conditioning"][:, :, :, 1:2] - ) - return contact_conditioning - - -class InputEmbedder(nn.Module): - def __init__( - self, - atom_s: int, - atom_z: int, - token_s: int, - token_z: int, - atoms_per_window_queries: int, - atoms_per_window_keys: int, - atom_feature_dim: int, - atom_encoder_depth: int, - atom_encoder_heads: int, - activation_checkpointing: bool = False, - add_method_conditioning: bool = False, - add_modified_flag: bool = False, - add_cyclic_flag: bool = False, - add_mol_type_feat: bool = False, - use_no_atom_char: bool = False, - use_atom_backbone_feat: bool = False, - use_residue_feats_atoms: bool = False, - ) -> None: - """Initialize the input embedder. - - Parameters - ---------- - atom_s : int - The atom embedding size. - atom_z : int - The atom pairwise embedding size. - token_s : int - The token embedding size. - - """ - super().__init__() - self.token_s = token_s - self.add_method_conditioning = add_method_conditioning - self.add_modified_flag = add_modified_flag - self.add_cyclic_flag = add_cyclic_flag - self.add_mol_type_feat = add_mol_type_feat - - self.atom_encoder = AtomEncoder( - atom_s=atom_s, - atom_z=atom_z, - token_s=token_s, - token_z=token_z, - atoms_per_window_queries=atoms_per_window_queries, - atoms_per_window_keys=atoms_per_window_keys, - atom_feature_dim=atom_feature_dim, - structure_prediction=False, - use_no_atom_char=use_no_atom_char, - use_atom_backbone_feat=use_atom_backbone_feat, - use_residue_feats_atoms=use_residue_feats_atoms, - ) - - self.atom_enc_proj_z = nn.Sequential( - nn.LayerNorm(atom_z), - nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False), - ) - - self.atom_attention_encoder = AtomAttentionEncoder( - atom_s=atom_s, - token_s=token_s, - atoms_per_window_queries=atoms_per_window_queries, - atoms_per_window_keys=atoms_per_window_keys, - atom_encoder_depth=atom_encoder_depth, - atom_encoder_heads=atom_encoder_heads, - structure_prediction=False, - activation_checkpointing=activation_checkpointing, - ) - - self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False) - self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False) - - if add_method_conditioning: - self.method_conditioning_init = nn.Embedding( - const.num_method_types, token_s - ) - self.method_conditioning_init.weight.data.fill_(0) - if add_modified_flag: - self.modified_conditioning_init = nn.Embedding(2, token_s) - self.modified_conditioning_init.weight.data.fill_(0) - if add_cyclic_flag: - self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False) - self.cyclic_conditioning_init.weight.data.fill_(0) - if add_mol_type_feat: - self.mol_type_conditioning_init = nn.Embedding( - len(const.chain_type_ids), token_s - ) - self.mol_type_conditioning_init.weight.data.fill_(0) - - def forward(self, feats: Dict[str, Tensor], affinity: bool = False) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - feats : dict[str, Tensor] - Input features - - Returns - ------- - Tensor - The embedded tokens. - - """ - # Load relevant features - res_type = feats["res_type"].float() - if affinity: - profile = feats["profile_affinity"] - deletion_mean = feats["deletion_mean_affinity"].unsqueeze(-1) - else: - profile = feats["profile"] - deletion_mean = feats["deletion_mean"].unsqueeze(-1) - - # Compute input embedding - q, c, p, to_keys = self.atom_encoder(feats) - atom_enc_bias = self.atom_enc_proj_z(p) - a, _, _, _ = self.atom_attention_encoder( - feats=feats, - q=q, - c=c, - atom_enc_bias=atom_enc_bias, - to_keys=to_keys, - ) - - s = ( - a - + self.res_type_encoding(res_type) - + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1)) - ) - - if self.add_method_conditioning: - s = s + self.method_conditioning_init(feats["method_feature"]) - if self.add_modified_flag: - s = s + self.modified_conditioning_init(feats["modified"]) - if self.add_cyclic_flag: - cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1) - s = s + self.cyclic_conditioning_init(cyclic) - if self.add_mol_type_feat: - s = s + self.mol_type_conditioning_init(feats["mol_type"]) - - return s - - -class TemplateModule(nn.Module): - """Template module.""" - - def __init__( - self, - token_z: int, - template_dim: int, - template_blocks: int, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - activation_checkpointing: bool = False, - min_dist: float = 3.25, - max_dist: float = 50.75, - num_bins: int = 38, - **kwargs, - ) -> None: - """Initialize the template module. - - Parameters - ---------- - token_z : int - The token pairwise embedding size. - - """ - super().__init__() - self.min_dist = min_dist - self.max_dist = max_dist - self.num_bins = num_bins - self.relu = nn.ReLU() - self.z_norm = nn.LayerNorm(token_z) - self.v_norm = nn.LayerNorm(template_dim) - self.z_proj = nn.Linear(token_z, template_dim, bias=False) - self.a_proj = nn.Linear( - const.num_tokens * 2 + num_bins + 5, - template_dim, - bias=False, - ) - self.u_proj = nn.Linear(template_dim, token_z, bias=False) - self.pairformer = PairformerNoSeqModule( - template_dim, - num_blocks=template_blocks, - dropout=dropout, - pairwise_head_width=pairwise_head_width, - pairwise_num_heads=pairwise_num_heads, - post_layer_norm=post_layer_norm, - activation_checkpointing=activation_checkpointing, - ) - - def forward( - self, - z: Tensor, - feats: Dict[str, Tensor], - pair_mask: Tensor, - use_kernels: bool = False, - ) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - z : Tensor - The pairwise embeddings - feats : dict[str, Tensor] - Input features - pair_mask : Tensor - The pair mask - - Returns - ------- - Tensor - The updated pairwise embeddings. - - """ - # Load relevant features - asym_id = feats["asym_id"] - res_type = feats["template_restype"] - frame_rot = feats["template_frame_rot"] - frame_t = feats["template_frame_t"] - frame_mask = feats["template_mask_frame"] - cb_coords = feats["template_cb"] - ca_coords = feats["template_ca"] - cb_mask = feats["template_mask_cb"] - template_mask = feats["template_mask"].any(dim=2).float() - num_templates = template_mask.sum(dim=1) - num_templates = num_templates.clamp(min=1) - - # Compute pairwise masks - b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :] - b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :] - - b_cb_mask = b_cb_mask[..., None] - b_frame_mask = b_frame_mask[..., None] - - # Compute asym mask, template features only attend within the same chain - B, T = res_type.shape[:2] # noqa: N806 - asym_mask = (asym_id[:, :, None] == asym_id[:, None, :]).float() - asym_mask = asym_mask[:, None].expand(-1, T, -1, -1) - - # Compute template features - with torch.autocast(device_type="cuda", enabled=False): - # Compute distogram - cb_dists = torch.cdist(cb_coords, cb_coords) - boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1) - boundaries = boundaries.to(cb_dists.device) - distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long() - distogram = one_hot(distogram, num_classes=self.num_bins) - - # Compute unit vector in each frame - frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2) - frame_t = frame_t.unsqueeze(2).unsqueeze(-1) - ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1) - vector = torch.matmul(frame_rot, (ca_coords - frame_t)) - norm = torch.norm(vector, dim=-1, keepdim=True) - unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector)) - unit_vector = unit_vector.squeeze(-1) - - # Concatenate input features - a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask] - a_tij = torch.cat(a_tij, dim=-1) - a_tij = a_tij * asym_mask.unsqueeze(-1) - - res_type_i = res_type[:, :, :, None] - res_type_j = res_type[:, :, None, :] - res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1) - res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1) - a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1) - a_tij = self.a_proj(a_tij) - - # Expand mask - pair_mask = pair_mask[:, None].expand(-1, T, -1, -1) - pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:]) - - # Compute input projections - v = self.z_proj(self.z_norm(z[:, None])) + a_tij - v = v.view(B * T, *v.shape[2:]) - v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels) - v = self.v_norm(v) - v = v.view(B, T, *v.shape[1:]) - - # Aggregate templates - template_mask = template_mask[:, :, None, None, None] - num_templates = num_templates[:, None, None, None] - u = (v * template_mask).sum(dim=1) / num_templates.to(v) - - # Compute output projection - u = self.u_proj(self.relu(u)) - return u - - -class TemplateV2Module(nn.Module): - """Template module.""" - - def __init__( - self, - token_z: int, - template_dim: int, - template_blocks: int, - dropout: float = 0.25, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - post_layer_norm: bool = False, - activation_checkpointing: bool = False, - min_dist: float = 3.25, - max_dist: float = 50.75, - num_bins: int = 38, - **kwargs, - ) -> None: - """Initialize the template module. - - Parameters - ---------- - token_z : int - The token pairwise embedding size. - - """ - super().__init__() - self.min_dist = min_dist - self.max_dist = max_dist - self.num_bins = num_bins - self.relu = nn.ReLU() - self.z_norm = nn.LayerNorm(token_z) - self.v_norm = nn.LayerNorm(template_dim) - self.z_proj = nn.Linear(token_z, template_dim, bias=False) - self.a_proj = nn.Linear( - const.num_tokens * 2 + num_bins + 5, - template_dim, - bias=False, - ) - self.u_proj = nn.Linear(template_dim, token_z, bias=False) - self.pairformer = PairformerNoSeqModule( - template_dim, - num_blocks=template_blocks, - dropout=dropout, - pairwise_head_width=pairwise_head_width, - pairwise_num_heads=pairwise_num_heads, - post_layer_norm=post_layer_norm, - activation_checkpointing=activation_checkpointing, - ) - - def forward( - self, - z: Tensor, - feats: Dict[str, Tensor], - pair_mask: Tensor, - use_kernels: bool = False, - ) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - z : Tensor - The pairwise embeddings - feats : dict[str, Tensor] - Input features - pair_mask : Tensor - The pair mask - - Returns - ------- - Tensor - The updated pairwise embeddings. - - """ - # Load relevant features - res_type = feats["template_restype"] - frame_rot = feats["template_frame_rot"] - frame_t = feats["template_frame_t"] - frame_mask = feats["template_mask_frame"] - cb_coords = feats["template_cb"] - ca_coords = feats["template_ca"] - cb_mask = feats["template_mask_cb"] - visibility_ids = feats["visibility_ids"] - template_mask = feats["template_mask"].any(dim=2).float() - num_templates = template_mask.sum(dim=1) - num_templates = num_templates.clamp(min=1) - - # Compute pairwise masks - b_cb_mask = cb_mask[:, :, :, None] * cb_mask[:, :, None, :] - b_frame_mask = frame_mask[:, :, :, None] * frame_mask[:, :, None, :] - - b_cb_mask = b_cb_mask[..., None] - b_frame_mask = b_frame_mask[..., None] - - # Compute asym mask, template features only attend within the same chain - B, T = res_type.shape[:2] # noqa: N806 - tmlp_pair_mask = ( - visibility_ids[:, :, :, None] == visibility_ids[:, :, None, :] - ).float() - - # Compute template features - with torch.autocast(device_type="cuda", enabled=False): - # Compute distogram - cb_dists = torch.cdist(cb_coords, cb_coords) - boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1) - boundaries = boundaries.to(cb_dists.device) - distogram = (cb_dists[..., None] > boundaries).sum(dim=-1).long() - distogram = one_hot(distogram, num_classes=self.num_bins) - - # Compute unit vector in each frame - frame_rot = frame_rot.unsqueeze(2).transpose(-1, -2) - frame_t = frame_t.unsqueeze(2).unsqueeze(-1) - ca_coords = ca_coords.unsqueeze(3).unsqueeze(-1) - vector = torch.matmul(frame_rot, (ca_coords - frame_t)) - norm = torch.norm(vector, dim=-1, keepdim=True) - unit_vector = torch.where(norm > 0, vector / norm, torch.zeros_like(vector)) - unit_vector = unit_vector.squeeze(-1) - - # Concatenate input features - a_tij = [distogram, b_cb_mask, unit_vector, b_frame_mask] - a_tij = torch.cat(a_tij, dim=-1) - a_tij = a_tij * tmlp_pair_mask.unsqueeze(-1) - - res_type_i = res_type[:, :, :, None] - res_type_j = res_type[:, :, None, :] - res_type_i = res_type_i.expand(-1, -1, -1, res_type.size(2), -1) - res_type_j = res_type_j.expand(-1, -1, res_type.size(2), -1, -1) - a_tij = torch.cat([a_tij, res_type_i, res_type_j], dim=-1) - a_tij = self.a_proj(a_tij) - - # Expand mask - pair_mask = pair_mask[:, None].expand(-1, T, -1, -1) - pair_mask = pair_mask.reshape(B * T, *pair_mask.shape[2:]) - - # Compute input projections - v = self.z_proj(self.z_norm(z[:, None])) + a_tij - v = v.view(B * T, *v.shape[2:]) - v = v + self.pairformer(v, pair_mask, use_kernels=use_kernels) - v = self.v_norm(v) - v = v.view(B, T, *v.shape[1:]) - - # Aggregate templates - template_mask = template_mask[:, :, None, None, None] - num_templates = num_templates[:, None, None, None] - u = (v * template_mask).sum(dim=1) / num_templates.to(v) - - # Compute output projection - u = self.u_proj(self.relu(u)) - return u - - -class MSAModule(nn.Module): - """MSA module.""" - - def __init__( - self, - msa_s: int, - token_z: int, - token_s: int, - msa_blocks: int, - msa_dropout: float, - z_dropout: float, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - activation_checkpointing: bool = False, - use_paired_feature: bool = True, - subsample_msa: bool = False, - num_subsampled_msa: int = 1024, - **kwargs, - ) -> None: - """Initialize the MSA module. - - Parameters - ---------- - token_z : int - The token pairwise embedding size. - - """ - super().__init__() - self.msa_blocks = msa_blocks - self.msa_dropout = msa_dropout - self.z_dropout = z_dropout - self.use_paired_feature = use_paired_feature - self.activation_checkpointing = activation_checkpointing - self.subsample_msa = subsample_msa - self.num_subsampled_msa = num_subsampled_msa - - self.s_proj = nn.Linear(token_s, msa_s, bias=False) - self.msa_proj = nn.Linear( - const.num_tokens + 2 + int(use_paired_feature), - msa_s, - bias=False, - ) - self.layers = nn.ModuleList() - for i in range(msa_blocks): - self.layers.append( - MSALayer( - msa_s, - token_z, - msa_dropout, - z_dropout, - pairwise_head_width, - pairwise_num_heads, - ) - ) - - def forward( - self, - z: Tensor, - emb: Tensor, - feats: Dict[str, Tensor], - use_kernels: bool = False, - ) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - z : Tensor - The pairwise embeddings - emb : Tensor - The input embeddings - feats : dict[str, Tensor] - Input features - use_kernels: bool - Whether to use kernels for triangular updates - - Returns - ------- - Tensor - The output pairwise embeddings. - - """ - # Set chunk sizes - if not self.training: - if z.shape[1] > const.chunk_size_threshold: - chunk_heads_pwa = True - chunk_size_transition_z = 64 - chunk_size_transition_msa = 32 - chunk_size_outer_product = 4 - chunk_size_tri_attn = 128 - else: - chunk_heads_pwa = False - chunk_size_transition_z = None - chunk_size_transition_msa = None - chunk_size_outer_product = None - chunk_size_tri_attn = 512 - else: - chunk_heads_pwa = False - chunk_size_transition_z = None - chunk_size_transition_msa = None - chunk_size_outer_product = None - chunk_size_tri_attn = None - - # Load relevant features - msa = feats["msa"] - if msa.dtype in (torch.long, torch.int32, torch.int64): - msa = torch.nn.functional.one_hot(msa, num_classes=const.num_tokens).float() - # else: already float one-hot (soft/differentiable path) - has_deletion = feats["has_deletion"].unsqueeze(-1) - deletion_value = feats["deletion_value"].unsqueeze(-1) - is_paired = feats["msa_paired"].unsqueeze(-1) - msa_mask = feats["msa_mask"] - token_mask = feats["token_pad_mask"].float() - token_mask = token_mask[:, :, None] * token_mask[:, None, :] - - # Compute MSA embeddings - if self.use_paired_feature: - m = torch.cat([msa, has_deletion, deletion_value, is_paired], dim=-1) - else: - m = torch.cat([msa, has_deletion, deletion_value], dim=-1) - - # Subsample the MSA - if self.subsample_msa: - msa_indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa] - m = m[:, msa_indices] - msa_mask = msa_mask[:, msa_indices] - - # Compute input projections - m = self.msa_proj(m) - m = m + self.s_proj(emb).unsqueeze(1) - - # Perform MSA blocks - for i in range(self.msa_blocks): - if self.activation_checkpointing: - z, m = torch.utils.checkpoint.checkpoint( - self.layers[i], - z, - m, - token_mask, - msa_mask, - chunk_heads_pwa, - chunk_size_transition_z, - chunk_size_transition_msa, - chunk_size_outer_product, - chunk_size_tri_attn, - use_kernels, - use_reentrant=False, - ) - else: - z, m = self.layers[i]( - z, - m, - token_mask, - msa_mask, - chunk_heads_pwa, - chunk_size_transition_z, - chunk_size_transition_msa, - chunk_size_outer_product, - chunk_size_tri_attn, - use_kernels, - ) - return z - - -class MSALayer(nn.Module): - """MSA module.""" - - def __init__( - self, - msa_s: int, - token_z: int, - msa_dropout: float, - z_dropout: float, - pairwise_head_width: int = 32, - pairwise_num_heads: int = 4, - ) -> None: - """Initialize the MSA module. - - Parameters - ---------- - token_z : int - The token pairwise embedding size. - - """ - super().__init__() - self.msa_dropout = msa_dropout - self.msa_transition = Transition(dim=msa_s, hidden=msa_s * 4) - self.pair_weighted_averaging = PairWeightedAveraging( - c_m=msa_s, - c_z=token_z, - c_h=32, - num_heads=8, - ) - - self.pairformer_layer = PairformerNoSeqLayer( - token_z=token_z, - dropout=z_dropout, - pairwise_head_width=pairwise_head_width, - pairwise_num_heads=pairwise_num_heads, - ) - self.outer_product_mean = OuterProductMean( - c_in=msa_s, - c_hidden=32, - c_out=token_z, - ) - - def forward( - self, - z: Tensor, - m: Tensor, - token_mask: Tensor, - msa_mask: Tensor, - chunk_heads_pwa: bool = False, - chunk_size_transition_z: int = None, - chunk_size_transition_msa: int = None, - chunk_size_outer_product: int = None, - chunk_size_tri_attn: int = None, - use_kernels: bool = False, - ) -> Tuple[Tensor, Tensor]: - """Perform the forward pass. - - Parameters - ---------- - z : Tensor - The pairwise embeddings - emb : Tensor - The input embeddings - feats : dict[str, Tensor] - Input features - - Returns - ------- - Tensor - The output pairwise embeddings. - - """ - # Communication to MSA stack - msa_dropout = get_dropout_mask(self.msa_dropout, m, self.training) - m = m + msa_dropout * self.pair_weighted_averaging( - m, z, token_mask, chunk_heads_pwa - ) - m = m + self.msa_transition(m, chunk_size_transition_msa) - - z = z + self.outer_product_mean(m, msa_mask, chunk_size_outer_product) - - # Compute pairwise stack - z = self.pairformer_layer( - z, token_mask, chunk_size_tri_attn, use_kernels=use_kernels - ) - - return z, m - - -class BFactorModule(nn.Module): - """BFactor Module.""" - - def __init__(self, token_s: int, num_bins: int) -> None: - """Initialize the bfactor module. - - Parameters - ---------- - token_s : int - The token embedding size. - - """ - super().__init__() - self.bfactor = nn.Linear(token_s, num_bins) - self.num_bins = num_bins - - def forward(self, s: Tensor) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - s : Tensor - The sequence embeddings - - Returns - ------- - Tensor - The predicted bfactor histogram. - - """ - return self.bfactor(s) - - -class DistogramModule(nn.Module): - """Distogram Module.""" - - def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None: - """Initialize the distogram module. - - Parameters - ---------- - token_z : int - The token pairwise embedding size. - - """ - super().__init__() - self.distogram = nn.Linear(token_z, num_distograms * num_bins) - self.num_distograms = num_distograms - self.num_bins = num_bins - - def forward(self, z: Tensor) -> Tensor: - """Perform the forward pass. - - Parameters - ---------- - z : Tensor - The pairwise embeddings - - Returns - ------- - Tensor - The predicted distogram. - - """ - z = z + z.transpose(1, 2) - return self.distogram(z).reshape( - z.shape[0], z.shape[1], z.shape[2], self.num_distograms, self.num_bins - ) diff --git a/fastplms/boltz/vb_modules_utils.py b/fastplms/boltz/vb_modules_utils.py deleted file mode 100644 index a5a1f2e..0000000 --- a/fastplms/boltz/vb_modules_utils.py +++ /dev/null @@ -1,303 +0,0 @@ -# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang - -from functools import partial -from typing import Optional - -import torch -import torch.nn.functional as F -from torch.nn import ( - Linear, - Module, -) -from torch.types import Device - -LinearNoBias = partial(Linear, bias=False) - - -def exists(v): - return v is not None - - -def default(v, d): - return v if exists(v) else d - - -def log(t, eps=1e-20): - return torch.log(t.clamp(min=eps)) - - -class SwiGLU(Module): - def forward( - self, - x, #: Float['... d'] - ): # -> Float[' ... (d//2)']: - x, gates = x.chunk(2, dim=-1) - return F.silu(gates) * x - - -def center(atom_coords, atom_mask): - atom_mean = torch.sum( - atom_coords * atom_mask[:, :, None], dim=1, keepdim=True - ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True) - atom_coords = atom_coords - atom_mean - return atom_coords - - -def compute_random_augmentation( - multiplicity, s_trans=1.0, device=None, dtype=torch.float32 -): - R = random_rotations(multiplicity, dtype=dtype, device=device) - random_trans = ( - torch.randn((multiplicity, 1, 3), dtype=dtype, device=device) * s_trans - ) - return R, random_trans - - -def randomly_rotate(coords, return_second_coords=False, second_coords=None): - R = random_rotations(len(coords), coords.dtype, coords.device) - - if return_second_coords: - return torch.einsum("bmd,bds->bms", coords, R), torch.einsum( - "bmd,bds->bms", second_coords, R - ) if second_coords is not None else None - - return torch.einsum("bmd,bds->bms", coords, R) - - -def center_random_augmentation( - atom_coords, - atom_mask, - s_trans=1.0, - augmentation=True, - centering=True, - return_second_coords=False, - second_coords=None, -): - """Algorithm 19""" - if centering: - atom_mean = torch.sum( - atom_coords * atom_mask[:, :, None], dim=1, keepdim=True - ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True) - atom_coords = atom_coords - atom_mean - - if second_coords is not None: - # apply same transformation also to this input - second_coords = second_coords - atom_mean - - if augmentation: - atom_coords, second_coords = randomly_rotate( - atom_coords, return_second_coords=True, second_coords=second_coords - ) - random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans - atom_coords = atom_coords + random_trans - - if second_coords is not None: - second_coords = second_coords + random_trans - - if return_second_coords: - return atom_coords, second_coords - - return atom_coords - - -class ExponentialMovingAverage: - """from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ema.py, Apache-2.0 license - Maintains (exponential) moving average of a set of parameters.""" - - def __init__(self, parameters, decay, use_num_updates=True): - """ - Args: - parameters: Iterable of `torch.nn.Parameter`; usually the result of - `model.parameters()`. - decay: The exponential decay. - use_num_updates: Whether to use number of updates when computing - averages. - """ - if decay < 0.0 or decay > 1.0: - raise ValueError("Decay must be between 0 and 1") - self.decay = decay - self.num_updates = 0 if use_num_updates else None - self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad] - self.collected_params = [] - - def update(self, parameters): - """ - Update currently maintained parameters. - Call this every time the parameters are updated, such as the result of - the `optimizer.step()` call. - Args: - parameters: Iterable of `torch.nn.Parameter`; usually the same set of - parameters used to initialize this object. - """ - decay = self.decay - if self.num_updates is not None: - self.num_updates += 1 - decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates)) - one_minus_decay = 1.0 - decay - with torch.no_grad(): - parameters = [p for p in parameters if p.requires_grad] - for s_param, param in zip(self.shadow_params, parameters): - s_param.sub_(one_minus_decay * (s_param - param)) - - def compatible(self, parameters): - if len(self.shadow_params) != len(parameters): - print( - f"Model has {len(self.shadow_params)} parameter tensors, the incoming ema {len(parameters)}" - ) - return False - - for s_param, param in zip(self.shadow_params, parameters): - if param.data.shape != s_param.data.shape: - print( - f"Model has parameter tensor of shape {s_param.data.shape} , the incoming ema {param.data.shape}" - ) - return False - return True - - def copy_to(self, parameters): - """ - Copy current parameters into given collection of parameters. - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored moving averages. - """ - parameters = [p for p in parameters if p.requires_grad] - for s_param, param in zip(self.shadow_params, parameters): - if param.requires_grad: - param.data.copy_(s_param.data) - - def store(self, parameters): - """ - Save the current parameters for restoring later. - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - temporarily stored. - """ - self.collected_params = [param.clone() for param in parameters] - - def restore(self, parameters): - """ - Restore the parameters stored with the `store` method. - Useful to validate the model with EMA parameters without affecting the - original optimization process. Store the parameters before the - `copy_to` method. After validation (or model saving), use this to - restore the former parameters. - Args: - parameters: Iterable of `torch.nn.Parameter`; the parameters to be - updated with the stored parameters. - """ - for c_param, param in zip(self.collected_params, parameters): - param.data.copy_(c_param.data) - - def state_dict(self): - return dict( - decay=self.decay, - num_updates=self.num_updates, - shadow_params=self.shadow_params, - ) - - def load_state_dict(self, state_dict, device): - self.decay = state_dict["decay"] - self.num_updates = state_dict["num_updates"] - self.shadow_params = [ - tensor.to(device) for tensor in state_dict["shadow_params"] - ] - - def to(self, device): - self.shadow_params = [tensor.to(device) for tensor in self.shadow_params] - - -# the following is copied from Torch3D, BSD License, Copyright (c) Meta Platforms, Inc. and affiliates. - - -def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """ - Return a tensor where each element has the absolute value taken from the, - corresponding element of a, with sign taken from the corresponding - element of b. This is like the standard copysign floating-point operation, - but is not careful about negative 0 and NaN. - - Args: - a: source tensor. - b: tensor whose signs will be used, of the same shape as a. - - Returns: - Tensor of the same shape as a with the signs of b. - """ - signs_differ = (a < 0) != (b < 0) - return torch.where(signs_differ, -a, a) - - -def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: - """ - Convert rotations given as quaternions to rotation matrices. - - Args: - quaternions: quaternions with real part first, - as tensor of shape (..., 4). - - Returns: - Rotation matrices as tensor of shape (..., 3, 3). - """ - r, i, j, k = torch.unbind(quaternions, -1) - # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. - two_s = 2.0 / (quaternions * quaternions).sum(-1) - - o = torch.stack( - ( - 1 - two_s * (j * j + k * k), - two_s * (i * j - k * r), - two_s * (i * k + j * r), - two_s * (i * j + k * r), - 1 - two_s * (i * i + k * k), - two_s * (j * k - i * r), - two_s * (i * k - j * r), - two_s * (j * k + i * r), - 1 - two_s * (i * i + j * j), - ), - -1, - ) - return o.reshape(quaternions.shape[:-1] + (3, 3)) - - -def random_quaternions( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None -) -> torch.Tensor: - """ - Generate random quaternions representing rotations, - i.e. versors with nonnegative real part. - - Args: - n: Number of quaternions in a batch to return. - dtype: Type to return. - device: Desired device of returned tensor. Default: - uses the current device for the default tensor type. - - Returns: - Quaternions as tensor of shape (N, 4). - """ - if isinstance(device, str): - device = torch.device(device) - o = torch.randn((n, 4), dtype=dtype, device=device) - s = (o * o).sum(1) - o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] - return o - - -def random_rotations( - n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None -) -> torch.Tensor: - """ - Generate random rotations as 3x3 rotation matrices. - - Args: - n: Number of rotation matrices in a batch to return. - dtype: Type to return. - device: Device of returned tensor. Default: if None, - uses the current device for the default tensor type. - - Returns: - Rotation matrices as tensor of shape (n, 3, 3). - """ - quaternions = random_quaternions(n, dtype=dtype, device=device) - return quaternion_to_matrix(quaternions) diff --git a/fastplms/boltz/vb_potentials_potentials.py b/fastplms/boltz/vb_potentials_potentials.py deleted file mode 100644 index a92fadb..0000000 --- a/fastplms/boltz/vb_potentials_potentials.py +++ /dev/null @@ -1,787 +0,0 @@ -from abc import ABC, abstractmethod -from typing import Optional, Dict, Any, Set, List, Union - -import torch -import numpy as np -from . import vb_const as const -from .vb_potentials_schedules import ( - ParameterSchedule, - ExponentialInterpolation, - PiecewiseStepFunction, -) -from .vb_loss_diffusionv2 import weighted_rigid_align - - -class Potential(ABC): - def __init__( - self, - parameters: Optional[ - Dict[str, Union[ParameterSchedule, float, int, bool]] - ] = None, - ): - self.parameters = parameters - - def compute(self, coords, feats, parameters): - index, args, com_args, ref_args, operator_args = self.compute_args( - feats, parameters - ) - - if index.shape[1] == 0: - return torch.zeros(coords.shape[:-2], device=coords.device) - - if com_args is not None: - com_index, atom_pad_mask = com_args - unpad_com_index = com_index[atom_pad_mask] - unpad_coords = coords[..., atom_pad_mask, :] - coords = torch.zeros( - (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3), - device=coords.device, - ).scatter_reduce( - -2, - unpad_com_index.unsqueeze(-1).expand_as(unpad_coords), - unpad_coords, - "mean", - ) - else: - com_index, atom_pad_mask = None, None - - if ref_args is not None: - ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args - coords = coords[..., ref_atom_index, :] - else: - ref_coords, ref_mask, ref_atom_index, ref_token_index = ( - None, - None, - None, - None, - ) - - if operator_args is not None: - negation_mask, union_index = operator_args - else: - negation_mask, union_index = None, None - - value = self.compute_variable( - coords, - index, - ref_coords=ref_coords, - ref_mask=ref_mask, - compute_gradient=False, - ) - energy = self.compute_function( - value, *args, negation_mask=negation_mask, compute_derivative=False - ) - - if union_index is not None: - neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy) - Z = torch.zeros( - (*energy.shape[:-1], union_index.max() + 1), device=union_index.device - ).scatter_reduce( - -1, - union_index.expand_as(neg_exp_energy), - neg_exp_energy, - "sum", - ) - softmax_energy = neg_exp_energy / Z[..., union_index] - softmax_energy[Z[..., union_index] == 0] = 0 - return (energy * softmax_energy).sum(dim=-1) - - return energy.sum(dim=tuple(range(1, energy.dim()))) - - def compute_gradient(self, coords, feats, parameters): - index, args, com_args, ref_args, operator_args = self.compute_args( - feats, parameters - ) - if index.shape[1] == 0: - return torch.zeros_like(coords) - - if com_args is not None: - com_index, atom_pad_mask = com_args - unpad_coords = coords[..., atom_pad_mask, :] - unpad_com_index = com_index[atom_pad_mask] - coords = torch.zeros( - (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3), - device=coords.device, - ).scatter_reduce( - -2, - unpad_com_index.unsqueeze(-1).expand_as(unpad_coords), - unpad_coords, - "mean", - ) - com_counts = torch.bincount(com_index[atom_pad_mask]) - else: - com_index, atom_pad_mask = None, None - - if ref_args is not None: - ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args - coords = coords[..., ref_atom_index, :] - else: - ref_coords, ref_mask, ref_atom_index, ref_token_index = ( - None, - None, - None, - None, - ) - - if operator_args is not None: - negation_mask, union_index = operator_args - else: - negation_mask, union_index = None, None - - value, grad_value = self.compute_variable( - coords, - index, - ref_coords=ref_coords, - ref_mask=ref_mask, - compute_gradient=True, - ) - energy, dEnergy = self.compute_function( - value, - *args, negation_mask=negation_mask, compute_derivative=True - ) - if union_index is not None: - neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy) - Z = torch.zeros( - (*energy.shape[:-1], union_index.max() + 1), device=union_index.device - ).scatter_reduce( - -1, - union_index.expand_as(energy), - neg_exp_energy, - "sum", - ) - softmax_energy = neg_exp_energy / Z[..., union_index] - softmax_energy[Z[..., union_index] == 0] = 0 - f = torch.zeros( - (*energy.shape[:-1], union_index.max() + 1), device=union_index.device - ).scatter_reduce( - -1, - union_index.expand_as(energy), - energy * softmax_energy, - "sum", - ) - dSoftmax = ( - dEnergy - * softmax_energy - * (1 + parameters["union_lambda"] * (energy - f[..., union_index])) - ) - prod = dSoftmax.tile(grad_value.shape[-3]).unsqueeze( - -1 - ) * grad_value.flatten(start_dim=-3, end_dim=-2) - if prod.dim() > 3: - prod = prod.sum(dim=list(range(1, prod.dim() - 2))) - grad_atom = torch.zeros_like(coords).scatter_reduce( - -2, - index.flatten(start_dim=0, end_dim=1) - .unsqueeze(-1) - .expand((*coords.shape[:-2], -1, 3)), - prod, - "sum", - ) - else: - prod = dEnergy.tile(grad_value.shape[-3]).unsqueeze( - -1 - ) * grad_value.flatten(start_dim=-3, end_dim=-2) - if prod.dim() > 3: - prod = prod.sum(dim=list(range(1, prod.dim() - 2))) - grad_atom = torch.zeros_like(coords).scatter_reduce( - -2, - index.flatten(start_dim=0, end_dim=1) - .unsqueeze(-1) - .expand((*coords.shape[:-2], -1, 3)), # 9 x 516 x 3 - prod, - "sum", - ) - - if com_index is not None: - grad_atom = grad_atom[..., com_index, :] - elif ref_token_index is not None: - grad_atom = grad_atom[..., ref_token_index, :] - - return grad_atom - - def compute_parameters(self, t): - if self.parameters is None: - return None - parameters = { - name: parameter - if not isinstance(parameter, ParameterSchedule) - else parameter.compute(t) - for name, parameter in self.parameters.items() - } - return parameters - - @abstractmethod - def compute_function( - self, value, *args, negation_mask=None, compute_derivative=False - ): - raise NotImplementedError - - @abstractmethod - def compute_variable(self, coords, index, compute_gradient=False): - raise NotImplementedError - - @abstractmethod - def compute_args(self, t, feats, **parameters): - raise NotImplementedError - - def get_reference_coords(self, feats, parameters): - return None, None - - -class FlatBottomPotential(Potential): - def compute_function( - self, - value, - k, - lower_bounds, - upper_bounds, - negation_mask=None, - compute_derivative=False, - ): - if lower_bounds is None: - lower_bounds = torch.full_like(value, float("-inf")) - if upper_bounds is None: - upper_bounds = torch.full_like(value, float("inf")) - lower_bounds = lower_bounds.expand_as(value).clone() - upper_bounds = upper_bounds.expand_as(value).clone() - - if negation_mask is not None: - unbounded_below_mask = torch.isneginf(lower_bounds) - unbounded_above_mask = torch.isposinf(upper_bounds) - unbounded_mask = unbounded_below_mask + unbounded_above_mask - assert torch.all(unbounded_mask + negation_mask) - lower_bounds[~unbounded_above_mask * ~negation_mask] = upper_bounds[ - ~unbounded_above_mask * ~negation_mask - ] - upper_bounds[~unbounded_above_mask * ~negation_mask] = float("inf") - upper_bounds[~unbounded_below_mask * ~negation_mask] = lower_bounds[ - ~unbounded_below_mask * ~negation_mask - ] - lower_bounds[~unbounded_below_mask * ~negation_mask] = float("-inf") - - neg_overflow_mask = value < lower_bounds - pos_overflow_mask = value > upper_bounds - - energy = torch.zeros_like(value) - energy[neg_overflow_mask] = (k * (lower_bounds - value))[neg_overflow_mask] - energy[pos_overflow_mask] = (k * (value - upper_bounds))[pos_overflow_mask] - if not compute_derivative: - return energy - - dEnergy = torch.zeros_like(value) - dEnergy[neg_overflow_mask] = ( - -1 * k.expand_as(neg_overflow_mask)[neg_overflow_mask] - ) - dEnergy[pos_overflow_mask] = ( - 1 * k.expand_as(pos_overflow_mask)[pos_overflow_mask] - ) - - return energy, dEnergy - - -class ReferencePotential(Potential): - def compute_variable( - self, coords, index, ref_coords, ref_mask, compute_gradient=False - ): - aligned_ref_coords = weighted_rigid_align( - ref_coords.float(), - coords[:, index].float(), - ref_mask, - ref_mask, - ) - - r = coords[:, index] - aligned_ref_coords - r_norm = torch.linalg.norm(r, dim=-1) - - if not compute_gradient: - return r_norm - - r_hat = r / r_norm.unsqueeze(-1) - grad = (r_hat * ref_mask.unsqueeze(-1)).unsqueeze(1) - return r_norm, grad - - -class DistancePotential(Potential): - def compute_variable( - self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False - ): - r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1]) - r_ij_norm = torch.linalg.norm(r_ij, dim=-1) - r_hat_ij = r_ij / r_ij_norm.unsqueeze(-1) - - if not compute_gradient: - return r_ij_norm - - grad_i = r_hat_ij - grad_j = -1 * r_hat_ij - grad = torch.stack((grad_i, grad_j), dim=1) - return r_ij_norm, grad - - -class DihedralPotential(Potential): - def compute_variable( - self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False - ): - r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1]) - r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1]) - r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3]) - - n_ijk = torch.cross(r_ij, r_kj, dim=-1) - n_jkl = torch.cross(r_kj, r_kl, dim=-1) - - r_kj_norm = torch.linalg.norm(r_kj, dim=-1) - n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1) - n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1) - - sign_phi = torch.sign( - r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1) - ).squeeze(-1, -2) - phi = sign_phi * torch.arccos( - torch.clamp( - (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2) - / (n_ijk_norm * n_jkl_norm), - -1 + 1e-8, - 1 - 1e-8, - ) - ) - - if not compute_gradient: - return phi - - a = ( - (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) - ).unsqueeze(-1) - b = ( - (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) - ).unsqueeze(-1) - - grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1) - grad_l = -1 * n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1) - grad_j = (a - 1) * grad_i - b * grad_l - grad_k = (b - 1) * grad_l - a * grad_i - grad = torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1) - return phi, grad - - -class AbsDihedralPotential(DihedralPotential): - def compute_variable( - self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False - ): - if not compute_gradient: - phi = super().compute_variable( - coords, index, compute_gradient=compute_gradient - ) - phi = torch.abs(phi) - return phi - - phi, grad = super().compute_variable( - coords, index, compute_gradient=compute_gradient - ) - grad[(phi < 0)[..., None, :, None].expand_as(grad)] *= -1 - phi = torch.abs(phi) - - return phi, grad - - -class PoseBustersPotential(FlatBottomPotential, DistancePotential): - def compute_args(self, feats, parameters): - pair_index = feats["rdkit_bounds_index"][0] - lower_bounds = feats["rdkit_lower_bounds"][0].clone() - upper_bounds = feats["rdkit_upper_bounds"][0].clone() - bond_mask = feats["rdkit_bounds_bond_mask"][0] - angle_mask = feats["rdkit_bounds_angle_mask"][0] - - lower_bounds[bond_mask * ~angle_mask] *= 1.0 - parameters["bond_buffer"] - upper_bounds[bond_mask * ~angle_mask] *= 1.0 + parameters["bond_buffer"] - lower_bounds[~bond_mask * angle_mask] *= 1.0 - parameters["angle_buffer"] - upper_bounds[~bond_mask * angle_mask] *= 1.0 + parameters["angle_buffer"] - lower_bounds[bond_mask * angle_mask] *= 1.0 - min( - parameters["bond_buffer"], parameters["angle_buffer"] - ) - upper_bounds[bond_mask * angle_mask] *= 1.0 + min( - parameters["bond_buffer"], parameters["angle_buffer"] - ) - lower_bounds[~bond_mask * ~angle_mask] *= 1.0 - parameters["clash_buffer"] - upper_bounds[~bond_mask * ~angle_mask] = float("inf") - - vdw_radii = torch.zeros( - const.num_elements, dtype=torch.float32, device=pair_index.device - ) - vdw_radii[1:119] = torch.tensor( - const.vdw_radii, dtype=torch.float32, device=pair_index.device - ) - atom_vdw_radii = ( - feats["ref_element"].float() @ vdw_radii.unsqueeze(-1) - ).squeeze(-1)[0] - bond_cutoffs = 0.35 + atom_vdw_radii[pair_index].mean(dim=0) - lower_bounds[~bond_mask] = torch.max(lower_bounds[~bond_mask], bond_cutoffs[~bond_mask]) - upper_bounds[bond_mask] = torch.min(upper_bounds[bond_mask], bond_cutoffs[bond_mask]) - - k = torch.ones_like(lower_bounds) - - return pair_index, (k, lower_bounds, upper_bounds), None, None, None - - -class ConnectionsPotential(FlatBottomPotential, DistancePotential): - def compute_args(self, feats, parameters): - pair_index = feats["connected_atom_index"][0] - lower_bounds = None - upper_bounds = torch.full( - (pair_index.shape[1],), parameters["buffer"], device=pair_index.device - ) - k = torch.ones_like(upper_bounds) - - return pair_index, (k, lower_bounds, upper_bounds), None, None, None - - -class VDWOverlapPotential(FlatBottomPotential, DistancePotential): - def compute_args(self, feats, parameters): - atom_chain_id = ( - torch.bmm( - feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float() - ) - .squeeze(-1) - .long() - )[0] - atom_pad_mask = feats["atom_pad_mask"][0].bool() - chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask]) - single_ion_mask = (chain_sizes > 1)[atom_chain_id] - - vdw_radii = torch.zeros( - const.num_elements, dtype=torch.float32, device=atom_chain_id.device - ) - vdw_radii[1:119] = torch.tensor( - const.vdw_radii, dtype=torch.float32, device=atom_chain_id.device - ) - atom_vdw_radii = ( - feats["ref_element"].float() @ vdw_radii.unsqueeze(-1) - ).squeeze(-1)[0] - - pair_index = torch.triu_indices( - atom_chain_id.shape[0], - atom_chain_id.shape[0], - 1, - device=atom_chain_id.device, - ) - - pair_pad_mask = atom_pad_mask[pair_index].all(dim=0) - pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]] - - num_chains = atom_chain_id.max() + 1 - connected_chain_index = feats["connected_chain_index"][0] - connected_chain_matrix = torch.eye( - num_chains, device=atom_chain_id.device, dtype=torch.bool - ) - connected_chain_matrix[connected_chain_index[0], connected_chain_index[1]] = ( - True - ) - connected_chain_matrix[connected_chain_index[1], connected_chain_index[0]] = ( - True - ) - connected_chain_mask = connected_chain_matrix[ - atom_chain_id[pair_index[0]], atom_chain_id[pair_index[1]] - ] - - pair_index = pair_index[ - :, pair_pad_mask * pair_ion_mask * ~connected_chain_mask - ] - - lower_bounds = atom_vdw_radii[pair_index].sum(dim=0) * ( - 1.0 - parameters["buffer"] - ) - upper_bounds = None - k = torch.ones_like(lower_bounds) - - return pair_index, (k, lower_bounds, upper_bounds), None, None, None - - -class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential): - def compute_args(self, feats, parameters): - atom_chain_id = ( - torch.bmm( - feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float() - ) - .squeeze(-1) - .long() - )[0] - atom_pad_mask = feats["atom_pad_mask"][0].bool() - chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask]) - single_ion_mask = chain_sizes > 1 - - pair_index = feats["symmetric_chain_index"][0] - pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]] - pair_index = pair_index[:, pair_ion_mask] - lower_bounds = torch.full( - (pair_index.shape[1],), - parameters["buffer"], - dtype=torch.float32, - device=pair_index.device, - ) - upper_bounds = None - k = torch.ones_like(lower_bounds) - - return ( - pair_index, - (k, lower_bounds, upper_bounds), - (atom_chain_id, atom_pad_mask), - None, - None, - ) - - -class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential): - def compute_args(self, feats, parameters): - stereo_bond_index = feats["stereo_bond_index"][0] - stereo_bond_orientations = feats["stereo_bond_orientations"][0].bool() - - lower_bounds = torch.zeros( - stereo_bond_orientations.shape, device=stereo_bond_orientations.device - ) - upper_bounds = torch.zeros( - stereo_bond_orientations.shape, device=stereo_bond_orientations.device - ) - lower_bounds[stereo_bond_orientations] = torch.pi - parameters["buffer"] - upper_bounds[stereo_bond_orientations] = float("inf") - lower_bounds[~stereo_bond_orientations] = float("-inf") - upper_bounds[~stereo_bond_orientations] = parameters["buffer"] - - k = torch.ones_like(lower_bounds) - - return stereo_bond_index, (k, lower_bounds, upper_bounds), None, None, None - - -class ChiralAtomPotential(FlatBottomPotential, DihedralPotential): - def compute_args(self, feats, parameters): - chiral_atom_index = feats["chiral_atom_index"][0] - chiral_atom_orientations = feats["chiral_atom_orientations"][0].bool() - - lower_bounds = torch.zeros( - chiral_atom_orientations.shape, device=chiral_atom_orientations.device - ) - upper_bounds = torch.zeros( - chiral_atom_orientations.shape, device=chiral_atom_orientations.device - ) - lower_bounds[chiral_atom_orientations] = parameters["buffer"] - upper_bounds[chiral_atom_orientations] = float("inf") - upper_bounds[~chiral_atom_orientations] = -1 * parameters["buffer"] - lower_bounds[~chiral_atom_orientations] = float("-inf") - - k = torch.ones_like(lower_bounds) - return chiral_atom_index, (k, lower_bounds, upper_bounds), None, None, None - - -class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential): - def compute_args(self, feats, parameters): - double_bond_index = feats["planar_bond_index"][0].T - double_bond_improper_index = torch.tensor( - [ - [1, 2, 3, 0], - [4, 5, 0, 3], - ], - device=double_bond_index.device, - ).T - improper_index = ( - double_bond_index[:, double_bond_improper_index] - .swapaxes(0, 1) - .flatten(start_dim=1) - ) - lower_bounds = None - upper_bounds = torch.full( - (improper_index.shape[1],), - parameters["buffer"], - device=improper_index.device, - ) - k = torch.ones_like(upper_bounds) - - return improper_index, (k, lower_bounds, upper_bounds), None, None, None - - -class TemplateReferencePotential(FlatBottomPotential, ReferencePotential): - def compute_args(self, feats, parameters): - if "template_mask_cb" not in feats or "template_force" not in feats: - return torch.empty([1, 0]), None, None, None, None - - template_mask = feats["template_mask_cb"][feats["template_force"]] - if template_mask.shape[0] == 0: - return torch.empty([1, 0]), None, None, None, None - - ref_coords = feats["template_cb"][feats["template_force"]].clone() - ref_mask = feats["template_mask_cb"][feats["template_force"]].clone() - ref_atom_index = ( - torch.bmm( - feats["token_to_rep_atom"].float(), - torch.arange( - feats["atom_pad_mask"].shape[1], - device=feats["atom_pad_mask"].device, - dtype=torch.float32, - )[None, :, None], - ) - .squeeze(-1) - .long() - )[0] - ref_token_index = ( - torch.bmm( - feats["atom_to_token"].float(), - feats["token_index"].unsqueeze(-1).float(), - ) - .squeeze(-1) - .long() - )[0] - - index = torch.arange( - template_mask.shape[-1], dtype=torch.long, device=template_mask.device - )[None] - upper_bounds = torch.full( - template_mask.shape, float("inf"), device=index.device, dtype=torch.float32 - ) - ref_idxs = torch.argwhere(template_mask).T - upper_bounds[ref_idxs.unbind()] = feats["template_force_threshold"][ - feats["template_force"] - ][ref_idxs[0]] - - lower_bounds = None - k = torch.ones_like(upper_bounds) - return ( - index, - (k, lower_bounds, upper_bounds), - None, - (ref_coords, ref_mask, ref_atom_index, ref_token_index), - None, - ) - - -class ContactPotentital(FlatBottomPotential, DistancePotential): - def compute_args(self, feats, parameters): - index = feats["contact_pair_index"][0] - union_index = feats["contact_union_index"][0] - negation_mask = feats["contact_negation_mask"][0] - lower_bounds = None - upper_bounds = feats["contact_thresholds"][0].clone() - k = torch.ones_like(upper_bounds) - return ( - index, - (k, lower_bounds, upper_bounds), - None, - None, - (negation_mask, union_index), - ) - - -def get_potentials(steering_args, boltz2=False): - potentials = [] - if steering_args["fk_steering"] or steering_args["physical_guidance_update"]: - potentials.extend( - [ - SymmetricChainCOMPotential( - parameters={ - "guidance_interval": 4, - "guidance_weight": 0.5 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 0.5, - "buffer": ExponentialInterpolation( - start=1.0, end=5.0, alpha=-2.0 - ), - } - ), - VDWOverlapPotential( - parameters={ - "guidance_interval": 5, - "guidance_weight": ( - PiecewiseStepFunction(thresholds=[0.4], values=[0.125, 0.0]) - if steering_args["physical_guidance_update"] - else 0.0 - ), - "resampling_weight": PiecewiseStepFunction( - thresholds=[0.6], values=[0.01, 0.0] - ), - "buffer": 0.225, - } - ), - ConnectionsPotential( - parameters={ - "guidance_interval": 1, - "guidance_weight": 0.15 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 1.0, - "buffer": 2.0, - } - ), - PoseBustersPotential( - parameters={ - "guidance_interval": 1, - "guidance_weight": 0.01 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 0.1, - "bond_buffer": 0.125, - "angle_buffer": 0.125, - "clash_buffer": 0.10, - } - ), - ChiralAtomPotential( - parameters={ - "guidance_interval": 1, - "guidance_weight": 0.1 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 1.0, - "buffer": 0.52360, - } - ), - StereoBondPotential( - parameters={ - "guidance_interval": 1, - "guidance_weight": 0.05 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 1.0, - "buffer": 0.52360, - } - ), - PlanarBondPotential( - parameters={ - "guidance_interval": 1, - "guidance_weight": 0.05 - if steering_args["physical_guidance_update"] - else 0.0, - "resampling_weight": 1.0, - "buffer": 0.26180, - } - ), - ] - ) - if boltz2 and ( - steering_args["fk_steering"] or steering_args["contact_guidance_update"] - ): - potentials.extend( - [ - ContactPotentital( - parameters={ - "guidance_interval": 4, - "guidance_weight": ( - PiecewiseStepFunction( - thresholds=[0.25, 0.75], values=[0.0, 0.5, 1.0] - ) - if steering_args["contact_guidance_update"] - else 0.0 - ), - "resampling_weight": 1.0, - "union_lambda": ExponentialInterpolation( - start=8.0, end=0.0, alpha=-2.0 - ), - } - ), - TemplateReferencePotential( - parameters={ - "guidance_interval": 2, - "guidance_weight": 0.1 - if steering_args["contact_guidance_update"] - else 0.0, - "resampling_weight": 1.0, - } - ), - ] - ) - return potentials diff --git a/fastplms/boltz/vb_potentials_schedules.py b/fastplms/boltz/vb_potentials_schedules.py deleted file mode 100644 index 564249a..0000000 --- a/fastplms/boltz/vb_potentials_schedules.py +++ /dev/null @@ -1,37 +0,0 @@ -import math -from abc import ABC - - -class ParameterSchedule(ABC): - def compute(self, t): - raise NotImplementedError - - -class ExponentialInterpolation(ParameterSchedule): - def __init__(self, start, end, alpha): - self.start = start - self.end = end - self.alpha = alpha - - def compute(self, t): - if self.alpha != 0: - return self.start + (self.end - self.start) * ( - math.exp(self.alpha * t) - 1 - ) / (math.exp(self.alpha) - 1) - else: - return self.start + (self.end - self.start) * t - - -class PiecewiseStepFunction(ParameterSchedule): - def __init__(self, thresholds, values): - self.thresholds = thresholds - self.values = values - - def compute(self, t): - assert len(self.thresholds) > 0 - assert len(self.values) == len(self.thresholds) + 1 - - idx = 0 - while idx < len(self.thresholds) and t > self.thresholds[idx]: - idx += 1 - return self.values[idx] diff --git a/fastplms/boltz/vb_tri_attn_attention.py b/fastplms/boltz/vb_tri_attn_attention.py deleted file mode 100644 index 6fba28b..0000000 --- a/fastplms/boltz/vb_tri_attn_attention.py +++ /dev/null @@ -1,189 +0,0 @@ -# Copyright 2021 AlQuraishi Laboratory -# Copyright 2021 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from functools import partial, partialmethod -from typing import Optional - -import torch -import torch.nn as nn - -from .vb_tri_attn_primitives import ( - Attention, - LayerNorm, - Linear, -) -from .vb_tri_attn_utils import ( - chunk_layer, - permute_final_dims, -) - - -class TriangleAttention(nn.Module): - """Implement Algorithm 12.""" - - def __init__( - self, - c_in: int, - c_hidden: int, - no_heads: int, - starting: bool = True, - inf: float = 1e9, - ) -> None: - super().__init__() - - self.c_in = c_in - self.c_hidden = c_hidden - self.no_heads = no_heads - self.starting = starting - self.inf = inf - - self.layer_norm = LayerNorm(self.c_in) - - self.linear = Linear(c_in, self.no_heads, bias=False, init="normal") - - self.mha = Attention( - self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads - ) - - @torch.jit.ignore - def _chunk( - self, - x: torch.Tensor, - tri_bias: torch.Tensor, - mask_bias: torch.Tensor, - mask: torch.Tensor, - chunk_size: int, - use_kernels: bool = False, - ) -> torch.Tensor: - """Compute triangle attention. - - Parameters - ---------- - x : torch.Tensor - Input tensor of shape [*, I, J, C_in] - biases : list[torch.Tensor] - List of bias tensors of shape [*, H, I, J] - chunk_size : int - Size of chunks for memory efficient computation - use_kernels : bool, default=False - Whether to use optimized CUDA kernels - - Returns - ------- - torch.Tensor - Output tensor of shape [*, I, J, C_in] - - """ - mha_inputs = { - "q_x": x, - "kv_x": x, - "tri_bias": tri_bias, - "mask_bias": mask_bias, - "mask": mask, - } - - return chunk_layer( - partial( - self.mha, - use_kernels=use_kernels, - ), - mha_inputs, - chunk_size=chunk_size, - no_batch_dims=len(x.shape[:-2]), - _out=None, - ) - - def forward( - self, - x: torch.Tensor, - mask: Optional[torch.Tensor] = None, - chunk_size: Optional[int] = None, - use_kernels: bool = False, - ) -> torch.Tensor: - """Compute triangle attention. - - Parameters - ---------- - x : torch.Tensor - Input tensor of shape [*, I, J, C_in] - mask : torch.Tensor, optional - Attention mask of shape [*, I, J] - chunk_size : int, optional - Size of chunks for memory efficient computation - use_kernels : bool, default=False - Whether to use optimized CUDA kernels - - Returns - ------- - torch.Tensor - Output tensor of shape [*, I, J, C_in] - - """ - if mask is None: - # [*, I, J] - mask = x.new_ones( - x.shape[:-1], - ) - - if not self.starting: - x = x.transpose(-2, -3) - mask = mask.transpose(-1, -2) - - # [*, I, J, C_in] - x = self.layer_norm(x) - - # [*, I, 1, 1, J] - mask = mask[..., :, None, None, :] - mask_bias = self.inf * (mask - 1) - - # [*, H, I, J] - triangle_bias = permute_final_dims(self.linear(x), (2, 0, 1)) - - # [*, 1, H, I, J] - triangle_bias = triangle_bias.unsqueeze(-4) - - if chunk_size is not None and not use_kernels: - x = self._chunk( - x, - triangle_bias, - mask_bias, - mask, - chunk_size, - use_kernels=use_kernels, - ) - else: - x = self.mha( - x, - x, - triangle_bias, - mask_bias, - mask, - use_kernels=use_kernels, - ) - - if not self.starting: - x = x.transpose(-2, -3) - - return x - - -# Implements Algorithm 13 -TriangleAttentionStartingNode = TriangleAttention - - -class TriangleAttentionEndingNode(TriangleAttention): - """Implement Algorithm 14.""" - - __init__ = partialmethod(TriangleAttention.__init__, starting=False) diff --git a/fastplms/boltz/vb_tri_attn_primitives.py b/fastplms/boltz/vb_tri_attn_primitives.py deleted file mode 100644 index e8b2cee..0000000 --- a/fastplms/boltz/vb_tri_attn_primitives.py +++ /dev/null @@ -1,411 +0,0 @@ -# Copyright 2021 AlQuraishi Laboratory -# Copyright 2021 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -import importlib -from typing import Callable, List, Optional, Tuple - -import torch -from einops import rearrange -from torch import nn - -from . import vb_layers_initialize as initialize -from .vb_tri_attn_utils import ( - flatten_final_dims, - permute_final_dims, -) - - -class Linear(nn.Linear): - """ - A Linear layer with built-in nonstandard initializations. Called just - like torch.nn.Linear. - - Implements the initializers in 1.11.4, plus some additional ones found - in the code. - """ - - def __init__( - self, - in_dim: int, - out_dim: int, - bias: bool = True, - init: str = "default", - init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None, - precision=None, - ): - """Initialize the linear layer. - - Parameters - ---------- - in_dim : int - The final dimension of inputs to the layer - out_dim : int - The final dimension of layer outputs - bias : bool, default=True - Whether to learn an additive bias - init : str, default='default' - The initializer to use. Choose from: - - - "default": LeCun fan-in truncated normal initialization - - "relu": He initialization w/ truncated normal distribution - - "glorot": Fan-average Glorot uniform initialization - - "gating": Weights=0, Bias=1 - - "normal": Normal initialization with std=1/sqrt(fan_in) - - "final": Weights=0, Bias=0 - - Overridden by init_fn if the latter is not None. - init_fn : callable, optional - A custom initializer taking weight and bias as inputs. - Overrides init if not None. - - """ - super().__init__(in_dim, out_dim, bias=bias) - - if bias: - with torch.no_grad(): - self.bias.fill_(0) - - with torch.no_grad(): - if init_fn is not None: - init_fn(self.weight, self.bias) - else: - if init == "default": - initialize.lecun_normal_init_(self.weight) - elif init == "relu": - initialize.he_normal_init_(self.weight) - elif init == "glorot": - initialize.glorot_uniform_init_(self.weight) - elif init == "gating": - initialize.gating_init_(self.weight) - if bias: - self.bias.fill_(1.0) - elif init == "normal": - initialize.normal_init_(self.weight) - elif init == "final": - initialize.final_init_(self.weight) - else: - raise ValueError("Invalid init string.") - - self.precision = precision - - def forward(self, input: torch.Tensor) -> torch.Tensor: - d = input.dtype - if self.precision is not None: - with torch.autocast("cuda", enabled=False): - bias = ( - self.bias.to(dtype=self.precision) - if self.bias is not None - else None - ) - return nn.functional.linear( - input.to(dtype=self.precision), - self.weight.to(dtype=self.precision), - bias, - ).to(dtype=d) - - if d is torch.bfloat16: - with torch.autocast("cuda", enabled=False): - bias = self.bias.to(dtype=d) if self.bias is not None else None - return nn.functional.linear(input, self.weight.to(dtype=d), bias) - - return nn.functional.linear(input, self.weight, self.bias) - - -class LayerNorm(nn.Module): - def __init__(self, c_in, eps=1e-5): - super(LayerNorm, self).__init__() - - self.c_in = (c_in,) - self.eps = eps - - self.weight = nn.Parameter(torch.ones(c_in)) - self.bias = nn.Parameter(torch.zeros(c_in)) - - def forward(self, x): - d = x.dtype - if d is torch.bfloat16: - with torch.autocast("cuda", enabled=False): - out = nn.functional.layer_norm( - x, - self.c_in, - self.weight.to(dtype=d), - self.bias.to(dtype=d), - self.eps, - ) - else: - out = nn.functional.layer_norm( - x, - self.c_in, - self.weight, - self.bias, - self.eps, - ) - - return out - - -@torch.jit.ignore -def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor: - """ - Softmax, but without automatic casting to fp32 when the input is of - type bfloat16 - """ - d = t.dtype - if d is torch.bfloat16: - with torch.autocast("cuda", enabled=False): - s = torch.nn.functional.softmax(t, dim=dim) - else: - s = torch.nn.functional.softmax(t, dim=dim) - - return s - - -# @torch.jit.script -def _attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - biases: List[torch.Tensor], -) -> torch.Tensor: - # [*, H, C_hidden, K] - key = permute_final_dims(key, (1, 0)) - - # [*, H, Q, K] - a = torch.matmul(query, key) - - for b in biases: - a += b - - a = softmax_no_cast(a, -1) - - # [*, H, Q, C_hidden] - a = torch.matmul(a, value) - - return a - - -@torch.compiler.disable -def kernel_triangular_attn(q, k, v, tri_bias, mask, scale): - triangle_module = importlib.import_module("cuequivariance_torch.primitives.triangle") - triangle_attention = triangle_module.triangle_attention - return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale) - - -class Attention(nn.Module): - """ - Standard multi-head attention using AlphaFold's default layer - initialization. Allows multiple bias vectors. - """ - - def __init__( - self, - c_q: int, - c_k: int, - c_v: int, - c_hidden: int, - no_heads: int, - gating: bool = True, - ): - """Initialize the attention layer. - - Parameters - ---------- - c_q : int - Input dimension of query data - c_k : int - Input dimension of key data - c_v : int - Input dimension of value data - c_hidden : int - Per-head hidden dimension - no_heads : int - Number of attention heads - gating : bool, default=True - Whether the output should be gated using query data - - """ - super().__init__() - - self.c_q = c_q - self.c_k = c_k - self.c_v = c_v - self.c_hidden = c_hidden - self.no_heads = no_heads - self.gating = gating - - # DISCREPANCY: c_hidden is not the per-head channel dimension, as - # stated in the supplement, but the overall channel dimension. - - self.linear_q = Linear( - self.c_q, self.c_hidden * self.no_heads, bias=False, init="glorot" - ) - self.linear_k = Linear( - self.c_k, self.c_hidden * self.no_heads, bias=False, init="glorot" - ) - self.linear_v = Linear( - self.c_v, self.c_hidden * self.no_heads, bias=False, init="glorot" - ) - self.linear_o = Linear( - self.c_hidden * self.no_heads, self.c_q, bias=False, init="final" - ) - - self.linear_g = None - if self.gating: - self.linear_g = Linear( - self.c_q, self.c_hidden * self.no_heads, bias=False, init="gating" - ) - - self.sigmoid = nn.Sigmoid() - - def _prep_qkv( - self, q_x: torch.Tensor, kv_x: torch.Tensor, apply_scale: bool = True - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # [*, Q/K/V, H * C_hidden] - q = self.linear_q(q_x) - k = self.linear_k(kv_x) - v = self.linear_v(kv_x) - - # [*, Q/K, H, C_hidden] - q = q.view(q.shape[:-1] + (self.no_heads, -1)) - k = k.view(k.shape[:-1] + (self.no_heads, -1)) - v = v.view(v.shape[:-1] + (self.no_heads, -1)) - - # [*, H, Q/K, C_hidden] - q = q.transpose(-2, -3) - k = k.transpose(-2, -3) - v = v.transpose(-2, -3) - - if apply_scale: - q /= math.sqrt(self.c_hidden) - - return q, k, v - - def _wrap_up(self, o: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor: - if self.linear_g is not None: - g = self.sigmoid(self.linear_g(q_x)) - - # [*, Q, H, C_hidden] - g = g.view(g.shape[:-1] + (self.no_heads, -1)) - o = o * g - - # [*, Q, H * C_hidden] - o = flatten_final_dims(o, 2) - - # [*, Q, C_q] - o = self.linear_o(o) - - return o - - def forward( - self, - q_x: torch.Tensor, - kv_x: torch.Tensor, - tri_bias: torch.Tensor, - mask_bias: torch.Tensor, - mask: torch.Tensor, - use_kernels: bool = False, - ) -> torch.Tensor: - """Compute attention. - - Parameters - ---------- - q_x : torch.Tensor - [*, Q, C_q] query data - kv_x : torch.Tensor - [*, K, C_k] key data - tri_bias : torch.Tensor - [*, H, Q, K] triangular bias - mask_bias : torch.Tensor - [*, H, Q, K] mask bias - mask : torch.Tensor - [*, Q, K] mask - use_kernels : bool, default=False - Whether to use optimized CUDA kernels - - Returns - ------- - [*, Q, C_q] attention update - - """ - # Attention kernel applies scaling internally - q, k, v = self._prep_qkv( - q_x, - kv_x, - apply_scale=not use_kernels, - ) - - if use_kernels: - scale = 1.0 / math.sqrt(self.c_hidden) - o = kernel_triangular_attn( - q, - k, - v, - tri_bias=tri_bias, - mask=mask.bool(), - scale=scale, - ) - o = o.transpose(-2, -3) - else: - biases = [mask_bias, tri_bias] - o = _attention(q, k, v, biases) - o = o.transpose(-2, -3) - - o = self._wrap_up(o, q_x) - - return o - - -def _trifast_attn(q, k, v, biases): - orig_n_dims = len(q.shape) - - if len(biases) != 2: - raise ValueError(f"Trifast expects two bias terms, found {len(biases)}") - - mask, b = biases - - if len(b.shape) == 5: - # Sometimes there is an extra batch dim -- why? - b = b.squeeze(1) - - if orig_n_dims == 4: - # add fake batch dim - q = q.unsqueeze(0) - k = k.unsqueeze(0) - v = v.unsqueeze(0) - # b = b.unsqueeze(0) not sure why this and only this has a batch dim? - mask = mask.unsqueeze(0) - - if len(q.shape) != 5: - raise ValueError(f"Trifast expects q/k/v to be 5D, found {len(q.shape)}") - - # Reorder q/k/v - q = rearrange(q, "b i h j d -> b h i j d") - k = rearrange(k, "b i h j d -> b h i j d") - v = rearrange(v, "b i h j d -> b h i j d") - - # Make mask the right shape. - mask = rearrange(mask, "b i () () j -> b i j").bool() - - # Delay import to here to avoid initializing cuda too early - from trifast import triangle_attention - - o = triangle_attention(q, k, v, b, mask) - o = rearrange(o, "b h i j d -> b i j h d") - - # Remove the batch dim if we added it. - if orig_n_dims == 4: - o = o.squeeze(0) - return o diff --git a/fastplms/boltz/vb_tri_attn_utils.py b/fastplms/boltz/vb_tri_attn_utils.py deleted file mode 100644 index 89899da..0000000 --- a/fastplms/boltz/vb_tri_attn_utils.py +++ /dev/null @@ -1,380 +0,0 @@ -# Copyright 2021 AlQuraishi Laboratory -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from functools import partial -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple - -import torch - - -def add(m1, m2, inplace): - # The first operation in a checkpoint can't be in-place, but it's - # nice to have in-place addition during inference. Thus... - if not inplace: - m1 = m1 + m2 - else: - m1 += m2 - - return m1 - - -def permute_final_dims(tensor: torch.Tensor, inds: List[int]): - zero_index = -1 * len(inds) - first_inds = list(range(len(tensor.shape[:zero_index]))) - return tensor.permute(first_inds + [zero_index + i for i in inds]) - - -def is_fp16_enabled(): - # Autocast world - fp16_enabled = torch.get_autocast_gpu_dtype() == torch.float16 - fp16_enabled = fp16_enabled and torch.is_autocast_enabled() - - return fp16_enabled - - -# With tree_map, a poor man's JAX tree_map -def dict_map(fn, dic, leaf_type): - new_dict = {} - for k, v in dic.items(): - if type(v) is dict: - new_dict[k] = dict_map(fn, v, leaf_type) - else: - new_dict[k] = tree_map(fn, v, leaf_type) - - return new_dict - - -def tree_map(fn, tree, leaf_type): - if isinstance(tree, dict): - return dict_map(fn, tree, leaf_type) - elif isinstance(tree, list): - return [tree_map(fn, x, leaf_type) for x in tree] - elif isinstance(tree, tuple): - return tuple([tree_map(fn, x, leaf_type) for x in tree]) - elif isinstance(tree, leaf_type): - return fn(tree) - else: - raise ValueError(f"Tree of type {type(tree)} not supported") - - -tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor) - - -def flatten_final_dims(t: torch.Tensor, no_dims: int): - return t.reshape(t.shape[:-no_dims] + (-1,)) - - -def _fetch_dims(tree): - shapes = [] - tree_type = type(tree) - if tree_type is dict: - for v in tree.values(): - shapes.extend(_fetch_dims(v)) - elif tree_type is list or tree_type is tuple: - for t in tree: - shapes.extend(_fetch_dims(t)) - elif tree_type is torch.Tensor: - shapes.append(tree.shape) - else: - raise ValueError("Not supported") - - return shapes - - -@torch.jit.ignore -def _flat_idx_to_idx( - flat_idx: int, - dims: Tuple[int], -) -> Tuple[int]: - idx = [] - for d in reversed(dims): - idx.append(flat_idx % d) - flat_idx = flat_idx // d - - return tuple(reversed(idx)) - - -@torch.jit.ignore -def _get_minimal_slice_set( - start: Sequence[int], - end: Sequence[int], - dims: int, - start_edges: Optional[Sequence[bool]] = None, - end_edges: Optional[Sequence[bool]] = None, -) -> Sequence[Tuple[int]]: - """ - Produces an ordered sequence of tensor slices that, when used in - sequence on a tensor with shape dims, yields tensors that contain every - leaf in the contiguous range [start, end]. Care is taken to yield a - short sequence of slices, and perhaps even the shortest possible (I'm - pretty sure it's the latter). - - end is INCLUSIVE. - """ - - # start_edges and end_edges both indicate whether, starting from any given - # dimension, the start/end index is at the top/bottom edge of the - # corresponding tensor, modeled as a tree - def reduce_edge_list(l): - tally = 1 - for i in range(len(l)): - reversed_idx = -1 * (i + 1) - l[reversed_idx] *= tally - tally = l[reversed_idx] - - if start_edges is None: - start_edges = [s == 0 for s in start] - reduce_edge_list(start_edges) - if end_edges is None: - end_edges = [e == (d - 1) for e, d in zip(end, dims)] - reduce_edge_list(end_edges) - - # Base cases. Either start/end are empty and we're done, or the final, - # one-dimensional tensor can be simply sliced - if len(start) == 0: - return [tuple()] - elif len(start) == 1: - return [(slice(start[0], end[0] + 1),)] - - slices = [] - path = [] - - # Dimensions common to start and end can be selected directly - for s, e in zip(start, end): - if s == e: - path.append(slice(s, s + 1)) - else: - break - - path = tuple(path) - divergence_idx = len(path) - - # start == end, and we're done - if divergence_idx == len(dims): - return [tuple(path)] - - def upper(): - sdi = start[divergence_idx] - return [ - path + (slice(sdi, sdi + 1),) + s - for s in _get_minimal_slice_set( - start[divergence_idx + 1 :], - [d - 1 for d in dims[divergence_idx + 1 :]], - dims[divergence_idx + 1 :], - start_edges=start_edges[divergence_idx + 1 :], - end_edges=[1 for _ in end_edges[divergence_idx + 1 :]], - ) - ] - - def lower(): - edi = end[divergence_idx] - return [ - path + (slice(edi, edi + 1),) + s - for s in _get_minimal_slice_set( - [0 for _ in start[divergence_idx + 1 :]], - end[divergence_idx + 1 :], - dims[divergence_idx + 1 :], - start_edges=[1 for _ in start_edges[divergence_idx + 1 :]], - end_edges=end_edges[divergence_idx + 1 :], - ) - ] - - # If both start and end are at the edges of the subtree rooted at - # divergence_idx, we can just select the whole subtree at once - if start_edges[divergence_idx] and end_edges[divergence_idx]: - slices.append(path + (slice(start[divergence_idx], end[divergence_idx] + 1),)) - # If just start is at the edge, we can grab almost all of the subtree, - # treating only the ragged bottom edge as an edge case - elif start_edges[divergence_idx]: - slices.append(path + (slice(start[divergence_idx], end[divergence_idx]),)) - slices.extend(lower()) - # Analogous to the previous case, but the top is ragged this time - elif end_edges[divergence_idx]: - slices.extend(upper()) - slices.append( - path + (slice(start[divergence_idx] + 1, end[divergence_idx] + 1),) - ) - # If both sides of the range are ragged, we need to handle both sides - # separately. If there's contiguous meat in between them, we can index it - # in one big chunk - else: - slices.extend(upper()) - middle_ground = end[divergence_idx] - start[divergence_idx] - if middle_ground > 1: - slices.append( - path + (slice(start[divergence_idx] + 1, end[divergence_idx]),) - ) - slices.extend(lower()) - - return [tuple(s) for s in slices] - - -@torch.jit.ignore -def _chunk_slice( - t: torch.Tensor, - flat_start: int, - flat_end: int, - no_batch_dims: int, -) -> torch.Tensor: - """ - Equivalent to - - t.reshape((-1,) + t.shape[no_batch_dims:])[flat_start:flat_end] - - but without the need for the initial reshape call, which can be - memory-intensive in certain situations. The only reshape operations - in this function are performed on sub-tensors that scale with - (flat_end - flat_start), the chunk size. - """ - - batch_dims = t.shape[:no_batch_dims] - start_idx = list(_flat_idx_to_idx(flat_start, batch_dims)) - # _get_minimal_slice_set is inclusive - end_idx = list(_flat_idx_to_idx(flat_end - 1, batch_dims)) - - # Get an ordered list of slices to perform - slices = _get_minimal_slice_set( - start_idx, - end_idx, - batch_dims, - ) - - sliced_tensors = [t[s] for s in slices] - - return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors]) - - -def chunk_layer( - layer: Callable, - inputs: Dict[str, Any], - chunk_size: int, - no_batch_dims: int, - low_mem: bool = False, - _out: Any = None, - _add_into_out: bool = False, -) -> Any: - """ - Implements the "chunking" procedure described in section 1.11.8. - - Layer outputs and inputs are assumed to be simple "pytrees," - consisting only of (arbitrarily nested) lists, tuples, and dicts with - torch.Tensor leaves. - - Args: - layer: - The layer to be applied chunk-wise - inputs: - A (non-nested) dictionary of keyworded inputs. All leaves must - be tensors and must share the same batch dimensions. - chunk_size: - The number of sub-batches per chunk. If multiple batch - dimensions are specified, a "sub-batch" is defined as a single - indexing of all batch dimensions simultaneously (s.t. the - number of sub-batches is the product of the batch dimensions). - no_batch_dims: - How many of the initial dimensions of each input tensor can - be considered batch dimensions. - low_mem: - Avoids flattening potentially large input tensors. Unnecessary - in most cases, and is ever so slightly slower than the default - setting. - Returns: - The reassembled output of the layer on the inputs. - """ - if not (len(inputs) > 0): - raise ValueError("Must provide at least one input") - - initial_dims = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)] - orig_batch_dims = tuple([max(s) for s in zip(*initial_dims)]) - - def _prep_inputs(t): - if not low_mem: - if not sum(t.shape[:no_batch_dims]) == no_batch_dims: - t = t.expand(orig_batch_dims + t.shape[no_batch_dims:]) - t = t.reshape(-1, *t.shape[no_batch_dims:]) - else: - t = t.expand(orig_batch_dims + t.shape[no_batch_dims:]) - return t - - prepped_inputs = tensor_tree_map(_prep_inputs, inputs) - prepped_outputs = None - if _out is not None: - reshape_fn = lambda t: t.view([-1] + list(t.shape[no_batch_dims:])) - prepped_outputs = tensor_tree_map(reshape_fn, _out) - - flat_batch_dim = 1 - for d in orig_batch_dims: - flat_batch_dim *= d - - no_chunks = flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0) - - i = 0 - out = prepped_outputs - for _ in range(no_chunks): - # Chunk the input - if not low_mem: - select_chunk = lambda t: t[i : i + chunk_size] if t.shape[0] != 1 else t - else: - select_chunk = partial( - _chunk_slice, - flat_start=i, - flat_end=min(flat_batch_dim, i + chunk_size), - no_batch_dims=len(orig_batch_dims), - ) - - chunks = tensor_tree_map(select_chunk, prepped_inputs) - - # Run the layer on the chunk - output_chunk = layer(**chunks) - - # Allocate space for the output - if out is None: - allocate = lambda t: t.new_zeros((flat_batch_dim,) + t.shape[1:]) - out = tensor_tree_map(allocate, output_chunk) - - # Put the chunk in its pre-allocated space - out_type = type(output_chunk) - if out_type is dict: - - def assign(d1, d2): - for k, v in d1.items(): - if type(v) is dict: - assign(v, d2[k]) - else: - if _add_into_out: - v[i : i + chunk_size] += d2[k] - else: - v[i : i + chunk_size] = d2[k] - - assign(out, output_chunk) - elif out_type is tuple: - for x1, x2 in zip(out, output_chunk): - if _add_into_out: - x1[i : i + chunk_size] += x2 - else: - x1[i : i + chunk_size] = x2 - elif out_type is torch.Tensor: - if _add_into_out: - out[i : i + chunk_size] += output_chunk - else: - out[i : i + chunk_size] = output_chunk - else: - raise ValueError("Not supported") - - i += chunk_size - - reshape = lambda t: t.view(orig_batch_dims + t.shape[1:]) - out = tensor_tree_map(reshape, out) - - return out diff --git a/fastplms/dplm/README.md b/fastplms/dplm/README.md deleted file mode 100644 index 4f62607..0000000 --- a/fastplms/dplm/README.md +++ /dev/null @@ -1,124 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements can be found [here](https://github.com/Synthyra/FastPLMs.git). - -# DPLM -Synthyra DPLM checkpoints are HuggingFace AutoModel compatible and include FastPLMs embedding helpers. - -## Supported models -```python -model_dict = { - "Synthyra/DPLM-150M": "airkingbd/dplm_150m", - "Synthyra/DPLM-650M": "airkingbd/dplm_650m", - "Synthyra/DPLM-3B": "airkingbd/dplm_3b", -} -``` - -## Use with transformers -```python -import torch -from transformers import AutoModel, AutoModelForMaskedLM - -model_path = "Synthyra/DPLM-150M" -model = AutoModel.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval() -tokenizer = model.tokenizer - -batch = tokenizer(["MPRTEIN", "MSEQWENCE"], padding=True, return_tensors="pt") -with torch.no_grad(): - hidden = model(**batch).last_hidden_state - -mlm = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval() -with torch.no_grad(): - logits = mlm(**batch).logits -``` - -## Experimental test-time training - -TTT is disabled by default. Normal DPLM inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the PLM backbone with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. - -```python -metrics = mlm.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -mlm.ttt_reset() -print(metrics["losses"]) -``` - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask — faster on variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30–120 s). Best combined with `torch.compile`. | -| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. | - -Set via config before loading, or change on the model after loading (DPLM propagates the change to all attention layers immediately): - -```python -from transformers import AutoConfig, AutoModel - -# Option 1: set before loading -config = AutoConfig.from_pretrained("Synthyra/DPLM-150M", trust_remote_code=True) -config.attn_backend = "flex" -model = AutoModel.from_pretrained("Synthyra/DPLM-150M", config=config, trust_remote_code=True) - -# Option 2: set after loading -model = AutoModel.from_pretrained("Synthyra/DPLM-150M", trust_remote_code=True) -model.attn_backend = "flex" # propagates to all attention layers in-place -``` - -## Embed datasets -All DPLM models inherit `EmbeddingMixin`, so you can call `model.embed_dataset(...)` directly. - -## Citations - -```bibtex -@article{wang2024dplm, - title={Diffusion Language Models Are Versatile Protein Learners}, - author={Wang, Xinyou and Ye, Zaixiang and Huang, Fei and Cao, Dongyan and Liang, Shujian and Huang, Liang}, - journal={Proceedings of the 41st International Conference on Machine Learning}, - year={2024} -} -``` - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/dplm/get_weights.py b/fastplms/dplm/get_weights.py deleted file mode 100644 index 2be5720..0000000 --- a/fastplms/dplm/get_weights.py +++ /dev/null @@ -1,139 +0,0 @@ -import copy -import os -import torch -from typing import List, Optional, Tuple - -from huggingface_hub import HfApi, login -from transformers import AutoModelForMaskedLM, EsmTokenizer - -from fastplms.dplm.modeling_dplm import DPLMConfig as FastDPLMConfig, DPLMForMaskedLM -from fastplms.weight_parity_utils import assert_model_parameters_fp32 - - -MODEL_DICT = { - "Synthyra/DPLM-150M": "airkingbd/dplm_150m", - "Synthyra/DPLM-650M": "airkingbd/dplm_650m", - "Synthyra/DPLM-3B": "airkingbd/dplm_3b", -} -SHARDED_REPO_IDS = {"Synthyra/DPLM-3B"} -SHARD_SIZE = "5GB" - - -def _delete_legacy_unsharded_weights_if_present(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - if "model.safetensors" in repo_files: - print(f"Deleting legacy unified model.safetensors from {repo_id}") - api.delete_file( - path_in_repo="model.safetensors", - repo_id=repo_id, - repo_type="model", - ) - - -def _assert_repo_has_sharded_weights(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - has_index_file = "model.safetensors.index.json" in repo_files - has_shard_file = any( - repo_file.startswith("model-") and repo_file.endswith(".safetensors") - for repo_file in repo_files - ) - assert has_index_file, f"{repo_id} is missing model.safetensors.index.json." - assert has_shard_file, f"{repo_id} has no model shard files." - assert "model.safetensors" not in repo_files, f"{repo_id} still has unified model.safetensors." - - -def _push_model_with_expected_format(model: DPLMForMaskedLM, api: HfApi, repo_id: str) -> None: - if repo_id in SHARDED_REPO_IDS: - print(f"Pushing sharded weights for {repo_id} with max_shard_size={SHARD_SIZE}") - model.push_to_hub(repo_id, max_shard_size=SHARD_SIZE) - _delete_legacy_unsharded_weights_if_present(api, repo_id) - _assert_repo_has_sharded_weights(api, repo_id) - return - model.push_to_hub(repo_id) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return list(MODEL_DICT.items()) - - selected_items: List[Tuple[str, str]] = [] - for repo_id in repo_ids: - assert repo_id in MODEL_DICT, ( - f"Unknown repo_id {repo_id}. " - f"Valid options: {sorted(MODEL_DICT.keys())}" - ) - selected_items.append((repo_id, MODEL_DICT[repo_id])) - return selected_items - - -if __name__ == "__main__": - # py -m fastplms.dplm.get_weights - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - api = HfApi() - - if args.hf_token is not None: - assert len(args.hf_token) > 0, "--hf_token cannot be empty." - login(token=args.hf_token) - - for repo_id, source_repo in _resolve_repo_items(args.repo_ids): - config = FastDPLMConfig.from_pretrained(source_repo) - config.auto_map = { - "AutoConfig": "modeling_dplm.DPLMConfig", - "AutoModel": "modeling_dplm.DPLMModel", - "AutoModelForMaskedLM": "modeling_dplm.DPLMForMaskedLM", - "AutoModelForSequenceClassification": "modeling_dplm.DPLMForSequenceClassification", - "AutoModelForTokenClassification": "modeling_dplm.DPLMForTokenClassification", - } - config.tie_word_embeddings = False - if args.skip_weights: - tokenizer = EsmTokenizer.from_pretrained(source_repo) - config.push_to_hub(repo_id) - tokenizer.push_to_hub(repo_id) - print(f"[skip-weights] uploaded config+tokenizer for {repo_id}") - continue - model = DPLMForMaskedLM.from_pretrained(source_repo, config=config).eval().cpu().to(torch.float32) - model.tokenizer = EsmTokenizer.from_pretrained(source_repo) - - # Break any potential embedding/LM-head parameter aliasing before export. - model.lm_head.dense.weight = copy.deepcopy(model.lm_head.dense.weight) - model.lm_head.dense.bias = copy.deepcopy(model.lm_head.dense.bias) - model.lm_head.decoder.weight = copy.deepcopy(model.lm_head.decoder.weight) - model.lm_head.decoder.bias = copy.deepcopy(model.lm_head.decoder.bias) - model.lm_head.layer_norm.weight = copy.deepcopy(model.lm_head.layer_norm.weight) - model.lm_head.layer_norm.bias = copy.deepcopy(model.lm_head.layer_norm.bias) - - assert_model_parameters_fp32( - model=model, - model_name=f"DPLM model ({source_repo})", - ) - - tokenizer = model.tokenizer - tokenizer.push_to_hub(repo_id) - _push_model_with_expected_format(model, api, repo_id) - api.upload_file( - path_or_fileobj=os.path.join(os.path.dirname(os.path.abspath(__file__)), "modeling_dplm.py"), - path_in_repo="modeling_dplm.py", - repo_id=repo_id, - repo_type="model", - ) - downloaded_model = AutoModelForMaskedLM.from_pretrained( - repo_id, - dtype=torch.float32, - device_map="cpu", - force_download=True, - trust_remote_code=True, - ) - assert_model_parameters_fp32( - model=downloaded_model, - model_name=f"downloaded DPLM model ({repo_id})", - ) diff --git a/fastplms/dplm/modeling_dplm.py b/fastplms/dplm/modeling_dplm.py deleted file mode 100644 index a955037..0000000 --- a/fastplms/dplm/modeling_dplm.py +++ /dev/null @@ -1,1053 +0,0 @@ -from __future__ import annotations -# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates -# SPDX-License-Identifier: Apache-2.0 -""" -FastPLMs-compatible DPLM implementation. -""" - -import torch -import torch.nn as nn -from torch.nn import functional as F -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union -from einops import rearrange - -from transformers import EsmTokenizer -from transformers.modeling_outputs import ( - BaseModelOutputWithPastAndCrossAttentions, - BaseModelOutputWithPoolingAndCrossAttentions, - ModelOutput, - SequenceClassifierOutput, - TokenClassifierOutput, -) -from transformers.models.esm.configuration_esm import EsmConfig -from transformers.models.esm.modeling_esm import ( - EsmAttention, - EsmClassificationHead, - EsmContactPredictionHead, - EsmEmbeddings, - EsmEncoder, - EsmIntermediate, - EsmLayer, - EsmLMHead, - EsmOutput, - EsmPooler, - EsmPreTrainedModel, - EsmSelfAttention, - EsmSelfOutput, -) - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - kernels_flash_attention_func, - index_first_axis, index_put_first_axis, pad_input, _unpad_input, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -@dataclass -class DPLMMaskedLMOutput(ModelOutput): - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class DPLMEncoderOutput(ModelOutput): - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -class DPLMConfig(EsmConfig): - model_type = "dplm" - - def __init__( - self, - attn_backend: str = "sdpa", - **kwargs, - ): - super().__init__(**kwargs) - self.attn_backend = attn_backend - self.tie_word_embeddings = False - - -class DPLMPreTrainedModel(EsmPreTrainedModel): - config_class = DPLMConfig - base_model_prefix = "dplm" - supports_gradient_checkpointing = True - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - all_tied_weights_keys = {} - - @classmethod - def is_remote_code(cls) -> bool: - # Prevent post-load reinitialization of tensors already loaded from checkpoints. - return True - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - for module in self.modules(): - if isinstance(module, ModifiedEsmEncoder): - module.attention_backend = resolved - elif isinstance(module, ModifiedEsmSelfAttention): - module.attn_backend = resolved - - -class ModifiedEsmSelfAttention(EsmSelfAttention): - def __init__(self, config, position_embedding_type=None): - super().__init__(config, position_embedding_type) - self.config = config - self.scale = self.attention_head_size**-0.5 - self.attn_backend = resolve_attention_backend(config.attn_backend) - - def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: - new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) - x = x.view(new_x_shape) - return x.permute(0, 2, 1, 3) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[object] = None, - head_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - output_attentions: Optional[bool] = False, - output_s_max: Optional[bool] = False, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - if past_key_values is not None: - past_key_value = past_key_values - - mixed_query_layer = self.query(hidden_states) - is_cross_attention = encoder_hidden_states is not None - - if is_cross_attention and past_key_value is not None: - key_layer = past_key_value[0] - value_layer = past_key_value[1] - cross_attn_mask = encoder_attention_mask - elif is_cross_attention: - key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) - value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) - cross_attn_mask = encoder_attention_mask - elif past_key_value is not None: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - key_layer = torch.cat([past_key_value[0], key_layer], dim=2) - value_layer = torch.cat([past_key_value[1], value_layer], dim=2) - cross_attn_mask = None - else: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - cross_attn_mask = None - - query_layer = self.transpose_for_scores(mixed_query_layer) * self.scale - - if self.position_embedding_type == "rotary": - query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) - - if self.position_embedding_type in ["relative_key", "relative_key_query"]: - raise NotImplementedError - - query_layer = query_layer.contiguous() - key_layer = key_layer.contiguous() - value_layer = value_layer.contiguous() - - if is_cross_attention: - if output_attentions: - attn_output, attn_weights, s_max = self._manual_attn( - query_layer, key_layer, value_layer, cross_attn_mask, output_s_max, - ) - else: - attn_output, attn_weights = self._sdpa_attn( - query_layer, key_layer, value_layer, cross_attn_mask, - ) - s_max = self._compute_s_max(query_layer, key_layer) if output_s_max else None - else: - attn_output, attn_weights, s_max = self._attn( - query_layer, key_layer, value_layer, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - if head_mask is not None and torch.is_tensor(head_mask): - batch_size, seq_len, _ = attn_output.shape - attn_output = attn_output.view(batch_size, seq_len, self.num_attention_heads, self.attention_head_size) - attn_output = attn_output.permute(0, 2, 1, 3) * head_mask - attn_output = rearrange(attn_output, "b h s d -> b s (h d)") - - return attn_output, attn_weights, s_max - - def _attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - if output_attentions: - return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d, output_s_max) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - attn_output, attn_weights = self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) - elif self.attn_backend == AttentionBackend.FLEX: - attn_output, attn_weights = self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output, attn_weights = self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - @torch.no_grad() - def _compute_s_max(self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor) -> List[torch.Tensor]: - q_norm = torch.linalg.vector_norm(query_BHLD, dim=-1) - k_norm = torch.linalg.vector_norm(key_BHLD, dim=-1) - s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values - return [s_max_bound[h] for h in range(self.num_attention_heads)] - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[torch.Tensor]]]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - def _kernels_flash_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - query_BLHD = query_BHLD.transpose(1, 2).contiguous() - key_BLHD = key_BHLD.transpose(1, 2).contiguous() - value_BLHD = value_BHLD.transpose(1, 2).contiguous() - # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward(). - # Pass softmax_scale=1.0 to prevent double-scaling by the kernel. - attn_output = kernels_flash_attention_func( - query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, - attention_mask_2d=attention_mask_2d, causal=False, - softmax_scale=1.0, - ) - return rearrange(attn_output, "b s h d -> b s (h d)"), None - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask] = None, - ) -> Tuple[torch.Tensor, None]: - assert flex_attention is not None, "Flex attention is not available in this environment." - fn = _get_flex_attention_fn() - context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, - attn_mask=attention_mask_4d, - scale=1.0, - ) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - -class ModifiedEsmAttention(EsmAttention): - def __init__(self, config): - nn.Module.__init__(self) - self.self = ModifiedEsmSelfAttention(config) - self.output = EsmSelfOutput(config) - self.pruned_heads = set() - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[object] = None, - head_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - hidden_states_ln = self.LayerNorm(hidden_states) - attn_output, attn_weights, s_max = self.self( - hidden_states_ln, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_value=past_key_value, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - attention_output = self.output(attn_output, hidden_states) - return attention_output, attn_weights, s_max - - -class ModifiedEsmLayer(EsmLayer): - def __init__(self, config): - nn.Module.__init__(self) - self.chunk_size_feed_forward = config.chunk_size_feed_forward - self.seq_len_dim = 1 - self.attention = ModifiedEsmAttention(config) - self.is_decoder = config.is_decoder - self.add_cross_attention = config.add_cross_attention - if self.add_cross_attention: - if self.is_decoder is False: - raise RuntimeError(f"{self} should be used as a decoder model if cross attention is added") - self.crossattention = ModifiedEsmAttention(config) - self.intermediate = EsmIntermediate(config) - self.output = EsmOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[object] = None, - head_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - attention_output, attn_weights, s_max = self.attention( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - head_mask=head_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - past_key_value=past_key_value[:2] if past_key_value is not None else None, - ) - - if self.is_decoder and encoder_hidden_states is not None: - if self.add_cross_attention is False: - raise AttributeError( - f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention " - "layers by setting `config.add_cross_attention=True`" - ) - cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None - cross_attention_output, _, _ = self.crossattention( - attention_output, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_value=cross_attn_past_key_value, - output_attentions=output_attentions, - output_s_max=False, - ) - attention_output = cross_attention_output - - layer_output = self.feed_forward_chunk(attention_output) - return layer_output, attn_weights, s_max - - -class ModifiedEsmEncoder(EsmEncoder): - def __init__(self, config): - nn.Module.__init__(self) - self.config = config - self.attention_backend = resolve_attention_backend(config.attn_backend) - self.layer = nn.ModuleList([ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers)]) - self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[Tuple[Tuple[torch.FloatTensor]]]] = None, - use_cache: Optional[bool] = None, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - ) -> DPLMEncoderOutput: - all_hidden_states = () if output_hidden_states else None - all_self_attentions = () if output_attentions else None - full_s_max = () if output_s_max else None - - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=hidden_states.shape[0], - seq_len=hidden_states.shape[1], - device=hidden_states.device, - attention_mask=attention_mask, - ) - - for i, layer_module in enumerate(self.layer): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_head_mask = head_mask[i] if head_mask is not None else None - past_key_value = past_key_values[i] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask_2d, - attention_mask_4d, - flex_block_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - output_s_max, - ) - else: - hidden_states, attn_weights, s_max = layer_module( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - head_mask=layer_head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_value=past_key_value, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - if all_self_attentions is not None: - all_self_attentions = all_self_attentions + (attn_weights,) - if full_s_max is not None: - full_s_max = full_s_max + (s_max,) - - if self.emb_layer_norm_after: - hidden_states = self.emb_layer_norm_after(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - return DPLMEncoderOutput( - last_hidden_state=hidden_states, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - s_max=full_s_max, - ) - - -class FAST_DPLM_ENCODER(DPLMPreTrainedModel, EmbeddingMixin): - """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder, - contact_head) so that the weight keys are prefixed with 'esm.' in the outer DPLMModel, - matching pretrained DPLM checkpoints.""" - - def __init__(self, config, **kwargs): - DPLMPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.embeddings = EsmEmbeddings(config) - self.encoder = ModifiedEsmEncoder(config) - self.contact_head = EsmContactPredictionHead( - in_features=config.num_hidden_layers * config.num_attention_heads, - bias=True, - ) - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - if attention_mask is None: - attention_mask = input_ids.ne(self.config.pad_token_id) - embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - encoder_outputs = self.encoder( - embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - attns = self(input_ids, attention_mask=attention_mask, output_attentions=True).attentions - attns = torch.stack(attns, dim=1) - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) - return self.contact_head(input_ids, attns) - - def _convert_head_mask_to_5d(self, head_mask: torch.Tensor, num_hidden_layers: int) -> torch.Tensor: - if head_mask.dim() == 1: - head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) - head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) - elif head_mask.dim() == 2: - head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) - assert head_mask.dim() == 5, f"head_mask.dim != 5, got {head_mask.dim()}" - head_mask = head_mask.to(dtype=self.dtype) - return head_mask - - def get_head_mask( - self, - head_mask: Optional[torch.Tensor], - num_hidden_layers: int, - is_attention_chunked: bool = False, - ) -> Union[torch.Tensor, List[None]]: - if head_mask is None: - return [None] * num_hidden_layers - head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) - if is_attention_chunked: - head_mask = head_mask.unsqueeze(-1) - return head_mask - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.Tensor], DPLMEncoderOutput]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if self.config.is_decoder: - use_cache = use_cache if use_cache is not None else self.config.use_cache - else: - use_cache = False - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - if input_ids is not None: - input_shape = input_ids.size() - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - batch_size, seq_length = input_shape - device = input_ids.device if input_ids is not None else inputs_embeds.device - - if attention_mask is None: - attention_mask_2d = torch.ones((batch_size, seq_length), device=device).bool() - elif attention_mask.dim() == 2: - attention_mask_2d = attention_mask.bool() - elif attention_mask.dim() == 4: - assert input_ids is not None, "4D attention_mask requires input_ids to infer token-level mask." - attention_mask_2d = input_ids.ne(self.config.pad_token_id) - else: - raise ValueError(f"Unsupported attention_mask shape: {attention_mask.shape}") - - encoder_extended_attention_mask = encoder_attention_mask - if self.config.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - - head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - - embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask_2d, - inputs_embeds=inputs_embeds, - ) - encoder_outputs = self.encoder( - embedding_output, - attention_mask=attention_mask_2d, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - sequence_output = encoder_outputs.last_hidden_state - - if return_dict is False: - return (sequence_output,) + encoder_outputs[1:] - - return DPLMEncoderOutput( - last_hidden_state=sequence_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - s_max=encoder_outputs.s_max, - ) - - -class DPLMModel(DPLMPreTrainedModel, EmbeddingMixin): - config_class = DPLMConfig - - def __init__(self, config, add_pooling_layer=True): - DPLMPreTrainedModel.__init__(self, config) - self.config = config - self.esm = FAST_DPLM_ENCODER(config) - self.pooler = EsmPooler(config) if add_pooling_layer else None - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.esm.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.esm.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.Tensor], DPLMEncoderOutput]: - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - return_dict=return_dict, - ) - sequence_output = outputs[0] - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if return_dict is False: - return (sequence_output, pooled_output) + outputs[1:] - - return DPLMEncoderOutput( - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLMForMaskedLM(FastPLMTestTimeTrainingMixin, DPLMPreTrainedModel, EmbeddingMixin): - config_class = DPLMConfig - - def __init__(self, config, dropout: float = 0.1): - config.hidden_dropout_prob = dropout - DPLMPreTrainedModel.__init__(self, config) - self.esm = FAST_DPLM_ENCODER(config) - self.lm_head = EsmLMHead(config) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - - self.tokenizer = self.__class__.tokenizer - if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0: - try: - self.tokenizer = EsmTokenizer.from_pretrained(config._name_or_path) - except Exception: - self.tokenizer = self.__class__.tokenizer - - self.mask_id = self.tokenizer.mask_token_id - self.pad_id = self.tokenizer.pad_token_id - self.bos_id = self.tokenizer.cls_token_id - self.eos_id = self.tokenizer.eos_token_id - self.x_id = self.tokenizer.convert_tokens_to_ids("X") - self.contact_head = None - self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"}) - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def get_output_embeddings(self): - return self.lm_head.decoder - - def set_output_embeddings(self, new_embeddings): - self.lm_head.decoder = new_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.esm] - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - decoder_input_ids: Optional[torch.Tensor] = None, - decoder_attention_mask: Optional[torch.Tensor] = None, - decoder_inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - ) -> Union[Tuple[torch.Tensor], DPLMMaskedLMOutput]: - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if attention_mask is None and input_ids is not None: - attention_mask = input_ids.ne(self.pad_id) - - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - return_dict=True, - ) - sequence_output = outputs.last_hidden_state - logits = self.lm_head(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) - - if return_dict is False: - output = (logits, sequence_output, outputs.hidden_states, outputs.attentions) - if loss is not None: - return (loss,) + output - return output - - return DPLMMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLMForSequenceClassification(DPLMPreTrainedModel, EmbeddingMixin): - config_class = DPLMConfig - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def __init__(self, config): - DPLMPreTrainedModel.__init__(self, config) - self.num_labels = config.num_labels - self.esm = FAST_DPLM_ENCODER(config) - self.classifier = EsmClassificationHead(config) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - self.post_init() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> Union[Tuple[torch.Tensor], DPLMMaskedLMOutput]: - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - return_dict=True, - ) - sequence_output = outputs.last_hidden_state - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - if self.num_labels == 1: - loss = self.mse(logits.squeeze(), labels.squeeze()) - else: - loss = self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return DPLMMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLMForTokenClassification(DPLMPreTrainedModel, EmbeddingMixin): - config_class = DPLMConfig - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def __init__(self, config): - DPLMPreTrainedModel.__init__(self, config) - self.num_labels = config.num_labels - self.esm = FAST_DPLM_ENCODER(config) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - self.classifier = nn.Linear(config.hidden_size, config.num_labels) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> Union[Tuple[torch.Tensor], DPLMMaskedLMOutput]: - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - return_dict=True, - ) - sequence_output = self.dropout(outputs.last_hidden_state) - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return DPLMMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -if __name__ == "__main__": - import random - - import torch - - from torch import Tensor - from transformers import EsmTokenizer - - def print_tensor_shapes(prefix: str, obj): - if isinstance(obj, Tensor): - print(f"{prefix}{obj.shape}") - elif isinstance(obj, dict): - for name, value in obj.items(): - print_tensor_shapes(f"{prefix}{name}.", value) - elif isinstance(obj, list): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif isinstance(obj, tuple): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif hasattr(obj, "__dict__"): - for name, value in vars(obj).items(): - if name.startswith("_"): - continue - print_tensor_shapes(f"{prefix}{name}.", value) - else: - print(f"{prefix}{type(obj)}") - - random.seed(0) - torch.manual_seed(0) - - num_attention_heads = random.choice([2, 4]) - config = DPLMConfig( - hidden_size=16 * num_attention_heads, - num_attention_heads=num_attention_heads, - num_hidden_layers=random.choice([1, 2]), - attention_probs_dropout_prob=0.0, - hidden_dropout_prob=0.0, - attn_backend="sdpa", - ) - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - batch = tokenizer(["ACDEFG", "MKTW"], return_tensors="pt", padding="longest") - batch["labels"] = batch["input_ids"].clone() - model = DPLMForMaskedLM(config=config).eval() - - with torch.no_grad(): - output = model(**batch, return_dict=True) - - print("Batch shape:") - print_tensor_shapes("", batch) - print("Output shape:") - print_tensor_shapes("", output) diff --git a/fastplms/dplm2/README.md b/fastplms/dplm2/README.md deleted file mode 100644 index 53c8648..0000000 --- a/fastplms/dplm2/README.md +++ /dev/null @@ -1,127 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements can be found [here](https://github.com/Synthyra/FastPLMs.git). - -# DPLM2 -Synthyra DPLM2 checkpoints are HuggingFace AutoModel compatible and include FastPLMs embedding helpers. - -## Supported models -```python -model_dict = { - "Synthyra/DPLM2-150M": "airkingbd/dplm2_150m", - "Synthyra/DPLM2-650M": "airkingbd/dplm2_650m", - "Synthyra/DPLM2-3B": "airkingbd/dplm2_3b", -} -``` - -## Use with transformers -```python -import torch -from transformers import AutoModel, AutoModelForMaskedLM - -model_path = "Synthyra/DPLM2-150M" -model = AutoModel.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval() -tokenizer = model.tokenizer - -batch = tokenizer(["MPRTEIN", "MSEQWENCE"], padding=True, return_tensors="pt") -with torch.no_grad(): - hidden = model(**batch).last_hidden_state - -mlm = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval() -with torch.no_grad(): - logits = mlm(**batch).logits -``` - -## Experimental test-time training - -TTT is disabled by default. Normal DPLM2 inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the PLM backbone with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. - -```python -metrics = mlm.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -mlm.ttt_reset() -print(metrics["losses"]) -``` - -## DPLM2 modality types -DPLM2 infers `type_ids` automatically from `input_ids` and `attention_mask` when they are not provided. - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask — faster on variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30–120 s). Best combined with `torch.compile`. | -| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. | - -Set via config before loading, or change on the model after loading (DPLM2 propagates the change to all attention layers immediately): - -```python -from transformers import AutoConfig, AutoModel - -# Option 1: set before loading -config = AutoConfig.from_pretrained("Synthyra/DPLM2-150M", trust_remote_code=True) -config.attn_backend = "flex" -model = AutoModel.from_pretrained("Synthyra/DPLM2-150M", config=config, trust_remote_code=True) - -# Option 2: set after loading -model = AutoModel.from_pretrained("Synthyra/DPLM2-150M", trust_remote_code=True) -model.attn_backend = "flex" # propagates to all attention layers in-place -``` - -## Embed datasets -All DPLM2 models inherit `EmbeddingMixin`, so you can call `model.embed_dataset(...)` directly. - -## Citations - -```bibtex -@article{wang2024dplm2, - title={DPLM-2: A Multimodal Diffusion Protein Language Model}, - author={Wang, Xinyou and Ye, Zaixiang and Huang, Fei and Cao, Dongyan and Liang, Shujian and Huang, Liang}, - journal={arXiv preprint arXiv:2410.13782}, - year={2024} -} -``` - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/dplm2/get_weights.py b/fastplms/dplm2/get_weights.py deleted file mode 100644 index 7c94d02..0000000 --- a/fastplms/dplm2/get_weights.py +++ /dev/null @@ -1,139 +0,0 @@ -import copy -import os -import torch -from typing import List, Optional, Tuple - -from huggingface_hub import HfApi, login -from transformers import AutoModelForMaskedLM, EsmTokenizer - -from fastplms.dplm2.modeling_dplm2 import DPLM2Config as FastDPLM2Config, DPLM2ForMaskedLM -from fastplms.weight_parity_utils import assert_model_parameters_fp32 - - -MODEL_DICT = { - "Synthyra/DPLM2-150M": "airkingbd/dplm2_150m", - "Synthyra/DPLM2-650M": "airkingbd/dplm2_650m", - "Synthyra/DPLM2-3B": "airkingbd/dplm2_3b", -} -SHARDED_REPO_IDS = {"Synthyra/DPLM2-3B"} -SHARD_SIZE = "5GB" - - -def _delete_legacy_unsharded_weights_if_present(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - if "model.safetensors" in repo_files: - print(f"Deleting legacy unified model.safetensors from {repo_id}") - api.delete_file( - path_in_repo="model.safetensors", - repo_id=repo_id, - repo_type="model", - ) - - -def _assert_repo_has_sharded_weights(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - has_index_file = "model.safetensors.index.json" in repo_files - has_shard_file = any( - repo_file.startswith("model-") and repo_file.endswith(".safetensors") - for repo_file in repo_files - ) - assert has_index_file, f"{repo_id} is missing model.safetensors.index.json." - assert has_shard_file, f"{repo_id} has no model shard files." - assert "model.safetensors" not in repo_files, f"{repo_id} still has unified model.safetensors." - - -def _push_model_with_expected_format(model: DPLM2ForMaskedLM, api: HfApi, repo_id: str) -> None: - if repo_id in SHARDED_REPO_IDS: - print(f"Pushing sharded weights for {repo_id} with max_shard_size={SHARD_SIZE}") - model.push_to_hub(repo_id, max_shard_size=SHARD_SIZE) - _delete_legacy_unsharded_weights_if_present(api, repo_id) - _assert_repo_has_sharded_weights(api, repo_id) - return - model.push_to_hub(repo_id) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return list(MODEL_DICT.items()) - - selected_items: List[Tuple[str, str]] = [] - for repo_id in repo_ids: - assert repo_id in MODEL_DICT, ( - f"Unknown repo_id {repo_id}. " - f"Valid options: {sorted(MODEL_DICT.keys())}" - ) - selected_items.append((repo_id, MODEL_DICT[repo_id])) - return selected_items - - -if __name__ == "__main__": - # py -m fastplms.dplm2.get_weights - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - api = HfApi() - - if args.hf_token is not None: - assert len(args.hf_token) > 0, "--hf_token cannot be empty." - login(token=args.hf_token) - - for repo_id, source_repo in _resolve_repo_items(args.repo_ids): - config = FastDPLM2Config.from_pretrained(source_repo) - config.auto_map = { - "AutoConfig": "modeling_dplm2.DPLM2Config", - "AutoModel": "modeling_dplm2.DPLM2Model", - "AutoModelForMaskedLM": "modeling_dplm2.DPLM2ForMaskedLM", - "AutoModelForSequenceClassification": "modeling_dplm2.DPLM2ForSequenceClassification", - "AutoModelForTokenClassification": "modeling_dplm2.DPLM2ForTokenClassification", - } - config.tie_word_embeddings = False - if args.skip_weights: - tokenizer = EsmTokenizer.from_pretrained(source_repo) - config.push_to_hub(repo_id) - tokenizer.push_to_hub(repo_id) - print(f"[skip-weights] uploaded config+tokenizer for {repo_id}") - continue - model = DPLM2ForMaskedLM.from_pretrained(source_repo, config=config).eval().cpu().to(torch.float32) - model.tokenizer = EsmTokenizer.from_pretrained(source_repo) - - # Break any potential embedding/LM-head parameter aliasing before export. - model.lm_head.dense.weight = copy.deepcopy(model.lm_head.dense.weight) - model.lm_head.dense.bias = copy.deepcopy(model.lm_head.dense.bias) - model.lm_head.decoder.weight = copy.deepcopy(model.lm_head.decoder.weight) - model.lm_head.decoder.bias = copy.deepcopy(model.lm_head.decoder.bias) - model.lm_head.layer_norm.weight = copy.deepcopy(model.lm_head.layer_norm.weight) - model.lm_head.layer_norm.bias = copy.deepcopy(model.lm_head.layer_norm.bias) - - assert_model_parameters_fp32( - model=model, - model_name=f"DPLM2 model ({source_repo})", - ) - - tokenizer = model.tokenizer - tokenizer.push_to_hub(repo_id) - _push_model_with_expected_format(model, api, repo_id) - api.upload_file( - path_or_fileobj=os.path.join(os.path.dirname(os.path.abspath(__file__)), "modeling_dplm2.py"), - path_in_repo="modeling_dplm2.py", - repo_id=repo_id, - repo_type="model", - ) - downloaded_model = AutoModelForMaskedLM.from_pretrained( - repo_id, - dtype=torch.float32, - device_map="cpu", - force_download=True, - trust_remote_code=True, - ) - assert_model_parameters_fp32( - model=downloaded_model, - model_name=f"downloaded DPLM2 model ({repo_id})", - ) diff --git a/fastplms/dplm2/modeling_dplm2.py b/fastplms/dplm2/modeling_dplm2.py deleted file mode 100644 index 608b78e..0000000 --- a/fastplms/dplm2/modeling_dplm2.py +++ /dev/null @@ -1,1022 +0,0 @@ -from __future__ import annotations -""" -FastPLMs-compatible DPLM2 implementation. -""" - -import torch -import torch.nn as nn -from torch.nn import functional as F -from dataclasses import dataclass -from einops import rearrange -from enum import Enum -from typing import List, Optional, Tuple, Union - -from transformers import EsmTokenizer -from transformers.modeling_outputs import ( - BaseModelOutputWithPastAndCrossAttentions, - BaseModelOutputWithPoolingAndCrossAttentions, - ModelOutput, - SequenceClassifierOutput, - TokenClassifierOutput, -) -from transformers.models.esm.configuration_esm import EsmConfig -from transformers.models.esm.modeling_esm import ( - EsmAttention, - EsmClassificationHead, - EsmEmbeddings, - EsmEncoder, - EsmIntermediate, - EsmLayer, - EsmLMHead, - EsmOutput, - EsmPooler, - EsmPreTrainedModel, - EsmSelfAttention, - EsmSelfOutput, - RotaryEmbedding, - apply_rotary_pos_emb, -) - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - kernels_flash_attention_func, - index_first_axis, index_put_first_axis, pad_input, _unpad_input, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -def _infer_modality_type(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - input_mask = attention_mask.bool() - modality_type = ((input_ids < 33) & input_mask).int() - modality_type[~input_mask] = 2 - return modality_type - - -def _normalize_dplm2_input_ids(input_ids: torch.Tensor, vocab_size: int) -> torch.Tensor: - if input_ids.numel() == 0: - return input_ids - - normalized_input_ids = input_ids.clone() - generic_to_aa_special_ids = { - vocab_size: 2, - vocab_size + 1: 3, - vocab_size + 2: 0, - vocab_size + 3: 32, - } - for generic_id, aa_id in generic_to_aa_special_ids.items(): - normalized_input_ids[input_ids == generic_id] = aa_id - - valid_token_mask = normalized_input_ids.ge(0) - if valid_token_mask.any(): - max_token_id = int(normalized_input_ids[valid_token_mask].max().item()) - assert max_token_id < vocab_size, ( - f"Found token id {max_token_id} outside the DPLM2 embedding table (vocab_size={vocab_size}). " - "Tokenizer special tokens must be normalized before embedding." - ) - return normalized_input_ids - - -def _has_packed_multimodal_layout( - type_ids: Optional[torch.Tensor], - aa_type: int, - struct_type: int, - pad_type: int, -) -> bool: - if type_ids is None: - return False - assert type_ids.ndim == 2, f"Expected type_ids to have shape (batch, seq_len), got {tuple(type_ids.shape)}" - seq_len = type_ids.shape[-1] - if seq_len % 2 != 0: - return False - - half_len = seq_len // 2 - first_half = type_ids[:, :half_len] - second_half = type_ids[:, half_len:] - - first_half_valid = ((first_half == aa_type) | (first_half == pad_type)).all(dim=-1) - second_half_valid = ((second_half == struct_type) | (second_half == pad_type)).all(dim=-1) - aa_count = (first_half == aa_type).sum(dim=-1) - struct_count = (second_half == struct_type).sum(dim=-1) - packed_rows = first_half_valid & second_half_valid & aa_count.gt(0) & aa_count.eq(struct_count) - return bool(packed_rows.all()) - - -@dataclass -class DPLM2MaskedLMOutput(ModelOutput): - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class DPLM2EncoderOutput(ModelOutput): - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -class DPLM2Config(EsmConfig): - model_type = "dplm2" - - def __init__( - self, - attn_backend: str = "sdpa", - aa_type: int = 1, - struct_type: int = 0, - pad_type: int = 2, - **kwargs, - ): - super().__init__(**kwargs) - self.attn_backend = attn_backend - self.aa_type = aa_type - self.struct_type = struct_type - self.pad_type = pad_type - self.tie_word_embeddings = False - - -class DPLM2PreTrainedModel(EsmPreTrainedModel): - config_class = DPLM2Config - base_model_prefix = "dplm2" - supports_gradient_checkpointing = True - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - all_tied_weights_keys = {} - - @classmethod - def is_remote_code(cls) -> bool: - # Prevent post-load reinitialization of tensors already loaded from checkpoints. - return True - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - for module in self.modules(): - if isinstance(module, ModifiedEsmEncoder): - module.attention_backend = resolved - elif isinstance(module, ModifiedEsmSelfAttention): - module.attn_backend = resolved - - - -class ModifiedRotaryEmbedding(RotaryEmbedding): - def __init__(self, dim: int, aa_type: int, struct_type: int, pad_type: int): - super().__init__(dim) - self.aa_type = aa_type - self.struct_type = struct_type - self.pad_type = pad_type - - def _has_multimodal_tokens(self, type_ids: Optional[torch.Tensor]) -> bool: - # The split rotary path only works when the sequence tensor is already packed - # as [AA half | structure half]. Plain protein batches can still contain - # high-ID special tokens, so mere modality presence is not enough. - return _has_packed_multimodal_layout( - type_ids=type_ids, - aa_type=self.aa_type, - struct_type=self.struct_type, - pad_type=self.pad_type, - ) - - def _update_cos_sin_tables( - self, - x: torch.Tensor, - type_ids: Optional[torch.Tensor], - seq_dimension: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: - seq_len = x.shape[seq_dimension] - if self._has_multimodal_tokens(type_ids): - seq_len = seq_len // 2 - - cache_is_stale = ( - self._cos_cached is None - or self._sin_cached is None - or seq_len != self._seq_len_cached - or self._cos_cached.device != x.device - or self._cos_cached.dtype != x.dtype - ) - if cache_is_stale: - self._seq_len_cached = seq_len - t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq) - freqs = torch.outer(t, self.inv_freq) - emb = torch.cat((freqs, freqs), dim=-1).to(device=x.device, dtype=x.dtype) - self._cos_cached = emb.cos()[None, None, :, :] - self._sin_cached = emb.sin()[None, None, :, :] - - return self._cos_cached, self._sin_cached - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - type_ids: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - self._cos_cached, self._sin_cached = self._update_cos_sin_tables( - k, - type_ids=type_ids, - seq_dimension=-2, - ) - - if self._has_multimodal_tokens(type_ids): - q_1, q_2 = q.chunk(2, dim=-2) - k_1, k_2 = k.chunk(2, dim=-2) - q_1 = apply_rotary_pos_emb(q_1, self._cos_cached, self._sin_cached) - q_2 = apply_rotary_pos_emb(q_2, self._cos_cached, self._sin_cached) - k_1 = apply_rotary_pos_emb(k_1, self._cos_cached, self._sin_cached) - k_2 = apply_rotary_pos_emb(k_2, self._cos_cached, self._sin_cached) - return torch.cat((q_1, q_2), dim=-2), torch.cat((k_1, k_2), dim=-2) - - return ( - apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), - apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), - ) - - -class ModifiedEsmSelfAttention(EsmSelfAttention): - def __init__(self, config, position_embedding_type=None): - super().__init__(config, position_embedding_type) - self.config = config - self.scale = self.attention_head_size**-0.5 - self.dropout_prob = config.attention_probs_dropout_prob - self.attn_backend = resolve_attention_backend(config.attn_backend) - self.rotary_embeddings = ModifiedRotaryEmbedding( - dim=self.attention_head_size, - aa_type=config.aa_type, - struct_type=config.struct_type, - pad_type=config.pad_type, - ) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - type_ids: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - batch_size, seq_length = hidden_states.shape[:-1] - hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) - query_BHLD = self.query(hidden_states).view(hidden_shape).transpose(1, 2) - key_BHLD = self.key(hidden_states).view(hidden_shape).transpose(1, 2) - value_BHLD = self.value(hidden_states).view(hidden_shape).transpose(1, 2) - - query_BHLD = query_BHLD * self.scale - - if self.position_embedding_type == "rotary": - query_BHLD, key_BHLD = self.rotary_embeddings(query_BHLD, key_BHLD, type_ids) - - attn_output, attn_weights, s_max = self._attn( - query_BHLD, key_BHLD, value_BHLD, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - return attn_output, attn_weights, s_max - - def _attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - if output_attentions: - return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d, output_s_max) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - attn_output, attn_weights = self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) - elif self.attn_backend == AttentionBackend.FLEX: - attn_output, attn_weights = self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output, attn_weights = self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - @torch.no_grad() - def _compute_s_max(self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor) -> List[torch.Tensor]: - q_norm = torch.linalg.vector_norm(query_BHLD, dim=-1) - k_norm = torch.linalg.vector_norm(key_BHLD, dim=-1) - s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values - return [s_max_bound[h] for h in range(self.num_attention_heads)] - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[torch.Tensor]]]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - if self.dropout_prob > 0 and self.training: - attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - def _kernels_flash_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - query_BLHD = query_BHLD.transpose(1, 2).contiguous() - key_BLHD = key_BHLD.transpose(1, 2).contiguous() - value_BLHD = value_BHLD.transpose(1, 2).contiguous() - # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0 - # to prevent the kernel from applying its default 1/sqrt(head_dim). - attn_output = kernels_flash_attention_func( - query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, - attention_mask_2d=attention_mask_2d, causal=False, - softmax_scale=1.0, - ) - return rearrange(attn_output, "b s h d -> b s (h d)"), None - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask] = None, - ) -> Tuple[torch.Tensor, None]: - assert flex_attention is not None, "Flex attention is not available in this environment." - fn = _get_flex_attention_fn() - context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, - attn_mask=attention_mask_4d, - dropout_p=self.dropout_prob if self.training else 0.0, - scale=1.0, - ) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - -class ModifiedEsmAttention(EsmAttention): - def __init__(self, config): - nn.Module.__init__(self) - self.self = ModifiedEsmSelfAttention(config) - self.output = EsmSelfOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - type_ids: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - hidden_states_ln = self.LayerNorm(hidden_states) - attn_output, attn_weights, s_max = self.self( - hidden_states_ln, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - type_ids=type_ids, - ) - attention_output = self.output(attn_output, hidden_states) - return attention_output, attn_weights, s_max - - -class ModifiedEsmLayer(EsmLayer): - def __init__(self, config): - nn.Module.__init__(self) - self.chunk_size_feed_forward = config.chunk_size_feed_forward - self.seq_len_dim = 1 - self.attention = ModifiedEsmAttention(config) - self.intermediate = EsmIntermediate(config) - self.output = EsmOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - type_ids: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - attention_output, attn_weights, s_max = self.attention( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - type_ids=type_ids, - ) - layer_output = self.feed_forward_chunk(attention_output) - return layer_output, attn_weights, s_max - - -class ModifiedEsmEncoder(EsmEncoder): - def __init__(self, config): - nn.Module.__init__(self) - self.config = config - self.attention_backend = resolve_attention_backend(config.attn_backend) - self.layer = nn.ModuleList([ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers)]) - self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: bool = False, - output_attentions: bool = False, - output_s_max: bool = False, - type_ids: Optional[torch.Tensor] = None, - ) -> DPLM2EncoderOutput: - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - full_s_max = () if output_s_max else None - - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=hidden_states.shape[0], - seq_len=hidden_states.shape[1], - device=hidden_states.device, - attention_mask=attention_mask, - ) - - for layer_module in self.layer: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask_2d, - attention_mask_4d, - flex_block_mask, - output_attentions, - output_s_max, - type_ids, - ) - else: - hidden_states, attn_weights, s_max = layer_module( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - type_ids=type_ids, - ) - - if all_attentions is not None: - all_attentions = all_attentions + (attn_weights,) - if full_s_max is not None: - full_s_max = full_s_max + (s_max,) - - if self.emb_layer_norm_after: - hidden_states = self.emb_layer_norm_after(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - return DPLM2EncoderOutput( - last_hidden_state=hidden_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - s_max=full_s_max, - ) - - -class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin): - """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder) - so that the weight keys are prefixed with 'esm.' in the outer DPLM2Model, - matching pretrained DPLM2 checkpoints.""" - - def __init__(self, config, **kwargs): - DPLM2PreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.embeddings = EsmEmbeddings(config) - self.encoder = ModifiedEsmEncoder(config) - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) - if attention_mask is None: - attention_mask = input_ids.ne(self.config.pad_token_id) - type_ids = _infer_modality_type(input_ids, attention_mask) - token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - encoder_outputs = self.encoder( - token_embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - type_ids=type_ids, - ) - return select_hidden_state_embeddings( - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - type_ids: Optional[torch.Tensor] = None, - ) -> DPLM2EncoderOutput: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) - elif inputs_embeds is None: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - token_embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - ) - encoder_outputs = self.encoder( - token_embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - type_ids=type_ids, - ) - - return DPLM2EncoderOutput( - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - s_max=encoder_outputs.s_max, - ) - - -class DPLM2Model(DPLM2PreTrainedModel, EmbeddingMixin): - config_class = DPLM2Config - def __init__(self, config, add_pooling_layer=True): - DPLM2PreTrainedModel.__init__(self, config) - self.config = config - self.esm = FAST_DPLM2_ENCODER(config) - self.pooler = EsmPooler(config) if add_pooling_layer else None - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.esm.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.esm.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - type_ids: Optional[torch.Tensor] = None, - ) -> DPLM2EncoderOutput: - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - type_ids=type_ids, - ) - sequence_output = outputs.last_hidden_state - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - - return DPLM2EncoderOutput( - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLM2ForMaskedLM(FastPLMTestTimeTrainingMixin, DPLM2PreTrainedModel, EmbeddingMixin): - config_class = DPLM2Config - def __init__(self, config, dropout: float = 0.1, vocab_size: Optional[int] = None): - config.hidden_dropout_prob = dropout - config.tie_word_embeddings = False - if vocab_size is not None: - config.vocab_size = vocab_size - DPLM2PreTrainedModel.__init__(self, config) - self.esm = FAST_DPLM2_ENCODER(config) - self.lm_head = EsmLMHead(config) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - self.pad_id = config.pad_token_id - self.tokenizer = self.__class__.tokenizer - if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0: - self.tokenizer = EsmTokenizer.from_pretrained(config._name_or_path) - self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"}) - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def get_output_embeddings(self): - return self.lm_head.decoder - - def set_output_embeddings(self, new_embeddings): - self.lm_head.decoder = new_embeddings - - def _get_modality_type(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) - return _infer_modality_type(input_ids, attention_mask) - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - if attention_mask is None: - attention_mask = input_ids.ne(self.pad_id) - type_ids = self._get_modality_type(input_ids, attention_mask) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - type_ids=type_ids, - output_attentions=False, - output_hidden_states=output_hidden_states, - ) - return select_hidden_state_embeddings( - outputs.last_hidden_state, - outputs.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.esm] - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - type_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.Tensor], DPLM2MaskedLMOutput]: - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if attention_mask is None: - assert input_ids is not None - attention_mask = input_ids.ne(self.pad_id) - - if type_ids is None: - assert input_ids is not None - type_ids = self._get_modality_type(input_ids, attention_mask) - - outputs = self.esm( - input_ids=input_ids, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - type_ids=type_ids, - ) - - sequence_output = outputs.last_hidden_state - logits = self.lm_head(sequence_output) - loss = None - if labels is not None: - labels = _normalize_dplm2_input_ids(labels, self.config.vocab_size) - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) - - if return_dict is False: - output = (logits, sequence_output, outputs.hidden_states, outputs.attentions) - if loss is not None: - return (loss,) + output - return output - - return DPLM2MaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLM2ForSequenceClassification(DPLM2PreTrainedModel, EmbeddingMixin): - config_class = DPLM2Config - - def __init__(self, config): - DPLM2PreTrainedModel.__init__(self, config) - self.num_labels = config.num_labels - self.esm = FAST_DPLM2_ENCODER(config) - self.classifier = EsmClassificationHead(config) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - type_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> DPLM2MaskedLMOutput: - if type_ids is None and input_ids is not None: - if attention_mask is None: - attention_mask = input_ids.ne(self.config.pad_token_id) - input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) - type_ids = _infer_modality_type(input_ids, attention_mask) - - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - type_ids=type_ids, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - sequence_output = outputs.last_hidden_state - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - if self.num_labels == 1: - loss = self.mse(logits.squeeze(), labels.squeeze()) - else: - loss = self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return DPLM2MaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class DPLM2ForTokenClassification(DPLM2PreTrainedModel, EmbeddingMixin): - config_class = DPLM2Config - - def __init__(self, config): - DPLM2PreTrainedModel.__init__(self, config) - self.num_labels = config.num_labels - self.esm = FAST_DPLM2_ENCODER(config) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - self.classifier = nn.Linear(config.hidden_size, config.num_labels) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - - def get_input_embeddings(self) -> nn.Module: - return self.esm.get_input_embeddings() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - type_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> DPLM2MaskedLMOutput: - if type_ids is None and input_ids is not None: - if attention_mask is None: - attention_mask = input_ids.ne(self.config.pad_token_id) - input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) - type_ids = _infer_modality_type(input_ids, attention_mask) - - outputs = self.esm( - input_ids=input_ids, - attention_mask=attention_mask, - type_ids=type_ids, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - sequence_output = self.dropout(outputs.last_hidden_state) - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return DPLM2MaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -if __name__ == "__main__": - import random - - import torch - - from torch import Tensor - from transformers import EsmTokenizer - - def print_tensor_shapes(prefix: str, obj): - if isinstance(obj, Tensor): - print(f"{prefix}{obj.shape}") - elif isinstance(obj, dict): - for name, value in obj.items(): - print_tensor_shapes(f"{prefix}{name}.", value) - elif isinstance(obj, list): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif isinstance(obj, tuple): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif hasattr(obj, "__dict__"): - for name, value in vars(obj).items(): - if name.startswith("_"): - continue - print_tensor_shapes(f"{prefix}{name}.", value) - else: - print(f"{prefix}{type(obj)}") - - random.seed(0) - torch.manual_seed(0) - - num_attention_heads = random.choice([2, 4]) - config = DPLM2Config( - hidden_size=16 * num_attention_heads, - num_attention_heads=num_attention_heads, - num_hidden_layers=random.choice([1, 2]), - attention_probs_dropout_prob=0.0, - hidden_dropout_prob=0.0, - attn_backend="sdpa", - ) - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - batch = tokenizer(["ACDEFGH", "MKTW"], return_tensors="pt", padding="longest") - batch["labels"] = batch["input_ids"].clone() - model = DPLM2ForMaskedLM(config=config).eval() - - with torch.no_grad(): - output = model(**batch, return_dict=True) - - print("Batch shape:") - print_tensor_shapes("", batch) - print("Output shape:") - print_tensor_shapes("", output) diff --git a/fastplms/e1/README.md b/fastplms/e1/README.md deleted file mode 100644 index 0faaff2..0000000 --- a/fastplms/e1/README.md +++ /dev/null @@ -1,255 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# Profluent-E1 -[Synthyra's version of Profluent-E1](https://github.com/Synthyra/Profluent-E1-300M) is a faithful implementation of Profluent's [E1](https://www.profluent.bio/showcase/e1) models ([license](https://github.com/Profluent-AI/E1/tree/main?tab=License-1-ov-file)) that integrates Huggingface AutoModel compatability and nice embedding functionality. - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. The backend is set via `config.attn_backend` before loading. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Uses a block-causal mask that skips padding tokens. Near-exact numerics. First use compiles a Triton kernel (30–120 s). Best combined with `torch.compile`. | -| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. | - -```python -from transformers import AutoConfig, AutoModelForMaskedLM - -config = AutoConfig.from_pretrained("Synthyra/Profluent-E1-150M", trust_remote_code=True) -config.attn_backend = "flex" # or "kernels_flash", "sdpa", "auto" -model = AutoModelForMaskedLM.from_pretrained("Synthyra/Profluent-E1-150M", config=config, trust_remote_code=True) -``` - -`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention. - - -## Use with 🤗 transformers -### Supported models -```python -model_dict = { - # Synthyra/Profluent-E1-150M - 'Profluent-E1-150M': 'Profluent-Bio/E1-150m', - # Synthyra/Profluent-E1-150M - 'Profluent-E1-300M': 'Profluent-Bio/E1-300m', - # Synthyra/Profluent-E1-150M - 'Profluent-E1-600M': 'Profluent-Bio/E1-600m', -} -``` - -```python -import torch -from transformers import AutoModelForMaskedLM - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -model = AutoModelForMaskedLM.from_pretrained('Synthyra/Profluent-E1-150M', trust_remote_code=True, dtype=torch.bfloat16).eval().to(device) - -sequences = ['MPRTEIN', 'MSEQWENCE'] -batch = model.prep_tokens.get_batch_kwargs(sequences, device=device) - -output = model(**batch) # get all hidden states with output_hidden_states=True -print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 34) -print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 768) -print(output.loss) # language modeling loss if you passed labels -#print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple) -#print(outout.attentions) # all attention matrices if you passed output_attentions=True (in tuple) -``` - -Our E1 implementation also supports sequence and token level classification tasks like ESM2. Simply pass the number of labels during initialization. - -```python -from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification - -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/Profluent-E1-150M', num_labels=2, trust_remote_code=True) -logits = model(**batch, labels=labels).logits -print(logits.shape) # (batch_size, num_labels), (2, 2) -``` - -E1 weights were trained in bf16 and are in bf16 by default. You can load them in the precision of your choosing by leveraging the dtype parameter: -```python -import torch -model = AutoModelForMaskedLM.from_pretrained('Synthyra/Profluent-E1-150M', trust_remote_code=True, dtype=torch.float) # fp32 -``` - -## Experimental test-time training - -TTT is disabled by default. Normal E1 inference, MSA-context utilities, -embeddings, and `state_dict()` keys are unchanged unless you explicitly call -`model.ttt(...)`. The current implementation is experimental and trains only -local LoRA adapters with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -## Embed entire datasets with no new code -To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take. - -Example: -```python -embedding_dict = model.embed_dataset( - sequences=[ - 'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences - ], - batch_size=2, # adjust for your GPU memory - max_len=512, # adjust for your needs - full_embeddings=False, # if True, no pooling is performed - embed_dtype=torch.float32, # cast to what dtype you want - pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together - sql=False, # if True, embeddings will be stored in SQLite database - sql_db_path='embeddings.db', - save=True, # if True, embeddings will be saved as a .pth file - save_path='embeddings.pth', -) -# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql -``` - -``` -model.embed_dataset() -Args: - sequences: List of protein sequences - batch_size: Batch size for processing - max_len: Maximum sequence length - full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False) - pooling_type: Type of pooling ('mean' or 'cls') - sql: Whether to store embeddings in SQLite database - will be stored in float32 - sql_db_path: Path to SQLite database - -Returns: - Dictionary mapping sequences to embeddings, or None if sql=True - -Note: - - If sql=True, embeddings can only be stored in float32 - - sql is ideal if you need to stream a very large dataset for training in real-time - - save=True is ideal if you can store the entire embedding dictionary in RAM - - sql will be used if it is True and save is True or False - - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences - - Sequences will be truncated to max_len and sorted by length in descending order for faster processing -``` - -## MSA context, PPLL scoring, and RAG embeddings - -FastPLMs exposes E1 retrieval-augmented MSA context utilities directly on the model object: - -```python -a3m_path = model.search_homologues( - sequence="MALWMRLLPLLALLALWGPDPAAA", - output_dir="msas", - provider="colabfold", -) - -contexts = model.sample_msa_contexts( - a3m_path=a3m_path, - max_context_tokens=[6144, 12288, 24576], - similarity_thresholds=[1.0, 0.95, 0.9, 0.7, 0.5], -) - -scores = model.score_ppll( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - a3m_path=a3m_path, - ensemble=True, -) - -embeddings = model.embed_with_msa( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - a3m_path=a3m_path, - pooling_types=["mean"], -) -``` - -The MSA parsing and context sampling follow Profluent's official E1 `msa_sampling` behavior, including A3M insertion stripping, neighbor reweighting, query-similarity filtering, seeded sampling, and context token budgets. - -`score_ppll()` is intentionally different from Profluent's official `E1Scorer`. The official scorer computes mutant scores against a parent sequence with wildtype or masked marginal log-probability deltas. FastPLMs uses a PPLL-style mean correct-token probability over each scored sequence, then optionally averages over sampled contexts. We prefer this API because it is much cheaper while remaining comparable for our use cases. - -For dataset embeddings with precomputed MSAs: - -```python -embedding_dict = model.embed_dataset_with_msa( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - msa_dir="msas", - batch_size=2, - pooling_types=["mean"], -) -``` - -The standard `embed()` and `embed_dataset()` paths are unchanged. Use `embed_with_msa()` or `embed_dataset_with_msa()` when you want retrieval context included. - -## Fine-tuning with 🤗 peft -```python -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/Profluent-E1-150M', num_labels=2, trust_remote_code=True) -# these modules handle E1 attention layers -target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] - -lora_config = LoraConfig( - r=8, # choose lora parameters to your liking - lora_alpha=16, - lora_dropout=0.01, - bias="none", - target_modules=target_modules, -) - -# Apply LoRA to the model -model = get_peft_model(model, lora_config) - -# Unfreeze the classifier head -for param in model.classifier.parameters(): - param.requires_grad = True -``` - -For a more thourough example of fine-tuning, check out our example script [here](https://github.com/Synthyra/FastPLMs/blob/main/fine_tuning_example.py). - - -### Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{jain2025e1, - title={E1: Retrieval-Augmented Protein Encoder Models}, - author={Jain, Sarthak and Beazer, Joel and Ruffolo, Jeffrey A and Bhatnagar, Aadyot and Madani, Ali}, - journal={bioRxiv}, - DOI={10.1101/2025.11.12.688125}, - year={2025} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/e1/get_weights.py b/fastplms/e1/get_weights.py deleted file mode 100644 index 42cc4c3..0000000 --- a/fastplms/e1/get_weights.py +++ /dev/null @@ -1,41 +0,0 @@ -import argparse -import torch -from huggingface_hub import login - -from fastplms.e1.modeling_e1 import E1ForMaskedLM, E1Config - - -model_dict = { - 'Profluent-E1-150M': 'Profluent-Bio/E1-150m', - 'Profluent-E1-300M': 'Profluent-Bio/E1-300m', - 'Profluent-E1-600M': 'Profluent-Bio/E1-600m', -} - - -if __name__ == "__main__": - # py -m fastplms.e1.get_weights - - parser = argparse.ArgumentParser() - parser.add_argument('--hf_token', type=str, default=None) - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - - if args.hf_token: - login(token=args.hf_token) - - for model_name in model_dict: - repo_id = "Synthyra/" + model_name - config = E1Config.from_pretrained(model_dict[model_name]) - config.auto_map = { - "AutoConfig": "modeling_e1.E1Config", - "AutoModel": "modeling_e1.E1Model", - "AutoModelForMaskedLM": "modeling_e1.E1ForMaskedLM", - "AutoModelForSequenceClassification": "modeling_e1.E1ForSequenceClassification", - "AutoModelForTokenClassification": "modeling_e1.E1ForTokenClassification" - } - if args.skip_weights: - config.push_to_hub(repo_id) - print(f"[skip-weights] uploaded config for {repo_id}") - continue - model = E1ForMaskedLM.from_pretrained(model_dict[model_name], config=config, dtype=torch.float32) - model.push_to_hub(repo_id) \ No newline at end of file diff --git a/fastplms/e1/modeling_e1.py b/fastplms/e1/modeling_e1.py deleted file mode 100644 index 3afdb00..0000000 --- a/fastplms/e1/modeling_e1.py +++ /dev/null @@ -1,3433 +0,0 @@ -from __future__ import annotations - -import hashlib -import itertools -import os -import pickle -import random -import shutil -import subprocess -import tarfile -import tempfile -import time -from collections import defaultdict, namedtuple -from collections.abc import Iterator, Sequence -from dataclasses import dataclass -from enum import Enum -from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.nn.utils.rnn import pad_sequence -from tokenizers import Tokenizer -from tqdm.auto import tqdm -from transformers import PretrainedConfig, PreTrainedModel -from transformers.activations import ACT2FN -from transformers.modeling_outputs import ModelOutput -from transformers.utils import logging - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - index_first_axis, index_put_first_axis, pad_input, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -logger = logging.get_logger(__name__) - -from torch.nn.attention.flex_attention import _create_sparse_block_from_block_mask - -try: - from kernels import get_kernel - layer_norm = get_kernel("kernels-community/triton-layer-norm") -except Exception as e: - logger.warning(f"Failed to load triton layer norm kernel: {e}; Will be using PyTorch RMSNorm instead") - layer_norm = None - - -@torch.compiler.disable -def create_block_causal_mask_optimized(sequence_ids: torch.Tensor) -> BlockMask: - # Assumes sequence_ids is sorted in increasing order for each batch item, except for - # the -1 values, which are used to indicate the padding tokens. - def document_mask(b, h, q_idx, kv_idx): # type: ignore[no-untyped-def] - return ( - (sequence_ids[b, q_idx] >= sequence_ids[b, kv_idx]) - & (sequence_ids[b, q_idx] != -1) - & (sequence_ids[b, kv_idx] != -1) - ) - - batch_size, seqlen = sequence_ids.shape - return create_block_mask(document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device) - - -@torch.compiler.disable -def create_within_seq_block_mask(sequence_ids: torch.Tensor) -> BlockMask: - def document_mask(b, h, q_idx, kv_idx): # type: ignore[no-untyped-def] - return ( - (sequence_ids[b, q_idx] == sequence_ids[b, kv_idx]) - & (sequence_ids[b, q_idx] != -1) - & (sequence_ids[b, kv_idx] != -1) - ) - - batch_size, seqlen = sequence_ids.shape - return create_block_mask(document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device) - - -def build_within_seq_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor: - not_pad = (sequence_ids != -1) - same_seq = sequence_ids.unsqueeze(-1) == sequence_ids.unsqueeze(-2) - valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2) - return (same_seq & valid).unsqueeze(1) - - -def build_block_causal_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor: - not_pad = (sequence_ids != -1) - causal = sequence_ids.unsqueeze(-1) >= sequence_ids.unsqueeze(-2) - valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2) - return (causal & valid).unsqueeze(1) - - -def flex_attention_func( - query_states: torch.Tensor, # (bs, seqlen, nh, hs) - key_states: torch.Tensor, # (bs, seqlen, nkv, hs) - value_states: torch.Tensor, # (bs, seqlen, nkv, hs) - score_mod: Optional[Callable] = None, - block_mask: Optional[BlockMask] = None, -) -> torch.Tensor: - assert flex_attention is not None, "Flex Attention is not available in this environment" - assert score_mod is None, "Score mod is not supported yet" - query_states = query_states.transpose(1, 2).contiguous() # (bs, nh, seqlen, hs) - key_states = key_states.transpose(1, 2).contiguous() # (bs, nkv, seqlen, hs) - value_states = value_states.transpose(1, 2).contiguous() # (bs, nkv, seqlen, hs) - - fn = _get_flex_attention_fn() - outputs = fn( - query_states, - key_states, - value_states, - block_mask=block_mask, - score_mod=score_mod, - enable_gqa=query_states.shape[1] != key_states.shape[1], # if nkv != nh - ) - - outputs = outputs.transpose(1, 2) # (bs, seqlen, nh, hs) - return outputs - - -def kernels_flash_attention_func( - query_states: torch.Tensor, # (bs, seqlen, nh, hs) - key_states: torch.Tensor, # (bs, seqlen, nkv, hs) - value_states: torch.Tensor, # (bs, seqlen, nkv, hs) - q_sequence_ids: torch.Tensor, - k_sequence_ids: torch.Tensor, - causal: bool = False, -) -> torch.Tensor: # (bs, seqlen, nh, hs) - assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." - - if not causal: - batch_size, q_len = query_states.shape[0], query_states.shape[1] - ( - query_states, - key_states, - value_states, - indices_q, - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_in_batch_q, max_seqlen_in_batch_k), - ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids) - - attn_output_unpad = _kernels_flash_varlen_forward( - query_states, - key_states, - value_states, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_in_batch_q=max_seqlen_in_batch_q, - max_seqlen_in_batch_k=max_seqlen_in_batch_k, - causal=False, - ) - attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len) - - else: - attn_output = _kernels_flash_forward(query_states, key_states, value_states, causal=True) - - return attn_output - - -def block_min_max_seq_ids(SLEN: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: - device = SLEN.device - total_tokens = torch.sum(SLEN) - B = (total_tokens + block_size - 1) // block_size - padding_tokens = B * block_size - total_tokens - SLEN = torch.cat([SLEN, padding_tokens.reshape(1).to(device=device, dtype=SLEN.dtype)], dim=0) - - assert torch.sum(SLEN) == B * block_size - - # Cumulative ends (exclusive) for each sequence; cum[i] == end offset of seq i - cum = torch.cumsum(SLEN.to(torch.long), dim=0) # (N,) - total_tokens = cum[-1].item() - - # Block start/end offsets [start, end) in token index space - block_starts = torch.arange(0, B * block_size, block_size, device=device, dtype=torch.long) # (B,) - block_ends = torch.minimum(block_starts + block_size, torch.tensor(total_tokens, device=device)) # (B,) - - # MIN_SEQ_ID[i] = first sequence whose end > block_start - # searchsorted with right=True returns first index where cum > value - MIN_SEQ_ID = torch.searchsorted(cum, block_starts, right=True) - - # MAX_SEQ_ID[i] = sequence containing the last token in the block (block_end - 1) - # For empty tail beyond total_tokens we already clipped block_ends. - last_token_in_block = torch.clamp(block_ends - 1, min=0) # valid only if block has at least 1 token - MAX_SEQ_ID = torch.searchsorted(cum, last_token_in_block, right=True) - - return MIN_SEQ_ID, MAX_SEQ_ID - - -def get_overlapping_blocks(SLEN_Q: torch.Tensor, SLEN_K: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - MIN_Q, MAX_Q = block_min_max_seq_ids(SLEN_Q) - MIN_K, MAX_K = block_min_max_seq_ids(SLEN_K) - - cond1 = MIN_Q.unsqueeze(1) <= MAX_K.unsqueeze(0) - cond2 = MIN_K.unsqueeze(0) <= MAX_Q.unsqueeze(1) - overlap = cond1 & cond2 - - cond1 = (MIN_Q == MAX_Q).unsqueeze(1) - cond2 = (MIN_K == MAX_K).unsqueeze(0) - same_seq_in_qk = cond1 & cond2 - - full_blocks = overlap & same_seq_in_qk - partial_blocks = overlap & ~same_seq_in_qk - - return full_blocks, partial_blocks - - -@torch.compiler.disable -def direct_block_mask(SLEN_Q: torch.Tensor, SLEN_K: torch.Tensor) -> BlockMask: - full_blocks, partial_blocks = get_overlapping_blocks(SLEN_Q, SLEN_K) - partial_blocks = partial_blocks[None, None] - full_blocks = full_blocks[None, None] - - q_doc_id = torch.repeat_interleave(SLEN_Q) - k_doc_id = torch.repeat_interleave(SLEN_K) - - def doc_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor: - return q_doc_id[q_idx] == k_doc_id[kv_idx] - - total_q_len = q_doc_id.shape[0] - total_k_len = k_doc_id.shape[0] - - return _create_sparse_block_from_block_mask( - (partial_blocks, full_blocks), - doc_mask, - seq_lengths=(total_q_len, total_k_len), - Q_BLOCK_SIZE=128, - KV_BLOCK_SIZE=128, - ) - - -@torch.compiler.disable -def doc_id_mask(SLEN_Q: torch.Tensor, SLEN_K: torch.Tensor) -> BlockMask: - q_doc_id = torch.repeat_interleave(SLEN_Q) - k_doc_id = torch.repeat_interleave(SLEN_K) - - def doc_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor: - return q_doc_id[q_idx] == k_doc_id[kv_idx] - - total_q_len = q_doc_id.shape[0] - total_k_len = k_doc_id.shape[0] - - return create_block_mask(doc_mask, 1, 1, total_q_len, total_k_len, BLOCK_SIZE=128, device=SLEN_Q.device) - - -def varlen_flex_attention_func( - query_states: torch.Tensor, - key_states: torch.Tensor, - value_states: torch.Tensor, - q_sequence_ids: torch.Tensor, - k_sequence_ids: torch.Tensor, -) -> torch.Tensor: - batch_size, q_len = query_states.shape[0], query_states.shape[1] - ( - query_states, - key_states, - value_states, - indices_q, - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_in_batch_q, max_seqlen_in_batch_k), - ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids) - - query_states = query_states.unsqueeze(0).transpose(1, 2).contiguous() - key_states = key_states.unsqueeze(0).transpose(1, 2).contiguous() - value_states = value_states.unsqueeze(0).transpose(1, 2).contiguous() - - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1] - block_mask = block_mask_creator(seqlens_q, seqlens_k) - - fn = _get_flex_attention_fn() - attn_output_unpad = fn( - query_states, - key_states, - value_states, - block_mask=block_mask, - enable_gqa=query_states.shape[1] != key_states.shape[1], - ) - - attn_output = pad_input(attn_output_unpad.transpose(1, 2).squeeze(0), indices_q, batch_size, q_len) - - return attn_output - - -def _get_unpad_data(sequence_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, int]: - non_pad_indices = sequence_ids != -1 - non_pad_indices = torch.nonzero(non_pad_indices.flatten(), as_tuple=False).flatten() - sequence_ids = sequence_ids + torch.arange(len(sequence_ids), device=sequence_ids.device)[:, None] * 1e5 - sequence_ids = sequence_ids.flatten()[non_pad_indices] - _, seqlens_in_batch = torch.unique_consecutive(sequence_ids, return_counts=True) - max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) - return non_pad_indices, cu_seqlens, max_seqlen_in_batch - - -def _unpad_input( - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - q_sequence_ids: torch.Tensor, - k_sequence_ids: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]: - batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape - query_length, num_q_heads = query_layer.shape[1], query_layer.shape[2] - assert query_layer.shape[:2] == q_sequence_ids.shape, ( - f"Shape mismatch between query layer and query sequence ids: {query_layer.shape[:2]} != {q_sequence_ids.shape}" - ) - assert key_layer.shape[:2] == k_sequence_ids.shape, ( - f"Shape mismatch between key layer and key sequence ids: {key_layer.shape[:2]} != {k_sequence_ids.shape}" - ) - assert query_length <= kv_seq_len, ( - f"Query length should be less than or equal to KV sequence length: {query_length} <= {kv_seq_len}" - ) - - indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(k_sequence_ids) - - key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) - value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) - - if torch.equal(q_sequence_ids, k_sequence_ids): - indices_q = indices_k - cu_seqlens_q = cu_seqlens_k - max_seqlen_in_batch_q = max_seqlen_in_batch_k - else: - indices_q, cu_seqlens_q, max_seqlen_in_batch_q = _get_unpad_data(q_sequence_ids) - - query_layer = index_first_axis(query_layer.reshape(batch_size * query_length, num_q_heads, head_dim), indices_q) - - assert cu_seqlens_q.shape == cu_seqlens_k.shape, ( - f"Query and KV should have the same number of sequences: {cu_seqlens_q.shape} != {cu_seqlens_k.shape}" - ) - - return ( - query_layer, - key_layer, - value_layer, - indices_q, - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_in_batch_q, max_seqlen_in_batch_k), - ) - - -block_mask_creator = direct_block_mask if os.getenv("FAST_BLOCK_MASK", "1") == "1" else doc_id_mask -PAD_TOKEN_ID = 0 -BOS_TOKEN_ID = 1 -EOS_TOKEN_ID = 2 -E1_VOCAB_SIZE = 34 -E1_TOKENIZER_REPO_ID = "Synthyra/Profluent-E1-150M" - - -def _load_tokenizer_file(fname: str) -> Tokenizer: - tokenizer: Tokenizer = Tokenizer.from_file(fname) - assert tokenizer.padding["pad_id"] == PAD_TOKEN_ID, ( - f"Padding token id must be {PAD_TOKEN_ID}, but got {tokenizer.padding['pad_id']}" - ) - return tokenizer - - -def get_tokenizer( - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - *, - local_files_only: bool = False, - cache_dir: Optional[Union[str, os.PathLike]] = None, - revision: Optional[str] = None, - token: Optional[Union[str, bool]] = None, -) -> Tokenizer: - source_path = None - checked_local_source = False - if pretrained_model_name_or_path is not None: - source_path = os.fspath(pretrained_model_name_or_path) - if os.path.isdir(source_path): - checked_local_source = True - fname = os.path.join(source_path, "tokenizer.json") - if os.path.isfile(fname): - return _load_tokenizer_file(fname) - - fname = os.path.join(os.path.dirname(__file__), "tokenizer.json") - if os.path.isfile(fname): - return _load_tokenizer_file(fname) - - if local_files_only and checked_local_source: - raise FileNotFoundError( - f"E1 tokenizer.json was not found in {source_path} or next to {__file__}." - ) - - from huggingface_hub import hf_hub_download - - repo_id = E1_TOKENIZER_REPO_ID - if source_path is not None and not checked_local_source: - repo_id = source_path - try: - fname = hf_hub_download( - repo_id=repo_id, - filename="tokenizer.json", - cache_dir=os.fspath(cache_dir) if cache_dir is not None else None, - revision=revision, - token=token, - local_files_only=local_files_only, - ) - except Exception as error: - raise FileNotFoundError( - f"E1 tokenizer.json was not found locally and could not be loaded from {repo_id}." - ) from error - return _load_tokenizer_file(fname) - - -@dataclass -class DataPrepConfig: - max_num_sequences: int = 512 - max_num_positions_within_seq: int = 8192 - remove_X_tokens: bool = False - - -def get_context(sequence: str) -> Optional[str]: - if "," in sequence: - return sequence.rsplit(",", 1)[0] - return None - - -class E1BatchPreparer: - def __init__( - self, - data_prep_config: Optional[DataPrepConfig] = None, - tokenizer: Optional[Tokenizer] = None, - tokenizer_source: Optional[Union[str, os.PathLike]] = None, - local_files_only: bool = False, - cache_dir: Optional[Union[str, os.PathLike]] = None, - revision: Optional[str] = None, - token: Optional[Union[str, bool]] = None, - preserve_context_labels: bool = False, - ): - self.tokenizer = tokenizer or get_tokenizer( - tokenizer_source, - local_files_only=local_files_only, - cache_dir=cache_dir, - revision=revision, - token=token, - ) - self.data_prep_config = data_prep_config or DataPrepConfig() - self.pad_token_id = self.tokenizer.token_to_id("") - self.preserve_context_labels = preserve_context_labels - device = torch.cuda.current_device() if torch.cuda.is_available() else torch.device("cpu") - self.boundary_token_ids = torch.tensor( - [self.tokenizer.token_to_id(token) for token in ["", "", "1", "2", ""]], device=device - ).long() - self.mask_token = "?" # nosec - self.mask_token_id = self.tokenizer.token_to_id(self.mask_token) - self.X_token_id = self.tokenizer.token_to_id("X") - self.vocab = self.tokenizer.get_vocab() - - def get_batch_kwargs( # type: ignore[override] - self, sequences: List[str], device: torch.device = torch.device("cpu"), non_blocking: bool = False - ) -> Dict[str, Union[torch.Tensor, List[str], List[int]]]: - sequence_encodings = [self.prepare_multiseq(sequence) for sequence in sequences] - return self.pad_encodings(sequence_encodings, device, non_blocking) - - def pad_encodings( - self, - sequence_encodings: List[Dict[str, torch.Tensor]], - device: torch.device = torch.device("cpu"), - non_blocking: bool = False, - ) -> Dict[str, Union[torch.Tensor, List[str], List[int]]]: - non_blocking = non_blocking and device.type == "cuda" - padded_encodings = {} - # Note: We use -1 as the padding value for sequence and position ids because the 0 value - # is a valid value for sequence and position ids. -1 is then used to distinguish valid - # tokens from padding tokens, for example, when doing padding/unpadding for flash attention. - for key, padding_value in { - "input_ids": self.pad_token_id, - "sequence_ids": -1, - "within_seq_position_ids": -1, - "global_position_ids": -1, - "labels": self.pad_token_id, - }.items(): - padded_encodings[key] = pad_sequence( - [enc[key] for enc in sequence_encodings], batch_first=True, padding_value=padding_value - ).to(device=device, dtype=torch.long, non_blocking=non_blocking) - - padded_encodings["context"] = [enc["context"] for enc in sequence_encodings] - padded_encodings["context_len"] = [enc["context_len"] for enc in sequence_encodings] - - return padded_encodings - - def prepare_multiseq(self, sequence: str) -> Dict[str, Union[torch.Tensor, str, int]]: - single_sequences = sequence.split(",") - if len(single_sequences) > self.data_prep_config.max_num_sequences: - raise ValueError( - f"Number of sequences {len(single_sequences)} exceeds max number of sequences {self.data_prep_config.max_num_sequences}" - " in the provided multi-sequence instance. Please remove some homologous sequences before trying again." - ) - - single_sequence_encodings = [self.prepare_singleseq(sequence) for sequence in single_sequences] - - num_tokens = [len(x["input_ids"]) for x in single_sequence_encodings] - input_ids = torch.cat([x["input_ids"] for x in single_sequence_encodings]) - labels = torch.cat([x["labels"] for x in single_sequence_encodings]) - - within_seq_position_ids = torch.cat([encoding["position_ids"] for encoding in single_sequence_encodings]) - global_position_ids, ctx_len = [], 0 - for encoding in single_sequence_encodings: - global_position_ids.append(encoding["position_ids"] + ctx_len) - ctx_len = max(ctx_len, encoding["position_ids"].max().item() + ctx_len + 1) - global_position_ids = torch.cat(global_position_ids) - - sequence_ids = torch.repeat_interleave(torch.tensor(num_tokens)) - - # Get multi-seq context & mask out all but last sequence in multi-seq instance if desired - context_len = sum(num_tokens[:-1]) - context = self.tokenizer.decode(input_ids[:context_len].tolist(), skip_special_tokens=False) - if not self.preserve_context_labels: - labels[:context_len] = self.pad_token_id - - assert ( - input_ids.shape - == sequence_ids.shape - == within_seq_position_ids.shape - == global_position_ids.shape - == labels.shape - ), "Input ids, sequence ids, within seq position ids, global position ids, and labels must have the same shape" - - assert input_ids.shape[0] >= context_len, "Input ids must have at least as many tokens as the context length" - - return { - "input_ids": input_ids, - "sequence_ids": sequence_ids, - "within_seq_position_ids": within_seq_position_ids, - "global_position_ids": global_position_ids, - "labels": labels, - "context": context, - "context_len": context_len, - } - - def prepare_singleseq(self, sequence: str) -> Dict[str, torch.Tensor]: - if not self.validate_sequence(sequence): - raise ValueError(f"Invalid sequence: {sequence}; Input sequence should contain [A-Z] or ? characters only") - - if len(sequence) > self.data_prep_config.max_num_positions_within_seq: - raise ValueError( - f"Sequence length {len(sequence)} exceeds max length {self.data_prep_config.max_num_positions_within_seq}" - ) - - # Can also use `tokens = torch.tensor(self.tokenizer.encode(f"1{sequence}2").ids)` - # but following is faster since our vocabulary is simple. - tokens = torch.tensor([self.vocab[token] for token in ["", "1", *sequence, "2", ""]]) - position_ids = torch.arange(len(tokens)) - - if self.data_prep_config.remove_X_tokens: - X_positions = torch.where(tokens != self.X_token_id)[0] - tokens = tokens[X_positions] - position_ids = position_ids[X_positions] - - return {"input_ids": tokens, "labels": tokens, "position_ids": position_ids} - - def get_boundary_token_mask(self, tokens: torch.Tensor) -> torch.BoolTensor: - return torch.isin(tokens, self.boundary_token_ids.to(tokens.device)) - - def get_mask_positions_mask(self, tokens: torch.Tensor) -> torch.BoolTensor: - return tokens == self.mask_token_id - - def validate_sequence(self, sequence: str) -> bool: - assert isinstance(sequence, str), "Sequence must be a string" - sequence = sequence.replace(self.mask_token, "") - return sequence.isalpha() and sequence.isupper() - - -class E1Config(PretrainedConfig): - model_type = "E1" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( # type: ignore - self, - # Model architecture/initialization - vocab_size=None, - hidden_size=4096, - intermediate_size=16384, - gated_mlp=False, - num_hidden_layers=40, - num_attention_heads=32, - num_key_value_heads=8, - hidden_act="silu", - rms_norm_eps=1e-5, - initializer_range=0.02, - dtype="bfloat16", - gradient_checkpointing=False, - no_ffn_gradient_checkpointing=False, - # Tokenization - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - tie_word_embeddings=False, - # Attention implementation & rotary positional embeddings - global_attention_every_n_layers=0, - max_num_sequences=512, - max_num_positions_within_seq=8192, - max_num_positions_global=1024 * 128, - rope_theta_within_seq=10000.0, - rope_theta_global=100000.0, - clip_qkv=None, - attn_backend="sdpa", - **kwargs, - ) -> None: - super().__init__( - pad_token_id=PAD_TOKEN_ID, - bos_token_id=BOS_TOKEN_ID, - eos_token_id=EOS_TOKEN_ID, - tie_word_embeddings=tie_word_embeddings, - dtype=dtype, - **kwargs, - ) - - self.hidden_size = hidden_size - if intermediate_size is None: - intermediate_size = 3 * hidden_size if gated_mlp else 4 * hidden_size - self.intermediate_size = intermediate_size - self.gated_mlp = gated_mlp - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.max_num_positions_within_seq = max_num_positions_within_seq - self.max_num_positions_global = max_num_positions_global - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.rope_theta_within_seq = rope_theta_within_seq - self.rope_theta_global = rope_theta_global - self.max_num_sequences = max_num_sequences - assert clip_qkv is None or clip_qkv > 0 - self.clip_qkv = clip_qkv - self.global_attention_every_n_layers = global_attention_every_n_layers - - self.vocab_size = E1_VOCAB_SIZE - self.gradient_checkpointing = gradient_checkpointing - self.no_ffn_gradient_checkpointing = no_ffn_gradient_checkpointing - self.attn_backend = attn_backend - - if vocab_size is not None: - if vocab_size < self.vocab_size: - logger.warning( - f"Using vocab_size {vocab_size} smaller than {self.vocab_size} from tokenizer. MAKE SURE THIS IS INTENTIONAL." - ) - self.vocab_size = vocab_size - elif vocab_size > self.vocab_size: - logger.warning( - f"Using vocab_size {vocab_size} instead of smaller {self.vocab_size} " - "from E1 tokenizer contract." - ) - self.vocab_size = vocab_size - if pad_token_id is not None and pad_token_id != self.pad_token_id: - logger.warning(f"Ignoring pad_token_id. Using {self.pad_token_id} from E1 tokenizer contract") - if bos_token_id is not None and bos_token_id != self.bos_token_id: - logger.warning(f"Ignoring bos_token_id. Using {self.bos_token_id} from E1 tokenizer contract") - if eos_token_id is not None and eos_token_id != self.eos_token_id: - logger.warning(f"Ignoring eos_token_id. Using {self.eos_token_id} from E1 tokenizer contract") - - -class DynamicCache: - """ - A cache layer that grows dynamically as more tokens are generated. This is the default for generative models. - It stores the key and value states as tensors of shape `[batch_size, seq_len, num_heads, head_dim]`. - - Args: - key_cache (`list[torch.Tensor]`): The list of key states. - value_cache (`list[torch.Tensor]`): The list of value states. - """ - - def __init__(self) -> None: - self.key_cache: List[torch.Tensor] = [] - self.value_cache: List[torch.Tensor] = [] - - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int - ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Update the key and value caches in-place, and return the necessary keys and value states. - - Args: - key_states (`torch.Tensor`): The new key states to cache of shape [batch_size, seq_len, num_heads, head_dim] - value_states (`torch.Tensor`): The new value states to cache of shape [batch_size, seq_len, num_heads, head_dim] - layer_idx (`int`): The index of the layer to update. - - Returns: - tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states of shape [batch_size, seq_len, num_heads, head_dim]. - """ - # Lazy initialization - if len(self.key_cache) <= layer_idx: - # There may be skipped layers, fill them with empty lists - for _ in range(len(self.key_cache), layer_idx): - self.key_cache.append(torch.tensor([])) - self.value_cache.append(torch.tensor([])) - self.key_cache.append(key_states) - self.value_cache.append(value_states) - elif ( - not self.key_cache[layer_idx].numel() # prefers not t.numel() to len(t) == 0 to export the model - ): # fills previously skipped layers; checking for tensor causes errors - self.key_cache[layer_idx] = key_states - self.value_cache[layer_idx] = value_states - else: - self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=1) - self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=1) - - return self.key_cache[layer_idx], self.value_cache[layer_idx] - - def get_seq_length(self, layer_idx: int = 0) -> int: - """Returns the sequence length of the cached states. A layer index can be optionally passed.""" - is_empty_layer = ( - len(self.key_cache) == 0 # no cache in any layer - or len(self.key_cache) <= layer_idx # skipped `layer_idx` and hasn't run a layer with cache after it - or not self.key_cache[layer_idx].numel() # the layer has no cache - ) - layer_seq_length = self.key_cache[layer_idx].shape[1] if not is_empty_layer else 0 - return layer_seq_length - - def crop(self, max_length: int) -> None: - """Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be - negative to remove `max_length` tokens. This is used in assisted decoding and contrastive search.""" - assert max_length > 0, "max_length must be positive" - - if self.get_seq_length() <= max_length: - return - - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx].numel(): - self.key_cache[layer_idx] = self.key_cache[layer_idx][:, :max_length, ...] - self.value_cache[layer_idx] = self.value_cache[layer_idx][:, :max_length, ...] - - def batch_repeat_interleave(self, repeats: int) -> None: - """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx].numel(): - self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(repeats, dim=0) - self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(repeats, dim=0) - - def batch_select_indices(self, indices: torch.Tensor) -> None: - """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.""" - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx].numel(): - self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...] - self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...] - - -class KVCache: - def __init__(self, cache_size: int = 4) -> None: - self.cache_size = cache_size - self.tensor_input_field_names = [ - "input_ids", - "within_seq_position_ids", - "global_position_ids", - "sequence_ids", - "labels", - ] - self.tensor_output_field_names = ["logits", "embeddings"] - self.cache_dict: Dict[str, DynamicCache] = {} - self.cache_queue: List[str] = [] - - def reset(self) -> None: - for k in list(self.cache_dict.keys()): - del self.cache_dict[k] - del self.cache_dict - self.cache_dict = {} - self.cache_queue = [] - - torch.cuda.empty_cache() - - def before_forward(self, batch: Dict[str, torch.Tensor]) -> None: - contexts: Optional[List[str]] = batch.get("context", None) - if contexts is None or "context_len" not in batch: - logger.warning_once( - "KVCache requires the batch dict to have both `context` and `context_len` keys to trigger. Skipping." - ) - return - - context_lens: List[int] = list(set(batch["context_len"])) - contexts: List[str] = list(set(contexts)) # type: ignore[no-redef] - if len(contexts) != 1 or len(context_lens) != 1: - logger.warning( - "SingleContextKVCache requires a single context and context length. " - "Multiple contexts or context lengths found in a single batch. Skipping." - ) - return - - batch_size = batch["input_ids"].shape[0] - - unique_context = contexts[0] - unique_context_len = context_lens[0] - batch["use_cache"] = True - - if unique_context not in self.cache_dict: - return - - self.cache_dict[unique_context].batch_repeat_interleave(batch_size) - past_key_values = self.cache_dict[unique_context] - batch["past_key_values"] = past_key_values - - # Remove context from the input fields - for field_name in self.tensor_input_field_names: - if batch.get(field_name, None) is not None: - batch[field_name] = batch[field_name][:, unique_context_len:] - - def after_forward(self, batch: Dict[str, Any], outputs: ModelOutput) -> None: - contexts = batch.get("context", None) - context_lens = batch.get("context_len", []) - if contexts is None or len(set(contexts)) != 1 or len(set(context_lens)) != 1 or context_lens[0] == 0: - return - - assert batch["use_cache"] - unique_context = contexts[0] - unique_context_len = context_lens[0] - - past_key_values = getattr(outputs, "past_key_values", None) - if not isinstance(past_key_values, DynamicCache): - logger.warning_once("KVCache is incompatible with models that don't return a DynamicCache. Skipping.") - return - - if "past_key_values" not in batch: - if len(self.cache_queue) == self.cache_size: - last_context = self.cache_queue.pop(0) - if last_context not in self.cache_queue: - del self.cache_dict[last_context] - torch.cuda.empty_cache() - - self.cache_dict[unique_context] = past_key_values - self.cache_queue.append(unique_context) - - # Remove context from the input fields - for field_name in self.tensor_input_field_names: - if field_name in batch and batch[field_name] is not None: - batch[field_name] = batch[field_name][:, unique_context_len:] - - # Remove context from the output fields - for field_name in self.tensor_output_field_names: - if field_name in outputs and outputs[field_name] is not None: - outputs[field_name] = outputs[field_name][:, unique_context_len:] - if "hidden_states" in outputs and outputs["hidden_states"] is not None: - outputs["hidden_states"] = [h[:, unique_context_len:] for h in outputs["hidden_states"]] - - self.cache_dict[unique_context].crop(unique_context_len) - self.cache_dict[unique_context].batch_select_indices([0]) - - -DOCKER_IMAGE = "ghcr.io/soedinglab/mmseqs2" -COLABFOLD_HOST = "https://api.colabfold.com" -LOWERCASE_CHARS = b"abcdefghijklmnopqrstuvwxyz" -DEFAULT_MAX_CONTEXT_TOKENS = [6144, 12288, 24576] -DEFAULT_SIMILARITY_THRESHOLDS = [1.0, 0.95, 0.9, 0.7, 0.5] -DEFAULT_EMBED_MAX_TOKENS = 8192 -DEFAULT_EMBED_SIMILARITY = 0.95 - -IdSequence = namedtuple("IdSequence", ["id", "sequence"]) -IndexedSequence = Tuple[int, str] - - -@dataclass -class ContextSpecification: - max_num_samples: int = 511 - max_token_length: int = 32768 - max_query_similarity: float = 1.0 - min_query_similarity: float = 0.0 - neighbor_similarity_lower_bound: float = 0.8 - - -class E1Prediction(TypedDict, total=False): - id: str | int - context_id: str | int | None - logits: torch.Tensor - token_embeddings: torch.Tensor - mean_token_embeddings: torch.Tensor - - -def read_fasta_sequences(path: str) -> Dict[str, str]: - sequences: Dict[str, str] = {} - header: Optional[str] = None - parts: List[str] = [] - with open(path, "r", encoding="utf-8") as handle: - for raw_line in handle: - line = raw_line.strip() - if not line: - continue - if line.startswith(">"): - if header is not None: - sequences[header] = "".join(parts) - header = line[1:].strip() - parts = [] - else: - assert header is not None, f"FASTA sequence found before header in {path}" - parts.append(line) - if header is not None: - sequences[header] = "".join(parts) - return sequences - - -def write_fasta_sequences(path: str, sequences: Dict[str, str]) -> None: - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path, "w", encoding="utf-8") as handle: - for header, sequence in sequences.items(): - handle.write(f">{header}\n{sequence}\n") - - -def parse_msa(path: str) -> List[IdSequence]: - records = read_fasta_sequences(path) - sequences = [] - for record_id, record_seq in records.items(): - sequence = str(record_seq).replace("\x00", "").replace(".", "-") - sequences.append(IdSequence(record_id, sequence)) - assert len(sequences) > 0, f"No sequences found in MSA file: {path}" - return sequences - - -def convert_to_tensor(sequences: List[IdSequence], device: Optional[torch.device] = None) -> torch.ByteTensor: - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - byte_sequences = [ - sequence.sequence.encode("ascii").translate(None, LOWERCASE_CHARS) - for sequence in sequences - ] - lengths = {len(byte_sequence) for byte_sequence in byte_sequences} - assert len(lengths) == 1, f"MSA rows must have equal aligned lengths after removing insertions: {sorted(lengths)}" - array = np.vstack([np.frombuffer(byte_sequence, dtype=np.uint8) for byte_sequence in byte_sequences]) - return torch.from_numpy(array).to(device) - - -def get_num_neighbors(byte_seqs: torch.ByteTensor, sim_threshold: float = 0.8) -> List[int]: - gap_token_id = np.frombuffer(b"-", np.uint8)[0].item() - seq_lens = (byte_seqs != gap_token_id).sum(dim=1) - num_neighbors: List[int] = [] - for i in range(byte_seqs.shape[0]): - query_non_gaps = byte_seqs[i] != gap_token_id - seqs_sim = (byte_seqs[:, query_non_gaps] == byte_seqs[i, query_non_gaps]).sum(dim=1) / seq_lens - num_neighbors.append(int((seqs_sim >= sim_threshold).sum().item())) - return num_neighbors - - -def get_similarity_to_query(byte_seqs: torch.ByteTensor) -> torch.FloatTensor: - return (byte_seqs == byte_seqs[0, :]).sum(dim=1) / byte_seqs.shape[1] - - -def sample_context( - msa_path: str, - max_num_samples: int, - max_token_length: int, - max_query_similarity: float = 1.0, - min_query_similarity: float = 0.0, - neighbor_similarity_lower_bound: float = 0.8, - use_full_sequences_in_context: bool = False, - full_sequences_path: Optional[str] = None, - seed: int = 0, - device: Optional[torch.device] = None, - cache_num_neighbors_path: Optional[str] = None, -) -> Tuple[str, List[str]]: - msa_sequences = parse_msa(msa_path) - msa_as_byte_tensor = convert_to_tensor(msa_sequences, device) - if cache_num_neighbors_path is not None and os.path.exists(cache_num_neighbors_path): - num_neighbors = np.load(cache_num_neighbors_path) - else: - num_neighbors = np.array(get_num_neighbors(msa_as_byte_tensor, neighbor_similarity_lower_bound)) - if cache_num_neighbors_path is not None: - np.save(cache_num_neighbors_path, num_neighbors) - - sampling_weights = 1.0 / num_neighbors - query_similarity = get_similarity_to_query(msa_as_byte_tensor) - filtered_mask = (query_similarity <= max_query_similarity) & (query_similarity >= min_query_similarity) - assert filtered_mask.sum() >= 1, ( - f"No sequences found with similarity to query within range " - f"{min_query_similarity} <= query_similarity <= {max_query_similarity}." - ) - - filtered_weights = np.where(filtered_mask.cpu().numpy(), sampling_weights, 0.0) - sampled_indices = np.random.default_rng(seed).choice( - len(filtered_weights), - size=min(max_num_samples, int(filtered_mask.sum())), - p=filtered_weights / filtered_weights.sum(), - replace=False, - shuffle=True, - ) - - if use_full_sequences_in_context: - assert full_sequences_path is not None, "full_sequences_path is required when use_full_sequences_in_context=True" - full_sequences = parse_msa(full_sequences_path) - assert len(full_sequences) == len(msa_sequences), "Number of full sequences must match number of MSA sequences" - for i, (full_seq, msa_seq) in enumerate(zip(full_sequences, msa_sequences)): - assert full_seq.id == msa_seq.id, ( - "Full sequences and MSA sequences must be in the same order and have the same ids. " - f"Found differing id for sample {i}: {full_seq.id} != {msa_seq.id}" - ) - sampled_sequences = [full_sequences[int(i)] for i in sampled_indices] - else: - sampled_sequences = [msa_sequences[int(i)] for i in sampled_indices] - - context_sequences: List[str] = [] - context_ids: List[str] = [] - context_length = 0 - for seq in sampled_sequences: - seq_str = seq.sequence.upper().encode("ascii").translate(None, b"-").decode("ascii") - if context_length + len(seq_str) > max_token_length: - break - context_sequences.append(seq_str) - context_ids.append(seq.id) - context_length += len(seq_str) - return ",".join(context_sequences), context_ids - - -def sample_multiple_contexts( - msa_path: str, - context_specifications: List[ContextSpecification], - use_full_sequences_in_context: bool = False, - full_sequences_path: Optional[str] = None, - seed: int = 0, - device: Optional[torch.device] = None, - cache_num_neighbors_path: Optional[str] = None, -) -> Tuple[List[str], List[List[str]]]: - with tempfile.TemporaryDirectory() as temp_dir: - if cache_num_neighbors_path is None: - cache_num_neighbors_path = os.path.join(temp_dir, "num_neighbors.npy") - - contexts: List[str] = [] - context_ids: List[List[str]] = [] - for i, context_specification in enumerate(context_specifications): - context, ids = sample_context( - msa_path=msa_path, - max_num_samples=context_specification.max_num_samples, - max_token_length=context_specification.max_token_length, - max_query_similarity=context_specification.max_query_similarity, - min_query_similarity=context_specification.min_query_similarity, - neighbor_similarity_lower_bound=context_specification.neighbor_similarity_lower_bound, - use_full_sequences_in_context=use_full_sequences_in_context, - full_sequences_path=full_sequences_path, - seed=seed + i, - device=device, - cache_num_neighbors_path=cache_num_neighbors_path, - ) - contexts.append(context) - context_ids.append(ids) - return contexts, context_ids - - -def get_context_id(max_tokens: int, sim_threshold: float) -> str: - return f"identity_{sim_threshold}_tokens_{max_tokens}" - - -def build_context_specifications( - max_context_tokens: Optional[List[int]] = None, - similarity_thresholds: Optional[List[float]] = None, - min_query_similarity: float = 0.3, -) -> List[Tuple[ContextSpecification, str]]: - if max_context_tokens is None: - max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS - if similarity_thresholds is None: - similarity_thresholds = DEFAULT_SIMILARITY_THRESHOLDS - - specs = [] - for max_tokens in max_context_tokens: - for sim_threshold in similarity_thresholds: - spec = ContextSpecification( - max_num_samples=511, - max_token_length=max_tokens, - max_query_similarity=sim_threshold, - min_query_similarity=min_query_similarity, - neighbor_similarity_lower_bound=0.8, - ) - specs.append((spec, get_context_id(max_tokens, sim_threshold))) - return specs - - -def sample_contexts_for_msa( - a3m_path: str, - context_specs: List[Tuple[ContextSpecification, str]], - seed: int = 42, -) -> Dict[str, str]: - specs_only = [spec for spec, _ in context_specs] - context_ids = [context_id for _, context_id in context_specs] - contexts, _ = sample_multiple_contexts( - msa_path=a3m_path, - context_specifications=specs_only, - seed=seed, - ) - return dict(zip(context_ids, contexts)) - - -def _strip_a3m_insertions(sequence: str) -> str: - uppercase_or_gap = [char for char in sequence if char.isupper() or char in "-."] - return "".join(uppercase_or_gap).replace("-", "").replace(".", "") - - -def get_query_from_a3m(path: str) -> str: - header_found = False - seq_parts: List[str] = [] - with open(path, "r", encoding="utf-8") as handle: - for raw_line in handle: - line = raw_line.strip() - if not line: - continue - if line.startswith(">"): - if header_found: - break - header_found = True - continue - if header_found: - seq_parts.append(line) - assert header_found, f"No FASTA header found in A3M file: {path}" - return _strip_a3m_insertions("".join(seq_parts)) - - -def load_msa_dir(msa_dir: str) -> Dict[str, str]: - msa_lookup: Dict[str, str] = {} - a3m_files = list(Path(msa_dir).rglob("*.a3m")) - if not a3m_files: - raise FileNotFoundError(f"No .a3m files found in {msa_dir}") - for a3m_path in tqdm(a3m_files, desc="Loading MSAs"): - query_seq = get_query_from_a3m(str(a3m_path)) - msa_lookup[query_seq] = str(a3m_path) - logger.info("Loaded %d MSAs from %s", len(msa_lookup), msa_dir) - return msa_lookup - - -def _safe_extract_tar(tar: tarfile.TarFile, output_dir: str) -> None: - output_root = Path(output_dir).resolve() - for member in tar.getmembers(): - target = (output_root / member.name).resolve() - if output_root != target and output_root not in target.parents: - raise ValueError(f"Unsafe tar member path: {member.name}") - tar.extractall(output_root) - - -def load_msa_from_hf( - hf_path: str, - cache_dir: Optional[str] = None, - token: Optional[str] = None, -) -> Dict[str, str]: - from huggingface_hub import snapshot_download - - if cache_dir is None: - cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "fastplms_msa") - os.makedirs(cache_dir, exist_ok=True) - local_dir = os.path.join(cache_dir, hf_path.replace("/", "_")) - if not os.path.exists(local_dir) or not any(Path(local_dir).rglob("*.a3m")): - local_dir = snapshot_download( - repo_id=hf_path, - repo_type="dataset", - local_dir=local_dir, - token=token, - ) - for tar_path in Path(local_dir).rglob("*.tar.gz"): - with tarfile.open(tar_path) as tar: - _safe_extract_tar(tar, str(tar_path.parent)) - return load_msa_dir(local_dir) - - -def get_msa_for_sequence(sequence: str, msa_lookup: Dict[str, str], min_identity: float = 0.95) -> Optional[str]: - if sequence in msa_lookup: - return msa_lookup[sequence] - - best_match_path: Optional[str] = None - best_identity = 0.0 - for query_seq, a3m_path in msa_lookup.items(): - if abs(len(query_seq) - len(sequence)) > 10: - continue - min_len = min(len(query_seq), len(sequence)) - if min_len == 0: - continue - matches = sum(a == b for a, b in zip(query_seq[:min_len], sequence[:min_len])) - identity = matches / min_len - if identity > best_identity: - best_identity = identity - best_match_path = a3m_path - - if best_identity >= min_identity: - return best_match_path - return None - - -class ContextCache: - def __init__(self, cache_dir: str, specs_hash: str, seed: int) -> None: - self.cache_dir = cache_dir - self.specs_hash = specs_hash - self.seed = seed - os.makedirs(cache_dir, exist_ok=True) - - def _cache_path(self, key: str) -> str: - safe_key = hashlib.md5(key.encode()).hexdigest()[:16] - return os.path.join(self.cache_dir, f"{safe_key}_seed{self.seed}_{self.specs_hash}.pkl") - - def load(self, key: str) -> Optional[Dict[str, str]]: - path = self._cache_path(key) - if os.path.exists(path): - with open(path, "rb") as handle: - return pickle.load(handle) - return None - - def store(self, key: str, contexts: Dict[str, str]) -> None: - path = self._cache_path(key) - with open(path, "wb") as handle: - pickle.dump(contexts, handle) - - -def compute_ppll(logits: torch.Tensor, token_ids: torch.Tensor) -> float: - assert token_ids.numel() > 0, "Cannot score an empty token sequence" - if token_ids.device != logits.device: - token_ids = token_ids.to(logits.device) - if logits.shape[0] != token_ids.shape[0]: - raise ValueError(f"Logits length {logits.shape[0]} != token_ids length {token_ids.shape[0]}") - probs = logits.softmax(dim=-1) - token_probs = probs.gather(dim=1, index=token_ids.unsqueeze(1)).squeeze(1) - return float(token_probs.mean().item()) - - -class _E1ContextPredictor: - def __init__( - self, - model: PreTrainedModel, - data_prep_config: Optional[DataPrepConfig] = None, - max_batch_tokens: int = 65536, - use_cache: bool = True, - cache_size: int = 4, - save_masked_positions_only: bool = False, - fields_to_save: Optional[List[str]] = None, - keep_predictions_in_gpu: bool = False, - progress: bool = True, - ) -> None: - self.model = model - self.max_batch_tokens = max_batch_tokens - self.batch_preparer = E1BatchPreparer(data_prep_config=data_prep_config) - self.model.eval() - self.kv_cache = KVCache(cache_size=cache_size) if use_cache else None - self.fields_to_save = fields_to_save or ["logits", "token_embeddings", "mean_token_embeddings"] - self.save_masked_positions_only = save_masked_positions_only - self.keep_predictions_in_gpu = keep_predictions_in_gpu - self.progress = progress - - @property - def device(self) -> torch.device: - return next(self.model.parameters()).device - - def group_by_length(self, indexed_sequences: List[IndexedSequence]) -> List[List[IndexedSequence]]: - batches: List[List[IndexedSequence]] = [[]] - for idx, seq in sorted(indexed_sequences, key=lambda idx_seq: (len(idx_seq[1]), idx_seq[0])): - if len(batches[-1]) > 0 and len(seq) * (len(batches[-1]) + 1) > self.max_batch_tokens: - batches.append([]) - batches[-1].append((idx, seq)) - return batches - - def group_by_context(self, indexed_sequences: List[IndexedSequence]) -> List[List[IndexedSequence]]: - batches: Dict[Optional[str], List[IndexedSequence]] = defaultdict(list) - for idx, seq in indexed_sequences: - batches[get_context(seq)].append((idx, seq)) - return list(batches.values()) - - def batch_sequences(self, sequences: List[str]) -> List[List[int]]: - indexed_sequences: List[IndexedSequence] = list(enumerate(sequences)) - indexed_batches = self.group_by_context(indexed_sequences) - indexed_batches = list( - itertools.chain.from_iterable([self.group_by_length(batch) for batch in indexed_batches]) - ) - batches = [[item[0] for item in batch] for batch in indexed_batches] - assert sorted(sum(batches, [])) == list(range(len(sequences))), ( - "Batches must contain all indices with no repetition" - ) - return batches - - @torch.no_grad() - def predict_batch(self, sequences: List[str], sequence_metadata: List[Dict[str, str | int]]) -> List[E1Prediction]: - outputs = self.predict_batch_padded(sequences) - outputs["logits"] = outputs["logits"].float() - outputs["embeddings"] = outputs["embeddings"].float() - - token_mask = outputs["non_boundary_token_mask"] & outputs["last_sequence_mask"] - if self.save_masked_positions_only: - token_mask = token_mask & outputs["mask_positions_mask"] - - predictions: List[E1Prediction] = [] - for i in range(len(sequences)): - pred: E1Prediction = {"id": sequence_metadata[i]["id"]} - if "context_id" in sequence_metadata[i]: - pred["context_id"] = sequence_metadata[i]["context_id"] - if "logits" in self.fields_to_save: - pred["logits"] = outputs["logits"][i, token_mask[i]] - if not self.keep_predictions_in_gpu: - pred["logits"] = pred["logits"].to("cpu") - if "token_embeddings" in self.fields_to_save: - pred["token_embeddings"] = outputs["embeddings"][i, token_mask[i]] - if not self.keep_predictions_in_gpu: - pred["token_embeddings"] = pred["token_embeddings"].to("cpu") - if "mean_token_embeddings" in self.fields_to_save: - pred["mean_token_embeddings"] = outputs["embeddings"][i, token_mask[i]].mean(dim=0) - if not self.keep_predictions_in_gpu: - pred["mean_token_embeddings"] = pred["mean_token_embeddings"].to("cpu") - predictions.append(pred) - return predictions - - @torch.no_grad() - def predict_batch_padded(self, sequences: List[str]) -> Dict[str, torch.Tensor]: - device = self.device - autocast_enabled = device.type == "cuda" - with torch.autocast(device.type, torch.bfloat16, enabled=autocast_enabled): - batch = self.batch_preparer.get_batch_kwargs(sequences, device=device) - if self.kv_cache is not None: - self.kv_cache.before_forward(batch) - - past_key_values = batch["past_key_values"] if "past_key_values" in batch else None - use_cache = bool(batch["use_cache"]) if "use_cache" in batch else False - output: E1MaskedLMOutputWithPast = self.model( - input_ids=batch["input_ids"], - within_seq_position_ids=batch["within_seq_position_ids"], - global_position_ids=batch["global_position_ids"], - sequence_ids=batch["sequence_ids"], - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=False, - output_hidden_states=False, - ) - if self.kv_cache is not None: - self.kv_cache.after_forward(batch, output) - - padding_mask = batch["input_ids"] == self.batch_preparer.pad_token_id - last_sequence_mask = batch["sequence_ids"] == batch["sequence_ids"].max(dim=1).values[:, None] - boundary_token_mask = self.batch_preparer.get_boundary_token_mask(batch["input_ids"]) - mask_positions_mask = self.batch_preparer.get_mask_positions_mask(batch["input_ids"]) - return { - "logits": output.logits, - "embeddings": output.last_hidden_state, - "last_sequence_mask": last_sequence_mask, - "non_boundary_token_mask": ~boundary_token_mask, - "mask_positions_mask": mask_positions_mask, - "valid_token_mask": ~padding_mask, - } - - @torch.no_grad() - def predict( - self, - sequences: Sequence[str], - sequence_ids: Optional[Sequence[int | str]] = None, - context_seqs: Optional[Dict[str, str]] = None, - ) -> Iterator[E1Prediction]: - if sequence_ids is None: - sequence_ids = list(range(len(sequences))) - if context_seqs: - sequences_with_context = [ - (ctx + "," + seq, {"context_id": ctx_id, "id": sequence_id}) - for ctx_id, ctx in context_seqs.items() - for seq, sequence_id in zip(sequences, sequence_ids) - ] - else: - sequences_with_context = [(seq, {"id": sequence_id}) for seq, sequence_id in zip(sequences, sequence_ids)] - - batched_sequences, sequence_metadata = tuple(zip(*sequences_with_context)) - batches = self.batch_sequences(list(batched_sequences)) - iterator = tqdm(batches, desc="Predicting batches", disable=not self.progress) - for indices in iterator: - sequence_batch = [batched_sequences[i] for i in indices] - sequence_batch_metadata = [sequence_metadata[i] for i in indices] - yield from self.predict_batch(sequence_batch, sequence_batch_metadata) - - -def _pool_hidden_states( - hidden_list: List[torch.Tensor], - pooling_types: List[str], - device: torch.device, -) -> torch.Tensor: - pooler = Pooler(pooling_types) - max_len = max(hidden.shape[0] for hidden in hidden_list) - hidden_dim = hidden_list[0].shape[1] - batch_size = len(hidden_list) - padded = torch.zeros(batch_size, max_len, hidden_dim, device=device) - attention_mask = torch.zeros(batch_size, max_len, device=device) - for i, hidden in enumerate(hidden_list): - seq_len = hidden.shape[0] - padded[i, :seq_len] = hidden - attention_mask[i, :seq_len] = 1.0 - return pooler(padded, attention_mask) - - -def _forward_for_embedding( - model: PreTrainedModel, - sequences: List[str], - context: Optional[str], - max_batch_tokens: int, - progress: bool, -) -> List[torch.Tensor]: - predictor = _E1ContextPredictor( - model=model, - data_prep_config=DataPrepConfig(remove_X_tokens=True), - max_batch_tokens=max_batch_tokens, - fields_to_save=["token_embeddings"], - keep_predictions_in_gpu=True, - use_cache=False, - cache_size=1, - progress=progress, - ) - context_seqs = {"embed_ctx": context} if context else None - predictions = list( - predictor.predict( - sequences=sequences, - sequence_ids=list(range(len(sequences))), - context_seqs=context_seqs, - ) - ) - predictions.sort(key=lambda prediction: prediction["id"]) - return [prediction["token_embeddings"] for prediction in predictions] - - -class HomologueSearcher: - def __init__( - self, - target_db: str, - docker_image: str = DOCKER_IMAGE, - sensitivity: float = 7.5, - max_seqs: int = 1000, - min_seq_id: float = 0.0, - coverage: float = 0.8, - split_memory_limit: Optional[str] = None, - use_gpu: bool = True, - ) -> None: - self.target_db = target_db - self.docker_image = docker_image - self.sensitivity = sensitivity - self.max_seqs = max_seqs - self.min_seq_id = min_seq_id - self.coverage = coverage - self.split_memory_limit = split_memory_limit - self.use_gpu = use_gpu - - @staticmethod - def _seq_hash(sequence: str) -> str: - return hashlib.md5(sequence.encode()).hexdigest()[:12] - - def _run_docker_command(self, cmd: List[str], **kwargs) -> subprocess.CompletedProcess: - return subprocess.run(cmd, **kwargs) - - def _validate_paths_under_cwd(self, *paths: str) -> None: - cwd = os.path.abspath(os.getcwd()) - for path in paths: - absolute_path = os.path.abspath(path) - if not (absolute_path == cwd or absolute_path.startswith(cwd + os.sep)): - raise ValueError( - "Path must be under the current working directory for docker volume mount. " - f"cwd={cwd!r}, path={absolute_path!r}" - ) - - def _path_in_container(self, local_path: str) -> str: - self._validate_paths_under_cwd(local_path) - rel = os.path.relpath(os.path.abspath(local_path), start=os.path.abspath(os.getcwd())) - return rel.replace(os.sep, "/") - - def _docker_base_cmd(self) -> List[str]: - cmd = ["docker", "run", "--rm", "-v", f"{os.getcwd()}:/app", "-w", "/app"] - if self.use_gpu and torch.cuda.is_available(): - cmd.extend(["--gpus", "all"]) - cmd.append(self.docker_image) - return cmd - - def _ensure_docker_image(self) -> None: - subprocess.run(["docker", "version"], capture_output=True, text=True, check=True) - inspect = subprocess.run(["docker", "image", "inspect", self.docker_image], capture_output=True, text=True) - if inspect.returncode == 0: - return - self._run_docker_command(["docker", "pull", self.docker_image], check=True, text=True) - - def create_db(self, fasta_path: str, db_path: str) -> str: - os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) - if os.path.exists(f"{db_path}.dbtype"): - return db_path - self._ensure_docker_image() - self._validate_paths_under_cwd(fasta_path, db_path) - self._run_docker_command( - self._docker_base_cmd() + [ - "createdb", - self._path_in_container(fasta_path), - self._path_in_container(db_path), - ], - check=True, - capture_output=True, - text=True, - ) - return db_path - - def create_index(self, db_path: str, tmp_dir: Optional[str] = None) -> None: - if tmp_dir is None: - tmp_dir = os.path.join(os.path.dirname(db_path), "tmp_index") - os.makedirs(tmp_dir, exist_ok=True) - self._ensure_docker_image() - self._validate_paths_under_cwd(db_path, tmp_dir) - self._run_docker_command( - self._docker_base_cmd() + [ - "createindex", - self._path_in_container(db_path), - self._path_in_container(tmp_dir), - ], - check=True, - capture_output=True, - text=True, - ) - - def search(self, sequence: str, output_dir: str, seq_id: Optional[str] = None) -> str: - if seq_id is None: - seq_id = self._seq_hash(sequence) - seq_output_dir = os.path.join(output_dir, seq_id) - a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m") - if os.path.exists(a3m_output): - return a3m_output - - self._ensure_docker_image() - os.makedirs(seq_output_dir, exist_ok=True) - query_fasta = os.path.join(seq_output_dir, "query.fasta") - write_fasta_sequences(query_fasta, {seq_id: sequence}) - query_db = os.path.join(seq_output_dir, "queryDB") - result_db = os.path.join(seq_output_dir, "resultDB") - tmp_dir = os.path.join(seq_output_dir, "tmp") - os.makedirs(tmp_dir, exist_ok=True) - self._validate_paths_under_cwd(query_fasta, query_db, self.target_db, seq_output_dir, result_db, tmp_dir) - - docker_base = self._docker_base_cmd() - self._run_docker_command( - docker_base + [ - "createdb", - self._path_in_container(query_fasta), - self._path_in_container(query_db), - ], - check=True, - capture_output=True, - text=True, - ) - search_cmd = docker_base + [ - "search", - self._path_in_container(query_db), - self._path_in_container(self.target_db), - self._path_in_container(result_db), - self._path_in_container(tmp_dir), - "-s", - str(self.sensitivity), - "--max-seqs", - str(self.max_seqs), - "--min-seq-id", - str(self.min_seq_id), - "-c", - str(self.coverage), - ] - if self.split_memory_limit is not None: - search_cmd.extend(["--split-memory-limit", self.split_memory_limit]) - if self.use_gpu and torch.cuda.is_available(): - search_cmd.extend(["--gpu", "1"]) - self._run_docker_command(search_cmd, check=True, capture_output=True, text=True) - self._run_docker_command( - docker_base + [ - "result2msa", - self._path_in_container(query_db), - self._path_in_container(self.target_db), - self._path_in_container(result_db), - self._path_in_container(a3m_output), - "--msa-format-mode", - "6", - ], - check=True, - capture_output=True, - text=True, - ) - for pattern in ["queryDB*", "resultDB*"]: - for path in Path(seq_output_dir).glob(pattern): - path.unlink(missing_ok=True) - tmp_path = Path(tmp_dir) - if tmp_path.exists(): - shutil.rmtree(tmp_path, ignore_errors=True) - return a3m_output - - def batch_search( - self, - sequences: List[str], - output_dir: str, - seq_ids: Optional[List[str]] = None, - continue_on_error: bool = True, - ) -> Dict[str, str]: - if seq_ids is None: - seq_ids = [self._seq_hash(seq) for seq in sequences] - os.makedirs(output_dir, exist_ok=True) - results: Dict[str, str] = {} - for seq, sid in tqdm(list(zip(sequences, seq_ids)), desc="Searching homologues"): - try: - results[seq] = self.search(seq, output_dir, sid) - except Exception: - if not continue_on_error: - raise - return results - - -class ColabFoldSearcher: - def __init__( - self, - host_url: str = COLABFOLD_HOST, - user_agent: str = "", - mode: str = "env", - timeout: float = 30.0, - max_retries: int = 10, - base_delay: float = 1.0, - max_delay: float = 60.0, - inter_request_delay: Tuple[float, float] = (1.0, 3.0), - max_wait_time: int = 600, - ) -> None: - import requests - - self.requests = requests - self.host_url = host_url.rstrip("/") - self.mode = mode - self.timeout = timeout - self.max_retries = max_retries - self.base_delay = base_delay - self.max_delay = max_delay - self.inter_request_delay = inter_request_delay - self.max_wait_time = max_wait_time - self.session = requests.Session() - if user_agent: - self.session.headers["User-Agent"] = user_agent - - @staticmethod - def _seq_hash(sequence: str) -> str: - return hashlib.md5(sequence.encode()).hexdigest()[:12] - - def _backoff_delay(self, attempt: int) -> float: - delay = min(self.base_delay * (2 ** attempt), self.max_delay) - return delay + random.uniform(0, delay * 0.5) - - def _request_with_retries(self, method: str, url: str, **kwargs): - for attempt in range(self.max_retries): - try: - if method.upper() == "GET": - response = self.session.get(url, timeout=self.timeout, **kwargs) - else: - response = self.session.post(url, timeout=self.timeout, **kwargs) - if response.status_code == 429: - retry_after = ( - float(response.headers["Retry-After"]) - if "Retry-After" in response.headers - else self._backoff_delay(attempt) - ) - time.sleep(retry_after) - continue - if response.status_code >= 500: - time.sleep(self._backoff_delay(attempt)) - continue - return response - except (self.requests.exceptions.Timeout, self.requests.exceptions.ConnectionError): - time.sleep(self._backoff_delay(attempt)) - raise RuntimeError(f"Request to {url} failed after {self.max_retries} attempts") - - def _submit(self, sequence: str, mode: Optional[str] = None) -> Dict[str, Any]: - mode = mode or self.mode - query = f">101\n{sequence}\n" - for attempt in range(self.max_retries): - response = self._request_with_retries( - "POST", - f"{self.host_url}/ticket/msa", - data={"q": query, "mode": mode}, - ) - data = response.json() - status = data["status"] if "status" in data else "UNKNOWN" - if status in ("RATELIMIT", "UNKNOWN"): - time.sleep(self._backoff_delay(attempt)) - continue - return data - raise RuntimeError(f"Failed to submit sequence after {self.max_retries} attempts") - - def _poll(self, ticket_id: str) -> Dict[str, Any]: - total_wait = 0.0 - poll_interval = 1.0 - while True: - response = self._request_with_retries("GET", f"{self.host_url}/ticket/{ticket_id}") - data = response.json() - status = data["status"] if "status" in data else "ERROR" - if status in ("COMPLETE", "ERROR"): - return data - if status not in ("RUNNING", "PENDING", "UNKNOWN"): - return data - wait = min(poll_interval + random.uniform(0, 0.5), 5.0) - time.sleep(wait) - total_wait += wait - poll_interval = min(poll_interval + 1.0, 5.0) - if total_wait > self.max_wait_time: - raise TimeoutError(f"Job {ticket_id} did not complete within {self.max_wait_time}s") - - def _download(self, ticket_id: str, output_path: str) -> None: - response = self._request_with_retries("GET", f"{self.host_url}/result/download/{ticket_id}") - os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) - with open(output_path, "wb") as handle: - handle.write(response.content) - - def _extract_a3m(self, tar_path: str, output_dir: str, seq_id: str) -> str: - with tarfile.open(tar_path) as tar: - _safe_extract_tar(tar, output_dir) - - uniref_a3m = os.path.join(output_dir, "uniref.a3m") - env_a3m = os.path.join(output_dir, "bfd.mgnify30.metaeuk30.smag30.a3m") - a3m_files: List[str] = [] - if os.path.exists(uniref_a3m): - a3m_files.append(uniref_a3m) - if "env" in self.mode and os.path.exists(env_a3m): - a3m_files.append(env_a3m) - combined_path = os.path.join(output_dir, f"{seq_id}.a3m") - if len(a3m_files) == 1: - os.replace(a3m_files[0], combined_path) - elif len(a3m_files) > 1: - with open(combined_path, "w", encoding="utf-8") as out_handle: - for a3m_file in a3m_files: - with open(a3m_file, "r", encoding="utf-8") as in_handle: - out_handle.write(in_handle.read()) - else: - raise RuntimeError("No .a3m files found in downloaded archive") - if os.path.exists(tar_path): - os.remove(tar_path) - for a3m_file in a3m_files: - if os.path.exists(a3m_file) and a3m_file != combined_path: - os.remove(a3m_file) - return combined_path - - def search(self, sequence: str, output_dir: str, seq_id: Optional[str] = None) -> str: - if seq_id is None: - seq_id = self._seq_hash(sequence) - seq_output_dir = os.path.join(output_dir, seq_id) - a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m") - if os.path.exists(a3m_output): - return a3m_output - os.makedirs(seq_output_dir, exist_ok=True) - result = self._submit(sequence) - status = result["status"] if "status" in result else "UNKNOWN" - if status == "ERROR": - raise RuntimeError(f"ColabFold API error for {seq_id}") - if status == "MAINTENANCE": - raise RuntimeError("ColabFold API is under maintenance") - ticket_id = result["id"] - result = self._poll(ticket_id) - status = result["status"] if "status" in result else "UNKNOWN" - if status != "COMPLETE": - raise RuntimeError(f"Job failed for {seq_id}: {status}") - tar_path = os.path.join(seq_output_dir, f"{seq_id}.tar.gz") - self._download(ticket_id, tar_path) - return self._extract_a3m(tar_path, seq_output_dir, seq_id) - - def batch_search( - self, - sequences: List[str], - output_dir: str, - seq_ids: Optional[List[str]] = None, - continue_on_error: bool = True, - ) -> Dict[str, str]: - if seq_ids is None: - seq_ids = [self._seq_hash(seq) for seq in sequences] - os.makedirs(output_dir, exist_ok=True) - results: Dict[str, str] = {} - pairs = list(zip(sequences, seq_ids)) - for i, (seq, sid) in enumerate(tqdm(pairs, desc="ColabFold search")): - try: - results[seq] = self.search(seq, output_dir, sid) - except Exception: - if not continue_on_error: - raise - if i < len(pairs) - 1: - time.sleep(random.uniform(*self.inter_request_delay)) - return results - - -def _make_homologue_searcher(provider: str, target_db: Optional[str], **kwargs) -> HomologueSearcher | ColabFoldSearcher: - if provider == "mmseqs2": - assert target_db is not None, "target_db is required for MMseqs2 homologue search" - return HomologueSearcher(target_db=target_db, **kwargs) - if provider == "colabfold": - return ColabFoldSearcher(**kwargs) - raise ValueError(f"Unknown homologue search provider: {provider}") - - -class AttentionLayerType(Enum): - WITHIN_SEQ = "within_seq" - GLOBAL = "global" - - -class AttentionArgs(TypedDict, total=False): - within_seq_block_mask: Optional[BlockMask] - block_causal_block_mask: Optional[BlockMask] - within_seq_mask_4d: Optional[torch.Tensor] - block_causal_mask_4d: Optional[torch.Tensor] - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). - - The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, - num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -class RotaryPositionalEmbedding(nn.Module): - def __init__( - self, dim: int, max_position_embeddings: int = 2048, base: int = 10000, device: Optional[torch.device] = None - ): - super().__init__() - - self.dim = dim - self.base = base - self.max_position_embeddings = max_position_embeddings - inv_freq = base ** -(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - # Build here to make `torch.jit.trace` work. - self._set_sin_cos_cache(seq_len=max_position_embeddings, device=self.inv_freq.device) - - @staticmethod - def rotate_half(x: torch.Tensor) -> torch.Tensor: - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - def _set_sin_cos_cache(self, seq_len: int, device: torch.device) -> None: - # Different from paper, but it uses a different permutation in order to obtain the same calculation - self.max_seq_len_cached = seq_len - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - angles = torch.outer(t, self.inv_freq.to(device)) - angles = torch.cat((angles, angles), dim=1) - self.register_buffer("cos_cached", angles.cos(), persistent=False) - self.register_buffer("sin_cached", angles.sin(), persistent=False) - - def forward( - self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.LongTensor, seq_len: Optional[int] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: - # x: [bsz, seq_len, num_attention_heads, head_size] - device, dtype = q.device, q.dtype - seq_len = position_ids.max().item() + 1 if seq_len is None else seq_len - - if seq_len > self.max_seq_len_cached: - self._set_sin_cos_cache(seq_len=seq_len, device=device) - - # angles_cached[position_ids] gets us something of shape (batch_size, seq_len, head_dim), - # so unsqueeze dimension -2 to broadcast to (batch_size, seq_len, n_heads, head_dim). - idxs = position_ids.to(device) - cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] - sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] - - # Apply rotary positional embeddings to q and k (treating them as complex numbers). The first half is - # Re[x exp(it)] = Re[x] cos(t) - Im[x] sin(t), while the second half is - # Im[x exp(it)] = Im[x] cos(t) + Re[x] sin(t). This works b/c both halves of cos/sin are the same. - q_embed = (q * cos) + (self.rotate_half(q) * sin) - k_embed = (k * cos) + (self.rotate_half(k) * sin) - return q_embed, k_embed - - -class Attention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper.""" - - def __init__(self, config: E1Config, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.hidden_size // self.num_heads - self.num_kv_heads = config.num_key_value_heads - self.num_key_value_groups = self.num_heads // self.num_kv_heads - self.max_num_seqs = config.max_num_sequences - self.clip_qkv = config.clip_qkv - - if (self.head_dim * self.num_heads) != self.hidden_size: - raise ValueError( - f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" - f" and `num_heads`: {self.num_heads})." - ) - self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) - self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) - self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) - self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) - - if self.config.global_attention_every_n_layers > 0: - self.layer_type = ( - AttentionLayerType.GLOBAL - if (self.layer_idx + 1) % self.config.global_attention_every_n_layers == 0 - else AttentionLayerType.WITHIN_SEQ - ) - else: - self.layer_type = AttentionLayerType.WITHIN_SEQ - - self.rope_theta = ( - config.rope_theta_within_seq - if self.layer_type == AttentionLayerType.WITHIN_SEQ - else config.rope_theta_global - ) - self.max_position_embeddings = ( - config.max_num_positions_within_seq - if self.layer_type == AttentionLayerType.WITHIN_SEQ - else config.max_num_positions_global - ) - - self.rotary_emb = RotaryPositionalEmbedding( - self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta - ) - - self.attn_backend = resolve_attention_backend(config.attn_backend) - - def prepare_qkv( - self, - hidden_states: torch.Tensor, - position_ids: torch.LongTensor, - past_key_value: Optional[DynamicCache] = None, - use_cache: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - bsz, q_len, _ = hidden_states.size() - query_states: torch.Tensor = self.q_proj(hidden_states) - key_states: torch.Tensor = self.k_proj(hidden_states) - val_states: torch.Tensor = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) - key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim) - val_states = val_states.view(bsz, q_len, self.num_kv_heads, self.head_dim) - - if self.clip_qkv is not None: - query_states = query_states.clamp(-self.clip_qkv, self.clip_qkv) - key_states = key_states.clamp(-self.clip_qkv, self.clip_qkv) - val_states = val_states.clamp(-self.clip_qkv, self.clip_qkv) - - query_states, key_states = self.rotary_emb(query_states, key_states, position_ids) - - if use_cache and past_key_value is not None: - key_states, val_states = past_key_value.update(key_states, val_states, self.layer_idx) - - input_dtype = query_states.dtype - if torch.is_autocast_enabled(): - target_dtype = torch.get_autocast_gpu_dtype() - else: - target_dtype = self.q_proj.weight.dtype - if input_dtype != target_dtype: - logger.warning_once( - f"The input hidden states seems to be silently casted in {input_dtype}. " - f"This might be because you have upcasted embedding or layer norm layers " - f"in {input_dtype}. We will cast back the input in {target_dtype}." - ) - query_states = query_states.to(target_dtype) - key_states = key_states.to(target_dtype) - val_states = val_states.to(target_dtype) - - return query_states, key_states, val_states - - def forward( - self, - hidden_states: torch.Tensor, - within_seq_position_ids: torch.LongTensor, - global_position_ids: torch.LongTensor, - sequence_ids: torch.LongTensor, - attention_args: Optional[AttentionArgs] = None, - past_key_value: Optional[DynamicCache] = None, - output_attentions: bool = False, - output_s_max: bool = False, - use_cache: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[DynamicCache], Optional[List[torch.Tensor]]]: - is_cache_prefilled = ( - use_cache and past_key_value is not None and past_key_value.get_seq_length(self.layer_idx) > 0 - ) - - query_states, key_states, val_states = self.prepare_qkv( - hidden_states=hidden_states, - position_ids=within_seq_position_ids - if self.layer_type == AttentionLayerType.WITHIN_SEQ - else global_position_ids, - past_key_value=past_key_value, - use_cache=use_cache, - ) - - attn_output, attn_weights, s_max = self._attn( - query_states=query_states, - key_states=key_states, - val_states=val_states, - sequence_ids=sequence_ids, - attention_args=attention_args, - output_attentions=output_attentions, - output_s_max=output_s_max, - is_cache_prefilled=is_cache_prefilled, - ) - - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights, past_key_value, s_max - - def _attn( - self, - query_states: torch.Tensor, - key_states: torch.Tensor, - val_states: torch.Tensor, - sequence_ids: torch.Tensor, - attention_args: Optional[AttentionArgs] = None, - output_attentions: bool = False, - output_s_max: bool = False, - is_cache_prefilled: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - effective_layer_type = self.layer_type - if is_cache_prefilled and self.layer_type == AttentionLayerType.GLOBAL: - effective_layer_type = AttentionLayerType.WITHIN_SEQ - - if output_attentions: - return self._manual_attn( - query_states, key_states, val_states, - sequence_ids=sequence_ids, - attention_args=attention_args, - effective_layer_type=effective_layer_type, - output_s_max=output_s_max, - is_cache_prefilled=is_cache_prefilled, - ) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - attn_output, attn_weights = self._kernels_flash_attn( - query_states, key_states, val_states, - sequence_ids=sequence_ids, - is_cache_prefilled=is_cache_prefilled, - ) - else: - attn_output, attn_weights = self._flex_attn( - query_states, key_states, val_states, - attention_args=attention_args, - effective_layer_type=effective_layer_type, - ) - elif self.attn_backend == AttentionBackend.FLEX: - attn_output, attn_weights = self._flex_attn( - query_states, key_states, val_states, - attention_args=attention_args, - effective_layer_type=effective_layer_type, - ) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output, attn_weights = self._sdpa_attn( - query_states, key_states, val_states, - sequence_ids=sequence_ids, - attention_args=attention_args, - effective_layer_type=effective_layer_type, - is_cache_prefilled=is_cache_prefilled, - ) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - s_max = self._compute_s_max(query_states, key_states) if output_s_max else None - return attn_output, attn_weights, s_max - - @torch.no_grad() - def _compute_s_max( - self, - query_states: torch.Tensor, # (B, L, H, D) - key_states: torch.Tensor, # (B, L, Hkv, D) - ) -> List[torch.Tensor]: - query_BHLD = query_states.transpose(1, 2).contiguous() - key_BHLD = key_states.transpose(1, 2).contiguous() - key_BHLD = repeat_kv(key_BHLD, self.num_key_value_groups) - scale = 1.0 / (self.head_dim ** 0.5) - q_norm = torch.linalg.vector_norm(query_BHLD, dim=-1) - k_norm = torch.linalg.vector_norm(key_BHLD, dim=-1) - s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values * scale - return [s_max_bound[h] for h in range(self.num_heads)] - - def _kernels_flash_attn( - self, - query_states: torch.Tensor, - key_states: torch.Tensor, - val_states: torch.Tensor, - sequence_ids: torch.Tensor, - is_cache_prefilled: bool = False, - ) -> Tuple[torch.Tensor, None]: - bsz, q_len = query_states.shape[0], query_states.shape[1] - _, kv_len = key_states.shape[0], key_states.shape[1] - - if self.layer_type == AttentionLayerType.GLOBAL and not is_cache_prefilled: - q_sequence_ids = sequence_ids - if q_len < kv_len: - first_token_id = sequence_ids[:, 0].unsqueeze(1) - k_sequence_ids = torch.cat([first_token_id.expand(bsz, kv_len - q_len), sequence_ids], dim=-1) - else: - k_sequence_ids = sequence_ids - else: - if q_len < kv_len: - key_states = key_states[:, -q_len:] - val_states = val_states[:, -q_len:] - q_sequence_ids = k_sequence_ids = sequence_ids - - attn_output = kernels_flash_attention_func( - query_states, key_states, val_states, - q_sequence_ids=q_sequence_ids, - k_sequence_ids=k_sequence_ids, - causal=False, - ) - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() - return attn_output, None - - def _flex_attn( - self, - query_states: torch.Tensor, - key_states: torch.Tensor, - val_states: torch.Tensor, - attention_args: Optional[AttentionArgs] = None, - effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, - ) -> Tuple[torch.Tensor, None]: - bsz, q_len = query_states.shape[0], query_states.shape[1] - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - block_mask = attention_args["within_seq_block_mask"] if attention_args is not None else None - else: - block_mask = attention_args["block_causal_block_mask"] if attention_args is not None else None - outputs = flex_attention_func(query_states, key_states, val_states, block_mask=block_mask) - outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous() - return outputs, None - - def _sdpa_attn( - self, - query_states: torch.Tensor, # (B, L, H, D) - key_states: torch.Tensor, # (B, L, Hkv, D) - val_states: torch.Tensor, # (B, L, Hkv, D) - sequence_ids: torch.Tensor, - attention_args: Optional[AttentionArgs] = None, - effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, - is_cache_prefilled: bool = False, - ) -> Tuple[torch.Tensor, None]: - bsz, q_len = query_states.shape[:2] - kv_len = key_states.shape[1] - - if is_cache_prefilled and q_len < kv_len: - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - key_states = key_states[:, -q_len:] - val_states = val_states[:, -q_len:] - attention_mask_4d = build_within_seq_mask_4d(sequence_ids) if effective_layer_type == AttentionLayerType.WITHIN_SEQ else None - elif attention_args is not None: - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - attention_mask_4d = attention_args["within_seq_mask_4d"] - else: - attention_mask_4d = attention_args["block_causal_mask_4d"] - else: - attention_mask_4d = None - - query_BHLD = query_states.transpose(1, 2).contiguous() - key_BHLD = key_states.transpose(1, 2).contiguous() - val_BHLD = val_states.transpose(1, 2).contiguous() - key_BHLD = repeat_kv(key_BHLD, self.num_key_value_groups) - val_BHLD = repeat_kv(val_BHLD, self.num_key_value_groups) - context_BHLD = F.scaled_dot_product_attention(query_BHLD, key_BHLD, val_BHLD, attn_mask=attention_mask_4d) - attn_output = context_BHLD.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() - return attn_output, None - - def _manual_attn( - self, - query_states: torch.Tensor, # (B, L, H, D) - key_states: torch.Tensor, # (B, L, Hkv, D) - val_states: torch.Tensor, # (B, L, Hkv, D) - sequence_ids: torch.Tensor, - attention_args: Optional[AttentionArgs] = None, - effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, - output_s_max: bool = False, - is_cache_prefilled: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[torch.Tensor]]]: - bsz, q_len = query_states.shape[:2] - kv_len = key_states.shape[1] - - if is_cache_prefilled and q_len < kv_len: - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - key_states = key_states[:, -q_len:] - val_states = val_states[:, -q_len:] - attention_mask_4d = build_within_seq_mask_4d(sequence_ids) if effective_layer_type == AttentionLayerType.WITHIN_SEQ else None - elif attention_args is not None: - if effective_layer_type == AttentionLayerType.WITHIN_SEQ: - attention_mask_4d = attention_args["within_seq_mask_4d"] - else: - attention_mask_4d = attention_args["block_causal_mask_4d"] - else: - attention_mask_4d = None - - query_BHLD = query_states.transpose(1, 2).contiguous() - key_BHLD = key_states.transpose(1, 2).contiguous() - val_BHLD = val_states.transpose(1, 2).contiguous() - key_BHLD = repeat_kv(key_BHLD, self.num_key_value_groups) - val_BHLD = repeat_kv(val_BHLD, self.num_key_value_groups) - scale = 1.0 / (self.head_dim ** 0.5) - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-2, -1)) * scale - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - context_BHLD = torch.matmul(attn_weights, val_BHLD) - attn_output = context_BHLD.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() - s_max = self._compute_s_max(query_states, key_states) if output_s_max else None - return attn_output, attn_weights, s_max - - -class MLP(nn.Module): - def __init__(self, config: E1Config): - super().__init__() - self.ffn_dim = config.intermediate_size - self.hidden_dim = config.hidden_size - self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) - self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.w2(self.act_fn(self.w1(hidden_states))) - - -class GLUMLP(nn.Module): - def __init__(self, config: E1Config): - super().__init__() - self.ffn_dim = config.intermediate_size - self.hidden_dim = config.hidden_size - self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) - self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) - self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states) - hidden_states = self.w2(hidden_states) - return hidden_states - - -class FFN(nn.Module): - def __init__(self, config: E1Config): - super().__init__() - mlp_cls = GLUMLP if config.gated_mlp else MLP - self.mlp = mlp_cls(config) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.mlp(hidden_states) - - -@dataclass -class E1ModelOutputWithPast(ModelOutput): - """Base class for model's outputs, with potential hidden states and attentions. - - Attributes: - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): - Sequence of hidden-states at the output of the last layer of the model. - past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape - `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if - `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, - encoder_sequence_length, embed_size_per_head)`. - - Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if - `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` - input) to speed up sequential decoding. - hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): - Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + - one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. - - Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. - attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): - Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, - sequence_length)`. - - Attentions weights after the attention softmax, used to compute the weighted average in the self-attention - heads. - """ - - last_hidden_state: Optional[torch.FloatTensor] = None - past_key_values: Optional[DynamicCache] = None - hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None - attentions: Optional[Tuple[torch.FloatTensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class E1MaskedLMOutputWithPast(ModelOutput): - loss: Optional[torch.FloatTensor] = None - mlm_loss: Optional[torch.FloatTensor] = None - logits: Optional[torch.FloatTensor] = None - last_hidden_state: Optional[torch.FloatTensor] = None - past_key_values: Optional[DynamicCache] = None - hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None - attentions: Optional[Tuple[torch.FloatTensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class E1ClassificationOutputWithPast(ModelOutput): - loss: Optional[torch.FloatTensor] = None - logits: Optional[torch.FloatTensor] = None - last_hidden_state: Optional[torch.FloatTensor] = None - past_key_values: Optional[DynamicCache] = None - hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None - attentions: Optional[Tuple[torch.FloatTensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -class RMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - self.hidden_size = hidden_size - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - if layer_norm is None: - return torch.nn.functional.rms_norm( - hidden_states, (self.hidden_size,), self.weight, self.variance_epsilon - ).to(input_dtype) - else: - return layer_norm.rms_norm_fn( - x=hidden_states, - weight=self.weight, - bias=None, # no bias - residual=None, - eps=self.variance_epsilon, - dropout_p=0.0, # no dropout by default - prenorm=False, - residual_in_fp32=False, - ).to(input_dtype) - - -class NormAttentionNorm(nn.Module): - def __init__(self, config: E1Config, layer_idx: int): - super().__init__() - self.self_attn = Attention(config, layer_idx) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - within_seq_position_ids: torch.LongTensor, - global_position_ids: torch.LongTensor, - sequence_ids: torch.LongTensor, - attention_args: Optional[AttentionArgs] = None, - past_key_value: Optional[DynamicCache] = None, - output_attentions: bool = False, - output_s_max: bool = False, - use_cache: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[DynamicCache], Optional[List[torch.Tensor]]]: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states, self_attn_weights, present_key_value, s_max = self.self_attn( - hidden_states=hidden_states, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - attention_args=attention_args, - past_key_value=past_key_value, - output_attentions=output_attentions, - output_s_max=output_s_max, - use_cache=use_cache, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - return hidden_states, residual, self_attn_weights, present_key_value, s_max - - -class DecoderLayer(nn.Module): - def __init__(self, config: E1Config, layer_idx: int): - super().__init__() - self.initializer_range = config.initializer_range - self.hidden_size = config.hidden_size - self.norm_attn_norm = NormAttentionNorm(config, layer_idx) - self.ffn = FFN(config) - - def forward( - self, - hidden_states: torch.Tensor, - within_seq_position_ids: torch.LongTensor, - global_position_ids: torch.LongTensor, - sequence_ids: torch.LongTensor, - attention_args: Optional[AttentionArgs] = None, - past_key_value: Optional[DynamicCache] = None, - output_attentions: bool = False, - output_s_max: bool = False, - use_cache: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[DynamicCache], Optional[List[torch.Tensor]]]: - hidden_states, residual, self_attn_weights, present_key_value, s_max = self.norm_attn_norm( - hidden_states=hidden_states, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - attention_args=attention_args, - past_key_value=past_key_value, - output_attentions=output_attentions, - output_s_max=output_s_max, - use_cache=use_cache, - ) - - # Fully Connected - hidden_states = self.ffn(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states, self_attn_weights, present_key_value, s_max - - -class E1PreTrainedModel(PreTrainedModel): - config_class = E1Config - config: E1Config - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["DecoderLayer"] - _transformer_layer_cls = [DecoderLayer] - _skip_keys_device_placement = "past_key_values" - all_tied_weights_keys = {} - _tokenizer_source: Optional[Union[str, os.PathLike]] = None - _tokenizer_local_files_only = False - _tokenizer_cache_dir: Optional[Union[str, os.PathLike]] = None - _tokenizer_revision: Optional[str] = None - _tokenizer_token: Optional[Union[str, bool]] = None - - @classmethod - def from_pretrained( # type: ignore[override] - cls, - pretrained_model_name_or_path: Union[str, os.PathLike], - *model_args: Any, - **kwargs: Any, - ) -> "E1PreTrainedModel": - previous_source = E1PreTrainedModel._tokenizer_source - previous_local_files_only = E1PreTrainedModel._tokenizer_local_files_only - previous_cache_dir = E1PreTrainedModel._tokenizer_cache_dir - previous_revision = E1PreTrainedModel._tokenizer_revision - previous_token = E1PreTrainedModel._tokenizer_token - E1PreTrainedModel._tokenizer_source = pretrained_model_name_or_path - E1PreTrainedModel._tokenizer_local_files_only = ( - bool(kwargs["local_files_only"]) if "local_files_only" in kwargs else False - ) - E1PreTrainedModel._tokenizer_cache_dir = kwargs["cache_dir"] if "cache_dir" in kwargs else None - E1PreTrainedModel._tokenizer_revision = kwargs["revision"] if "revision" in kwargs else None - if "token" in kwargs: - E1PreTrainedModel._tokenizer_token = kwargs["token"] - elif "use_auth_token" in kwargs: - E1PreTrainedModel._tokenizer_token = kwargs["use_auth_token"] - else: - E1PreTrainedModel._tokenizer_token = None - try: - return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) - finally: - E1PreTrainedModel._tokenizer_source = previous_source - E1PreTrainedModel._tokenizer_local_files_only = previous_local_files_only - E1PreTrainedModel._tokenizer_cache_dir = previous_cache_dir - E1PreTrainedModel._tokenizer_revision = previous_revision - E1PreTrainedModel._tokenizer_token = previous_token - - @staticmethod - def _tokenizer_kwargs_from_config(config: E1Config) -> Dict[str, Any]: - tokenizer_source = E1PreTrainedModel._tokenizer_source - if tokenizer_source is None and isinstance(config._name_or_path, str) and len(config._name_or_path) > 0: - tokenizer_source = config._name_or_path - return { - "tokenizer_source": tokenizer_source, - "local_files_only": E1PreTrainedModel._tokenizer_local_files_only, - "cache_dir": E1PreTrainedModel._tokenizer_cache_dir, - "revision": E1PreTrainedModel._tokenizer_revision, - "token": E1PreTrainedModel._tokenizer_token, - } - - def _init_weights(self, module: nn.Module) -> None: - std = self.config.initializer_range - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - elif isinstance(module, RMSNorm): - module.weight.data.fill_(1.0) - - def _backward_compatibility_gradient_checkpointing(self) -> None: - if self.supports_gradient_checkpointing and getattr(self.config, "gradient_checkpointing", False): - self.gradient_checkpointing_enable(dict(use_reentrant=False)) - - def post_init(self) -> None: - super().post_init() - - @property - def _device(self) -> torch.device: - return next(self.parameters()).device - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, ( - f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - ) - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - for module in self.modules(): - if isinstance(module, FAST_E1_ENCODER): - module._attn_backend = resolved - elif isinstance(module, Attention): - module.attn_backend = resolved - - -class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin): - config: E1Config - config_class = E1Config - def __init__(self, config: E1Config, **kwargs): - E1PreTrainedModel.__init__(self, config, **kwargs) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.embed_seq_id = nn.Embedding(config.max_num_sequences, config.hidden_size) - self.layers = nn.ModuleList([DecoderLayer(config, i) for i in range(config.num_hidden_layers)]) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.gradient_checkpointing = config.gradient_checkpointing - self.prep_tokens = E1BatchPreparer( - **E1PreTrainedModel._tokenizer_kwargs_from_config(config) - ) - self._attn_backend = resolve_attention_backend(config.attn_backend) - self.post_init() - - def get_input_embeddings(self) -> nn.Embedding: - return self.embed_tokens - - def set_input_embeddings(self, value: nn.Embedding) -> None: - self.embed_tokens = value - - def _embed( - self, - sequences: List[str], - return_attention_mask: bool = False, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - **kwargs, - ) -> torch.Tensor: - batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - output = self.forward( - **batch, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - embeddings = select_hidden_state_embeddings( - output.last_hidden_state, - output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - if return_attention_mask: - attention_mask = (batch['sequence_ids'] != -1).long() - return embeddings, attention_mask - else: - return embeddings - - # Ignore copy - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - within_seq_position_ids: Optional[torch.LongTensor] = None, - global_position_ids: Optional[torch.LongTensor] = None, - sequence_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - past_key_values: Optional[DynamicCache] = None, - use_cache: bool = False, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - **kwargs - ) -> E1ModelOutputWithPast: - """ - Args: - input_ids: (batch_size, seq_length) - within_seq_position_ids: (batch_size, seq_length) - This tensor contains the position of each residue within the sequence itself. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0,1,2,3,4,5,6,0,1,2,3,4,5,6], [0,1,2,3,4,5,0,1,2,3,4,5,6,-1]] - global_position_ids: (batch_size, seq_length) - This tensor contains the position of each residue within the global sequence. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1]] - sequence_ids: (batch_size, seq_length) - This tensor contains the sequence id of each residue. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0,0,0,0,0,0,0,1,1,1,1,1,1,1], [0,0,0,0,0,0,1,1,1,1,1,1,1,-1]] - inputs_embeds: (batch_size, seq_length, hidden_size) - pre-computed embeddings, - bypasses embed_tokens and embed_seq_id when provided. Used by PDE for - differentiable soft sequence optimization. - past_key_values: DynamicCache - use_cache: bool - output_attentions: bool - output_hidden_states: bool - output_s_max: bool - - Returns: - E1ModelOutputWithPast: Model Outputs - """ - assert not (input_ids is not None and inputs_embeds is not None), ( - "Cannot specify both input_ids and inputs_embeds" - ) - assert input_ids is not None or inputs_embeds is not None, ( - "Must specify either input_ids or inputs_embeds" - ) - - if input_ids is not None: - batch_size, seq_length = input_ids.shape - else: - batch_size, seq_length = inputs_embeds.shape[:2] - - if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - if use_cache and past_key_values is None: - past_key_values = DynamicCache() - elif not use_cache: - past_key_values = None - - # Synthesize positional IDs for soft embedding path (single-sequence) - if inputs_embeds is not None: - device = inputs_embeds.device - if within_seq_position_ids is None: - within_seq_position_ids = torch.arange(seq_length, device=device).unsqueeze(0).expand(batch_size, -1) - if global_position_ids is None: - global_position_ids = torch.arange(seq_length, device=device).unsqueeze(0).expand(batch_size, -1) - if sequence_ids is None: - sequence_ids = torch.zeros(batch_size, seq_length, device=device, dtype=torch.long) - - global_position_ids = global_position_ids.view(-1, seq_length).long() - within_seq_position_ids = within_seq_position_ids.view(-1, seq_length).long() - sequence_ids = sequence_ids.view(-1, seq_length).long() - - max_position_id = torch.max(within_seq_position_ids).item() - min_position_id = torch.min(within_seq_position_ids).item() - assert max_position_id < self.config.max_num_positions_within_seq and min_position_id >= -1, ( - f"Position ids must be in the range [-1, {self.config.max_num_positions_within_seq}); got max {max_position_id} and min {min_position_id}" - ) - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - inputs_embeds = inputs_embeds + self.embed_seq_id(sequence_ids.clamp(min=0)) - - if torch.is_autocast_enabled(): - target_dtype = torch.get_autocast_gpu_dtype() - else: - target_dtype = self.layers[0].norm_attn_norm.self_attn.q_proj.weight.dtype - hidden_states = inputs_embeds.to(target_dtype) - - past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 - - attn_backend = self._attn_backend - has_global_layers = self.config.global_attention_every_n_layers > 0 - needs_4d_masks = (attn_backend == AttentionBackend.SDPA) or output_attentions - needs_block_causal_flex = ( - (attn_backend == AttentionBackend.FLEX and has_global_layers) - or (attn_backend == AttentionBackend.KERNELS_FLASH and has_global_layers) - ) - needs_within_seq_flex = (attn_backend == AttentionBackend.FLEX) - - attention_args: Optional[AttentionArgs] = None - if past_key_values_length == 0: - attention_args = AttentionArgs( - block_causal_block_mask=create_block_causal_mask_optimized(sequence_ids) if needs_block_causal_flex else None, - within_seq_block_mask=create_within_seq_block_mask(sequence_ids) if needs_within_seq_flex else None, - within_seq_mask_4d=build_within_seq_mask_4d(sequence_ids) if needs_4d_masks else None, - block_causal_mask_4d=build_block_causal_mask_4d(sequence_ids) if needs_4d_masks else None, - ) - - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - full_s_max = () if output_s_max else None - next_decoder_cache = None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) # type: ignore[operator] - - if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - within_seq_position_ids, - global_position_ids, - sequence_ids, - attention_args, - past_key_values, - output_attentions, - output_s_max, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - attention_args=attention_args, - past_key_value=past_key_values, - output_attentions=output_attentions, - output_s_max=output_s_max, - use_cache=use_cache, - ) - - hidden_states, self_attn_weights, present_key_value, s_max = layer_outputs - - if use_cache: - next_decoder_cache = past_key_values = present_key_value - - if output_attentions: - all_self_attns += (self_attn_weights,) # type: ignore[operator] - - if full_s_max is not None: - full_s_max += (s_max,) # type: ignore[operator] - - hidden_states = self.norm(hidden_states) - - if output_hidden_states: - all_hidden_states += (hidden_states,) # type: ignore[operator] - - next_cache = next_decoder_cache if use_cache else None - - return E1ModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - s_max=full_s_max, - ) - - -class E1Model(E1PreTrainedModel, EmbeddingMixin): - config: E1Config - config_class = E1Config - - def __init__(self, config: E1Config, **kwargs): - E1PreTrainedModel.__init__(self, config, **kwargs) - self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) - self.prep_tokens = self.model.prep_tokens - self.post_init() - - def get_input_embeddings(self) -> nn.Embedding: - return self.model.get_input_embeddings() - - def set_input_embeddings(self, value: nn.Embedding) -> None: - self.model.set_input_embeddings(value) - - def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor: - return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - within_seq_position_ids: Optional[torch.LongTensor] = None, - global_position_ids: Optional[torch.LongTensor] = None, - sequence_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - past_key_values: Optional[DynamicCache] = None, - use_cache: bool = False, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - **kwargs, - ) -> E1ModelOutputWithPast: - return self.model( - input_ids=input_ids, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - inputs_embeds=inputs_embeds, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - **kwargs, - ) - - -class E1ForMaskedLM(FastPLMTestTimeTrainingMixin, E1PreTrainedModel, EmbeddingMixin): - config: E1Config - config_class = E1Config - def __init__(self, config: E1Config, **kwargs): - E1PreTrainedModel.__init__(self, config, **kwargs) - self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) - self.vocab_size = config.vocab_size - self.mlm_head = torch.nn.Sequential( - nn.Linear(config.hidden_size, config.hidden_size, bias=True), - nn.GELU(), - nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps), - nn.Linear(config.hidden_size, config.vocab_size, bias=True), - ) - self.gradient_checkpointing = config.gradient_checkpointing - self.prep_tokens = self.model.prep_tokens - self.post_init() - self.init_ttt({"lora_target_replace_module": "Attention"}) - - @property - def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: - return self.model.device_mesh - - def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor: - return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.model] - - def _ttt_tokenize( - self, - seq: str | list[str] | None = None, - input_ids: torch.Tensor | None = None, - **kwargs, - ) -> dict[str, torch.Tensor]: - if input_ids is not None: - return { - "input_ids": input_ids, - "within_seq_position_ids": kwargs["within_seq_position_ids"], - "global_position_ids": kwargs["global_position_ids"], - "sequence_ids": kwargs["sequence_ids"], - } - assert seq is not None, "Pass either seq or E1 token tensors for TTT." - sequences = [seq] if isinstance(seq, str) else seq - batch = self.prep_tokens.get_batch_kwargs(sequences, device=torch.device("cpu")) - return { - "input_ids": batch["input_ids"], - "within_seq_position_ids": batch["within_seq_position_ids"], - "global_position_ids": batch["global_position_ids"], - "sequence_ids": batch["sequence_ids"], - } - - def _ttt_mask_token(self) -> int: - return int(self.prep_tokens.mask_token_id) - - def _ttt_padding_token(self) -> int: - return int(self.prep_tokens.pad_token_id) - - def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: - amino_acids = "ACDEFGHIKLMNPQRSTVWY" - ids = [self.prep_tokens.vocab[aa] for aa in amino_acids] - return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) - - def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: - return ~self.prep_tokens.get_boundary_token_mask(input_ids) - - def _ttt_predict_logits( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - **kwargs, - ) -> torch.Tensor: - del kwargs - assert isinstance(batch, dict), "E1 TTT expects a tensor dictionary." - output = self( - input_ids=batch["input_ids"], - within_seq_position_ids=batch["within_seq_position_ids"], - global_position_ids=batch["global_position_ids"], - sequence_ids=batch["sequence_ids"], - ) - return output.logits - - def search_homologues( - self, - sequence: str, - output_dir: str, - provider: str = "colabfold", - target_db: Optional[str] = None, - seq_id: Optional[str] = None, - **kwargs, - ) -> str: - searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) - return searcher.search(sequence=sequence, output_dir=output_dir, seq_id=seq_id) - - def batch_search_homologues( - self, - sequences: List[str], - output_dir: str, - provider: str = "colabfold", - target_db: Optional[str] = None, - seq_ids: Optional[List[str]] = None, - continue_on_error: bool = True, - **kwargs, - ) -> Dict[str, str]: - searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) - return searcher.batch_search( - sequences=sequences, - output_dir=output_dir, - seq_ids=seq_ids, - continue_on_error=continue_on_error, - ) - - def sample_msa_contexts( - self, - a3m_path: str, - seed: int = 42, - max_context_tokens: Optional[List[int]] = None, - similarity_thresholds: Optional[List[float]] = None, - min_query_similarity: float = 0.3, - context_cache_dir: Optional[str] = None, - ) -> Dict[str, str]: - context_specs = build_context_specifications( - max_context_tokens=max_context_tokens, - similarity_thresholds=similarity_thresholds, - min_query_similarity=min_query_similarity, - ) - cache = None - if context_cache_dir is not None: - key = repr((max_context_tokens, similarity_thresholds, min_query_similarity)) - specs_hash = hashlib.md5(key.encode()).hexdigest()[:8] - cache = ContextCache(context_cache_dir, specs_hash, seed) - cached = cache.load(a3m_path) - if cached is not None: - return cached - contexts = sample_contexts_for_msa(a3m_path, context_specs, seed=seed) - if cache is not None: - cache.store(a3m_path, contexts) - return contexts - - @torch.inference_mode() - def score_ppll( - self, - sequences: List[str], - a3m_path: str, - ensemble: bool = True, - seed: int = 42, - max_context_tokens: Optional[List[int]] = None, - similarity_thresholds: Optional[List[float]] = None, - min_query_similarity: float = 0.3, - max_batch_tokens: int = 131072, - cache_size: int = 1, - context_cache_dir: Optional[str] = None, - progress: bool = True, - ) -> List[float] | List[List[float]]: - """Score sequences with FastPLMs PPLL reduction over sampled E1 MSA contexts. - - This intentionally differs from Profluent's official E1Scorer, which scores - mutants against a parent sequence with wildtype or masked marginal log-prob - deltas. Here each sequence is scored by mean correct-token probability and - optionally averaged across sampled contexts. - """ - contexts = self.sample_msa_contexts( - a3m_path=a3m_path, - seed=seed, - max_context_tokens=max_context_tokens, - similarity_thresholds=similarity_thresholds, - min_query_similarity=min_query_similarity, - context_cache_dir=context_cache_dir, - ) - assert len(contexts) > 0, "At least one sampled MSA context is required for PPLL scoring" - - predictor = _E1ContextPredictor( - model=self, - data_prep_config=DataPrepConfig(remove_X_tokens=True), - max_batch_tokens=max_batch_tokens, - fields_to_save=["logits"], - save_masked_positions_only=False, - keep_predictions_in_gpu=False, - use_cache=True, - cache_size=cache_size, - progress=progress, - ) - vocab = predictor.batch_preparer.vocab - seq_token_ids = [ - torch.tensor([vocab[aa] for aa in seq if aa != "X"], device=self.device) - for seq in sequences - ] - context_ids = list(contexts.keys()) - all_scores = torch.zeros(len(sequences), len(context_ids), device=self.device) - - iterator = tqdm(context_ids, desc="Scoring with contexts", disable=not progress) - for ctx_idx, ctx_id in enumerate(iterator): - predictions = list( - predictor.predict( - sequences=sequences, - sequence_ids=list(range(len(sequences))), - context_seqs={ctx_id: contexts[ctx_id]}, - ) - ) - for prediction in predictions: - seq_idx = prediction["id"] - assert isinstance(seq_idx, int), "Expected integer sequence ids for score aggregation" - all_scores[seq_idx, ctx_idx] = compute_ppll(prediction["logits"], seq_token_ids[seq_idx]) - if predictor.kv_cache is not None: - predictor.kv_cache.reset() - - if ensemble: - return all_scores.mean(dim=1).tolist() - return all_scores.tolist() - - @torch.inference_mode() - def embed_with_msa( - self, - sequences: List[str], - a3m_path: Optional[str] = None, - context: Optional[str] = None, - pooling_types: Optional[List[str]] = None, - pooling: str = "mean", - matrix_embed: bool = False, - seed: int = 42, - max_batch_tokens: int = 131072, - embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, - embed_similarity: float = DEFAULT_EMBED_SIMILARITY, - min_query_similarity: float = 0.3, - progress: bool = True, - ) -> torch.Tensor | List[torch.Tensor]: - if a3m_path is not None and context is None: - spec = ContextSpecification( - max_num_samples=511, - max_token_length=embed_max_tokens, - max_query_similarity=embed_similarity, - min_query_similarity=min_query_similarity, - ) - contexts, _ = sample_multiple_contexts( - msa_path=a3m_path, - context_specifications=[spec], - seed=seed, - ) - context = contexts[0] if contexts else None - - hidden_list = _forward_for_embedding( - model=self, - sequences=sequences, - context=context, - max_batch_tokens=max_batch_tokens, - progress=progress, - ) - if matrix_embed: - return hidden_list - if pooling_types is not None: - return _pool_hidden_states(hidden_list, pooling_types, self.device) - if pooling not in ("mean", "cls"): - raise ValueError("pooling must be 'mean' or 'cls' when pooling_types is not provided") - embeddings = [hidden.mean(dim=0) if pooling == "mean" else hidden[0] for hidden in hidden_list] - return torch.stack(embeddings) - - @torch.inference_mode() - def embed_dataset_with_msa( - self, - sequences: List[str], - msa_lookup: Optional[Dict[str, str]] = None, - msa_dir: Optional[str] = None, - msa_hf_path: Optional[str] = None, - batch_size: int = 2, - max_len: int = 2048, - pooling_types: Optional[List[str]] = None, - pooling: str = "mean", - matrix_embed: bool = False, - embed_dtype: torch.dtype = torch.bfloat16, - embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, - embed_similarity: float = DEFAULT_EMBED_SIMILARITY, - min_query_similarity: float = 0.3, - seed: int = 42, - progress: bool = True, - ) -> Dict[str, torch.Tensor]: - if msa_lookup is None: - if msa_dir is not None: - msa_lookup = load_msa_dir(msa_dir) - elif msa_hf_path is not None: - msa_lookup = load_msa_from_hf(msa_hf_path) - else: - msa_lookup = {} - - truncated_map = {seq: seq[:max_len] for seq in sequences} - unique_seqs = sorted(set(truncated_map.values()), key=len, reverse=True) - context_map: Dict[str, Optional[str]] = {} - spec = ContextSpecification( - max_num_samples=511, - max_token_length=embed_max_tokens, - max_query_similarity=embed_similarity, - min_query_similarity=min_query_similarity, - ) - for seq in unique_seqs: - a3m_path = get_msa_for_sequence(seq, msa_lookup) - if a3m_path is None: - context_map[seq] = None - continue - contexts, _ = sample_multiple_contexts( - msa_path=a3m_path, - context_specifications=[spec], - seed=seed, - ) - context_map[seq] = contexts[0] if contexts else None - - context_groups: Dict[Optional[str], List[str]] = defaultdict(list) - for seq in unique_seqs: - context_groups[context_map[seq]].append(seq) - - embeddings_dict: Dict[str, torch.Tensor] = {} - total_batches = sum((len(seqs) + batch_size - 1) // batch_size for seqs in context_groups.values()) - pbar = tqdm(total=total_batches, desc="Embedding with MSA", disable=not progress) - for ctx, ctx_seqs in context_groups.items(): - for i in range(0, len(ctx_seqs), batch_size): - batch_seqs = ctx_seqs[i:i + batch_size] - batch_embeddings = self.embed_with_msa( - sequences=batch_seqs, - context=ctx, - pooling_types=pooling_types, - pooling=pooling, - matrix_embed=matrix_embed, - seed=seed, - embed_max_tokens=embed_max_tokens, - embed_similarity=embed_similarity, - min_query_similarity=min_query_similarity, - progress=False, - ) - if matrix_embed: - assert isinstance(batch_embeddings, list) - for seq, hidden in zip(batch_seqs, batch_embeddings): - embeddings_dict[seq] = hidden.to(embed_dtype).cpu() - else: - assert isinstance(batch_embeddings, torch.Tensor) - for j, seq in enumerate(batch_seqs): - embeddings_dict[seq] = batch_embeddings[j].to(embed_dtype).cpu() - pbar.update(1) - pbar.close() - - result: Dict[str, torch.Tensor] = {} - for seq in sequences: - trunc = truncated_map[seq] - if trunc in embeddings_dict: - result[seq] = embeddings_dict[trunc] - return result - - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - within_seq_position_ids: Optional[torch.LongTensor] = None, - global_position_ids: Optional[torch.LongTensor] = None, - sequence_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - past_key_values: Optional[DynamicCache] = None, - use_cache: bool = False, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - **kwargs, - ) -> E1MaskedLMOutputWithPast: - """ - Args: - input_ids: (batch_size, seq_length) - within_seq_position_ids: (batch_size, seq_length) - This tensor contains the position of each residue within the sequence itself. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0,1,2,3,4,5,6,0,1,2,3,4,5,6], [0,1,2,3,4,5,0,1,2,3,4,5,6,-1]] - global_position_ids: (batch_size, seq_length) - This tensor contains the position of each residue within the global sequence. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1]] - sequence_ids: (batch_size, seq_length) - This tensor contains the sequence id of each residue. - For example, if the input is ["1ABC21DEF2", "1GH21JKL2"], - the tensor would be [[0,0,0,0,0,0,0,1,1,1,1,1,1,1], [0,0,0,0,0,0,1,1,1,1,1,1,1,-1]] - inputs_embeds: (batch_size, seq_length, hidden_size) - pre-computed embeddings - labels: (batch_size, seq_length) - past_key_values: DynamicCache - use_cache: bool - output_attentions: bool - output_hidden_states: bool - output_s_max: bool - - Returns: - E1MaskedLMOutputWithPast: Model Outputs - """ - outputs: E1ModelOutputWithPast = self.model( - input_ids=input_ids, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - inputs_embeds=inputs_embeds, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - - last_hidden_state = outputs.last_hidden_state - loss = None - - mlm_logits = self.mlm_head(last_hidden_state).float() - mlm_loss = 0.0 - if labels is not None: - mlm_logits_flat = mlm_logits.contiguous().view(-1, self.config.vocab_size) - mlm_labels_flat = labels.to(mlm_logits_flat.device).contiguous().view(-1) - mlm_loss = F.cross_entropy(mlm_logits_flat, mlm_labels_flat, reduction="none") - mask = mlm_labels_flat != self.model.padding_idx - n_mlm = mask.sum() - mlm_loss = (mlm_loss * mask.to(mlm_loss)).sum() / (1 if n_mlm == 0 else n_mlm) - loss = 0.0 - loss += mlm_loss - - return E1MaskedLMOutputWithPast( - loss=loss, - mlm_loss=mlm_loss, - logits=mlm_logits, - last_hidden_state=last_hidden_state, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin): - config: E1Config - config_class = E1Config - def __init__(self, config: E1Config, **kwargs): - E1PreTrainedModel.__init__(self, config, **kwargs) - self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) - self.vocab_size = config.vocab_size - self.num_labels = config.num_labels - self.classifier = nn.Sequential( - nn.Linear(config.hidden_size * 2, config.hidden_size * 4), - nn.GELU(), - nn.LayerNorm(config.hidden_size * 4), - nn.Linear(config.hidden_size * 4, config.num_labels), - ) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - self.gradient_checkpointing = config.gradient_checkpointing - self.prep_tokens = self.model.prep_tokens - - if 'pooling_types' in kwargs and isinstance(kwargs['pooling_types'], List[str]) and len(kwargs['pooling_types']) > 0: - pooling_types = kwargs['pooling_types'] - else: - pooling_types = ['mean', 'var'] - self.pooler = Pooler(pooling_types) - self.post_init() - - @property - def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: - return self.model.device_mesh - - def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor: - return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - within_seq_position_ids: Optional[torch.LongTensor] = None, - global_position_ids: Optional[torch.LongTensor] = None, - sequence_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - past_key_values: Optional[DynamicCache] = None, - use_cache: bool = False, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - **kwargs, - ) -> E1ClassificationOutputWithPast: - outputs: E1ModelOutputWithPast = self.model( - input_ids=input_ids, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - inputs_embeds=inputs_embeds, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - - attention_mask = (sequence_ids != -1).long() if sequence_ids is not None else torch.ones(outputs.last_hidden_state.shape[:2], device=outputs.last_hidden_state.device, dtype=torch.long) - x = outputs.last_hidden_state - features = self.pooler(x, attention_mask) - logits = self.classifier(features) - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - if self.num_labels == 1: - loss = self.mse(logits.flatten(), labels.flatten()) - else: - loss = self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return E1ClassificationOutputWithPast( - loss=loss, - logits=logits, - last_hidden_state=x, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin): - config: E1Config - config_class = E1Config - def __init__(self, config: E1Config, **kwargs): - E1PreTrainedModel.__init__(self, config, **kwargs) - self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) - self.vocab_size = config.vocab_size - self.num_labels = config.num_labels - self.classifier = nn.Sequential( - nn.Linear(config.hidden_size * 2, config.hidden_size * 4), - nn.GELU(), - nn.LayerNorm(config.hidden_size * 4), - nn.Linear(config.hidden_size * 4, config.num_labels), - ) - self.loss_fct = nn.CrossEntropyLoss() - self.gradient_checkpointing = config.gradient_checkpointing - self.prep_tokens = self.model.prep_tokens - self.post_init() - - @property - def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: - return self.model.device_mesh - - def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor: - return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - within_seq_position_ids: Optional[torch.LongTensor] = None, - global_position_ids: Optional[torch.LongTensor] = None, - sequence_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - past_key_values: Optional[DynamicCache] = None, - use_cache: bool = False, - output_attentions: bool = False, - output_hidden_states: bool = False, - output_s_max: bool = False, - **kwargs, - ) -> E1ClassificationOutputWithPast: - outputs: E1ModelOutputWithPast = self.model( - input_ids=input_ids, - within_seq_position_ids=within_seq_position_ids, - global_position_ids=global_position_ids, - sequence_ids=sequence_ids, - inputs_embeds=inputs_embeds, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - - x = outputs.last_hidden_state - logits = self.classifier(x) - loss = None - if labels is not None: - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return E1ClassificationOutputWithPast( - loss=loss, - logits=logits, - last_hidden_state=x, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -if __name__ == "__main__": - import random - - import torch - - from torch import Tensor - - def print_tensor_shapes(prefix: str, obj): - if isinstance(obj, Tensor): - print(f"{prefix}{obj.shape}") - elif isinstance(obj, dict): - for name, value in obj.items(): - print_tensor_shapes(f"{prefix}{name}.", value) - elif isinstance(obj, list): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif isinstance(obj, tuple): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif hasattr(obj, "__dict__"): - for name, value in vars(obj).items(): - if name.startswith("_"): - continue - print_tensor_shapes(f"{prefix}{name}.", value) - else: - print(f"{prefix}{type(obj)}") - - def get_e1_batch(tokenizer, sequences: List[str], device: torch.device): - preparer = E1BatchPreparer(data_prep_config=DataPrepConfig(max_num_positions_within_seq=64), tokenizer=tokenizer) - return preparer.get_batch_kwargs(sequences=sequences, device=device) - - random.seed(0) - torch.manual_seed(0) - - num_attention_heads = random.choice([2, 4]) - config = E1Config( - hidden_size=16 * num_attention_heads, - intermediate_size=64 * num_attention_heads, - num_hidden_layers=random.choice([1, 2]), - num_attention_heads=num_attention_heads, - num_key_value_heads=num_attention_heads, - max_num_positions_within_seq=128, - max_num_positions_global=256, - max_num_sequences=8, - dtype="float32", - ) - model = E1ForMaskedLM(config=config).eval() - tokenizer = get_tokenizer() - batch = get_e1_batch(tokenizer=tokenizer, sequences=["ACDEFG", "MKTW"], device=torch.device("cpu")) - batch["labels"] = batch["labels"].clone() - - with torch.no_grad(): - output = model( - input_ids=batch["input_ids"], - within_seq_position_ids=batch["within_seq_position_ids"], - global_position_ids=batch["global_position_ids"], - sequence_ids=batch["sequence_ids"], - labels=batch["labels"], - ) - - print("Batch shape:") - print_tensor_shapes("", batch) - print("Output shape:") - print_tensor_shapes("", output) diff --git a/fastplms/embedding_mixin.py b/fastplms/embedding_mixin.py deleted file mode 100644 index 5828827..0000000 --- a/fastplms/embedding_mixin.py +++ /dev/null @@ -1,750 +0,0 @@ -import io -import os -import queue -import sqlite3 -import struct -import threading -import time - -import networkx as nx -import numpy as np -import torch -from tqdm.auto import tqdm -from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple -from torch.utils.data import DataLoader -from torch.utils.data import Dataset as TorchDataset -from transformers import PreTrainedTokenizerBase - - -# SQLite stores tensors as compact blobs. Keep this header format compatible -# with Protify readers that share the same dtype/version codes. -_COMPACT_VERSION = 0x01 -_DTYPE_TO_CODE = {torch.float16: 0, torch.bfloat16: 1, torch.float32: 2} -_CODE_TO_DTYPE = {0: torch.float16, 1: torch.bfloat16, 2: torch.float32} -_CODE_TO_NP_DTYPE = {0: np.float16, 1: np.float16, 2: np.float32} - - -def tensor_to_embedding_blob(tensor: torch.Tensor) -> bytes: - """Serialize a tensor to compact binary format for SQLite blob storage. - - Format: [version:1][dtype_code:1][ndim:4][shape:4*ndim][raw_bytes] - bfloat16 tensors are stored as float16 bytes (numpy lacks bfloat16) - but tagged with dtype_code=1 so they can be cast back on read. - Falls back to torch.save for unsupported dtypes. - """ - t = tensor.cpu() - if t.dtype not in _DTYPE_TO_CODE: - buffer = io.BytesIO() - torch.save(t, buffer) - return buffer.getvalue() - dtype_code = _DTYPE_TO_CODE[t.dtype] - - if t.dtype == torch.bfloat16: - raw = t.half().numpy().tobytes() - else: - raw = t.numpy().tobytes() - - shape = t.shape - header = struct.pack(f' bytes: - """Build just the compact header for a given dtype and shape.""" - dtype_code = _DTYPE_TO_CODE[dtype] - return struct.pack(f' List[bytes]: - """Serialize a batch of same-shape tensors to compact blobs (fast path for vectors). - - Builds the header once and slices raw bytes per row. Much faster than - per-row tensor_to_embedding_blob calls for uniform-shape batches. - """ - assert batch.ndim >= 2, f"Expected batch with >= 2 dims, got {batch.ndim}" - t = batch.cpu() - store_dtype = t.dtype - if t.dtype not in _DTYPE_TO_CODE: - return [tensor_to_embedding_blob(t[i]) for i in range(t.shape[0])] - - if t.dtype == torch.bfloat16: - arr = t.half().numpy() - store_dtype = torch.bfloat16 - else: - arr = t.numpy() - - row_shape = tuple(t.shape[1:]) - header = _compact_header(store_dtype, row_shape) - raw = arr.tobytes() - stride = len(raw) // t.shape[0] - return [header + raw[i * stride:(i + 1) * stride] for i in range(t.shape[0])] - - -def embedding_blob_to_tensor(blob: bytes, fallback_shape: Optional[Tuple[int, ...]] = None) -> torch.Tensor: - """Deserialize a blob back to a tensor. Auto-detects compact vs legacy formats.""" - if len(blob) >= 6 and blob[0] == _COMPACT_VERSION: - dtype_code = blob[1] - ndim = struct.unpack_from(' torch.Tensor: - assert isinstance(hidden_state_index, int), "hidden_state_index must be an integer." - if store_all_hidden_states: - assert hidden_states is not None, "store_all_hidden_states requires output_hidden_states=True." - assert len(hidden_states) > 0, "Model returned no hidden states." - return torch.stack(tuple(hidden_states), dim=1) - if hidden_state_index == -1: - return last_hidden_state - assert hidden_states is not None, "hidden_state_index selection requires output_hidden_states=True." - return hidden_states[hidden_state_index] - - -def _trim_full_embedding(embedding: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - mask = attention_mask.bool() - if embedding.ndim == 2: - return embedding[mask].reshape(-1, embedding.shape[-1]) - if embedding.ndim == 3: - return embedding[:, mask, :].reshape(embedding.shape[0], -1, embedding.shape[-1]) - raise AssertionError(f"Expected full embedding tensor with 2 or 3 dims, got {embedding.ndim}.") - - -def pool_embeddings( - embeddings: Dict[str, torch.Tensor], - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, -) -> Dict[str, torch.Tensor]: - pooler = Pooler(pooling_types) - pooled: Dict[str, torch.Tensor] = {} - for sequence, embedding in embeddings.items(): - assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)." - assert isinstance(embedding, torch.Tensor), "Expected embedding dictionary values to be tensors." - if embedding.ndim == 1: - pooled[sequence] = embedding.cpu() - continue - if embedding.ndim == 3: - embedding = embedding[hidden_state_index] - assert embedding.ndim == 2, f"Expected token-wise embedding with 2 dims, got {embedding.ndim}." - pooled[sequence] = pooler(embedding.unsqueeze(0)).squeeze(0).cpu() - return pooled - - -def load_pooled_embeddings_from_pth( - save_path: str, - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, -) -> Dict[str, torch.Tensor]: - assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}" - payload = torch.load(save_path, map_location="cpu", weights_only=True) - assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary." - return pool_embeddings(payload, pooling_types=pooling_types, hidden_state_index=hidden_state_index) - - -def load_pooled_embeddings_from_db( - db_path: str, - sequences: Optional[List[str]] = None, - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, -) -> Dict[str, torch.Tensor]: - assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}" - loaded: Dict[str, torch.Tensor] = {} - with sqlite3.connect(db_path, timeout=30) as conn: - cursor = conn.cursor() - if sequences is None: - cursor.execute("SELECT sequence, embedding FROM embeddings") - else: - if len(sequences) == 0: - return loaded - placeholders = ",".join(["?"] * len(sequences)) - cursor.execute( - f"SELECT sequence, embedding FROM embeddings WHERE sequence IN ({placeholders})", - tuple(sequences), - ) - for sequence, embedding_bytes in cursor.fetchall(): - loaded[sequence] = embedding_blob_to_tensor(embedding_bytes) - return pool_embeddings(loaded, pooling_types=pooling_types, hidden_state_index=hidden_state_index) - - -def maybe_compile(model: torch.nn.Module, dynamic: bool = False) -> torch.nn.Module: - """Compile model with torch.compile if possible. - - Skips compilation when dynamic=True (padding='longest') because - flex attention's create_block_mask is incompatible with dynamic shapes - under torch.compile, causing CUDA illegal memory access. - """ - if dynamic: - print("Skipping torch.compile (dynamic shapes + flex attention incompatible)") - return model - try: - model = torch.compile(model) - print("Model compiled") - except Exception as e: - print(f"Skipping torch.compile: {e}") - return model - - -def build_collator( - tokenizer: PreTrainedTokenizerBase, - padding: str = 'max_length', - max_length: int = 512, -) -> Callable[[List[str]], Dict[str, torch.Tensor]]: - def _collate_fn(sequences: List[str]) -> Dict[str, torch.Tensor]: - kwargs: Dict[str, Any] = dict( - return_tensors="pt", padding=padding, truncation=True, max_length=max_length, - ) - if padding != 'max_length': - kwargs['pad_to_multiple_of'] = 8 - return tokenizer(sequences, **kwargs) - return _collate_fn - - -def _make_embedding_progress( - dataloader: DataLoader, - padding: str, - n_warmup: int = 3, - n_calibration: int = 5, -) -> Iterator[Tuple[int, Any]]: - """Progress-bar wrapper for embedding loops. Drop-in replacement for enumerate(dataloader). - - When padding='max_length', all batches have uniform cost so plain tqdm works. - When padding='longest' (sorted longest-first), batch times vary dramatically. - In that case: yield warmup batches first (compiler warmup + OOM check on longest - sequences), then time mid-length calibration batches to estimate total ETA. - - Keep in sync with protify/embedder.py and core/atlas/precomputed.py. - """ - total = len(dataloader) - if padding == 'max_length' or total <= n_warmup + n_calibration: - for i, batch in tqdm(enumerate(dataloader), total=total, desc='Embedding batches'): - yield i, batch - return - - dl_iter = iter(dataloader) - - # Warm up on the longest batches first; sorted inputs make these the OOM-risk - # and compile-stabilization cases. - warmup_bar = tqdm(range(n_warmup), desc='Warmup (longest batches)', leave=False) - for i in warmup_bar: - batch = next(dl_iter) - yield i, batch - warmup_bar.close() - - # Move toward mid-length batches for ETA calibration, yielding every real - # batch on the way so no sequences are skipped. - mid_start = total // 2 - intermediate_bar = tqdm( - range(n_warmup, mid_start), desc='Embedding batches', leave=False, - ) - for i in intermediate_bar: - batch = next(dl_iter) - yield i, batch - intermediate_bar.close() - - # Mid-length batches give a better remaining-time estimate than the longest - # warmup batches. - calibration_times: List[float] = [] - cal_bar = tqdm(range(n_calibration), desc='Calibrating ETA', leave=False) - for j in cal_bar: - t0 = time.perf_counter() - batch = next(dl_iter) - yield mid_start + j, batch - calibration_times.append(time.perf_counter() - t0) - cal_bar.close() - - avg_time = sum(calibration_times) / len(calibration_times) - remaining_start = mid_start + n_calibration - remaining_count = total - remaining_start - estimated_total_seconds = avg_time * remaining_count - - # Finish the tail with the calibrated ETA shown in the progress bar. - main_bar = tqdm( - range(remaining_count), - desc='Embedding batches', - bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]', - ) - main_bar.set_postfix_str(f'ETA ~{estimated_total_seconds:.0f}s (calibrated)') - for k in main_bar: - batch = next(dl_iter) - yield remaining_start + k, batch - main_bar.close() - - -class _SQLWriter: - """Context manager for async SQL embedding writes. Matches core/embed/storage.SQLEmbeddingWriter.""" - - def __init__(self, conn: sqlite3.Connection, queue_maxsize: int = 4) -> None: - self._conn = conn - self._queue: queue.Queue = queue.Queue(maxsize=queue_maxsize) - self._thread: Optional[threading.Thread] = None - - def __enter__(self) -> "_SQLWriter": - self._thread = threading.Thread(target=self._writer_loop, daemon=True) - self._thread.start() - return self - - def write_batch(self, rows: List[Tuple[str, bytes]]) -> None: - self._queue.put(rows) - - def _writer_loop(self) -> None: - cursor = self._conn.cursor() - while True: - item = self._queue.get() - if item is None: - break - cursor.executemany("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", item) - if self._queue.qsize() == 0: - self._conn.commit() - self._conn.commit() - - def __exit__(self, *exc) -> None: - if self._thread is not None: - self._queue.put(None) - self._thread.join() - self._thread = None - - -class Pooler: - def __init__(self, pooling_types: List[str]) -> None: - self.pooling_types = pooling_types - self.pooling_options: Dict[str, Callable] = { - 'mean': self.mean_pooling, - 'max': self.max_pooling, - 'norm': self.norm_pooling, - 'median': self.median_pooling, - 'std': self.std_pooling, - 'var': self.var_pooling, - 'cls': self.cls_pooling, - 'parti': self._pool_parti, - } - - def _create_pooled_matrices_across_layers(self, attentions: torch.Tensor) -> torch.Tensor: - assert isinstance(attentions, torch.Tensor) - maxed_attentions = torch.max(attentions, dim=1)[0] - return maxed_attentions - - def _page_rank(self, attention_matrix: np.ndarray, personalization: Optional[dict] = None, nstart: Optional[dict] = None, prune_type: str = "top_k_outdegree") -> Dict[int, float]: - G = self._convert_to_graph(attention_matrix) - if G.number_of_nodes() != attention_matrix.shape[0]: - raise Exception( - f"The number of nodes in the graph should be equal to the number of tokens in sequence! You have {G.number_of_nodes()} nodes for {attention_matrix.shape[0]} tokens.") - if G.number_of_edges() == 0: - raise Exception(f"You don't seem to have any attention edges left in the graph.") - - return nx.pagerank(G, alpha=0.85, tol=1e-06, weight='weight', personalization=personalization, nstart=nstart, max_iter=100) - - def _convert_to_graph(self, matrix: np.ndarray) -> nx.DiGraph: - G = nx.from_numpy_array(matrix, create_using=nx.DiGraph) - return G - - def _calculate_importance_weights(self, dict_importance: Dict[int, float], attention_mask: Optional[torch.Tensor] = None) -> np.ndarray: - if attention_mask is not None: - for k in list(dict_importance.keys()): - if attention_mask[k] == 0: - del dict_importance[k] - - total = sum(dict_importance.values()) - return np.array([v / total for _, v in dict_importance.items()]) - - def _pool_parti(self, emb: torch.Tensor, attentions: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: - maxed_attentions = self._create_pooled_matrices_across_layers(attentions).numpy() - emb_pooled = [] - for e, a, mask in zip(emb, maxed_attentions, attention_mask): - dict_importance = self._page_rank(a) - importance_weights = self._calculate_importance_weights(dict_importance, mask) - num_tokens = int(mask.sum().item()) - emb_pooled.append(np.average(e[:num_tokens], weights=importance_weights, axis=0)) - pooled = torch.tensor(np.array(emb_pooled)) - return pooled - - def mean_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.mean(dim=1) - else: - attention_mask = attention_mask.unsqueeze(-1) - return (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1) - - def max_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.max(dim=1).values - else: - mask = attention_mask.unsqueeze(-1).bool() - return emb.masked_fill(~mask, float('-inf')).max(dim=1).values - - def norm_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.norm(dim=1, p=2) - else: - attention_mask = attention_mask.unsqueeze(-1) - return (emb * attention_mask).norm(dim=1, p=2) - - def median_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.median(dim=1).values - else: - mask = attention_mask.unsqueeze(-1).bool() - return emb.masked_fill(~mask, float('nan')).nanmedian(dim=1).values - - def std_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.std(dim=1) - else: - var = self.var_pooling(emb, attention_mask, **kwargs) - return torch.sqrt(var) - - def var_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - if attention_mask is None: - return emb.var(dim=1) - else: - attention_mask = attention_mask.unsqueeze(-1) - mean = (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1) - mean = mean.unsqueeze(1) - squared_diff = (emb - mean) ** 2 - var = (squared_diff * attention_mask).sum(dim=1) / attention_mask.sum(dim=1) - return var - - def cls_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: - return emb[:, 0, :] - - def __call__( - self, - emb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - attentions: Optional[torch.Tensor] = None - ) -> torch.Tensor: - if attention_mask is not None: - assert attention_mask.sum(dim=-1).min() > 0, ( - "Pooler received samples with all-zero attention masks. " - "This causes NaN from division by zero. Filter empty inputs before pooling." - ) - final_emb: List[torch.Tensor] = [] - for pooling_type in self.pooling_types: - final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions)) - return torch.cat(final_emb, dim=-1) - - -class ProteinDataset(TorchDataset): - """Simple dataset for protein sequences.""" - def __init__(self, sequences: List[str]) -> None: - self.sequences = sequences - - def __len__(self) -> int: - return len(self.sequences) - - def __getitem__(self, idx: int) -> str: - return self.sequences[idx] - - -def parse_fasta(fasta_path: str) -> List[str]: - assert os.path.exists(fasta_path), f"FASTA file does not exist: {fasta_path}" - sequences = [] - current_seq = [] - with open(fasta_path, 'r') as f: - for line in f: - line = line.strip() - if not line: - continue - if line.startswith('>'): - if current_seq: - sequences.append(''.join(current_seq)) - current_seq = [] - else: - current_seq.append(line) - if current_seq: - sequences.append(''.join(current_seq)) - return sequences - - -class EmbeddingMixin: - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - raise NotImplementedError - - @property - def device(self) -> torch.device: - """Get the device of the model.""" - return next(self.parameters()).device - - def _read_sequences_from_db(self, db_path: str) -> Set[str]: - """Read sequences from SQLite database.""" - with sqlite3.connect(db_path, timeout=30) as conn: - c = conn.cursor() - c.execute("SELECT sequence FROM embeddings") - return {row[0] for row in c.fetchall()} - - def _ensure_embeddings_table(self, conn: sqlite3.Connection) -> None: - cursor = conn.cursor() - cursor.execute( - "CREATE TABLE IF NOT EXISTS embeddings (" - "sequence TEXT PRIMARY KEY, " - "embedding BLOB NOT NULL" - ")" - ) - conn.commit() - - def load_embeddings_from_pth(self, save_path: str) -> Dict[str, torch.Tensor]: - assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}" - payload = torch.load(save_path, map_location="cpu", weights_only=True) - assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary." - for sequence, tensor in payload.items(): - assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)." - assert isinstance(tensor, torch.Tensor), "Expected embedding dictionary values to be tensors." - return payload - - def load_embeddings_from_db(self, db_path: str, sequences: Optional[List[str]] = None) -> Dict[str, torch.Tensor]: - assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}" - loaded: Dict[str, torch.Tensor] = {} - with sqlite3.connect(db_path, timeout=30) as conn: - self._ensure_embeddings_table(conn) - cursor = conn.cursor() - if sequences is None: - cursor.execute("SELECT sequence, embedding FROM embeddings") - else: - if len(sequences) == 0: - return loaded - placeholders = ",".join(["?"] * len(sequences)) - cursor.execute( - f"SELECT sequence, embedding FROM embeddings WHERE sequence IN ({placeholders})", - tuple(sequences), - ) - - rows = cursor.fetchall() - for row in rows: - sequence = row[0] - embedding_bytes = row[1] - loaded[sequence] = embedding_blob_to_tensor(embedding_bytes) - return loaded - - def pool_embeddings( - self, - embeddings: Dict[str, torch.Tensor], - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, - ) -> Dict[str, torch.Tensor]: - return pool_embeddings(embeddings, pooling_types=pooling_types, hidden_state_index=hidden_state_index) - - def load_pooled_embeddings_from_pth( - self, - save_path: str, - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, - ) -> Dict[str, torch.Tensor]: - return load_pooled_embeddings_from_pth( - save_path, - pooling_types=pooling_types, - hidden_state_index=hidden_state_index, - ) - - def load_pooled_embeddings_from_db( - self, - db_path: str, - sequences: Optional[List[str]] = None, - pooling_types: List[str] = ['mean'], - hidden_state_index: int = -1, - ) -> Dict[str, torch.Tensor]: - return load_pooled_embeddings_from_db( - db_path, - sequences=sequences, - pooling_types=pooling_types, - hidden_state_index=hidden_state_index, - ) - - def embed_dataset( - self, - sequences: Optional[List[str]] = None, - tokenizer: Optional[PreTrainedTokenizerBase] = None, - batch_size: int = 2, - max_len: int = 512, - truncate: bool = True, - full_embeddings: bool = False, - embed_dtype: torch.dtype = torch.float32, - pooling_types: List[str] = ['mean'], - num_workers: int = 0, - sql: bool = False, - save: bool = True, - sql_db_path: str = 'embeddings.db', - save_path: str = 'embeddings.pth', - fasta_path: Optional[str] = None, - padding: str = 'max_length', - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - **kwargs, - ) -> Optional[Dict[str, torch.Tensor]]: - """ - Embed a dataset of protein sequences. - - Supports two modes: - - Tokenizer mode (ESM2/ESM++): provide `tokenizer` or use `self.tokenizer`. - - Sequence mode (E1): pass `tokenizer=None`, `_embed(sequences, return_attention_mask=True, **kwargs)` is used. - - Sequences can be supplied as a list via `sequences`, parsed from a FASTA file via - `fasta_path`, or both (the two sources are combined). At least one must be provided. - """ - if fasta_path is not None: - fasta_sequences = parse_fasta(fasta_path) - sequences = list(sequences or []) + fasta_sequences - assert sequences is not None and len(sequences) > 0, \ - "Must provide at least one sequence via `sequences` or `fasta_path`." - assert isinstance(hidden_state_index, int), "hidden_state_index must be an integer." - assert full_embeddings or not store_all_hidden_states, \ - "store_all_hidden_states=True requires full_embeddings=True." - sequences = list(set([seq[:max_len] if truncate else seq for seq in sequences])) - sequences = sorted(sequences, key=len, reverse=True) - pooler = Pooler(pooling_types) if not full_embeddings else None - if tokenizer is None and self.config.model_type != "E1": - tokenizer = self.tokenizer - tokenizer_mode = tokenizer is not None - - # Resolve padding and compilation - dynamic = padding == 'longest' - compiled_model = maybe_compile(self, dynamic=dynamic) - - if tokenizer_mode: - collate_fn = build_collator(tokenizer, padding=padding, max_length=max_len) - device = self.device - else: - collate_fn = None - device = None - - def get_embeddings(residue_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: - assert isinstance(residue_embeddings, torch.Tensor) - if full_embeddings or residue_embeddings.ndim == 2: - return residue_embeddings - return pooler(residue_embeddings, attention_mask) - - def iter_batches(to_embed: List[str]): - if tokenizer_mode: - assert collate_fn is not None - assert device is not None - dataset = ProteinDataset(to_embed) - dataloader = DataLoader( - dataset, - batch_size=batch_size, - num_workers=num_workers, - prefetch_factor=2 if num_workers > 0 else None, - collate_fn=collate_fn, - shuffle=False, - pin_memory=True, - ) - for i, batch in _make_embedding_progress(dataloader, padding): - seqs = to_embed[i * batch_size:(i + 1) * batch_size] - input_ids = batch['input_ids'].to(device) - attention_mask = batch['attention_mask'].to(device) - residue_embeddings = compiled_model._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - yield seqs, residue_embeddings, attention_mask - else: - for batch_start in tqdm(range(0, len(to_embed), batch_size), desc='Embedding batches'): - seqs = to_embed[batch_start:batch_start + batch_size] - batch_output = compiled_model._embed( - seqs, - return_attention_mask=True, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - **kwargs, - ) - assert isinstance(batch_output, tuple), "Sequence mode _embed must return (last_hidden_state, attention_mask)." - assert len(batch_output) == 2, "Sequence mode _embed must return exactly two values." - residue_embeddings, attention_mask = batch_output - assert isinstance(attention_mask, torch.Tensor), "Sequence mode _embed must return attention_mask as a torch.Tensor." - yield seqs, residue_embeddings, attention_mask - - if sql: - # Resume safely: skip sequences already present in the SQLite table. - conn = sqlite3.connect(sql_db_path, timeout=30, check_same_thread=False) - conn.execute('PRAGMA journal_mode=WAL') - conn.execute('PRAGMA busy_timeout=30000') - conn.execute('PRAGMA synchronous=OFF') - conn.execute('PRAGMA cache_size=-64000') - self._ensure_embeddings_table(conn) - already_embedded = self._read_sequences_from_db(sql_db_path) - to_embed = [seq for seq in sequences if seq not in already_embedded] - print(f"Found {len(already_embedded)} already embedded sequences in {sql_db_path}") - print(f"Embedding {len(to_embed)} new sequences") - if len(to_embed) > 0: - # Embed batches synchronously; serialize/write them on the SQL writer thread. - with _SQLWriter(conn) as writer: - with torch.inference_mode(): - for seqs, residue_embeddings, attention_mask in iter_batches(to_embed): - embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype) - if full_embeddings: - batch_rows = [] - for seq, emb, mask in zip(seqs, embeddings, attention_mask): - batch_rows.append((seq, tensor_to_embedding_blob(_trim_full_embedding(emb, mask)))) - else: - blobs = batch_tensor_to_blobs(embeddings) - batch_rows = list(zip(seqs, blobs)) - writer.write_batch(batch_rows) - conn.close() - return None - - embeddings_dict = {} - if os.path.exists(save_path): - embeddings_dict = self.load_embeddings_from_pth(save_path) - to_embed = [seq for seq in sequences if seq not in embeddings_dict] - print(f"Found {len(embeddings_dict)} already embedded sequences in {save_path}") - print(f"Embedding {len(to_embed)} new sequences") - else: - to_embed = sequences - print(f"Embedding {len(to_embed)} new sequences") - - if len(to_embed) > 0: - with torch.inference_mode(): - for seqs, residue_embeddings, attention_mask in iter_batches(to_embed): - embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype) - for seq, emb, mask in zip(seqs, embeddings, attention_mask): - if full_embeddings: - emb = _trim_full_embedding(emb, mask) - embeddings_dict[seq] = emb.cpu() - - if save: - torch.save(embeddings_dict, save_path) - - return embeddings_dict - - -if __name__ == "__main__": - # Manual smoke test for pooling shape behavior. - pooler = Pooler(pooling_types=['max', 'parti']) - batch_size = 8 - seq_len = 64 - hidden_size = 128 - num_layers = 12 - emb = torch.randn(batch_size, seq_len, hidden_size) - attentions = torch.randn(batch_size, num_layers, seq_len, seq_len) - attention_mask = torch.ones(batch_size, seq_len) - y = pooler(emb=emb, attention_mask=attention_mask, attentions=attentions) - print(y.shape) diff --git a/fastplms/entrypoint_setup.py b/fastplms/entrypoint_setup.py deleted file mode 100644 index 6701378..0000000 --- a/fastplms/entrypoint_setup.py +++ /dev/null @@ -1,22 +0,0 @@ -import torch -import torch._inductor.config as inductor_config -import torch._dynamo as dynamo - -# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs) -# Provides significant speedup with minimal precision loss -torch.set_float32_matmul_precision('high') - -# Enable TF32 for matrix multiplications and cuDNN operations -torch.backends.cuda.matmul.allow_tf32 = True -torch.backends.cudnn.allow_tf32 = True - -# Enable cuDNN autotuner - finds fastest algorithms for your hardware -# Best when input sizes are consistent; may slow down first iterations -torch.backends.cudnn.benchmark = True - -# Deterministic operations off for speed (set True if reproducibility needed) -torch.backends.cudnn.deterministic = False -inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM" - -dynamo.config.capture_scalar_outputs = True -torch._dynamo.config.recompile_limit = 16 diff --git a/fastplms/esm2/README.md b/fastplms/esm2/README.md deleted file mode 100644 index c8cdc10..0000000 --- a/fastplms/esm2/README.md +++ /dev/null @@ -1,214 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# FastESM -FastESM is a Huggingface compatible plug in version of ESM2 rewritten with a newer PyTorch attention implementation. - -Load any ESM2 models into a FastEsm model to dramatically speed up training and inference without **ANY** cost in performance. - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. It is fast, memory-efficient, and numerically equivalent to naive attention. The backend is set via `config.attn_backend` before loading. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask — faster on variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30–120 s). | -| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. | - -```python -from transformers import AutoConfig, AutoModel - -config = AutoConfig.from_pretrained("Synthyra/ESM2-150M", trust_remote_code=True) -config.attn_backend = "flex" # or "kernels_flash", "sdpa", "auto" -model = AutoModel.from_pretrained("Synthyra/ESM2-150M", config=config, trust_remote_code=True) -``` - -`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention. - -Attention maps (`output_attentions=True`) are supported with all backends. For SDPA, Flash, and Flex, the attention weights are computed via a separate naive pass, so there is no memory benefit to enabling it during normal inference. -Various other optimizations also make the base implementation slightly different than the one in transformers. - -## Use with 🤗 transformers - -### Supported models -```python -model_dict = { - # Synthyra/ESM2-8M - 'ESM2-8M': 'facebook/esm2_t6_8M_UR50D', - # Synthyra/ESM2-35M - 'ESM2-35M': 'facebook/esm2_t12_35M_UR50D', - # Synthyra/ESM2-150M - 'ESM2-150M': 'facebook/esm2_t30_150M_UR50D', - # Synthyra/ESM2-650M - 'ESM2-650M': 'facebook/esm2_t33_650M_UR50D', - # Synthyra/ESM2-3B - 'ESM2-3B': 'facebook/esm2_t36_3B_UR50D', -} -``` - -### For working with embeddings -```python -import torch -from transformers import AutoModel, AutoTokenizer - -model_path = 'Synthyra/ESM2-8M' -model = AutoModel.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -tokenizer = model.tokenizer - -sequences = ['MPRTEIN', 'MSEQWENCE'] -tokenized = tokenizer(sequences, padding=True, return_tensors='pt') -with torch.no_grad(): - embeddings = model(**tokenized).last_hidden_state - -print(embeddings.shape) # (2, 11, 1280) -``` - -### For working with sequence logits -```python -import torch -from transformers import AutoModelForMaskedLM, AutoTokenizer - -model = AutoModelForMaskedLM.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -with torch.no_grad(): - logits = model(**tokenized).logits - -print(logits.shape) # (2, 11, 33) -``` - -### Experimental test-time training - -TTT is disabled by default. Normal inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the ESM2 backbone using masked language modeling on the test protein. It can -help difficult proteins, but it adds test-time compute and can hurt already -confident predictions. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -### For working with attention maps -```python -import torch -from transformers import AutoModel, AutoTokenizer - -model = AutoModel.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -with torch.no_grad(): - attentions = model(**tokenized, output_attentions).attentions # tuples of (batch_size, num_heads, seq_len, seq_len) - -print(attentions[-1].shape) # (2, 20, 11, 11) -``` - -### Contact prediction -Because we can output attentions using the naive attention implementation, the contact prediction is also supported -```python -with torch.no_grad(): - contact_map = model.predict_contacts(**tokenized).squeeze().cpu().numpy() # (seq_len, seq_len) -``` -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/9707OSXZ3Wdgn0Ni-55T-.png) - -## Embed entire datasets with no new code -To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take. - -Example: -```python -embedding_dict = model.embed_dataset( - sequences=[ - 'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences - ], - tokenizer=model.tokenizer, - batch_size=2, # adjust for your GPU memory - max_len=512, # adjust for your needs - full_embeddings=False, # if True, no pooling is performed - embed_dtype=torch.float32, # cast to what dtype you want - pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together - num_workers=0, # if you have many cpu cores, we find that num_workers = 4 is fast for large datasets - sql=False, # if True, embeddings will be stored in SQLite database - sql_db_path='embeddings.db', - save=True, # if True, embeddings will be saved as a .pth file - save_path='embeddings.pth', -) -# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql -``` - -``` -model.embed_dataset() -Args: - sequences: List of protein sequences - batch_size: Batch size for processing - max_len: Maximum sequence length - full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False) - pooling_type: Type of pooling ('mean' or 'cls') - num_workers: Number of workers for data loading, 0 for the main process - sql: Whether to store embeddings in SQLite database - will be stored in float32 - sql_db_path: Path to SQLite database - -Returns: - Dictionary mapping sequences to embeddings, or None if sql=True - -Note: - - If sql=True, embeddings can only be stored in float32 - - sql is ideal if you need to stream a very large dataset for training in real-time - - save=True is ideal if you can store the entire embedding dictionary in RAM - - sql will be used if it is True and save is True or False - - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences - - Sequences will be truncated to max_len and sorted by length in descending order for faster processing -``` - - -### Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{lin2023esm2, - title={Evolutionary-scale prediction of atomic-level protein structure with a language model}, - author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smestad, Nikita and Verkuil, Robert and Kabeli, Ori and Shmueli, Yaniv and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Salvatore and Rives, Alexander}, - journal={Science}, - volume={379}, - number={6637}, - pages={1123--1130}, - year={2023}, - DOI={10.1126/science.ade2574} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/esm2/README_650.md b/fastplms/esm2/README_650.md deleted file mode 100644 index 927899c..0000000 --- a/fastplms/esm2/README_650.md +++ /dev/null @@ -1,192 +0,0 @@ ---- -library_name: transformers -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# FastESM -FastESM is a Huggingface compatible plug in version of ESM2 rewritten with a newer PyTorch attention implementation. - -Load any ESM2 models into a FastEsm model to dramatically speed up training and inference without **ANY** cost in performance. - -The default attention backend is `sdpa`. See the [FastPLMs README](https://github.com/Synthyra/FastPLMs) for a full breakdown of available backends (`sdpa`, `kernels_flash`, `flex`, `auto`) and how to switch between them. Attention maps (`output_attentions=True`) are supported on all backends via a separate naive computation. -Various other optimizations also make the base implementation slightly different than the one in transformers. - -# FastESM2-650 - -## A faster half-precision version of ESM2-650 with FlashAttention2 and longer context -To enhance the weights with longer context and better fp16 support, we trained ESM2-650 50000 additional steps with a traditional MLM objective (20% masking) in fp16 mixed precision on [OMGprot50](https://huggingface.co/datasets/tattabio/OMG_prot50) up to sequence length of **2048**. - -## Use with 🤗 transformers - -### For working with embeddings -```python -import torch -from transformers import AutoModel, AutoTokenizer - -model_path = 'Synthyra/FastESM2_650' -model = AutoModel.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -tokenizer = model.tokenizer - -sequences = ['MPRTEIN', 'MSEQWENCE'] -tokenized = tokenizer(sequences, padding=True, return_tensors='pt') -with torch.no_grad(): - embeddings = model(**tokenized).last_hidden_state - -print(embeddings.shape) # (2, 11, 1280) -``` - -### For working with sequence logits -```python -import torch -from transformers import AutoModelForMaskedLM, AutoTokenizer - -model = AutoModelForMaskedLM.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -with torch.no_grad(): - logits = model(**tokenized).logits - -print(logits.shape) # (2, 11, 33) -``` - -### Experimental test-time training - -TTT is disabled by default. Normal inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the ESM2 backbone using masked language modeling on the test protein. It can -help difficult proteins, but it adds test-time compute and can hurt already -confident predictions. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -### For working with attention maps -```python -import torch -from transformers import AutoModel, AutoTokenizer - -model = AutoModel.from_pretrained(model_path, dtype=torch.float16, trust_remote_code=True).eval() -with torch.no_grad(): - attentions = model(**tokenized, output_attentions).attentions # tuples of (batch_size, num_heads, seq_len, seq_len) - -print(attentions[-1].shape) # (2, 20, 11, 11) -``` - -## Embed entire datasets with no new code -To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take. - -Example: -```python -embedding_dict = model.embed_dataset( - sequences=[ - 'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences - ], - tokenizer=model.tokenizer, - batch_size=2, # adjust for your GPU memory - max_len=512, # adjust for your needs - full_embeddings=False, # if True, no pooling is performed - embed_dtype=torch.float32, # cast to what dtype you want - pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together - num_workers=0, # if you have many cpu cores, we find that num_workers = 4 is fast for large datasets - sql=False, # if True, embeddings will be stored in SQLite database - sql_db_path='embeddings.db', - save=True, # if True, embeddings will be saved as a .pth file - save_path='embeddings.pth', -) -# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql -``` - -``` -model.embed_dataset() -Args: - sequences: List of protein sequences - batch_size: Batch size for processing - max_len: Maximum sequence length - full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False) - pooling_type: Type of pooling ('mean' or 'cls') - num_workers: Number of workers for data loading, 0 for the main process - sql: Whether to store embeddings in SQLite database - will be stored in float32 - sql_db_path: Path to SQLite database - -Returns: - Dictionary mapping sequences to embeddings, or None if sql=True - -Note: - - If sql=True, embeddings can only be stored in float32 - - sql is ideal if you need to stream a very large dataset for training in real-time - - save=True is ideal if you can store the entire embedding dictionary in RAM - - sql will be used if it is True and save is True or False - - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences - - Sequences will be truncated to max_len and sorted by length in descending order for faster processing -``` - -## Model probes -We employ linear probing techniques on various PLMs and standard datasets, similar our previous [paper](https://www.biorxiv.org/content/10.1101/2024.07.30.605924v1), to assess the intrinsic correlation between pooled hidden states and valuable properties. FastESM performs very well. - -The plot below showcases performance normalized between the negative control (random vector embeddings) and the best performer. Classification task scores are averaged between MCC and F1 (or F1max for multilabel) and regression tasks are averaged between Spearman rho and R2. -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/d1Xi6k1Q4-9By_MtzTvdV.png) - -## Comparison of half precisions -Presumabely because we trained in mixed-precision fp16, fp16 has closer outputs to the fp32 weights then bf16. Therefore, we recommend loading in fp16. - -When summing the MSE of 1000 sequences vs. the fp32 weights: - -Average MSE for FP16: 0.00000140 - -Average MSE for BF16: 0.00004125 - -### Inference speed -We look at various ESM models and their throughput on an H100. FastESM is over twice as fast as ESM2-650 with longer sequences. Requires PyTorch 2.5+ for the most savings, see [SDPA](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html). -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/PvaBGfuJXEW2v_WLkt63y.png) - -### Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@article{lin2023esm2, - title={Evolutionary-scale prediction of atomic-level protein structure with a language model}, - author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smestad, Nikita and Verkuil, Robert and Kabeli, Ori and Shmueli, Yaniv and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Salvatore and Rives, Alexander}, - journal={Science}, - volume={379}, - number={6637}, - pages={1123--1130}, - year={2023}, - DOI={10.1126/science.ade2574} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/esm2/__init__.py b/fastplms/esm2/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/fastplms/esm2/get_weights.py b/fastplms/esm2/get_weights.py deleted file mode 100644 index 2567452..0000000 --- a/fastplms/esm2/get_weights.py +++ /dev/null @@ -1,191 +0,0 @@ -import copy -import os -import torch -from typing import List, Optional, Tuple - -from huggingface_hub import HfApi, login -from transformers import EsmConfig, EsmForMaskedLM, AutoModelForMaskedLM, AutoTokenizer - -from fastplms.esm2.modeling_fastesm import FastEsmConfig, FastEsmForMaskedLM -from fastplms.weight_parity_utils import assert_state_dict_equal, assert_model_parameters_fp32 - - -MODEL_DICT = { - # Synthyra/ESM2-8M - "ESM2-8M": "facebook/esm2_t6_8M_UR50D", - # Synthyra/ESM2-35M - "ESM2-35M": "facebook/esm2_t12_35M_UR50D", - # Synthyra/ESM2-150M - "ESM2-150M": "facebook/esm2_t30_150M_UR50D", - # Synthyra/ESM2-650M - "ESM2-650M": "facebook/esm2_t33_650M_UR50D", - # Synthyra/ESM2-3B - "ESM2-3B": "facebook/esm2_t36_3B_UR50D", -} -SHARDED_REPO_IDS = {"Synthyra/ESM2-3B"} -SHARD_SIZE = "5GB" - - -def _delete_legacy_unsharded_weights_if_present(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - if "model.safetensors" in repo_files: - print(f"Deleting legacy unified model.safetensors from {repo_id}") - api.delete_file( - path_in_repo="model.safetensors", - repo_id=repo_id, - repo_type="model", - ) - - -def _assert_repo_has_sharded_weights(api: HfApi, repo_id: str) -> None: - if repo_id not in SHARDED_REPO_IDS: - return - repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") - has_index_file = "model.safetensors.index.json" in repo_files - has_shard_file = any( - repo_file.startswith("model-") and repo_file.endswith(".safetensors") - for repo_file in repo_files - ) - assert has_index_file, f"{repo_id} is missing model.safetensors.index.json." - assert has_shard_file, f"{repo_id} has no model shard files." - assert "model.safetensors" not in repo_files, f"{repo_id} still has unified model.safetensors." - - -def _push_model_with_expected_format(model: FastEsmForMaskedLM, api: HfApi, repo_id: str) -> None: - if repo_id in SHARDED_REPO_IDS: - print(f"Pushing sharded weights for {repo_id} with max_shard_size={SHARD_SIZE}") - model.push_to_hub(repo_id, max_shard_size=SHARD_SIZE) - _delete_legacy_unsharded_weights_if_present(api, repo_id) - _assert_repo_has_sharded_weights(api, repo_id) - return - model.push_to_hub(repo_id) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return list(MODEL_DICT.items()) - - selected_items: List[Tuple[str, str]] = [] - for repo_id in repo_ids: - # Check if repo_id is a key in MODEL_DICT - if repo_id in MODEL_DICT: - selected_items.append((repo_id, MODEL_DICT[repo_id])) - else: - assert repo_id in MODEL_DICT, ( - f"Unknown model name {repo_id}. " - f"Valid options: {sorted(MODEL_DICT.keys())}" - ) - return selected_items - - -if __name__ == "__main__": - # py -m fastplms.esm2.get_weights - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--dry_run", action="store_true") - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - api = HfApi() - - if args.hf_token is not None: - assert len(args.hf_token) > 0, "--hf_token cannot be empty." - login(token=args.hf_token) - - script_root = os.path.dirname(os.path.abspath(__file__)) - - for model_name, source_repo in _resolve_repo_items(args.repo_ids): - repo_id = "Synthyra/" + model_name - official_config = EsmConfig.from_pretrained(source_repo) - # Makes sure the esm2 word and lm head are correctly loaded - official_config.tie_word_embeddings = True - - official_model = EsmForMaskedLM.from_pretrained( - source_repo, - config=official_config, - dtype=torch.float32, - device_map="cpu", - force_download=True - ) - - config = FastEsmConfig.from_pretrained(source_repo) - config.auto_map = { - "AutoConfig": "modeling_fastesm.FastEsmConfig", - "AutoModel": "modeling_fastesm.FastEsmModel", - "AutoModelForMaskedLM": "modeling_fastesm.FastEsmForMaskedLM", - "AutoModelForSequenceClassification": "modeling_fastesm.FastEsmForSequenceClassification", - "AutoModelForTokenClassification": "modeling_fastesm.FastEsmForTokenClassification", - } - config.tie_word_embeddings = False - if args.skip_weights: - if args.dry_run: - print(f"[skip-weights][dry-run] validated config for {repo_id} <- {source_repo}") - continue - tokenizer = AutoTokenizer.from_pretrained(source_repo) - config.push_to_hub(repo_id) - tokenizer.push_to_hub(repo_id) - print(f"[skip-weights] uploaded config+tokenizer for {repo_id}") - continue - model = FastEsmForMaskedLM.from_pretrained( - source_repo, - config=config, - dtype=torch.float32, - device_map="cpu", - ) - model.load_state_dict(official_model.state_dict(), strict=True) - - # Manually load LM head to prevent weight tying issues - model.lm_head.dense.weight = copy.deepcopy(official_model.lm_head.dense.weight) - model.lm_head.dense.bias = copy.deepcopy(official_model.lm_head.dense.bias) - model.lm_head.decoder.weight = copy.deepcopy(official_model.lm_head.decoder.weight) - model.lm_head.decoder.bias = copy.deepcopy(official_model.lm_head.decoder.bias) - model.lm_head.layer_norm.weight = copy.deepcopy(official_model.lm_head.layer_norm.weight) - model.lm_head.layer_norm.bias = copy.deepcopy(official_model.lm_head.layer_norm.bias) - - assert_model_parameters_fp32( - model=official_model, - model_name=f"official ESM2 model ({source_repo})", - ) - assert_model_parameters_fp32( - model=model, - model_name=f"mapped ESM2 model ({source_repo})", - ) - assert_state_dict_equal( - reference_state_dict=official_model.state_dict(), - candidate_state_dict=model.state_dict(), - context=f"ESM2 weight parity ({source_repo})", - ) - - if args.dry_run: - print(f"[dry_run] validated ESM2 parity for {repo_id} <- {source_repo}") - continue - - tokenizer = model.tokenizer - tokenizer.push_to_hub(repo_id) - - _push_model_with_expected_format(model, api, repo_id) - - api.upload_file( - path_or_fileobj=os.path.join(script_root, "modeling_fastesm.py"), - path_in_repo="modeling_fastesm.py", - repo_id=repo_id, - repo_type="model", - ) - downloaded_model = AutoModelForMaskedLM.from_pretrained( - repo_id, - dtype=torch.float32, - device_map="cpu", - force_download=True, - trust_remote_code=True, - ) - assert_state_dict_equal( - reference_state_dict=official_model.state_dict(), - candidate_state_dict=downloaded_model.state_dict(), - context=f"ESM2 weight parity post-download ({repo_id})", - ) - - \ No newline at end of file diff --git a/fastplms/esm2/modeling_fastesm.py b/fastplms/esm2/modeling_fastesm.py deleted file mode 100644 index 72fc5e0..0000000 --- a/fastplms/esm2/modeling_fastesm.py +++ /dev/null @@ -1,913 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, Dict, List, Optional, Tuple -from einops import rearrange -from dataclasses import dataclass -from transformers import PreTrainedModel, PretrainedConfig, EsmTokenizer -from transformers.modeling_outputs import ModelOutput -from transformers.models.esm.modeling_esm import ( - EsmIntermediate, - EsmOutput, - EsmPooler, - EsmLMHead, - EsmSelfOutput, - EsmClassificationHead, - EsmContactPredictionHead, - EsmEmbeddings, - RotaryEmbedding, -) - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - kernels_flash_attention_func, - index_first_axis, index_put_first_axis, pad_input, _unpad_input, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -@dataclass -class FastEsmEncoderOutput(ModelOutput): - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class EsmMaskedLMOutput(ModelOutput): - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -class FastEsmConfig(PretrainedConfig): - model_type = "fast_esm" - def __init__( - self, - vocab_size: int = None, - mask_token_id: int = None, - pad_token_id: int = None, - hidden_size: int = 768, - num_hidden_layers: int = 12, - num_attention_heads: int = 12, - intermediate_size: int = 3072, - hidden_dropout_prob: float = 0.1, - attention_probs_dropout_prob: float = 0.1, - max_position_embeddings: int = 1026, - initializer_range: float = 0.02, - layer_norm_eps: float = 1e-12, - position_embedding_type: str = "rotary", - emb_layer_norm_before: bool = None, - token_dropout: bool = True, - attn_backend: str = "sdpa", - **kwargs, - ): - super().__init__( - pad_token_id=pad_token_id, - mask_token_id=mask_token_id, - **kwargs, - ) - - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.intermediate_size = intermediate_size - self.hidden_dropout_prob = hidden_dropout_prob - self.attention_probs_dropout_prob = attention_probs_dropout_prob - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.layer_norm_eps = layer_norm_eps - self.position_embedding_type = position_embedding_type - self.emb_layer_norm_before = emb_layer_norm_before - self.tie_word_embeddings = False - self.token_dropout = token_dropout - self.attn_backend = attn_backend - - def to_dict(self) -> Dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. - - Returns: - `Dict[str, any]`: Dictionar y of all the attributes that make up this configuration instance, - """ - output = super().to_dict() - return output - - -class EsmSelfAttention(nn.Module): - def __init__(self, config, position_embedding_type: Optional[str] = None): - super().__init__() - assert config.hidden_size % config.num_attention_heads == 0, ( - f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " - f"heads ({config.num_attention_heads})" - ) - - self.num_attention_heads = config.num_attention_heads - self.attention_head_size = int(config.hidden_size / config.num_attention_heads) - self.all_head_size = self.num_attention_heads * self.attention_head_size - - self.query = nn.Linear(config.hidden_size, self.all_head_size) - self.key = nn.Linear(config.hidden_size, self.all_head_size) - self.value = nn.Linear(config.hidden_size, self.all_head_size) - self.scale = self.attention_head_size**-0.5 - - self.dropout_prob = config.attention_probs_dropout_prob - self.config = config - self.attn_backend = resolve_attention_backend(config.attn_backend) - self.position_embedding_type = position_embedding_type or config.position_embedding_type - self.rotary_embeddings = None - if self.position_embedding_type == "rotary": - self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - batch_size, seq_length = hidden_states.shape[:-1] - hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) - query_BHLD = self.query(hidden_states).view(hidden_shape).transpose(1, 2) - key_BHLD = self.key(hidden_states).view(hidden_shape).transpose(1, 2) - value_BHLD = self.value(hidden_states).view(hidden_shape).transpose(1, 2) - - query_BHLD = query_BHLD * self.scale - - if self.position_embedding_type == "rotary": - query_BHLD, key_BHLD = self.rotary_embeddings(query_BHLD, key_BHLD) - - attn_output, attn_weights, s_max = self._attn( - query_BHLD, key_BHLD, value_BHLD, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - return attn_output, attn_weights, s_max - - def _attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - if output_attentions: - return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d, output_s_max) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - attn_output, attn_weights = self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) - elif self.attn_backend == AttentionBackend.FLEX: - attn_output, attn_weights = self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output, attn_weights = self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - @torch.no_grad() - def _compute_s_max(self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor) -> List[torch.Tensor]: - q_norm = torch.linalg.vector_norm(query_BHLD, dim=-1) - k_norm = torch.linalg.vector_norm(key_BHLD, dim=-1) - s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values - return [s_max_bound[h] for h in range(self.num_attention_heads)] - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[torch.Tensor]]]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - if self.dropout_prob > 0 and self.training: - attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - def _kernels_flash_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - query_BLHD = query_BHLD.transpose(1, 2).contiguous() - key_BLHD = key_BHLD.transpose(1, 2).contiguous() - value_BLHD = value_BHLD.transpose(1, 2).contiguous() - # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward(). - # Pass softmax_scale=1.0 to prevent the kernel from applying its default - # 1/sqrt(head_dim) scale on top (which would yield effective scale - # 1/head_dim and break parity vs sdpa). - attn_output = kernels_flash_attention_func( - query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, - attention_mask_2d=attention_mask_2d, causal=False, - softmax_scale=1.0, - ) - return rearrange(attn_output, "b s h d -> b s (h d)"), None - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask] = None, - ) -> Tuple[torch.Tensor, None]: - assert flex_attention is not None, "Flex attention is not available in this environment." - fn = _get_flex_attention_fn() - context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, - attn_mask=attention_mask_4d, - dropout_p=self.dropout_prob if self.training else 0.0, - scale=1.0, - ) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - -class EsmAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.self = EsmSelfAttention(config) - self.output = EsmSelfOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - hidden_states_ln = self.LayerNorm(hidden_states) - attn_output, attn_weights, s_max = self.self( - hidden_states_ln, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - attention_output = self.output(attn_output, hidden_states) - return attention_output, attn_weights, s_max - - -class EsmLayer(nn.Module): - def __init__(self, config): - super().__init__() - self.chunk_size_feed_forward = config.chunk_size_feed_forward - self.seq_len_dim = 1 - self.attention = EsmAttention(config) - self.intermediate = EsmIntermediate(config) - self.output = EsmOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - attention_output, attn_weights, s_max = self.attention( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - layer_output = self.feed_forward_chunk(attention_output) - return layer_output, attn_weights, s_max - - def feed_forward_chunk(self, attention_output): - attention_output_ln = self.LayerNorm(attention_output) - intermediate_output = self.intermediate(attention_output_ln) - layer_output = self.output(intermediate_output, attention_output) - return layer_output - - -class EsmEncoder(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.attention_backend = resolve_attention_backend(config.attn_backend) - self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)]) - self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: bool = False, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> FastEsmEncoderOutput: - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - full_s_max = () if output_s_max else None - - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=hidden_states.shape[0], - seq_len=hidden_states.shape[1], - device=hidden_states.device, - attention_mask=attention_mask, - ) - - for layer_module in self.layer: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask_2d, - attention_mask_4d, - flex_block_mask, - output_attentions, - output_s_max, - ) - else: - hidden_states, attn_weights, s_max = layer_module( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - if all_attentions is not None: - all_attentions = all_attentions + (attn_weights,) - if full_s_max is not None: - full_s_max = full_s_max + (s_max,) - - if self.emb_layer_norm_after: - hidden_states = self.emb_layer_norm_after(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - return FastEsmEncoderOutput( - last_hidden_state=hidden_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - s_max=full_s_max, - ) - - -class FastEsmPreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - config_class = FastEsmConfig - base_model_prefix = "fastesm" - supports_gradient_checkpointing = True - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - all_tied_weights_keys = {} - - @classmethod - def is_remote_code(cls) -> bool: - # Prevent post-load reinitialization of tensors already loaded from checkpoints. - return True - - @torch.no_grad() - def _init_weights(self, module: nn.Module) -> None: - std = self.config.initializer_range - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - def post_init(self) -> None: - super().post_init() - - def get_output_embeddings(self): - # NOTE: get_output_embeddings() must return None to prevent accidental weight tying. - # See e.g. https://github.com/huggingface/transformers/pull/39339#discussion_r2219126400 - return None - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - for module in self.modules(): - if isinstance(module, EsmEncoder): - module.attention_backend = resolved - elif isinstance(module, EsmSelfAttention): - module.attn_backend = resolved - - -class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin): - def __init__(self, config, add_pooling_layer: Optional[bool] = True, **kwargs): - FastEsmPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.embeddings = EsmEmbeddings(config) - self.encoder = EsmEncoder(config) - self.contact_head = EsmContactPredictionHead( - in_features=config.num_hidden_layers * config.num_attention_heads, bias=True - ) - # Initialize weights and apply final processing - self.post_init() - - def get_input_embeddings(self): - return self.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - encoder_outputs = self.encoder( - token_embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - attns = self(input_ids, attention_mask=attention_mask, output_attentions=True).attentions - attns = torch.stack(attns, dim=1) - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) - return self.contact_head(input_ids, attns) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - ) -> FastEsmEncoderOutput: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) - elif inputs_embeds is None: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - token_embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - ) - encoder_outputs = self.encoder( - token_embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - return FastEsmEncoderOutput( - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - s_max=encoder_outputs.s_max, - ) - - -class FastEsmModel(FastEsmPreTrainedModel, EmbeddingMixin): - def __init__(self, config, add_pooling_layer: Optional[bool] = True, **kwargs): - FastEsmPreTrainedModel.__init__(self, config, **kwargs) - self.config = config - self.esm = FAST_ESM_ENCODER(config) - self.pooler = EsmPooler(config) if add_pooling_layer else None - self.post_init() - - def get_input_embeddings(self): - return self.esm.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.esm.embeddings.word_embeddings = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> FastEsmEncoderOutput: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - - outputs = self.esm( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - sequence_output = outputs.last_hidden_state - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - - return FastEsmEncoderOutput( - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class FastEsmForMaskedLM(FastPLMTestTimeTrainingMixin, FastEsmPreTrainedModel, EmbeddingMixin): - def __init__(self, config, **kwargs): - FastEsmPreTrainedModel.__init__(self, config, **kwargs) - self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) - self.lm_head = EsmLMHead(config) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - self.init_ttt({"lora_target_replace_module": "EsmAttention"}) - - def get_input_embeddings(self): - return self.esm.embeddings.word_embeddings - - def get_output_embeddings(self): - return self.lm_head.decoder - - def set_output_embeddings(self, new_embeddings): - self.lm_head.decoder = new_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.esm] - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> EsmMaskedLMOutput: - outputs = self.esm( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - sequence_output = outputs.last_hidden_state - prediction_scores = self.lm_head(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(prediction_scores.device) - loss = self.loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) - - return EsmMaskedLMOutput( - loss=loss, - logits=prediction_scores, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin): - def __init__(self, config, **kwargs): - FastEsmPreTrainedModel.__init__(self, config, **kwargs) - self.num_labels = config.num_labels - self.config = config - self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) - self.classifier = EsmClassificationHead(config) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - self.post_init() - - def get_input_embeddings(self): - return self.esm.embeddings.word_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> EsmMaskedLMOutput: - outputs = self.esm( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - sequence_output = outputs.last_hidden_state - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - if self.num_labels == 1: - loss = self.mse(logits.squeeze(), labels.squeeze()) - else: - loss = self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return EsmMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -class FastEsmForTokenClassification(FastEsmPreTrainedModel, EmbeddingMixin): - def __init__(self, config, **kwargs): - FastEsmPreTrainedModel.__init__(self, config, **kwargs) - self.num_labels = config.num_labels - self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - self.classifier = nn.Linear(config.hidden_size, config.num_labels) - self.loss_fct = nn.CrossEntropyLoss() - self.post_init() - - def get_input_embeddings(self): - return self.esm.embeddings.word_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - return self.esm._embed( - input_ids, - attention_mask, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def predict_contacts(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> EsmMaskedLMOutput: - outputs = self.esm( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - sequence_output = outputs.last_hidden_state - sequence_output = self.dropout(sequence_output) - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return EsmMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=sequence_output, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - s_max=outputs.s_max, - ) - - -if __name__ == "__main__": - import random - - import torch - - from torch import Tensor - from transformers import EsmTokenizer - - def print_tensor_shapes(prefix: str, obj): - if isinstance(obj, Tensor): - print(f"{prefix}{obj.shape}") - elif isinstance(obj, dict): - for name, value in obj.items(): - print_tensor_shapes(f"{prefix}{name}.", value) - elif isinstance(obj, list): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif isinstance(obj, tuple): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif hasattr(obj, "__dict__"): - for name, value in vars(obj).items(): - if name.startswith("_"): - continue - print_tensor_shapes(f"{prefix}{name}.", value) - else: - print(f"{prefix}{type(obj)}") - - random.seed(0) - torch.manual_seed(0) - - tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - num_attention_heads = random.choice([2, 4]) - config = FastEsmConfig( - vocab_size=tokenizer.vocab_size, - hidden_size=16 * num_attention_heads, - num_attention_heads=num_attention_heads, - num_hidden_layers=random.choice([1, 2]), - intermediate_size=64 * num_attention_heads, - hidden_dropout_prob=0.0, - attention_probs_dropout_prob=0.0, - mask_token_id=tokenizer.mask_token_id, - pad_token_id=tokenizer.pad_token_id, - max_position_embeddings=256, - emb_layer_norm_before=False, - position_embedding_type="rotary", - attn_backend="sdpa", - ) - batch = tokenizer(["ACDEFG", "MKTW"], return_tensors="pt", padding="longest") - batch["labels"] = batch["input_ids"].clone() - model = FastEsmForMaskedLM(config=config).eval() - - with torch.no_grad(): - output = model(**batch, return_dict=True) - - print("Batch shape:") - print_tensor_shapes("", batch) - print("Output shape:") - print_tensor_shapes("", output) diff --git a/fastplms/esm3/README.md b/fastplms/esm3/README.md deleted file mode 100644 index 36698b1..0000000 --- a/fastplms/esm3/README.md +++ /dev/null @@ -1,176 +0,0 @@ ---- -library_name: transformers -license: mit -tags: - - biology - - protein-language-model - - esm3 - - multimodal-protein-model ---- - -# FastPLMs ESM3 Small - -FastPLMs ESM3 Small is a Hugging Face compatible implementation of Biohub's open ESM3 small model. It loads through `AutoModel`, supports sequence-only inference by default, and exposes ESM3's additional tensor tracks directly through normal keyword arguments. - -This repository includes the Biohub ESM MIT license in `LICENSE`. - -## Use With Transformers - -```python -import torch -from transformers import AutoModel - -model = AutoModel.from_pretrained( - "Synthyra/ESM3_small", - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="cuda", -).eval() - -sequences = ["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"] -tokens = model.tokenize_sequences(sequences, device=model.device) - -with torch.inference_mode(): - output = model(**tokens) - -print(output.logits.shape) # sequence logits, (batch_size, seq_len, 64) -print(output.last_hidden_state.shape) # ESM3 embeddings, (batch_size, seq_len, hidden_size) -print(output.function_logits.shape) # function logits, (batch_size, seq_len, 8, 260) -``` - -You can also call sequence inference directly: - -```python -output = model.forward_sequence(["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"]) -``` - -## Experimental Test-Time Training - -TTT is disabled by default. No LoRA adapters are injected during normal -`forward_sequence`, `forward`, or `embed_dataset` calls. Calling `model.ttt(...)` -opts in to experimental masked-LM adaptation of the ESM3 sequence track through -local LoRA weights. It can improve some difficult proteins, but it adds -test-time compute and can degrade already confident predictions. - -```python -metrics = model.ttt( - seq="MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -Switch between SDPA and Flex Attention after loading: - -```python -model.attn_backend = "flex" -output = model.forward_sequence(["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"]) -model.attn_backend = "sdpa" -``` - -## Embed Entire Datasets - -To embed a list of protein sequences, call `embed_dataset`. Sequences are deduplicated, sorted by length, optionally truncated, and embedded in batches. - -```python -embedding_dict = model.embed_dataset( - sequences=[ - "MALWMRLLPLLALLALWGPDPAAA", - "MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP", - ], - batch_size=2, - max_len=512, - full_embeddings=False, - embed_dtype=torch.float32, - pooling_types=["mean", "cls"], - save=True, - save_path="esm3_embeddings.pth", -) - -# embedding_dict maps sequence strings to pooled tensors. -print(embedding_dict["MALWMRLLPLLALLALWGPDPAAA"].shape) -``` - -Residue-wise embeddings are available by setting `full_embeddings=True`: - -```python -residue_embeddings = model.embed_dataset( - sequences=["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"], - batch_size=1, - max_len=512, - full_embeddings=True, - save=False, -) - -print(residue_embeddings["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"].shape) -``` - -FASTA input is also supported: - -```python -embedding_dict = model.embed_dataset( - fasta_path="proteins.fasta", - batch_size=4, - pooling_types=["mean"], - save_path="esm3_fasta_embeddings.pth", -) -``` - -`embed_dataset` currently supports pooled `mean`, `cls`, and `max` embeddings, plus unpooled residue embeddings. It supports `.pth` saves; SQLite streaming is not enabled for the ESM3 wrapper yet. - -## Multimodal Track Arguments - -The default path is amino acid sequence inference. Additional ESM3 tracks can be supplied directly using the same tensor shapes as Biohub ESM3: - -```python -tokens = model.tokenize_sequences( - ["MKTAYIAKQRQISFVKSHFSRQDILDLWIYHTQGYFP"], - device=model.device, -) - -function_tokens = tokens["input_ids"].new_zeros((*tokens["input_ids"].shape, 8)) - -with torch.inference_mode(): - output = model( - **tokens, - function_tokens=function_tokens, - ) - -print(output.sequence_logits.shape) -print(output.function_logits.shape) -``` - -Accepted track arguments include `sequence_tokens`, `structure_tokens`, `ss8_tokens`, `sasa_tokens`, `function_tokens`, `residue_annotation_tokens`, `average_plddt`, `per_res_plddt`, `structure_coords`, `chain_id`, and `sequence_id`. `input_ids` aliases `sequence_tokens`, and `attention_mask` is converted into `sequence_id` if no explicit `sequence_id` is provided. - -## Loading Biohub Checkpoints Locally - -You can build the FastPLMs wrapper from the Biohub checkpoint directly: - -```python -from fastplms.esm3.modeling_esm3 import FastESM3Model - -model = FastESM3Model.from_pretrained_esm("esm3-sm-open-v1", device="cuda") -``` - -This requires Hugging Face access to the gated `biohub/esm3-sm-open-v1` source repo. - -## Biohub SDK Compatibility - -The core forward path is self-contained. Higher-level Biohub SDK workflows are delegated lazily to the official `esm` submodule when available: - -```python -# These methods use Biohub SDK dataclasses and generation configs. -encoded = model.encode(esm_protein) -decoded = model.decode(encoded) -generated = model.generate(esm_protein, generation_config) -``` - -Available delegated methods include `encode`, `decode`, `generate`, `batch_generate`, `logits`, and `forward_and_sample`. - -## Source - -- Biohub ESM repository: https://github.com/Biohub/esm -- Biohub ESM license: https://github.com/Biohub/esm/blob/main/LICENSE.md -- Paper: https://biohub.ai/papers/esm_protein.pdf -- Official model source: https://huggingface.co/biohub/esm3-sm-open-v1 diff --git a/fastplms/esm3/__init__.py b/fastplms/esm3/__init__.py deleted file mode 100644 index 57ed3c4..0000000 --- a/fastplms/esm3/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from fastplms.esm3.modeling_esm3 import FastESM3Config, FastESM3Model - -__all__ = ["FastESM3Config", "FastESM3Model"] diff --git a/fastplms/esm3/get_weights.py b/fastplms/esm3/get_weights.py deleted file mode 100644 index 4a738e9..0000000 --- a/fastplms/esm3/get_weights.py +++ /dev/null @@ -1,167 +0,0 @@ -import argparse -import os -from pathlib import Path -from typing import Dict, List, Optional, Tuple - -import torch -from huggingface_hub import HfApi, login -from transformers import AutoModel - -from fastplms.esm3.modeling_esm3 import ( - ESM3_OPEN_SMALL, - FastESM3Config, - FastESM3Model, - _ESM3_CHECKPOINT_SPECS, - _resolve_esm3_checkpoint_key, -) -from fastplms.weight_parity_utils import assert_model_parameters_fp32, assert_state_dict_equal - - -MODEL_DICT: Dict[str, Tuple[str, str]] = { - "Synthyra/ESM3_small": ("esm3-sm-open-v1", "README.md"), -} - -HUB_AUTO_MAP = { - "AutoConfig": "modeling_esm3.FastESM3Config", - "AutoModel": "modeling_esm3.FastESM3Model", - "AutoModelForMaskedLM": "modeling_esm3.FastESM3Model", -} - - -def _build_config(model_name: str) -> FastESM3Config: - key = _resolve_esm3_checkpoint_key(model_name) - spec = _ESM3_CHECKPOINT_SPECS[key] - config = FastESM3Config( - hidden_size=spec["hidden_size"], - num_attention_heads=spec["num_attention_heads"], - num_vector_heads=spec["num_vector_heads"], - num_hidden_layers=spec["num_hidden_layers"], - model_name=key, - ) - config.architectures = ["FastESM3Model"] - config.auto_map = HUB_AUTO_MAP - config.tie_word_embeddings = False - return config - - -def _upload_repo_files(api: HfApi, repo_id: str, script_root: Path, readme_name: str) -> None: - readme_path = script_root / readme_name - assert readme_path.exists(), f"Missing model card: {readme_path}" - api.upload_file( - path_or_fileobj=str(script_root / "modeling_esm3.py"), - path_in_repo="modeling_esm3.py", - repo_id=repo_id, - repo_type="model", - ) - api.upload_file( - path_or_fileobj=str(readme_path), - path_in_repo="README.md", - repo_id=repo_id, - repo_type="model", - ) - license_path = script_root / "LICENSE" - assert license_path.exists(), f"Missing license: {license_path}" - api.upload_file( - path_or_fileobj=str(license_path), - path_in_repo="LICENSE", - repo_id=repo_id, - repo_type="model", - ) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return [ - (repo_id, esm3_model_key, readme_name) - for repo_id, (esm3_model_key, readme_name) in MODEL_DICT.items() - ] - - selected_items: List[Tuple[str, str, str]] = [] - for repo_id in repo_ids: - assert repo_id in MODEL_DICT, ( - f"Unknown repo_id {repo_id}. " - f"Valid options: {sorted(MODEL_DICT.keys())}" - ) - esm3_model_key, readme_name = MODEL_DICT[repo_id] - selected_items.append((repo_id, esm3_model_key, readme_name)) - return selected_items - - -def _token_from_environment() -> Optional[str]: - for key in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN"): - if key in os.environ and len(os.environ[key]) > 0: - return os.environ[key] - return None - - -def _login_if_requested(args: argparse.Namespace) -> None: - token = args.hf_token - if token is None: - token = _token_from_environment() - if token is not None: - assert len(token) > 0, "HF token cannot be empty." - login(token=token) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--hf_token", - type=str, - default=None, - help="Deprecated. Prefer HF_TOKEN in the environment so tokens are not in shell history.", - ) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--dry_run", action="store_true") - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - - script_root = Path(__file__).resolve().parent - _login_if_requested(args) - api = HfApi() - - for repo_id, esm3_model_key, readme_name in _resolve_repo_items(args.repo_ids): - config = _build_config(esm3_model_key) - if args.skip_weights: - if args.dry_run: - print(f"[skip-weights][dry-run] validated config metadata for {repo_id}") - continue - config.push_to_hub(repo_id) - _upload_repo_files(api, repo_id, script_root, readme_name) - print(f"[skip-weights] uploaded config for {repo_id}") - continue - - model = FastESM3Model.from_pretrained_esm( - ESM3_OPEN_SMALL, - device=torch.device("cpu"), - dtype=torch.float32, - ) - model.config.architectures = ["FastESM3Model"] - model.config.auto_map = HUB_AUTO_MAP - model.config.tie_word_embeddings = False - tokenizer = model.tokenizer - - assert_model_parameters_fp32( - model=model, - model_name=f"mapped ESM3 model ({esm3_model_key})", - ) - - if args.dry_run: - print(f"[dry_run] validated ESM3 conversion for {repo_id} <- {esm3_model_key}") - continue - - tokenizer.push_to_hub(repo_id) - model.push_to_hub(repo_id) - _upload_repo_files(api, repo_id, script_root, readme_name) - downloaded_model = AutoModel.from_pretrained( - repo_id, - dtype=torch.float32, - device_map="cpu", - force_download=True, - trust_remote_code=True, - ) - assert_state_dict_equal( - reference_state_dict=model.state_dict(), - candidate_state_dict=downloaded_model.state_dict(), - context=f"ESM3 weight parity post-download ({repo_id})", - ) diff --git a/fastplms/esm3/modeling_esm3.py b/fastplms/esm3/modeling_esm3.py deleted file mode 100644 index 80ea342..0000000 --- a/fastplms/esm3/modeling_esm3.py +++ /dev/null @@ -1,1783 +0,0 @@ -from __future__ import annotations - -""" -Hugging Face compatible ESM3 wrapper. - -This module keeps Biohub's ESM3 implementation as the execution core and adds -the FastPLMs conventions around it: AutoModel loading, sequence-only -`input_ids` forwarding, and direct multimodal track arguments. -""" - -import sys -import math -import functools -import importlib -import os -from dataclasses import dataclass -from enum import Enum -from pathlib import Path -from typing import Dict, List, Optional, Union - -import einops -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange -from huggingface_hub import snapshot_download -from tokenizers import Tokenizer -from tokenizers.models import BPE -from tokenizers.processors import TemplateProcessing -from transformers import PreTrainedModel, PreTrainedTokenizerFast, PretrainedConfig -from transformers.modeling_outputs import ModelOutput - -try: - from torch.nn.attention.flex_attention import ( - BlockMask, - create_block_mask, - flex_attention, - ) -except ImportError: - BlockMask = None - create_block_mask = None - flex_attention = None - -try: - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -ESM3_OPEN_SMALL = "esm3_sm_open_v1" -ESM3_OPEN_SMALL_ALIASES = { - "ESM3_small", - "esm3_small", - "esm3_sm_open_v1", - "esm3-open-2024-03", - "esm3-sm-open-v1", - "esm3-open", -} - -SEQUENCE_BOS_TOKEN = 0 -SEQUENCE_PAD_TOKEN = 1 -SEQUENCE_EOS_TOKEN = 2 -SEQUENCE_CHAINBREAK_TOKEN = 31 -SEQUENCE_MASK_TOKEN = 32 - -VQVAE_CODEBOOK_SIZE = 4096 -STRUCTURE_MASK_TOKEN = VQVAE_CODEBOOK_SIZE -STRUCTURE_EOS_TOKEN = VQVAE_CODEBOOK_SIZE + 1 -STRUCTURE_BOS_TOKEN = VQVAE_CODEBOOK_SIZE + 2 -STRUCTURE_PAD_TOKEN = VQVAE_CODEBOOK_SIZE + 3 -STRUCTURE_CHAINBREAK_TOKEN = VQVAE_CODEBOOK_SIZE + 4 - -SASA_PAD_TOKEN = 0 -SS8_PAD_TOKEN = 0 -INTERPRO_PAD_TOKEN = 0 -RESIDUE_PAD_TOKEN = 0 -MAX_RESIDUE_ANNOTATIONS = 16 -FUNCTION_TOKENS_DEPTH = 8 - -SEQUENCE_VOCAB = [ - "", - "", - "", - "", - "L", - "A", - "G", - "V", - "S", - "E", - "R", - "T", - "I", - "D", - "P", - "K", - "Q", - "N", - "F", - "Y", - "M", - "H", - "W", - "C", - "X", - "B", - "U", - "Z", - "O", - ".", - "-", - "|", - "", -] - -_SUPPORTED_ATTENTION_BACKENDS = ("auto", "flex", "sdpa") -_compiled_flex_attention = None - - -class AttentionBackend(Enum): - AUTO = "auto" - FLEX = "flex" - SDPA = "sdpa" - - -def _get_flex_attention_fn(): - global _compiled_flex_attention - if flex_attention is None: - return None - flex_mod = torch.nn.attention.flex_attention - if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False): - return flex_attention - if _compiled_flex_attention is None: - _compiled_flex_attention = torch.compile( - flex_attention, - dynamic=False, - ) - return _compiled_flex_attention - - -def resolve_attention_backend(requested_backend: str) -> AttentionBackend: - assert requested_backend in _SUPPORTED_ATTENTION_BACKENDS, ( - f"Unsupported ESM3 attention backend: {requested_backend}. " - f"Expected one of {_SUPPORTED_ATTENTION_BACKENDS}." - ) - if requested_backend == AttentionBackend.AUTO.value: - if flex_attention is not None: - return AttentionBackend.FLEX - return AttentionBackend.SDPA - if requested_backend == AttentionBackend.FLEX.value: - assert flex_attention is not None, "Flex Attention is not available in this environment." - return AttentionBackend.FLEX - if requested_backend == AttentionBackend.SDPA.value: - return AttentionBackend.SDPA - raise AssertionError(f"Unsupported ESM3 attention backend: {requested_backend}") - -_ESM3_CHECKPOINT_SPECS = { - ESM3_OPEN_SMALL: { - "repo_id": "biohub/esm3-sm-open-v1", - "hidden_size": 1536, - "num_attention_heads": 24, - "num_vector_heads": 256, - "num_hidden_layers": 48, - }, -} - - -class FastESM3Config(PretrainedConfig): - model_type = "fast_esm3" - - def __init__( - self, - vocab_size: int = 64, - hidden_size: int = 1536, - num_attention_heads: int = 24, - num_vector_heads: int = 256, - num_hidden_layers: int = 48, - initializer_range: float = 0.02, - attn_backend: str = "sdpa", - model_name: str = ESM3_OPEN_SMALL, - **kwargs, - ): - super().__init__(**kwargs) - assert hidden_size % FUNCTION_TOKENS_DEPTH == 0 - assert hidden_size % num_attention_heads == 0 - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads - self.num_vector_heads = num_vector_heads - self.num_hidden_layers = num_hidden_layers - self.initializer_range = initializer_range - self.attn_backend = attn_backend - self.model_name = _resolve_esm3_checkpoint_key(model_name) - self.tie_word_embeddings = False - - -@dataclass -class FastESM3Output(ModelOutput): - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - sequence_logits: Optional[torch.Tensor] = None - structure_logits: Optional[torch.Tensor] = None - secondary_structure_logits: Optional[torch.Tensor] = None - sasa_logits: Optional[torch.Tensor] = None - function_logits: Optional[torch.Tensor] = None - residue_logits: Optional[torch.Tensor] = None - embeddings: Optional[torch.Tensor] = None - hidden_states: Optional[tuple[torch.Tensor, ...]] = None - attentions: Optional[tuple[torch.Tensor, ...]] = None - - -class EsmSequenceTokenizer(PreTrainedTokenizerFast): - model_input_names = ["input_ids", "attention_mask"] - - def __init__( - self, - unk_token: str = "", - cls_token: str = "", - pad_token: str = "", - mask_token: str = "", - eos_token: str = "", - chain_break_token: str = "|", - **kwargs, - ): - token_to_id = {token: index for index, token in enumerate(SEQUENCE_VOCAB)} - bpe = BPE(token_to_id, merges=[], unk_token=unk_token) - tokenizer = Tokenizer(bpe) - special_tokens = [ - cls_token, - pad_token, - mask_token, - eos_token, - chain_break_token, - ] - self.cb_token = chain_break_token - tokenizer.add_special_tokens(special_tokens) - tokenizer.post_processor = TemplateProcessing( - single=" $A ", - pair=":0 $A:0 :0 $B:1 :1", - special_tokens=[ - ("", tokenizer.token_to_id("")), - ("", tokenizer.token_to_id("")), - ], - ) - super().__init__( - tokenizer_object=tokenizer, - unk_token=unk_token, - cls_token=cls_token, - pad_token=pad_token, - mask_token=mask_token, - eos_token=eos_token, - additional_special_tokens=[chain_break_token], - **kwargs, - ) - - @property - def bos_token(self) -> str: - return self.cls_token - - @property - def bos_token_id(self) -> int: - return self.cls_token_id - - @property - def chain_break_token(self) -> str: - return self.cb_token - - @property - def chain_break_token_id(self) -> int: - token_id = self.convert_tokens_to_ids(self.chain_break_token) - assert isinstance(token_id, int) - return token_id - - @property - def all_token_ids(self) -> list[int]: - return list(range(self.vocab_size)) - - @property - def special_token_ids(self) -> list[int]: - return self.all_special_ids - - -@dataclass -class FastESM3TokenizerCollection: - sequence: EsmSequenceTokenizer - structure: Optional[object] = None - secondary_structure: Optional[object] = None - sasa: Optional[object] = None - function: Optional[object] = None - residue_annotations: Optional[object] = None - - -def rbf(values: torch.Tensor, v_min: float, v_max: float, n_bins: int = 16) -> torch.Tensor: - centers = torch.linspace( - v_min, - v_max, - n_bins, - device=values.device, - dtype=values.dtype, - ) - centers = centers.view([1] * len(values.shape) + [-1]) - std = (v_max - v_min) / n_bins - z = (values.unsqueeze(-1) - centers) / std - return torch.exp(-(z**2)) - - -def RegressionHead( - d_model: int, - output_dim: int, - hidden_dim: Optional[int] = None, -) -> nn.Module: - hidden_dim = hidden_dim if hidden_dim is not None else d_model - return nn.Sequential( - nn.Linear(d_model, hidden_dim), - nn.GELU(), - nn.LayerNorm(hidden_dim), - nn.Linear(hidden_dim, output_dim), - ) - - -def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange( - torch.stack((-x2, x1), dim=-1), - "... d two -> ... (d two)", - two=2, - ) - - -def apply_rotary_emb_torch( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - interleaved: bool = False, -) -> torch.Tensor: - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - seqlen = x.size(1) - cos = cos[:seqlen] - sin = sin[:seqlen] - cos = einops.repeat(cos, "s d -> s 1 (2 d)") - sin = einops.repeat(sin, "s d -> s 1 (2 d)") - return torch.cat( - [ - x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, - x[..., ro_dim:], - ], - dim=-1, - ) - - -class RotaryEmbedding(nn.Module): - def __init__( - self, - dim: int, - base: float = 10000.0, - interleaved: bool = False, - scale_base: Optional[float] = None, - scaling_factor: float = 1.0, - pos_idx_in_fp32: bool = True, - device: Optional[torch.device] = None, - ): - super().__init__() - self.dim = dim - self.base = float(base) - self.pos_idx_in_fp32 = pos_idx_in_fp32 - self.interleaved = interleaved - self.scale_base = scale_base - self.scaling_factor = scaling_factor - self.device = device - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self.reset_parameters() - - def reset_parameters(self) -> None: - inv_freq = self._compute_inv_freq(self.device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - arange = torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32) - scale = ( - (arange + 0.4 * self.dim) / (1.4 * self.dim) - if self.scale_base is not None - else None - ) - self.register_buffer("scale", scale) - - def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor: - return 1 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) - / self.dim - ) - ) - - def _update_cos_sin_cache( - self, - seqlen: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> None: - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - if self.pos_idx_in_fp32: - t = torch.arange(seqlen, device=device, dtype=torch.float32) - t /= self.scaling_factor - inv_freq = self.inv_freq.to(torch.float32) - else: - t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) - t /= self.scaling_factor - inv_freq = self.inv_freq - freqs = torch.outer(t, inv_freq) - - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - raise NotImplementedError("Scaled rotary embeddings are not used by ESM3.") - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - seqlen_offset: int = 0, - ) -> tuple[torch.Tensor, torch.Tensor]: - self._update_cos_sin_cache( - q.shape[1] + seqlen_offset, - device=q.device, - dtype=q.dtype, - ) - assert self._cos_cached is not None - assert self._sin_cached is not None - return ( - apply_rotary_emb_torch( - q, - self._cos_cached[seqlen_offset:], - self._sin_cached[seqlen_offset:], - self.interleaved, - ), - apply_rotary_emb_torch( - k, - self._cos_cached[seqlen_offset:], - self._sin_cached[seqlen_offset:], - self.interleaved, - ), - ) - - -def fp32_autocast_context(device_type: str): - if device_type == "cuda": - return torch.autocast(device_type="cuda", enabled=False) - return torch.autocast(device_type=device_type, enabled=False) - - -class RotationMatrix: - def __init__(self, rots: torch.Tensor): - if rots.shape[-1] == 9: - rots = rots.unflatten(-1, (3, 3)) - assert rots.shape[-1] == 3 - assert rots.shape[-2] == 3 - self._rots = rots.to(torch.float32) - - @classmethod - def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> "RotationMatrix": - rots = torch.eye(3, **tensor_kwargs) - rots = rots.view(*[1 for _ in range(len(shape))], 3, 3) - rots = rots.expand(*shape, -1, -1) - return cls(rots) - - def __getitem__(self, idx) -> "RotationMatrix": - indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) - return RotationMatrix(self._rots[indices + (slice(None), slice(None))]) - - @property - def shape(self) -> torch.Size: - return self._rots.shape[:-2] - - @property - def tensor(self) -> torch.Tensor: - return self._rots.flatten(-2) - - @property - def device(self) -> torch.device: - return self._rots.device - - def as_matrix(self) -> "RotationMatrix": - return self - - def apply(self, p: torch.Tensor) -> torch.Tensor: - with fp32_autocast_context(self.device.type): - p = p.to(self._rots.dtype) - if self._rots.shape[-3] == 1: - return p @ self._rots.transpose(-1, -2).squeeze(-3) - return torch.einsum("...ij,...j", self._rots, p) - - def invert(self) -> "RotationMatrix": - return RotationMatrix(self._rots.transpose(-1, -2)) - - @staticmethod - def from_graham_schmidt( - x_axis: torch.Tensor, - xy_plane: torch.Tensor, - eps: float = 1e-12, - ) -> "RotationMatrix": - with fp32_autocast_context(x_axis.device.type): - e1 = xy_plane - denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps) - x_axis = x_axis / denom - dot = (x_axis * e1).sum(dim=-1, keepdim=True) - e1 = e1 - x_axis * dot - denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps) - e1 = e1 / denom - e2 = torch.cross(x_axis, e1, dim=-1) - return RotationMatrix(torch.stack([x_axis, e1, e2], dim=-1)) - - -@dataclass(frozen=True) -class Affine3D: - trans: torch.Tensor - rot: RotationMatrix - - def __post_init__(self) -> None: - assert self.trans.shape[:-1] == self.rot.shape - - def __getitem__(self, idx) -> "Affine3D": - indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) - return Affine3D( - trans=self.trans[indices + (slice(None),)], - rot=self.rot[idx], - ) - - @property - def shape(self) -> torch.Size: - return self.trans.shape[:-1] - - @property - def dtype(self) -> torch.dtype: - return self.trans.dtype - - @property - def device(self) -> torch.device: - return self.trans.device - - @property - def tensor(self) -> torch.Tensor: - return torch.cat([self.rot.tensor, self.trans], dim=-1) - - def as_matrix(self) -> "Affine3D": - return Affine3D(trans=self.trans, rot=self.rot.as_matrix()) - - def apply(self, p: torch.Tensor) -> torch.Tensor: - return self.rot.apply(p) + self.trans - - @staticmethod - def from_tensor(t: torch.Tensor) -> "Affine3D": - match t.shape[-1]: - case 12: - trans = t[..., -3:] - rot = RotationMatrix(t[..., :-3].unflatten(-1, (3, 3))) - case _: - raise RuntimeError( - f"Cannot detect rotation format from {t.shape[-1] - 3}-d flat vector" - ) - return Affine3D(trans, rot) - - @staticmethod - def from_graham_schmidt( - neg_x_axis: torch.Tensor, - origin: torch.Tensor, - xy_plane: torch.Tensor, - eps: float = 1e-10, - ) -> "Affine3D": - x_axis = origin - neg_x_axis - xy_plane = xy_plane - origin - return Affine3D( - trans=origin, - rot=RotationMatrix.from_graham_schmidt(x_axis, xy_plane, eps), - ) - - -def build_affine3d_from_coordinates(coords: torch.Tensor) -> tuple[Affine3D, torch.Tensor]: - max_supported_distance = 1e6 - coord_mask = torch.all( - torch.all(torch.isfinite(coords) & (coords < max_supported_distance), dim=-1), - dim=-1, - ) - - def atom3_to_backbone_affine(bb_positions: torch.Tensor) -> Affine3D: - n_atom, ca_atom, c_atom = bb_positions.unbind(dim=-2) - return Affine3D.from_graham_schmidt(c_atom, ca_atom, n_atom) - - coords = coords.clone().float() - coords[~coord_mask] = 0 - average_per_n_ca_c = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / ( - coord_mask.sum(-1)[..., None, None] + 1e-8 - ) - affine_from_average = atom3_to_backbone_affine( - average_per_n_ca_c.float() - ).as_matrix() - - batch_size, seq_len, _, _ = coords.shape - affine_rot_mats = affine_from_average.rot.tensor[..., None, :].expand( - batch_size, - seq_len, - 9, - ) - affine_trans = affine_from_average.trans[..., None, :].expand(batch_size, seq_len, 3) - identity_rot = RotationMatrix.identity( - (batch_size, seq_len), - dtype=torch.float32, - device=coords.device, - requires_grad=False, - ) - affine_rot_mats = affine_rot_mats.where( - coord_mask.any(-1)[..., None, None], - identity_rot.tensor, - ) - black_hole_affine = Affine3D(affine_trans, RotationMatrix(affine_rot_mats)) - - affine = atom3_to_backbone_affine(coords.float()) - affine = Affine3D.from_tensor( - affine.tensor.where(coord_mask[..., None], black_hole_affine.tensor) - ) - return affine, coord_mask - - -class MultiHeadAttention(nn.Module): - def __init__( - self, - d_model: int, - n_heads: int, - bias: bool = False, - qk_layernorm: bool = True, - attn_backend: str = "sdpa", - ): - super().__init__() - self.d_model = d_model - self.n_heads = n_heads - self.d_head = self.d_model // self.n_heads - self.scale = self.d_head**-0.5 - self.attn_backend = resolve_attention_backend(attn_backend) - self.layernorm_qkv = nn.Sequential( - nn.LayerNorm(d_model), - nn.Linear(d_model, d_model * 3, bias=bias), - ) - self.out_proj = nn.Linear(d_model, d_model, bias=bias) - if qk_layernorm: - self.q_ln = nn.LayerNorm(d_model, bias=bias) - self.k_ln = nn.LayerNorm(d_model, bias=bias) - else: - self.q_ln = nn.Identity() - self.k_ln = nn.Identity() - self.rotary = RotaryEmbedding(d_model // n_heads) - - def _apply_rotary( - self, - q: torch.Tensor, - k: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - q = q.unflatten(-1, (self.n_heads, self.d_head)) - k = k.unflatten(-1, (self.n_heads, self.d_head)) - q, k = self.rotary(q, k) - q = q.flatten(-2, -1) - k = k.flatten(-2, -1) - return q, k - - def forward( - self, - x: torch.Tensor, - seq_id: Optional[torch.Tensor], - output_attentions: bool = False, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - qkv = self.layernorm_qkv(x) - query, key, value = torch.chunk(qkv, 3, dim=-1) - query = self.q_ln(query).to(query.dtype) - key = self.k_ln(key).to(query.dtype) - query, key = self._apply_rotary(query, key) - - reshaper = functools.partial( - einops.rearrange, - pattern="b s (h d) -> b h s d", - h=self.n_heads, - ) - query, key, value = map(reshaper, (query, key, value)) - - if seq_id is not None: - mask = seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2) - mask = mask.unsqueeze(1) - else: - mask = None - - if output_attentions: - attn_scores = torch.einsum("bhld,bhsd->bhls", query, key) * self.scale - if mask is not None: - attn_scores = attn_scores.masked_fill(~mask, float("-inf")) - attn_weights = torch.softmax(attn_scores, dim=-1) - context = torch.einsum("bhls,bhsd->bhld", attn_weights, value) - else: - attn_weights = None - if self.attn_backend == AttentionBackend.FLEX: - block_mask = self._create_flex_block_mask(seq_id, query) - fn = _get_flex_attention_fn() - assert fn is not None, "Flex Attention is not available in this environment." - context = fn( - query, - key, - value, - block_mask=block_mask, - scale=self.scale, - ) - elif self.attn_backend == AttentionBackend.SDPA: - context = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=mask, - scale=self.scale, - ) - else: - raise AssertionError(f"Unsupported resolved ESM3 backend: {self.attn_backend}") - - context = einops.rearrange(context, "b h s d -> b s (h d)") - return self.out_proj(context), attn_weights - - @staticmethod - def _create_flex_block_mask( - seq_id: Optional[torch.Tensor], - query: torch.Tensor, - ) -> Optional["BlockMask"]: - if seq_id is None: - return None - assert create_block_mask is not None, ( - "Flex Attention requested but torch.create_block_mask is unavailable." - ) - batch_size, _, seq_len, _ = query.shape - - def mask_mod(batch_idx, head_idx, q_idx, kv_idx): - return seq_id[batch_idx, q_idx] == seq_id[batch_idx, kv_idx] - - return create_block_mask( - mask_mod, - batch_size, - 1, - seq_len, - seq_len, - device=query.device, - ) - - -class GeometricReasoningOriginalImpl(nn.Module): - def __init__( - self, - c_s: int, - v_heads: int, - num_vector_messages: int = 1, - mask_and_zero_frameless: bool = True, - bias: bool = False, - ): - super().__init__() - self.c_s = c_s - self.v_heads = v_heads - self.num_vector_messages = num_vector_messages - self.mask_and_zero_frameless = mask_and_zero_frameless - self.s_norm = nn.LayerNorm(c_s, bias=bias) - dim_proj = 4 * self.v_heads * 3 + self.v_heads * 3 * self.num_vector_messages - self.proj = nn.Linear(c_s, dim_proj, bias=bias) - channels_out = self.v_heads * 3 * self.num_vector_messages - self.out_proj = nn.Linear(channels_out, c_s, bias=bias) - self.distance_scale_per_head = nn.Parameter(torch.zeros((self.v_heads))) - self.rotation_scale_per_head = nn.Parameter(torch.zeros((self.v_heads))) - - def forward( - self, - s: torch.Tensor, - affine: Affine3D, - affine_mask: torch.Tensor, - sequence_id: Optional[torch.Tensor], - chain_id: torch.Tensor, - ) -> torch.Tensor: - if sequence_id is None: - sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64) - attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2) - attn_bias = attn_bias.unsqueeze(1).float() - attn_bias = attn_bias.masked_fill( - ~affine_mask[:, None, None, :], - torch.finfo(attn_bias.dtype).min, - ) - chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2) - attn_bias = attn_bias.masked_fill( - chain_id_mask.unsqueeze(1), - torch.finfo(s.dtype).min, - ) - - ns = self.s_norm(s) - vec_rot, vec_dist = self.proj(ns).split( - [ - self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages, - self.v_heads * 2 * 3, - ], - dim=-1, - ) - - query_rot, key_rot, value = ( - affine.rot[..., None] - .apply(rearrange(vec_rot, "... (h c) -> ... h c", c=3)) - .split( - [self.v_heads, self.v_heads, self.v_heads * self.num_vector_messages], - dim=-2, - ) - ) - query_dist, key_dist = ( - affine[..., None] - .apply(rearrange(vec_dist, "... (h c) -> ... h c", c=3)) - .chunk(2, dim=-2) - ) - - query_dist = rearrange(query_dist, "b s h d -> b h s 1 d") - key_dist = rearrange(key_dist, "b s h d -> b h 1 s d") - query_rot = rearrange(query_rot, "b s h d -> b h s d") - key_rot = rearrange(key_rot, "b s h d -> b h d s") - value = rearrange( - value, - "b s (h m) d -> b h s (m d)", - m=self.num_vector_messages, - ) - - distance_term = (query_dist - key_dist).norm(dim=-1) / math.sqrt(3) - rotation_term = query_rot.matmul(key_rot) / math.sqrt(3) - distance_term_weight = rearrange( - F.softplus(self.distance_scale_per_head), - "h -> h 1 1", - ) - rotation_term_weight = rearrange( - F.softplus(self.rotation_scale_per_head), - "h -> h 1 1", - ) - attn_weight = ( - rotation_term * rotation_term_weight - distance_term * distance_term_weight - ) - - s_q = attn_weight.size(2) - s_k = attn_weight.size(3) - offset_q = max(0, attn_bias.size(2) - s_q) - offset_k = max(0, attn_bias.size(3) - s_k) - attn_bias = attn_bias[:, :, offset_q:, offset_k:] - attn_weight = torch.softmax(attn_weight + attn_bias, dim=-1) - - attn_out = attn_weight.matmul(value) - attn_out = ( - affine.rot[..., None] - .invert() - .apply( - rearrange( - attn_out, - "b h s (m d) -> b s (h m) d", - m=self.num_vector_messages, - ) - ) - ) - attn_out = rearrange( - attn_out, - "b s (h m) d -> b s (h m d)", - m=self.num_vector_messages, - ) - if self.mask_and_zero_frameless: - attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0) - attn_out = attn_out.to(self.out_proj.weight.dtype) - return self.out_proj(attn_out) - - -def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int: - return int(((expansion_ratio * d_model) + 255) // 256 * 256) - - -class SwiGLU(nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return F.silu(x1) * x2 - - -def swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module: - return nn.Sequential( - nn.LayerNorm(d_model), - nn.Linear( - d_model, - swiglu_correction_fn(expansion_ratio, d_model) * 2, - bias=bias, - ), - SwiGLU(), - nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=bias), - ) - - -def gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module: - hidden_dim = int(expansion_ratio * d_model) - return nn.Sequential( - nn.LayerNorm(d_model), - nn.Linear(d_model, hidden_dim, bias=bias), - nn.GELU(), - nn.Linear(hidden_dim, d_model, bias=bias), - ) - - -class UnifiedTransformerBlock(nn.Module): - def __init__( - self, - d_model: int, - n_heads: int, - use_geom_attn: bool = False, - use_plain_attn: bool = True, - v_heads: Optional[int] = None, - bias: bool = False, - expansion_ratio: float = 4.0, - residue_scaling_factor: float = 1.0, - mask_and_zero_frameless: bool = False, - qk_layernorm: bool = True, - ffn_type: str = "swiglu", - attn_backend: str = "sdpa", - ): - super().__init__() - self.use_plain_attn = use_plain_attn - if self.use_plain_attn: - self.attn = MultiHeadAttention( - d_model, - n_heads, - bias, - qk_layernorm=qk_layernorm, - attn_backend=attn_backend, - ) - self.use_geom_attn = use_geom_attn - if self.use_geom_attn: - assert v_heads is not None - self.geom_attn = GeometricReasoningOriginalImpl( - c_s=d_model, - v_heads=v_heads, - bias=bias, - mask_and_zero_frameless=mask_and_zero_frameless, - ) - if ffn_type == "swiglu": - self.ffn = swiglu_ln_ffn(d_model, expansion_ratio, bias) - elif ffn_type == "gelu": - self.ffn = gelu_ln_ffn(d_model, expansion_ratio, bias) - else: - raise ValueError(f"Unknown ffn_type: {ffn_type}") - self.scaling_factor = residue_scaling_factor - - def forward( - self, - x: torch.Tensor, - sequence_id: Optional[torch.Tensor], - frames: Affine3D, - frames_mask: torch.Tensor, - chain_id: torch.Tensor, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - attn_weights = None - if self.use_plain_attn: - r1, attn_weights = self.attn( - x, - sequence_id, - output_attentions=output_attentions, - ) - x = x + r1 / self.scaling_factor - - if self.use_geom_attn: - r2 = self.geom_attn(x, frames, frames_mask, sequence_id, chain_id) - x = x + r2 / self.scaling_factor - - r3 = self.ffn(x) / self.scaling_factor - x = x + r3 - return x, attn_weights - - -class TransformerStack(nn.Module): - def __init__( - self, - d_model: int, - n_heads: int, - v_heads: Optional[int], - n_layers: int, - n_layers_geom: int = 1, - scale_residue: bool = True, - mask_and_zero_frameless: bool = False, - bias: bool = False, - qk_layernorm: bool = True, - ffn_type: str = "swiglu", - expansion_ratio: float = 8 / 3, - attn_backend: str = "sdpa", - ): - super().__init__() - self.blocks = nn.ModuleList( - [ - UnifiedTransformerBlock( - d_model, - n_heads, - v_heads=v_heads, - use_geom_attn=index < n_layers_geom, - residue_scaling_factor=( - math.sqrt(n_layers / 36) if scale_residue else 1.0 - ), - expansion_ratio=expansion_ratio, - mask_and_zero_frameless=mask_and_zero_frameless, - bias=bias, - qk_layernorm=qk_layernorm, - ffn_type=ffn_type, - attn_backend=attn_backend, - ) - for index in range(n_layers) - ] - ) - self.norm = nn.LayerNorm(d_model, bias=False) - - def forward( - self, - x: torch.Tensor, - sequence_id: Optional[torch.Tensor] = None, - affine: Optional[Affine3D] = None, - affine_mask: Optional[torch.Tensor] = None, - chain_id: Optional[torch.Tensor] = None, - output_attentions: bool = False, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - tuple[torch.Tensor, ...], - Optional[tuple[torch.Tensor, ...]], - ]: - *batch_dims, _ = x.shape - if chain_id is None: - chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device) - assert affine is not None - assert affine_mask is not None - all_hidden_states = [] - all_attentions = [] - for block in self.blocks: - x, attn_weights = block( - x, - sequence_id, - affine, - affine_mask, - chain_id, - output_attentions=output_attentions, - ) - all_hidden_states.append(x) - if output_attentions and attn_weights is not None: - all_attentions.append(attn_weights) - hidden_states = tuple(all_hidden_states) - attentions = tuple(all_attentions) if output_attentions else None - return self.norm(x), x, hidden_states, attentions - - -class EncodeInputs(nn.Module): - def __init__(self, d_model: int): - super().__init__() - self.sequence_embed = nn.Embedding(64, d_model) - self.plddt_projection = nn.Linear(16, d_model) - self.structure_per_res_plddt_projection = nn.Linear(16, d_model) - self.structure_tokens_embed = nn.Embedding(4096 + 5, d_model) - self.ss8_embed = nn.Embedding(8 + 3, d_model) - self.sasa_embed = nn.Embedding(16 + 3, d_model) - self.function_embed = nn.ModuleList( - [nn.Embedding(260, d_model // 8, padding_idx=0) for _ in range(8)] - ) - self.residue_embed = nn.EmbeddingBag(1478, d_model, mode="sum", padding_idx=0) - - def forward( - self, - sequence_tokens: torch.Tensor, - structure_tokens: torch.Tensor, - average_plddt: torch.Tensor, - per_res_plddt: torch.Tensor, - ss8_tokens: torch.Tensor, - sasa_tokens: torch.Tensor, - function_tokens: torch.Tensor, - residue_annotation_tokens: torch.Tensor, - ) -> torch.Tensor: - sequence_embed = self.sequence_embed(sequence_tokens) - rbf_16_fn = functools.partial(rbf, v_min=0.0, v_max=1.0, n_bins=16) - plddt_embed = self.plddt_projection( - rbf_16_fn(average_plddt).to(self.plddt_projection.weight.dtype) - ) - structure_per_res_plddt = self.structure_per_res_plddt_projection( - rbf_16_fn(per_res_plddt).to( - self.structure_per_res_plddt_projection.weight.dtype - ) - ) - structure_embed = self.structure_tokens_embed(structure_tokens) - ss8_embed = self.ss8_embed(ss8_tokens) - sasa_embed = self.sasa_embed(sasa_tokens) - function_embed = torch.cat( - [ - embed_fn(funcs) - for embed_fn, funcs in zip( - self.function_embed, - function_tokens.unbind(-1), - ) - ], - -1, - ) - - batch_size, seq_len, num_annotations = residue_annotation_tokens.shape - residue_embed = self.residue_embed( - rearrange( - residue_annotation_tokens, - "b l n -> (b l) n", - b=batch_size, - l=seq_len, - n=num_annotations, - ) - ) - residue_embed = rearrange( - residue_embed, - "(b l) d -> b l d", - b=batch_size, - l=seq_len, - ) - - return ( - sequence_embed - + plddt_embed - + structure_per_res_plddt - + structure_embed - + ss8_embed - + sasa_embed - + function_embed - + residue_embed - ) - - -@dataclass -class ESM3CoreOutput: - sequence_logits: torch.Tensor - structure_logits: torch.Tensor - secondary_structure_logits: torch.Tensor - sasa_logits: torch.Tensor - function_logits: torch.Tensor - residue_logits: torch.Tensor - embeddings: torch.Tensor - hidden_states: tuple[torch.Tensor, ...] - attentions: Optional[tuple[torch.Tensor, ...]] = None - - -class OutputHeads(nn.Module): - def __init__(self, d_model: int): - super().__init__() - self.sequence_head = RegressionHead(d_model, 64) - self.structure_head = RegressionHead(d_model, 4096) - self.ss8_head = RegressionHead(d_model, 8 + 3) - self.sasa_head = RegressionHead(d_model, 16 + 3) - self.function_head = RegressionHead(d_model, 260 * 8) - self.residue_head = RegressionHead(d_model, 1478) - - def forward( - self, - x: torch.Tensor, - embed: torch.Tensor, - hidden_states: tuple[torch.Tensor, ...], - attentions: Optional[tuple[torch.Tensor, ...]] = None, - ) -> ESM3CoreOutput: - function_logits = self.function_head(x) - function_logits = rearrange(function_logits, "... (k v) -> ... k v", k=8) - return ESM3CoreOutput( - sequence_logits=self.sequence_head(x), - structure_logits=self.structure_head(x), - secondary_structure_logits=self.ss8_head(x), - sasa_logits=self.sasa_head(x), - function_logits=function_logits, - residue_logits=self.residue_head(x), - embeddings=embed, - hidden_states=hidden_states, - attentions=attentions, - ) - - -class ESM3Core(nn.Module): - def __init__( - self, - d_model: int, - n_heads: int, - v_heads: int, - n_layers: int, - tokenizers: FastESM3TokenizerCollection, - attn_backend: str = "sdpa", - ): - super().__init__() - self.encoder = EncodeInputs(d_model) - self.transformer = TransformerStack( - d_model, - n_heads, - v_heads, - n_layers, - mask_and_zero_frameless=True, - attn_backend=attn_backend, - ) - self.output_heads = OutputHeads(d_model) - self.tokenizers = tokenizers - - def forward( - self, - *, - sequence_tokens: Optional[torch.Tensor] = None, - structure_tokens: Optional[torch.Tensor] = None, - ss8_tokens: Optional[torch.Tensor] = None, - sasa_tokens: Optional[torch.Tensor] = None, - function_tokens: Optional[torch.Tensor] = None, - residue_annotation_tokens: Optional[torch.Tensor] = None, - average_plddt: Optional[torch.Tensor] = None, - per_res_plddt: Optional[torch.Tensor] = None, - structure_coords: Optional[torch.Tensor] = None, - chain_id: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - ) -> ESM3CoreOutput: - output_attentions = bool(output_attentions) - present_inputs = [ - sequence_tokens, - structure_tokens, - ss8_tokens, - sasa_tokens, - structure_coords, - function_tokens, - residue_annotation_tokens, - ] - try: - seq_len, device = next( - (x.shape[1], x.device) for x in present_inputs if x is not None - ) - except StopIteration: - raise ValueError("At least one of the inputs must be non-None") - - def defaults(x: Optional[torch.Tensor], token: int) -> torch.Tensor: - if x is None: - return torch.full( - (1, seq_len), - token, - dtype=torch.long, - device=device, - ) - return x - - sequence_tokens = defaults(sequence_tokens, self.tokenizers.sequence.mask_token_id) - ss8_tokens = defaults(ss8_tokens, SS8_PAD_TOKEN) - sasa_tokens = defaults(sasa_tokens, SASA_PAD_TOKEN) - average_plddt = defaults(average_plddt, 1).float() - per_res_plddt = defaults(per_res_plddt, 0).float() - chain_id = defaults(chain_id, 0) - - if residue_annotation_tokens is None: - residue_annotation_tokens = torch.full( - (1, seq_len, MAX_RESIDUE_ANNOTATIONS), - RESIDUE_PAD_TOKEN, - dtype=torch.long, - device=device, - ) - if function_tokens is None: - function_tokens = torch.full( - (1, seq_len, FUNCTION_TOKENS_DEPTH), - INTERPRO_PAD_TOKEN, - dtype=torch.long, - device=device, - ) - if structure_coords is None: - structure_coords = torch.full( - (1, seq_len, 3, 3), - float("nan"), - dtype=torch.float, - device=device, - ) - - structure_coords = structure_coords[..., :3, :] - affine, affine_mask = build_affine3d_from_coordinates(structure_coords) - - structure_tokens = defaults(structure_tokens, STRUCTURE_MASK_TOKEN) - structure_tokens = ( - structure_tokens.masked_fill(structure_tokens == -1, STRUCTURE_MASK_TOKEN) - .masked_fill(sequence_tokens == SEQUENCE_BOS_TOKEN, STRUCTURE_BOS_TOKEN) - .masked_fill(sequence_tokens == SEQUENCE_PAD_TOKEN, STRUCTURE_PAD_TOKEN) - .masked_fill(sequence_tokens == SEQUENCE_EOS_TOKEN, STRUCTURE_EOS_TOKEN) - .masked_fill( - sequence_tokens == SEQUENCE_CHAINBREAK_TOKEN, - STRUCTURE_CHAINBREAK_TOKEN, - ) - ) - - x = self.encoder( - sequence_tokens, - structure_tokens, - average_plddt, - per_res_plddt, - ss8_tokens, - sasa_tokens, - function_tokens, - residue_annotation_tokens, - ) - x, embedding, hidden_states, attentions = self.transformer( - x, - sequence_id, - affine, - affine_mask, - chain_id, - output_attentions=output_attentions, - ) - return self.output_heads( - x, - embedding, - hidden_states=hidden_states, - attentions=attentions, - ) - - -def _resolve_esm3_checkpoint_key(model_name: str) -> str: - if model_name in ESM3_OPEN_SMALL_ALIASES: - return ESM3_OPEN_SMALL - raise ValueError( - f"Unsupported ESM3 checkpoint {model_name}. " - f"Supported names: {sorted(ESM3_OPEN_SMALL_ALIASES)}" - ) - - -def parse_fasta(fasta_path: str) -> list[str]: - assert os.path.exists(fasta_path), f"FASTA file does not exist: {fasta_path}" - sequences = [] - current_seq = [] - with open(fasta_path, "r", encoding="utf-8") as handle: - for line in handle: - stripped = line.strip() - if len(stripped) == 0: - continue - if stripped.startswith(">"): - if len(current_seq) > 0: - sequences.append("".join(current_seq)) - current_seq = [] - else: - current_seq.append(stripped) - if len(current_seq) > 0: - sequences.append("".join(current_seq)) - return sequences - - -def _ensure_official_esm_on_path() -> None: - for parent in Path(__file__).resolve().parents: - candidate = parent / "official" / "esm" - if (candidate / "esm" / "models" / "esm3.py").exists(): - candidate_str = str(candidate) - if candidate_str not in sys.path: - sys.path.insert(0, candidate_str) - return - - -def _make_structure_encoder(device: Union[torch.device, str]) -> nn.Module: - _ensure_official_esm_on_path() - pretrained = importlib.import_module("esm.pretrained") - return pretrained.ESM3_structure_encoder_v0(device) - - -def _make_structure_decoder(device: Union[torch.device, str]) -> nn.Module: - _ensure_official_esm_on_path() - pretrained = importlib.import_module("esm.pretrained") - return pretrained.ESM3_structure_decoder_v0(device) - - -def _make_function_decoder(device: Union[torch.device, str]) -> nn.Module: - _ensure_official_esm_on_path() - pretrained = importlib.import_module("esm.pretrained") - return pretrained.ESM3_function_decoder_v0(device) - - -def _build_official_esm3(config: FastESM3Config) -> nn.Module: - return ESM3Core( - d_model=config.hidden_size, - n_heads=config.num_attention_heads, - v_heads=config.num_vector_heads, - n_layers=config.num_hidden_layers, - tokenizers=FastESM3TokenizerCollection(sequence=EsmSequenceTokenizer()), - attn_backend=config.attn_backend, - ) - - -class FastESM3PreTrainedModel(PreTrainedModel): - config_class = FastESM3Config - base_model_prefix = "esm3" - main_input_name = "input_ids" - supports_gradient_checkpointing = False - all_tied_weights_keys = {} - - @classmethod - def is_remote_code(cls) -> bool: - return True - - def _init_weights(self, module: nn.Module) -> None: - for parameter in module.parameters(recurse=False): - if "_is_hf_initialized" in parameter.__dict__ and parameter.__dict__["_is_hf_initialized"]: - return - - if isinstance(module, nn.Linear): - nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) - if module.bias is not None: - nn.init.zeros_(module.bias) - elif isinstance(module, nn.Embedding): - nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) - if module.padding_idx is not None: - with torch.no_grad(): - module.weight[module.padding_idx].zero_() - elif isinstance(module, nn.LayerNorm): - if module.bias is not None: - nn.init.zeros_(module.bias) - nn.init.ones_(module.weight) - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in _SUPPORTED_ATTENTION_BACKENDS, ( - f"ESM3 currently supports only {_SUPPORTED_ATTENTION_BACKENDS}; got {backend}." - ) - self.config.attn_backend = backend - resolved = resolve_attention_backend(backend) - for module in self.modules(): - if isinstance(module, MultiHeadAttention): - module.attn_backend = resolved - - @classmethod - def from_pretrained_esm( - cls, - model_name: str = ESM3_OPEN_SMALL, - device: Union[torch.device, str] = "cpu", - dtype: Optional[torch.dtype] = None, - ) -> "FastESM3Model": - key = _resolve_esm3_checkpoint_key(model_name) - spec = _ESM3_CHECKPOINT_SPECS[key] - config = FastESM3Config( - hidden_size=spec["hidden_size"], - num_attention_heads=spec["num_attention_heads"], - num_vector_heads=spec["num_vector_heads"], - num_hidden_layers=spec["num_hidden_layers"], - model_name=key, - ) - model = FastESM3Model(config) - checkpoint_root = Path( - snapshot_download( - repo_id=spec["repo_id"], - allow_patterns=["data/weights/esm3_sm_open_v1.pth"], - ) - ) - state_dict = torch.load( - checkpoint_root / "data" / "weights" / "esm3_sm_open_v1.pth", - map_location=torch.device(device), - ) - load_result = model.esm3.load_state_dict(state_dict, strict=True) - assert len(load_result.missing_keys) == 0, load_result.missing_keys - assert len(load_result.unexpected_keys) == 0, load_result.unexpected_keys - model = model.to(device) - if dtype is not None: - model = model.to(dtype=dtype) - model.eval() - return model - - -class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel): - config_class = FastESM3Config - - def __init__(self, config: FastESM3Config, **kwargs): - super().__init__(config, **kwargs) - self.tokenizer = EsmSequenceTokenizer() - self.esm3 = _build_official_esm3(config) - self.__dict__["_official_sdk_model"] = None - self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) - - @property - def device(self) -> torch.device: - return next(self.parameters()).device - - @property - def raw_model(self) -> nn.Module: - return self.esm3 - - def _get_official_sdk_model(self) -> nn.Module: - cached_model = self.__dict__["_official_sdk_model"] - if cached_model is not None: - return cached_model - _ensure_official_esm_on_path() - esm3_module = importlib.import_module("esm.models.esm3") - tokenization = importlib.import_module("esm.tokenization") - - sdk_model = esm3_module.ESM3( - d_model=self.config.hidden_size, - n_heads=self.config.num_attention_heads, - v_heads=self.config.num_vector_heads, - n_layers=self.config.num_hidden_layers, - structure_encoder_fn=_make_structure_encoder, - structure_decoder_fn=_make_structure_decoder, - function_decoder_fn=_make_function_decoder, - tokenizers=tokenization.get_esm3_model_tokenizers(self.config.model_name), - ) - load_result = sdk_model.load_state_dict(self.esm3.state_dict(), strict=True) - assert len(load_result.missing_keys) == 0, load_result.missing_keys - assert len(load_result.unexpected_keys) == 0, load_result.unexpected_keys - dtype = next(self.esm3.parameters()).dtype - sdk_model = sdk_model.to(self.device).to(dtype=dtype).eval() - self.__dict__["_official_sdk_model"] = sdk_model - return sdk_model - - def get_input_embeddings(self) -> nn.Module: - return self.esm3.encoder.sequence_embed - - def set_input_embeddings(self, value: nn.Module) -> None: - self.esm3.encoder.sequence_embed = value - - def tokenize_sequences( - self, - sequences: Union[str, list[str]], - padding: bool = True, - return_tensors: str = "pt", - device: Optional[Union[torch.device, str]] = None, - add_special_tokens: bool = True, - ) -> dict[str, torch.Tensor]: - tokenized = self.tokenizer( - sequences, - padding=padding, - return_tensors=return_tensors, - add_special_tokens=add_special_tokens, - ) - if device is None: - return tokenized - return {name: tensor.to(device) for name, tensor in tokenized.items()} - - def forward_sequence( - self, - sequences: Union[str, list[str]], - device: Optional[Union[torch.device, str]] = None, - **kwargs, - ) -> FastESM3Output: - if device is None: - device = self.device - tokenized = self.tokenize_sequences(sequences, device=device) - return self(**tokenized, **kwargs) - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - **kwargs, - ) -> torch.Tensor: - output = self( - input_ids=input_ids, - attention_mask=attention_mask, - **kwargs, - ) - if store_all_hidden_states: - assert output.hidden_states is not None, "store_all_hidden_states requires hidden states." - return torch.stack(tuple(output.hidden_states), dim=1) - if hidden_state_index == -1: - return output.last_hidden_state - assert output.hidden_states is not None, "hidden_state_index selection requires hidden states." - return output.hidden_states[hidden_state_index] - - def _pool_embeddings( - self, - embeddings: torch.Tensor, - attention_mask: torch.Tensor, - pooling_types: list[str], - ) -> torch.Tensor: - pooled = [] - mask = attention_mask.to(dtype=embeddings.dtype).unsqueeze(-1) - for pooling_type in pooling_types: - if pooling_type == "mean": - pooled.append((embeddings * mask).sum(dim=1) / mask.sum(dim=1)) - elif pooling_type == "cls": - pooled.append(embeddings[:, 0, :]) - elif pooling_type == "max": - bool_mask = attention_mask.unsqueeze(-1).bool() - pooled.append( - embeddings.masked_fill(~bool_mask, float("-inf")).max(dim=1).values - ) - else: - raise ValueError( - f"Unsupported ESM3 pooling type {pooling_type}. " - "Supported values are 'mean', 'cls', and 'max'." - ) - return torch.cat(pooled, dim=-1) - - def embed_dataset( - self, - sequences: Optional[List[str]] = None, - tokenizer: Optional[PreTrainedTokenizerFast] = None, - batch_size: int = 2, - max_len: int = 512, - truncate: bool = True, - full_embeddings: bool = False, - embed_dtype: torch.dtype = torch.float32, - pooling_types: List[str] = ["mean"], - num_workers: int = 0, - sql: bool = False, - save: bool = True, - sql_db_path: str = "embeddings.db", - save_path: str = "embeddings.pth", - fasta_path: Optional[str] = None, - padding: str = "longest", - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - **kwargs, - ) -> Dict[str, torch.Tensor]: - del num_workers, sql_db_path - assert not sql, "ESM3 embed_dataset currently supports .pth saves, not SQLite." - assert isinstance(hidden_state_index, int), "hidden_state_index must be an integer." - assert full_embeddings or not store_all_hidden_states, ( - "store_all_hidden_states=True requires full_embeddings=True." - ) - if tokenizer is None: - tokenizer = self.tokenizer - if fasta_path is not None: - fasta_sequences = parse_fasta(fasta_path) - sequences = list(sequences or []) + fasta_sequences - assert sequences is not None and len(sequences) > 0, ( - "Must provide at least one sequence via `sequences` or `fasta_path`." - ) - - unique_sequences = [] - seen_sequences = set() - for sequence in sequences: - prepared_sequence = sequence[:max_len] if truncate else sequence - if prepared_sequence not in seen_sequences: - unique_sequences.append(prepared_sequence) - seen_sequences.add(prepared_sequence) - unique_sequences = sorted(unique_sequences, key=len, reverse=True) - - embeddings_by_sequence: Dict[str, torch.Tensor] = {} - was_training = self.training - self.eval() - for batch_start in range(0, len(unique_sequences), batch_size): - batch_sequences = unique_sequences[batch_start : batch_start + batch_size] - tokenized = tokenizer( - batch_sequences, - padding=padding, - truncation=truncate, - max_length=max_len + 2, - return_tensors="pt", - ) - tokenized = { - name: tensor.to(self.device) for name, tensor in tokenized.items() - } - with torch.inference_mode(): - residue_embeddings = self._embed( - **tokenized, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - **kwargs, - ) - attention_mask = tokenized["attention_mask"] - if full_embeddings: - batch_embeddings = residue_embeddings.to(embed_dtype).cpu() - for sequence, embedding, mask in zip( - batch_sequences, - batch_embeddings, - attention_mask.cpu(), - ): - if embedding.ndim == 3: - embeddings_by_sequence[sequence] = embedding[:, mask.bool(), :] - else: - embeddings_by_sequence[sequence] = embedding[mask.bool()] - else: - pooled_embeddings = self._pool_embeddings( - residue_embeddings, - attention_mask, - pooling_types, - ) - pooled_embeddings = pooled_embeddings.to(embed_dtype).cpu() - for sequence, embedding in zip(batch_sequences, pooled_embeddings): - embeddings_by_sequence[sequence] = embedding - - if was_training: - self.train() - if save: - torch.save(embeddings_by_sequence, save_path) - return embeddings_by_sequence - - def encode(self, input): - return self._get_official_sdk_model().encode(input) - - def decode(self, input): - return self._get_official_sdk_model().decode(input) - - def generate(self, input, config): - return self._get_official_sdk_model().generate(input, config) - - def batch_generate(self, inputs, configs): - return self._get_official_sdk_model().batch_generate(inputs, configs) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.esm3] - - def forward_and_sample(self, input, sampling_configuration): - return self._get_official_sdk_model().forward_and_sample( - input, - sampling_configuration, - ) - - def logits(self, input=None, config=None, **kwargs): - if input is None: - return self.forward(**kwargs) - if isinstance(input, torch.Tensor): - return self.forward(sequence_tokens=input, **kwargs) - if config is None: - return self._get_official_sdk_model().logits(input) - return self._get_official_sdk_model().logits(input, config) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_tokens: Optional[torch.Tensor] = None, - structure_tokens: Optional[torch.Tensor] = None, - ss8_tokens: Optional[torch.Tensor] = None, - sasa_tokens: Optional[torch.Tensor] = None, - function_tokens: Optional[torch.Tensor] = None, - residue_annotation_tokens: Optional[torch.Tensor] = None, - average_plddt: Optional[torch.Tensor] = None, - per_res_plddt: Optional[torch.Tensor] = None, - structure_coords: Optional[torch.Tensor] = None, - chain_id: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, - ) -> FastESM3Output: - del output_hidden_states, return_dict, kwargs - if sequence_tokens is None: - sequence_tokens = input_ids - if sequence_id is None and attention_mask is not None: - sequence_id = attention_mask.to(dtype=torch.bool) - - output = self.esm3( - sequence_tokens=sequence_tokens, - structure_tokens=structure_tokens, - ss8_tokens=ss8_tokens, - sasa_tokens=sasa_tokens, - function_tokens=function_tokens, - residue_annotation_tokens=residue_annotation_tokens, - average_plddt=average_plddt, - per_res_plddt=per_res_plddt, - structure_coords=structure_coords, - chain_id=chain_id, - sequence_id=sequence_id, - output_attentions=output_attentions, - ) - - loss = None - if labels is not None: - loss = F.cross_entropy( - output.sequence_logits.view(-1, output.sequence_logits.shape[-1]), - labels.view(-1), - ignore_index=-100, - ) - - return FastESM3Output( - loss=loss, - logits=output.sequence_logits, - last_hidden_state=output.embeddings, - sequence_logits=output.sequence_logits, - structure_logits=output.structure_logits, - secondary_structure_logits=output.secondary_structure_logits, - sasa_logits=output.sasa_logits, - function_logits=output.function_logits, - residue_logits=output.residue_logits, - embeddings=output.embeddings, - hidden_states=output.hidden_states, - attentions=output.attentions, - ) diff --git a/fastplms/esm_plusplus/LICENSE_6B b/fastplms/esm_plusplus/LICENSE_6B deleted file mode 100644 index 1b34dc7..0000000 --- a/fastplms/esm_plusplus/LICENSE_6B +++ /dev/null @@ -1,25 +0,0 @@ -License for FastPLMs ESM++ 6B - -This derivative is built from Biohub ESMC. Biohub ESM is released under the MIT License. - -MIT License - -Copyright 2026 Chan Zuckerberg Biohub, Inc. - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/fastplms/esm_plusplus/LICENSE_large b/fastplms/esm_plusplus/LICENSE_large deleted file mode 100644 index 0a170c3..0000000 --- a/fastplms/esm_plusplus/LICENSE_large +++ /dev/null @@ -1,25 +0,0 @@ -License for FastPLMs ESM++ large - -This derivative is built from Biohub ESMC. Biohub ESM is released under the MIT License. - -MIT License - -Copyright 2026 Chan Zuckerberg Biohub, Inc. - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/fastplms/esm_plusplus/LICENSE_small b/fastplms/esm_plusplus/LICENSE_small deleted file mode 100644 index e00a7cf..0000000 --- a/fastplms/esm_plusplus/LICENSE_small +++ /dev/null @@ -1,25 +0,0 @@ -License for FastPLMs ESM++ small - -This derivative is built from Biohub ESMC. Biohub ESM is released under the MIT License. - -MIT License - -Copyright 2026 Chan Zuckerberg Biohub, Inc. - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/fastplms/esm_plusplus/README_6B.md b/fastplms/esm_plusplus/README_6B.md deleted file mode 100644 index d9656f0..0000000 --- a/fastplms/esm_plusplus/README_6B.md +++ /dev/null @@ -1,323 +0,0 @@ ---- -library_name: transformers -license: mit -tags: - - biology - - esm - - protein - - protein-language-model - - masked-language-modeling ---- - -# ESM++ 6B - -[ESM++](https://github.com/Synthyra/FastPLMs) is a Hugging Face compatible implementation of [Biohub ESMC](https://biohub.ai/esm/protein) ([license](https://github.com/Biohub/esm/blob/main/LICENSE.md)). -This checkpoint corresponds to the 6 billion parameter ESMC model released as [`biohub/ESMC-6B`](https://huggingface.co/biohub/ESMC-6B). - -This repository includes the Biohub ESM MIT license in `LICENSE`. - -The 6B model has 80 transformer layers, hidden size 2560, and 40 attention heads. It is large enough that `dtype=torch.bfloat16` or `torch.float16` plus `device_map="auto"` is usually the practical loading path. - -## Attention Backends - -`sdpa` is the default backend. Set `config.attn_backend` before loading if you want a different attention implementation. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics and stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs when `kernels` is installed. Outputs are not bitwise identical to SDPA. | -| Flex Attention | `"flex"` | Skips padding tokens via block masks. First use compiles a Triton kernel. | -| Auto | `"auto"` | Picks the best available backend: `kernels_flash`, then `flex`, then `sdpa`. | - -```python -import torch -from transformers import AutoConfig, AutoModelForMaskedLM - -config = AutoConfig.from_pretrained( - "Synthyra/ESMplusplus_6B", - trust_remote_code=True, -) -config.attn_backend = "auto" - -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_6B", - config=config, - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="auto", -) -``` - -## Masked Language Modeling - -```python -import torch -from transformers import AutoModelForMaskedLM - -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_6B", - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="auto", -) -tokenizer = model.tokenizer - -sequences = ["MPRTEIN", "MSEQWENCE"] -inputs = tokenizer(sequences, padding=True, return_tensors="pt") -inputs = inputs.to(model.device) - -with torch.no_grad(): - output = model(**inputs) - -print(output.logits.shape) -print(output.last_hidden_state.shape) -``` - -Pass `output_hidden_states=True` if you need all intermediate hidden states. - -## Experimental Test-Time Training - -TTT is disabled by default. Normal ESM++ inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the ESMC backbone with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. The 6B checkpoint is large, so start with small -`steps`, `ags`, and `batch_size` values. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 1, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -## Binder Design Regularizer - -The FastPLMs binder design tutorial uses `Synthyra/ESMplusplus_6B` as the -ESMC-style masked-LM regularizer while FastPLMs ESMFold2 experimental models -provide differentiable folding losses and final critics. The script lives at -`cookbook/tutorials/binder_design_fastplms.py` and supports local CUDA Docker -runs plus Modal deployment. - -Run the verified EGFR 128 amino acid de novo minibinder example: - -```bash -cd /home/ubuntu/FastPLMs - -sudo -n docker run --gpus all --rm \ - -v /home/ubuntu/FastPLMs:/app \ - -v /home/ubuntu/FastPLMs:/workspace \ - -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \ - -w /workspace fastplms-esmfold2 \ - python /app/cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli -``` - -The run writes `trajectory.jsonl`, `best_sequences.fasta`, `results.parquet`, -`selection.parquet`, and per-critic PDB/CIF/logit files. The verified candidate -had hero mean iPTM `0.913870`, hero min iPTM `0.904600`, and all four ESMFold2 -hero critics above `0.9`. - -Binder sequence: - -```text -SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ -``` - -See [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md) -for the full strategy, Modal backend, official pI and selection scoring, -per-critic metrics, and caveats. - -## Embed Datasets - -All FastPLMs sequence models include `embed_dataset`, which handles batching, length sorting, pooling, FASTA parsing, optional resume from existing outputs, and `.pth` or SQLite storage. - -```python -import torch -from transformers import AutoModelForMaskedLM - -model = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_6B", - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="auto", -) - -embedding_dict = model.embed_dataset( - sequences=[ - "MALWMRLLPLLALLALWGPDPAAA", - "MSEQWENCE", - "MPRTEIN", - ], - batch_size=1, - max_len=1024, - full_embeddings=False, - embed_dtype=torch.float32, - pooling_types=["mean", "cls"], - num_workers=0, - save=True, - save_path="esmplusplus_6b_embeddings.pth", -) - -print(embedding_dict["MPRTEIN"].shape) -``` - -For residue-level embeddings, set `full_embeddings=True`: - -```python -residue_embeddings = model.embed_dataset( - sequences=["MALWMRLLPLLALLALWGPDPAAA"], - batch_size=1, - max_len=1024, - full_embeddings=True, - embed_dtype=torch.float32, - save=False, -) -``` - -For very large datasets, write embeddings directly to SQLite: - -```python -model.embed_dataset( - fasta_path="proteins.fasta", - batch_size=1, - max_len=1024, - pooling_types=["mean"], - sql=True, - sql_db_path="esmplusplus_6b_embeddings.db", - save=False, -) -``` - -`embed_dataset` returns a dictionary when `sql=False`. With `sql=True`, embeddings are written to the database and loaded as needed. - -## Classification Heads - -ESM++ supports sequence-level and token-level classification through the standard Transformers auto classes. - -```python -import torch -from transformers import AutoModelForSequenceClassification - -model = AutoModelForSequenceClassification.from_pretrained( - "Synthyra/ESMplusplus_6B", - num_labels=2, - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="auto", -) - -tokenized = model.tokenizer( - ["MPRTEIN", "MSEQWENCE"], - padding=True, - return_tensors="pt", -).to(model.device) - -with torch.no_grad(): - logits = model(**tokenized).logits - -print(logits.shape) -``` - -## LoRA Fine-Tuning - -```python -from peft import LoraConfig, get_peft_model -from transformers import AutoModelForSequenceClassification - -model = AutoModelForSequenceClassification.from_pretrained( - "Synthyra/ESMplusplus_6B", - num_labels=2, - trust_remote_code=True, - dtype=torch.bfloat16, - device_map="auto", -) - -lora_config = LoraConfig( - r=8, - lora_alpha=16, - lora_dropout=0.01, - bias="none", - target_modules=[ - "layernorm_qkv.1", - "out_proj", - "query", - "key", - "value", - "dense", - ], -) - -model = get_peft_model(model, lora_config) -``` - -## Attention Maps - -Optimized attention backends do not return attention maps directly. ESM++ can compute them manually with `output_attentions=True`, but this is much slower and memory-heavy for the 6B model. - -```python -with torch.no_grad(): - output = model(**inputs, output_attentions=True) - -attentions = output.attentions -print(len(attentions)) -print(attentions[0].shape) -``` - -## Load Biohub Source Weights - -You can also load the Biohub source weights directly through FastPLMs: - -```python -from fastplms.esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM - -model = ESMplusplusForMaskedLM.from_pretrained_esm("esmc-6b") -``` - -The source repository is [`biohub/ESMC-6B`](https://huggingface.co/biohub/ESMC-6B). -The Biohub ESM license is available at https://github.com/Biohub/esm/blob/main/LICENSE.md. - -## Citation - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Hugging Face AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_6B}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@misc{candido2026language, - title = {Language Modeling Materializes a World Model of Protein Biology}, - author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan - and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub - and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo - and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson - and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian - and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil - and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth - and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius - and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz - and Sandora, McCullen and Hla, Michael and Jia, Tina Z. - and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi - and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S. - and Sercu, Tom and Rives, Alexander}, - year = {2026}, - url = {https://biohub.ai/papers/esm_protein.pdf}, - note = {Preprint} -} -``` diff --git a/fastplms/esm_plusplus/README_large.md b/fastplms/esm_plusplus/README_large.md deleted file mode 100644 index 969c36e..0000000 --- a/fastplms/esm_plusplus/README_large.md +++ /dev/null @@ -1,277 +0,0 @@ ---- -library_name: transformers -license: mit -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# ESM++ -[ESM++](https://github.com/Synthyra/FastPLMs) is a faithful implementation of [ESMC](https://biohub.ai/esm/protein) ([license](https://github.com/Biohub/esm/blob/main/LICENSE.md)) that allows for batching and standard Hugging Face compatibility without requiring the ESM Python package. -The large version corresponds to the 600 million parameter version of ESMC. - -This repository includes the Biohub ESM MIT license in `LICENSE`. - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. The backend is set via `config.attn_backend` before loading. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built, no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential, so use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask for faster variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30-120 s). Best combined with `torch.compile`. | -| Auto | `"auto"` | Picks the best available: `kernels_flash`, then `flex`, then `sdpa`. | - -```python -from transformers import AutoConfig, AutoModelForMaskedLM - -config = AutoConfig.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True) -config.attn_backend = "flex" # or "kernels_flash", "sdpa", "auto" -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', config=config, trust_remote_code=True) -``` - -`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention. - -## Binder Design Regularizer - -The FastPLMs binder design tutorial uses the ESM++ model family as the -masked-LM pseudoperplexity regularizer while FastPLMs ESMFold2 experimental -models provide differentiable folding losses and final critics. The verified -EGFR example defaults to `Synthyra/ESMplusplus_6B`; this 600M checkpoint exposes -the same `AutoModelForMaskedLM` API and can be used as a lower-memory -regularizer by editing `FastPLMsBinderDesign.lm_name` in -`cookbook/tutorials/binder_design_fastplms.py`. - -Default verified run: - -```bash -python cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir binder_design_egfr_len128_seed103 -``` - -The verified 6B-regularized result had hero mean iPTM `0.913870`, hero min iPTM -`0.904600`, and all four ESMFold2 hero critics above `0.9`. - -See [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md) -for the complete workflow, output files, metrics, and Modal/local compute -options. - -## Use with Hugging Face Transformers -```python -from transformers import AutoModelForMaskedLM -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True) -tokenizer = model.tokenizer - -sequences = ['MPRTEIN', 'MSEQWENCE'] -tokenized = tokenizer(sequences, padding=True, return_tensors='pt') - -# tokenized['labels'] = tokenized['input_ids'].clone() # correctly mask input_ids and set unmasked instances of labels to -100 for MLM training - -output = model(**tokenized) # get all hidden states with output_hidden_states=True -print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 64) -print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 1152) -print(output.loss) # language modeling loss if you passed labels -#print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple) -``` - -ESM++ also supports sequence and token level classification tasks like ESM2. Simply pass the number of labels during initialization. - -```python -from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification - -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True) -logits = model(**tokenized).logits -print(logits.shape) # (batch_size, num_labels), (2, 2) -``` - -ESM++ weights are fp32 by default. You can load them in fp16 or bf16 like this: -```python -import torch -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True, dtype=torch.float16) # or torch.bfloat16 -``` - -## Experimental test-time training - -TTT is disabled by default. Normal ESM++ inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the ESMC backbone with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -## Embed entire datasets with no new code -To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take. - -Example: -```python -embedding_dict = model.embed_dataset( - sequences=[ - 'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences - ], - batch_size=2, # adjust for your GPU memory - max_len=512, # adjust for your needs - full_embeddings=False, # if True, no pooling is performed - embed_dtype=torch.float32, # cast to what dtype you want - pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together - num_workers=0, # if you have many cpu cores, we find that num_workers = 4 is fast for large datasets - sql=False, # if True, embeddings will be stored in SQLite database - sql_db_path='embeddings.db', - save=True, # if True, embeddings will be saved as a .pth file - save_path='embeddings.pth', -) -# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql -``` - -``` -model.embed_dataset() -Args: - sequences: List of protein sequences - batch_size: Batch size for processing - max_len: Maximum sequence length - full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False) - pooling_type: Type of pooling ('mean' or 'cls') - num_workers: Number of workers for data loading, 0 for the main process - sql: Whether to store embeddings in SQLite database - will be stored in float32 - sql_db_path: Path to SQLite database - -Returns: - Dictionary mapping sequences to embeddings, or None if sql=True - -Note: - - If sql=True, embeddings can only be stored in float32 - - sql is ideal if you need to stream a very large dataset for training in real-time - - save=True is ideal if you can store the entire embedding dictionary in RAM - - sql will be used if it is True and save is True or False - - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences - - Sequences will be truncated to max_len and sorted by length in descending order for faster processing -``` - -## Fine-tuning with Hugging Face PEFT -```python -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True) -# these modules handle ESM++ and ESM2 attention layers -target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"] - -lora_config = LoraConfig( - r=8, # choose lora parameters to your liking - lora_alpha=16, - lora_dropout=0.01, - bias="none", - target_modules=target_modules, -) - -# Apply LoRA to the model -model = get_peft_model(model, lora_config) - -# Unfreeze the classifier head -for param in model.classifier.parameters(): - param.requires_grad = True -``` - -For a more thorough example of fine-tuning, check out our example script [here](https://github.com/Synthyra/FastPLMs/blob/main/fine_tuning_example.py). - - -## Returning attention maps -When `attn_backend="flex"`, Flex Attention with a pad-token block mask is used for attention calculations. Optimized attention paths do not return attention maps directly. -ESM++ has the option to ```output_attentions```, which will calculate attention manually. This is much slower, so do not use unless you need the attention maps. - -```python -output = model(**tokenized, output_attentions=True) -att = output.attentions -len(att) # 33, one for each layer, size (batch_size, num_heads, seq_len, seq_len) each -``` - -## Comparison across floating-point precision and implementations -We measured the difference of the last hidden states of the fp32 weights vs. fp16 or bf16. We find that the fp16 is closer to the fp32 outputs, so we recommend loading in fp16. -Please note that the ESM package also loads ESMC in fp32 but casts to bf16 by default, which has its share of advantages and disadvantages in inference / training - so load whichever you like for half precision. - -Average MSE for FP16: 0.00000003 - -Average MSE for BF16: 0.00000122 - -We also measured the difference between the outputs of ESM++ vs. ESMC (both in bfloat16) on 1000 random sequences to ensure compliance with the ESM package. - -Average MSE of last hidden state: 2.46e-09 - -You can load the weights from the ESM package instead of transformers by replacing .from_pretrained(...) to .from_pretrained_esm('esmc_600m') - -## Model probes -We employ linear probing techniques on various PLMs and standard datasets, similar our previous [paper](https://www.biorxiv.org/content/10.1101/2024.07.30.605924v1), to assess the intrinsic correlation between pooled hidden states and valuable properties. ESMC (and thus ESM++) perform very well. - -The plot below showcases performance normalized between the negative control (random vector embeddings) and the best performer. Classification task scores are averaged between MCC and F1 (or F1max for multilabel) and regression tasks are averaged between Spearman rho and R2. -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/uRAHYQcwkbgajylTIFbUb.png) - -## Inference speeds -We look at various ESM models and their throughput on an H100. Adding efficient batching between ESMC and ESM++ significantly improves the throughput, although ESM++ is also faster than ESMC for batch size one. ESM++ small is even faster than ESM2-35M with long sequences! The most gains will be seen with PyTorch > 2.5 on linux machines. -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/Lu6nWB9Fc-7YTql3Z1hVB.png) - -### Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Hugging Face AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_large}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@misc{candido2026language, - title = {Language Modeling Materializes a World Model of Protein Biology}, - author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan - and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub - and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo - and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson - and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian - and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil - and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth - and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius - and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz - and Sandora, McCullen and Hla, Michael and Jia, Tina Z. - and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi - and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S. - and Sercu, Tom and Rives, Alexander}, - year = {2026}, - url = {https://biohub.ai/papers/esm_protein.pdf}, - note = {Preprint} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/esm_plusplus/README_small.md b/fastplms/esm_plusplus/README_small.md deleted file mode 100644 index 30c33d3..0000000 --- a/fastplms/esm_plusplus/README_small.md +++ /dev/null @@ -1,278 +0,0 @@ ---- -library_name: transformers -license: mit -tags: [] ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# ESM++ -[ESM++](https://github.com/Synthyra/FastPLMs) is a faithful implementation of [ESMC](https://biohub.ai/esm/protein) ([license](https://github.com/Biohub/esm/blob/main/LICENSE.md)) that allows for batching and standard Hugging Face compatibility without requiring the ESM Python package. -The small version corresponds to the 300 million parameter version of ESMC. - -This repository includes the Biohub ESM MIT license in `LICENSE`. - -## Attention backends - -`sdpa` (PyTorch Scaled Dot Product Attention) is the default. The backend is set via `config.attn_backend` before loading. - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built, no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential, so use `"sdpa"` if exact numerics matter. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask for faster variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30-120 s). Best combined with `torch.compile`. | -| Auto | `"auto"` | Picks the best available: `kernels_flash`, then `flex`, then `sdpa`. | - -```python -from transformers import AutoConfig, AutoModelForMaskedLM - -config = AutoConfig.from_pretrained('Synthyra/ESMplusplus_small', trust_remote_code=True) -config.attn_backend = "flex" # or "kernels_flash", "sdpa", "auto" -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_small', config=config, trust_remote_code=True) -``` - -`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention. - -## Binder Design Regularizer - -The FastPLMs binder design tutorial uses the ESM++ model family as the -masked-LM pseudoperplexity regularizer while FastPLMs ESMFold2 experimental -models provide differentiable folding losses and final critics. The verified -EGFR example defaults to `Synthyra/ESMplusplus_6B`; this 300M checkpoint exposes -the same `AutoModelForMaskedLM` API and can be used as a lower-memory -regularizer by editing `FastPLMsBinderDesign.lm_name` in -`cookbook/tutorials/binder_design_fastplms.py`. - -Default verified run: - -```bash -python cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir binder_design_egfr_len128_seed103 -``` - -The verified 6B-regularized result had hero mean iPTM `0.913870`, hero min iPTM -`0.904600`, and all four ESMFold2 hero critics above `0.9`. - -See [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md) -for the complete workflow, output files, metrics, and Modal/local compute -options. - -## Use with Hugging Face Transformers -```python -from transformers import AutoModelForMaskedLM -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_small', trust_remote_code=True) -tokenizer = model.tokenizer - -sequences = ['MPRTEIN', 'MSEQWENCE'] -tokenized = tokenizer(sequences, padding=True, return_tensors='pt') - -# tokenized['labels'] = tokenized['input_ids'].clone() # correctly mask input_ids and set unmasked instances of labels to -100 for MLM training - -output = model(**tokenized) # get all hidden states with output_hidden_states=True -print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 64) -print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 960) -print(output.loss) # language modeling loss if you passed labels -#print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple) -``` - -ESM++ also supports sequence and token level classification tasks like ESM2. Simply pass the number of labels during initialization. - -```python -from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification - -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_small', num_labels=2, trust_remote_code=True) -logits = model(**tokenized).logits -print(logits.shape) # (batch_size, num_labels), (2, 2) -``` - -ESM++ weights are fp32 by default. You can load them in fp16 or bf16 like this: -```python -import torch -model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_small', trust_remote_code=True, dtype=torch.float16) # or torch.bfloat16 -``` - -## Experimental test-time training - -TTT is disabled by default. Normal ESM++ inference, embeddings, logits, and -`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`. -The current implementation is experimental and trains only local LoRA adapters -on the ESMC backbone with masked language modeling on the test protein. It can -help some difficult proteins, but it adds test-time compute and can degrade -already confident predictions. - -```python -metrics = model.ttt( - seq="MSTNPKPQRKTKRNT", - ttt_config={"steps": 3, "ags": 1, "batch_size": 1}, -) -model.ttt_reset() -print(metrics["losses"]) -``` - -## Embed entire datasets with no new code -To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take. - -Example: -```python -embedding_dict = model.embed_dataset( - sequences=[ - 'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences - ], - batch_size=2, # adjust for your GPU memory - max_len=512, # adjust for your needs - full_embeddings=False, # if True, no pooling is performed - embed_dtype=torch.float32, # cast to what dtype you want - pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together - num_workers=0, # if you have many cpu cores, we find that num_workers = 4 is fast for large datasets - sql=False, # if True, embeddings will be stored in SQLite database - sql_db_path='embeddings.db', - save=True, # if True, embeddings will be saved as a .pth file - save_path='embeddings.pth', -) -# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql -``` - -``` -model.embed_dataset() -Args: - sequences: List of protein sequences - batch_size: Batch size for processing - max_len: Maximum sequence length - full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False) - pooling_type: Type of pooling ('mean' or 'cls') - num_workers: Number of workers for data loading, 0 for the main process - sql: Whether to store embeddings in SQLite database - will be stored in float32 - sql_db_path: Path to SQLite database - -Returns: - Dictionary mapping sequences to embeddings, or None if sql=True - -Note: - - If sql=True, embeddings can only be stored in float32 - - sql is ideal if you need to stream a very large dataset for training in real-time - - save=True is ideal if you can store the entire embedding dictionary in RAM - - sql will be used if it is True and save is True or False - - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences - - Sequences will be truncated to max_len and sorted by length in descending order for faster processing -``` - -## Fine-tuning with Hugging Face PEFT -```python -model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_small', num_labels=2, trust_remote_code=True) -# these modules handle ESM++ and ESM2 attention layers -target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"] - -lora_config = LoraConfig( - r=8, # choose lora parameters to your liking - lora_alpha=16, - lora_dropout=0.01, - bias="none", - target_modules=target_modules, -) - -# Apply LoRA to the model -model = get_peft_model(model, lora_config) - -# Unfreeze the classifier head -for param in model.classifier.parameters(): - param.requires_grad = True -``` - -For a more thorough example of fine-tuning, check out our example script [here](https://github.com/Synthyra/FastPLMs/blob/main/fine_tuning_example.py). - - -## Returning attention maps -When `attn_backend="flex"`, Flex Attention with a pad-token block mask is used for attention calculations. Optimized attention paths do not return attention maps directly. -ESM++ has the option to ```output_attentions```, which will calculate attention manually. This is much slower, so do not use unless you need the attention maps. - -```python -output = model(**tokenized, output_attentions=True) -att = output.attentions -len(att) # 30, one for each layer, size (batch_size, num_heads, seq_len, seq_len) each -``` - -## Comparison across floating-point precision and implementations -We measured the difference of the last hidden states of the fp32 weights vs. fp16 or bf16. We find that the fp16 is closer to the fp32 outputs, so we recommend loading in fp16. -Please note that the ESM package also loads ESMC in fp32 but casts to bf16 by default, which has its share of advantages and disadvantages in inference / training - so load whichever you like for half precision. - -Average MSE FP32 vs. FP16: 0.00000003 - -Average MSE FP32 vs. BF16: 0.00000140 - -We also measured the difference between the outputs of ESM++ vs. ESMC (both in bfloat16) on 1000 random sequences to ensure compliance with the ESM package. - -Average MSE of last hidden state: 7.74e-10 - -You can load the weights from the ESM package instead of transformers by replacing .from_pretrained(...) to .from_pretrained_esm('esmc_300m') - -## Model probes -We employ linear probing techniques on various PLMs and standard datasets, similar our previous [paper](https://www.biorxiv.org/content/10.1101/2024.07.30.605924v1), to assess the intrinsic correlation between pooled hidden states and valuable properties. ESMC (and thus ESM++) perform very well. - -The plot below showcases performance normalized between the negative control (random vector embeddings) and the best performer. Classification task scores are averaged between MCC and F1 (or F1max for multilabel) and regression tasks are averaged between Spearman rho and R2. -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/2zyUZeHyOgCR_twvPF2Wy.png) - -## Inference speeds -We look at various ESM models and their throughput on an H100. Adding efficient batching between ESMC and ESM++ significantly improves the throughput, although ESM++ is also faster than ESMC for batch size one. ESM++ small is even faster than ESM2-35M with long sequences! -The most gains will be seen with PyTorch > 2.5 on linux machines. -![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/RfLRSchFivdsqJrWMh4bo.png) - -### Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Hugging Face AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@misc{candido2026language, - title = {Language Modeling Materializes a World Model of Protein Biology}, - author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan - and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub - and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo - and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson - and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian - and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil - and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth - and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius - and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz - and Sandora, McCullen and Hla, Michael and Jia, Tina Z. - and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi - and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S. - and Sercu, Tom and Rives, Alexander}, - year = {2026}, - url = {https://biohub.ai/papers/esm_protein.pdf}, - note = {Preprint} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/esm_plusplus/__init__.py b/fastplms/esm_plusplus/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/fastplms/esm_plusplus/get_weights.py b/fastplms/esm_plusplus/get_weights.py deleted file mode 100644 index 22e7862..0000000 --- a/fastplms/esm_plusplus/get_weights.py +++ /dev/null @@ -1,184 +0,0 @@ -import argparse -import os -import tempfile -from pathlib import Path -from typing import Dict, List, Optional, Tuple - -import torch -from huggingface_hub import HfApi, login -from transformers import AutoModelForMaskedLM - -from fastplms.esm_plusplus.modeling_esm_plusplus import ( - ESMplusplusConfig, - ESMplusplusForMaskedLM, - _ESMC_CHECKPOINT_SPECS, - _resolve_esmc_checkpoint_key, -) -from fastplms.weight_parity_utils import assert_state_dict_equal, assert_model_parameters_fp32 - - -MODEL_DICT: Dict[str, Tuple[str, str]] = { - "Synthyra/ESMplusplus_small": ("biohub/ESMC-300M", "README_small.md"), - "Synthyra/ESMplusplus_large": ("biohub/ESMC-600M", "README_large.md"), - "Synthyra/ESMplusplus_6B": ("biohub/ESMC-6B", "README_6B.md"), -} - -LICENSE_DICT: Dict[str, str] = { - "Synthyra/ESMplusplus_small": "LICENSE_small", - "Synthyra/ESMplusplus_large": "LICENSE_large", - "Synthyra/ESMplusplus_6B": "LICENSE_6B", -} - -HUB_AUTO_MAP = { - "AutoConfig": "modeling_esm_plusplus.ESMplusplusConfig", - "AutoModel": "modeling_esm_plusplus.ESMplusplusModel", - "AutoModelForMaskedLM": "modeling_esm_plusplus.ESMplusplusForMaskedLM", - "AutoModelForSequenceClassification": "modeling_esm_plusplus.ESMplusplusForSequenceClassification", - "AutoModelForTokenClassification": "modeling_esm_plusplus.ESMplusplusForTokenClassification", -} - - -def _build_config(esmc_model_key: str) -> ESMplusplusConfig: - spec = _ESMC_CHECKPOINT_SPECS[_resolve_esmc_checkpoint_key(esmc_model_key)] - config = ESMplusplusConfig( - hidden_size=spec["hidden_size"], - num_attention_heads=spec["num_attention_heads"], - num_hidden_layers=spec["num_hidden_layers"], - ) - config.architectures = ["ESMplusplusForMaskedLM"] - config.auto_map = HUB_AUTO_MAP - config.tie_word_embeddings = False - return config - - -def _upload_repo_files(api: HfApi, repo_id: str, script_root: str, readme_name: str) -> None: - readme_path = os.path.join(script_root, readme_name) - assert os.path.exists(readme_path), f"Missing model card: {readme_path}" - from update_HF import build_composite - - composite_code = build_composite( - "fastplms/esm_plusplus/modeling_esm_plusplus.py", - include_embedding_mixin=True, - ) - compile(composite_code, "modeling_esm_plusplus.py", "exec") - with tempfile.TemporaryDirectory() as tmpdir: - composite_path = Path(tmpdir) / "modeling_esm_plusplus.py" - composite_path.write_text(composite_code, encoding="utf-8") - api.upload_file( - path_or_fileobj=str(composite_path), - path_in_repo="modeling_esm_plusplus.py", - repo_id=repo_id, - repo_type="model", - ) - api.upload_file( - path_or_fileobj=readme_path, - path_in_repo="README.md", - repo_id=repo_id, - repo_type="model", - ) - license_path = os.path.join(script_root, LICENSE_DICT[repo_id]) - assert os.path.exists(license_path), f"Missing license: {license_path}" - api.upload_file( - path_or_fileobj=license_path, - path_in_repo="LICENSE", - repo_id=repo_id, - repo_type="model", - ) - - -def _resolve_repo_items(repo_ids: Optional[List[str]]) -> List[Tuple[str, str, str]]: - if repo_ids is None or len(repo_ids) == 0: - return [ - (repo_id, esmc_model_key, readme_name) - for repo_id, (esmc_model_key, readme_name) in MODEL_DICT.items() - ] - - selected_items: List[Tuple[str, str, str]] = [] - for repo_id in repo_ids: - assert repo_id in MODEL_DICT, ( - f"Unknown repo_id {repo_id}. " - f"Valid options: {sorted(MODEL_DICT.keys())}" - ) - esmc_model_key, readme_name = MODEL_DICT[repo_id] - selected_items.append((repo_id, esmc_model_key, readme_name)) - return selected_items - - -def _token_from_environment() -> Optional[str]: - for key in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN"): - if key in os.environ and len(os.environ[key]) > 0: - return os.environ[key] - return None - - -def _login_if_requested(args: argparse.Namespace) -> None: - token = args.hf_token - if token is None: - token = _token_from_environment() - if token is not None: - assert len(token) > 0, "HF token cannot be empty." - login(token=token) - - -if __name__ == "__main__": - # py -m fastplms.esm_plusplus.get_weights - parser = argparse.ArgumentParser() - parser.add_argument( - "--hf_token", - type=str, - default=None, - help="Deprecated. Prefer HF_TOKEN in the environment so tokens are not in shell history.", - ) - parser.add_argument("--repo_ids", nargs="*", type=str, default=None) - parser.add_argument("--dry_run", action="store_true") - parser.add_argument("--skip-weights", action="store_true") - args = parser.parse_args() - _login_if_requested(args) - api = HfApi() - - script_root = os.path.dirname(os.path.abspath(__file__)) - - for repo_id, esmc_model_key, readme_name in _resolve_repo_items(args.repo_ids): - config = _build_config(esmc_model_key) - if args.skip_weights: - if args.dry_run: - print(f"[skip-weights][dry-run] validated config metadata for {repo_id}") - continue - config.push_to_hub(repo_id) - _upload_repo_files(api, repo_id, script_root, readme_name) - print(f"[skip-weights] uploaded config for {repo_id}") - continue - - model = ESMplusplusForMaskedLM.from_pretrained_esm( - esmc_model_key, - device=torch.device("cpu"), - ).eval().cpu().to(torch.float32) - model.config.architectures = ["ESMplusplusForMaskedLM"] - model.config.auto_map = HUB_AUTO_MAP - model.config.tie_word_embeddings = False - tokenizer = model.tokenizer - - assert_model_parameters_fp32( - model=model, - model_name=f"mapped ESM++ model ({esmc_model_key})", - ) - - if args.dry_run: - print(f"[dry_run] validated ESM++ conversion for {repo_id} <- {esmc_model_key}") - continue - - tokenizer.push_to_hub(repo_id) - model.push_to_hub(repo_id) - _upload_repo_files(api, repo_id, script_root, readme_name) - downloaded_model = AutoModelForMaskedLM.from_pretrained( - repo_id, - dtype=torch.float32, - device_map="cpu", - force_download=True, - trust_remote_code=True, - ) - assert_state_dict_equal( - reference_state_dict=model.state_dict(), - candidate_state_dict=downloaded_model.state_dict(), - context=f"ESMC/ESM++ weight parity post-download ({repo_id})", - ) diff --git a/fastplms/esm_plusplus/modeling_esm_plusplus.py b/fastplms/esm_plusplus/modeling_esm_plusplus.py deleted file mode 100644 index 7966124..0000000 --- a/fastplms/esm_plusplus/modeling_esm_plusplus.py +++ /dev/null @@ -1,1511 +0,0 @@ -from __future__ import annotations -""" -ESM++ model implementation. - -ESM++ is a faithful implementation of ESMC that allows for batching and standard Huggingface compatibility -The ESM Python package is not required - -Modified from https://github.com/Biohub/esm -License: https://github.com/Biohub/esm/blob/main/LICENSE.md -""" - -import math -import os -import json -import torch -import torch.nn as nn -import torch.nn.functional as F -from dataclasses import dataclass -from functools import cache, partial -from pathlib import Path -from typing import Optional, Tuple, Union, List -from einops import rearrange, repeat -from huggingface_hub import snapshot_download -from safetensors.torch import load_file as load_safetensors_file -from tokenizers import Tokenizer -from tokenizers.models import BPE -from tokenizers.processors import TemplateProcessing -from transformers import PreTrainedModel, PreTrainedTokenizerFast, PretrainedConfig -from transformers.modeling_outputs import ModelOutput - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - kernels_flash_attention_func, - index_first_axis, index_put_first_axis, pad_input, _unpad_input, - create_block_mask, flex_attention, BlockMask, - ) - from fastplms.embedding_mixin import ( - Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator, - select_hidden_state_embeddings, - ) - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -class ESMplusplusConfig(PretrainedConfig): - """Configuration class for ESM++ model. - - Args: - vocab_size: Size of the vocabulary - hidden_size: Dimension of hidden layers - num_attention_heads: Number of attention heads - num_hidden_layers: Number of transformer layers - num_labels: Number of output labels for classification - problem_type: Type of problem - regression, single/multi label classification - """ - model_type = "ESMplusplus" - def __init__( - self, - vocab_size: int = 64, - hidden_size: int = 960, - num_attention_heads: int = 15, - num_hidden_layers: int = 30, - num_labels: int = 2, - problem_type: Optional[str] = None, - dropout: float = 0.0, - initializer_range: float = 0.02, - attn_backend: str = "sdpa", - **kwargs, - ): - super().__init__(**kwargs) - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads - self.num_hidden_layers = num_hidden_layers - self.num_labels = num_labels - self.problem_type = problem_type - self.dropout = dropout - self.initializer_range = initializer_range - self.tie_word_embeddings = False - self.attn_backend = attn_backend - - -### Rotary Embeddings -def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: - """Rotates half the hidden dims of the input.""" - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange( - torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2 - ) - - -def apply_rotary_emb_torch( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - interleaved: bool = False, - _inplace: bool = False, -) -> torch.Tensor: - """Apply rotary embeddings to input based on cos and sin.""" - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - seqlen = x.size(1) - cos = cos[:seqlen] - sin = sin[:seqlen] - cos = repeat(cos, "s d -> s 1 (2 d)") - sin = repeat(sin, "s d -> s 1 (2 d)") - return torch.cat( - [ - x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, - x[..., ro_dim:], - ], - dim=-1, - ) - - -class RotaryEmbedding(torch.nn.Module): - """Rotary position embeddings. - - Based on the paper "RoFormer: Enhanced Transformer with Rotary Position Embedding" - - Args: - dim: Dimension of the embedding - base: Base for computing angular frequencies - interleaved: Whether to use interleaved rotations - scale_base: Base for scaling - scaling_factor: Factor for scaling positions - pos_idx_in_fp32: Whether to compute position indices in fp32 - device: Computation device - """ - def __init__( - self, - dim: int, - base: float = 10000.0, - interleaved: bool = False, - scale_base: Optional[float] = None, - scaling_factor: float = 1.0, - pos_idx_in_fp32: bool = True, - device: Optional[torch.device] = None, - ): - super().__init__() - self.dim = dim - self.base = float(base) - self.pos_idx_in_fp32 = pos_idx_in_fp32 - self.interleaved = interleaved - self.scale_base = scale_base - self.scaling_factor = scaling_factor - self.device = device - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - self.reset_parameters() - - def reset_parameters(self): - """Reset the parameters of the embedding.""" - if "inv_freq" in self._buffers and isinstance(self._buffers["inv_freq"], torch.Tensor): - buffer_device = self._buffers["inv_freq"].device - else: - buffer_device = self.device - inv_freq = self._compute_inv_freq(buffer_device) - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - self.register_buffer("inv_freq", inv_freq, persistent=False) - arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32) - scale = ( - (arange + 0.4 * self.dim) / (1.4 * self.dim) - if self.scale_base is not None - else None - ) - self.register_buffer("scale", scale) - - def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor: - """Compute inverse frequency bands. - - Always computes on CPU then moves to the requested device. This matches - native Biohub ESMC, which computes inv_freq on CPU at - `__init__` and migrates via `.to(device)`. Computing directly on GPU - gives a ~3.7e-9 bit-level difference in inv_freq (fp32 transcendental - precision differs between CPU and GPU), which compounds through the 30 - attention layers to ~1e-3 mse divergence from native at - `hidden_states[-2]`. See testing/parity_debug_rotary.py. - """ - cpu_inv_freq = 1 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, device="cpu", dtype=torch.float32) - / self.dim - ) - ) - if device is not None and torch.device(device).type != "cpu": - return cpu_inv_freq.to(device) - return cpu_inv_freq - - def _update_cos_sin_cache(self, seqlen: int, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None): - """Update the cached cosine and sine values.""" - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - if self.pos_idx_in_fp32: - t = torch.arange(seqlen, device=device, dtype=torch.float32) - t /= self.scaling_factor - if self.inv_freq.dtype != torch.float32: - inv_freq = self.inv_freq.to(torch.float32) - else: - inv_freq = self.inv_freq - else: - t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) - t /= self.scaling_factor - inv_freq = self.inv_freq - freqs = torch.outer(t, inv_freq) - - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange( - seqlen, dtype=self.scale.dtype, device=self.scale.device - ) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** power.unsqueeze(-1) - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """Apply rotary embeddings to queries and keys. - - Args: - q: Query tensor of shape (batch, seqlen, nheads, headdim) - k: Key tensor of shape (batch, seqlen, nheads, headdim) - - Returns: - Tuple of rotated query and key tensors - """ - # NOTE: do NOT recompute inv_freq here if device has changed. The native - # ESMC implementation computes inv_freq once on CPU at __init__ and - # relies on PyTorch's `.to(device)` to migrate the buffer. Recomputing - # the values directly on GPU gives a ~3.7e-9 bit-level difference vs the - # CPU-computed-then-moved values due to fp32 transcendental precision, - # which compounds through 30 attention layers to ~1e-3 mse divergence - # from native at `hidden_states[-2]`. See testing/parity_debug_rotary.py. - self._update_cos_sin_cache(q.shape[1], device=q.device, dtype=q.dtype) - assert self._cos_cached is not None - assert self._sin_cached is not None - if self.scale is None: - return ( - apply_rotary_emb_torch( - q, - self._cos_cached, - self._sin_cached, - self.interleaved, - True, # inplace=True - ), - apply_rotary_emb_torch( - k, - self._cos_cached, - self._sin_cached, - self.interleaved, - True, # inplace=True - ), - ) # type: ignore - else: - assert False - - -### Feedforward Network Components -def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int: - """Compute corrected dimension for SwiGLU.""" - return int(((expansion_ratio * d_model) + 255) // 256 * 256) - - -class SwiGLU(nn.Module): - """SwiGLU activation function.""" - def __init__(self): - super(SwiGLU, self).__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return F.silu(x1) * x2 - - -def swiglu_ln_ffn(d_model: int, expansion_ratio: float) -> nn.Sequential: - """Create SwiGLU feedforward network with layer normalization.""" - return nn.Sequential( - nn.LayerNorm(d_model), - nn.Linear( - d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=False - ), - SwiGLU(), - nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=False), - ) - - -### Attention -class MultiHeadAttention(nn.Module): - """Multi-head attention with rotary embeddings and configurable backend. - - Args: - d_model: Model dimension - n_heads: Number of attention heads - attn_backend: One of "auto", "kernels_flash", "flex", "sdpa" - """ - def __init__( - self, - d_model: int, - n_heads: int, - attn_backend: str = "sdpa", - ): - super().__init__() - self.d_model = d_model - self.n_heads = n_heads - self.d_head = self.d_model // self.n_heads - self.scale = 1.0 / math.sqrt(self.d_head) - self.attn_backend = resolve_attention_backend(attn_backend) - self.layernorm_qkv = nn.Sequential( - nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=False) - ) - self.out_proj = nn.Linear(d_model, d_model, bias=False) - self.q_ln = nn.LayerNorm(d_model, bias=False) - self.k_ln = nn.LayerNorm(d_model, bias=False) - self.reshaper = partial(rearrange, pattern="b s (h d) -> b h s d", h=n_heads) - self.rotary = RotaryEmbedding(d_model // n_heads) - - def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - q = q.unflatten(-1, (self.n_heads, self.d_head)) - k = k.unflatten(-1, (self.n_heads, self.d_head)) - q, k = self.rotary(q, k) - q = q.flatten(-2, -1) - k = k.flatten(-2, -1) - return q, k - - def forward( - self, - x: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - qkv_BLD3 = self.layernorm_qkv(x) - query_BLD, key_BLD, value_BLD = torch.chunk(qkv_BLD3, 3, dim=-1) - query_BLD, key_BLD = ( - self.q_ln(query_BLD).to(query_BLD.dtype), - self.k_ln(key_BLD).to(query_BLD.dtype), - ) - query_BLD, key_BLD = self._apply_rotary(query_BLD, key_BLD) - query_BHLD, key_BHLD, value_BHLD = map(self.reshaper, (query_BLD, key_BLD, value_BLD)) - - attn_output, attn_weights, s_max = self._attn( - query_BHLD, key_BHLD, value_BHLD, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - output = self.out_proj(attn_output) - return output, attn_weights, s_max - - def _attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - if output_attentions: - return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d, output_s_max) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - attn_output, attn_weights = self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) - elif self.attn_backend == AttentionBackend.FLEX: - attn_output, attn_weights = self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) - elif self.attn_backend == AttentionBackend.SDPA: - attn_output, attn_weights = self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - @torch.no_grad() - def _compute_s_max(self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor) -> List[torch.Tensor]: - q_norm = torch.linalg.vector_norm(query_BHLD, dim=-1) - k_norm = torch.linalg.vector_norm(key_BHLD, dim=-1) - s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values * self.scale - return [s_max_bound[h] for h in range(self.n_heads)] - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, Optional[List[torch.Tensor]]]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-2, -1)) * self.scale - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") - s_max = self._compute_s_max(query_BHLD, key_BHLD) if output_s_max else None - return attn_output, attn_weights, s_max - - def _kernels_flash_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - query_BLHD = query_BHLD.transpose(1, 2).contiguous() - key_BLHD = key_BHLD.transpose(1, 2).contiguous() - value_BLHD = value_BHLD.transpose(1, 2).contiguous() - attn_output = kernels_flash_attention_func( - query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, - attention_mask_2d=attention_mask_2d, causal=False, - ) - return rearrange(attn_output, "b s h d -> b s (h d)"), None - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask] = None, - ) -> Tuple[torch.Tensor, None]: - assert flex_attention is not None, "Flex attention is not available in this environment." - fn = _get_flex_attention_fn() - context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=self.scale) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, attn_mask=attention_mask_4d, scale=self.scale, - ) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - -### Regression Head -def RegressionHead(d_model: int, output_dim: int, hidden_dim: Optional[int] = None) -> nn.Module: - """Create a regression head with optional hidden dimension. - - Args: - d_model: Input dimension - output_dim: Output dimension - hidden_dim: Optional hidden dimension (defaults to d_model) - """ - hidden_dim = hidden_dim if hidden_dim is not None else d_model - return nn.Sequential( - nn.Linear(d_model, hidden_dim), - nn.GELU(), - nn.LayerNorm(hidden_dim), - nn.Linear(hidden_dim, output_dim), - ) - - -### Transformer Block -class UnifiedTransformerBlock(nn.Module): - """Transformer block with attention and feedforward layers.""" - def __init__( - self, - d_model: int, - n_heads: int, - residue_scaling_factor: float = 1, - expansion_ratio: float = 8 / 3, - dropout: float = 0.0, - attn_backend: str = "sdpa", - ): - super().__init__() - self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, attn_backend=attn_backend) - self.ffn = swiglu_ln_ffn(d_model, expansion_ratio) - self.scaling_factor = residue_scaling_factor - self.dropout = nn.Dropout(dropout) - - def forward( - self, - x: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - output_s_max: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[torch.Tensor]]]: - attn_output, attn_weights, s_max = self.attn( - x, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - x = x + self.dropout(attn_output) / self.scaling_factor - x = x + self.dropout(self.ffn(x)) / self.scaling_factor - return x, attn_weights, s_max - - -### Model Outputs -@dataclass -class TransformerOutput(ModelOutput): - """Output type for transformer encoder.""" - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor]] = None - attentions: Optional[Tuple[torch.Tensor]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -@dataclass -class ESMplusplusOutput(ModelOutput): - """Output type for ESM++ models.""" - loss: Optional[torch.Tensor] = None - logits: Optional[torch.Tensor] = None - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor]] = None - attentions: Optional[Tuple[torch.Tensor]] = None - s_max: Optional[Tuple[List[torch.Tensor], ...]] = None - - -### Transformer Stack -class TransformerStack(nn.Module): - """Stack of transformer blocks.""" - def __init__( - self, - d_model: int, - n_heads: int, - n_layers: int, - dropout: float = 0.0, - attn_backend: str = "sdpa", - ): - super().__init__() - self.attention_backend = resolve_attention_backend(attn_backend) - self.blocks = nn.ModuleList( - [ - UnifiedTransformerBlock( - d_model, - n_heads, - residue_scaling_factor=math.sqrt(n_layers / 36), - dropout=dropout, - attn_backend=attn_backend, - ) - for i in range(n_layers) - ] - ) - self.norm = nn.LayerNorm(d_model, bias=False) - self.gradient_checkpointing = False - - @property - def attn_backend(self) -> AttentionBackend: - return self.attention_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - resolved = resolve_attention_backend(backend) - self.attention_backend = resolved - for block in self.blocks: - block.attn.attn_backend = resolved - - def forward( - self, - x: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = False, - output_attentions: Optional[bool] = False, - output_s_max: Optional[bool] = False, - esmfold2_hidden_states: bool = False, - ) -> TransformerOutput: - hidden_states = () if output_hidden_states else None - attentions = () if output_attentions else None - full_s_max = () if output_s_max else None - - if sequence_id is None: - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=x.shape[0], - seq_len=x.shape[1], - device=x.device, - attention_mask=attention_mask, - ) - else: - attention_mask_2d, attention_mask_4d, flex_block_mask = self._sequence_id_attention_masks( - sequence_id=sequence_id, - batch_size=x.shape[0], - seq_len=x.shape[1], - device=x.device, - ) - - if output_hidden_states and esmfold2_hidden_states: - assert hidden_states is not None - hidden_states += (x,) - - for block_index, block in enumerate(self.blocks): - if self.gradient_checkpointing and self.training: - x, attn_weights, s_max = self._gradient_checkpointing_func( - block.__call__, - x=x, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - else: - x, attn_weights, s_max = block( - x=x, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - output_s_max=output_s_max, - ) - - if attentions is not None: - attentions += (attn_weights,) - if output_hidden_states: - assert hidden_states is not None - if not esmfold2_hidden_states or block_index < len(self.blocks) - 1: - hidden_states += (x,) - if full_s_max is not None: - full_s_max += (s_max,) - - last_hidden_state = self.norm(x) - if output_hidden_states: - hidden_states += (last_hidden_state,) - - return TransformerOutput( - last_hidden_state=last_hidden_state, - hidden_states=hidden_states, - attentions=attentions, - s_max=full_s_max, - ) - - def _sequence_id_attention_masks( - self, - sequence_id: torch.Tensor, - batch_size: int, - seq_len: int, - device: torch.device, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]: - assert sequence_id.shape == (batch_size, seq_len), ( - f"sequence_id shape must be {(batch_size, seq_len)}, got {tuple(sequence_id.shape)}" - ) - if sequence_id.dtype == torch.bool: - attention_mask_2d = sequence_id - attention_mask_4d = ( - attention_mask_2d[:, None, :, None] - & attention_mask_2d[:, None, None, :] - ) - else: - attention_mask_2d = sequence_id != -1 - attention_mask_4d = ( - attention_mask_2d[:, None, :, None] - & attention_mask_2d[:, None, None, :] - & (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze(1) - ) - - if self.attention_backend == AttentionBackend.KERNELS_FLASH: - assert sequence_id.dtype == torch.bool, ( - "ESM++ kernels_flash only supports boolean sequence_id padding masks. " - "Use sdpa or flex for chain-aware integer sequence_id masks." - ) - return attention_mask_2d, attention_mask_4d, None - - if self.attention_backend == AttentionBackend.FLEX: - assert create_block_mask is not None, ( - "Flex attention backend requested but torch.create_block_mask is unavailable." - ) - - if sequence_id.dtype == torch.bool: - - def mask_mod(batch_idx, head_idx, q_idx, kv_idx): - return ( - sequence_id[batch_idx, q_idx] - & sequence_id[batch_idx, kv_idx] - ) - - else: - - def mask_mod(batch_idx, head_idx, q_idx, kv_idx): - q_id = sequence_id[batch_idx, q_idx] - kv_id = sequence_id[batch_idx, kv_idx] - return (q_id != -1) & (q_id == kv_id) - - flex_block_mask = create_block_mask( - mask_mod, - batch_size, - 1, - seq_len, - seq_len, - device=device, - ) - return attention_mask_2d, attention_mask_4d, flex_block_mask - - return attention_mask_2d, attention_mask_4d, None - - -class PreTrainedESMplusplusModel(PreTrainedModel): - """ - init weights for ESM++ models - """ - config_class = ESMplusplusConfig - base_model_prefix = "esm++" - supports_gradient_checkpointing = True - all_tied_weights_keys = {} - - @classmethod - def is_remote_code(cls) -> bool: - # Prevent post-load reinitialization of tensors already loaded from checkpoints. - return True - - def _init_weights(self, module): - """Initialize the weights""" - # HF from_pretrained marks loaded parameters with `_is_hf_initialized`. - # Skip this module if any local parameter is already marked as loaded. - for parameter in module.parameters(recurse=False): - if "_is_hf_initialized" in parameter.__dict__ and parameter.__dict__["_is_hf_initialized"]: - return - - if isinstance(module, nn.Linear): - nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) - if module.bias is not None: - nn.init.zeros_(module.bias) - elif isinstance(module, nn.Embedding): - nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) - if module.padding_idx is not None: - with torch.no_grad(): - module.weight[module.padding_idx].zero_() - elif isinstance(module, nn.LayerNorm): - if module.bias is not None: - nn.init.zeros_(module.bias) - nn.init.ones_(module.weight) - - @property - def attn_backend(self) -> str: - return self.config.attn_backend - - @attn_backend.setter - def attn_backend(self, backend: str) -> None: - assert backend in VALID_ATTENTION_BACKENDS, f"Unsupported attn_backend: {backend}. Expected one of {VALID_ATTENTION_BACKENDS}." - self.config.attn_backend = backend - for module in self.modules(): - if isinstance(module, TransformerStack): - module.attn_backend = backend - - def _reset_rotary_embeddings(self): - """Refresh non-persistent rotary buffers after checkpoint loading.""" - for module in self.modules(): - if isinstance(module, RotaryEmbedding): - module.reset_parameters() - - @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): - output_loading_info = bool(kwargs["output_loading_info"]) if "output_loading_info" in kwargs else False - loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) - if output_loading_info: - model, loading_info = loaded - model._reset_rotary_embeddings() - return model, loading_info - loaded._reset_rotary_embeddings() - return loaded - - @classmethod - def from_pretrained_esm( - cls, - model_name: str, - device: Union[torch.device, str] = "cpu", - ): - """Load a pretrained ESM++ model.""" - key = _resolve_esmc_checkpoint_key(model_name) - if key == "esmc-300": - return ESMplusplus_300M(device=device) - if key == "esmc-600": - return ESMplusplus_600M(device=device) - if key == "esmc-6b": - return ESMplusplus_6B(device=device) - raise ValueError(f"Invalid model name: {model_name}") - - -### ESM++ Models -class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin): - """ - ESM++ model. transformer model with no heads - """ - config_class = ESMplusplusConfig - def __init__(self, config: ESMplusplusConfig, **kwargs): - PreTrainedESMplusplusModel.__init__(self, config, **kwargs) - self.config = config - self.vocab_size = config.vocab_size - self.embed = nn.Embedding(self.vocab_size, config.hidden_size) - self.transformer = TransformerStack( - d_model=config.hidden_size, - n_heads=config.num_attention_heads, - n_layers=config.num_hidden_layers, - dropout=config.dropout, - attn_backend=config.attn_backend, - ) - self.tokenizer = EsmSequenceTokenizer() - self.init_weights() - - def get_input_embeddings(self): - return self.embed - - def set_input_embeddings(self, value): - self.embed = value - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - x = self.embed(input_ids) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - output = self.transformer( - x=x, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - output.last_hidden_state, - output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - esmfold2_hidden_states: bool = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> ESMplusplusOutput: - assert input_ids is not None or inputs_embeds is not None, "You have to specify either input_ids or inputs_embeds" - assert not (input_ids is not None and inputs_embeds is not None), "You cannot specify both input_ids and inputs_embeds at the same time" - - if inputs_embeds is None: - x = self.embed(input_ids) - else: - x = inputs_embeds - - transformer_output = self.transformer( - x=x, - attention_mask=attention_mask, - sequence_id=sequence_id, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - esmfold2_hidden_states=esmfold2_hidden_states, - ) - return ESMplusplusOutput( - last_hidden_state=transformer_output.last_hidden_state, - hidden_states=transformer_output.hidden_states, - attentions=transformer_output.attentions, - s_max=transformer_output.s_max, - ) - -class ESMplusplusForMaskedLM(FastPLMTestTimeTrainingMixin, PreTrainedESMplusplusModel, EmbeddingMixin): - """ - ESM++ model for masked language modeling. - Implements the base ESM++ architecture with a masked language modeling head. - """ - config_class = ESMplusplusConfig - def __init__(self, config: ESMplusplusConfig, **kwargs): - PreTrainedESMplusplusModel.__init__(self, config, **kwargs) - self.config = config - self.vocab_size = config.vocab_size - self.embed = nn.Embedding(self.vocab_size, config.hidden_size) - self.transformer = TransformerStack( - d_model=config.hidden_size, - n_heads=config.num_attention_heads, - n_layers=config.num_hidden_layers, - dropout=config.dropout, - attn_backend=config.attn_backend, - ) - self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size) - self.ce_loss = nn.CrossEntropyLoss() - self.tokenizer = EsmSequenceTokenizer() - self.init_weights() - self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) - - def get_input_embeddings(self): - return self.embed - - def set_input_embeddings(self, value): - self.embed = value - - def get_output_embeddings(self): - return self.sequence_head[-1] - - def set_output_embeddings(self, new_embeddings): - self.sequence_head[-1] = new_embeddings - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - x = self.embed(input_ids) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - output = self.transformer( - x=x, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - output.last_hidden_state, - output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.transformer] - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - esmfold2_hidden_states: bool = False, - return_dict: Optional[bool] = None, - compute_logits: bool = True, - **kwargs, - ) -> ESMplusplusOutput: - if inputs_embeds is None: - x = self.embed(input_ids) - else: - x = inputs_embeds - - output = self.transformer( - x=x, - attention_mask=attention_mask, - sequence_id=sequence_id, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - output_s_max=output_s_max, - esmfold2_hidden_states=esmfold2_hidden_states, - ) - - last_hidden_state = output.last_hidden_state - logits = self.sequence_head(last_hidden_state) if compute_logits else None - loss = None - if labels is not None: - assert logits is not None, "labels require compute_logits=True." - loss = self.ce_loss(logits.view(-1, self.vocab_size), labels.view(-1)) - - return ESMplusplusOutput( - loss=loss, - logits=logits, - last_hidden_state=last_hidden_state, - hidden_states=output.hidden_states, - attentions=output.attentions, - s_max=output.s_max, - ) - - -class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixin): - """ - ESM++ model for sequence classification. - Extends the base ESM++ model with a classification head. - """ - def __init__(self, config: ESMplusplusConfig, **kwargs): - ESMplusplusForMaskedLM.__init__(self, config, **kwargs) - self.config = config - self.num_labels = config.num_labels - self.classifier = RegressionHead(config.hidden_size * 2, config.num_labels, config.hidden_size * 4) - # Large intermediate projections help with sequence classification tasks (*4) - self.mse = nn.MSELoss() - self.ce = nn.CrossEntropyLoss() - self.bce = nn.BCEWithLogitsLoss() - # if kwargs has pooling_types, use them, otherwise use ['cls', 'mean'] - if 'pooling_types' in kwargs and isinstance(kwargs['pooling_types'], List[str]) and len(kwargs['pooling_types']) > 0: - pooling_types = kwargs['pooling_types'] - else: - pooling_types = ['mean', 'var'] - self.pooler = Pooler(pooling_types) - self.init_weights() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - x = self.embed(input_ids) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - output = self.transformer( - x=x, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - output.last_hidden_state, - output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> ESMplusplusOutput: - output = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - sequence_id=sequence_id, - inputs_embeds=inputs_embeds, - labels=None, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - - last_hidden_state = output.last_hidden_state - features = self.pooler(last_hidden_state, attention_mask) - logits = self.classifier(features) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - if self.num_labels == 1: - loss = self.mse(logits.flatten(), labels.flatten()) - else: - loss = self.mse(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss = self.bce(logits, labels) - - return ESMplusplusOutput( - loss=loss, - logits=logits, - last_hidden_state=last_hidden_state, - hidden_states=output.hidden_states, - attentions=output.attentions, - s_max=output.s_max, - ) - - -class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin): - """ - ESM++ model for token classification. - Extends the base ESM++ model with a token classification head. - """ - def __init__(self, config: ESMplusplusConfig, **kwargs): - ESMplusplusForMaskedLM.__init__(self, config, **kwargs) - self.config = config - self.num_labels = config.num_labels - self.classifier = RegressionHead(config.hidden_size, config.num_labels, config.hidden_size * 4) - # Large intermediate projections help with sequence classification tasks (*4) - self.loss_fct = nn.CrossEntropyLoss() - self.init_weights() - - def _embed( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - ) -> torch.Tensor: - x = self.embed(input_ids) - output_hidden_states = store_all_hidden_states or hidden_state_index != -1 - output = self.transformer( - x, - attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=False, - ) - return select_hidden_state_embeddings( - output.last_hidden_state, - output.hidden_states, - hidden_state_index=hidden_state_index, - store_all_hidden_states=store_all_hidden_states, - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_s_max: Optional[bool] = False, - return_dict: Optional[bool] = None, - **kwargs, - ) -> ESMplusplusOutput: - output = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - sequence_id=sequence_id, - inputs_embeds=inputs_embeds, - labels=None, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_s_max=output_s_max, - ) - - last_hidden_state = output.last_hidden_state - logits = self.classifier(last_hidden_state) - loss = None - if labels is not None: - loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - return ESMplusplusOutput( - loss=loss, - logits=logits, - last_hidden_state=last_hidden_state, - hidden_states=output.hidden_states, - attentions=output.attentions, - s_max=output.s_max, - ) - - -### Loading from Biohub -_ESMC_CHECKPOINT_SPECS = { - "esmc-300": { - "repo_id": "biohub/ESMC-300M", - "hidden_size": 960, - "num_attention_heads": 15, - "num_hidden_layers": 30, - }, - "esmc-600": { - "repo_id": "biohub/ESMC-600M", - "hidden_size": 1152, - "num_attention_heads": 18, - "num_hidden_layers": 36, - }, - "esmc-6b": { - "repo_id": "biohub/ESMC-6B", - "hidden_size": 2560, - "num_attention_heads": 40, - "num_hidden_layers": 80, - }, -} - - -def _resolve_esmc_checkpoint_key(model: str) -> str: - normalized = model.lower().replace("_", "-") - if "300" in normalized: - return "esmc-300" - if "600" in normalized: - return "esmc-600" - if "6b" in normalized: - return "esmc-6b" - raise ValueError(f"{model=} is an invalid ESMC model name.") - - -@staticmethod -@cache -def data_root(model: str): - if "INFRA_PROVIDER" in os.environ: - return Path("") - key = _resolve_esmc_checkpoint_key(model) - return Path(snapshot_download(repo_id=_ESMC_CHECKPOINT_SPECS[key]["repo_id"])) - - -def get_esmc_checkpoint_path(model: str) -> Path: - key = _resolve_esmc_checkpoint_key(model) - spec = _ESMC_CHECKPOINT_SPECS[key] - if "weights_relpath" in spec: - return data_root(key) / spec["weights_relpath"] - checkpoint_dir = data_root(key) - if (checkpoint_dir / "model.safetensors").exists(): - return checkpoint_dir / "model.safetensors" - if (checkpoint_dir / "model.safetensors.index.json").exists(): - return checkpoint_dir / "model.safetensors.index.json" - raise FileNotFoundError(f"No ESMC checkpoint found under {checkpoint_dir}.") - - -def _normalize_esmc_state_key(key: str) -> Optional[str]: - if key.endswith("._extra_state"): - return None - if key.startswith("esmc."): - key = key[len("esmc."):] - if key.startswith("lm_head."): - key = f"sequence_head.{key[len('lm_head.'):]}" - replacements = ( - (".attn.layernorm_qkv.layer_norm_bias", ".attn.layernorm_qkv.0.bias"), - (".attn.layernorm_qkv.layer_norm_weight", ".attn.layernorm_qkv.0.weight"), - (".attn.layernorm_qkv.weight", ".attn.layernorm_qkv.1.weight"), - (".ffn.layer_norm_bias", ".ffn.0.bias"), - (".ffn.layer_norm_weight", ".ffn.0.weight"), - (".ffn.fc1_weight", ".ffn.1.weight"), - (".ffn.fc2_weight", ".ffn.3.weight"), - ) - for old, new in replacements: - key = key.replace(old, new) - return key - - -def _normalize_esmc_state_dict(state_dict: dict) -> dict: - normalized = {} - for key, tensor in state_dict.items(): - normalized_key = _normalize_esmc_state_key(key) - if normalized_key is None: - continue - normalized[normalized_key] = tensor - return normalized - - -def _safetensors_checkpoint_files(checkpoint_path: Path) -> List[Path]: - if checkpoint_path.name == "model.safetensors": - return [checkpoint_path] - with checkpoint_path.open("r", encoding="utf-8") as f: - index = json.load(f) - return [ - checkpoint_path.parent / filename - for filename in sorted(set(index["weight_map"].values())) - ] - - -def _load_safetensors_state_dict( - model_obj: ESMplusplusForMaskedLM, - checkpoint_path: Path, - device: Union[torch.device, str], -) -> None: - expected_keys = set(model_obj.state_dict().keys()) - loaded_keys = set() - device_string = str(torch.device(device)) - for shard_path in _safetensors_checkpoint_files(checkpoint_path): - shard_state_dict = load_safetensors_file(shard_path, device=device_string) - normalized = _normalize_esmc_state_dict(shard_state_dict) - unexpected = set(normalized.keys()) - expected_keys - assert len(unexpected) == 0, ( - f"Unexpected ESMC checkpoint keys in {shard_path.name}: " - f"{sorted(unexpected)[:10]}" - ) - model_obj.load_state_dict(normalized, strict=False) - loaded_keys.update(normalized.keys()) - - missing = expected_keys - loaded_keys - assert len(missing) == 0, ( - f"ESMC checkpoint did not provide all expected keys: {sorted(missing)[:10]}" - ) - - -def _load_esmc_checkpoint_model( - config: ESMplusplusConfig, - model: str, - device: Union[torch.device, str] = "cpu", -) -> ESMplusplusForMaskedLM: - key = _resolve_esmc_checkpoint_key(model) - spec = _ESMC_CHECKPOINT_SPECS[key] - assert config.hidden_size == spec["hidden_size"], ( - f"ESMC loader expected hidden_size={spec['hidden_size']} for {key}, " - f"but got {config.hidden_size}." - ) - assert config.num_attention_heads == spec["num_attention_heads"], ( - f"ESMC loader expected num_attention_heads={spec['num_attention_heads']} for {key}, " - f"but got {config.num_attention_heads}." - ) - assert config.num_hidden_layers == spec["num_hidden_layers"], ( - f"ESMC loader expected num_hidden_layers={spec['num_hidden_layers']} for {key}, " - f"but got {config.num_hidden_layers}." - ) - with torch.device(device): - model_obj = ESMplusplusForMaskedLM(config) - checkpoint_path = get_esmc_checkpoint_path(key) - if checkpoint_path.suffix == ".safetensors" or checkpoint_path.name == "model.safetensors.index.json": - _load_safetensors_state_dict( - model_obj=model_obj, - checkpoint_path=checkpoint_path, - device=device, - ) - else: - state_dict = torch.load(checkpoint_path, map_location=device) - model_obj.load_state_dict(_normalize_esmc_state_dict(state_dict)) - return model_obj - - -def ESMplusplus_300M(device: Union[torch.device, str] = "cpu"): - config = ESMplusplusConfig( - hidden_size=960, - num_attention_heads=15, - num_hidden_layers=30, - ) - return _load_esmc_checkpoint_model(config=config, model="esmc-300", device=device) - - -def ESMplusplus_600M(device: Union[torch.device, str] = "cpu"): - config = ESMplusplusConfig( - hidden_size=1152, - num_attention_heads=18, - num_hidden_layers=36, - ) - return _load_esmc_checkpoint_model(config=config, model="esmc-600", device=device) - - -def ESMplusplus_6B(device: Union[torch.device, str] = "cpu"): - config = ESMplusplusConfig( - hidden_size=2560, - num_attention_heads=40, - num_hidden_layers=80, - ) - return _load_esmc_checkpoint_model(config=config, model="esmc-6b", device=device) - - -### Tokenization -SEQUENCE_VOCAB = [ - "", "", "", "", - "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", - "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", - "O", ".", "-", "|", - "", -] - -class EsmSequenceTokenizer(PreTrainedTokenizerFast): - model_input_names = ["input_ids", "attention_mask"] - - def __init__( - self, - unk_token="", - cls_token="", - pad_token="", - mask_token="", - eos_token="", - chain_break_token="|", - **kwargs, - ): - all_tokens = SEQUENCE_VOCAB - token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)} - - # a character-level tokenizer is the same as BPE with no token merges - bpe = BPE(token_to_id, merges=[], unk_token=unk_token) - tokenizer = Tokenizer(bpe) - special_tokens = [ - cls_token, - pad_token, - mask_token, - eos_token, - chain_break_token, - ] - self.cb_token = chain_break_token - additional_special_tokens = [chain_break_token] - - tokenizer.add_special_tokens(special_tokens) - - # This is where we configure the automatic addition of special tokens when we call - # tokenizer(text, add_special_tokens=True). Note that you can also configure how two - # sequences are merged if you want. - tokenizer.post_processor = TemplateProcessing( # type: ignore - single=" $A ", - pair=":0 $A:0 :0 $B:1 :1", - special_tokens=[ - ("", tokenizer.token_to_id("")), - ("", tokenizer.token_to_id("")), - ], - ) - super().__init__( - tokenizer_object=tokenizer, - unk_token=unk_token, - cls_token=cls_token, - pad_token=pad_token, - mask_token=mask_token, - eos_token=eos_token, - additional_special_tokens=additional_special_tokens, - **kwargs, - ) - - # These are a footgun, we never use the `bos` token anywhere so we're just overriding it here. - @property - def bos_token(self): - return self.cls_token - - @property - def bos_token_id(self): - return self.cls_token_id - - @property - def chain_break_token(self): - return self.cb_token - - @property - def chain_break_token_id(self): - return self.convert_tokens_to_ids(self.chain_break_token) - - @property - def all_token_ids(self): - return list(range(self.vocab_size)) - - @property - def special_token_ids(self): - return self.all_special_ids - - -if __name__ == "__main__": - import random - - import torch - - from torch import Tensor - - def print_tensor_shapes(prefix: str, obj): - if isinstance(obj, Tensor): - print(f"{prefix}{obj.shape}") - elif isinstance(obj, dict): - for name, value in obj.items(): - print_tensor_shapes(f"{prefix}{name}.", value) - elif isinstance(obj, list): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif isinstance(obj, tuple): - for idx, value in enumerate(obj): - print_tensor_shapes(f"{prefix}[{idx}].", value) - elif hasattr(obj, "__dict__"): - for name, value in vars(obj).items(): - if name.startswith("_"): - continue - print_tensor_shapes(f"{prefix}{name}.", value) - else: - print(f"{prefix}{type(obj)}") - - random.seed(0) - torch.manual_seed(0) - - tokenizer = EsmSequenceTokenizer() - num_attention_heads = random.choice([2, 4]) - config = ESMplusplusConfig( - vocab_size=tokenizer.vocab_size, - hidden_size=16 * num_attention_heads, - num_attention_heads=num_attention_heads, - num_hidden_layers=random.choice([1, 2]), - num_labels=2, - dropout=0.0, - ) - - batch = tokenizer(["ACDEFG", "MKTW"], return_tensors="pt", padding=True) - batch["labels"] = batch["input_ids"].clone() - model = ESMplusplusForMaskedLM(config=config).eval() - - with torch.no_grad(): - output = model(**batch, return_dict=True) - - print("Batch shape:") - print_tensor_shapes("", batch) - print("Output shape:") - print_tensor_shapes("", output) diff --git a/fastplms/esmfold/README.md b/fastplms/esmfold/README.md deleted file mode 100644 index 9d0abe2..0000000 --- a/fastplms/esmfold/README.md +++ /dev/null @@ -1,234 +0,0 @@ ---- -library_name: transformers -tags: - - protein - - structure-prediction - - esmfold - - test-time-training ---- - -# NOTE -The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git) - -# FastESMFold - -FastESMFold is a self-contained, HuggingFace-compatible reimplementation of ESMFold with optional experimental **Test-Time Training (TTT)** and multi-backend attention (SDPA, Flash, Flex). - -No dependency on `fair-esm`, `proteinttt`, or `openfold`. Just `transformers`, `torch`, and `einops`. - -## Why Test-Time Training? - -Protein language models like ESM2 are trained on millions of sequences, but at inference time they process each new protein in a single forward pass with no adaptation. This is a missed opportunity: the input sequence itself contains structural signal that the model could learn from. - -**Test-Time Training (TTT)** adapts the model to each individual protein before predicting its structure. The idea is simple: before folding, we briefly train the ESM2 backbone on the input sequence using masked language modeling (the same objective it was pretrained with). This forces the model to "study" the specific sequence, strengthening its internal representation of that protein's structural features. - -TTT is disabled by default. Standard `fold_protein(...)`, `infer(...)`, and -`state_dict()` behavior are unchanged unless you explicitly pass `ttt=True` or -call `fold_protein_ttt(...)`. - -The adaptation uses **LoRA** (Low-Rank Adaptation) for efficiency: only small -adapter weights are trained (~4.4M parameters out of 3.5B), and the base model -is restored after each prediction. This takes 20-45 seconds per sequence on an -A10G GPU. It can improve structure prediction quality on difficult targets -where standard ESMFold produces low-confidence predictions, but it is -experimental and can degrade predictions that already have high confidence. - -**When is TTT most useful?** -- Sequences with low baseline pLDDT (< 0.5): TTT can improve pLDDT by 10-30+ points -- Novel proteins with limited homology in training data -- Disordered or multi-domain proteins where ESMFold struggles - -**When is TTT unnecessary?** -- Sequences that already fold well (baseline pLDDT > 0.7): TTT rarely helps and may slightly degrade predictions -- High-throughput screening where speed matters more than accuracy - -## Key Features - -- **Standard ESMFold**: Full ESMFold v1 structure prediction, loadable via `AutoModel` -- **Optional experimental TTT**: Enable test-time training for difficult sequences with explicit `ttt=True` -- **Best structure selection**: When TTT is enabled, folds after each step and returns the structure with the highest pLDDT -- **FastESM2 attention**: SDPA/Flash/Flex backends for the 3B ESM2 backbone -- **Self-contained LoRA**: lora_diffusion-compatible implementation (no peft dependency) -- **3.5B parameters**: Full ESMFold v1 architecture (ESM2-3B backbone + folding trunk) - -## Use with transformers - -### Standard structure prediction (no TTT) - -```python -import torch -from transformers import AutoModel - -model = AutoModel.from_pretrained( - "Synthyra/FastESMFold", - trust_remote_code=True, - dtype=torch.float32, -).cuda().eval() - -# Standard fold (no TTT) -with torch.no_grad(): - output = model.infer("MKTLLILAVVAAALA...") -pdb_strings = model.output_to_pdb(output) -plddt = output["plddt"].mean().item() -print(f"pLDDT: {plddt:.3f}") -``` - -### Structure prediction with experimental TTT - -TTT adapts the ESM2 backbone to a specific input sequence via masked language -modeling before folding. It can improve pLDDT on difficult sequences, but it is -experimental, adds test-time compute, and should not be assumed to improve every -sequence. - -```python -# Configure TTT -model._ttt_cfg.steps = 10 # 10 optimizer steps (default) -model._ttt_cfg.lora_rank = 8 # LoRA rank (default) -model._ttt_cfg.lora_alpha = 32 # LoRA scale (default) - -# ttt=True runs TTT, folds after each step, returns best structure -result = model.fold_protein("MKTLLILAVVAAALA...", ttt=True) -# Equivalent: -# result = model.fold_protein_ttt("MKTLLILAVVAAALA...") -print(f"pLDDT: {result['plddt']:.3f}") -print(f"Best step: {result['best_step']} (0=baseline, 1-10=TTT steps)") -print(f"Step pLDDTs: {[f'{p:.2f}' for p in result['step_plddts']]}") - -# Save PDB -with open("structure.pdb", "w") as f: - f.write(result["pdb_string"]) -``` - -### Return values - -`fold_protein(sequence)` returns a dict. Without `ttt=True`, `step_plddts` -contains only the baseline pLDDT and `best_step` is `0`. - -| Key | Type | Description | -|-----|------|-------------| -| `plddt` | float | Mean pLDDT for the selected structure | -| `ptm` | float | Predicted TM-score for the selected structure | -| `pdb_string` | str | PDB format structure | -| `step_plddts` | list[float] | Baseline pLDDT, plus per-step pLDDT when TTT is enabled | -| `best_step` | int | Which step produced the selected structure (0=baseline) | - -### TTT default behavior - -TTT is disabled by default. Use FastESMFold as a standard ESMFold by calling -`fold_protein(...)` or `infer(...)` without `ttt=True`: - -```python -# Baseline fold, no TTT -result = model.fold_protein("MKTLLILAVVAAALA...") -print(result["best_step"]) # 0 - -# Raw ESMFold output -with torch.no_grad(): - output = model.infer("MKTLLILAVVAAALA...") -pdb_strings = model.output_to_pdb(output) -``` - -## Experimental TTT Benchmark - -This benchmark is provided as an example of where TTT can help. It is not a -guarantee of improvement on every sequence. - -Tested on 10 difficult sequences on A10G GPU: - -| Metric | Value | -|--------|-------| -| Mean baseline pLDDT | 0.549 | -| Mean best TTT pLDDT | 0.637 | -| Mean improvement | +0.088 | -| Sequences improved >5pt | 5/10 | -| Time per sequence | ~20-45s | -| GPU memory peak | 18.3 GB | - -On the hardest sequence (baseline pLDDT 0.38), TTT improves to 0.72 (+34 points). - -## Attention backends - -The ESM2 backbone supports multiple attention backends via `config.attn_backend`: - -| Backend | Key | Notes | -| :--- | :--- | :--- | -| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. | -| Flash Attention | `"kernels_flash"` | Fastest. Requires `pip install kernels`. | -| Flex Attention | `"flex"` | Skips padding tokens via block mask. First use compiles a Triton kernel. | -| Auto | `"auto"` | Picks best available: `kernels_flash` > `flex` > `sdpa`. | - -```python -from transformers import AutoConfig, AutoModel - -config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True) -config.attn_backend = "kernels_flash" -model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True) -``` - -## TTT Configuration - -TTT parameters are set via `config.ttt_config` (a dict) or by modifying `model._ttt_cfg` after loading: - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `lr` | 4e-4 | Learning rate for SGD optimizer | -| `steps` | 10 | Number of optimizer steps when TTT is explicitly enabled | -| `ags` | 4 | Gradient accumulation steps per optimizer step | -| `batch_size` | 4 | Batch size for masked language model training | -| `mask_ratio` | 0.15 | Fraction of tokens to mask | -| `lora_rank` | 8 | LoRA rank (0 for full backbone fine-tuning) | -| `lora_alpha` | 32.0 | LoRA scaling factor (applied as `scale=alpha`, matching lora_diffusion) | -| `seed` | 0 | Random seed for reproducible LoRA initialization and masking | -| `lora_target_class` | `"EsmSelfAttention"` | Which module class to inject LoRA into | - -## How TTT Works - -1. **Baseline fold** (step 0): Standard ESMFold prediction -2. **LoRA injection**: Rank-8 LoRA adapters on all `nn.Linear` layers inside ESM2 attention modules -3. **Masked LM training**: 10 optimizer steps (each with 4 gradient accumulation sub-steps) of BERT-style masked language modeling on the input sequence -4. **Per-step folding**: After each optimizer step, fold the sequence and record pLDDT -5. **Best selection**: Return the structure with the highest pLDDT -6. **Reset**: Restore LoRA weights to initial state for the next sequence - -## Citations - -```bibtex -@misc{FastPLMs, - author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.}, - title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.}, - year={2024}, - url={https://huggingface.co/Synthyra/ESMplusplus_small}, - DOI={10.57967/hf/3726}, - publisher={Hugging Face} -} -``` - -```bibtex -@misc{bushuiev2026proteinneed, - title={One protein is all you need}, - author={Anton Bushuiev and Roman Bushuiev and Olga Pimenova and Nikola Zadorozhny and Raman Samusevich and Elisabet Manaskova and Rachel Seongeun Kim and Hannes St\"ark and Jiri Sedlar and Martin Steinegger and Tom\'a\v{s} Pluskal and Josef Sivic}, - year={2026}, - eprint={2411.02109}, - archivePrefix={arXiv}, - primaryClass={cs.LG}, - url={https://arxiv.org/abs/2411.02109} -} -``` - -```bibtex -@article{dong2024flexattention, - title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels}, - author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace}, - journal={arXiv preprint arXiv:2412.05496}, - year={2024} -} -``` - -```bibtex -@inproceedings{paszke2019pytorch, - title={PyTorch: An Imperative Style, High-Performance Deep Learning Library}, - author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, - booktitle={Advances in Neural Information Processing Systems 32}, - year={2019} -} -``` diff --git a/fastplms/esmfold/__init__.py b/fastplms/esmfold/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/fastplms/esmfold/get_weights.py b/fastplms/esmfold/get_weights.py deleted file mode 100644 index c0a00b0..0000000 --- a/fastplms/esmfold/get_weights.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Convert Synthyra/ESMFold-v1 weights to FastESMFold format and push to HuggingFace. - -Usage: - py -m fastplms.esmfold.get_weights - py -m fastplms.esmfold.get_weights --dry_run - py -m fastplms.esmfold.get_weights --skip-weights - py -m fastplms.esmfold.get_weights --hf_token -""" -import os -import sys - -import torch - -from huggingface_hub import HfApi, login -from transformers import AutoTokenizer, EsmConfig - -from fastplms.esmfold.modeling_fast_esmfold import FastEsmFoldConfig, FastEsmForProteinFolding - -SOURCE_REPO = "Synthyra/ESMFold-v1" -TARGET_REPO = "Synthyra/FastESMFold" -SHARD_SIZE = "5GB" - - -def convert_and_push( - hf_token: str = None, - dry_run: bool = False, - skip_weights: bool = False, - attn_backend: str = "sdpa", -) -> None: - if hf_token is not None: - login(token=hf_token) - - api = HfApi() - script_root = os.path.dirname(os.path.abspath(__file__)) - - if skip_weights: - config = FastEsmFoldConfig.from_pretrained(TARGET_REPO, trust_remote_code=True) - config.auto_map = { - "AutoConfig": "modeling_fast_esmfold.FastEsmFoldConfig", - "AutoModel": "modeling_fast_esmfold.FastEsmForProteinFolding", - } - if dry_run: - print(f"[skip-weights][dry-run] validated config for {TARGET_REPO}") - return - config.push_to_hub(TARGET_REPO) - print(f"[skip-weights] uploaded config for {TARGET_REPO}") - return - - print(f"Loading source model from {SOURCE_REPO}...") - source_config = EsmConfig.from_pretrained(SOURCE_REPO) - - # Build FastEsmFoldConfig from the source ESMFold config - config_dict = source_config.to_dict() - config_dict["attn_backend"] = attn_backend - config_dict["ttt_config"] = { - "lr": 4e-4, - "steps": 10, - "ags": 4, - "batch_size": 4, - "mask_ratio": 0.15, - "lora_rank": 8, - "lora_alpha": 32.0, - } - config_dict["auto_map"] = { - "AutoConfig": "modeling_fast_esmfold.FastEsmFoldConfig", - "AutoModel": "modeling_fast_esmfold.FastEsmForProteinFolding", - } - config = FastEsmFoldConfig(**config_dict) - - print("Creating FastEsmForProteinFolding model...") - model = FastEsmForProteinFolding(config) - - # Load source weights (EsmFoldWithLMHead = EsmForProteinFolding + mlm_head) - print(f"Loading weights from {SOURCE_REPO}...") - source_state_dict = torch.hub.load_state_dict_from_url( - f"https://huggingface.co/{SOURCE_REPO}/resolve/main/model.safetensors", - map_location="cpu", - ) if False else None - - # Use from_pretrained to load weights into a temporary model, then transfer - from transformers import EsmForProteinFolding as HFEsmFold - - # Load the source model state dict - source_model = HFEsmFold.from_pretrained( - SOURCE_REPO, - dtype=torch.float32, - device_map="cpu", - ) - source_sd = source_model.state_dict() - - # Load what we can (backbone weights map directly since FastEsmBackbone - # has the same module structure as EsmModel) - missing, unexpected = model.load_state_dict(source_sd, strict=False) - - print(f"Missing keys: {len(missing)}") - for k in sorted(missing): - print(f" {k}") - print(f"Unexpected keys: {len(unexpected)}") - for k in sorted(unexpected): - print(f" {k}") - - if dry_run: - print(f"[dry_run] Validated weight transfer for {TARGET_REPO} <- {SOURCE_REPO}") - print(f" Missing: {len(missing)}, Unexpected: {len(unexpected)}") - return - - print(f"Pushing model to {TARGET_REPO}...") - tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - tokenizer.push_to_hub(TARGET_REPO) - model.push_to_hub(TARGET_REPO, max_shard_size=SHARD_SIZE) - - # Upload modeling file - api.upload_file( - path_or_fileobj=os.path.join(script_root, "modeling_fast_esmfold.py"), - path_in_repo="modeling_fast_esmfold.py", - repo_id=TARGET_REPO, - repo_type="model", - ) - - print(f"Done. Model available at https://huggingface.co/{TARGET_REPO}") - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument( - "--hf_token", - type=str, - default=None, - ) - parser.add_argument( - "--dry_run", - action="store_true", - ) - parser.add_argument( - "--skip-weights", - action="store_true", - ) - parser.add_argument( - "--attn_backend", - type=str, - default="sdpa", - ) - args = parser.parse_args() - - convert_and_push( - hf_token=args.hf_token, - dry_run=args.dry_run, - skip_weights=args.skip_weights, - attn_backend=args.attn_backend, - ) diff --git a/fastplms/esmfold/modeling_fast_esmfold.py b/fastplms/esmfold/modeling_fast_esmfold.py deleted file mode 100644 index 556a49a..0000000 --- a/fastplms/esmfold/modeling_fast_esmfold.py +++ /dev/null @@ -1,1047 +0,0 @@ -from __future__ import annotations -"""FastESMFold: self-contained ESMFold with FastESM2 attention and opt-in TTT. - -Usage: - from transformers import AutoModel - model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True).cuda() - - # Basic folding, no TTT - result = model.fold_protein("MKTLLILAVVA...") - print(result["plddt"], result["pdb_string"][:100]) - - # Experimental folding with TTT - result = model.fold_protein("MKTLLILAVVA...", ttt=True) - -Dependencies: torch, transformers, einops -No dependency on: esm (fair-esm), proteinttt, openfold -""" -import copy -from dataclasses import dataclass, field -from functools import wraps -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import torch -import torch.nn as nn -from torch.nn import functional as F - -from einops import rearrange -from transformers import EsmTokenizer, PretrainedConfig, PreTrainedModel -from transformers.modeling_outputs import ModelOutput -from transformers.models.esm.configuration_esm import EsmConfig -from transformers.models.esm.modeling_esm import ( - EsmContactPredictionHead, - EsmEmbeddings, - EsmIntermediate, - EsmLMHead, - EsmOutput, - EsmSelfOutput, - RotaryEmbedding, -) -from transformers.models.esm.modeling_esmfold import EsmForProteinFolding - - -try: - from fastplms.attention import ( - AttentionBackend, VALID_ATTENTION_BACKENDS, - resolve_attention_backend, get_attention_mask, - _get_flex_attention_fn, - _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT, - _kernels_flash_forward, _kernels_flash_varlen_forward, - kernels_flash_attention_func, - index_first_axis, index_put_first_axis, pad_input, _unpad_input, - create_block_mask, flex_attention, BlockMask, - ) -except ImportError: - pass # Running as HF Hub composite; shared definitions are above - - -# ============================================================================= -# Output Dataclass -# ============================================================================= - -@dataclass -class FastEsmEncoderOutput(ModelOutput): - last_hidden_state: Optional[torch.Tensor] = None - hidden_states: Optional[Tuple[torch.Tensor, ...]] = None - attentions: Optional[Tuple[torch.Tensor, ...]] = None - - -# ============================================================================= -# FastESM2 Attention Layers (multi-backend: SDPA, Flash, Flex) -# ============================================================================= - -class EsmSelfAttention(nn.Module): - def __init__(self, config, position_embedding_type: Optional[str] = None): - super().__init__() - assert config.hidden_size % config.num_attention_heads == 0, ( - f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " - f"heads ({config.num_attention_heads})" - ) - self.num_attention_heads = config.num_attention_heads - self.attention_head_size = int(config.hidden_size / config.num_attention_heads) - self.all_head_size = self.num_attention_heads * self.attention_head_size - - self.query = nn.Linear(config.hidden_size, self.all_head_size) - self.key = nn.Linear(config.hidden_size, self.all_head_size) - self.value = nn.Linear(config.hidden_size, self.all_head_size) - self.scale = self.attention_head_size**-0.5 - - self.dropout_prob = config.attention_probs_dropout_prob - self.config = config - self.attn_backend = resolve_attention_backend(config.attn_backend) - self.position_embedding_type = position_embedding_type or config.position_embedding_type - self.rotary_embeddings = None - if self.position_embedding_type == "rotary": - self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - batch_size, seq_length = hidden_states.shape[:-1] - hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) - query_BHLD = self.query(hidden_states).view(hidden_shape).transpose(1, 2) - key_BHLD = self.key(hidden_states).view(hidden_shape).transpose(1, 2) - value_BHLD = self.value(hidden_states).view(hidden_shape).transpose(1, 2) - - query_BHLD = query_BHLD * self.scale - - if self.position_embedding_type == "rotary": - query_BHLD, key_BHLD = self.rotary_embeddings(query_BHLD, key_BHLD) - - attn_output, attn_weights = self._attn( - query_BHLD, key_BHLD, value_BHLD, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - ) - return attn_output, attn_weights - - def _attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - if output_attentions: - return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - - if self.attn_backend == AttentionBackend.KERNELS_FLASH: - return self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) - elif self.attn_backend == AttentionBackend.FLEX: - return self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) - elif self.attn_backend == AttentionBackend.SDPA: - return self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) - else: - raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") - - def _manual_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) - if attention_mask_4d is not None: - attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) - attn_weights = F.softmax(attn_weights, dim=-1) - if self.dropout_prob > 0 and self.training: - attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) - context_BHLD = torch.matmul(attn_weights, value_BHLD) - attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") - return attn_output, attn_weights - - def _kernels_flash_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - query_BLHD = query_BHLD.transpose(1, 2).contiguous() - key_BLHD = key_BHLD.transpose(1, 2).contiguous() - value_BLHD = value_BHLD.transpose(1, 2).contiguous() - # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0 - # to prevent the kernel from applying its default 1/sqrt(head_dim). - attn_output = kernels_flash_attention_func( - query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, - attention_mask_2d=attention_mask_2d, causal=False, - softmax_scale=1.0, - ) - return rearrange(attn_output, "b s h d -> b s (h d)"), None - - def _flex_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - flex_block_mask: Optional[BlockMask] = None, - ) -> Tuple[torch.Tensor, None]: - assert flex_attention is not None, "Flex attention is not available in this environment." - fn = _get_flex_attention_fn() - context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - def _sdpa_attn( - self, - query_BHLD: torch.Tensor, - key_BHLD: torch.Tensor, - value_BHLD: torch.Tensor, - attention_mask_4d: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, None]: - context_BHLD = F.scaled_dot_product_attention( - query_BHLD, key_BHLD, value_BHLD, - attn_mask=attention_mask_4d, - dropout_p=self.dropout_prob if self.training else 0.0, - scale=1.0, - ) - return rearrange(context_BHLD, "b h s d -> b s (h d)"), None - - -class EsmAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.self = EsmSelfAttention(config) - self.output = EsmSelfOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - hidden_states_ln = self.LayerNorm(hidden_states) - attn_output, attn_weights = self.self( - hidden_states_ln, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - ) - attention_output = self.output(attn_output, hidden_states) - return attention_output, attn_weights - - -class EsmLayer(nn.Module): - def __init__(self, config): - super().__init__() - self.attention = EsmAttention(config) - self.intermediate = EsmIntermediate(config) - self.output = EsmOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask_2d: Optional[torch.Tensor] = None, - attention_mask_4d: Optional[torch.Tensor] = None, - flex_block_mask: Optional[BlockMask] = None, - output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - attention_output, attn_weights = self.attention( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - ) - layer_output = self._feed_forward(attention_output) - return layer_output, attn_weights - - def _feed_forward(self, attention_output: torch.Tensor) -> torch.Tensor: - attention_output_ln = self.LayerNorm(attention_output) - intermediate_output = self.intermediate(attention_output_ln) - return self.output(intermediate_output, attention_output) - - -class FastEsmEncoder(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.attention_backend = resolve_attention_backend(config.attn_backend) - self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)]) - self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: bool = False, - output_attentions: bool = False, - ) -> FastEsmEncoderOutput: - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( - effective_backend=self.attention_backend, - batch_size=hidden_states.shape[0], - seq_len=hidden_states.shape[1], - device=hidden_states.device, - attention_mask=attention_mask, - ) - - for layer_module in self.layer: - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - hidden_states, attn_weights = layer_module( - hidden_states, - attention_mask_2d=attention_mask_2d, - attention_mask_4d=attention_mask_4d, - flex_block_mask=flex_block_mask, - output_attentions=output_attentions, - ) - - if all_attentions is not None: - all_attentions = all_attentions + (attn_weights,) - - if self.emb_layer_norm_after: - hidden_states = self.emb_layer_norm_after(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - return FastEsmEncoderOutput( - last_hidden_state=hidden_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - ) - - -# ============================================================================= -# FastESM Backbone (replaces EsmModel inside ESMFold) -# ============================================================================= - -class FastEsmBackbone(nn.Module): - """FastESM2 backbone with multi-backend attention. Drop-in replacement for - transformers.EsmModel inside EsmForProteinFolding. - - State dict keys match HuggingFace EsmModel exactly, so pretrained weights - load without any key remapping. - """ - - def __init__(self, config): - super().__init__() - self.config = config - self.embeddings = EsmEmbeddings(config) - self.encoder = FastEsmEncoder(config) - self.contact_head = EsmContactPredictionHead( - in_features=config.num_hidden_layers * config.num_attention_heads, bias=True - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - **kwargs, - ) -> FastEsmEncoderOutput: - output_attentions = output_attentions if output_attentions is not None else False - output_hidden_states = output_hidden_states if output_hidden_states is not None else False - - token_embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - ) - encoder_outputs = self.encoder( - token_embedding_output, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - ) - return FastEsmEncoderOutput( - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - -# ============================================================================= -# TTT (Test-Time Training) Configuration and Utilities -# ============================================================================= - -_ESM_STANDARD_AA = list("ACDEFGHIKLMNPQRSTVWY") - - -class LoraInjectedLinear(nn.Module): - """LoRA-augmented linear layer matching lora_diffusion's behavior. - - Replaces an existing nn.Linear with base(x) + lora_up(lora_down(x)) * scale. - Initialization follows cloneofsimo/lora: down=Normal(0, 1/r), up=zeros. - """ - - def __init__(self, original_linear: nn.Linear, r: int = 4, scale: float = 1.0): - super().__init__() - self.linear = original_linear - in_features = original_linear.in_features - out_features = original_linear.out_features - assert r <= min(in_features, out_features), f"LoRA rank {r} exceeds dimensions ({in_features}, {out_features})" - self.lora_down = nn.Linear(in_features, r, bias=False) - self.lora_up = nn.Linear(r, out_features, bias=False) - self.scale = scale - nn.init.normal_(self.lora_down.weight, std=1.0 / r) - nn.init.zeros_(self.lora_up.weight) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.linear(x) + self.lora_up(self.lora_down(x)) * self.scale - - -def inject_trainable_lora( - model: nn.Module, - target_class_name: str, - r: int, - scale: float, -) -> List[nn.Parameter]: - """Replace nn.Linear layers inside modules matching target_class_name with LoRA. - - Matches lora_diffusion's inject_trainable_lora behavior: finds all modules whose - class name matches target_class_name, then replaces their nn.Linear children with - LoraInjectedLinear. Returns the list of trainable LoRA parameters. - """ - lora_params: List[nn.Parameter] = [] - for _parent_name, parent_module in model.named_modules(): - if parent_module.__class__.__name__ != target_class_name: - continue - for child_name, child_module in list(parent_module.named_children()): - if not isinstance(child_module, nn.Linear): - continue - lora_linear = LoraInjectedLinear(child_module, r=r, scale=scale) - lora_linear = lora_linear.to( - device=child_module.weight.device, - dtype=child_module.weight.dtype, - ) - setattr(parent_module, child_name, lora_linear) - lora_params.extend(lora_linear.lora_down.parameters()) - lora_params.extend(lora_linear.lora_up.parameters()) - return lora_params - - -@dataclass -class TTTConfig: - lr: float = 4e-4 - ags: int = 4 - steps: int = 10 - batch_size: int = 4 - mask_ratio: float = 0.15 - crop_size: int = 1024 - bert_leave_prob: float = 0.1 - bert_replace_prob: float = 0.1 - optimizer: str = "sgd" - momentum: float = 0.0 - weight_decay: float = 0.0 - seed: Optional[int] = 0 - initial_state_reset: bool = True - freeze_embeddings: bool = True - lora_rank: int = 8 - lora_alpha: float = 32.0 - lora_target_class: str = "EsmSelfAttention" - - def verify(self) -> None: - assert self.lr > 0.0, "TTT learning rate must be positive." - assert self.ags > 0, "TTT ags must be positive." - assert self.steps >= 0, "TTT steps must be non-negative." - assert self.batch_size > 0, "TTT batch_size must be positive." - assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]." - assert self.crop_size > 0, "TTT crop_size must be positive." - assert 0.0 <= self.bert_leave_prob <= 1.0 - assert 0.0 <= self.bert_replace_prob <= 1.0 - assert self.bert_leave_prob + self.bert_replace_prob <= 1.0 - assert self.optimizer in {"sgd", "adamw"} - assert self.lora_rank >= 0 - assert self.lora_alpha > 0.0 - - -def preserve_model_state(func: Callable[..., Any]) -> Callable[..., Any]: - @wraps(func) - def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: - was_training = self.training - original_device = next(self.parameters()).device - original_requires_grad = { - name: parameter.requires_grad - for name, parameter in self.named_parameters() - } - try: - return func(self, *args, **kwargs) - finally: - self.train(was_training) - self.to(original_device) - for name, parameter in self.named_parameters(): - if name in original_requires_grad: - parameter.requires_grad = original_requires_grad[name] - else: - parameter.requires_grad = False - return wrapper - - -# ============================================================================= -# FastEsmFoldConfig -# ============================================================================= - -class FastEsmFoldConfig(EsmConfig): - model_type = "fast_esmfold" - - def __init__(self, attn_backend: str = "sdpa", ttt_config: Optional[Dict[str, Any]] = None, **kwargs): - super().__init__(**kwargs) - self.attn_backend = attn_backend - self.ttt_config = ttt_config or { - "lr": 4e-4, - "steps": 10, - "lora_rank": 8, - "lora_alpha": 32.0, - } - - -# ============================================================================= -# FastEsmForProteinFolding -# ============================================================================= - -class FastEsmForProteinFolding(EsmForProteinFolding): - """ESMFold with FastESM2 attention backends and opt-in experimental TTT. - - Inherits all folding logic (trunk, structure module, output_to_pdb, infer) - from transformers.EsmForProteinFolding. Replaces the ESM2 backbone with - FastESM2 for optimized attention and adds opt-in TTT for difficult targets. - - Key API: - result = model.fold_protein("MKTL...", ttt=True) - # result = {"plddt": float, "ptm": float, "pdb_string": str} - """ - config_class = FastEsmFoldConfig - - def __init__(self, config: FastEsmFoldConfig): - super().__init__(config) - - # Replace standard ESM2 backbone with FastESM2 (multi-backend attention) - # unless use_standard_backbone is set (for TTT debugging/compatibility) - if not config.ttt_config.get("use_standard_backbone", False): - self.esm = FastEsmBackbone(config) - self.esm.requires_grad_(False) - if config.esmfold_config.fp16_esm: - self.esm.half() - - # MLM head for TTT (pretrained EsmLMHead: Dense -> GELU -> LN -> Linear) - self.mlm_head = EsmLMHead(config) - - # TTT state (lazy initialization) - ttt_kwargs = {k: v for k, v in config.ttt_config.items() if k != "use_standard_backbone"} - self._ttt_cfg = TTTConfig(**ttt_kwargs) - self._ttt_cfg.verify() - self._ttt_initialized = False - self._ttt_initial_state = None - self._ttt_generator = torch.Generator() - if self._ttt_cfg.seed is not None: - self._ttt_generator.manual_seed(self._ttt_cfg.seed) - self._non_special_tokens_cache = None - self._ttt_tokenizer = None - - def _get_ttt_tokenizer(self) -> EsmTokenizer: - if self._ttt_tokenizer is None: - self._ttt_tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") - return self._ttt_tokenizer - - def _ensure_ttt_ready(self) -> None: - """Lazy TTT initialization. Injects LoRA adapters and saves initial state. - Must be called after weights are loaded (not in __init__).""" - if self._ttt_initialized: - return - self._ttt_initialized = True - - tokenizer = self._get_ttt_tokenizer() - vocab = tokenizer.get_vocab() - self._non_special_tokens_cache = [vocab[c] for c in _ESM_STANDARD_AA if c in vocab] - - if self._ttt_cfg.lora_rank > 0: - self.mlm_head.eval() - for p in self.mlm_head.parameters(): - p.requires_grad = False - # Seed global state before LoRA init for reproducible weight initialization - if self._ttt_cfg.seed is not None: - torch.manual_seed(self._ttt_cfg.seed) - self._inject_lora() - else: - # Legacy path: jointly-trained random linear projection head - H = self.config.hidden_size - V = self.config.vocab_size - device = next(self.esm.parameters()).device - self._ttt_lm_proj = nn.Linear(H, V, bias=True).to(device) - - if self._ttt_cfg.initial_state_reset: - self._ttt_initial_state = self._ttt_get_state() - - @property - def _uses_lora(self) -> bool: - return self._ttt_cfg.lora_rank > 0 - - def _inject_lora(self) -> None: - """Inject LoRA adapters into ESM2 attention layers (matching lora_diffusion behavior).""" - self._lora_params = inject_trainable_lora( - self.esm, - target_class_name=self._ttt_cfg.lora_target_class, - r=self._ttt_cfg.lora_rank, - scale=self._ttt_cfg.lora_alpha, - ) - assert len(self._lora_params) > 0, ( - f"No LoRA params injected. Check target_class_name='{self._ttt_cfg.lora_target_class}' " - f"matches attention modules in the backbone." - ) - - # ---- TTT State Management ---- - - def _get_lora_modules(self) -> List[LoraInjectedLinear]: - """Find all LoraInjectedLinear modules in the backbone.""" - return [m for m in self.esm.modules() if isinstance(m, LoraInjectedLinear)] - - def _ttt_get_state(self) -> Dict[str, Any]: - if self._uses_lora: - lora_state = [] - for m in self._get_lora_modules(): - lora_state.append({ - "down": m.lora_down.weight.data.clone(), - "up": m.lora_up.weight.data.clone(), - }) - return {"_lora_state": lora_state} - return { - "esm": copy.deepcopy(self.esm), - "_ttt_lm_proj": copy.deepcopy(self._ttt_lm_proj), - } - - def _ttt_set_state(self, state: Dict[str, Any]) -> None: - if "_lora_state" in state: - modules = self._get_lora_modules() - assert len(modules) == len(state["_lora_state"]) - for m, saved in zip(modules, state["_lora_state"]): - m.lora_down.weight.data.copy_(saved["down"]) - m.lora_up.weight.data.copy_(saved["up"]) - return - if "esm" in state: - self.esm = copy.deepcopy(state["esm"]) - if "_ttt_lm_proj" in state: - self._ttt_lm_proj = copy.deepcopy(state["_ttt_lm_proj"]) - - def ttt_reset(self) -> None: - """Reset model to pre-TTT state (restore initial LoRA or backbone weights).""" - assert self._ttt_initial_state is not None, "TTT reset requires initial_state_reset=True." - self._ttt_set_state(self._ttt_initial_state) - - # ---- TTT Core ---- - - def _ttt_tokenize(self, seq: str) -> torch.Tensor: - tokenizer = self._get_ttt_tokenizer() - out = tokenizer( - seq, - return_tensors="pt", - add_special_tokens=self._uses_lora, - padding=False, - truncation=False, - ) - return out["input_ids"] - - def _ttt_mask_token(self) -> int: - return self._get_ttt_tokenizer().mask_token_id - - def _ttt_get_non_special_tokens(self) -> List[int]: - if self._non_special_tokens_cache is not None: - return self._non_special_tokens_cache - tokenizer = self._get_ttt_tokenizer() - vocab = tokenizer.get_vocab() - self._non_special_tokens_cache = [vocab[c] for c in _ESM_STANDARD_AA if c in vocab] - return self._non_special_tokens_cache - - def _ttt_predict_logits(self, batch: torch.Tensor) -> torch.Tensor: - """Run ESM2 backbone + LM head to get MLM logits.""" - # Temporarily unfreeze backbone for gradient flow during TTT - output = self.esm(input_ids=batch) - hidden = output.last_hidden_state - if self._uses_lora: - return self.mlm_head(hidden) - return self._ttt_lm_proj(hidden) - - def _ttt_sample_batch( - self, - x: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - _, seq_len = x.shape - batch_size = self._ttt_cfg.batch_size - crop_size = min(self._ttt_cfg.crop_size, seq_len) - - x_expanded = x.expand(batch_size, -1) - if seq_len == crop_size: - start_indices = torch.zeros(batch_size, dtype=torch.long) - else: - start_indices = torch.randint( - 0, seq_len - crop_size + 1, (batch_size,), - generator=self._ttt_generator, - ).to(torch.long) - - batch_cropped = torch.stack([ - x_expanded[index, start : start + crop_size] - for index, start in enumerate(start_indices) - ]) - - non_special_tokens = set(self._ttt_get_non_special_tokens()) - mask = torch.zeros((batch_size, crop_size), dtype=torch.bool) - mask_token_id = self._ttt_mask_token() - - for row_index in range(batch_size): - non_special_positions = [ - col for col in range(crop_size) - if batch_cropped[row_index, col].item() in non_special_tokens - ] - assert len(non_special_positions) > 0, "Sequence must contain at least one non-special token." - num_to_mask = max(1, int(round(len(non_special_positions) * self._ttt_cfg.mask_ratio))) - sampled_indices = torch.randperm( - len(non_special_positions), generator=self._ttt_generator, - )[:num_to_mask] - positions_to_mask = torch.tensor(non_special_positions, dtype=torch.long)[sampled_indices] - mask[row_index, positions_to_mask] = True - - batch_masked = batch_cropped.clone() - for row_index in range(batch_size): - masked_positions = torch.nonzero(mask[row_index], as_tuple=True)[0] - for masked_position in masked_positions: - probability = float(torch.rand(1, generator=self._ttt_generator).item()) - if probability < 1.0 - self._ttt_cfg.bert_leave_prob - self._ttt_cfg.bert_replace_prob: - batch_masked[row_index, masked_position] = mask_token_id - continue - if probability < 1.0 - self._ttt_cfg.bert_leave_prob: - replacement_candidates = self._ttt_get_non_special_tokens() - replacement_index = int(torch.randint( - 0, len(replacement_candidates), (1,), generator=self._ttt_generator, - ).item()) - batch_masked[row_index, masked_position] = replacement_candidates[replacement_index] - - return batch_masked, batch_cropped, mask, start_indices - - def _ttt_cross_entropy_loss( - self, - logits: torch.Tensor, - targets: torch.Tensor, - mask: torch.Tensor, - ) -> torch.Tensor: - assert logits.ndim == 3, "Logits must be [batch, seq, vocab]." - _, _, vocab_size = logits.shape - logits_flat = logits.reshape(-1, vocab_size) - targets_flat = targets.reshape(-1) - mask_flat = mask.reshape(-1) - assert int(mask_flat.sum().item()) > 0, "TTT mask must select at least one token." - loss = F.cross_entropy( - logits_flat[mask_flat], - targets_flat[mask_flat], - reduction="none", - ) - masked_tokens_per_seq = mask.sum(dim=1).tolist() - per_sequence_losses = torch.split(loss, masked_tokens_per_seq) - return torch.stack([sl.mean() for sl in per_sequence_losses]).mean() - - def _ttt_get_optimizer(self, parameters) -> torch.optim.Optimizer: - if self._ttt_cfg.optimizer == "sgd": - return torch.optim.SGD( - parameters, - lr=self._ttt_cfg.lr, - momentum=self._ttt_cfg.momentum, - weight_decay=self._ttt_cfg.weight_decay, - ) - return torch.optim.AdamW( - parameters, - lr=self._ttt_cfg.lr, - weight_decay=self._ttt_cfg.weight_decay, - ) - - def _lora_ttt(self, seq: str) -> Dict[str, List[float]]: - """LoRA TTT: only LoRA adapter weights are trained, mlm_head is frozen.""" - x = self._ttt_tokenize(seq) - device = next(self.parameters()).device - non_blocking = device.type == "cuda" - losses = [] - - if self._ttt_cfg.steps == 0: - return {"losses": losses} - - for parameter in self.parameters(): - parameter.requires_grad = False - for p in self._lora_params: - p.requires_grad = True - optimizer = self._ttt_get_optimizer(self._lora_params) - optimizer.zero_grad(set_to_none=True) - - self.eval() - for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): - batch_masked, targets, mask, start_indices = self._ttt_sample_batch(x) - batch_masked = batch_masked.to(device, non_blocking=non_blocking) - targets = targets.to(device, non_blocking=non_blocking) - mask = mask.to(device, non_blocking=non_blocking) - - self.train() - logits = self._ttt_predict_logits(batch_masked) - loss = self._ttt_cross_entropy_loss(logits, targets, mask) - loss.backward() - losses.append(float(loss.detach().cpu().item())) - - if (step + 1) % self._ttt_cfg.ags == 0: - optimizer.step() - optimizer.zero_grad(set_to_none=True) - - self.eval() - return {"losses": losses} - - def _legacy_ttt(self, seq: str) -> Dict[str, List[float]]: - """Legacy TTT: full fine-tuning of ESM2 backbone with random linear projection head.""" - x = self._ttt_tokenize(seq) - device = next(self.parameters()).device - non_blocking = device.type == "cuda" - losses = [] - - if self._ttt_cfg.steps == 0: - return {"losses": losses} - - # Full fine-tune: all backbone params trainable - for parameter in self.parameters(): - parameter.requires_grad = False - for parameter in self.esm.parameters(): - parameter.requires_grad = True - if self._ttt_cfg.freeze_embeddings: - for parameter in self.esm.embeddings.parameters(): - parameter.requires_grad = False - for parameter in self._ttt_lm_proj.parameters(): - parameter.requires_grad = True - - trainable_params = filter(lambda p: p.requires_grad, self.parameters()) - optimizer = self._ttt_get_optimizer(trainable_params) - optimizer.zero_grad(set_to_none=True) - - self.eval() - for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): - batch_masked, targets, mask, start_indices = self._ttt_sample_batch(x) - batch_masked = batch_masked.to(device, non_blocking=non_blocking) - targets = targets.to(device, non_blocking=non_blocking) - mask = mask.to(device, non_blocking=non_blocking) - - self.train() - logits = self._ttt_predict_logits(batch_masked) - loss = self._ttt_cross_entropy_loss(logits, targets, mask) - loss.backward() - losses.append(float(loss.detach().cpu().item())) - - if (step + 1) % self._ttt_cfg.ags == 0: - optimizer.step() - optimizer.zero_grad(set_to_none=True) - - self.eval() - return {"losses": losses} - - @preserve_model_state - def ttt(self, seq: str) -> Dict[str, List[float]]: - """Run test-time training on a single sequence using masked language modeling. - - Adapts the ESM2 backbone (via LoRA or full fine-tuning) to the input sequence - before structure prediction. Call fold_protein(seq, ttt=True) for the full pipeline. - - Args: - seq: Protein sequence (single-letter amino acid codes) - - Returns: - Dict with "losses" key containing per-step MLM loss values - """ - self._ensure_ttt_ready() - # TTT requires fp32 for stable gradient computation. ESMFold typically - # runs the backbone in fp16, but small LoRA updates vanish in half precision. - esm_dtype = next(self.esm.parameters()).dtype - if esm_dtype != torch.float32: - self.esm.float() - self.mlm_head.float() - if self._uses_lora: - result = self._lora_ttt(seq) - else: - result = self._legacy_ttt(seq) - # Restore original dtype (backbone back to fp16 for inference) - if esm_dtype != torch.float32: - self.esm.to(esm_dtype) - self.mlm_head.to(esm_dtype) - return result - - # ---- High-Level API ---- - - def _fold_single(self, sequence: str, return_pdb_string: bool = True) -> Dict[str, Any]: - """Fold a sequence once and return pLDDT, ptm, and optionally PDB string.""" - with torch.no_grad(): - output = self.infer(sequence) - plddt = output["plddt"] - # plddt shape is (batch, L, 37) - per-atom across atom37 types. - # Use CA atom (index 1) only, matching PDB B-factor output. - if plddt.dim() == 3: - mean_plddt = float(plddt[:, :, 1].mean().item()) - elif plddt.dim() == 2: - mean_plddt = float(plddt[:, 1].mean().item()) - else: - mean_plddt = float(plddt.mean().item()) - result = { - "plddt": mean_plddt, - "ptm": float(output["ptm"].item()) if "ptm" in output else None, - } - if return_pdb_string: - pdb_strings = self.output_to_pdb(output) - result["pdb_string"] = pdb_strings[0] if isinstance(pdb_strings, list) else pdb_strings - return result - - def fold_protein( - self, - sequence: str, - return_pdb_string: bool = True, - ttt: bool = False, - ) -> Dict[str, Any]: - """Fold a protein sequence. - - Test-time training is disabled by default. Pass ``ttt=True`` or call - ``fold_protein_ttt`` to opt in to experimental TTT. - - Args: - sequence: Protein sequence (single-letter amino acid codes) - return_pdb_string: If True, include PDB string in output - ttt: If True, run experimental LoRA TTT before returning the best fold - - Returns: - Dict with keys: - - plddt: float, mean pLDDT - - ptm: float, predicted TM-score - - pdb_string: str (if return_pdb_string=True), PDB from best step - - step_plddts: list[float], baseline pLDDT when TTT is disabled - - best_step: int, 0 when TTT is disabled - """ - if ttt: - return self.fold_protein_ttt( - sequence=sequence, - return_pdb_string=return_pdb_string, - ) - result = self._fold_single(sequence, return_pdb_string=return_pdb_string) - return { - "plddt": result["plddt"], - "ptm": result["ptm"], - "pdb_string": result.get("pdb_string"), - "step_plddts": [result["plddt"]], - "best_step": 0, - } - - def fold_protein_ttt( - self, - sequence: str, - return_pdb_string: bool = True, - ) -> Dict[str, Any]: - """Fold a protein sequence with experimental test-time training. - - Runs TTT (masked language model adaptation via LoRA) for the configured - number of steps, folding after each optimizer step to track pLDDT. Returns - the structure with the highest pLDDT across all steps (including baseline). - - Args: - sequence: Protein sequence (single-letter amino acid codes) - return_pdb_string: If True, include PDB string in output - - Returns: - Dict with keys: - - plddt: float, best mean pLDDT across all TTT steps - - ptm: float, predicted TM-score from best step - - pdb_string: str (if return_pdb_string=True), PDB from best step - - step_plddts: list[float], pLDDT at each step [baseline, s1, ..., s10] - - best_step: int, which step produced best structure (0=baseline) - """ - self._ensure_ttt_ready() - - # Cast to fp32 for TTT stability - esm_dtype = next(self.esm.parameters()).dtype - if esm_dtype != torch.float32: - self.esm.float() - self.mlm_head.float() - - device = next(self.parameters()).device - non_blocking = device.type == "cuda" - - # Step 0: baseline fold (no TTT adaptation) - best = self._fold_single(sequence, return_pdb_string=return_pdb_string) - step_plddts = [best["plddt"]] - - if self._ttt_cfg.steps > 0: - # Tokenize for masked LM training - x = self._ttt_tokenize(sequence) - - # Freeze all, unfreeze LoRA - for p in self.parameters(): - p.requires_grad = False - if self._uses_lora: - for p in self._lora_params: - p.requires_grad = True - optimizer = self._ttt_get_optimizer(self._lora_params) - else: - for p in self.esm.parameters(): - p.requires_grad = True - if self._ttt_cfg.freeze_embeddings: - for p in self.esm.embeddings.parameters(): - p.requires_grad = False - for p in self._ttt_lm_proj.parameters(): - p.requires_grad = True - trainable = [p for p in self.parameters() if p.requires_grad] - optimizer = self._ttt_get_optimizer(trainable) - optimizer.zero_grad(set_to_none=True) - - self.eval() - for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): - batch_masked, targets, mask, _start = self._ttt_sample_batch(x) - batch_masked = batch_masked.to(device, non_blocking=non_blocking) - targets = targets.to(device, non_blocking=non_blocking) - mask = mask.to(device, non_blocking=non_blocking) - - self.train() - logits = self._ttt_predict_logits(batch_masked) - loss = self._ttt_cross_entropy_loss(logits, targets, mask) - loss.backward() - - if (step + 1) % self._ttt_cfg.ags == 0: - optimizer.step() - optimizer.zero_grad(set_to_none=True) - - # Fold after this optimizer step - self.eval() - current = self._fold_single(sequence, return_pdb_string=return_pdb_string) - step_plddts.append(current["plddt"]) - if current["plddt"] > best["plddt"]: - best = current - - self.eval() - - # Restore requires_grad - for p in self.parameters(): - p.requires_grad = False - - # Reset LoRA weights for next sequence - self.ttt_reset() - - # Restore dtype - if esm_dtype != torch.float32: - self.esm.to(esm_dtype) - self.mlm_head.to(esm_dtype) - - return { - "plddt": best["plddt"], - "ptm": best["ptm"], - "pdb_string": best.get("pdb_string"), - "step_plddts": step_plddts, - "best_step": step_plddts.index(max(step_plddts)), - } diff --git a/fastplms/esmfold2/README.md b/fastplms/esmfold2/README.md deleted file mode 100644 index a8e80f5..0000000 --- a/fastplms/esmfold2/README.md +++ /dev/null @@ -1,247 +0,0 @@ ---- -library_name: transformers -tags: - - biology - - protein-structure - - esmfold2 - - multimodal-protein-model ---- - -# FastPLMs ESMFold2 - -FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for -Biohub's ESMFold2, ESMFold2-Fast, and experimental ESMFold2 structure -predictors. It vendors the released Biohub ESMFold2 model code, input builder, -MSA helpers, and structure export utilities, while loading the PLM backbone -through FastPLMs ESM++. - -## Load With AutoModel - -```python -import torch -from transformers import AutoModel - -model = AutoModel.from_pretrained( - "Synthyra/ESMFold2-Fast", - trust_remote_code=True, - dtype=torch.float32, -).eval().cuda() -``` - -Use `Synthyra/ESMFold2` for the full model, `Synthyra/ESMFold2-Fast` for the -faster release variant, and the `Synthyra/ESMFold2-Experimental*` checkpoints -for differentiable binder design and experimental critic ensembles. -The folding trunk runs in fp32; the 6B FastPLMs ESM++ backbone is loaded in -bf16 by default via `esmc_precision="bf16"` and uses the flex attention backend -by default inside ESMFold2. - -## Fold One Protein - -```python -sequence = "MKTLLILAVVAAALA" - -result = model.fold_protein( - sequence, - num_loops=3, - num_sampling_steps=50, - num_diffusion_samples=1, - seed=0, -) - -print(float(result.plddt.mean())) -print(float(result.ptm)) -``` - -## Experimental Test-Time Training - -TTT is disabled by default. Standard `fold_protein(...)`, `fold(...)`, raw tensor -inference, and `state_dict()` keys are unchanged unless you explicitly pass -`ttt=True` or call `fold_protein_ttt(...)`. - -The ESMFold2 TTT path is experimental and protein-only in v1. It trains local -LoRA adapters only on `_esmc` with a masked language modeling objective. The -folding trunk, confidence head, diffusion head, and structure input pipeline are -frozen. TTT can improve difficult low-confidence folds, but it adds substantial -test-time compute and can degrade already confident predictions. - -```python -result = model.fold_protein( - "MSTNPKPQRKTKRNT", - num_loops=1, - num_sampling_steps=10, - num_diffusion_samples=1, - seed=0, - ttt=True, - ttt_config={ - "steps": 1, - "ags": 1, - "batch_size": 1, - "lora_rank": 8, - "lora_alpha": 32.0, - }, -) - -print(result.ttt_metrics["losses"]) -print(result.ttt_metrics["step_plddts"]) -print(result.ttt_metrics["best_step"]) -``` - -`load_esmc=True` is required for TTT because the ESM++ MLM head is loaded lazily -from `config.esmc_id`. If that pretrained MLM head cannot be loaded, TTT raises -an assertion instead of silently using a random head. - -## Save mmCIF or PDB - -```python -model.save_as_cif(result, "prediction.cif") -model.save_as_pdb(result, "prediction.pdb") - -cif_text = model.result_to_cif(result) -pdb_text = model.result_to_pdb(result) -``` - -`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids. - -## Fold Complexes - -```python -types = model.input_types - -complex_input = types.StructurePredictionInput( - sequences=[ - types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"), - types.DNAInput(id="B", sequence="GATAGC"), - types.LigandInput(id="L", ccd=["SAH"]), - ] -) - -result = model.fold( - complex_input, - num_loops=3, - num_sampling_steps=50, - num_diffusion_samples=1, - seed=0, -) - -model.save_as_cif(result, "complex_prediction.cif") -``` - -## Binder Design With FastPLMs ESMFold2 - -FastPLMs includes a FastPLMs-only port of the Biohub ESMFold2 binder design -tutorial at `cookbook/tutorials/binder_design_fastplms.py`. The workflow uses -ESMFold2 experimental checkpoints for differentiable folding losses, ESM++ for -sequence regularization, and ESMFold2 hero critics for final confidence scoring. - -![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png) - -The optimizer follows the official strategy: - -1. Optimize mutable `#` residues as continuous amino acid logits. -2. Suppress cysteine design by masking cysteine logits and gradients. -3. Backpropagate through ESMFold2 `res_type_soft` using intra-contact, - inter-contact, and globularity losses from the distogram. -4. Add an ESM++ masked-LM pseudoperplexity regularizer on mutable binder - residues. -5. Keep the late-trajectory sequence with the best iPTM. -6. Fold the selected sequence with the final critic ensemble and write - `results.parquet`, `selection.parquet`, `trajectory.jsonl`, - `best_sequences.fasta`, and per-critic PDB/CIF/logit files. - -Run the verified EGFR 128 amino acid de novo minibinder example: - -```bash -cd /home/ubuntu/FastPLMs - -sudo -n docker run --gpus all --rm \ - -v /home/ubuntu/FastPLMs:/app \ - -v /home/ubuntu/FastPLMs:/workspace \ - -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \ - -w /workspace fastplms-esmfold2 \ - python /app/cookbook/tutorials/binder_design_fastplms.py \ - --backend local \ - --target-name egfr \ - --binder-sequence '################################################################################################################################' \ - --not-antibody \ - --steps 150 \ - --batch-size 1 \ - --seed 103 \ - --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli -``` - -Verified result: - -| Metric | Value | -| :--- | :--- | -| Binder length | `128` | -| Seed | `103` | -| Steps | `150` | -| Hero mean iPTM | `0.913870` | -| Hero min iPTM | `0.904600` | -| All four hero critics above 0.9 | `True` | - -Binder sequence: - -```text -SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ -``` - -See the full guide in [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md) -for Modal execution, official pI and selection scoring, per-critic metrics, and -the tested cheaper step-count boundary. - -## Use MSAs - -```python -types = model.input_types - -msa = types.MSA.from_a3m("query.a3m", max_sequences=128) -input_with_msa = types.StructurePredictionInput( - sequences=[ - types.ProteinInput(id="A", sequence=msa.query, msa=msa), - ] -) - -result = model.fold(input_with_msa, num_sampling_steps=50, seed=0) -``` - -## Raw Tensor Inference - -```python -features, chain_infos = model.prepare_structure_input(complex_input, seed=0) - -with torch.inference_mode(): - output = model( - **features, - num_loops=3, - num_sampling_steps=50, - num_diffusion_samples=1, - ) - -decoded = model.input_builder.decode(output, features, chain_infos) -``` - -Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESM++ backbone: - -```python -model = AutoModel.from_pretrained( - "Synthyra/ESMFold2-Fast", - trust_remote_code=True, - load_esmc=False, -).cuda().eval() -``` - -For FP8 LM inference, install `transformer_engine.pytorch` in a CUDA -environment with FP8-capable hardware and load the shared FastPLMs ESM++ -backbone with: - -```python -model = AutoModel.from_pretrained( - "Synthyra/ESMFold2-Fast", - trust_remote_code=True, - esmc_precision="fp8", -).cuda().eval() -``` - -FP8 is inference-only for the ESMFold2 LM backbone. TTT remains a bf16/fp32 -path. diff --git a/fastplms/esmfold2/__init__.py b/fastplms/esmfold2/__init__.py deleted file mode 100644 index 83263a0..0000000 --- a/fastplms/esmfold2/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .configuration_esmfold2 import ESMFold2Config -from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel -from .modeling_esmfold2 import ESMFold2Model - -__all__ = ["ESMFold2Config", "ESMFold2ExperimentalModel", "ESMFold2Model"] diff --git a/fastplms/esmfold2/configuration_esmc.py b/fastplms/esmfold2/configuration_esmc.py deleted file mode 100644 index 629680d..0000000 --- a/fastplms/esmfold2/configuration_esmc.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""ESMC model configuration.""" - -from transformers.configuration_utils import PretrainedConfig - - -class ESMCConfig(PretrainedConfig): - """ - This is the configuration class to store the configuration of a [`ESMCModel`]. It is used to - instantiate an ESMC model according to the specified arguments, defining the model architecture. - - Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model - outputs. Read the documentation from [`PretrainedConfig`] for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 64): - Vocabulary size of the ESMC model. Defines the number of different amino acid tokens that - can be represented by the ``input_ids`` passed to [`ESMCModel`]. - d_model (`int`, *optional*, defaults to 2560): - Dimensionality of the encoder layers and the pooler layer. - n_heads (`int`, *optional*, defaults to 40): - Number of attention heads for each attention layer in the Transformer encoder. - n_layers (`int`, *optional*, defaults to 80): - Number of hidden layers in the Transformer encoder. - pad_token_id (`int`, *optional*, defaults to 1): - Index of the padding token in the vocabulary (``""``). - mask_token_id (`int`, *optional*, defaults to 32): - Index of the mask token in the vocabulary (``""``), used for masked language modelling. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated normal initialiser for weight matrix initialisation. - classifier_dropout (`float`, *optional*, defaults to 0.1): - Dropout ratio for the classification head. - - Examples: - - ```python - >>> from transformers import ESMCConfig, ESMCModel - - >>> # Initializing an ESMC EvolutionaryScale/esmc-600m-2024-12 style configuration - >>> configuration = ESMCConfig() - - >>> # Initializing a model (with random weights) from the EvolutionaryScale/esmc-600m-2024-12 style configuration - >>> model = ESMCModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - """ - - model_type = "esmc" - - def __init__( - self, - vocab_size: int = 64, - d_model: int = 2560, - n_heads: int = 40, - n_layers: int = 80, - pad_token_id: int = 1, - mask_token_id: int = 32, - initializer_range: float = 0.02, - classifier_dropout: float = 0.1, - **kwargs, - ): - super().__init__( - pad_token_id=pad_token_id, mask_token_id=mask_token_id, **kwargs - ) - - self.vocab_size = vocab_size - self.d_model = d_model - self.n_heads = n_heads - self.n_layers = n_layers - self.initializer_range = initializer_range - self.classifier_dropout = classifier_dropout - self.tie_word_embeddings = False - - -__all__ = ["ESMCConfig"] diff --git a/fastplms/esmfold2/configuration_esmc_sae.py b/fastplms/esmfold2/configuration_esmc_sae.py deleted file mode 100644 index 217feab..0000000 --- a/fastplms/esmfold2/configuration_esmc_sae.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""ESMC sparse autoencoder (SAE) configuration.""" - -from dataclasses import dataclass - -from transformers.configuration_utils import PretrainedConfig - - -@dataclass -class ESMCSAEParams: - """Parameters for one backbone layer's SAE inside :class:`ESMCSAEModel`. - - The SAE itself is an internal ``nn.Module``; this dataclass just bundles - the handful of fields needed to instantiate one. - """ - - d_model: int = 2560 - codebook_dim: int = 65536 - k: int = 64 - layer: int = 0 - - -class ESMCSAEConfig(PretrainedConfig): - """ - Configuration class for [`ESMCSAEModel`] — a container that holds one - SAE per backbone layer for a fixed ``(model, codebook_dim, k)`` group. - - All SAEs in a container share ``d_model``, ``codebook_dim``, and ``k``; - they differ only in the backbone layer they were trained on. - ``available_layers`` lists the backbone-layer indices the repo ships; - each entry ``i`` is stored on disk as ``layer_{i}.safetensors`` (the - filename index *is* the backbone layer, so a single-layer repo for - layer 23 stores ``layer_23.safetensors``). - - Args: - d_model (`int`, *optional*, defaults to 2560): - Dimensionality of the ESMC hidden states fed into the SAEs. - codebook_dim (`int`, *optional*, defaults to 65536): - Number of sparse features in each SAE's codebook. - k (`int`, *optional*, defaults to 64): - Top-k sparsity per SAE. - available_layers (`list[int]`, *optional*, defaults to ``[0]``): - Which backbone-layer indices the repo ships. - """ - - model_type = "esmc_sae" - - def __init__( - self, - d_model: int = 2560, - codebook_dim: int = 65536, - k: int = 64, - available_layers: list[int] | None = None, - **kwargs, - ): - super().__init__(**kwargs) - self.d_model = d_model - self.codebook_dim = codebook_dim - self.k = k - self.available_layers = ( - list(available_layers) if available_layers is not None else [0] - ) - - -__all__ = ["ESMCSAEConfig", "ESMCSAEParams"] diff --git a/fastplms/esmfold2/configuration_esmfold2.py b/fastplms/esmfold2/configuration_esmfold2.py deleted file mode 100644 index c4467ac..0000000 --- a/fastplms/esmfold2/configuration_esmfold2.py +++ /dev/null @@ -1,328 +0,0 @@ -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""ESMFold2 model configuration.""" - -from __future__ import annotations - -from dataclasses import asdict, dataclass, field - -from transformers.configuration_utils import PretrainedConfig - -# --------------------------------------------------------------------------- -# Nested dataclass configs -# --------------------------------------------------------------------------- - -_DEFAULT_ESMC_HF_REPO = "Synthyra/ESMplusplus_6B" -_DEFAULT_ESMC_ATTN_BACKEND = "flex" -_ESMC_ID_ALIASES = { - "biohub/ESMC-300M": "Synthyra/ESMplusplus_small", - "biohub/ESMC-600M": "Synthyra/ESMplusplus_large", - "biohub/ESMC-6B": "Synthyra/ESMplusplus_6B", -} - - -def normalize_esmc_id(esmc_id: str) -> str: - if esmc_id in _ESMC_ID_ALIASES: - return _ESMC_ID_ALIASES[esmc_id] - return esmc_id - - -@dataclass -class MSAEncoderConfig: - """Config for the optional MSA encoder module (Large MSA models only).""" - - enabled: bool = False - d_msa: int = 128 - d_hidden: int = 32 - n_layers: int = 4 - n_heads_msa: int = 8 - msa_head_width: int = 32 - - -@dataclass -class ParcaeConfig: - """Release-only config for the parcae diffusion-loop scheduler.""" - - enabled: bool = True - poisson_mean: float = 3.0 - min_steps: int = 1 - max_steps: int | None = 6 - coda_n_layers: int = 2 - - -@dataclass -class LMEncoderConfig: - """Release-only config for the LM-side pair encoder.""" - - enabled: bool = True - n_layers: int = 4 - lm_dropout: float = 0.25 - per_loop_lm_dropout: bool = True - - -@dataclass -class AtomAttentionConfig: - """Config for SWA atom encoder/decoder with 3D RoPE.""" - - d_atom: int = 128 - d_token: int = 768 - n_blocks: int = 3 - n_heads: int = 4 - swa_window_size: int = 128 - expansion_ratio: int = 2 - # 3D RoPE config - spatial_rope_base_frequency: float = 20.0 - n_spatial_rope_pairs_per_axis: int = 2 - n_uid_rope_pairs: int = 10 - uid_rope_base_frequency: float = 10000.0 - - -@dataclass -class FoldingTrunkConfig: - n_layers: int = 24 - n_heads: int = 8 - dropout: float = 0.0 - - -@dataclass -class InputsEmbedderConfig: - d_inputs: int = 451 - atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig) - - def __post_init__(self): - if isinstance(self.atom_encoder, dict): - self.atom_encoder = AtomAttentionConfig(**self.atom_encoder) - - -@dataclass -class DiffusionModuleConfig: - """Config for the DiffusionModule.""" - - sigma_data: float = 16.0 - c_atom: int = 128 - c_token: int = 768 - c_z: int = 256 - c_s_inputs: int = 451 - fourier_dim: int = 256 - relpos_r_max: int = 32 - relpos_s_max: int = 2 - atom_num_blocks: int = 3 - atom_num_heads: int = 4 - token_num_blocks: int = 12 - token_num_heads: int = 16 - transition_multiplier: int = 2 - - -@dataclass -class DiffusionStructureHeadConfig: - """Config for the diffusion-based structure prediction head.""" - - diffusion_module: DiffusionModuleConfig = field( - default_factory=DiffusionModuleConfig - ) - distogram_bins: int = 128 - - # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1)) - train_noise_log_mean: float = -1.2 - train_noise_log_std: float = 1.5 - - # Sampling defaults (ODE) - gamma_0: float = 0.605 - gamma_min: float = 1.107 - noise_scale: float = 0.0 - step_scale: float = 1.0 - - # Inference schedule defaults - inference_s_max: float = 160.0 - inference_s_min: float = 4e-4 - inference_p: float = 8.0 - inference_num_steps: int = 68 - - def __post_init__(self): - if isinstance(self.diffusion_module, dict): - self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module) - - -@dataclass -class ConfidenceHeadConfig: - enabled: bool = True - num_plddt_bins: int = 50 - num_pde_bins: int = 64 - num_pae_bins: int = 64 - min_dist: float = 2.0 - max_dist: float = 52.0 - distogram_bins: int = 128 - folding_trunk: FoldingTrunkConfig = field( - default_factory=lambda: FoldingTrunkConfig(n_layers=4) - ) - - def __post_init__(self): - if isinstance(self.folding_trunk, dict): - self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk) - - -# --------------------------------------------------------------------------- -# Top-level config -# --------------------------------------------------------------------------- - - -class ESMFold2Config(PretrainedConfig): - """ - Configuration for the ESMFold2 structure prediction model. - - Uses SWA atom encoders with 3D RoPE, a diffusion transformer, - a folding trunk, and an ESMC 6B PLM backbone. - - Configuration objects inherit from [`PretrainedConfig`] and can be used to control - the model outputs. Read the documentation from [`PretrainedConfig`] for more - information. - - Args: - d_single (`int`, defaults to 384): - Dimensionality of single (per-residue) representations. - d_pair (`int`, defaults to 256): - Dimensionality of pair (residue-residue) representations. - n_relative_residx_bins (`int`, defaults to 32): - Number of bins for relative residue index encoding. - n_relative_chain_bins (`int`, defaults to 2): - Number of bins for relative chain encoding. - num_loops (`int`, defaults to 10): - Number of trunk loops for iterative refinement. - num_diffusion_samples (`int`, defaults to 8): - Number of parallel structure predictions to generate. - lm_dropout (`float`, defaults to 0.0): - Dropout probability on LM pair embeddings. When > 0, dropout is - applied with ``training=True`` (including at inference) to match - the experimental training recipe used by binder design. - force_lm_dropout_during_inference (`bool`, defaults to False): - When True, apply ``lm_dropout`` even when ``model.eval()`` and - ``lm_dropout`` > 0. Binder-design loads set this to True. - lm_mask_pct (`float`, defaults to 0.0): - Fraction of LM residue tokens randomly replaced with the LM mask - token before running the PLM backbone. - disable_msa_features (`bool`, defaults to False): - When True, zero out MSA-derived ``profile`` and ``deletion_mean`` - before the inputs embedder (experimental medium/large checkpoints). - inputs (`InputsEmbedderConfig`): - Configuration for the inputs embedder module. - folding_trunk (`FoldingTrunkConfig`): - Configuration for the folding trunk. - structure_head (`DiffusionStructureHeadConfig`): - Configuration for the diffusion-based structure prediction head. - confidence_head (`ConfidenceHeadConfig`): - Configuration for the confidence prediction head. - - Examples: - - ```python - >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel - - >>> # Initializing an ESMFold2 configuration - >>> configuration = ESMFold2Config(type="experimental") - - >>> # Initializing a model (with random weights) from the configuration - >>> model = ESMFold2ExperimentalModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - """ - - model_type = "esmfold2" - has_no_defaults_at_init = True - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - self.type: str = kwargs.get("type", "release") - if self.type not in ("release", "experimental"): - raise ValueError( - f"ESMFold2Config.type must be 'release' or 'experimental', " - f"got {self.type!r}" - ) - - # Top-level scalar fields - self.d_single: int = kwargs.get("d_single", 384) - self.d_pair: int = kwargs.get("d_pair", 256) - self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32) - self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2) - self.num_loops: int = kwargs.get("num_loops", 10) - self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8) - # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs - # embedder. - self.disable_msa_features: bool = kwargs.get("disable_msa_features", False) - self.lm_dropout: float = kwargs.get("lm_dropout", 0.0) - self.force_lm_dropout_during_inference: bool = kwargs.get( - "force_lm_dropout_during_inference", False - ) - self.lm_mask_pct: float = kwargs.get("lm_mask_pct", 0.0) - - self.lm_d_model: int = kwargs.get("lm_d_model", 2560) - self.lm_num_layers: int = kwargs.get("lm_num_layers", 80) - # Backward-compatible field name; values now point to FastPLMs ESM++. - raw_esmc_id = ( - kwargs["esmc_id"] if "esmc_id" in kwargs else _DEFAULT_ESMC_HF_REPO - ) - self.esmc_id: str = normalize_esmc_id(raw_esmc_id) - self.esmc_attn_backend: str = ( - kwargs["esmc_attn_backend"] - if "esmc_attn_backend" in kwargs - else _DEFAULT_ESMC_ATTN_BACKEND - ) - - def _init_nested(cls, val): - if isinstance(val, cls): - return val - if isinstance(val, dict): - return cls(**val) - return cls() - - self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs")) - self.folding_trunk = _init_nested( - FoldingTrunkConfig, kwargs.get("folding_trunk") - ) - self.structure_head = _init_nested( - DiffusionStructureHeadConfig, kwargs.get("structure_head") - ) - self.confidence_head = _init_nested( - ConfidenceHeadConfig, kwargs.get("confidence_head") - ) - self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder")) - # Release-only modules — ignored when ``type == "experimental"``. - self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae")) - self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder")) - # If True, MSA encoder output replaces the pair stream; if False, it is added. - self.msa_encoder_overwrite: bool = bool( - kwargs.get("msa_encoder_overwrite", True) - ) - - def to_dict(self): - output = super().to_dict() - output["inputs"] = asdict(self.inputs) - output["folding_trunk"] = asdict(self.folding_trunk) - output["structure_head"] = asdict(self.structure_head) - output["confidence_head"] = asdict(self.confidence_head) - output["msa_encoder"] = asdict(self.msa_encoder) - output["parcae"] = asdict(self.parcae) - output["lm_encoder"] = asdict(self.lm_encoder) - return output - - -__all__ = [ - "ESMFold2Config", - "MSAEncoderConfig", - "ParcaeConfig", - "LMEncoderConfig", - "normalize_esmc_id", -] diff --git a/fastplms/esmfold2/esmfold2_affine3d.py b/fastplms/esmfold2/esmfold2_affine3d.py deleted file mode 100644 index 799c53d..0000000 --- a/fastplms/esmfold2/esmfold2_affine3d.py +++ /dev/null @@ -1,560 +0,0 @@ -from __future__ import annotations - -import typing as T -from abc import ABC -from dataclasses import dataclass - -import torch -from torch.nn import functional as F -from typing_extensions import Self - -from .esmfold2_misc import fp32_autocast_context - - -class Rotation(ABC): - @classmethod - def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... - - @classmethod - def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... - - def __getitem__(self, idx: T.Any) -> Self: ... - - @property - def tensor(self) -> torch.Tensor: - # We claim that this should be zero-cost abstraction that returns the raw tensor backing this - # object. The raw tensor should always have exactly 1 more dim than self.shape, which should be - # implemented using reshaping - ... - - @property - def shape(self) -> torch.Size: - # The "shape" of the rotation, as if it was a torch.tensor object - # This means that 1x4 quaternions are treated as size (1,) for example - ... - - def as_matrix(self) -> RotationMatrix: ... - - def as_quat(self, normalize: bool = False) -> RotationQuat: ... - - def compose(self, other: Self) -> Self: - # To be safe, we force users to explicitly convert between rotation types. - ... - - def convert_compose(self, other: Self) -> Self: - # This function will automatically convert between types of rotations - ... - - def apply(self, p: torch.Tensor) -> torch.Tensor: - # rotates points by this rotation object - ... - - def invert(self) -> Self: ... - - @property - def dtype(self) -> torch.dtype: - return self.tensor.dtype - - @property - def device(self) -> torch.device: - return self.tensor.device - - @property - def requires_grad(self) -> bool: - return self.tensor.requires_grad - - @classmethod - def _from_tensor(cls, t: torch.Tensor) -> Self: - # This function exists to simplify the below functions, esp type signatures - # Its implementation is different from Affine3D.from_tensor and does not - # autodetect rotation types. - return cls(t) # type: ignore - - def to(self, **kwargs) -> Self: - return self._from_tensor(self.tensor.to(**kwargs)) - - def detach(self, *args, **kwargs) -> Self: - return self._from_tensor(self.tensor.detach(**kwargs)) - - def tensor_apply(self, func) -> Self: - # Applys a function to the underlying tensor - return self._from_tensor( - torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1) - ) - - -class RotationMatrix(Rotation): - def __init__(self, rots: torch.Tensor): - if rots.shape[-1] == 9: - rots = rots.unflatten(-1, (3, 3)) - assert rots.shape[-1] == 3 - assert rots.shape[-2] == 3 - # Force full precision - rots = rots.to(torch.float32) - self._rots = rots - - @classmethod - def identity(cls, shape, **tensor_kwargs): - rots = torch.eye(3, **tensor_kwargs) - rots = rots.view(*[1 for _ in range(len(shape))], 3, 3) - rots = rots.expand(*shape, -1, -1) - return cls(rots) - - @classmethod - def random(cls, shape, **tensor_kwargs): - return RotationQuat.random(shape, **tensor_kwargs).as_matrix() - - def __getitem__(self, idx: T.Any) -> RotationMatrix: - indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) - return RotationMatrix(self._rots[indices + (slice(None), slice(None))]) - - @property - def shape(self) -> torch.Size: - return self._rots.shape[:-2] - - def as_matrix(self) -> RotationMatrix: - return self - - def as_quat(self, normalize: bool = False) -> RotationQuat: - m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( - self._rots.flatten(-2), dim=-1 - ) - q_abs = _sqrt_subgradient( - torch.stack( - [ - 1.0 + m00 + m11 + m22, - 1.0 + m00 - m11 - m22, - 1.0 - m00 + m11 - m22, - 1.0 - m00 - m11 + m22, - ], - dim=-1, - ) - ) - # we produce the desired quaternion multiplied by each of r, i, j, k - quat_by_rijk = torch.stack( - [ - x - for lst in [ - [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], - [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], - [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], - [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], - ] - for x in lst - ], - dim=-1, - ).unflatten(-1, (4, 4)) - - # We floor here at 0.1 but the exact level is not important; if q_abs is small, - # the candidate won't be picked. - flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) - quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) - - # if not for numerical problems, quat_candidates[i] should be same (up to a sign), - # forall i; we pick the best-conditioned one (with the largest denominator) - # We manually implement one_hot so torch.compile works - one_hot = torch.zeros_like(q_abs, dtype=torch.bool) - one_hot.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True) - quat = quat_candidates[one_hot, :].reshape(q_abs.shape) - return RotationQuat(quat) - - def compose(self, other: RotationMatrix) -> RotationMatrix: - with fp32_autocast_context(self._rots.device.type): - return RotationMatrix(self._rots @ other._rots) - - def convert_compose(self, other: Rotation): - return self.compose(other.as_matrix()) - - def apply(self, p: torch.Tensor) -> torch.Tensor: - with fp32_autocast_context(self.device.type): - if self._rots.shape[-3] == 1: - # This is a slight speedup over einsum for batched rotations - return p @ self._rots.transpose(-1, -2).squeeze(-3) - else: - # einsum way faster than bmm! - return torch.einsum("...ij,...j", self._rots, p) - - def invert(self) -> RotationMatrix: - return RotationMatrix(self._rots.transpose(-1, -2)) - - @property - def tensor(self) -> torch.Tensor: - return self._rots.flatten(-2) - - def to_3x3(self) -> torch.Tensor: - return self._rots - - @staticmethod - def from_graham_schmidt( - x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12 - ) -> RotationMatrix: - # A low eps here is necessary for good stability! - return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps)) - - -class RotationQuat(Rotation): - def __init__(self, quats: torch.Tensor, normalized=False): - assert quats.shape[-1] == 4 - self._normalized = normalized - # Force float32 as well - if normalized: - self._quats = F.normalize(quats.to(torch.float32), dim=-1) - self._quats = self._quats.where(self._quats[..., :1] >= 0, -self._quats) - else: - self._quats = quats.to(torch.float32) - - @classmethod - def identity(cls, shape, **tensor_kwargs): - q = torch.ones((*shape, 4), **tensor_kwargs) - mult = torch.tensor([1, 0, 0, 0], device=q.device) - return RotationQuat(q * mult) - - @classmethod - def random(cls, shape, **tensor_kwargs): - quat = torch.randn((*shape, 4), **tensor_kwargs) - return RotationQuat(quat, normalized=True) - - def __getitem__(self, idx: T.Any) -> RotationQuat: - indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) - return RotationQuat(self._quats[indices + (slice(None),)]) - - @property - def shape(self) -> torch.Size: - return self._quats.shape[:-1] - - def compose(self, other: RotationQuat) -> RotationQuat: - with fp32_autocast_context(self._quats.device.type): - return RotationQuat(_quat_mult(self._quats, other._quats)) - - def convert_compose(self, other: Rotation): - return self.compose(other.as_quat()) - - def as_matrix(self) -> RotationMatrix: - q = self.normalized().tensor - r, i, j, k = torch.unbind(q, -1) - two_s = 2.0 / torch.linalg.norm(q, dim=-1) - - o = torch.stack( - ( - 1 - two_s * (j * j + k * k), - two_s * (i * j - k * r), - two_s * (i * k + j * r), - two_s * (i * j + k * r), - 1 - two_s * (i * i + k * k), - two_s * (j * k - i * r), - two_s * (i * k - j * r), - two_s * (j * k + i * r), - 1 - two_s * (i * i + j * j), - ), - -1, - ) - return RotationMatrix(o.reshape(q.shape[:-1] + (3, 3))) - - def as_quat(self, normalize: bool = False) -> RotationQuat: - return self - - def apply(self, p: torch.Tensor) -> torch.Tensor: - return _quat_rotation(self.normalized()._quats, p) - - def invert(self) -> RotationQuat: - return RotationQuat(_quat_invert(self._quats)) - - @property - def tensor(self) -> torch.Tensor: - return self._quats - - def normalized(self) -> RotationQuat: - return self if self._normalized else RotationQuat(self._quats, normalized=True) - - -@dataclass(frozen=True) -class Affine3D: - trans: torch.Tensor - rot: Rotation - - def __post_init__(self): - assert self.trans.shape[:-1] == self.rot.shape - - @staticmethod - def identity( - shape_or_affine: T.Union[tuple[int, ...], "Affine3D"], - rotation_type: T.Type[Rotation] = RotationMatrix, - **tensor_kwargs, - ): - # Creates a new identity Affine3D object with a specified shape - # or the same shape as another Affine3D object. - if isinstance(shape_or_affine, Affine3D): - kwargs = {"dtype": shape_or_affine.dtype, "device": shape_or_affine.device} - kwargs.update(tensor_kwargs) - shape = shape_or_affine.shape - rotation_type = type(shape_or_affine.rot) - else: - kwargs = tensor_kwargs - shape = shape_or_affine - return Affine3D( - torch.zeros((*shape, 3), **kwargs), rotation_type.identity(shape, **kwargs) - ) - - @staticmethod - def random( - shape: tuple[int, ...], - std: float = 1, - rotation_type: T.Type[Rotation] = RotationMatrix, - **tensor_kwargs, - ) -> "Affine3D": - return Affine3D( - trans=torch.randn((*shape, 3), **tensor_kwargs).mul(std), - rot=rotation_type.random(shape, **tensor_kwargs), - ) - - def __getitem__(self, idx: T.Any) -> "Affine3D": - indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) - return Affine3D(trans=self.trans[indices + (slice(None),)], rot=self.rot[idx]) - - @property - def shape(self) -> torch.Size: - return self.trans.shape[:-1] - - @property - def dtype(self) -> torch.dtype: - return self.trans.dtype - - @property - def device(self) -> torch.device: - return self.trans.device - - @property - def requires_grad(self) -> bool: - return self.trans.requires_grad - - def to(self, **kwargs) -> "Affine3D": - return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs)) - - def detach(self, *args, **kwargs) -> "Affine3D": - return Affine3D(self.trans.detach(**kwargs), self.rot.detach(**kwargs)) - - def tensor_apply(self, func) -> "Affine3D": - # Applys a function to the underlying tensor - return self.from_tensor( - torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1) - ) - - def as_matrix(self): - return Affine3D(trans=self.trans, rot=self.rot.as_matrix()) - - def as_quat(self, normalize: bool = False): - return Affine3D(trans=self.trans, rot=self.rot.as_quat(normalize)) - - def compose(self, other: "Affine3D", autoconvert: bool = False): - rot = self.rot - new_rot = (rot.convert_compose if autoconvert else rot.compose)(other.rot) - new_trans = rot.apply(other.trans) + self.trans - return Affine3D(trans=new_trans, rot=new_rot) - - def compose_rotation(self, other: Rotation, autoconvert: bool = False): - return Affine3D( - trans=self.trans, - rot=(self.rot.convert_compose if autoconvert else self.rot.compose)(other), - ) - - def scale(self, v: torch.Tensor | float): - return Affine3D(self.trans * v, self.rot) - - def mask(self, mask: torch.Tensor, with_zero=False): - # Returns a transform where True positions in mask is identity - if with_zero: - tensor = self.tensor - return Affine3D.from_tensor( - torch.zeros_like(tensor).where(mask[..., None], tensor) - ) - else: - identity = self.identity( - self.shape, - rotation_type=type(self.rot), - device=self.device, - dtype=self.dtype, - ).tensor - return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor)) - - def apply(self, p: torch.Tensor) -> torch.Tensor: - return self.rot.apply(p) + self.trans - - def invert(self): - inv_rot = self.rot.invert() - return Affine3D(trans=-inv_rot.apply(self.trans), rot=inv_rot) - - @property - def tensor(self) -> torch.Tensor: - return torch.cat([self.rot.tensor, self.trans], dim=-1) - - @staticmethod - def from_tensor(t: torch.Tensor) -> "Affine3D": - match t.shape[-1]: - case 4: - # Assume tensor 4x4 for backward compat with alphafold - trans = t[..., :3, 3] - rot = RotationMatrix(t[..., :3, :3]) - case 6: - # Assume quaternion representation with real part = 1 - trans = t[..., -3:] - rot = RotationQuat(F.pad(t[..., :3], (1, 0), value=1)) - case 7: - trans = t[..., -3:] - rot = RotationQuat(t[..., :4]) - case 12: - trans = t[..., -3:] - rot = RotationMatrix(t[..., :-3].unflatten(-1, (3, 3))) - case _: - raise RuntimeError( - f"Cannot detect rotation fromat from {t.shape[-1] -3}-d flat vector" - ) - return Affine3D(trans, rot) - - @staticmethod - def from_tensor_pair(t: torch.Tensor, r: torch.Tensor) -> "Affine3D": - return Affine3D(t, RotationMatrix(r)) - - @staticmethod - def from_graham_schmidt( - neg_x_axis: torch.Tensor, - origin: torch.Tensor, - xy_plane: torch.Tensor, - eps: float = 1e-10, - ): - # The arguments of this function is for parity with AlphaFold - x_axis = origin - neg_x_axis - xy_plane = xy_plane - origin - return Affine3D( - trans=origin, rot=RotationMatrix.from_graham_schmidt(x_axis, xy_plane, eps) - ) - - @staticmethod - def cat(affines: list["Affine3D"], dim: int = 0): - if dim < 0: - dim = len(affines[0].shape) + dim - return Affine3D.from_tensor(torch.cat([x.tensor for x in affines], dim=dim)) - - -def _quat_mult(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """ - Multiply two quaternions. - Usual torch rules for broadcasting apply. - - Args: - a: Quaternions as tensor of shape (..., 4), real part first. - b: Quaternions as tensor of shape (..., 4), real part first. - - Returns: - The product of a and b, a tensor of quaternions shape (..., 4). - """ - aw, ax, ay, az = torch.unbind(a, -1) - bw, bx, by, bz = torch.unbind(b, -1) - ow = aw * bw - ax * bx - ay * by - az * bz - ox = aw * bx + ax * bw + ay * bz - az * by - oy = aw * by - ax * bz + ay * bw + az * bx - oz = aw * bz + ax * by - ay * bx + az * bw - return torch.stack((ow, ox, oy, oz), -1) - - -def _quat_rotation(q: torch.Tensor, p: torch.Tensor) -> torch.Tensor: - """ - Rotates p by quaternion q. Usual torch rules for broadcasting apply. - - Args: - q: Quaternions as tensor of shape (..., 4), real part first. - p: Points as tensor of shape (..., 3) - - Returns: - The rotated version of p, of shape (..., 3) - """ - aw, ax, ay, az = torch.unbind(q, -1) - bx, by, bz = torch.unbind(p, -1) - # fmt: off - ow = - ax * bx - ay * by - az * bz - ox = aw * bx + ay * bz - az * by - oy = aw * by - ax * bz + az * bx - oz = aw * bz + ax * by - ay * bx - # fmt: on - q_mul_pts = torch.stack((ow, ox, oy, oz), -1) - return _quat_mult(q_mul_pts, _quat_invert(q))[..., 1:] - - -def _quat_invert(q: torch.Tensor): - return q * torch.tensor([1, -1, -1, -1], device=q.device) - - -def _sqrt_subgradient(x: torch.Tensor) -> torch.Tensor: - # Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0. - ret = torch.zeros_like(x) - positive_mask = x > 0 - ret[positive_mask] = torch.sqrt(x[positive_mask]) - return ret - - -def _graham_schmidt(x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12): - # A low eps here is necessary for good stability! - with fp32_autocast_context(x_axis.device.type): - e1 = xy_plane - - denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps) - x_axis = x_axis / denom - dot = (x_axis * e1).sum(dim=-1, keepdim=True) - e1 = e1 - x_axis * dot - denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps) - e1 = e1 / denom - e2 = torch.cross(x_axis, e1, dim=-1) - - rots = torch.stack([x_axis, e1, e2], dim=-1) - - return rots - - -def build_affine3d_from_coordinates( - coords: torch.Tensor, # (N, CA, C). -) -> tuple[Affine3D, torch.Tensor]: - _MAX_SUPPORTED_DISTANCE = 1e6 - coord_mask = torch.all( - torch.all(torch.isfinite(coords) & (coords < _MAX_SUPPORTED_DISTANCE), dim=-1), - dim=-1, - ) - - def atom3_to_backbone_affine(bb_positions: torch.Tensor) -> Affine3D: - N, CA, C = bb_positions.unbind(dim=-2) - return Affine3D.from_graham_schmidt(C, CA, N) - - coords = coords.clone().float() - coords[~coord_mask] = 0 - - # NOTE(thayes): If you have already normalized the coordinates, then - # the black hole affine translations will be zeros and the rotations will be - # the identity. - average_per_n_ca_c = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / ( - coord_mask.sum(-1)[..., None, None] + 1e-8 - ) - affine_from_average = atom3_to_backbone_affine( - average_per_n_ca_c.float() - ).as_matrix() - - B, S, _, _ = coords.shape - assert isinstance(B, int) - assert isinstance(S, int) - affine_rot_mats = affine_from_average.rot.tensor[..., None, :].expand(B, S, 9) - affine_trans = affine_from_average.trans[..., None, :].expand(B, S, 3) - - # We use the identity rotation whereever we have no coordinates. This is - # important because otherwise the rotation matrices will be all zeros, which - # will cause collapse in the distance/direction attention mechanism. - identity_rot = RotationMatrix.identity( - (B, S), dtype=torch.float32, device=coords.device, requires_grad=False - ) - affine_rot_mats = affine_rot_mats.where( - coord_mask.any(-1)[..., None, None], identity_rot.tensor - ) - black_hole_affine = Affine3D(affine_trans, RotationMatrix(affine_rot_mats)) - - affine = atom3_to_backbone_affine(coords.float()) - affine = Affine3D.from_tensor( - affine.tensor.where(coord_mask[..., None], black_hole_affine.tensor) - ) - - return affine, coord_mask diff --git a/fastplms/esmfold2/esmfold2_aligner.py b/fastplms/esmfold2/esmfold2_aligner.py deleted file mode 100644 index 26a3d31..0000000 --- a/fastplms/esmfold2/esmfold2_aligner.py +++ /dev/null @@ -1,101 +0,0 @@ -from __future__ import annotations - -from dataclasses import Field, replace -from typing import Any, ClassVar, Protocol, TypeVar - -import numpy as np -import torch - -from .esmfold2_protein_structure import compute_affine_and_rmsd - - -class Alignable(Protocol): - # Trick to detect whether an object is a dataclass - __dataclass_fields__: ClassVar[dict[str, Field[Any]]] - - @property - def atom37_positions(self) -> np.ndarray: # type: ignore - pass - - @property - def atom37_mask(self) -> np.ndarray: # type: ignore - pass - - def __len__(self) -> int: ... - - -T = TypeVar("T", bound=Alignable) - - -class Aligner: - def __init__( - self, - mobile: Alignable, - target: Alignable, - only_use_backbone: bool = False, - use_reflection: bool = False, - ): - """ - Aligns a mobile protein chain against a target protein chain. - - Args: - mobile (ProteinChain): Protein chain to be aligned. - target (ProteinChain): Protein chain target. - only_use_backbone (bool): Whether to only use backbone atoms. - use_reflection (bool): Whether to align to target reflection. - """ - # Check proteins must have same number of residues - assert len(mobile) == len(target) - - # Determine overlapping atoms - joint_atom37_mask = mobile.atom37_mask.astype(bool) & target.atom37_mask.astype( - bool - ) - - # Backbone atoms are first sites in atom37 representation - if only_use_backbone: - joint_atom37_mask[:, 3:] = False - - # Extract matching atom positions and convert to batched tensors - mobile_atom_tensor = ( - torch.from_numpy(mobile.atom37_positions).type(torch.double).unsqueeze(0) - ) - target_atom_tensor = ( - torch.from_numpy(target.atom37_positions).type(torch.double).unsqueeze(0) - ) - joint_atom37_mask = ( - torch.from_numpy(joint_atom37_mask).type(torch.bool).unsqueeze(0) - ) - - # If using reflection flip target - if use_reflection: - target_atom_tensor = -target_atom_tensor - - # Compute alignment and rmsd - affine3D, rmsd = compute_affine_and_rmsd( - mobile_atom_tensor, target_atom_tensor, atom_exists_mask=joint_atom37_mask - ) - self._affine3D = affine3D - self._rmsd = rmsd.item() - - @property - def rmsd(self): - return self._rmsd - - def apply(self, mobile: T) -> T: - """Apply alignment to a protein chain""" - # Extract atom positions and convert to batched tensors - mobile_atom_tensor = ( - torch.from_numpy(mobile.atom37_positions[mobile.atom37_mask]) - .type(torch.float32) - .unsqueeze(0) - ) - - # Transform atom arrays - aligned_atom_tensor = self._affine3D.apply(mobile_atom_tensor).squeeze(0) - - # Rebuild atom37 positions - aligned_atom37_positions = np.full_like(mobile.atom37_positions, np.nan) - aligned_atom37_positions[mobile.atom37_mask] = aligned_atom_tensor - - return replace(mobile, atom37_positions=aligned_atom37_positions) diff --git a/fastplms/esmfold2/esmfold2_atom_indexer.py b/fastplms/esmfold2/esmfold2_atom_indexer.py deleted file mode 100644 index e135c33..0000000 --- a/fastplms/esmfold2/esmfold2_atom_indexer.py +++ /dev/null @@ -1,15 +0,0 @@ -import numpy as np - -from .esmfold2_protein_structure import index_by_atom_name - - -class AtomIndexer: - def __init__(self, structure, property: str, dim: int): - self.structure = structure - self.property = property - self.dim = dim - - def __getitem__(self, atom_names: str | list[str]) -> np.ndarray: - return index_by_atom_name( - getattr(self.structure, self.property), atom_names, self.dim - ) diff --git a/fastplms/esmfold2/esmfold2_conformers.py b/fastplms/esmfold2/esmfold2_conformers.py deleted file mode 100644 index 6399b3c..0000000 --- a/fastplms/esmfold2/esmfold2_conformers.py +++ /dev/null @@ -1,291 +0,0 @@ -"""CCD conformer loading utilities. - -Loads idealized conformer coordinates from a CCD pickle file containing RDKit molecules. -Conformer priority follows AF3 Section 2.8: Computed > Ideal > first available. -""" - -from __future__ import annotations - -import os -import pickle -from pathlib import Path - -import numpy as np -from huggingface_hub import hf_hub_download - -from .esmfold2_constants import RES_TYPE_TO_CCD - -if os.environ.get("ESMCFOLD_CCD_PATH"): - CCD_PICKLE_PATH = Path(os.environ["ESMCFOLD_CCD_PATH"]) -else: - CCD_PICKLE_PATH = None - - -# Lazily loaded CCD dictionary -_CCD_MOLECULES: dict | None = None - -# Caches -_CCD_CONFORMERS: dict[str, dict[str, np.ndarray]] = {} -_CCD_ATOM_CACHE: dict[str, list[tuple[str, str, int]]] = {} -_CCD_BONDS_CACHE: dict[str, list[tuple[str, str]]] = {} -_CCD_LEAVING_ATOMS_CACHE: dict[str, set[str]] = {} -_IDEALIZED_POS_CACHE: dict[tuple[int, str], np.ndarray | None] = {} -_LIGAND_IDEALIZED_POS_CACHE: dict[tuple[str, str], np.ndarray | None] = {} - - -def load_ccd(cache_dir: Path | str | None = None) -> dict: - """Load CCD molecules from pickle file, downloading if needed. - - Args: - cache_dir: Directory to cache the downloaded CCD pickle. - If None, uses CCD_PICKLE_PATH env var or downloads to ~/.cache/esmcfold/. - """ - global _CCD_MOLECULES - if _CCD_MOLECULES is not None: - return _CCD_MOLECULES - - # Determine pickle path - if CCD_PICKLE_PATH is not None and CCD_PICKLE_PATH.exists(): - pkl_path = CCD_PICKLE_PATH - elif cache_dir is not None: - cache_dir = Path(cache_dir) - cache_dir.mkdir(parents=True, exist_ok=True) - pkl_path = cache_dir / "ccd.pkl" - else: - try: - pkl_path = Path( - hf_hub_download(repo_id="biohub/ESMFold2", filename="ccd.pkl") - ) - except Exception as e: - raise FileNotFoundError( - f"Failed to download CCD pickle file from Hugging Face repository: {e}" - ) - - if not pkl_path.exists(): - raise FileNotFoundError( - f"CCD pickle file not found: {pkl_path}. Please set the ESMCFOLD_CCD_PATH environment variable to the path of a valid CCD pickle file or download the file from the Hugging Face repository." - ) - - print(f"Loading CCD dictionary from {pkl_path}") - with open(pkl_path, "rb") as f: - _CCD_MOLECULES = pickle.load(f) - - if _CCD_MOLECULES is None: - _CCD_MOLECULES = {} - - return _CCD_MOLECULES - - -def _get_ccd_molecules() -> dict: - """Get CCD molecules, loading lazily on first call.""" - global _CCD_MOLECULES - if _CCD_MOLECULES is None: - return load_ccd() - return _CCD_MOLECULES - - -def _get_ccd_mol_with_significant_h(comp_id: str): - """Get CCD molecule with only chemically significant hydrogens. - - Returns (mol, conformer) tuple or (None, None) if not available. - """ - ccd = _get_ccd_molecules() - if comp_id not in ccd: - return None, None - - mol = ccd[comp_id] - if mol.GetNumConformers() == 0: - return None, None - - # Find the "Computed" conformer (RDKit ETKDGv3), fall back to "Ideal" - conf_idx = 0 - for i, c in enumerate(mol.GetConformers()): - props = c.GetPropsAsDict() - if props.get("name") == "Computed": - conf_idx = i - break - else: - for i, c in enumerate(mol.GetConformers()): - props = c.GetPropsAsDict() - if props.get("name") == "Ideal": - conf_idx = i - break - - from rdkit import Chem - - mol_no_h = Chem.RemoveHs(mol, sanitize=False) - - if mol_no_h.GetNumConformers() == 0: - return None, None - - return mol_no_h, mol_no_h.GetConformer( - min(conf_idx, mol_no_h.GetNumConformers() - 1) - ) - - -def get_ccd_conformer(comp_id: str) -> dict[str, np.ndarray] | None: - """Get idealized conformer as dict of atom_name -> position [3]. - - Conformer priority: Computed > Ideal > first available. - """ - if comp_id in _CCD_CONFORMERS: - cached = _CCD_CONFORMERS[comp_id] - return cached if cached else None - - mol, conf = _get_ccd_mol_with_significant_h(comp_id) - if mol is None or conf is None: - _CCD_CONFORMERS[comp_id] = {} - return None - - conformer: dict[str, np.ndarray] = {} - for atom in mol.GetAtoms(): - props = atom.GetPropsAsDict() - atom_name = props.get("name") - if not isinstance(atom_name, str) or not atom_name: - continue - idx = atom.GetIdx() - pos = conf.GetAtomPosition(idx) - conformer[atom_name] = np.array([pos.x, pos.y, pos.z], dtype=np.float32) - - _CCD_CONFORMERS[comp_id] = conformer - return conformer if conformer else None - - -def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None: - """Get idealized position for a standard residue atom. - - Uses res_type index to look up CCD component, then returns position. - Returns None if not found. - """ - cache_key = (res_type, atom_name) - if cache_key in _IDEALIZED_POS_CACHE: - return _IDEALIZED_POS_CACHE[cache_key] - - comp_id = RES_TYPE_TO_CCD.get(res_type) - if comp_id: - ccd_conformer = get_ccd_conformer(comp_id) - if ccd_conformer and atom_name in ccd_conformer: - pos = ccd_conformer[atom_name] - _IDEALIZED_POS_CACHE[cache_key] = pos - return pos - - _IDEALIZED_POS_CACHE[cache_key] = None - return None - - -def get_ligand_idealized_atom_pos(res_name: str, atom_name: str) -> np.ndarray | None: - """Get idealized position for a ligand/modified residue atom. - - Returns None if not found. - """ - cache_key = (res_name, atom_name) - if cache_key in _LIGAND_IDEALIZED_POS_CACHE: - return _LIGAND_IDEALIZED_POS_CACHE[cache_key] - - ccd_conformer = get_ccd_conformer(res_name) - if ccd_conformer and atom_name in ccd_conformer: - pos = ccd_conformer[atom_name] - _LIGAND_IDEALIZED_POS_CACHE[cache_key] = pos - return pos - - _LIGAND_IDEALIZED_POS_CACHE[cache_key] = None - return None - - -def get_ligand_ccd_atoms_with_charges( - comp_id: str, -) -> list[tuple[str, str, int]] | None: - """Get list of (atom_name, element, charge) for a CCD component. - - Uses RDKit RemoveHs(sanitize=False) to keep chemically significant hydrogens. - Returns None if CCD data not available. - """ - if comp_id in _CCD_ATOM_CACHE: - cached = _CCD_ATOM_CACHE[comp_id] - return cached if cached else None - - mol, _ = _get_ccd_mol_with_significant_h(comp_id) - if mol is None: - _CCD_ATOM_CACHE[comp_id] = [] - return None - - atoms: list[tuple[str, str, int]] = [] - for atom in mol.GetAtoms(): - props = atom.GetPropsAsDict() - atom_name = props.get("name") - if not isinstance(atom_name, str) or not atom_name: - continue - element = atom.GetSymbol() - charge = atom.GetFormalCharge() - atoms.append((atom_name, element, charge)) - - _CCD_ATOM_CACHE[comp_id] = atoms - return atoms if atoms else None - - -def get_ligand_ccd_bonds(comp_id: str) -> list[tuple[str, str]] | None: - """Get list of (atom1_name, atom2_name) bonds for a CCD component. - - Returns None if CCD data not available. - """ - if comp_id in _CCD_BONDS_CACHE: - cached = _CCD_BONDS_CACHE[comp_id] - return cached if cached else None - - mol, _ = _get_ccd_mol_with_significant_h(comp_id) - if mol is None: - _CCD_BONDS_CACHE[comp_id] = [] - return None - - # Get included atom names - included_atoms = set() - for atom in mol.GetAtoms(): - props = atom.GetPropsAsDict() - atom_name = props.get("name") - if isinstance(atom_name, str) and atom_name: - included_atoms.add(atom_name) - - bonds: list[tuple[str, str]] = [] - for bond in mol.GetBonds(): - a1 = bond.GetBeginAtom() - a2 = bond.GetEndAtom() - n1 = a1.GetPropsAsDict().get("name") - n2 = a2.GetPropsAsDict().get("name") - if ( - isinstance(n1, str) - and isinstance(n2, str) - and n1 - and n2 - and n1 in included_atoms - and n2 in included_atoms - ): - bonds.append((n1, n2)) - - _CCD_BONDS_CACHE[comp_id] = bonds - return bonds if bonds else None - - -def get_ccd_leaving_atoms(comp_id: str) -> set[str]: - """Get set of atom names marked as leaving atoms in CCD. - - Leaving atoms are removed during polymerization (e.g., OP3 in nucleotides). - """ - if comp_id in _CCD_LEAVING_ATOMS_CACHE: - return _CCD_LEAVING_ATOMS_CACHE[comp_id] - - ccd = _get_ccd_molecules() - if comp_id not in ccd: - _CCD_LEAVING_ATOMS_CACHE[comp_id] = set() - return set() - - mol = ccd[comp_id] - leaving_atoms = set() - for atom in mol.GetAtoms(): - if atom.HasProp("leaving_atom"): - if atom.GetProp("leaving_atom") == "1": - name = atom.GetProp("name") if atom.HasProp("name") else "" - if name: - leaving_atoms.add(name) - - _CCD_LEAVING_ATOMS_CACHE[comp_id] = leaving_atoms - return leaving_atoms diff --git a/fastplms/esmfold2/esmfold2_constants.py b/fastplms/esmfold2/esmfold2_constants.py deleted file mode 100644 index eb41c5d..0000000 --- a/fastplms/esmfold2/esmfold2_constants.py +++ /dev/null @@ -1,562 +0,0 @@ -"""Constants for the ESMFold2 input pipeline. - -Includes molecule types, residue types, vocabularies, atom lists, and element data. -""" - -# ============================================================================= -# Molecule types -# ============================================================================= - -MOL_TYPE_PROTEIN = 0 -MOL_TYPE_DNA = 1 -MOL_TYPE_RNA = 2 -MOL_TYPE_NONPOLYMER = 3 - -# ============================================================================= -# Residue type indices -# ============================================================================= - -# Standard amino acids (indices 2-21), MSE mapped to MET -PROTEIN_RESIDUE_TO_RES_TYPE = { - "ALA": 2, - "ARG": 3, - "ASN": 4, - "ASP": 5, - "CYS": 6, - "GLN": 7, - "GLU": 8, - "GLY": 9, - "HIS": 10, - "ILE": 11, - "LEU": 12, - "LYS": 13, - "MET": 14, - "PHE": 15, - "PRO": 16, - "SER": 17, - "THR": 18, - "TRP": 19, - "TYR": 20, - "VAL": 21, - "MSE": 14, # Selenomethionine -> MET -} -PROTEIN_UNK_RES_TYPE = 22 - -# RNA nucleotides (indices 23-26, unknown=27) -RNA_RESIDUE_TO_RES_TYPE = {"A": 23, "G": 24, "C": 25, "U": 26} -RNA_UNK_RES_TYPE = 27 - -# DNA nucleotides (indices 28-31, unknown=32) -DNA_RESIDUE_TO_RES_TYPE = {"DA": 28, "DG": 29, "DC": 30, "DT": 31} -DNA_UNK_RES_TYPE = 32 - -GAP_RES_TYPE = 32 - -# ============================================================================= -# Vocabularies -# ============================================================================= - -# 3-letter to 1-letter codes for proteins -PROTEIN_3TO1 = { - "ALA": "A", - "ARG": "R", - "ASN": "N", - "ASP": "D", - "CYS": "C", - "GLN": "Q", - "GLU": "E", - "GLY": "G", - "HIS": "H", - "ILE": "I", - "LEU": "L", - "LYS": "K", - "MET": "M", - "PHE": "F", - "PRO": "P", - "SER": "S", - "THR": "T", - "TRP": "W", - "TYR": "Y", - "VAL": "V", - "MSE": "M", -} - -# 1-letter to 3-letter codes -PROTEIN_1TO3 = {v: k for k, v in PROTEIN_3TO1.items() if k != "MSE"} -PROTEIN_1TO3["X"] = "UNK" - -# DNA 1-letter to CCD code -DNA_1TO3 = {"A": "DA", "T": "DT", "C": "DC", "G": "DG"} - -# RNA 1-letter to CCD code -RNA_1TO3 = {"A": "A", "U": "U", "C": "C", "G": "G"} - -# ESM-2 input_ids vocabulary for proteins -ESM_PROTEIN_VOCAB = { - "L": 4, - "A": 5, - "G": 6, - "V": 7, - "S": 8, - "E": 9, - "R": 10, - "T": 11, - "I": 12, - "D": 13, - "P": 14, - "K": 15, - "Q": 16, - "N": 17, - "F": 18, - "Y": 19, - "M": 20, - "H": 21, - "W": 22, - "C": 23, - "X": 3, # Unknown -} - -# For DNA/RNA/ligands -DNA_RNA_LIGAND_INPUT_ID = 24 - -# MSA tokens -MSA_PAD_TOKEN_ID = 0 -MSA_GAP_TOKEN_ID = 1 # Gap/insertion token for MSA - -# res_type int -> CCD component ID (for conformer lookup) -RES_TYPE_TO_CCD = { - # Proteins (2-22) - 2: "ALA", - 3: "ARG", - 4: "ASN", - 5: "ASP", - 6: "CYS", - 7: "GLN", - 8: "GLU", - 9: "GLY", - 10: "HIS", - 11: "ILE", - 12: "LEU", - 13: "LYS", - 14: "MET", - 15: "PHE", - 16: "PRO", - 17: "SER", - 18: "THR", - 19: "TRP", - 20: "TYR", - 21: "VAL", - 22: "UNK", - # RNA (23-27) - 23: "A", - 24: "G", - 25: "C", - 26: "U", - 27: "N", - # DNA (28-32) - 28: "DA", - 29: "DG", - 30: "DC", - 31: "DT", - 32: "DN", -} - -# ============================================================================= -# Charged atoms at physiological pH -# ============================================================================= - -CHARGED_ATOMS: dict[tuple[str, str], int] = { - ("LYS", "NZ"): 1, - ("ARG", "NH2"): 1, - ("HIS", "ND1"): 1, - ("PO4", "O2"): -1, - ("PO4", "O3"): -1, - ("PO4", "O4"): -1, - ("SO4", "O3"): -1, - ("SO4", "O4"): -1, - ("MG", "MG"): 2, - ("ZN", "ZN"): 2, - ("CA", "CA"): 2, - ("FE2", "FE"): 2, - ("MN", "MN"): 2, - ("CO", "CO"): 2, - ("NCO", "CO"): 3, - ("CU", "CU"): 2, - ("NI", "NI"): 2, - ("K", "K"): 1, - ("NA", "NA"): 1, - ("CD", "CD"): 2, - ("CL", "CL"): -1, - ("ACT", "OXT"): -1, - ("NAD", "O2N"): -1, - ("NAD", "N1N"): 1, - ("NAP", "O2N"): -1, - ("NAP", "N1N"): 1, - ("IMD", "N3"): 1, - ("SAM", "SD"): 1, - ("FE", "FE"): 3, - ("A1BH3", "N3"): 1, -} - -# ============================================================================= -# Element atomic numbers (Z=1 to 92) -# ============================================================================= - -ELEMENT_TO_ATOMIC_NUM = { - "H": 1, - "LI": 3, - "BE": 4, - "B": 5, - "C": 6, - "N": 7, - "O": 8, - "F": 9, - "NE": 10, - "NA": 11, - "MG": 12, - "AL": 13, - "SI": 14, - "P": 15, - "S": 16, - "CL": 17, - "AR": 18, - "K": 19, - "CA": 20, - "SC": 21, - "TI": 22, - "V": 23, - "CR": 24, - "MN": 25, - "FE": 26, - "CO": 27, - "NI": 28, - "CU": 29, - "ZN": 30, - "GA": 31, - "GE": 32, - "AS": 33, - "SE": 34, - "BR": 35, - "KR": 36, - "RB": 37, - "SR": 38, - "Y": 39, - "ZR": 40, - "NB": 41, - "MO": 42, - "TC": 43, - "RU": 44, - "RH": 45, - "PD": 46, - "AG": 47, - "CD": 48, - "IN": 49, - "SN": 50, - "SB": 51, - "TE": 52, - "I": 53, - "XE": 54, - "CS": 55, - "BA": 56, - "LA": 57, - "CE": 58, - "PR": 59, - "ND": 60, - "PM": 61, - "SM": 62, - "EU": 63, - "GD": 64, - "TB": 65, - "DY": 66, - "HO": 67, - "ER": 68, - "TM": 69, - "YB": 70, - "LU": 71, - "HF": 72, - "TA": 73, - "W": 74, - "RE": 75, - "OS": 76, - "IR": 77, - "PT": 78, - "AU": 79, - "HG": 80, - "TL": 81, - "PB": 82, - "BI": 83, - "PO": 84, - "AT": 85, - "RN": 86, - "FR": 87, - "RA": 88, - "AC": 89, - "TH": 90, - "PA": 91, - "U": 92, -} - -# Inverse mapping: atomic number → element symbol -ELEMENT_NUMBER_TO_SYMBOL = {v: k for k, v in ELEMENT_TO_ATOMIC_NUM.items()} - -# ============================================================================= -# Standard heavy atoms per residue type -# ============================================================================= - -PROTEIN_HEAVY_ATOMS = { - "ALA": ["N", "CA", "C", "O", "CB"], - "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"], - "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"], - "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"], - "CYS": ["N", "CA", "C", "O", "CB", "SG"], - "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"], - "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"], - "GLY": ["N", "CA", "C", "O"], - "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"], - "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"], - "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"], - "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"], - "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"], - "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"], - "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"], - "SER": ["N", "CA", "C", "O", "CB", "OG"], - "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"], - "TRP": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD1", - "CD2", - "NE1", - "CE2", - "CE3", - "CZ2", - "CZ3", - "CH2", - ], - "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"], - "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"], - "MSE": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"], - "UNK": ["N", "CA", "C", "O"], -} - -DNA_HEAVY_ATOMS = { - "DA": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "C1'", - "N9", - "C8", - "N7", - "C5", - "C6", - "N6", - "N1", - "C2", - "N3", - "C4", - ], - "DG": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "C1'", - "N9", - "C8", - "N7", - "C5", - "C6", - "O6", - "N1", - "C2", - "N2", - "N3", - "C4", - ], - "DC": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "C1'", - "N1", - "C2", - "O2", - "N3", - "C4", - "N4", - "C5", - "C6", - ], - "DT": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "C1'", - "N1", - "C2", - "O2", - "N3", - "C4", - "O4", - "C5", - "C7", - "C6", - ], -} - -RNA_HEAVY_ATOMS = { - "A": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "O2'", - "C1'", - "N9", - "C8", - "N7", - "C5", - "C6", - "N6", - "N1", - "C2", - "N3", - "C4", - ], - "G": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "O2'", - "C1'", - "N9", - "C8", - "N7", - "C5", - "C6", - "O6", - "N1", - "C2", - "N2", - "N3", - "C4", - ], - "C": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "O2'", - "C1'", - "N1", - "C2", - "O2", - "N3", - "C4", - "N4", - "C5", - "C6", - ], - "U": [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "O2'", - "C1'", - "N1", - "C2", - "O2", - "N3", - "C4", - "O4", - "C5", - "C6", - ], -} - -# Unknown nucleotide backbone atoms -DNA_BACKBONE_ATOMS = [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "C1'", -] -RNA_BACKBONE_ATOMS = [ - "P", - "OP1", - "OP2", - "O5'", - "C5'", - "C4'", - "O4'", - "C3'", - "O3'", - "C2'", - "O2'", - "C1'", -] diff --git a/fastplms/esmfold2/esmfold2_constants_esm3.py b/fastplms/esmfold2/esmfold2_constants_esm3.py deleted file mode 100644 index 30a631e..0000000 --- a/fastplms/esmfold2/esmfold2_constants_esm3.py +++ /dev/null @@ -1,137 +0,0 @@ -import os -from functools import cache -from pathlib import Path - -from huggingface_hub import snapshot_download - -SEQUENCE_BOS_TOKEN = 0 -SEQUENCE_PAD_TOKEN = 1 -SEQUENCE_EOS_TOKEN = 2 -SEQUENCE_CHAINBREAK_TOKEN = 31 -SEQUENCE_MASK_TOKEN = 32 - -VQVAE_CODEBOOK_SIZE = 4096 -VQVAE_SPECIAL_TOKENS = { - "MASK": VQVAE_CODEBOOK_SIZE, - "EOS": VQVAE_CODEBOOK_SIZE + 1, - "BOS": VQVAE_CODEBOOK_SIZE + 2, - "PAD": VQVAE_CODEBOOK_SIZE + 3, - "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4, -} -VQVAE_DIRECTION_LOSS_BINS = 16 -VQVAE_PAE_BINS = 64 -VQVAE_MAX_PAE_BIN = 31.0 -VQVAE_PLDDT_BINS = 50 - -STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"] -STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"] -STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"] -STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"] -STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"] -STRUCTURE_UNDEFINED_TOKEN = 955 - -SASA_PAD_TOKEN = 0 - -SS8_PAD_TOKEN = 0 - -INTERPRO_PAD_TOKEN = 0 - -RESIDUE_PAD_TOKEN = 0 - -CHAIN_BREAK_STR = "|" - -SEQUENCE_BOS_STR = "" -SEQUENCE_EOS_STR = "" - -MASK_STR_SHORT = "_" -SEQUENCE_MASK_STR = "" -SASA_MASK_STR = "" -SS8_MASK_STR = "" - -# fmt: off -SEQUENCE_VOCAB = [ - "", "", "", "", - "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", - "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", - "O", ".", "-", "|", - "", -] -# fmt: on - -SEQUENCE_STANDARD_AA_MIN_TOKEN = 4 # L -SEQUENCE_STANDARD_AA_MAX_TOKEN = 24 # X (exclusive) - -SSE_8CLASS_VOCAB = "GHITEBSC" -SSE_3CLASS_VOCAB = "HEC" -SSE_8CLASS_TO_3CLASS_MAP = { - "G": "H", - "H": "H", - "I": "H", - "T": "C", - "E": "E", - "B": "E", - "S": "C", - "C": "C", -} - -SASA_DISCRETIZATION_BOUNDARIES = [ - 0.8, - 4.0, - 9.6, - 16.4, - 24.5, - 32.9, - 42.0, - 51.5, - 61.2, - 70.9, - 81.6, - 93.3, - 107.2, - 125.4, - 151.4, -] - -MAX_RESIDUE_ANNOTATIONS = 16 - - -TFIDF_VECTOR_SIZE = 58641 - -FUNCTION_TOKENS_DEPTH = 8 - - -@staticmethod -@cache -def data_root(model: str): - if "INFRA_PROVIDER" in os.environ: - return Path("") - # Try to download from huggingface if it doesn't exist - if model.startswith("esm3"): - path = Path(snapshot_download(repo_id="biohub/esm3-sm-open-v1")) - elif model.startswith("esmc-300"): - path = Path(snapshot_download(repo_id="biohub/esmc-300m-2024-12")) - elif model.startswith("esmc-600"): - path = Path(snapshot_download(repo_id="biohub/esmc-600m-2024-12")) - elif model.startswith("esmc-6b"): - path = Path(snapshot_download(repo_id="biohub/esmc-6b-2024-12")) - else: - raise ValueError(f"{model=} is an invalid model name.") - return path - - -IN_REPO_DATA_FOLDER = Path(__file__).parents[2] / "data" - -INTERPRO_ENTRY = IN_REPO_DATA_FOLDER / "entry_list_safety_29026.list" -INTERPRO_HIERARCHY = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt" -INTERPRO2GO = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt" -INTERPRO_2ID = "data/tag_dict_4_safety_filtered.json" - -LSH_TABLE_PATHS = {"8bit": "data/hyperplanes_8bit_58641.npz"} - -KEYWORDS_VOCABULARY = ( - IN_REPO_DATA_FOLDER / "keyword_vocabulary_safety_filtered_58641.txt" -) -KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy" - -RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv" -INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv" diff --git a/fastplms/esmfold2/esmfold2_input_builder.py b/fastplms/esmfold2/esmfold2_input_builder.py deleted file mode 100644 index 2fdbe16..0000000 --- a/fastplms/esmfold2/esmfold2_input_builder.py +++ /dev/null @@ -1,254 +0,0 @@ -from dataclasses import dataclass -from typing import Any, Sequence, TypeAlias, Union - -import numpy as np - -from .esmfold2_msa import MSA - -# fmt: off -MSAInput: TypeAlias = Union[ - MSA, - None, -] -# fmt: on - - -@dataclass -class Modification: - position: int # zero-indexed - ccd: str - smiles: str | None = None # TODO(mlee): add smiles support - - -@dataclass -class ProteinInput: - id: str | list[str] - sequence: str - modifications: list[Modification] | None = None - msa: MSAInput = None - - -@dataclass -class RNAInput: - id: str | list[str] - sequence: str - modifications: list[Modification] | None = None - - -@dataclass -class DNAInput: - id: str | list[str] - sequence: str - modifications: list[Modification] | None = None - - -@dataclass -class LigandInput: - id: str | list[str] - smiles: str | None = None - ccd: list[str] | None = None - - -@dataclass -class DistogramConditioning: - chain_id: str - distogram: np.ndarray - - -@dataclass -class PocketConditioning: - binder_chain_id: str - contacts: list[tuple[str, int]] - - -@dataclass -class CovalentBond: - chain_id1: str - res_idx1: int - atom_idx1: int - chain_id2: str - res_idx2: int - atom_idx2: int - - -@dataclass -class StructurePredictionInput: - sequences: Sequence[ProteinInput | RNAInput | DNAInput | LigandInput] - pocket: PocketConditioning | None = None - distogram_conditioning: list[DistogramConditioning] | None = None - covalent_bonds: list[CovalentBond] | None = None - - -def serialize_structure_prediction_input(all_atom_input: StructurePredictionInput): - def create_chain_data(seq_input, chain_type: str) -> dict[str, Any]: - chain_data: dict[str, Any] = { - "sequence": seq_input.sequence, - "id": seq_input.id, - "type": chain_type, - } - if hasattr(seq_input, "modifications") and seq_input.modifications: - mods = [ - {"position": mod.position, "ccd": mod.ccd} - for mod in seq_input.modifications - ] - chain_data["modifications"] = mods - if not hasattr(seq_input, "msa"): - pass - elif seq_input.msa is None: - chain_data["msa"] = None - elif isinstance(seq_input.msa, MSA): - chain_data["msa"] = {"sequences": seq_input.msa.sequences} - else: - error_msg = f"MSA must be None or MSA. Got {seq_input.msa} instead." - raise AttributeError(error_msg) - return chain_data - - sequences = [] - for seq_input in all_atom_input.sequences: - if isinstance(seq_input, ProteinInput): - sequences.append(create_chain_data(seq_input, "protein")) - elif isinstance(seq_input, RNAInput): - sequences.append(create_chain_data(seq_input, "rna")) - elif isinstance(seq_input, DNAInput): - sequences.append(create_chain_data(seq_input, "dna")) - elif isinstance(seq_input, LigandInput): - sequences.append( - { - "smiles": seq_input.smiles, - "id": seq_input.id, - "ccd": seq_input.ccd, - "type": "ligand", - } - ) - else: - raise ValueError(f"Unsupported sequence input type: {type(seq_input)}") - - result: dict[str, Any] = {"sequences": sequences} - - if all_atom_input.covalent_bonds is not None: - result["covalent_bonds"] = [ - { - "chain_id1": bond.chain_id1, - "res_idx1": bond.res_idx1, - "atom_idx1": bond.atom_idx1, - "chain_id2": bond.chain_id2, - "res_idx2": bond.res_idx2, - "atom_idx2": bond.atom_idx2, - } - for bond in all_atom_input.covalent_bonds - ] - - if all_atom_input.pocket is not None: - result["pocket"] = { - "binder_chain_id": all_atom_input.pocket.binder_chain_id, - "contacts": all_atom_input.pocket.contacts, - } - - if all_atom_input.distogram_conditioning is not None: - result["distogram_conditioning"] = [ - {"chain_id": disto.chain_id, "distogram": disto.distogram.tolist()} - for disto in all_atom_input.distogram_conditioning - ] - - return result - - -def deserialize_structure_prediction_input( - data: dict[str, Any], -) -> StructurePredictionInput: - """Inverse of :func:`serialize_structure_prediction_input`. - - Reconstructs a :class:`StructurePredictionInput` from the JSON-safe dict - produced by ``serialize_structure_prediction_input``. Values round-trip; - ``DistogramConditioning.distogram`` dtype follows from JSON (``int64`` - for integer entries, ``float64`` for floats) — cast back to the original - dtype if downstream code requires a specific one. - """ - - def _mods(chain: dict[str, Any]) -> list[Modification] | None: - raw = chain.get("modifications") - if not raw: - return None - return [Modification(position=m["position"], ccd=m["ccd"]) for m in raw] - - def _msa(chain: dict[str, Any]) -> MSAInput: - if "msa" not in chain or chain["msa"] is None: - return None - msa_blk = chain["msa"] - if isinstance(msa_blk, str): - raise ValueError(f"Unexpected MSA string value: {msa_blk!r}") - return MSA.from_sequences(msa_blk["sequences"]) - - sequences: list[ProteinInput | RNAInput | DNAInput | LigandInput] = [] - for chain in data["sequences"]: - t = chain["type"] - if t == "protein": - sequences.append( - ProteinInput( - id=chain["id"], - sequence=chain["sequence"], - modifications=_mods(chain), - msa=_msa(chain), - ) - ) - elif t == "rna": - sequences.append( - RNAInput( - id=chain["id"], - sequence=chain["sequence"], - modifications=_mods(chain), - ) - ) - elif t == "dna": - sequences.append( - DNAInput( - id=chain["id"], - sequence=chain["sequence"], - modifications=_mods(chain), - ) - ) - elif t == "ligand": - sequences.append( - LigandInput( - id=chain["id"], smiles=chain.get("smiles"), ccd=chain.get("ccd") - ) - ) - else: - raise ValueError(f"Unsupported sequence type: {t!r}") - - pocket: PocketConditioning | None = None - if (pocket_blk := data.get("pocket")) is not None: - pocket = PocketConditioning( - binder_chain_id=pocket_blk["binder_chain_id"], - contacts=[tuple(c) for c in pocket_blk["contacts"]], - ) - - distogram_conditioning: list[DistogramConditioning] | None = None - if (disto_blk := data.get("distogram_conditioning")) is not None: - distogram_conditioning = [ - DistogramConditioning( - chain_id=d["chain_id"], distogram=np.asarray(d["distogram"]) - ) - for d in disto_blk - ] - - covalent_bonds: list[CovalentBond] | None = None - if (bonds_blk := data.get("covalent_bonds")) is not None: - covalent_bonds = [ - CovalentBond( - chain_id1=b["chain_id1"], - res_idx1=b["res_idx1"], - atom_idx1=b["atom_idx1"], - chain_id2=b["chain_id2"], - res_idx2=b["res_idx2"], - atom_idx2=b["atom_idx2"], - ) - for b in bonds_blk - ] - - return StructurePredictionInput( - sequences=sequences, - pocket=pocket, - distogram_conditioning=distogram_conditioning, - covalent_bonds=covalent_bonds, - ) diff --git a/fastplms/esmfold2/esmfold2_metrics.py b/fastplms/esmfold2/esmfold2_metrics.py deleted file mode 100644 index e1f8a6f..0000000 --- a/fastplms/esmfold2/esmfold2_metrics.py +++ /dev/null @@ -1,373 +0,0 @@ -import numpy as np -import torch -import torch.nn.functional as F -from einops import rearrange -from torch import Tensor -from torch.amp import autocast # type: ignore - -from . import esmfold2_residue_constants as residue_constants -from .esmfold2_misc import binpack, unbinpack -from .esmfold2_protein_structure import ( - compute_alignment_tensors, - compute_gdt_ts_no_alignment, - compute_rmsd_no_alignment, -) - - -def contact_precision( - predictions: Tensor, - targets: Tensor, - src_lengths: Tensor | None = None, - minsep: int = 6, - maxsep: int | None = None, - override_length: int | None = None, # for casp -): - """Computes contact precisions. - - For protein contact prediction, precision is measured for the top (L/K) highest confidence predictions, - with L being the length of the protein sequence and K generally being equal to 1 or 5. - - K = 5 measures the predictions of the very highest confidence contacts, while K = 1 is a more general measure - over all relatively high confidence predictions. - - Since there are roughly ~L true contacts in a protein, this is a reasonable cutoff. - - - Args: - predictions (Tensor): Tensor of probabilities of size (B, L, L) - targets (Tensor): Tensor of true contacts of size (B, L, L) - src_lengths (Tensor, optional): Lengths of each sample in the batch, if using variable lengths. - If not provided, inferred from the size of the predictions. - minsep (int): Minimum separation distance to consider. We often want to measure contacts at a - certain range. Typical ranges are short [6, 12), medium [12, 24), and long [24, inf). - maxsep (int, optional): Used in conjunction with minsep to specify a contact range. If not provided uses - assumes no maximum range - override_length (int, optional): Used for casp evaluation where sometimes the "true" length is not - the same as the length of the input. Kept for posterity, we probably don't need this argument. - """ - if predictions.dim() == 2: - predictions = predictions.unsqueeze(0) - if targets.dim() == 2: - targets = targets.unsqueeze(0) - - # Check sizes - if predictions.size() != targets.size(): - raise ValueError( - f"Size mismatch. Received predictions of size {predictions.size()}, " - f"targets of size {targets.size()}" - ) - device = predictions.device - - batch_size, seqlen, _ = predictions.size() - - # Step 1) Construct a mask of size [B, L, L] to mask invalid contacts - seqlen_range = torch.arange(seqlen, device=device) - sep = seqlen_range.unsqueeze(0) - seqlen_range.unsqueeze(1) - sep = sep.unsqueeze(0) - # Mask contacts that are closer than minsep - valid_mask = sep >= minsep - # Mask contacts where target is negative (padding or unknown) - valid_mask = valid_mask & (targets >= 0) # negative targets are invalid - - # Mask contacts that are farther than maxsep, if provided - if maxsep is not None: - valid_mask &= sep < maxsep - - if src_lengths is not None: - # If the lengths of the individual sequences are provided, mask positions - # that are farther than the end of the sequence. - valid = seqlen_range.unsqueeze(0) < src_lengths.unsqueeze(1) - valid_mask &= valid.unsqueeze(1) & valid.unsqueeze(2) - else: - src_lengths = torch.full([batch_size], seqlen, device=device, dtype=torch.long) - - # Fill in the logit tensor with -inf for all invalid positions - predictions = predictions.masked_fill(~valid_mask, float("-inf")) - - # Step 2) Select the top half of the prediction (should be symmetric) - x_ind, y_ind = np.triu_indices(seqlen, minsep) - predictions_upper = predictions[:, x_ind, y_ind] - targets_upper = targets[:, x_ind, y_ind] - - # Step 3) Select the topk values in each batch where k = L (length of sequence) - topk = seqlen if override_length is None else max(seqlen, override_length) - # Indices are the indices into the predictions corresponding to the most confident predictions - indices = predictions_upper.argsort(dim=-1, descending=True)[:, :topk] - # topk_targets are the target values corresponding to the above indices - topk_targets = targets_upper[torch.arange(batch_size).unsqueeze(1), indices] - if topk_targets.size(1) < topk: - # If there aren't enough targets, pad to the output. - topk_targets = F.pad(topk_targets, [0, topk - topk_targets.size(1)]) - - # Step 4) Sum the accuracy at of the top-i predictions for i in 1, L - # topk_targets => 1/0 true vs. false contact, sorted by confidence of prediction - # cmumulative sum => Number of correct answers for the top-i predictions. - cumulative_dist = topk_targets.type_as(predictions).cumsum(-1) - - # Step 5) Find the gather indices. This should be P@(L / K) for varous values of K - # The values will differ for each batch. - gather_lengths = src_lengths.unsqueeze(1) - if override_length is not None: - gather_lengths = override_length * torch.ones_like( - gather_lengths, device=device - ) - - # This gets you (0.1 * L, 0.2 * L, 0.3 * L, etc.) - gather_indices = ( - (torch.arange(0.1, 1.1, 0.1, device=device).unsqueeze(0) * gather_lengths).type( - torch.long - ) - - 1 - ).clamp_min(0) - - # Step 6) Gather the results and divide by the number of guesses to get the precision. - binned_cumulative_dist = cumulative_dist.gather(1, gather_indices) - binned_precisions = binned_cumulative_dist / (gather_indices + 1).type_as( - binned_cumulative_dist - ) - - # Select specific P@L/k. pl5 is index 1 b/c that corresponds to L * 0.2 in - # gather_indices above - pl5 = binned_precisions[:, 1] - # pl2 = binned_precisions[:, 4] - pl = binned_precisions[:, 9] - # AUC is the integral wrt K of P@L/K for K in range(1, L) - auc = binned_precisions.mean(-1) - - return {"AUC": auc, "P@L": pl, "P@L5": pl5} - - -def compute_lddt( - all_atom_pred_pos: torch.Tensor, - all_atom_positions: torch.Tensor, - all_atom_mask: torch.Tensor, - pairwise_all_atom_mask: torch.Tensor | None = None, - cutoff: float | torch.Tensor = 15.0, - eps: float = 1e-10, - per_residue: bool = True, - sequence_id: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Computes LDDT for a protein. Tensor sizes below include some optional dimensions. Specifically: - Nstates: - all_atom_pred_pos can contain multiple states in the first dimension which corresponds to outputs from different layers of a model (e.g. each IPA block). The return size will be [Nstates x Batch size] if this is included. - Natoms: - LDDT can be computed for all atoms or some atoms. The second to last dimension should contain the *FLATTENED* representation of L x Natoms. If you want to calculate for atom37, e.g., this will be of size (L * 37). If you are only calculating CA LDDT, it will be of size L. - - Args: - all_atom_pred_pos (Tensor[float], [(Nstates x) B x (L * Natoms x) 3]): Tensor of predicted positions - all_atom_positions (Tensor[float], [B x (L * Natoms x) 3]): Tensor of true positions - all_atom_mask (Tensor[float], [B x (L * Natoms)]): Tensor of masks, indicating whether an atom exists. - pairwise_all_atom_mask (Tensor[float], [B x (L * Natoms x L * Natoms)], optional): Tensor of masks, indicating whether a pair of atoms should be considered in the LDDT calculation. - cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids. - per_residue (bool): Whether to return per-residue or full-protein lddt. - sequence_id (Tensor, optional): Sequence id tensor for binpacking. NOTE: only supported for lddt_ca calculations, not when Natoms is passed! - - Returns: - LDDT Tensor: - if per_residue: - Tensor[float], [(Nstates x) B x (L * Natoms)] - else: - Tensor[float], [(Nstates x) B] - """ - all_atom_mask = all_atom_mask[..., None] # add a dimension for broadcasting - dmat_true = torch.sqrt( - eps - + torch.sum( - (all_atom_positions[..., None, :] - all_atom_positions[..., None, :, :]) - ** 2, - dim=-1, - ) - ) - - dmat_pred = torch.sqrt( - eps - + torch.sum( - (all_atom_pred_pos[..., None, :] - all_atom_pred_pos[..., None, :, :]) ** 2, - dim=-1, - ) - ) - mask = all_atom_mask * rearrange(all_atom_mask, "... a b -> ... b a") - if pairwise_all_atom_mask is not None: - mask = mask * pairwise_all_atom_mask - - if sequence_id is not None: - # TODO: This will work for lddt_ca, but not for regular lddt - # Problem is that regular lddt has natoms * nres scores, so would need to repeat this mask by natoms - # Leaving for now because it won't fail silently so should be ook. - seqid_mask = sequence_id[..., None] == sequence_id[..., None, :] - mask = mask * seqid_mask.type_as(mask) - - return compute_lddt_from_dmat( - dmat_pred, dmat_true, mask, cutoff=cutoff, eps=eps, per_residue=per_residue - ) - - -def compute_lddt_from_dmat( - dmat_pred: torch.Tensor, - dmat_true: torch.Tensor, - pairwise_mask: torch.Tensor, - cutoff: float | torch.Tensor = 15.0, - eps: float = 1e-10, - per_residue: bool = True, -): - """ - Compute LDDT from pre-computed distance matrices. - This is useful when you want to compute LDDT with multiple different masks or cutoffs, e.g. for different molecule types (protein, nucleic acid, etc.). - - Args: - dmat_pred (Tensor[float], [B x L x L]): Predicted distance matrix - dmat_true (Tensor[float], [B x L x L]): True distance matrix - pairwise_mask (Tensor[float], [B x L x L]): Pairwise mask indicating which pairs of atoms to consider - cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids. - per_residue (bool): Whether to return per-residue or full-protein lddt. - - Returns: - LDDT Tensor: - if per_residue: - Tensor[float], [B x L] - else: - Tensor[float], [B] - """ - n = dmat_true.size(-1) - dists_to_score = ( - (dmat_true < cutoff) - * pairwise_mask - * (1.0 - torch.eye(n, device=dmat_true.device)) - ) - - dist_l1 = torch.abs(dmat_true - dmat_pred) - score = ( - (dist_l1 < 0.5).type(dist_l1.dtype) - + (dist_l1 < 1.0).type(dist_l1.dtype) - + (dist_l1 < 2.0).type(dist_l1.dtype) - + (dist_l1 < 4.0).type(dist_l1.dtype) - ) - score = score * 0.25 - - dims = (-1,) if per_residue else (-2, -1) - norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims)) - score = norm * (eps + torch.sum(dists_to_score * score, dim=dims)) - return score - - -def compute_lddt_ca( - all_atom_pred_pos: torch.Tensor, - all_atom_positions: torch.Tensor, - all_atom_mask: torch.Tensor, - cutoff: float = 15.0, - eps: float = 1e-10, - per_residue: bool = True, - sequence_id: torch.Tensor | None = None, -) -> torch.Tensor: - ca_pos = residue_constants.atom_order["CA"] - if all_atom_pred_pos.dim() != 3: - all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :] - all_atom_positions = all_atom_positions[..., ca_pos, :] - all_atom_mask = all_atom_mask[..., ca_pos] - - return compute_lddt( - all_atom_pred_pos, - all_atom_positions, - all_atom_mask, - cutoff=cutoff, - eps=eps, - per_residue=per_residue, - sequence_id=sequence_id, - ) - - -# NOTE(roshan): no_grad required for stack_variable_length_tensors apparently... let's revisit if we want to backprop -@torch.no_grad() -@autocast("cuda", enabled=False) -def compute_rmsd( - mobile: torch.Tensor, - target: torch.Tensor, - atom_exists_mask: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, - reduction: str = "batch", -): - """ - Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N) - - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking. - - reduction (str): One of "batch", "per_sample", "per_residue". - - Returns: - If reduction == "batch": - (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch - If reduction == "per_sample": - (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch - If reduction == "per_residue": - (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch - """ - - (centered_mobile, _, centered_target, _, rotation_matrix, num_valid_atoms) = ( - compute_alignment_tensors( - mobile=mobile, - target=target, - atom_exists_mask=atom_exists_mask, - sequence_id=sequence_id, - ) - ) - - # Apply transformation to centered structure - rotated_mobile = torch.matmul(centered_mobile, rotation_matrix) - - # Compute rmsd for centered structures - rmsd = compute_rmsd_no_alignment( - rotated_mobile, centered_target, num_valid_atoms, reduction=reduction - ) - if reduction == "per_residue" and sequence_id is not None: - rmsd = binpack(rmsd, sequence_id, pad_value=0) - return rmsd - - -def compute_gdt_ts( - mobile: torch.Tensor, - target: torch.Tensor, - atom_exists_mask: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, - reduction: str = "per_sample", -): - """ - Compute GDT_TS between two batches of structures with support for masking invalid atoms using PyTorch. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N) - - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking. - - reduction (str): One of "batch", "per_sample", "per_residue". - - Returns: - If reduction == "batch": - (torch.Tensor): 0-dim, GDT_TS between the structures for each batch - If reduction == "per_sample": - (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch - """ - if atom_exists_mask is None: - atom_exists_mask = torch.isfinite(target).all(dim=-1) - (centered_mobile, _, centered_target, _, rotation_matrix, _) = ( - compute_alignment_tensors( - mobile=mobile, - target=target, - atom_exists_mask=atom_exists_mask, - sequence_id=sequence_id, - ) - ) - - # Apply transformation to centered structure - rotated_mobile = torch.matmul(centered_mobile, rotation_matrix) - - # the coordinate tensors returned by `compute_alignment_tensors` are unbinpacked and contain zeros for invalid positions - # so `compute_gdt_ts_no_alignment` requires `atom_exists_mask` to be passed and be unbinpacked - if sequence_id is not None: - atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False) - return compute_gdt_ts_no_alignment( - rotated_mobile, centered_target, atom_exists_mask, reduction - ) diff --git a/fastplms/esmfold2/esmfold2_misc.py b/fastplms/esmfold2/esmfold2_misc.py deleted file mode 100644 index 5e645c9..0000000 --- a/fastplms/esmfold2/esmfold2_misc.py +++ /dev/null @@ -1,504 +0,0 @@ -from __future__ import annotations - -import os -from collections import defaultdict -from contextlib import nullcontext -from dataclasses import is_dataclass -from io import BytesIO -from typing import ( - Any, - ContextManager, - Generator, - Iterable, - Protocol, - Sequence, - TypeVar, - runtime_checkable, -) -from warnings import warn - -import huggingface_hub -import numpy as np -import torch -import zstd - -from .esmfold2_constants_esm3 import CHAIN_BREAK_STR -from .esmfold2_utils_types import FunctionAnnotation - -MAX_SUPPORTED_DISTANCE = 1e6 - - -TSequence = TypeVar("TSequence", bound=Sequence) - - -@runtime_checkable -class Concatable(Protocol): - @classmethod - def concat(cls, objs: list[Concatable]) -> Concatable: ... - - -def slice_python_object_as_numpy( - obj: TSequence, idx: int | list[int] | slice | np.ndarray -) -> TSequence: - """ - Slice a python object (like a list, string, or tuple) as if it was a numpy object. - - Example: - >>> obj = "ABCDE" - >>> slice_python_object_as_numpy(obj, [1, 3, 4]) - "BDE" - - >>> obj = [1, 2, 3, 4, 5] - >>> slice_python_object_as_numpy(obj, np.arange(5) < 3) - [1, 2, 3] - """ - if np.isscalar(idx): - idx = [int(idx)] # type: ignore - - if isinstance(idx, np.ndarray) and idx.dtype == bool: - sliced_obj = [obj[i] for i in np.where(idx)[0]] - elif isinstance(idx, slice): - sliced_obj = obj[idx] - else: - sliced_obj = [obj[i] for i in idx] # type: ignore - - match obj, sliced_obj: - case str(), list(): - sliced_obj = "".join(sliced_obj) - case _: - sliced_obj = obj.__class__(sliced_obj) # type: ignore - - return sliced_obj # type: ignore - - -def slice_any_object( - obj: TSequence, idx: int | list[int] | slice | np.ndarray -) -> TSequence: - """ - Slice a arbitrary object (like a list, string, or tuple) as if it was a numpy object. Similar to `slice_python_object_as_numpy`, but detects if it's a numpy array or Tensor and uses the existing slice method if so. - - If the object is a dataclass, it will simply apply the index to the object, under the assumption that the object has correcty implemented numpy indexing. - - Example: - >>> obj = "ABCDE" - >>> slice_any_object(obj, [1, 3, 4]) - "BDE" - - >>> obj = np.array([1, 2, 3, 4, 5]) - >>> slice_any_object(obj, np.arange(5) < 3) - np.array([1, 2, 3]) - - >>> obj = ProteinChain.from_rcsb("1a3a", "A") - >>> slice_any_object(obj, np.arange(len(obj)) < 10) - # ProteinChain w/ length 10 - - """ - if isinstance(obj, (np.ndarray, torch.Tensor)): - return obj[idx] # type: ignore - elif is_dataclass(obj): - # if passing a dataclass, assume it implements a custom slice - return obj[idx] # type: ignore - else: - return slice_python_object_as_numpy(obj, idx) - - -def rbf(values, v_min, v_max, n_bins=16): - """ - Returns RBF encodings in a new dimension at the end. - """ - rbf_centers = torch.linspace( - v_min, v_max, n_bins, device=values.device, dtype=values.dtype - ) - rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1]) - rbf_std = (v_max - v_min) / n_bins - z = (values.unsqueeze(-1) - rbf_centers) / rbf_std - return torch.exp(-(z**2)) - - -def batched_gather(data, inds, dim=0, no_batch_dims=0): - ranges = [] - for i, s in enumerate(data.shape[:no_batch_dims]): - r = torch.arange(s) - r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1)))) - ranges.append(r) - - remaining_dims = [slice(None) for _ in range(len(data.shape) - no_batch_dims)] - remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds - ranges.extend(remaining_dims) - return data[ranges] - - -def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor: - return batched_gather(s.unsqueeze(-3), edges, -2, no_batch_dims=len(s.shape) - 1) - - -def knn_graph( - coords: torch.Tensor, - coord_mask: torch.Tensor, - padding_mask: torch.Tensor, - sequence_id: torch.Tensor, - *, - no_knn: int, -): - L = coords.shape[-2] - num_by_dist = min(no_knn, L) - device = coords.device - - coords = coords.nan_to_num() - coord_mask = ~(coord_mask[..., None, :] & coord_mask[..., :, None]) - padding_pairwise_mask = padding_mask[..., None, :] | padding_mask[..., :, None] - if sequence_id is not None: - padding_pairwise_mask |= torch.unsqueeze(sequence_id, 1) != torch.unsqueeze( - sequence_id, 2 - ) - dists = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1) - arange = torch.arange(L, device=device) - seq_dists = (arange.unsqueeze(-1) - arange.unsqueeze(-2)).abs() - # We only support up to a certain distance, above that, we use sequence distance - # instead. This is so that when a large portion of the structure is masked out, - # the edges are built according to sequence distance. - max_dist = MAX_SUPPORTED_DISTANCE - if not (dists[~coord_mask] < max_dist).all(): - raise ValueError( - f"Coordinate pairwise distances exceed max supported distance ({max_dist}). " - ) - struct_then_seq_dist = ( - seq_dists.to(dists.dtype) - .mul(1e2) - .add(max_dist) - .where(coord_mask, dists) - .masked_fill(padding_pairwise_mask, torch.inf) - ) - dists, edges = struct_then_seq_dist.sort(dim=-1, descending=False) - # This is a L x L tensor, where we index by rows first, - # and columns are the edges we should pick. - chosen_edges = edges[..., :num_by_dist] - chosen_mask = dists[..., :num_by_dist].isfinite() - return chosen_edges, chosen_mask - - -def stack_variable_length_tensors( - sequences: Sequence[torch.Tensor], - constant_value: int | float = 0, - dtype: torch.dtype | None = None, -) -> torch.Tensor: - """Automatically stack tensors together, padding variable lengths with the - value in constant_value. Handles an arbitrary number of dimensions. - - Examples: - >>> tensor1, tensor2 = torch.ones([2]), torch.ones([5]) - >>> stack_variable_length_tensors(tensor1, tensor2) - tensor of shape [2, 5]. First row is [1, 1, 0, 0, 0]. Second row is all ones. - - >>> tensor1, tensor2 = torch.ones([2, 4]), torch.ones([5, 3]) - >>> stack_variable_length_tensors(tensor1, tensor2) - tensor of shape [2, 5, 4] - """ - batch_size = len(sequences) - shape = [batch_size] + np.max([seq.shape for seq in sequences], 0).tolist() - - if dtype is None: - dtype = sequences[0].dtype - device = sequences[0].device - - array = torch.full(shape, constant_value, dtype=dtype, device=device) - for arr, seq in zip(array, sequences): - arrslice = tuple(slice(dim) for dim in seq.shape) - arr[arrslice] = seq - - return array - - -def binpack( - tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float -): - """ - Args: - tensor (Tensor): [B, L, ...] - - Returns: - Tensor: [B_binpacked, L_binpacked, ...] - """ - if sequence_id is None: - return tensor - - num_sequences = sequence_id.max(dim=-1).values + 1 - - dims = sequence_id.shape + tensor.shape[2:] - output_tensor = torch.full( - dims, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device - ) - - idx = 0 - for batch_idx, (batch_seqid, batch_num_sequences) in enumerate( - zip(sequence_id, num_sequences) - ): - for seqid in range(batch_num_sequences): - mask = batch_seqid == seqid - output_tensor[batch_idx, mask] = tensor[idx, : mask.sum()] - idx += 1 - return output_tensor - - -def unbinpack( - tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float -): - """ - Args: - tensor (Tensor): [B, L, ...] - - Returns: - Tensor: [B_unbinpacked, L_unbinpack, ...] - """ - if sequence_id is None: - return tensor - - unpacked_tensors = [] - num_sequences = sequence_id.max(dim=-1).values + 1 - for batch_idx, (batch_seqid, batch_num_sequences) in enumerate( - zip(sequence_id, num_sequences) - ): - for seqid in range(batch_num_sequences): - mask = batch_seqid == seqid - unpacked = tensor[batch_idx, mask] - unpacked_tensors.append(unpacked) - return stack_variable_length_tensors(unpacked_tensors, pad_value) - - -def fp32_autocast_context(device_type: str) -> ContextManager[Any]: # type: ignore - """ - Returns an autocast context manager that disables downcasting by AMP. - - Args: - device_type: The device type ('cpu' or 'cuda') - - Returns: - An autocast context manager with the specified behavior. - """ - if device_type == "cpu": - return torch.amp.autocast(device_type, enabled=False) # type: ignore - elif device_type == "mps": - # For MPS, just return a no-op context manager (nullcontext) since MPS does not support autocast. - return nullcontext() - elif device_type == "cuda": - return torch.amp.autocast(device_type, dtype=torch.float32) # type: ignore - else: - raise ValueError(f"Unsupported device type: {device_type}") - - -def merge_ranges(ranges: list[range], merge_gap_max: int | None = None) -> list[range]: - """Merge overlapping ranges into sorted, non-overlapping segments. - - Args: - ranges: collection of ranges to merge. - merge_gap_max: optionally merge neighboring ranges that are separated by a gap - no larger than this size. - Returns: - non-overlapping ranges merged from the inputs, sorted by position. - """ - ranges = sorted(ranges, key=lambda r: r.start) - merge_gap_max = merge_gap_max if merge_gap_max is not None else 0 - assert merge_gap_max >= 0, f"Invalid merge_gap_max: {merge_gap_max}" - - merged = [] - for r in ranges: - if not merged: - merged.append(r) - else: - last = merged[-1] - if last.stop + merge_gap_max >= r.start: - merged[-1] = range(last.start, max(last.stop, r.stop)) - else: - merged.append(r) - return merged - - -def merge_annotations( - annotations: list[FunctionAnnotation], merge_gap_max: int | None = None -) -> list[FunctionAnnotation]: - """Merges annotations into non-overlapping segments. - - Args: - annotations: annotations to merge. - merge_gap_max: optionally merge neighboring ranges that are separated by a gap - no larger than this size. - Returns: - non-overlapping annotations with gaps merged. - """ - grouped: dict[str, list[range]] = defaultdict(list) - for a in annotations: - # +1 since FunctionAnnotation.end is inlcusive. - grouped[a.label].append(range(a.start, a.end + 1)) - - merged = [] - for label, ranges in grouped.items(): - merged_ranges = merge_ranges(ranges, merge_gap_max=merge_gap_max) - for range_ in merged_ranges: - annotation = FunctionAnnotation( - label=label, - start=range_.start, - end=range_.stop - 1, # convert range.stop exclusive -> inclusive. - ) - merged.append(annotation) - return merged - - -def replace_inf(data): - if data is None: - return None - array = np.asarray(data, dtype=np.float32) - array = np.where(np.isinf(array), 1000, array) - return array.tolist() - - -def maybe_tensor(x, convert_none_to_nan: bool = False) -> torch.Tensor | None: - if x is None: - return None - if isinstance(x, torch.Tensor): - return x - if isinstance(x, list) and all(isinstance(t, torch.Tensor) for t in x): - return torch.stack(x) - if convert_none_to_nan: - x = np.asarray(x, dtype=np.float32) - x = np.where(x is None, np.nan, x) - return torch.tensor(x) - - -def maybe_list(x, convert_nan_to_none: bool = False) -> list | None: - if x is None: - return None - if not convert_nan_to_none: - return x.tolist() - - # Handle both torch.tensor and np.ndarray input. - if isinstance(x, torch.Tensor): - nan_mask = torch.isnan(x).cpu().numpy() - np_arr = x.cpu().numpy().astype(object) - elif isinstance(x, np.ndarray): - nan_mask = np.isnan(x) - np_arr = x.astype(object) - else: - raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.") - - np_arr[nan_mask] = None - return np_arr.tolist() - - -def huggingfacehub_login(): - """Authenticates with the Hugging Face Hub using the HF_TOKEN environment - variable, else by prompting the user""" - token = os.environ.get("HF_TOKEN") - huggingface_hub.login(token=token) - - -def get_chainbreak_boundaries_from_sequence(sequence: Sequence[str]) -> np.ndarray: - chain_boundaries = [0] - for i, aa in enumerate(sequence): - if aa == CHAIN_BREAK_STR: - if i == (len(sequence) - 1): - raise ValueError( - "Encountered chain break token at end of sequence, this is unexpected." - ) - if i == (len(sequence) - 2): - warn( - "Encountered chain break token at penultimate position, this is unexpected." - ) - chain_boundaries.append(i) - chain_boundaries.append(i + 1) - chain_boundaries.append(len(sequence)) - assert len(chain_boundaries) % 2 == 0 - chain_boundaries = np.array(chain_boundaries).reshape(-1, 2) - return chain_boundaries - - -def deserialize_tensors(b: bytes) -> Any: - buf = BytesIO(zstd.ZSTD_uncompress(b)) - d = torch.load(buf, map_location="cpu", weights_only=False) - return d - - -def join_lists( - lists: Sequence[Sequence[Any]], separator: Sequence[Any] | None = None -) -> list[Any]: - """Joins multiple lists with separator element. Like str.join but for lists. - - Example: [[1, 2], [3], [4]], separator=[0] -> [1, 2, 0, 3, 0, 4] - - Args: - lists: Lists of elements to chain - separator: separators to intsert between chained output. - Returns: - Joined lists. - """ - if not lists: - return [] - joined = [] - joined.extend(lists[0]) - for l in lists[1:]: - if separator: - joined.extend(separator) - joined.extend(l) - return joined - - -def iterate_with_intermediate( - lists: Iterable, intermediate -) -> Generator[Any, None, None]: - """ - Iterate over the iterable, yielding the intermediate value between - every element of the intermediate. Useful for joining objects with - separator tokens. - """ - it = iter(lists) - yield next(it) - for l in it: - yield intermediate - yield l - - -def concat_objects(objs: Sequence[Any], separator: Any | None = None): - """ - Concat objects with each other using a separator token. - - Supports: - - Concatable (objects that implement `concat` classmethod) - - strings - - lists - - numpy arrays - - torch Tensors - - Example: - >>> foo = "abc" - >>> bar = "def" - >>> concat_objects([foo, bar], "|") - "abc|def" - """ - match objs[0]: - case Concatable(): - return objs[0].__class__.concat(objs) # type: ignore - case str(): - assert isinstance( - separator, str - ), "Trying to join strings but separator is not a string" - return separator.join(objs) - case list(): - if separator is not None: - return join_lists(objs, [separator]) - else: - return join_lists(objs) - case np.ndarray(): - if separator is not None: - return np.concatenate( - list(iterate_with_intermediate(objs, np.array([separator]))) - ) - else: - return np.concatenate(objs) - case torch.Tensor(): - if separator is not None: - return torch.cat( - list(iterate_with_intermediate(objs, torch.tensor([separator]))) - ) - else: - return torch.cat(objs) # type: ignore - case _: - raise TypeError(type(objs[0])) diff --git a/fastplms/esmfold2/esmfold2_mmcif_parsing.py b/fastplms/esmfold2/esmfold2_mmcif_parsing.py deleted file mode 100644 index 1b8cc54..0000000 --- a/fastplms/esmfold2/esmfold2_mmcif_parsing.py +++ /dev/null @@ -1,469 +0,0 @@ -from __future__ import annotations - -import functools -import io -import os -from dataclasses import dataclass -from datetime import datetime -from typing import Union - -import biotite.structure as bs -import biotite.structure.io.pdbx as pdbx - -from . import esmfold2_residue_constants as residue_constants - -# Define PathOrBuffer for the opensource version -PathOrBuffer = Union[str, os.PathLike, io.StringIO] - - -class NoProteinError(Exception): - pass - - -@dataclass -class Residue: - residue_number: int | None = None - insertion_code: str = "" - hetflag: bool = False - - -@dataclass -class MmcifHeader: - release_date: datetime | None = None - resolution: float | None = None - structure_method: str = "UNKNOWN" - - -class MmcifWrapper: - def __init__(self, id: str | None = None): - self.id: str = id or "" - self.raw: pdbx.CIFFile | None = None - self.structure: bs.AtomArray - self.header: MmcifHeader = MmcifHeader() - self.entities: dict[int, list[str]] = {} - self.chain_to_seqres: dict[str, str] = {} - self.seqres_to_structure: dict[str, dict[int, Residue]] = {} - - @classmethod - def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper: - obj = cls(id=id) - obj._load(path) - return obj - - def _load(self, path: PathOrBuffer, fileid: str | None = None): - """Load mmCIF data from file.""" - self.raw = pdbx.CIFFile.read(path) - - self._parse_structure() - self._parse_header() - self._parse_entities() - self._parse_sequences() - - def _parse_structure(self): - """Parse the atomic structure from mmCIF.""" - try: - structure = pdbx.get_structure(self.raw, model=1) - if structure is None or not isinstance(structure, bs.AtomArray): - raise NoProteinError("No structure found in mmCIF file") - if len(structure) == 0: - raise NoProteinError("Empty structure in mmCIF file") - self.structure = structure - except Exception as e: - raise ValueError(f"Failed to parse structure: {e}") - - def _parse_header(self): - """Parse header information from mmCIF.""" - if not self.raw: - return - - try: - # Get the first (and usually only) block - block = self.raw.block - - # Parse release date - if "pdbx_database_status" in block: - status_cat = block["pdbx_database_status"] - if "recvd_initial_deposition_date" in status_cat: - date_str = status_cat["recvd_initial_deposition_date"].as_item() - if date_str and date_str != "?": - try: - self.header.release_date = datetime.strptime( - date_str, "%Y-%m-%d" - ) - except ValueError: - pass - - # Parse resolution - if "refine" in block: - refine_cat = block["refine"] - if "ls_d_res_high" in refine_cat: - res_str = refine_cat["ls_d_res_high"].as_item() - if res_str and res_str != "?": - try: - self.header.resolution = float(res_str) - except ValueError: - pass - - # Parse structure method - if "exptl" in block: - exptl_cat = block["exptl"] - if "method" in exptl_cat: - method = exptl_cat["method"].as_item() - if method and method != "?": - self.header.structure_method = method.upper() - - except Exception: - # If parsing fails, keep default values - pass - - def _parse_entities(self): - """Parse entity information and map to chains.""" - if not self.raw: - return - - try: - block = self.raw.block - - # Parse entity information - if "entity" in block: - entity_cat = block["entity"] - entity_ids = entity_cat["id"].as_array(str) - entity_types = entity_cat["type"].as_array(str) - - # Initialize entities dict with all entities (not just polymers) - for i, (entity_id, entity_type) in enumerate( - zip(entity_ids, entity_types) - ): - self.entities[int(entity_id)] = [] - - # Map polymer chains to entities using entity_poly - if "entity_poly" in block: - poly_cat = block["entity_poly"] - entity_ids = poly_cat["entity_id"].as_array(str) - chain_lists = poly_cat["pdbx_strand_id"].as_array(str) - - for entity_id, chain_list in zip(entity_ids, chain_lists): - entity_id = int(entity_id) - # Chain list is comma-separated - chains = [c.strip() for c in chain_list.split(",") if c.strip()] - if entity_id in self.entities: - self.entities[entity_id] = chains - - # Map non-polymer chains using struct_asym for entities not covered by entity_poly - if "struct_asym" in block: - asym_cat = block["struct_asym"] - asym_ids = asym_cat["id"].as_array(str) - entity_ids = asym_cat["entity_id"].as_array(str) - - for asym_id, entity_id in zip(asym_ids, entity_ids): - entity_id = int(entity_id) - # Only add if entity exists but has no chains yet (non-polymer entities) - if entity_id in self.entities and not self.entities[entity_id]: - self.entities[entity_id].append(asym_id) - - except Exception: - # If parsing fails, try to infer from structure - if ( - self.structure - and hasattr(self.structure, "chain_id") - and self.structure.chain_id is not None - and hasattr(self.structure.chain_id, "__iter__") - ): - chain_ids = list(set(self.structure.chain_id)) - self.entities = {1: chain_ids} - - def _parse_sequences(self): - """Parse sequence information from mmCIF.""" - if not self.raw: - return - - block = self.raw.block - - # Parse polymer sequences - if "entity_poly" in block: - poly_cat = block["entity_poly"] - entity_ids = poly_cat["entity_id"].as_array(str) - sequences = poly_cat["pdbx_seq_one_letter_code_can"].as_array(str) - chain_lists = poly_cat["pdbx_strand_id"].as_array(str) - - for entity_id, sequence, chain_list in zip( - entity_ids, sequences, chain_lists - ): - # Clean up sequence (remove whitespace and newlines) - clean_seq = "".join(sequence.split()) - chains = [c.strip() for c in chain_list.split(",") if c.strip()] - - for chain_id in chains: - self.chain_to_seqres[chain_id] = clean_seq - - # Parse sequence to structure mapping - if "pdbx_poly_seq_scheme" in block: - seq_cat = block["pdbx_poly_seq_scheme"] - asym_ids = seq_cat["asym_id"].as_array(str) # Internal chain IDs - seq_positions = seq_cat["seq_id"].as_array(str) - auth_seq_nums = seq_cat["auth_seq_num"].as_array(str) - ins_codes = ( - seq_cat["pdb_ins_code"].as_array(str) - if "pdb_ins_code" in seq_cat - else [""] * len(asym_ids) - ) - hetflags = ( - seq_cat["hetflag"].as_array(str) - if "hetflag" in seq_cat - else ["N"] * len(asym_ids) - ) - - # Get author chain IDs if available - auth_chain_ids = ( - seq_cat["pdb_strand_id"].as_array(str) - if "pdb_strand_id" in seq_cat - else asym_ids # Fallback to internal IDs - ) - - # Build mapping from internal chain ID to author chain ID - asym_to_auth_mapping = {} - for asym_id, auth_id in zip(asym_ids, auth_chain_ids): - asym_to_auth_mapping[asym_id] = auth_id - - # Group by internal chain ID first, then map to author chain ID - chain_data = {} - for asym_id, seq_pos, auth_seq, ins_code, hetflag in zip( - asym_ids, seq_positions, auth_seq_nums, ins_codes, hetflags - ): - if asym_id not in chain_data: - chain_data[asym_id] = {} - - try: - seq_index = int(seq_pos) - 1 # Convert to 0-based indexing - res_num = int(auth_seq) if auth_seq != "?" else None - except ValueError: - continue - - if res_num is not None: - # Convert mmCIF "." and "?" to empty string - clean_ins_code = "" if ins_code in [".", "?"] else ins_code - else: - clean_ins_code = "" - res_num = None - - is_het = hetflag.upper() == "Y" # type: ignore - chain_data[asym_id][seq_index] = Residue( - residue_number=res_num, - insertion_code=clean_ins_code, # type: ignore - hetflag=is_het, - ) - - # Handle cases where multiple residues have the same auth_seq_num - # by adjusting residue numbers to be unique within each chain - for asym_id, residue_data in chain_data.items(): - # Check if there are duplicate residue numbers in this chain - positions_with_same_num = {} - for seq_idx, res_at_pos in residue_data.items(): - if res_at_pos.residue_number is not None: - res_num = res_at_pos.residue_number - if res_num not in positions_with_same_num: - positions_with_same_num[res_num] = [] - positions_with_same_num[res_num].append(seq_idx) - - # Fix duplicate residue numbers by making them sequential - for res_num, seq_indices in positions_with_same_num.items(): - if len(seq_indices) > 1: - # Multiple residues have the same residue number - # Make them sequential starting from the original number - seq_indices.sort() # Ensure consistent ordering - for i, seq_idx in enumerate(seq_indices): - original_pos = residue_data[seq_idx] - new_pos = Residue( - residue_number=res_num + i, - insertion_code=original_pos.insertion_code, - hetflag=original_pos.hetflag, - ) - residue_data[seq_idx] = new_pos - - # Create ordered mappings using author chain IDs - for asym_id in chain_data: - auth_chain_id = asym_to_auth_mapping.get(asym_id, asym_id) - if auth_chain_id in self.chain_to_seqres: - seq_len = len(self.chain_to_seqres[auth_chain_id]) - ordered_mapping = {} - - for i in range(seq_len): - if i in chain_data[asym_id]: - ordered_mapping[i] = chain_data[asym_id][i] - else: - # Missing residue - no structure coordinates - ordered_mapping[i] = Residue( - residue_number=None, insertion_code="", hetflag=False - ) - - self.seqres_to_structure[auth_chain_id] = ordered_mapping - else: - # Handle case where auth_chain_id is not in chain_to_seqres - # This can happen if the chain is not a polymer or if there's a parsing issue - # Create a basic mapping based on the chain_data - if chain_data[asym_id]: - # Sort by sequence index to create ordered mapping - sorted_indices = sorted(chain_data[asym_id].keys()) - ordered_mapping = {} - for i, seq_idx in enumerate(sorted_indices): - ordered_mapping[i] = chain_data[asym_id][seq_idx] - self.seqres_to_structure[auth_chain_id] = ordered_mapping - - # Ensure all chains have complete mappings - for chain_id in self.chain_to_seqres: - if chain_id not in self.seqres_to_structure: - seq_len = len(self.chain_to_seqres[chain_id]) - self.seqres_to_structure[chain_id] = { - i: Residue(residue_number=None, insertion_code="", hetflag=False) - for i in range(seq_len) - } - else: - # Fill in any missing indices - seq_len = len(self.chain_to_seqres[chain_id]) - mapping = self.seqres_to_structure[chain_id] - for i in range(seq_len): - if i not in mapping: - mapping[i] = Residue( - residue_number=None, insertion_code="", hetflag=False - ) - - # Fallback: create basic mappings from structure for missing chains - if ( - self.structure - and hasattr(self.structure, "chain_id") - and self.structure.chain_id is not None - and hasattr(self.structure.chain_id, "__iter__") - ): - for chain_id in set(self.structure.chain_id): - if chain_id not in self.seqres_to_structure: - chain_structure = self.structure[ - self.structure.chain_id == chain_id - ] - if ( - hasattr(chain_structure, "res_id") - and chain_structure.res_id is not None - and hasattr(chain_structure.res_id, "__iter__") - ): - residue_ids = list(set(chain_structure.res_id)) - residue_ids.sort() - - self.seqres_to_structure[chain_id] = { - i: Residue( - residue_number=res_id, insertion_code="", hetflag=False - ) - for i, res_id in enumerate(residue_ids) - } - - def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]: - """Parse non-polymer coordinates from mmCIF block data.""" - nonpoly_coords = {} - - # Get non-polymer entities from the mmCIF block - assert self.raw is not None - block = self.raw.block - nonpoly_entities = set() - - # Find non-polymer entities - if "entity" in block: - entity_cat = block["entity"] - entity_ids = entity_cat["id"].as_array(str) - entity_types = entity_cat["type"].as_array(str) - - for entity_id, entity_type in zip(entity_ids, entity_types): - if entity_type.upper() in ["NON-POLYMER", "WATER", "BRANCHED"]: - nonpoly_entities.add(entity_id) - - # Map entities to chains for non-polymers - entity_to_chains = {} - if "pdbx_entity_nonpoly" in block: - nonpoly_cat = block["pdbx_entity_nonpoly"] - entity_ids = nonpoly_cat["entity_id"].as_array(str) - comp_ids = nonpoly_cat["comp_id"].as_array(str) - - for entity_id, comp_id in zip(entity_ids, comp_ids): - if entity_id in nonpoly_entities: - entity_to_chains[entity_id] = comp_id - - # Get atom site information for non-polymers - if "atom_site" in block: - atom_cat = block["atom_site"] - atom_chain_ids = atom_cat["label_asym_id"].as_array(str) - atom_entity_ids = atom_cat["label_entity_id"].as_array(str) - atom_comp_ids = atom_cat["label_comp_id"].as_array(str) - - # Group non-polymer atoms by entity and chain - nonpoly_atom_groups = {} - for i, (chain_id, entity_id, comp_id) in enumerate( - zip(atom_chain_ids, atom_entity_ids, atom_comp_ids) - ): - if entity_id in nonpoly_entities: - key = (comp_id, chain_id) - if key not in nonpoly_atom_groups: - nonpoly_atom_groups[key] = [] - nonpoly_atom_groups[key].append(i) - - # Extract coordinates for each non-polymer group - for (comp_id, chain_id), atom_indices in nonpoly_atom_groups.items(): - # Match atoms by comparing chain_id and residue name - structure_mask = (self.structure.chain_id == chain_id) & ( - self.structure.res_name == comp_id - ) - - if structure_mask.any(): - nonpoly_array = self.structure[structure_mask] - if ( - isinstance(nonpoly_array, (bs.AtomArray, bs.AtomArrayStack)) - and len(nonpoly_array) > 0 - ): - nonpoly_coords[(comp_id, chain_id)] = nonpoly_array - - return nonpoly_coords - - def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]: - """Fallback method to extract heteroatoms directly from structure.""" - nonpoly_coords = {} - - if not (self.structure and hasattr(self.structure, "chain_id")): - return nonpoly_coords - - # Create set of standard residues from residue_constants - standard_residues = set(residue_constants.resnames[:-1]) # Exclude 'UNK' - standard_residues.update({"A", "C", "G", "T", "U"}) # Add nucleic acids - - if hasattr(self.structure, "chain_id") and self.structure.chain_id is not None: - for chain_id in set(self.structure.chain_id): - chain_structure = self.structure[self.structure.chain_id == chain_id] - - # Find non-standard residues - if ( - hasattr(chain_structure, "res_name") - and chain_structure.res_name is not None - and hasattr(chain_structure.res_name, "__iter__") - ): - for res_name in set(chain_structure.res_name): - if res_name not in standard_residues: - res_mask = (chain_structure.chain_id == chain_id) & ( - chain_structure.res_name == res_name - ) - if res_mask.any() and isinstance( - chain_structure, (bs.AtomArray, bs.AtomArrayStack) - ): - nonpoly_array = chain_structure[res_mask] - nonpoly_coords[(res_name, chain_id)] = nonpoly_array - - return nonpoly_coords - - @functools.cached_property - def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]: - """ - Extract non-polymer coordinates (ligands, cofactors, etc.) from mmCIF structure. - - Returns a dictionary mapping (nonpolymer_info, chain_id) tuples to AtomArrays. - """ - if not self.structure or not self.raw: - return {} - - try: - return self._parse_nonpoly_from_mmcif() - except Exception: - return self._parse_nonpoly_fallback() diff --git a/fastplms/esmfold2/esmfold2_molecular_complex.py b/fastplms/esmfold2/esmfold2_molecular_complex.py deleted file mode 100644 index a50ce14..0000000 --- a/fastplms/esmfold2/esmfold2_molecular_complex.py +++ /dev/null @@ -1,1226 +0,0 @@ -from __future__ import annotations - -import io -import os -import re -from dataclasses import asdict, dataclass -from pathlib import Path -from subprocess import check_output -from tempfile import TemporaryDirectory -from typing import TYPE_CHECKING, Any - -import biotite.structure as bs -import biotite.structure.io.pdbx as pdbx -import brotli -import msgpack -import numpy as np -import torch -from biotite.structure.io.pdbx import ( - CIFCategory, - CIFColumn, - CIFData, - CIFFile, - set_structure, -) - -from . import esmfold2_residue_constants as residue_constants -from .esmfold2_metrics import compute_lddt, compute_rmsd -from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata - - -@dataclass -class MolecularComplexResult: - """Result of molecular complex folding""" - - complex: MolecularComplex - plddt: torch.Tensor | None = None - ptm: float | None = None - iptm: float | None = None - pae: torch.Tensor | None = None - distogram: torch.Tensor | None = None - pair_chains_iptm: torch.Tensor | None = None - output_embedding_sequence: torch.Tensor | None = None - output_embedding_pair_pooled: torch.Tensor | None = None - residue_index: torch.Tensor | None = None - entity_id: torch.Tensor | None = None - sae_features: np.ndarray | None = None # [L, n_features] - ttt_metrics: dict[str, Any] | None = None - - -@dataclass -class MolecularComplexMetadata: - """Metadata for MolecularComplex objects.""" - - entity_lookup: dict[int, str] - chain_lookup: dict[int, str] - assembly_composition: dict[str, list[str]] | None = None - - -@dataclass -class Molecule: - """Represents a single molecule/token within a MolecularComplex.""" - - token: str - token_idx: int - atom_positions: np.ndarray # [N_atoms, 3] - atom_elements: np.ndarray # [N_atoms] element strings - atom_names: np.ndarray | None = None # [N_atoms] atom names (optional) - atom_hetero: np.ndarray | None = None # [N_atoms] hetero flags (optional) - residue_type: int = 0 - molecule_type: int = 0 # PROTEIN=0, RNA=1, DNA=2, LIGAND=3 - confidence: float = 0.0 - - -@dataclass(frozen=True) -class MolecularComplex: - """ - Dataclass representing a molecular complex with support for proteins, nucleic acids, and ligands. - - Uses a flat atom representation with token-based sequence indexing, supporting all atom types - beyond the traditional atom37 protein representation. - """ - - id: str - sequence: list[str] # Token sequence like ['MET', 'LYS', 'A', 'G', 'ATP'] - - # Flat atom arrays - simplified representation - atom_positions: np.ndarray # [N_atoms, 3] 3D coordinates - atom_elements: np.ndarray # [N_atoms] element strings - - # Token-to-atom mapping for efficient access - token_to_atoms: np.ndarray # [N_tokens, 2] start/end indices into atoms array - - # Chain information - chain_id: np.ndarray # [N_tokens] chain identifier for each token - - # Confidence data - plddt: np.ndarray # Per-token confidence scores [N_tokens] - - # Metadata - metadata: MolecularComplexMetadata - - # Optional atom names and hetero flags (preserved from original structures) - atom_names: np.ndarray | None = None # [N_atoms] atom names (optional) - atom_hetero: np.ndarray | None = None # [N_atoms] hetero flags (optional) - - def __post_init__(self): - """Validate array dimensions.""" - n_tokens = len(self.sequence) - n_atoms = len(self.atom_positions) - assert ( - self.token_to_atoms.shape[0] == n_tokens - ), f"token_to_atoms shape {self.token_to_atoms.shape} != {n_tokens} tokens" - assert ( - self.chain_id.shape[0] == n_tokens - ), f"chain_id shape {self.chain_id.shape} != {n_tokens} tokens" - assert ( - self.plddt.shape[0] == n_tokens - ), f"plddt shape {self.plddt.shape} != {n_tokens} tokens" - if self.atom_names is not None: - assert ( - self.atom_names.shape[0] == n_atoms - ), f"atom_names shape {self.atom_names.shape} != {n_atoms} atoms" - if self.atom_hetero is not None: - assert ( - self.atom_hetero.shape[0] == n_atoms - ), f"atom_hetero shape {self.atom_hetero.shape} != {n_atoms} atoms" - - def __len__(self) -> int: - """Return number of tokens.""" - return len(self.sequence) - - def __getitem__(self, idx: int) -> Molecule: - """Access individual molecules/tokens by index.""" - if idx >= len(self.sequence) or idx < 0: - raise IndexError( - f"Token index {idx} out of range for {len(self.sequence)} tokens" - ) - - token = self.sequence[idx] - start_atom, end_atom = self.token_to_atoms[idx] - - # Extract atom data for this token - token_atom_positions = self.atom_positions[start_atom:end_atom] - token_atom_elements = self.atom_elements[start_atom:end_atom] - token_atom_names = None - if self.atom_names is not None: - token_atom_names = self.atom_names[start_atom:end_atom] - token_atom_hetero = None - if self.atom_hetero is not None: - token_atom_hetero = self.atom_hetero[start_atom:end_atom] - - # Default values for residue/molecule type (would be extended based on actual implementation) - residue_type = 0 # Default to standard residue - molecule_type = 0 # Default to protein - - return Molecule( - token=token, - token_idx=idx, - atom_positions=token_atom_positions, - atom_elements=token_atom_elements, - atom_names=token_atom_names, - atom_hetero=token_atom_hetero, - residue_type=residue_type, - molecule_type=molecule_type, - confidence=self.plddt[idx], - ) - - @property - def atom_coordinates(self) -> np.ndarray: - """Get flat array of all atom coordinates [N_atoms, 3].""" - return self.atom_positions - - # Conversion methods - @classmethod - def from_protein_complex(cls, pc: ProteinComplex) -> "MolecularComplex": - """Convert a ProteinComplex to MolecularComplex. - - Args: - pc: ProteinComplex object with atom37 representation - - Returns: - MolecularComplex with flat atom arrays and token-based indexing - """ - from . import esmfold2_residue_constants - - # Extract sequence without chain breaks - sequence_no_breaks = pc.sequence.replace("|", "") - sequence_tokens = [ - residue_constants.restype_1to3.get(aa, "UNK") for aa in sequence_no_breaks - ] - - # Convert atom37 to flat arrays - flat_positions = [] - flat_elements = [] - flat_names = [] - flat_hetero = [] - token_to_atoms = [] - - atom_idx = 0 - - for i, aa in enumerate(pc.sequence): - if aa == "|": - # Skip chain break tokens - continue - - # Get atom37 positions and mask for this residue. - # ProteinComplex arrays are indexed by sequence position (including |), - # so use `i` not a separate residue counter. - res_positions = pc.atom37_positions[i] # [37, 3] - res_mask = pc.atom37_mask[i] # [37] - - # Track start position for this token - token_start = atom_idx - - # Process each atom type in atom37 representation - for atom_type_idx, atom_name in enumerate(residue_constants.atom_types): - if res_mask[atom_type_idx]: # Atom is present - # Add position - flat_positions.append(res_positions[atom_type_idx]) - - # Determine element from atom name - element = ( - atom_name[0] if atom_name else "C" - ) # First character is element - flat_elements.append(element) - - # Add atom name - flat_names.append(atom_name) - - # Add hetero flag (all proteins are non-hetero) - flat_hetero.append(False) - - atom_idx += 1 - - # Record token-to-atom mapping [start_idx, end_idx) - token_to_atoms.append([token_start, atom_idx]) - - # Convert to numpy arrays - atom_positions = np.array(flat_positions, dtype=np.float32) - atom_elements = np.array(flat_elements, dtype=object) - atom_names = np.array(flat_names, dtype=object) - atom_hetero = np.array(flat_hetero, dtype=bool) - token_to_atoms_array = np.array(token_to_atoms, dtype=np.int32) - - # Extract confidence scores and chain_ids (skip chain breaks) - confidence_scores = [] - chain_ids = [] - for seq_idx, aa in enumerate(pc.sequence): - if aa != "|": - confidence_scores.append(pc.confidence[seq_idx]) - chain_ids.append(pc.chain_id[seq_idx]) - - confidence_array = np.array(confidence_scores, dtype=np.float32) - chain_id_array = np.array(chain_ids, dtype=np.int64) - - # Create metadata - convert entity IDs to strings for MolecularComplexMetadata - entity_lookup_str = {k: str(v) for k, v in pc.metadata.entity_lookup.items()} - metadata = MolecularComplexMetadata( - entity_lookup=entity_lookup_str, - chain_lookup=pc.metadata.chain_lookup, - assembly_composition=pc.metadata.assembly_composition, - ) - - return cls( - id=pc.id, - sequence=sequence_tokens, - atom_positions=atom_positions, - atom_elements=atom_elements, - token_to_atoms=token_to_atoms_array, - chain_id=chain_id_array, - plddt=confidence_array, - metadata=metadata, - atom_names=atom_names, - atom_hetero=atom_hetero, - ) - - def to_protein_complex(self) -> ProteinComplex: - """Convert MolecularComplex back to ProteinComplex format. - - Extracts only protein tokens and converts from flat atom representation - back to atom37 format used by ProteinComplex. - - Returns: - ProteinComplex with protein residues only, excluding ligands/nucleic acids - """ - from . import esmfold2_residue_constants - - # No need for element mapping - already using element characters - - # Filter for protein tokens only (skip ligands, nucleic acids) - protein_tokens = [] - protein_indices = [] - - for i, token in enumerate(self.sequence): - # Check if token is a standard 3-letter amino acid code - if token in residue_constants.restype_3to1: - protein_tokens.append(token) - protein_indices.append(i) - - if not protein_tokens: - raise ValueError("No protein tokens found in MolecularComplex") - - n_residues = len(protein_tokens) - - # Initialize atom37 arrays - atom37_positions = np.full((n_residues, 37, 3), np.nan, dtype=np.float32) - atom37_mask = np.zeros((n_residues, 37), dtype=bool) - - # Extract confidence scores and chain_ids for protein residues only - protein_confidence = self.plddt[protein_indices] - protein_chain_ids = self.chain_id[protein_indices] - - # Convert tokens back to single-letter sequence with chain breaks - single_letter_residues = [] - prev_chain_id = None - - for i, (token, chain_id_val) in enumerate( - zip(protein_tokens, protein_chain_ids) - ): - # Add chain break if we're switching to a new chain - if prev_chain_id is not None and chain_id_val != prev_chain_id: - single_letter_residues.append("|") - single_letter_residues.append(residue_constants.restype_3to1[token]) - prev_chain_id = chain_id_val - - single_letter_sequence = "".join(single_letter_residues) - - # Calculate final sequence length (includes chain breaks) - sequence_length = len(single_letter_sequence) - - # Convert flat atoms back to atom37 representation using atom names - for res_idx, token_idx in enumerate(protein_indices): - token = self.sequence[token_idx] - start_atom, end_atom = self.token_to_atoms[token_idx] - - res_atom_positions = self.atom_positions[start_atom:end_atom] - res_atom_names = ( - np.array(self.atom_names[start_atom:end_atom], dtype=str) - if self.atom_names is not None - else np.array([], dtype=str) - ) - - # Build a mapping from normalized atom name -> position for this residue - # Normalize to uppercase and strip whitespace for robust matching - name_to_pos: dict[str, np.ndarray] = {} - for i, nm in enumerate(res_atom_names): - key = nm.upper().strip() - # Prefer first occurrence; ignore duplicates/altlocs - if key not in name_to_pos: - name_to_pos[key] = res_atom_positions[i] - - # Place atoms into atom37 by matching stored atom names to atom37 indices. - # This handles all atoms present in the flat representation, not just - # the canonical residue_atoms for this residue type. This preserves - # atoms that were in the original atom37_mask even if they're atypical - # for the residue (e.g., from alternate conformations or data quirks). - for atom_name_str, pos in name_to_pos.items(): - idx37 = residue_constants.atom_order.get(atom_name_str) - if idx37 is not None: - atom37_positions[res_idx, idx37] = pos - atom37_mask[res_idx, idx37] = True - - # Create arrays that match sequence length (including chain breaks) - # Initialize arrays with proper size - chain_id_expanded = np.full(sequence_length, -1, dtype=np.int64) - entity_id_expanded = np.full(sequence_length, -1, dtype=np.int64) - sym_id_expanded = np.zeros(sequence_length, dtype=np.int64) - residue_index_expanded = np.zeros(sequence_length, dtype=np.int64) - insertion_code_expanded = np.array([""] * sequence_length, dtype=object) - confidence_expanded = np.zeros(sequence_length, dtype=np.float32) - atom37_positions_expanded = np.full( - (sequence_length, 37, 3), np.nan, dtype=np.float32 - ) - atom37_mask_expanded = np.zeros((sequence_length, 37), dtype=bool) - - # Map residue data to sequence positions (skipping chain breaks) - residue_idx = 0 - residue_counter_per_chain = {} - - for seq_pos, char in enumerate(single_letter_sequence): - if char != "|": - # This is a residue position - chain_id_val = protein_chain_ids[residue_idx] - - chain_id_expanded[seq_pos] = chain_id_val - entity_id_expanded[seq_pos] = chain_id_val # Simplified mapping - - # Track residue numbering per chain - if chain_id_val not in residue_counter_per_chain: - residue_counter_per_chain[chain_id_val] = 1 - else: - residue_counter_per_chain[chain_id_val] += 1 - - residue_index_expanded[seq_pos] = residue_counter_per_chain[ - chain_id_val - ] - confidence_expanded[seq_pos] = protein_confidence[residue_idx] - atom37_positions_expanded[seq_pos] = atom37_positions[residue_idx] - atom37_mask_expanded[seq_pos] = atom37_mask[residue_idx] - - residue_idx += 1 - # Chain break positions keep default values (-1, False, etc.) - - # Use the expanded arrays - chain_id = chain_id_expanded - entity_id = entity_id_expanded - sym_id = sym_id_expanded - residue_index = residue_index_expanded - insertion_code = insertion_code_expanded - protein_confidence = confidence_expanded - atom37_positions = atom37_positions_expanded - atom37_mask = atom37_mask_expanded - - # Create protein complex metadata preserving chain information - # Convert MolecularComplex metadata to ProteinComplex format - unique_chain_ids = np.unique(protein_chain_ids) - entity_lookup = {int(cid): int(cid) for cid in unique_chain_ids} - chain_lookup = { - int(cid): self.metadata.chain_lookup.get(int(cid), chr(65 + int(cid))) - for cid in unique_chain_ids - } - - protein_metadata = ProteinComplexMetadata( - entity_lookup=entity_lookup, - chain_lookup=chain_lookup, - assembly_composition=self.metadata.assembly_composition, - ) - - return ProteinComplex( - id=self.id, - sequence=single_letter_sequence, - entity_id=entity_id, - chain_id=chain_id, - sym_id=sym_id, - residue_index=residue_index, - insertion_code=insertion_code, - atom37_positions=atom37_positions, - atom37_mask=atom37_mask, - confidence=protein_confidence, - metadata=protein_metadata, - ) - - @classmethod - def from_mmcif(cls, inp: str, id: str | None = None) -> "MolecularComplex": - """Read MolecularComplex from mmcif file or string. - - Args: - inp: Path to mmCIF file or mmCIF content as string - id: Optional identifier to assign to the complex - - Returns: - MolecularComplex with all molecules (proteins, ligands, nucleic acids) - """ - from io import StringIO - - # Check if input is a file path or mmCIF string content - if os.path.exists(inp): - # Input is a file path - mmcif_file = pdbx.CIFFile.read(inp) - else: - # Input is mmCIF string content - mmcif_file = pdbx.CIFFile.read(StringIO(inp)) - - # Get structure - handle missing model information gracefully - try: - structure = pdbx.get_structure( - mmcif_file, model=1, extra_fields=["b_factor"] - ) - except (KeyError, ValueError): - # Fallback for mmCIF files without model information - try: - structure = pdbx.get_structure(mmcif_file) - except Exception: - # Last resort: use the first available model or all atoms - structure = pdbx.get_structure(mmcif_file, model=None) - # Type hint for pyright - structure is an AtomArray which is iterable - if TYPE_CHECKING: - structure: Any = structure - - # Read label_asym_id from the raw CIF atom_site category. - # Biotite's atom.chain_id uses auth_asym_id, which collapses ligands - # onto their parent protein chain. label_asym_id gives each entity a - # distinct chain identifier. - block = mmcif_file.block - label_asym_ids: list[str] | None = None - if "atom_site" in block: - atom_site = block["atom_site"] - if "label_asym_id" in atom_site: - _col = atom_site["label_asym_id"] - _raw = ( - _col.as_array(str) - if hasattr(_col, "as_array") - else np.array(list(_col), dtype=str) # type: ignore[arg-type] - ) - # biotite's get_structure(model=1) filters to model 1 AND - # removes alternate conformations. We must apply the same - # filters to label_asym_id to keep arrays aligned. - keep = np.ones(len(_raw), dtype=bool) - if "pdbx_PDB_model_num" in atom_site: - _mc = atom_site["pdbx_PDB_model_num"] - _models = ( - _mc.as_array(str) - if hasattr(_mc, "as_array") - else np.array(list(_mc), dtype=str) # type: ignore[arg-type] - ) - keep &= _models == "1" - if "label_alt_id" in atom_site: - _ac = atom_site["label_alt_id"] - _alts = ( - _ac.as_array(str) - if hasattr(_ac, "as_array") - else np.array(list(_ac), dtype=str) # type: ignore[arg-type] - ) - keep &= np.isin(_alts, [".", "?", "", "A"]) - filtered = _raw[keep] - if len(filtered) == len(structure): - label_asym_ids = filtered.tolist() - # If lengths still don't match, fall back to atom.chain_id - - # Get entity information from mmCIF - entity_info = {} - try: - if "entity" in block: - entity_category = block["entity"] - if "id" in entity_category and "type" in entity_category: - entity_ids = entity_category["id"] - entity_types = entity_category["type"] - # Convert CIFColumn to list for iteration - if hasattr(entity_ids, "__iter__") and hasattr( - entity_types, "__iter__" - ): - # Type annotation to help pyright understand these are iterable - entity_ids_list = list(entity_ids) # type: ignore - entity_types_list = list(entity_types) # type: ignore - for eid, etype in zip(entity_ids_list, entity_types_list): - entity_info[eid] = etype - except Exception: - pass - - # Initialize arrays for flat atom representation - sequence_tokens = [] - flat_positions = [] - flat_elements = [] - flat_names = [] - flat_hetero = [] - token_to_atoms = [] - confidence_scores = [] - chain_ids = [] # Track chain IDs for each token - - atom_idx = 0 - - # Group atoms by chain and residue. - # Use label_asym_id (distinct per entity) when available, otherwise - # fall back to biotite's chain_id (auth_asym_id). - chain_residue_groups: dict[str, dict[tuple[int, str], dict]] = {} - for atom_i, atom in enumerate(structure): - chain_id = ( - label_asym_ids[atom_i] if label_asym_ids is not None else atom.chain_id - ) - res_id = atom.res_id - res_name = atom.res_name - - if chain_id not in chain_residue_groups: - chain_residue_groups[chain_id] = {} - # Key by (res_id, res_name) to distinguish residues that share - # the same res_id but have different res_name (e.g. a protein - # residue and a ligand that were on the same auth chain). - res_key = (res_id, res_name) - if res_key not in chain_residue_groups[chain_id]: - chain_residue_groups[chain_id][res_key] = { - "atoms": [], - "res_name": res_name, - "is_hetero": atom.hetero, - } - chain_residue_groups[chain_id][res_key]["atoms"].append(atom) - - # Create a mapping from chain_id to numeric indices - chain_id_to_numeric = { - chain_id: idx - for idx, chain_id in enumerate(sorted(chain_residue_groups.keys())) - } - - # Process each chain and residue - for chain_id in sorted(chain_residue_groups.keys()): - residues = chain_residue_groups[chain_id] - numeric_chain_id = chain_id_to_numeric[chain_id] - - for res_key in sorted(residues.keys()): - residue_data = residues[res_key] - res_name = residue_data["res_name"] - atoms = residue_data["atoms"] - is_hetero = residue_data["is_hetero"] - - # Skip water molecules - if res_name == "HOH": - continue - - # Determine token name - if not is_hetero and res_name in residue_constants.restype_3to1: - # Standard amino acid - token_name = res_name - elif res_name in ["A", "T", "G", "C", "U", "DA", "DT", "DG", "DC"]: - # Nucleotide - token_name = res_name - else: - # Ligand or other molecule - token_name = res_name - - sequence_tokens.append(token_name) - chain_ids.append( - numeric_chain_id - ) # Store the numeric chain ID for this token - token_start = atom_idx - - # Add all atoms from this residue - for atom in atoms: - flat_positions.append(atom.coord) - - # Get element character - element = atom.element - flat_elements.append(element) - - # Get atom name - atom_name = atom.atom_name - flat_names.append(atom_name) - - # Get hetero flag - hetero_flag = atom.hetero - flat_hetero.append(hetero_flag) - - atom_idx += 1 - - # Record token-to-atom mapping - token_to_atoms.append([token_start, atom_idx]) - - # Add confidence score (B-factor if available, otherwise 1.0) - bfactor = getattr(atoms[0], "b_factor", 50.0) if atoms else 50.0 - confidence_scores.append(min(bfactor / 100.0, 1.0)) - - # Convert to numpy arrays - if not flat_positions: - # Create minimal arrays if no atoms found - atom_positions = np.zeros((0, 3), dtype=np.float32) - atom_elements = np.zeros(0, dtype=object) - atom_names = np.zeros(0, dtype=object) - atom_hetero = np.zeros(0, dtype=bool) - token_to_atoms_array = np.zeros((len(sequence_tokens), 2), dtype=np.int32) - chain_id_array = ( - np.array(chain_ids, dtype=np.int64) - if chain_ids - else np.zeros(len(sequence_tokens), dtype=np.int64) - ) - else: - atom_positions = np.array(flat_positions, dtype=np.float32) - atom_elements = np.array(flat_elements, dtype=object) - atom_names = np.array(flat_names, dtype=object) - atom_hetero = np.array(flat_hetero, dtype=bool) - token_to_atoms_array = np.array(token_to_atoms, dtype=np.int32) - chain_id_array = np.array(chain_ids, dtype=np.int64) - - confidence_array = np.array(confidence_scores, dtype=np.float32) - - # Create metadata using the chain_id_to_numeric mapping - if chain_residue_groups: - chain_lookup = { - numeric_id: chain_id - for chain_id, numeric_id in chain_id_to_numeric.items() - } - else: - chain_lookup = {} - - metadata = MolecularComplexMetadata( - entity_lookup=entity_info, - chain_lookup=chain_lookup, - assembly_composition=None, - ) - - # Set complex ID - if input was a path, use the stem; otherwise use default - if os.path.exists(inp): - complex_id = id or Path(inp).stem - else: - complex_id = id or "complex_from_string" - - return cls( - id=complex_id, - sequence=sequence_tokens, - atom_positions=atom_positions, - atom_elements=atom_elements, - token_to_atoms=token_to_atoms_array, - chain_id=chain_id_array, - plddt=confidence_array, - metadata=metadata, - atom_names=atom_names, - atom_hetero=atom_hetero, - ) - - def _get_entity_mapping( - self, - ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]: - """Compute chain→sequence, chain→entity_id, and entity_id→sequence mappings. - - Returns: - (chain_sequences, chain_to_entity, entity_sequences) - """ - chain_sequences: dict[str, list[str]] = {} - for token_idx in range(len(self.token_to_atoms)): - chain_id_numeric = self.chain_id[token_idx] - chain_id_str = self.metadata.chain_lookup.get( - int(chain_id_numeric), chr(65 + int(chain_id_numeric)) - ) - if chain_id_str not in chain_sequences: - chain_sequences[chain_id_str] = [] - chain_sequences[chain_id_str].append(self.sequence[token_idx]) - - sequence_to_entity: dict[tuple[str, ...], int] = {} - chain_to_entity: dict[str, int] = {} - entity_sequences: dict[int, tuple[str, ...]] = {} - entity_id_counter = 1 - for chain_id_str, sequence in chain_sequences.items(): - seq_tuple = tuple(sequence) - if seq_tuple not in sequence_to_entity: - sequence_to_entity[seq_tuple] = entity_id_counter - entity_sequences[entity_id_counter] = seq_tuple - entity_id_counter += 1 - chain_to_entity[chain_id_str] = sequence_to_entity[seq_tuple] - - return chain_sequences, chain_to_entity, entity_sequences - - def _add_entity_information( - self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]] - ) -> None: - """Add _entity category to CIF file so OST can identify ligands vs polymers.""" - - entity_ids: list[str] = [] - entity_types: list[str] = [] - entity_descriptions: list[str] = [] - for eid in sorted(entity_sequences.keys()): - seq = entity_sequences[eid] - entity_ids.append(str(eid)) - has_protein = any(t in residue_constants.restype_3to1 for t in seq) - has_na = any( - t in ("A", "T", "G", "C", "U", "DA", "DT", "DG", "DC") for t in seq - ) - if has_protein or has_na: - entity_types.append("polymer") - if has_protein: - entity_descriptions.append(f"Polymer entity {eid} (protein)") - else: - entity_descriptions.append(f"Polymer entity {eid} (nucleic acid)") - else: - entity_types.append("non-polymer") - entity_descriptions.append(f"Non-polymer entity {eid}") - - if entity_ids: - cif_file.block["entity"] = CIFCategory( - name="entity", - columns={ - "id": CIFColumn( - data=CIFData(array=np.array(entity_ids), dtype=np.str_) - ), - "type": CIFColumn( - data=CIFData(array=np.array(entity_types), dtype=np.str_) - ), - "pdbx_description": CIFColumn( - data=CIFData(array=np.array(entity_descriptions), dtype=np.str_) - ), - }, - ) - - # Add _struct_asym to map chain IDs to entity IDs - _, chain_to_entity, _ = self._get_entity_mapping() - if chain_to_entity: - asym_ids = sorted(chain_to_entity.keys()) - asym_entity_ids = [str(chain_to_entity[c]) for c in asym_ids] - cif_file.block["struct_asym"] = CIFCategory( - name="struct_asym", - columns={ - "id": CIFColumn( - data=CIFData(array=np.array(asym_ids), dtype=np.str_) - ), - "entity_id": CIFColumn( - data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_) - ), - }, - ) - - def to_mmcif(self) -> str: - """Write MolecularComplex to mmcif string using biotite. - - Returns: - String representation of the complex in mmCIF format - """ - # Pre-allocate AtomArray - n_atoms = len(self.atom_positions) - atom_array = bs.AtomArray(length=n_atoms) - - # Set coordinates directly (already vectorized) - atom_array.coord = self.atom_positions - - # Pre-allocate per-atom arrays - atom_res_ids = np.zeros(n_atoms, dtype=np.int32) - atom_chain_ids = np.empty(n_atoms, dtype=object) - atom_res_names = np.empty(n_atoms, dtype=object) - atom_hetero = np.zeros(n_atoms, dtype=bool) - atom_bfactors = np.zeros(n_atoms, dtype=np.float32) - atom_names = np.empty(n_atoms, dtype=object) - - # Build entity mappings: chains with identical sequences share entity ID - _, chain_to_entity, entity_sequences = self._get_entity_mapping() - - atom_entity_ids = np.zeros(n_atoms, dtype=np.int32) - - # Track residue IDs per chain - chain_res_counters: dict[int, int] = {} - - # Vectorized expansion of token-level to atom-level annotations - for token_idx, (start, end) in enumerate(self.token_to_atoms): - token = self.sequence[token_idx] - chain_id_numeric = self.chain_id[token_idx] - chain_id_str = self.metadata.chain_lookup.get( - int(chain_id_numeric), chr(65 + int(chain_id_numeric)) - ) - - # Track residue numbering per chain - if chain_id_numeric not in chain_res_counters: - chain_res_counters[chain_id_numeric] = 1 - res_id = chain_res_counters[chain_id_numeric] - chain_res_counters[chain_id_numeric] += 1 - - # Determine if protein - is_protein = token in residue_constants.restype_3to1 - - # Get atom names for this residue - if self.atom_names is not None: - # Use stored atom names (preserves original names from mmCIF) - names = list(self.atom_names[start:end]) - elif is_protein: - # Fallback: use standard protein atom names - standard_names = residue_constants.residue_atoms.get( - token, ["N", "CA", "C", "O"] - ) - names = standard_names[: end - start] - # Pad if needed - while len(names) < (end - start): - names.append(f"X{len(names)+1}") - else: - # Fallback: generate names for ligands/nucleic acids - names = [f"C{i+1}" for i in range(end - start)] - - # Vectorized assignment for this token's atoms - atom_res_ids[start:end] = res_id - atom_chain_ids[start:end] = chain_id_str - atom_res_names[start:end] = token - # Use stored hetero flags if available, otherwise guess based on protein status - if self.atom_hetero is not None: - atom_hetero[start:end] = self.atom_hetero[start:end] - else: - atom_hetero[start:end] = not is_protein - atom_bfactors[start:end] = self.plddt[token_idx] * 100.0 - atom_names[start:end] = names - atom_entity_ids[start:end] = chain_to_entity.get(chain_id_str, 1) - - # Set all AtomArray attributes at once (convert object arrays to proper string arrays) - # res_name uses U8 to accommodate CCD codes up to 5 characters (e.g., A1AZ2); - # chain_id uses U16 because chain names like ``ligand_1`` / ``ligand_2`` / - # auth-asym ids of arbitrary length are possible. - atom_array.res_id = atom_res_ids - atom_array.chain_id = np.array(atom_chain_ids, dtype="U16") - atom_array.res_name = np.array(atom_res_names, dtype="U8") - atom_array.hetero = atom_hetero - atom_array.atom_name = np.array(atom_names, dtype="U4") - atom_array.add_annotation("b_factor", dtype=float) - atom_array.b_factor = atom_bfactors - atom_array.add_annotation("entity_id", dtype=int) - atom_array.entity_id = atom_entity_ids - - # Use existing elements or infer them from atom names - if self.atom_elements is not None and len(self.atom_elements) == n_atoms: - # Convert object array to proper string array for biotite - atom_array.element = np.array(self.atom_elements, dtype="U4") - else: - # Use biotite's built-in element inference - atom_array.element = bs.infer_elements(atom_array) - - # Create CIF file and set structure - cif_file = CIFFile() - set_structure(cif_file, atom_array, data_block=self.id) - - # Manually fix label_entity_id (biotite doesn't use entity_id annotation correctly) - if "atom_site" in cif_file.block: - atom_site = cif_file.block["atom_site"] - if "label_asym_id" in atom_site and "label_entity_id" in atom_site: - label_asym_ids = atom_site["label_asym_id"] - if hasattr(label_asym_ids, "as_array"): - chain_ids_list = label_asym_ids.as_array(str).tolist() - elif hasattr(label_asym_ids, "__iter__"): - chain_ids_list = list(label_asym_ids) # type: ignore[arg-type] - else: - chain_ids_list = [] - updated_entity_ids = [ - str(chain_to_entity.get(cid, 1)) for cid in chain_ids_list - ] - if updated_entity_ids: - atom_site["label_entity_id"] = CIFColumn( - data=CIFData(array=np.array(updated_entity_ids), dtype=np.str_) - ) - - # Add _entity category for OST compatibility - self._add_entity_information(cif_file, entity_sequences) - - # Convert to string - output = io.StringIO() - cif_file.write(output) - return output.getvalue() - - def dockq(self, native: "MolecularComplex") -> Any: - """Compute DockQ score against native structure. - - Args: - native: Native MolecularComplex to compute DockQ against - - Returns: - DockQ result containing score and alignment information - """ - # Imports moved to top of file - - # Convert both complexes to ProteinComplex format for DockQ computation - # This extracts only the protein portion and converts to PDB format - try: - self_pc = self.to_protein_complex() - native_pc = native.to_protein_complex() - except ValueError as e: - raise ValueError( - f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {e}" - ) - - # Normalize chain IDs for PDB compatibility - self_pc = self_pc.normalize_chain_ids_for_pdb() - native_pc = native_pc.normalize_chain_ids_for_pdb() - - # Use the existing ProteinComplex.dockq() method - try: - dockq_result = self_pc.dockq(native_pc) - return dockq_result - except Exception: - # Fallback to manual DockQ computation if ProteinComplex.dockq() fails - return self._compute_dockq_manual(native) - - def _compute_dockq_manual(self, native: "MolecularComplex") -> Any: - """Manual DockQ computation fallback.""" - # Imports moved to top of file - - # Convert both complexes to ProteinComplex format - try: - self_pc = self.to_protein_complex() - native_pc = native.to_protein_complex() - except ValueError as e: - raise ValueError( - f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {e}" - ) - - # Normalize chain IDs for PDB compatibility - self_pc = self_pc.normalize_chain_ids_for_pdb() - native_pc = native_pc.normalize_chain_ids_for_pdb() - - # Write temporary PDB files and run DockQ - with TemporaryDirectory() as tdir: - dir_path = Path(tdir) - self_pdb = dir_path / "self.pdb" - native_pdb = dir_path / "native.pdb" - - # Write PDB files - self_pc.to_pdb(self_pdb) - native_pc.to_pdb(native_pdb) - - # Run DockQ - try: - output = check_output(["DockQ", str(self_pdb), str(native_pdb)]) - output_text = output.decode() - - # Parse DockQ output - lines = output_text.split("\n") - - # Find the total DockQ score - dockq_score = None - for line in lines: - if "Total DockQ" in line: - match = re.search(r"Total DockQ.*: ([\d.]+)", line) - if match: - dockq_score = float(match.group(1)) - break - - if dockq_score is None: - # Try to find individual DockQ scores - for line in lines: - if line.startswith("DockQ") and ":" in line: - try: - dockq_score = float(line.split(":")[1].strip()) - break - except (ValueError, IndexError): - continue - - if dockq_score is None: - raise ValueError("Could not parse DockQ score from output") - - # Return a simple result structure - return { - "total_dockq": dockq_score, - "raw_output": output_text, - "aligned": self, # Return self as aligned structure - } - - except FileNotFoundError: - raise RuntimeError( - "DockQ is not installed. Please install DockQ to use this method." - ) - except Exception as e: - raise RuntimeError(f"DockQ computation failed: {e}") - - def rmsd(self, target: "MolecularComplex", **kwargs) -> float: - """Compute RMSD against target structure. - - Args: - target: Target MolecularComplex to compute RMSD against - **kwargs: Additional arguments passed to compute_rmsd - - Returns: - float: RMSD value between the two structures - """ - # Imports moved to top of file - - # Ensure both complexes have the same number of tokens - if len(self) != len(target): - raise ValueError( - f"Complexes must have the same number of tokens: {len(self)} vs {len(target)}" - ) - - # Extract center positions for each token (using centroid of atoms) - mobile_coords = [] - target_coords = [] - atom_mask = [] - - for i in range(len(self)): - # Get atom positions for this token - mobile_start, mobile_end = self.token_to_atoms[i] - target_start, target_end = target.token_to_atoms[i] - - # Extract atom positions - mobile_atoms = self.atom_positions[mobile_start:mobile_end] - target_atoms = target.atom_positions[target_start:target_end] - - # Check if both tokens have atoms - if len(mobile_atoms) == 0 or len(target_atoms) == 0: - # Skip tokens with no atoms - continue - - # For simplicity, use the centroid of atoms as the representative position - mobile_center = mobile_atoms.mean(axis=0) - target_center = target_atoms.mean(axis=0) - - mobile_coords.append(mobile_center) - target_coords.append(target_center) - atom_mask.append(True) - - if len(mobile_coords) == 0: - raise ValueError("No valid atoms found for RMSD computation") - - # Convert to tensors - mobile_tensor = torch.from_numpy(np.stack(mobile_coords, axis=0)).unsqueeze( - 0 - ) # [1, N, 3] - target_tensor = torch.from_numpy(np.stack(target_coords, axis=0)).unsqueeze( - 0 - ) # [1, N, 3] - mask_tensor = torch.tensor(atom_mask, dtype=torch.bool).unsqueeze(0) # [1, N] - - # Compute RMSD using existing infrastructure - rmsd_value = compute_rmsd( - mobile=mobile_tensor, - target=target_tensor, - atom_exists_mask=mask_tensor, - reduction="batch", - **kwargs, - ) - - return float(rmsd_value) - - def lddt_ca(self, target: "MolecularComplex", **kwargs) -> float: - """Compute LDDT score against target structure. - - Args: - target: Target MolecularComplex to compute LDDT against - **kwargs: Additional arguments passed to compute_lddt - - Returns: - float: LDDT value between the two structures - """ - # Imports moved to top of file - - # Ensure both complexes have the same number of tokens - if len(self) != len(target): - raise ValueError( - f"Complexes must have the same number of tokens: {len(self)} vs {len(target)}" - ) - - # Extract center positions for each token (using centroid of atoms) - mobile_coords = [] - target_coords = [] - atom_mask = [] - - for i in range(len(self)): - # Get atom positions for this token - mobile_start, mobile_end = self.token_to_atoms[i] - target_start, target_end = target.token_to_atoms[i] - - # Extract atom positions - mobile_atoms = self.atom_positions[mobile_start:mobile_end] - target_atoms = target.atom_positions[target_start:target_end] - - # Check if both tokens have atoms - if len(mobile_atoms) == 0 or len(target_atoms) == 0: - # Skip tokens with no atoms - mobile_coords.append(np.full(3, np.nan)) - target_coords.append(np.full(3, np.nan)) - atom_mask.append(False) - continue - - # For simplicity, use the centroid of atoms as the representative position - mobile_center = mobile_atoms.mean(axis=0) - target_center = target_atoms.mean(axis=0) - - mobile_coords.append(mobile_center) - target_coords.append(target_center) - atom_mask.append(True) - - if not any(atom_mask): - raise ValueError("No valid atoms found for LDDT computation") - - # Convert to tensors - mobile_tensor = torch.from_numpy(np.stack(mobile_coords, axis=0)).unsqueeze( - 0 - ) # [1, N, 3] - target_tensor = torch.from_numpy(np.stack(target_coords, axis=0)).unsqueeze( - 0 - ) # [1, N, 3] - mask_tensor = torch.tensor(atom_mask, dtype=torch.bool).unsqueeze(0) # [1, N] - - # Compute LDDT using existing infrastructure - lddt_value = compute_lddt( - all_atom_pred_pos=mobile_tensor, - all_atom_positions=target_tensor, - all_atom_mask=mask_tensor, - per_residue=False, # Return overall LDDT score - **kwargs, - ) - - return float(lddt_value) - - def state_dict(self): - """This state dict is optimized for storage, so it turns things to fp16 whenever - possible and converts numpy arrays to lists for JSON serialization. - """ - dct = {k: v for k, v in vars(self).items()} - for k, v in dct.items(): - if isinstance(v, np.ndarray): - match v.dtype: - case np.int64: - dct[k] = v.astype(np.int32).tolist() - case np.float64 | np.float32: - dct[k] = v.astype(np.float16).tolist() - case _: - dct[k] = v.tolist() - elif isinstance(v, MolecularComplexMetadata): - dct[k] = asdict(v) - - return dct - - def to_blob(self) -> bytes: - return brotli.compress(msgpack.dumps(self.state_dict()), quality=5) - - @classmethod - def from_state_dict(cls, dct): - for k, v in dct.items(): - if isinstance(v, list) and k in [ - "atom_positions", - "atom_elements", - "atom_names", - "atom_hetero", - "token_to_atoms", - "chain_id", - "plddt", - ]: - dct[k] = np.array(v) - - for k, v in dct.items(): - if isinstance(v, np.ndarray): - if k in ["atom_positions", "plddt"]: - dct[k] = v.astype(np.float32) - elif k in ["token_to_atoms", "chain_id"]: - dct[k] = ( - v.astype(np.int32) - if k == "token_to_atoms" - else v.astype(np.int64) - ) - - dct["metadata"] = MolecularComplexMetadata(**dct["metadata"]) - - # Backward compatibility: if chain_id is missing, create default array - if "chain_id" not in dct: - # Default all tokens to chain 0 - dct["chain_id"] = np.zeros(len(dct["sequence"]), dtype=np.int64) - - return cls(**dct) - - @classmethod - def from_blob(cls, input: Path | str | io.BytesIO | bytes): - match input: - case Path() | str(): - bytes = Path(input).read_bytes() - case io.BytesIO(): - bytes = input.getvalue() - case _: - bytes = input - return cls.from_state_dict( - msgpack.loads(brotli.decompress(bytes), strict_map_key=False) - ) diff --git a/fastplms/esmfold2/esmfold2_msa.py b/fastplms/esmfold2/esmfold2_msa.py deleted file mode 100644 index 4bca6b5..0000000 --- a/fastplms/esmfold2/esmfold2_msa.py +++ /dev/null @@ -1,506 +0,0 @@ -from __future__ import annotations - -import dataclasses -import string -from dataclasses import dataclass -from functools import cached_property -from itertools import islice -from typing import Sequence - -import numpy as np -from Bio import SeqIO -from scipy.spatial.distance import cdist - -from .esmfold2_misc import slice_any_object -from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter -from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences -from .esmfold2_sequential_dataclass import SequentialDataclass -from .esmfold2_system import PathOrBuffer - -REMOVE_LOWERCASE_TRANSLATION = str.maketrans(dict.fromkeys(string.ascii_lowercase)) - - -def remove_insertions_from_sequence(seq: str) -> str: - return seq.translate(REMOVE_LOWERCASE_TRANSLATION) - - -@dataclass(frozen=True) -class MSA(SequentialDataclass): - """Object-oriented interface to an MSA. - - Args: - sequences (list[str]): List of protein sequences - headers (list[str]): List of headers describing the sequences - - """ - - entries: list[FastaEntry] - - @cached_property - def sequences(self) -> list[str]: - return [entry.sequence for entry in self.entries] - - @cached_property - def headers(self) -> list[str]: - return [entry.header for entry in self.entries] - - def __repr__(self): - return ( - f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})" - ) - - def to_fast_msa(self) -> FastMSA: - return FastMSA(self.array, self.headers) - - @classmethod - def from_a3m( - cls, - path: PathOrBuffer, - remove_insertions: bool = True, - max_sequences: int | None = None, - ) -> MSA: - entries = [] - for header, seq in islice(read_sequences(path), max_sequences): - if remove_insertions: - seq = remove_insertions_from_sequence(seq) - if entries: - assert ( - len(seq) == len(entries[0].sequence) - ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}" - entries.append(FastaEntry(header, seq)) - return cls(entries) - - def to_a3m(self, path: PathOrBuffer) -> None: - write_sequences(self.entries, path) - - @classmethod - def from_stockholm( - cls, - path: PathOrBuffer, - remove_insertions: bool = True, - max_sequences: int | None = None, - ) -> MSA: - entries = [] - for record in islice(SeqIO.parse(path, "stockholm"), max_sequences): - header = f"{record.id} {record.description}" - seq = str(record.seq) - if entries: - assert ( - len(seq) == len(entries[0].sequence) - ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}" - entries.append(FastaEntry(header, seq)) - msa = cls(entries) - if remove_insertions: - keep_inds = [i for i, aa in enumerate(msa.query) if aa != "-"] - msa = msa.select_positions(keep_inds) - return msa - - def to_bytes(self) -> bytes: - version = 1 - version_bytes = version.to_bytes(1, "little") - seqlen_bytes = self.seqlen.to_bytes(4, "little") - depth_bytes = self.depth.to_bytes(4, "little") - array_bytes = self.array.tobytes() - header_bytes = "\n".join(entry.header for entry in self.entries).encode() - all_bytes = ( - version_bytes + seqlen_bytes + depth_bytes + array_bytes + header_bytes - ) - return all_bytes - - @classmethod - def from_bytes(cls, data: bytes) -> MSA: - version_bytes, seqlen_bytes, depth_bytes, data = ( - data[:1], - data[1:5], - data[5:9], - data[9:], - ) - version = int.from_bytes(version_bytes, "little") - if version != 1: - raise ValueError(f"Unsupported version: {version}") - seqlen = int.from_bytes(seqlen_bytes, "little") - depth = int.from_bytes(depth_bytes, "little") - array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :] - array = np.frombuffer(array_bytes, dtype="|S1") - array = array.reshape(depth, seqlen) - headers = header_bytes.decode().split("\n") - # Sometimes the separation is two newlines, which results in an empty header. - headers = [header for header in headers if header] - # If all headers were empty (e.g., saved from from_sequences), use empty headers - if len(headers) == 0 and depth > 0: - headers = [""] * depth - entries = [ - FastaEntry(header, b"".join(row).decode()) - for header, row in zip(headers, array) - ] - return cls(entries) - - # TODO(jmaccarl): set remove_insertions to True by default here to match other utils - @classmethod - def from_sequences( - cls, sequences: list[str], remove_insertions: bool = False - ) -> MSA: - if remove_insertions: - entries = [ - FastaEntry("", remove_insertions_from_sequence(seq)) - for seq in sequences - ] - else: - entries = [FastaEntry("", seq) for seq in sequences] - return cls(entries) - - def to_sequence_bytes(self) -> bytes: - """Stores ONLY SEQUENCES in array format as bytes. Header information will be lost.""" - seqlen_bytes = self.seqlen.to_bytes(4, "little") - array_bytes = self.array.tobytes() - all_bytes = seqlen_bytes + array_bytes - return all_bytes - - @classmethod - def from_sequence_bytes(cls, data: bytes) -> MSA: - seqlen_bytes, array_bytes = data[:4], data[4:] - seqlen = int.from_bytes(seqlen_bytes, "little") - array = np.frombuffer(array_bytes, dtype="|S1") - array = array.reshape(-1, seqlen) - entries = [FastaEntry("", b"".join(row).decode()) for row in array] - return cls(entries) - - @property - def depth(self) -> int: - return len(self.entries) - - @property - def seqlen(self) -> int: - return len(self.entries[0].sequence) - - @cached_property - def array(self) -> np.ndarray: - return np.array([list(seq) for seq in self.sequences], dtype="|S1") - - @property - def query(self) -> str: - return self.entries[0].sequence - - def select_sequences(self, indices: Sequence[int] | np.ndarray) -> MSA: - """Subselect rows of the MSA.""" - entries = [self.entries[idx] for idx in indices] - return dataclasses.replace(self, entries=entries) - - def select_positions(self, indices: Sequence[int] | np.ndarray) -> MSA: - """Subselect columns of the MSA.""" - entries = [ - FastaEntry(header, "".join(seq[idx] for idx in indices)) - for header, seq in self.entries - ] - return dataclasses.replace(self, entries=entries) - - def __getitem__(self, indices: int | list[int] | slice | np.ndarray): - if isinstance(indices, int): - indices = [indices] - - entries = [ - FastaEntry(header, slice_any_object(seq, indices)) - for header, seq in self.entries - ] - return dataclasses.replace(self, entries=entries) - - def __len__(self): - return self.seqlen - - def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA: - """Greedily select sequences that either maximize or minimize hamming distance. - - Algorithm proposed in the MSA Transformer paper. Starting from the query sequence, - iteratively add sequences to the list with the maximum (minimum) average Hamming - distance to the existing set of sequences. - - Args: - num_seqs (int): Number of sequences to select. - mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless - you're doing it to prove a point for a paper. - - Returns: - MSA object w/ subselected sequences. - """ - assert mode in ("max", "min") - if self.depth <= num_seqs: - return self - - indices = greedy_select_indices(self.array, num_seqs, mode) - return self.select_sequences(indices) - - def hhfilter( - self, - seqid: int = 90, - diff: int = 0, - cov: int = 0, - qid: int = 0, - qsc: float = -20.0, - binary: str = "hhfilter", - ) -> MSA: - """Apply hhfilter to the sequences in the MSA and return a filtered MSA.""" - - indices = hhfilter( - self.sequences, - seqid=seqid, - diff=diff, - cov=cov, - qid=qid, - qsc=qsc, - binary=binary, - ) - return self.select_sequences(indices) - - def select_random_sequences(self, num_seqs: int) -> MSA: - """Uses random sampling to subselect sequences from the MSA. Always - keeps the query sequence. - """ - if num_seqs >= self.depth: - return self - - # Subselect random, always keeping the query sequence. - indices = np.sort( - np.append( - 0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1 - ) - ) - msa = self.select_sequences(indices) # type: ignore - return msa - - def select_diverse_sequences(self, num_seqs: int) -> MSA: - """Applies hhfilter to select ~num_seqs sequences, then uses random sampling - to subselect if necessary. - """ - if num_seqs >= self.depth: - return self - - msa = self.hhfilter(diff=num_seqs) - if num_seqs < msa.depth: - msa = msa.select_random_sequences(num_seqs) - return msa - - def pad_to_depth(self, depth: int) -> MSA: - if depth < self.depth: - raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") - elif depth == self.depth: - return self - - num_to_add = depth - self.depth - extra_entries = [FastaEntry("", "-" * self.seqlen) for _ in range(num_to_add)] - return dataclasses.replace(self, entries=self.entries + extra_entries) - - @classmethod - def stack( - cls, msas: Sequence[MSA], remove_query_from_later_msas: bool = True - ) -> MSA: - """Stack a series of MSAs. Optionally remove the query from msas after the first.""" - all_entries = [] - for i, msa in enumerate(msas): - entries = msa.entries - if i > 0 and remove_query_from_later_msas: - entries = entries[1:] - all_entries.extend(entries) - return cls(entries=all_entries) - - @cached_property - def seqid(self) -> np.ndarray: - array = self.array.view(np.uint8) - seqid = 1 - cdist(array[0][None], array, "hamming") - return seqid[0] - - @classmethod - def concat( - cls, - msas: Sequence[MSA], - join_token: str | None = "|", - allow_depth_mismatch: bool = False, - ) -> MSA: - """Concatenate a series of MSAs horizontally, along the sequence dimension.""" - if not msas: - raise ValueError("Cannot concatenate an empty list of MSAs") - msa_depths = [msa.depth for msa in msas] - if len(set(msa_depths)) != 1: - if not allow_depth_mismatch: - raise ValueError("Depth mismatch in concatenating MSAs") - else: - max_depth = max(msa_depths) - msas = [msa.pad_to_depth(max_depth) for msa in msas] - headers = [ - "|".join([str(h) for h in headers]) - for headers in zip(*(msa.headers for msa in msas)) - ] - - if join_token is None: - join_token = "" - - seqs = [join_token.join(vals) for vals in zip(*(msa.sequences for msa in msas))] - entries = [FastaEntry(header, seq) for header, seq in zip(headers, seqs)] - return cls(entries) - - -@dataclass(frozen=True) -class FastMSA(SequentialDataclass): - """Object-oriented interface to an MSA stored as a numpy uint8 array.""" - - array: np.ndarray - headers: list[str] | None = None - - def __post_init__(self): - if self.headers is not None: - assert ( - len(self.headers) == self.depth - ), "Number of headers must match depth." - - @classmethod - def from_bytes(cls, data: bytes) -> FastMSA: - version_bytes, seqlen_bytes, depth_bytes, data = ( - data[:1], - data[1:5], - data[5:9], - data[9:], - ) - version = int.from_bytes(version_bytes, "little") - if version != 1: - raise ValueError(f"Unsupported version: {version}") - seqlen = int.from_bytes(seqlen_bytes, "little") - depth = int.from_bytes(depth_bytes, "little") - array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :] - array = np.frombuffer(array_bytes, dtype="|S1") - array = array.reshape(depth, seqlen) - headers = header_bytes.decode().split("\n") - # Sometimes the separation is two newlines, which results in an empty header. - headers = [header for header in headers if header] - # If all headers were empty (e.g., saved from from_sequences), use empty headers - if len(headers) == 0 and depth > 0: - headers = [""] * depth - return cls(array, headers) - - @classmethod - def from_sequence_bytes(cls, data: bytes) -> FastMSA: - seqlen_bytes, array_bytes = data[:4], data[4:] - seqlen = int.from_bytes(seqlen_bytes, "little") - array = np.frombuffer(array_bytes, dtype="|S1") - array = array.reshape(-1, seqlen) - return cls(array) - - @property - def depth(self) -> int: - return self.array.shape[0] - - @property - def seqlen(self) -> int: - return self.array.shape[1] - - def __len__(self): - return self.seqlen - - def __getitem__(self, indices: int | list[int] | slice | np.ndarray): - if isinstance(indices, int): - indices = [indices] - - return dataclasses.replace(self, array=self.array[:, indices]) - - def select_sequences(self, indices: Sequence[int] | np.ndarray) -> FastMSA: - """Subselect rows of the MSA.""" - array = self.array[indices] - headers = ( - [self.headers[idx] for idx in indices] if self.headers is not None else None - ) - return dataclasses.replace(self, array=array, headers=headers) - - def select_random_sequences(self, num_seqs: int) -> FastMSA: - """Uses random sampling to subselect sequences from the MSA. Always - keeps the query sequence. - """ - if num_seqs >= self.depth: - return self - - # Subselect random, always keeping the query sequence. - indices = np.sort( - np.append( - 0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1 - ) - ) - msa = self.select_sequences(indices) # type: ignore - return msa - - def pad_to_depth(self, depth: int) -> FastMSA: - if depth < self.depth: - raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") - elif depth == self.depth: - return self - - num_to_add = depth - self.depth - array = np.pad( - self.array, - [(0, num_to_add), (0, 0)], - constant_values=ord("-") if self.array.dtype == np.uint8 else b"-", - ) - headers = self.headers - if headers is not None: - headers = headers + [""] * num_to_add - return dataclasses.replace(self, array=array, headers=headers) - - @classmethod - def concat( - cls, - msas: Sequence[FastMSA], - join_token: str | None = None, - allow_depth_mismatch: bool = False, - ) -> FastMSA: - """Concatenate a series of MSAs horizontally, along the sequence dimension.""" - if not msas: - raise ValueError("Cannot concatenate an empty list of MSAs") - if join_token is not None and join_token != "": - raise NotImplementedError("join_token is not supported for FastMSA") - - msa_depths = [msa.depth for msa in msas] - if len(set(msa_depths)) != 1: - if not allow_depth_mismatch: - raise ValueError("Depth mismatch in concatenating MSAs") - else: - max_depth = max(msa_depths) - msas = [msa.pad_to_depth(max_depth) for msa in msas] - headers = [ - "|".join([str(h) for h in headers]) - for headers in zip( - *( - msa.headers if msa.headers is not None else [""] * msa.depth - for msa in msas - ) - ) - ] - - array = np.concatenate([msa.array for msa in msas], axis=1) - return cls(array, headers) - - def to_msa(self) -> MSA: - headers = ( - self.headers - if self.headers is not None - else [f"seq{i}" for i in range(self.depth)] - ) - entries = [ - FastaEntry(header, b"".join(row).decode()) - for header, row in zip(headers, self.array) - ] - return MSA(entries) - - @classmethod - def stack( - cls, msas: Sequence[FastMSA], remove_query_from_later_msas: bool = True - ) -> FastMSA: - """Stack a series of MSAs. Optionally remove the query from msas after the first.""" - arrays = [] - all_headers = [] - for i, msa in enumerate(msas): - array = msa.array - headers = msa.headers - if i > 0 and remove_query_from_later_msas: - array = array[1:] - if headers is not None: - headers = headers[1:] - arrays.append(array) - if headers is not None: - all_headers.extend(headers) - return cls(np.concatenate(arrays, axis=0), all_headers) diff --git a/fastplms/esmfold2/esmfold2_msa_filter_sequences.py b/fastplms/esmfold2/esmfold2_msa_filter_sequences.py deleted file mode 100644 index 1adff75..0000000 --- a/fastplms/esmfold2/esmfold2_msa_filter_sequences.py +++ /dev/null @@ -1,82 +0,0 @@ -import os -import tempfile -from pathlib import Path - -import numpy as np -from scipy.spatial.distance import cdist - -from .esmfold2_system import run_subprocess_with_errorcheck - - -def greedy_select_indices(array, num_seqs: int, mode: str = "max") -> list[int]: - """Greedily select sequences that either maximize or minimize hamming distance. - - Algorithm proposed in the MSA Transformer paper. Starting from the query sequence, - iteratively add sequences to the list with the maximum (minimum) average Hamming - distance to the existing set of sequences. - - Args: - array (np.ndarray): Character array representing the sequences in the MSA - num_seqs (int): Number of sequences to select. - mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless - you're doing it to prove a point for a paper. - - Returns: - list[int]: List of indices to select from the array - """ - assert mode in ("max", "min") - depth = array.shape[0] - if depth <= num_seqs: - return list(range(depth)) - array = array.view(np.uint8) - - optfunc = np.argmax if mode == "max" else np.argmin - all_indices = np.arange(depth) - indices = [0] - pairwise_distances = np.zeros((0, depth)) - for _ in range(num_seqs - 1): - dist = cdist(array[indices[-1:]], array, "hamming") - pairwise_distances = np.concatenate([pairwise_distances, dist]) - shifted_distance = np.delete(pairwise_distances, indices, axis=1).mean(0) - shifted_index = optfunc(shifted_distance) - index = np.delete(all_indices, indices)[shifted_index] - indices.append(index) - indices = sorted(indices) - return indices - - -def hhfilter( - sequences: list[str], - seqid: int = 90, - diff: int = 0, - cov: int = 0, - qid: int = 0, - qsc: float = -20.0, - binary: str = "hhfilter", -) -> list[int]: - with tempfile.TemporaryDirectory( - dir="/dev/shm" if os.path.exists("/dev/shm") else None - ) as tempdirname: - tempdir = Path(tempdirname) - fasta_file = tempdir / "input.fasta" - fasta_file.write_text( - "\n".join(f">{i}\n{seq}" for i, seq in enumerate(sequences)) - ) - output_file = tempdir / "output.fasta" - command = " ".join( - [ - f"{binary}", - f"-i {fasta_file}", - "-M a3m", - f"-o {output_file}", - f"-id {seqid}", - f"-diff {diff}", - f"-cov {cov}", - f"-qid {qid}", - f"-qsc {qsc}", - ] - ).split(" ") - run_subprocess_with_errorcheck(command, capture_output=True) - with output_file.open() as f: - indices = [int(line[1:].strip()) for line in f if line.startswith(">")] - return indices diff --git a/fastplms/esmfold2/esmfold2_normalize_coordinates.py b/fastplms/esmfold2/esmfold2_normalize_coordinates.py deleted file mode 100644 index 1b184c8..0000000 --- a/fastplms/esmfold2/esmfold2_normalize_coordinates.py +++ /dev/null @@ -1,79 +0,0 @@ -from typing import TypeVar - -import numpy as np -import torch -from torch import Tensor - -from . import esmfold2_residue_constants as RC -from .esmfold2_affine3d import Affine3D - -ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) - - -def atom3_to_backbone_frames(bb_positions: torch.Tensor) -> Affine3D: - N, CA, C = bb_positions.unbind(dim=-2) - return Affine3D.from_graham_schmidt(C, CA, N) - - -def index_by_atom_name( - atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2 -) -> ArrayOrTensor: - squeeze = False - if isinstance(atom_names, str): - atom_names = [atom_names] - squeeze = True - indices = [RC.atom_order[atom_name] for atom_name in atom_names] - dim = dim % atom37.ndim - index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim)) - result = atom37[index] # type: ignore - if squeeze: - result = result.squeeze(dim) - return result - - -def get_protein_normalization_frame(coords: Tensor) -> Affine3D: - """Given a set of coordinates for a protein, compute a single frame that can be used to normalize the coordinates. - Specifically, we compute the average position of the N, CA, and C atoms use those 3 points to construct a frame - using the Gram-Schmidt algorithm. The average CA position is used as the origin of the frame. - - Args: - coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates - - Returns: - Affine3D: tensor of Affine3D frame - """ - bb_coords = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2) - coord_mask = torch.all(torch.all(torch.isfinite(bb_coords), dim=-1), dim=-1) - - average_position_per_n_ca_c = bb_coords.masked_fill( - ~coord_mask[..., None, None], 0 - ).sum(-3) / (coord_mask.sum(-1)[..., None, None] + 1e-8) - frame = atom3_to_backbone_frames(average_position_per_n_ca_c.float()) - - return frame - - -def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor: - """Given a set of coordinates and a single frame, apply the frame to the coordinates. - - Args: - coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates - frame (Affine3D): Affine3D frame - - Returns: - torch.FloatTensor: [L, 37, 3] tensor of transformed coordinates - """ - coords_trans_rot = frame[..., None, None].invert().apply(coords) - - # only transform coordinates with frame that have a valid rotation - valid_frame = frame.trans.norm(dim=-1) > 0 - - is_inf = torch.isinf(coords) - coords = coords_trans_rot.where(valid_frame[..., None, None, None], coords) - coords.masked_fill_(is_inf, torch.inf) - - return coords - - -def normalize_coordinates(coords: Tensor) -> Tensor: - return apply_frame_to_coords(coords, get_protein_normalization_frame(coords)) diff --git a/fastplms/esmfold2/esmfold2_output.py b/fastplms/esmfold2/esmfold2_output.py deleted file mode 100644 index 8db9caa..0000000 --- a/fastplms/esmfold2/esmfold2_output.py +++ /dev/null @@ -1,224 +0,0 @@ -from itertools import groupby -from typing import Any - -import numpy as np -import torch - -from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER -from .esmfold2_molecular_complex import ( - MolecularComplex, - MolecularComplexMetadata, -) - - -def get_element_symbol(atomic_num: int) -> str: - return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_num, "X") - - -def build_molecular_complex_from_features( - coords: torch.Tensor, - plddt: torch.Tensor, - atom_mask: torch.Tensor, - ref_element: torch.Tensor, - ref_atom_name_chars: torch.Tensor, - chain_infos: list, - complex_id: str, -) -> MolecularComplex: - """Construct a MolecularComplex from feature-dict tensors and chain metadata. - - Non-polymer chains (ligands) collapse all per-atom tokens into a single - residue token whose pLDDT is the per-token average and whose hetero flag - is True. - """ - mask_np = atom_mask.bool().cpu().numpy() - coords_np = coords.float().cpu().numpy() - name_chars_np = ref_atom_name_chars.cpu().numpy() - elements_np = ref_element.cpu().numpy() - plddt_np = plddt.float().cpu().numpy() - - sequence_tokens: list[str] = [] - chain_ids_per_token: list[int] = [] - token_to_atoms: list[list[int]] = [] - confidence: list[float] = [] - flat_positions: list[list[float]] = [] - flat_elements: list[str] = [] - flat_names: list[str] = [] - flat_hetero: list[bool] = [] - - chain_lookup: dict[int, str] = {} - entity_info: dict[int, str] = {} - out_atom_cursor = 0 - - for ci in chain_infos: - chain_lookup[ci.asym_id] = ci.chain_id - is_nonpolymer = ci.mol_type == MOL_TYPE_NONPOLYMER - entity_info[ci.entity_id] = "non-polymer" if is_nonpolymer else "polymer" - - if is_nonpolymer: - residue_name = ci.tokens[0].residue_name if ci.tokens else "LIG" - sequence_tokens.append(residue_name) - chain_ids_per_token.append(ci.asym_id) - avg_plddt = ( - float(np.mean([plddt_np[ti.token_index] for ti in ci.tokens])) - if ci.tokens - else 0.0 - ) - confidence.append(avg_plddt) - token_atom_start = out_atom_cursor - for ti in ci.tokens: - for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count): - if not mask_np[atom_idx]: - continue - flat_positions.append(coords_np[atom_idx].tolist()) - flat_elements.append(get_element_symbol(int(elements_np[atom_idx]))) - chars = name_chars_np[atom_idx] - name = "".join( - chr(int(c) + 32) for c in chars if int(c) != 0 - ).strip() - flat_names.append(name) - flat_hetero.append(True) - out_atom_cursor += 1 - token_to_atoms.append([token_atom_start, out_atom_cursor]) - continue - - # Atom-tokenized modified residues (HYP, MSE, ...) span multiple - # tokens per residue; collapse them back to one mmCIF residue. - for _residue_index, ti_iter in groupby( - ci.tokens, key=lambda t: t.residue_index - ): - ti_group = list(ti_iter) - sequence_tokens.append(ti_group[0].residue_name) - chain_ids_per_token.append(ci.asym_id) - confidence.append( - float(np.mean([plddt_np[ti.token_index] for ti in ti_group])) - ) - token_atom_start = out_atom_cursor - for ti in ti_group: - for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count): - if not mask_np[atom_idx]: - continue - flat_positions.append(coords_np[atom_idx].tolist()) - flat_elements.append(get_element_symbol(int(elements_np[atom_idx]))) - chars = name_chars_np[atom_idx] - name = "".join( - chr(int(c) + 32) for c in chars if int(c) != 0 - ).strip() - flat_names.append(name) - flat_hetero.append(False) - out_atom_cursor += 1 - token_to_atoms.append([token_atom_start, out_atom_cursor]) - - return MolecularComplex( - id=complex_id, - sequence=sequence_tokens, - atom_positions=np.array(flat_positions, dtype=np.float32).reshape(-1, 3), - atom_elements=np.array(flat_elements, dtype=object), - token_to_atoms=np.array(token_to_atoms, dtype=np.int32).reshape(-1, 2), - chain_id=np.array(chain_ids_per_token, dtype=np.int64), - plddt=np.array(confidence, dtype=np.float32), - atom_names=np.array(flat_names, dtype=object), - atom_hetero=np.array(flat_hetero, dtype=bool), - metadata=MolecularComplexMetadata( - entity_lookup=entity_info, - chain_lookup=chain_lookup, - assembly_composition=None, - ), - ) - - -def build_molecular_complex( - structure: Any, coords: torch.Tensor, plddt: torch.Tensor, complex_id: str -) -> MolecularComplex: - """Directly constructs a MolecularComplex from model outputs without intermediate files. - - Args: - structure: Object with .chains, .residues, .atoms numpy structured arrays. - coords: [N_atoms, 3] predicted atom coordinates. - plddt: [N_residues] per-residue confidence scores. - complex_id: Identifier string for the resulting complex. - """ - flat_positions = [] - flat_elements = [] - flat_names = [] - flat_hetero = [] - - sequence_tokens = [] - token_to_atoms = [] - chain_ids_per_token = [] - confidence_scores = [] - - chain_lookup = {} - entity_info = {} - - global_atom_cursor = 0 - global_res_cursor = 0 - atom_array_idx = 0 - - for chain in structure.chains: - chain_idx_numeric = chain["asym_id"] - chain_name_str = str(chain["name"]) - mol_type = chain["mol_type"] - - chain_lookup[chain_idx_numeric] = chain_name_str - entity_info[chain["entity_id"]] = ( - "polymer" if mol_type != MOL_TYPE_NONPOLYMER else "non-polymer" - ) - - res_start = chain["res_idx"] - res_end = chain["res_idx"] + chain["res_num"] - residues = structure.residues[res_start:res_end] - - for residue in residues: - res_name = str(residue["name"]) - - sequence_tokens.append(res_name) - chain_ids_per_token.append(chain_idx_numeric) - - score = plddt[global_res_cursor].item() - confidence_scores.append(score) - token_start_idx = atom_array_idx - - atom_start = residue["atom_idx"] - atom_end = residue["atom_idx"] + residue["atom_num"] - atoms = structure.atoms[atom_start:atom_end] - - for atom in atoms: - if not atom["is_present"]: - continue - - pos = coords[global_atom_cursor].tolist() - flat_positions.append(pos) - - elem = get_element_symbol(atom["element"].item()) - flat_elements.append(elem) - - raw_name = atom["name"] - if hasattr(raw_name, "tolist"): - raw_name = raw_name.tolist() - name_str = "".join([chr(c + 32) for c in raw_name if c != 0]) - flat_names.append(name_str) - - flat_hetero.append(mol_type == MOL_TYPE_NONPOLYMER) - - global_atom_cursor += 1 - atom_array_idx += 1 - - token_to_atoms.append([token_start_idx, atom_array_idx]) - global_res_cursor += 1 - - return MolecularComplex( - id=complex_id, - sequence=sequence_tokens, - atom_positions=np.array(flat_positions, dtype=np.float32), - atom_elements=np.array(flat_elements, dtype=object), - token_to_atoms=np.array(token_to_atoms, dtype=np.int32), - chain_id=np.array(chain_ids_per_token, dtype=np.int64), - plddt=np.array(confidence_scores, dtype=np.float32), - atom_names=np.array(flat_names, dtype=object), - atom_hetero=np.array(flat_hetero, dtype=bool), - metadata=MolecularComplexMetadata( - entity_lookup=entity_info, - chain_lookup=chain_lookup, - assembly_composition=None, - ), - ) diff --git a/fastplms/esmfold2/esmfold2_paired_msa.py b/fastplms/esmfold2/esmfold2_paired_msa.py deleted file mode 100644 index 2356746..0000000 --- a/fastplms/esmfold2/esmfold2_paired_msa.py +++ /dev/null @@ -1,245 +0,0 @@ -"""Taxonomy-paired MSA construction for ESMFold2 inference. - -Taxonomy IDs are read from FASTA headers as ``key=N`` tokens. Rows -where any chain has ``key=-1`` (or no ``key=`` at all) are treated as -unpaired and assigned to that chain's block-diagonal section after -the paired rows. -""" - -import re - -import numpy as np - -from .esmfold2_constants import ( - MSA_GAP_TOKEN_ID, - PROTEIN_3TO1, - PROTEIN_RESIDUE_TO_RES_TYPE, - PROTEIN_UNK_RES_TYPE, -) -from .esmfold2_msa import MSA - -_KEY_RE = re.compile(r"key=(-?\d+)") - - -def protein_letter_to_res_type() -> dict[str, int]: - """Return the protein 1-letter → res_type mapping used by the MSA encoder.""" - mapping: dict[str, int] = {} - for three, one in PROTEIN_3TO1.items(): - if three in PROTEIN_RESIDUE_TO_RES_TYPE: - mapping[one] = PROTEIN_RESIDUE_TO_RES_TYPE[three] - mapping["-"] = MSA_GAP_TOKEN_ID - mapping["X"] = PROTEIN_UNK_RES_TYPE - return mapping - - -def _taxonomy_from_header(header: str) -> int: - if not header: - return -1 - m = _KEY_RE.search(header) - return int(m.group(1)) if m else -1 - - -def msa_to_res_type_and_deletions( - msa: MSA, letter_to_res_type: dict[str, int] -) -> tuple[np.ndarray, np.ndarray]: - """Convert an :class:`MSA` to ``(res_type[M, L], deletion_count[M, L])``. - - Handles a3m insertion convention: lowercase letters and ``.`` are - insertions and are not emitted; their count is accumulated into the - next non-insertion position's deletion value. ``L`` is the query - length after stripping insertions from row 0. - """ - query = msa.entries[0].sequence - L = sum(1 for ch in query if not (ch.islower() or ch == ".")) - M = msa.depth - - res_type = np.full((M, L), MSA_GAP_TOKEN_ID, dtype=np.int64) - deletions = np.zeros((M, L), dtype=np.float32) - - for r, entry in enumerate(msa.entries): - col = 0 - ins = 0 - for ch in entry.sequence: - if ch == "." or (ch.islower() and ch != "-"): - ins += 1 - continue - if col >= L: - break - if ch == "-": - res_type[r, col] = MSA_GAP_TOKEN_ID - else: - res_type[r, col] = letter_to_res_type.get( - ch.upper(), PROTEIN_UNK_RES_TYPE - ) - if ins > 0: - deletions[r, col] = float(ins) - ins = 0 - col += 1 - return res_type, deletions - - -def _dummy_msa_residues(query_res_types: np.ndarray) -> np.ndarray: - """Single-row 'MSA' for chains without one — just the query.""" - return query_res_types[None, :] # [1, L] - - -def construct_paired_msa( - chain_msas: dict[int, MSA | None], - chain_query_res_types: dict[int, np.ndarray], - token_asym_ids: np.ndarray, - token_res_ids: np.ndarray, - letter_to_res_type: dict[str, int] | None = None, - *, - max_pairs: int = 8192, - max_total: int = 16384, - max_seqs: int = 16384, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Build paired MSA features. - - Parameters - ---------- - chain_msas - ``asym_id -> MSA`` (or ``None`` for chains without an MSA). - chain_query_res_types - ``asym_id -> np.ndarray[L_c]`` of res-type ids for the chain's - query. Used to build dummy MSAs when a chain has no MSA. - token_asym_ids - Per-token asym_id, length ``T``. Must be non-decreasing. - token_res_ids - Per-token residue index within chain, length ``T``. - letter_to_res_type - 1-letter → res-type mapping. Defaults to - :func:`protein_letter_to_res_type`. - - Returns - ------- - msa_residues : ``np.ndarray[M, T]`` int64 - deletion_value : ``np.ndarray[M, T]`` float32 (raw deletion counts; the - ``arctan(/3) * pi/2`` transform is applied by the caller) - is_paired : ``np.ndarray[M, T]`` float32 broadcast of per-row, - per-chain paired flags. - """ - if letter_to_res_type is None: - letter_to_res_type = protein_letter_to_res_type() - - chain_ids: list[int] = sorted(chain_msas.keys()) - - # Build per-chain (res_type, deletions, taxonomy) tables. - chain_res_type: dict[int, np.ndarray] = {} - chain_deletions: dict[int, np.ndarray] = {} - chain_taxonomies: dict[int, list[int]] = {} - for c in chain_ids: - m = chain_msas.get(c) - if m is None or m.depth == 0: - qres = chain_query_res_types[c] - chain_res_type[c] = _dummy_msa_residues(qres) - chain_deletions[c] = np.zeros((1, qres.shape[0]), dtype=np.float32) - chain_taxonomies[c] = [-1] - continue - rt, dl = msa_to_res_type_and_deletions(m, letter_to_res_type) - chain_res_type[c] = rt - chain_deletions[c] = dl - chain_taxonomies[c] = [_taxonomy_from_header(e.header) for e in m.entries] - - # Group by taxonomy, skip query row and unpaired (-1) entries. - taxonomy_map: dict[int, list[tuple[int, int]]] = {} - for c in chain_ids: - for seq_idx, taxon in enumerate(chain_taxonomies[c]): - if seq_idx == 0 or taxon == -1: - continue - taxonomy_map.setdefault(taxon, []).append((c, seq_idx)) - taxonomy_map = {k: v for k, v in taxonomy_map.items() if len(v) > 1} - # Order taxonomies by number of distinct chains, descending. - sorted_taxa = sorted( - taxonomy_map.items(), key=lambda kv: len({c for c, _ in kv[1]}), reverse=True - ) - - visited = {s for _, items in taxonomy_map.items() for s in items} - available: dict[int, list[int]] = { - c: [i for i in range(1, len(chain_taxonomies[c])) if (c, i) not in visited] - for c in chain_ids - } - - pairing: list[dict[int, int]] = [{c: 0 for c in chain_ids}] - is_paired: list[dict[int, int]] = [{c: 1 for c in chain_ids}] - - for _, pairs in sorted_taxa: - per_chain: dict[int, list[int]] = {} - for c, seq_idx in pairs: - per_chain.setdefault(c, []).append(seq_idx) - max_occ = max(len(v) for v in per_chain.values()) - for i in range(max_occ): - row_pairing: dict[int, int] = {} - row_is_paired: dict[int, int] = {} - for c, seq_idxs in per_chain.items(): - row_pairing[c] = seq_idxs[i % len(seq_idxs)] - row_is_paired[c] = 1 - for c in chain_ids: - if c in row_pairing: - continue - row_is_paired[c] = 0 - if available[c]: - row_pairing[c] = available[c].pop(0) - else: - row_pairing[c] = -1 - pairing.append(row_pairing) - is_paired.append(row_is_paired) - if len(pairing) >= max_pairs: - break - if len(pairing) >= max_pairs: - break - - max_left = max((len(v) for v in available.values()), default=0) - for _ in range(min(max_total - len(pairing), max_left)): - row_pairing = {} - row_is_paired = {} - for c in chain_ids: - row_is_paired[c] = 0 - if available[c]: - row_pairing[c] = available[c].pop(0) - else: - row_pairing[c] = -1 - pairing.append(row_pairing) - is_paired.append(row_is_paired) - if len(pairing) >= max_total: - break - - pairing = pairing[:max_seqs] - is_paired = is_paired[:max_seqs] - M = len(pairing) - T = len(token_asym_ids) - - msa_residues = np.full((M, T), MSA_GAP_TOKEN_ID, dtype=np.int64) - deletion_value = np.zeros((M, T), dtype=np.float32) - paired_mask = np.zeros((M, T), dtype=np.float32) - - # Vectorize per chain: gather chain rows according to pairing[c], then - # index into them by the chain's token residue ids. - for c in chain_ids: - rt = chain_res_type[c] - dl = chain_deletions[c] - Lc = rt.shape[1] - chain_pairing = np.array([row[c] for row in pairing], dtype=np.int64) - chain_paired = np.array([row[c] for row in is_paired], dtype=np.float32) - - token_mask = token_asym_ids == c - if not token_mask.any(): - continue - token_res_in_chain = token_res_ids[token_mask] - # Clamp residue indices to the MSA's column range. Modified-residue - # tokens that exceed the query length fall back to the last column. - cols = np.minimum(token_res_in_chain, Lc - 1) - - # Rows where pairing == -1 fall back to gap (already initialized). - valid_rows = chain_pairing >= 0 - if valid_rows.any(): - gathered_rt = rt[chain_pairing[valid_rows]][:, cols] - gathered_dl = dl[chain_pairing[valid_rows]][:, cols] - valid_idx = np.where(valid_rows)[0] - token_idx = np.where(token_mask)[0] - msa_residues[np.ix_(valid_idx, token_idx)] = gathered_rt - deletion_value[np.ix_(valid_idx, token_idx)] = gathered_dl - - paired_mask[:, token_mask] = chain_paired[:, None] - - return msa_residues, deletion_value, paired_mask diff --git a/fastplms/esmfold2/esmfold2_parsing.py b/fastplms/esmfold2/esmfold2_parsing.py deleted file mode 100644 index 3137af5..0000000 --- a/fastplms/esmfold2/esmfold2_parsing.py +++ /dev/null @@ -1,112 +0,0 @@ -import io -from pathlib import Path -from typing import Generator, Iterable, NamedTuple - -PathOrBuffer = str | Path | io.TextIOBase -FastaEntry = NamedTuple("FastaEntry", [("header", str), ("sequence", str)]) - - -def parse_fasta(fasta_string: str) -> Generator[FastaEntry, None, None]: - """ - Parses a fasta file and yields FastaEntry objects - - Args: - fasta_string: The fasta file as a string - Returns: - A generator of FastaEntry objects - """ - header = None - seq = [] - num_sequences = 0 - for line in fasta_string.splitlines(): - if not line or line[0] == "#": - continue - if line.startswith(">"): - if header is not None: - yield FastaEntry(header, "".join(seq)) - seq = [] - header = line[1:].strip() - else: - seq.append(line) - if header is not None: - num_sequences += 1 - yield FastaEntry(header, "".join(seq)) - - if num_sequences == 0: - raise ValueError("Found no sequences in input") - - -def read_sequences(path: PathOrBuffer) -> Generator[FastaEntry, None, None]: - # Uses duck typing to try and call the right method - # Doesn't use explicit isinstance check to support - # inputs that are not explicitly str/Path/TextIOBase but - # may support similar functionality - data = None # type: ignore - try: - if str(path).endswith(".gz"): - import gzip - - data = gzip.open(path, "rt") # type: ignore - else: - try: - data = open(path) # type: ignore - except TypeError: - data: io.TextIOBase = path # type: ignore - - yield from parse_fasta(data.read()) - finally: - if data is not None: - data.close() - - -def read_first_sequence(path: PathOrBuffer) -> FastaEntry: - return next(iter(read_sequences(path))) - - -def count_fasta_sequences(path: str | Path) -> int: - """Count sequences in a FASTA file by counting header lines. - - Faster than parsing the full file — only scans for '>' prefixes. - Returns 0 if the file does not exist. - """ - path = Path(path) - if not path.exists(): - return 0 - with open(path) as f: - return sum(1 for line in f if line.startswith(">")) - - -def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None: - """Append a single sequence to a FASTA file (creating it if needed).""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - # The existing file may not end with a newline (e.g., write_sequences() - # explicitly avoids writing a newline at the end), so we insert one before - # appending to avoid merging with the last line. - needs_newline = ( - path.exists() and path.stat().st_size > 0 and path.read_bytes()[-1:] != b"\n" - ) - with open(path, "a") as f: - if needs_newline: - f.write("\n") - f.write(f">{header}\n{sequence}\n") - - -def write_sequences(sequences: Iterable[tuple[str, str]], path: PathOrBuffer) -> None: - needs_closing = False - handle = None - try: - try: - handle = open(path, "w") # type: ignore - needs_closing = True - except TypeError: - handle = path - has_prev = False - for header, seq in sequences: - if has_prev: - handle.write("\n") # type: ignore - handle.write(f">{header}\n{seq}") # type: ignore - has_prev = True - finally: - if needs_closing: - handle.close() # type: ignore diff --git a/fastplms/esmfold2/esmfold2_predicted_aligned_error.py b/fastplms/esmfold2/esmfold2_predicted_aligned_error.py deleted file mode 100644 index 821df9c..0000000 --- a/fastplms/esmfold2/esmfold2_predicted_aligned_error.py +++ /dev/null @@ -1,104 +0,0 @@ -import torch -import torch.nn.functional as F - -from .esmfold2_affine3d import Affine3D - - -def masked_mean( - mask: torch.Tensor, - value: torch.Tensor, - dim: int | None | tuple[int, ...] = None, - eps=1e-10, -) -> torch.Tensor: - """Compute the mean of `value` where only positions where `mask == true` are - counted. - """ - mask = mask.expand(*value.shape) - return torch.sum(mask * value, dim=dim) / (eps + torch.sum(mask, dim=dim)) - - -def _pae_bins( - max_bin: float = 31, num_bins: int = 64, device: torch.device = torch.device("cpu") -): - bins = torch.linspace(0, max_bin, steps=(num_bins - 1), device=device) - step = max_bin / (num_bins - 2) - bin_centers = bins + step / 2 - bin_centers = torch.cat( - [bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0 - ) - return bin_centers - - -def _compute_pae_masks(mask: torch.Tensor): - square_mask = (mask.unsqueeze(-1) * mask.unsqueeze(-2)).bool() - return square_mask - - -def compute_predicted_aligned_error( - logits: torch.Tensor, - aa_mask: torch.Tensor, - sequence_id: torch.Tensor | None = None, - max_bin: float = 31, -) -> torch.Tensor: - bins = _pae_bins(max_bin, logits.shape[-1], logits.device) - square_mask = _compute_pae_masks(aa_mask) - min_v = torch.finfo(logits.dtype).min - probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1) - - return (probs * bins).sum(dim=-1) - - -@torch.no_grad -def compute_tm(logits: torch.Tensor, aa_mask: torch.Tensor, max_bin: float = 31.0): - square_mask = _compute_pae_masks(aa_mask) - seqlens = aa_mask.sum(-1, keepdim=True) - bins = _pae_bins(max_bin, logits.shape[-1], logits.device) - d0 = 1.24 * (seqlens.clamp_min(19) - 15) ** (1 / 3) - 1.8 - f_d = 1.0 / (1 + (bins / d0.unsqueeze(-1)) ** 2) - - min_v = torch.finfo(logits.dtype).min - probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1) - # This is the sum over bins - ptm = (probs * f_d.unsqueeze(-2)).sum(dim=-1) - # This is the mean over residues j - ptm = masked_mean(square_mask, ptm, dim=-1) - # The we do a max over residues i - return ptm.max(dim=-1).values - - -def tm_loss( - logits: torch.Tensor, - pred_affine: torch.Tensor, - targ_affine: torch.Tensor, - targ_mask: torch.Tensor, - tm_mask: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, - max_bin: float = 31, -): - pred = Affine3D.from_tensor(pred_affine) - targ = Affine3D.from_tensor(targ_affine) - - def transform(affine: Affine3D): - pts = affine.trans[..., None, :, :] - return affine.invert()[..., None].apply(pts) - - with torch.no_grad(): - sq_diff = (transform(pred) - transform(targ)).square().sum(dim=-1) - - num_bins = logits.shape[-1] - sq_bins = torch.linspace( - 0, max_bin, num_bins - 1, device=logits.device - ).square() - # Gets the bin id by using a sum. - true_bins = (sq_diff[..., None] > sq_bins).sum(dim=-1).long() - - errors = F.cross_entropy(logits.movedim(3, 1), true_bins, reduction="none") - square_mask = _compute_pae_masks(targ_mask) - loss = masked_mean(square_mask, errors, dim=(-1, -2)) - - if tm_mask is not None: - loss = masked_mean(tm_mask, loss, dim=None) - else: - loss = loss.mean() - - return loss diff --git a/fastplms/esmfold2/esmfold2_prepare_input.py b/fastplms/esmfold2/esmfold2_prepare_input.py deleted file mode 100644 index 5763707..0000000 --- a/fastplms/esmfold2/esmfold2_prepare_input.py +++ /dev/null @@ -1,1463 +0,0 @@ -"""Prepare ESMFold2 model inputs from sequence-level StructurePredictionInput. - -This module converts StructurePredictionInput (protein/DNA/RNA/ligand sequences) -into the tensor dict expected by the ESMFold2 model forward pass. -""" - -from __future__ import annotations - -import math -import warnings -from collections import defaultdict -from dataclasses import dataclass, field - -import numpy as np -import torch - -from .esmfold2_conformers import ( - get_ccd_leaving_atoms, - get_idealized_atom_pos, - get_ligand_ccd_atoms_with_charges, - get_ligand_ccd_bonds, - get_ligand_idealized_atom_pos, -) -from .esmfold2_constants import ( - CHARGED_ATOMS, - DNA_1TO3, - DNA_BACKBONE_ATOMS, - DNA_HEAVY_ATOMS, - DNA_RESIDUE_TO_RES_TYPE, - DNA_RNA_LIGAND_INPUT_ID, - DNA_UNK_RES_TYPE, - ELEMENT_TO_ATOMIC_NUM, - ESM_PROTEIN_VOCAB, - MOL_TYPE_DNA, - MOL_TYPE_NONPOLYMER, - MOL_TYPE_PROTEIN, - MOL_TYPE_RNA, - MSA_GAP_TOKEN_ID, - PROTEIN_1TO3, - PROTEIN_3TO1, - PROTEIN_HEAVY_ATOMS, - PROTEIN_RESIDUE_TO_RES_TYPE, - PROTEIN_UNK_RES_TYPE, - RNA_1TO3, - RNA_BACKBONE_ATOMS, - RNA_HEAVY_ATOMS, - RNA_RESIDUE_TO_RES_TYPE, - RNA_UNK_RES_TYPE, -) -from .esmfold2_types import ( - MSA, - DNAInput, - LigandInput, - Modification, - ProteinInput, - RNAInput, - StructurePredictionInput, -) - -# ============================================================================= -# Lightweight data model -# ============================================================================= - -_ZERO_POS = np.array([0.0, 0.0, 0.0], dtype=np.float32) - - -@dataclass -class AtomInfo: - name: str - element: str - charge: int - ref_pos: np.ndarray # Idealized position from CCD [3] - pos: np.ndarray # Experimental position [3] (zeros for inference) - token_index: int = -1 - atom_index: int = -1 - space_uid: int = -1 - is_valid: bool = True - - -@dataclass -class TokenInfo: - token_index: int - residue_index: int # Within chain (0-based) - residue_name: str # 3-letter code - mol_type: int # 0=protein, 1=DNA, 2=RNA, 3=nonpolymer - res_type: int # Residue type index (2-32) - input_id: int # ESM vocab ID - asym_id: int - sym_id: int - entity_id: int - atom_start: int # Index into atoms list - atom_count: int - - -@dataclass -class ChainInfo: - chain_id: str - asym_id: int - entity_id: int - sym_id: int - mol_type: int - tokens: list[TokenInfo] = field(default_factory=list) - - -# ============================================================================= -# Helper functions -# ============================================================================= - -# Caches for hot-path functions -_ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {} -_ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {} - - -def encode_atom_name(name: str) -> list[int]: - """Encode atom name as 4 character indices (offset by 32 from ASCII).""" - if name in _ENCODE_ATOM_NAME_CACHE: - return _ENCODE_ATOM_NAME_CACHE[name] - padded = name.ljust(4)[:4] - result = [ord(c) - 32 if c != " " else 0 for c in padded] - _ENCODE_ATOM_NAME_CACHE[name] = result - return result - - -def get_element_atomic_num(element: str) -> int: - """Get atomic number for an element symbol.""" - if element in _ELEMENT_ATOMIC_NUM_CACHE: - return _ELEMENT_ATOMIC_NUM_CACHE[element] - result = ELEMENT_TO_ATOMIC_NUM.get(element.upper(), 0) - _ELEMENT_ATOMIC_NUM_CACHE[element] = result - return result - - -def _infer_element(atom_name: str) -> str: - """Infer element from atom name.""" - name = atom_name.strip() - if not name: - return "C" - if name[0].isdigit(): - return name[1] if len(name) > 1 else "H" - if len(name) == 2 and name in ( - "FE", - "ZN", - "MG", - "MN", - "CO", - "NI", - "CU", - "SE", - "BR", - ): - return name - return name[0] - - -def _compute_res_type(name: str, mol_type: int) -> int: - """Compute residue type index from residue name and mol_type.""" - if mol_type == MOL_TYPE_PROTEIN: - return PROTEIN_RESIDUE_TO_RES_TYPE.get(name, PROTEIN_UNK_RES_TYPE) - elif mol_type == MOL_TYPE_DNA: - if name in DNA_RESIDUE_TO_RES_TYPE: - return DNA_RESIDUE_TO_RES_TYPE[name] - if name in RNA_RESIDUE_TO_RES_TYPE: - return RNA_RESIDUE_TO_RES_TYPE[name] - return DNA_UNK_RES_TYPE - elif mol_type == MOL_TYPE_RNA: - if name in RNA_RESIDUE_TO_RES_TYPE: - return RNA_RESIDUE_TO_RES_TYPE[name] - if name in DNA_RESIDUE_TO_RES_TYPE: - return DNA_RESIDUE_TO_RES_TYPE[name] - return RNA_UNK_RES_TYPE - return PROTEIN_UNK_RES_TYPE - - -def _compute_esm_input_id(name: str, mol_type: int) -> int: - """Compute ESM vocabulary input ID.""" - if mol_type == MOL_TYPE_PROTEIN: - letter = PROTEIN_3TO1.get(name) - if letter is None: - return DNA_RNA_LIGAND_INPUT_ID - return ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"]) - return DNA_RNA_LIGAND_INPUT_ID - - -# ============================================================================= -# Tokenization functions — build tokens and atoms from sequences -# ============================================================================= - - -def tokenize_protein( - sequence: str, - modifications: list[Modification] | None, - entity_id: int, - asym_id: int, - sym_id: int, - token_offset: int, - atom_offset: int, - space_uid_offset: int, -) -> tuple[list[TokenInfo], list[AtomInfo]]: - """Tokenize a protein sequence into tokens and atoms. - - Standard residues produce 1 token with all heavy atoms. - Modified residues (from modifications) are atom-tokenized (1 token per atom). - """ - tokens: list[TokenInfo] = [] - atoms: list[AtomInfo] = [] - - # Build 3-letter sequence, applying modifications - seq_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence] - modified_positions: set[int] = set() - if modifications: - for mod in modifications: - seq_3letter[mod.position] = mod.ccd - modified_positions.add(mod.position) - - token_idx = token_offset - atom_idx = atom_offset - space_uid = space_uid_offset - - for res_idx, res_name in enumerate(seq_3letter): - # MSE → MET for atom lookup - res_corrected = "MET" if res_name == "MSE" else res_name - is_modified = res_idx in modified_positions - - # Check if standard residue (has predefined atom list) - if not is_modified and res_corrected in PROTEIN_HEAVY_ATOMS: - # Standard residue: 1 token, multiple atoms - atom_names = PROTEIN_HEAVY_ATOMS[res_corrected] - res_type = _compute_res_type(res_corrected, MOL_TYPE_PROTEIN) - input_id = _compute_esm_input_id(res_corrected, MOL_TYPE_PROTEIN) - - atom_start = atom_idx - for a_name in atom_names: - ref_pos = get_idealized_atom_pos(res_type, a_name) - atoms.append( - AtomInfo( - name=a_name, - element=_infer_element(a_name), - charge=CHARGED_ATOMS.get((res_corrected, a_name), 0), - ref_pos=ref_pos.copy() - if ref_pos is not None - else _ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - atom_idx += 1 - - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=res_corrected, - mol_type=MOL_TYPE_PROTEIN, - res_type=res_type, - input_id=input_id, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_start, - atom_count=len(atom_names), - ) - ) - token_idx += 1 - space_uid += 1 - - else: - # Modified or unknown residue: atom-tokenized - ccd_atoms = get_ligand_ccd_atoms_with_charges(res_name) - if ccd_atoms is None: - # Fallback: backbone only - ccd_atoms = [ - (_infer_element(n), _infer_element(n), 0) - for n in ["N", "CA", "C", "O"] - ] - - # Filter leaving atoms if not terminal - is_terminal = res_idx == len(seq_3letter) - 1 - leaving_atoms = set() if is_terminal else get_ccd_leaving_atoms(res_name) - kept_atoms = [a for a in ccd_atoms if a[0] not in leaving_atoms] - # Single-atom residues (e.g. NH2 cap): the local frame is - # ill-defined with one atom; place at origin. - single_atom_residue = len(kept_atoms) == 1 - - for a_name, a_element, a_charge in kept_atoms: - ref_pos = get_ligand_idealized_atom_pos(res_name, a_name) - atoms.append( - AtomInfo( - name=a_name, - element=a_element, - charge=a_charge, - ref_pos=_ZERO_POS.copy() - if single_atom_residue - else ( - ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy() - ), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=res_name, - mol_type=MOL_TYPE_PROTEIN, - res_type=PROTEIN_UNK_RES_TYPE, - input_id=DNA_RNA_LIGAND_INPUT_ID, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_idx, - atom_count=1, - ) - ) - token_idx += 1 - atom_idx += 1 - - space_uid += 1 - - return tokens, atoms - - -def tokenize_nucleotide( - sequence: str, - modifications: list[Modification] | None, - mol_type: int, - entity_id: int, - asym_id: int, - sym_id: int, - token_offset: int, - atom_offset: int, - space_uid_offset: int, -) -> tuple[list[TokenInfo], list[AtomInfo]]: - """Tokenize a DNA or RNA sequence into tokens and atoms.""" - tokens: list[TokenInfo] = [] - atoms: list[AtomInfo] = [] - - letter_to_3 = DNA_1TO3 if mol_type == MOL_TYPE_DNA else RNA_1TO3 - heavy_atoms = DNA_HEAVY_ATOMS if mol_type == MOL_TYPE_DNA else RNA_HEAVY_ATOMS - backbone_atoms = ( - DNA_BACKBONE_ATOMS if mol_type == MOL_TYPE_DNA else RNA_BACKBONE_ATOMS - ) - unk_res_type = DNA_UNK_RES_TYPE if mol_type == MOL_TYPE_DNA else RNA_UNK_RES_TYPE - - seq_3letter = [letter_to_3.get(c, "UNK") for c in sequence] - modified_positions: set[int] = set() - if modifications: - for mod in modifications: - seq_3letter[mod.position] = mod.ccd - modified_positions.add(mod.position) - - token_idx = token_offset - atom_idx = atom_offset - space_uid = space_uid_offset - - for res_idx, res_name in enumerate(seq_3letter): - is_modified = res_idx in modified_positions - - if not is_modified and res_name in heavy_atoms: - # Standard nucleotide - atom_names = heavy_atoms[res_name] - res_type = _compute_res_type(res_name, mol_type) - input_id = DNA_RNA_LIGAND_INPUT_ID - - atom_start = atom_idx - for a_name in atom_names: - ref_pos = get_idealized_atom_pos(res_type, a_name) - atoms.append( - AtomInfo( - name=a_name, - element=_infer_element(a_name), - charge=CHARGED_ATOMS.get((res_name, a_name), 0), - ref_pos=ref_pos.copy() - if ref_pos is not None - else _ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - atom_idx += 1 - - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=res_name, - mol_type=mol_type, - res_type=res_type, - input_id=input_id, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_start, - atom_count=len(atom_names), - ) - ) - token_idx += 1 - space_uid += 1 - - elif not is_modified and res_name == "UNK": - # Unknown nucleotide: backbone only - atom_names = backbone_atoms - atom_start = atom_idx - for a_name in atom_names: - ref_pos = None # No idealized positions for UNK - atoms.append( - AtomInfo( - name=a_name, - element=_infer_element(a_name), - charge=0, - ref_pos=_ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - atom_idx += 1 - - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=res_name, - mol_type=mol_type, - res_type=unk_res_type, - input_id=DNA_RNA_LIGAND_INPUT_ID, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_start, - atom_count=len(atom_names), - ) - ) - token_idx += 1 - space_uid += 1 - - else: - # Modified nucleotide: atom-tokenized - ccd_atoms = get_ligand_ccd_atoms_with_charges(res_name) - if ccd_atoms is None: - ccd_atoms = [ - (_infer_element(n), _infer_element(n), 0) for n in backbone_atoms - ] - - is_terminal = res_idx == len(seq_3letter) - 1 - leaving_atoms = set() if is_terminal else get_ccd_leaving_atoms(res_name) - - for a_name, a_element, a_charge in ccd_atoms: - if a_name in leaving_atoms: - continue - ref_pos = get_ligand_idealized_atom_pos(res_name, a_name) - atoms.append( - AtomInfo( - name=a_name, - element=a_element, - charge=a_charge, - ref_pos=ref_pos.copy() - if ref_pos is not None - else _ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=res_name, - mol_type=mol_type, - res_type=PROTEIN_UNK_RES_TYPE, - input_id=DNA_RNA_LIGAND_INPUT_ID, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_idx, - atom_count=1, - ) - ) - token_idx += 1 - atom_idx += 1 - - space_uid += 1 - - return tokens, atoms - - -def tokenize_ligand_ccd( - ccd_codes: list[str], - entity_id: int, - asym_id: int, - sym_id: int, - token_offset: int, - atom_offset: int, - space_uid_offset: int, - has_covalent_bond: bool, -) -> tuple[list[TokenInfo], list[AtomInfo]]: - """Tokenize a ligand from CCD codes (1 token per atom).""" - tokens: list[TokenInfo] = [] - atoms: list[AtomInfo] = [] - - token_idx = token_offset - atom_idx = atom_offset - space_uid = space_uid_offset - - for res_idx, code in enumerate(ccd_codes): - ccd_atoms = get_ligand_ccd_atoms_with_charges(code) - if ccd_atoms is None: - raise ValueError(f"CCD component {code} not found") - - leaving_atoms = get_ccd_leaving_atoms(code) if has_covalent_bond else set() - - for a_name, a_element, a_charge in ccd_atoms: - if a_name in leaving_atoms: - continue - ref_pos = get_ligand_idealized_atom_pos(code, a_name) - atoms.append( - AtomInfo( - name=a_name, - element=a_element, - charge=a_charge, - ref_pos=ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=res_idx, - residue_name=code, - mol_type=MOL_TYPE_NONPOLYMER, - res_type=PROTEIN_UNK_RES_TYPE, - input_id=DNA_RNA_LIGAND_INPUT_ID, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_idx, - atom_count=1, - ) - ) - token_idx += 1 - atom_idx += 1 - - space_uid += 1 - - return tokens, atoms - - -def tokenize_ligand_smiles( - smiles: str, - entity_id: int, - asym_id: int, - sym_id: int, - token_offset: int, - atom_offset: int, - space_uid_offset: int, - seed: int | None = None, -) -> tuple[list[TokenInfo], list[AtomInfo]]: - """Tokenize a ligand from SMILES (1 token per heavy atom).""" - from rdkit import Chem - from rdkit.Chem import AllChem - - mol = Chem.MolFromSmiles(smiles) - if mol is None: - raise ValueError(f"Failed to parse SMILES: {smiles}") - mol = Chem.AddHs(mol) - - # Assign atom names using canonical ranking - canonical_order = AllChem.CanonicalRankAtoms(mol) # type: ignore[attr-defined] - for atom, can_idx in zip(mol.GetAtoms(), canonical_order): - atom_name = atom.GetSymbol().upper() + str(can_idx + 1) - if len(atom_name) > 4: - raise ValueError( - f"SMILES {smiles} has atom name longer than 4 chars: {atom_name}" - ) - atom.SetProp("name", atom_name) - - # Generate 3D conformer - options = AllChem.ETKDGv3() # type: ignore[attr-defined] - options.clearConfs = False - if seed is not None: - options.randomSeed = seed - conf_id = AllChem.EmbedMolecule(mol, options) # type: ignore[attr-defined] - if conf_id == -1: - options.useRandomCoords = True - conf_id = AllChem.EmbedMolecule(mol, options) # type: ignore[attr-defined] - if conf_id != -1: - try: - AllChem.UFFOptimizeMolecule(mol, confId=conf_id, maxIters=1000) # type: ignore[attr-defined] - except (RuntimeError, ValueError): - pass - - # Remove hydrogens - mol_no_h = Chem.RemoveHs(mol) - if mol_no_h.GetNumConformers() == 0: - raise ValueError(f"Failed to generate conformer for SMILES: {smiles}") - - conformer = mol_no_h.GetConformer(0) - - tokens: list[TokenInfo] = [] - atoms_list: list[AtomInfo] = [] - token_idx = token_offset - atom_idx = atom_offset - space_uid = space_uid_offset - - for atom in mol_no_h.GetAtoms(): - a_name = atom.GetProp("name") - a_element = atom.GetSymbol() - a_charge = atom.GetFormalCharge() - pos_3d = conformer.GetAtomPosition(atom.GetIdx()) - ref_pos = np.array([pos_3d.x, pos_3d.y, pos_3d.z], dtype=np.float32) - - atoms_list.append( - AtomInfo( - name=a_name, - element=a_element, - charge=a_charge, - ref_pos=ref_pos, - pos=_ZERO_POS.copy(), - token_index=token_idx, - atom_index=atom_idx, - space_uid=space_uid, - ) - ) - tokens.append( - TokenInfo( - token_index=token_idx, - residue_index=0, - residue_name="LIG", - mol_type=MOL_TYPE_NONPOLYMER, - res_type=PROTEIN_UNK_RES_TYPE, - input_id=DNA_RNA_LIGAND_INPUT_ID, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - atom_start=atom_idx, - atom_count=1, - ) - ) - token_idx += 1 - atom_idx += 1 - - return tokens, atoms_list - - -# ============================================================================= -# Build chains from StructurePredictionInput -# ============================================================================= - - -def _get_sequence_key(item) -> str: - """Get a hashable key for entity deduplication.""" - if isinstance(item, ProteinInput): - return f"PROTEIN:{item.sequence}" - elif isinstance(item, DNAInput): - return f"DNA:{item.sequence}" - elif isinstance(item, RNAInput): - return f"RNA:{item.sequence}" - elif isinstance(item, LigandInput): - if item.ccd: - return f"LIGAND_CCD:{','.join(item.ccd)}" - return f"LIGAND_SMILES:{item.smiles}" - raise ValueError(f"Unknown input type: {type(item)}") - - -def build_chains_from_input( - input: StructurePredictionInput, seed: int | None = None -) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]: - """Build chains, tokens, and atoms from StructurePredictionInput. - - Handles entity deduplication (identical sequences get same entity_id), - sym_id assignment, and delegates to type-specific tokenization functions. - """ - chains: list[ChainInfo] = [] - all_tokens: list[TokenInfo] = [] - all_atoms: list[AtomInfo] = [] - - # Entity deduplication - sequence_to_entity: dict[str, int] = {} - entity_sym_count: dict[int, int] = {} - next_entity_id = 0 - - # Gather chain IDs involved in covalent bonds - covalent_chain_ids: set[str] = set() - if input.covalent_bonds: - for cb in input.covalent_bonds: - covalent_chain_ids.update([cb.chain_id1, cb.chain_id2]) - - token_offset = 0 - atom_offset = 0 - space_uid_offset = 0 - asym_id = 0 - - for item in input.sequences: - # Entity deduplication - seq_key = _get_sequence_key(item) - if seq_key in sequence_to_entity: - entity_id = sequence_to_entity[seq_key] - else: - entity_id = next_entity_id - sequence_to_entity[seq_key] = entity_id - next_entity_id += 1 - - # Get all chain IDs for this item - ids = [item.id] if isinstance(item.id, str) else item.id - - for chain_id_str in ids: - # sym_id is the per-entity copy index; increment per chain so - # ProteinInput(id=['A','B']) gives chain A sym_id=0, chain B sym_id=1. - sym_id = entity_sym_count.get(entity_id, 0) - entity_sym_count[entity_id] = sym_id + 1 - if isinstance(item, ProteinInput): - if item.msa is None: - warnings.warn( - f"No MSA provided for {item.id}, using single sequence mode" - ) - - new_tokens, new_atoms = tokenize_protein( - sequence=item.sequence, - modifications=item.modifications, - entity_id=entity_id, - asym_id=asym_id, - sym_id=sym_id, - token_offset=token_offset, - atom_offset=atom_offset, - space_uid_offset=space_uid_offset, - ) - - elif isinstance(item, (DNAInput, RNAInput)): - mol_type = MOL_TYPE_DNA if isinstance(item, DNAInput) else MOL_TYPE_RNA - new_tokens, new_atoms = tokenize_nucleotide( - sequence=item.sequence, - modifications=item.modifications, - mol_type=mol_type, - entity_id=entity_id, - asym_id=asym_id, - sym_id=sym_id, - token_offset=token_offset, - atom_offset=atom_offset, - space_uid_offset=space_uid_offset, - ) - - elif isinstance(item, LigandInput): - has_cov = chain_id_str in covalent_chain_ids - if item.ccd is not None: - if item.smiles is not None: - warnings.warn("Both ccd and smiles provided, using ccd") - new_tokens, new_atoms = tokenize_ligand_ccd( - ccd_codes=item.ccd, - entity_id=entity_id, - asym_id=asym_id, - sym_id=sym_id, - token_offset=token_offset, - atom_offset=atom_offset, - space_uid_offset=space_uid_offset, - has_covalent_bond=has_cov, - ) - elif item.smiles is not None: - new_tokens, new_atoms = tokenize_ligand_smiles( - smiles=item.smiles, - entity_id=entity_id, - asym_id=asym_id, - sym_id=sym_id, - token_offset=token_offset, - atom_offset=atom_offset, - space_uid_offset=space_uid_offset, - seed=seed, - ) - else: - raise ValueError("LigandInput must have either ccd or smiles") - else: - raise ValueError(f"Unknown input type: {type(item)}") - - chain = ChainInfo( - chain_id=chain_id_str, - asym_id=asym_id, - entity_id=entity_id, - sym_id=sym_id, - mol_type=new_tokens[0].mol_type if new_tokens else MOL_TYPE_PROTEIN, - tokens=new_tokens, - ) - chains.append(chain) - all_tokens.extend(new_tokens) - all_atoms.extend(new_atoms) - - token_offset += len(new_tokens) - atom_offset += len(new_atoms) - space_uid_offset += len(set(a.space_uid for a in new_atoms)) - asym_id += 1 - - return chains, all_tokens, all_atoms - - -# ============================================================================= -# Feature tensor building -# ============================================================================= - - -def compute_frame_indices( - tokens: list[TokenInfo], atoms: list[AtomInfo] -) -> tuple[np.ndarray, np.ndarray]: - """Compute backbone frame indices for each token. - - Protein: [N, CA, C]; DNA/RNA: [C1', C3', C4']; Ligand: distance-based. - """ - # Build atom name -> atom_index lookup per token - token_atoms: dict[int, dict[str, int]] = defaultdict(dict) - for atom in atoms: - if atom.is_valid: - token_atoms[atom.token_index][atom.name] = atom.atom_index - - # Ligand-token frames come from CCD reference-conformer geometry, - # grouped per residue. For each token, the frame is the 3 atoms nearest - # to its own atom in the residue's ref-pos space, ordered - # (1st-nearest, self, 2nd-nearest). - ligand_token_to_atom: dict[int, int] = {} - ligand_tokens_by_res: dict[tuple[int, int], list[int]] = defaultdict(list) - for t in tokens: - if t.mol_type == MOL_TYPE_NONPOLYMER: - ad = token_atoms.get(t.token_index) - if ad: - ligand_token_to_atom[t.token_index] = next(iter(ad.values())) - ligand_tokens_by_res[(t.asym_id, t.residue_index)].append(t.token_index) - - ligand_token_frames: dict[int, tuple[int, int, int]] = {} - for tok_indices in ligand_tokens_by_res.values(): - atom_indices = [ - ligand_token_to_atom[ti] for ti in tok_indices if ti in ligand_token_to_atom - ] - if len(atom_indices) < 3: - for ti in tok_indices: - if ti in ligand_token_to_atom: - ai = ligand_token_to_atom[ti] - ligand_token_frames[ti] = (ai, ai, ai) - continue - - ref_pos_chain = np.array([atoms[ai].ref_pos for ai in atom_indices]) - dist_mat = np.sqrt( - ((ref_pos_chain[:, None] - ref_pos_chain[None]) ** 2).sum(-1) - ) - sort_indices = np.argsort(dist_mat, axis=1) - local_frames = np.column_stack( - [sort_indices[:, 1], sort_indices[:, 0], sort_indices[:, 2]] - ) - - for ti in tok_indices: - if ti not in ligand_token_to_atom: - continue - ai = ligand_token_to_atom[ti] - local_idx = atom_indices.index(ai) - fl = local_frames[local_idx] - ligand_token_frames[ti] = ( - atom_indices[fl[0]], - atom_indices[fl[1]], - atom_indices[fl[2]], - ) - - # Build frames for all tokens - frames_list: list[tuple[int, int, int]] = [] - for t in tokens: - ad = token_atoms.get(t.token_index, {}) - fallback = list(ad.values())[0] if ad else 0 - - if t.mol_type == MOL_TYPE_PROTEIN: - if t.res_type == PROTEIN_UNK_RES_TYPE: - frames_list.append((fallback, fallback, fallback)) - else: - frames_list.append((ad.get("N", 0), ad.get("CA", 0), ad.get("C", 0))) - elif t.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): - if t.res_type == PROTEIN_UNK_RES_TYPE: - frames_list.append((fallback, fallback, fallback)) - else: - frames_list.append( - (ad.get("C1'", 0), ad.get("C3'", 0), ad.get("C4'", 0)) - ) - elif t.mol_type == MOL_TYPE_NONPOLYMER: - if t.token_index in ligand_token_frames: - frames_list.append(ligand_token_frames[t.token_index]) - else: - frames_list.append((fallback, fallback, fallback)) - else: - frames_list.append((fallback, fallback, fallback)) - - frames = np.array(frames_list, dtype=np.int64) - - # Compute resolved mask (vectorized) - n_atoms = len(atoms) - atom_positions = ( - np.array([a.pos for a in atoms], dtype=np.float32) - if atoms - else np.zeros((0, 3), dtype=np.float32) - ) - atom_is_valid = ( - np.array([a.is_valid for a in atoms], dtype=bool) - if atoms - else np.zeros(0, dtype=bool) - ) - atom_is_resolved = ( - atom_is_valid & np.any(atom_positions != 0, axis=1) - if n_atoms > 0 - else np.zeros(0, dtype=bool) - ) - - n_tokens = len(tokens) - if n_tokens == 0: - return frames, np.zeros(0, dtype=bool) - - pos1 = atom_positions[frames[:, 0]] - pos2 = atom_positions[frames[:, 1]] - pos3 = atom_positions[frames[:, 2]] - - all_resolved = ( - atom_is_resolved[frames[:, 0]] - & atom_is_resolved[frames[:, 1]] - & atom_is_resolved[frames[:, 2]] - ) - all_same = (frames[:, 0] == frames[:, 1]) & (frames[:, 1] == frames[:, 2]) - - v1 = pos1 - pos2 - v2 = pos3 - pos2 - norm1 = np.linalg.norm(v1, axis=1) - norm2 = np.linalg.norm(v2, axis=1) - valid_norms = (norm1 >= 1e-6) & (norm2 >= 1e-6) - - cos_angle = np.zeros(n_tokens, dtype=np.float32) - mask = valid_norms - if np.any(mask): - cos_angle[mask] = np.sum(v1[mask] * v2[mask], axis=1) / ( - norm1[mask] * norm2[mask] - ) - cos_angle = np.clip(cos_angle, -1, 1) - angle_deg = np.degrees(np.arccos(np.abs(cos_angle))) - not_colinear = angle_deg >= 25 - - resolved_mask = all_resolved & ~all_same & valid_norms & not_colinear - return frames, resolved_mask - - -def compute_token_bonds( - tokens: list[TokenInfo], - atoms: list[AtomInfo], - input: StructurePredictionInput, - chains: list[ChainInfo], -) -> torch.Tensor: - """Compute dense token bond matrix [L, L, 1]. - - Includes ligand intra-residue bonds (from CCD) and covalent bonds. - """ - n_tokens = len(tokens) - edge_set: set[tuple[int, int]] = set() - - def add_bond(i: int, j: int) -> None: - if i != j: - edge_set.add((min(i, j), max(i, j))) - - # Build per-residue atom name -> token_index mapping for ligands and modified residues - # Key: (asym_id, residue_index, atom_name) -> token_index - atom_name_to_token: dict[tuple[int, int, str], int] = {} - for atom in atoms: - if atom.is_valid: - t = tokens[atom.token_index] if atom.token_index < len(tokens) else None - if t and ( - t.mol_type == MOL_TYPE_NONPOLYMER or t.res_type == PROTEIN_UNK_RES_TYPE - ): - atom_name_to_token[(t.asym_id, t.residue_index, atom.name)] = ( - atom.token_index - ) - - # Group atom-tokenized tokens by (asym_id, residue_index) - residue_tokens: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list) - for atom in atoms: - if not atom.is_valid: - continue - t = tokens[atom.token_index] if atom.token_index < len(tokens) else None - if t and ( - t.mol_type == MOL_TYPE_NONPOLYMER or t.res_type == PROTEIN_UNK_RES_TYPE - ): - residue_tokens[(t.asym_id, t.residue_index)].append( - (atom.name, atom.token_index) - ) - - # Add intra-residue bonds from CCD - for (asym_id_val, res_idx), atom_list in residue_tokens.items(): - if not atom_list: - continue - res_name = tokens[atom_list[0][1]].residue_name - ccd_bonds = get_ligand_ccd_bonds(res_name) - atom_to_tok = {name: ti for name, ti in atom_list} - - if ccd_bonds: - for a1, a2 in ccd_bonds: - if a1 in atom_to_tok and a2 in atom_to_tok: - add_bond(atom_to_tok[a1], atom_to_tok[a2]) - else: - # Fallback: fully connected within residue - tok_indices = [ti for _, ti in atom_list] - for i_idx in tok_indices: - for j_idx in tok_indices: - add_bond(i_idx, j_idx) - - # Add covalent bonds from input - if input.covalent_bonds: - # Build chain_id -> chain mapping - chain_by_id: dict[str, ChainInfo] = {c.chain_id: c for c in chains} - # Build (asym_id, residue_index) -> list of tokens for atom index lookup - chain_res_atoms: dict[tuple[int, int], list[AtomInfo]] = defaultdict(list) - for atom in atoms: - if atom.is_valid and atom.token_index < len(tokens): - t = tokens[atom.token_index] - chain_res_atoms[(t.asym_id, t.residue_index)].append(atom) - - for cb in input.covalent_bonds: - c1 = chain_by_id.get(cb.chain_id1) - c2 = chain_by_id.get(cb.chain_id2) - if c1 is None or c2 is None: - continue - - atoms_1 = chain_res_atoms.get((c1.asym_id, cb.res_idx1), []) - atoms_2 = chain_res_atoms.get((c2.asym_id, cb.res_idx2), []) - - if cb.atom_idx1 < len(atoms_1) and cb.atom_idx2 < len(atoms_2): - add_bond( - atoms_1[cb.atom_idx1].token_index, atoms_2[cb.atom_idx2].token_index - ) - - # Add peptide bonds at modified-residue boundaries: an atom-tokenized - # residue's N atom connects to the prev residue's C atom (and same for - # the C side to the next residue's N). - tokens_by_chain_res: dict[tuple[int, int], list[TokenInfo]] = defaultdict(list) - for t in tokens: - if t.mol_type == MOL_TYPE_PROTEIN: - tokens_by_chain_res[(t.asym_id, t.residue_index)].append(t) - - def _backbone_token(res_tokens: list[TokenInfo], atom_name: str) -> int | None: - # Standard residue (single token wrapping all atoms): return that token. - if len(res_tokens) == 1 and res_tokens[0].res_type != PROTEIN_UNK_RES_TYPE: - return res_tokens[0].token_index - for t in res_tokens: - for a_idx in range(t.atom_start, t.atom_start + t.atom_count): - if a_idx < len(atoms) and atoms[a_idx].name == atom_name: - return t.token_index - # Atom-tokenized residue without an atom of that name (e.g. ACE has - # no N, NH2 has no C). Fall back to the first atom-tokenized token. - return res_tokens[0].token_index if res_tokens else None - - for (asym_id_val, res_idx), res_tokens in tokens_by_chain_res.items(): - is_atom_tokenized = any(t.res_type == PROTEIN_UNK_RES_TYPE for t in res_tokens) - if not is_atom_tokenized: - continue # Standard residue — no peptide bond added here. - n_tok = _backbone_token(res_tokens, "N") - c_tok = _backbone_token(res_tokens, "C") - prev_tokens = tokens_by_chain_res.get((asym_id_val, res_idx - 1)) - if prev_tokens and n_tok is not None: - prev_c = _backbone_token(prev_tokens, "C") - if prev_c is not None: - add_bond(prev_c, n_tok) - next_tokens = tokens_by_chain_res.get((asym_id_val, res_idx + 1)) - if next_tokens and c_tok is not None: - next_n = _backbone_token(next_tokens, "N") - if next_n is not None: - add_bond(c_tok, next_n) - - # Expand to dense matrix - bonds = torch.zeros(n_tokens, n_tokens, 1, dtype=torch.float32) - for i, j in edge_set: - bonds[i, j, 0] = 1.0 - bonds[j, i, 0] = 1.0 - return bonds - - -def compute_representative_atoms( - tokens: list[TokenInfo], atoms: list[AtomInfo] -) -> torch.Tensor: - """Compute representative atom index per token (for token_to_rep_atom). - - Returns: - distogram_atom_idx: [L] — representative atom per token - Protein: CB (or CA for GLY), DNA/RNA: C4/C2/C1', Ligand: first atom. - """ - n_tokens = len(tokens) - - # Build atom name -> index lookup per token - token_atoms: dict[int, dict[str, int]] = defaultdict(dict) - for atom in atoms: - if atom.is_valid: - token_atoms[atom.token_index][atom.name] = atom.atom_index - - distogram_atom_idx = torch.zeros(n_tokens, dtype=torch.int64) - - for t in tokens: - ad = token_atoms.get(t.token_index, {}) - fallback_idx = list(ad.values())[0] if ad else 0 - - if t.mol_type == MOL_TYPE_PROTEIN: - rep_idx = ad.get("CB", ad.get("CA", fallback_idx)) - elif t.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): - if t.res_type in (27, 32): # Unknown nucleotides - rep_idx = ad.get("C1'", fallback_idx) - elif t.res_type in (23, 24, 28, 29): # Purines (A, G) - rep_idx = ad.get("C4", ad.get("C1'", fallback_idx)) - else: # Pyrimidines (C, U, T) - rep_idx = ad.get("C2", ad.get("C1'", fallback_idx)) - else: - rep_idx = fallback_idx - - distogram_atom_idx[t.token_index] = rep_idx - - return distogram_atom_idx - - -def compute_msa_features( - input: StructurePredictionInput, - chains: list[ChainInfo], - tokens: list[TokenInfo], - max_seqs: int = 16384, -) -> dict[str, torch.Tensor]: - """Compute MSA features from protein MSAs. - - Uses taxonomy-based pairing across chains - (:func:`paired_msa.construct_paired_msa`): rows whose FASTA header - contains ``key=N`` get paired across chains sharing the same ``N``. - - Output: msa [M, L], deletion_value [M, L], has_deletion [M, L], - deletion_mean [L], msa_mask [M, L] - """ - from .esmfold2_paired_msa import ( - construct_paired_msa, - protein_letter_to_res_type, - ) - - n_tokens = len(tokens) - - # A single ProteinInput with id=['A','B','C',...] yields one item but - # multiple chains (one per id); broadcast the MSA across all of them. - chain_msas: dict[int, MSA | None] = {} - item_idx = 0 - for item in input.sequences: - ids = [item.id] if isinstance(item.id, str) else list(item.id) - for _ in ids: - chain = chains[item_idx] - if isinstance(item, ProteinInput): - msa = item.msa - if msa is None: - msa = MSA.from_sequences([item.sequence]) - chain_msas[chain.asym_id] = msa - else: - chain_msas[chain.asym_id] = None - item_idx += 1 - - letter_to_res_type = protein_letter_to_res_type() - - # Build per-chain query res_types (used for chains without an MSA). - chain_query_res_types: dict[int, np.ndarray] = {} - for chain in chains: - chain_tokens = [t for t in tokens if t.asym_id == chain.asym_id] - chain_query_res_types[chain.asym_id] = np.array( - [t.res_type for t in chain_tokens], dtype=np.int64 - ) - - token_asym_ids = np.array([t.asym_id for t in tokens], dtype=np.int64) - token_res_ids = np.array([t.residue_index for t in tokens], dtype=np.int64) - - msa_res, del_counts, paired = construct_paired_msa( - chain_msas, - chain_query_res_types, - token_asym_ids, - token_res_ids, - letter_to_res_type=letter_to_res_type, - max_seqs=max_seqs, - ) - - # Tokens for chains without an MSA get their res_type at row 0 and gap - # elsewhere; this mirrors the prior non-protein-token branch. - for t in tokens: - if chain_msas.get(t.asym_id) is None: - msa_res[:, t.token_index] = MSA_GAP_TOKEN_ID - msa_res[0, t.token_index] = t.res_type - - if msa_res.shape[0] == 0: - msa_res = np.full((1, n_tokens), MSA_GAP_TOKEN_ID, dtype=np.int64) - del_counts = np.zeros((1, n_tokens), dtype=np.float32) - - msa_data = torch.from_numpy(msa_res) - del_data = torch.from_numpy(del_counts) - - has_deletion = del_data > 0 - deletion_value = (np.pi / 2) * torch.arctan(del_data / 3) - deletion_mean = deletion_value.mean(dim=0) - - msa_mask = torch.ones_like(msa_data, dtype=torch.bool) - - return { - "msa": msa_data, - "deletion_value": deletion_value, - "has_deletion": has_deletion, - "deletion_mean": deletion_mean, - "msa_attention_mask": msa_mask, - } - - -def compute_distogram_conditioning( - input: StructurePredictionInput, - chains: list[ChainInfo], - tokens: list[TokenInfo], - disto_center: torch.Tensor, - min_dist: float = 2.0, - max_dist: float = 22.0, - num_bins: int = 64, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute distogram conditioning from user-provided distograms. - - Returns: - disto_cond: [L, L] int64 (bin indices) - disto_cond_mask: [L, L] bool - """ - n_tokens = len(tokens) - disto_cond = torch.zeros(n_tokens, n_tokens, dtype=torch.long) - disto_cond_mask = torch.zeros(n_tokens, n_tokens, dtype=torch.bool) - - if not input.distogram_conditioning: - return disto_cond, disto_cond_mask - - # Build chain_id -> asym_id mapping - chain_id_to_asym: dict[str, int] = {c.chain_id: c.asym_id for c in chains} - - # Build asym_id -> token indices mapping - asym_to_tokens: dict[int, list[int]] = defaultdict(list) - for t in tokens: - asym_to_tokens[t.asym_id].append(t.token_index) - - boundaries = torch.linspace(min_dist, max_dist, num_bins + 1) - - for dc in input.distogram_conditioning: - asym_id_val = chain_id_to_asym.get(dc.chain_id) - if asym_id_val is None: - continue - tok_indices = asym_to_tokens[asym_id_val] - n_chain = len(tok_indices) - distogram = torch.tensor(dc.distogram, dtype=torch.float32) - - if distogram.shape != (n_chain, n_chain): - raise ValueError( - f"Distogram shape {distogram.shape} doesn't match chain length {n_chain}" - ) - - # Bin the distogram - binned = torch.bucketize(distogram, boundaries[:-1]) - 1 - binned = binned.clamp(0, num_bins - 1) - - for i, ti in enumerate(tok_indices): - for j, tj in enumerate(tok_indices): - disto_cond[ti, tj] = binned[i, j] - disto_cond_mask[ti, tj] = True - - return disto_cond, disto_cond_mask - - -def build_feature_tensors( - chains: list[ChainInfo], - tokens: list[TokenInfo], - atoms: list[AtomInfo], - input: StructurePredictionInput, -) -> dict[str, torch.Tensor]: - """Build all model input tensors from tokens and atoms.""" - n_tokens = len(tokens) - n_real_atoms = len(atoms) - - # Pad atoms to nearest multiple of 32 - target_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32 - n_padding = target_atoms - n_real_atoms - padding_atoms = [ - AtomInfo( - name="", - element="", - charge=0, - ref_pos=_ZERO_POS.copy(), - pos=_ZERO_POS.copy(), - token_index=0, - atom_index=n_real_atoms + i, - space_uid=0, - is_valid=False, - ) - for i in range(n_padding) - ] - all_atoms = atoms + padding_atoms - n_atoms = len(all_atoms) - - # --- Token-level tensors --- - token_index_arr = np.empty(n_tokens, dtype=np.int64) - residue_index_arr = np.empty(n_tokens, dtype=np.int64) - asym_id_arr = np.empty(n_tokens, dtype=np.int64) - sym_id_arr = np.empty(n_tokens, dtype=np.int64) - entity_id_arr = np.empty(n_tokens, dtype=np.int64) - mol_type_arr = np.empty(n_tokens, dtype=np.int64) - res_type_arr = np.empty(n_tokens, dtype=np.int64) - input_ids_arr = np.empty(n_tokens, dtype=np.int64) - - for i, t in enumerate(tokens): - token_index_arr[i] = t.token_index - residue_index_arr[i] = t.residue_index - asym_id_arr[i] = t.asym_id - sym_id_arr[i] = t.sym_id - entity_id_arr[i] = t.entity_id - mol_type_arr[i] = t.mol_type - res_type_arr[i] = t.res_type - input_ids_arr[i] = t.input_id - - token_index = torch.from_numpy(token_index_arr) - residue_index = torch.from_numpy(residue_index_arr) - asym_id = torch.from_numpy(asym_id_arr) - sym_id = torch.from_numpy(sym_id_arr) - entity_id = torch.from_numpy(entity_id_arr) - mol_type = torch.from_numpy(mol_type_arr) - res_type = torch.from_numpy(res_type_arr) - input_ids = torch.from_numpy(input_ids_arr) - token_pad_mask = torch.ones(n_tokens, dtype=torch.bool) - - # --- Atom-level tensors --- - ref_pos_arr = np.zeros((n_atoms, 3), dtype=np.float32) - ref_element_arr = np.zeros(n_atoms, dtype=np.int64) - ref_charge_arr = np.zeros(n_atoms, dtype=np.int8) - ref_atom_name_chars_arr = np.zeros((n_atoms, 4), dtype=np.int64) - ref_space_uid_arr = np.zeros(n_atoms, dtype=np.int64) - atom_pad_mask_arr = np.zeros(n_atoms, dtype=np.bool_) - atom_to_token_arr = np.zeros(n_atoms, dtype=np.int64) - all_positions = np.zeros((n_atoms, 3), dtype=np.float64) - is_valid_arr = np.zeros(n_atoms, dtype=np.bool_) - - for i, atom in enumerate(all_atoms): - if atom.ref_pos is not None: - ref_pos_arr[i] = atom.ref_pos - ref_charge_arr[i] = atom.charge - ref_space_uid_arr[i] = ( - atom.space_uid if atom.space_uid >= 0 else atom.token_index - ) - atom_pad_mask_arr[i] = atom.is_valid - is_valid_arr[i] = atom.is_valid - all_positions[i] = atom.pos - - if atom.is_valid: - ref_element_arr[i] = get_element_atomic_num(atom.element) - name_indices = encode_atom_name(atom.name) - ref_atom_name_chars_arr[i] = name_indices - atom_to_token_arr[i] = atom.token_index - - ref_pos = torch.from_numpy(ref_pos_arr) - ref_element = torch.from_numpy(ref_element_arr) - ref_charge = torch.from_numpy(ref_charge_arr) - ref_atom_name_chars = torch.from_numpy(ref_atom_name_chars_arr) - ref_space_uid = torch.from_numpy(ref_space_uid_arr) - atom_pad_mask = torch.from_numpy(atom_pad_mask_arr) - atom_to_token = torch.from_numpy(atom_to_token_arr) - - # Coordinates — center on resolved atoms - raw_coords = torch.from_numpy(all_positions) - is_nonzero = np.any(all_positions != 0, axis=1) - atom_resolved_arr = is_valid_arr & is_nonzero - resolved_mask = torch.from_numpy(atom_resolved_arr) - valid_mask = torch.from_numpy(is_valid_arr) - - if resolved_mask.any(): - centroid = raw_coords[resolved_mask].mean(dim=0, keepdim=True) - raw_coords = raw_coords - centroid - raw_coords[~valid_mask] = 0.0 - - coords = raw_coords.float().unsqueeze(0) # [1, A, 3] - atom_resolved_mask = torch.tensor(atom_resolved_arr, dtype=torch.bool) - - # --- Frames --- - frames, _ = compute_frame_indices(tokens, atoms) - frames_idx = torch.from_numpy(frames).to(torch.int64) - - # --- Token bonds --- - token_bonds = compute_token_bonds(tokens, atoms, input, chains) - - # --- Representative atoms --- - distogram_atom_idx = compute_representative_atoms(tokens, atoms) - - # --- MSA features --- - msa_features = compute_msa_features(input, chains, tokens) - - # --- Distogram conditioning --- - # disto_center is not needed for inference (no experimental coords) - disto_center = torch.zeros(n_tokens, 3, dtype=torch.float32) - disto_cond, disto_cond_mask = compute_distogram_conditioning( - input, chains, tokens, disto_center - ) - - # ref_pos: CCD conformer positions, used as-is for inference. - # No random rotation or masking — at inference there are no resolved - # experimental coordinates, so atom_resolved_mask is all False. - # The model uses ref_pos for atom feature embedding. - - # --- Pocket (dropped) --- - pocket_feature = torch.zeros(n_tokens, dtype=torch.long) - - return { - # Token-level - "token_index": token_index, - "residue_index": residue_index, - "asym_id": asym_id, - "entity_id": entity_id, - "sym_id": sym_id, - "mol_type": mol_type, - "res_type": res_type, - "input_ids": input_ids, - "token_bonds": token_bonds, - "token_attention_mask": token_pad_mask, - "pocket_feature": pocket_feature, - # Atom-level - "ref_pos": ref_pos, - "ref_element": ref_element, - "ref_charge": ref_charge, - "ref_atom_name_chars": ref_atom_name_chars, - "ref_space_uid": ref_space_uid, - "gt_coords": coords, - "atom_attention_mask": atom_pad_mask, - "atom_to_token": atom_to_token, - "is_resolved": atom_resolved_mask, - "distogram_atom_idx": distogram_atom_idx, - # Frames - "frames_idx": frames_idx, - # Distogram - "disto_cond": disto_cond, - "disto_cond_mask": disto_cond_mask, - # MSA - **msa_features, - } - - -# ============================================================================= -# Top-level entry point -# ============================================================================= - - -def prepare_esmfold2_input( - input: StructurePredictionInput, seed: int | None = None -) -> tuple[dict[str, torch.Tensor], list[ChainInfo]]: - """Prepare ESMFold2 model inputs from StructurePredictionInput. - - Args: - input: The structure prediction input (sequences, conditioning, etc.) - seed: Random seed for SMILES conformer generation and augmentation. - - Returns: - Tuple of (feature_dict, chain_infos) where feature_dict contains - all tensors for the model forward pass, and chain_infos contains - metadata for output processing. - """ - chains, tokens, atoms = build_chains_from_input(input, seed) - features = build_feature_tensors(chains, tokens, atoms, input) - return features, chains diff --git a/fastplms/esmfold2/esmfold2_processor.py b/fastplms/esmfold2/esmfold2_processor.py deleted file mode 100644 index d085462..0000000 --- a/fastplms/esmfold2/esmfold2_processor.py +++ /dev/null @@ -1,355 +0,0 @@ -import random -from contextlib import contextmanager, nullcontext -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from .esmfold2_conformers import load_ccd -from .esmfold2_output import build_molecular_complex_from_features -from .esmfold2_prepare_input import ChainInfo, prepare_esmfold2_input -from .esmfold2_types import ( - MSA, - Modification, - ProteinInput, - StructurePredictionInput, -) -from .esmfold2_molecular_complex import MolecularComplexResult - - -@contextmanager -def _seed_context(seed: int | None): - if seed is None: - yield - return - py_state = random.getstate() - np_state = np.random.get_state() - torch_state = torch.random.get_rng_state() - cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - try: - yield - finally: - random.setstate(py_state) - np.random.set_state(np_state) - torch.random.set_rng_state(torch_state) - if cuda_state is not None: - torch.cuda.set_rng_state_all(cuda_state) - - -def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePredictionInput: - """Group identical protein sequences into the same ProteinInput with multiple ids. - - Example: Passing a tetramer like [ProteinInput(id=["0"], seq="AAA|AAA|BBB|BBB")] - gets converted into [ProteinInput(id=["0_0", "0_1"], seq="AAA"), - ProteinInput(id=["0_2", "0_3"], seq="BBB")] - - Preserves the original order of unique sequences. Also converts "|" chainbreak - tokens to ":" in the sequence. - """ - cleaned_sequences: list = [] - chain_to_ids: dict[str, list[str]] = {} - chain_to_modifications: dict[str, list] = {} - chain_to_msa: dict[str, MSA | None] = {} - - for item in input.sequences: - if isinstance(item, ProteinInput): - sequence = ":".join(item.sequence.split("|")) - if ":" not in sequence: - cleaned_sequences.append(item) - continue - - if ":" in sequence and input.covalent_bonds is not None: - raise ValueError( - "Covalent bonds are not supported when using chainbreaks. " - "Chains must be separated into multiple ProteinInput objects." - ) - - base_id = item.id[0] if isinstance(item.id, list) else item.id - chain_to_ids = {} - chain_to_modifications = {} - chain_to_msa = {} - chains = sequence.split(":") - - chain_start_positions = [] - pos = 0 - for chain in chains: - chain_start_positions.append(pos) - pos += len(chain) + 1 - - if item.modifications is not None: - for chain_idx, chain in enumerate(chains): - chain_start = chain_start_positions[chain_idx] - chain_end = chain_start + len(chain) - chain_modifications = [] - for mod in item.modifications: - if chain_start <= mod.position < chain_end: - adjusted_mod = Modification( - position=mod.position - chain_start, ccd=mod.ccd - ) - chain_modifications.append(adjusted_mod) - if chain not in chain_to_modifications: - chain_to_modifications[chain] = chain_modifications - else: - chain_to_modifications[chain].extend(chain_modifications) - - if item.msa is not None: - for chain_idx, chain in enumerate(chains): - if chain not in chain_to_msa: - chain_start = chain_start_positions[chain_idx] - chain_end = chain_start + len(chain) - chain_msa = item.msa.select_positions( # type: ignore - np.arange(chain_start, chain_end) - ) - chain_to_msa[chain] = chain_msa - - for i, chain in enumerate(chains): - chain_id = base_id + "_" + str(i) - if chain in chain_to_ids: - chain_to_ids[chain].append(chain_id) - else: - chain_to_ids[chain] = [chain_id] - cleaned_sequences.append((item, chain)) - else: - cleaned_sequences.append(item) - - for i in range(len(cleaned_sequences)): - if isinstance(cleaned_sequences[i], tuple): - item, chain = cleaned_sequences[i] - chain_ids = chain_to_ids[chain] - chain_modifications = ( - chain_to_modifications.get(chain) if item.modifications else None - ) - chain_msa = chain_to_msa.get(chain) if item.msa else None - cleaned_sequences[i] = ProteinInput( - id=chain_ids, - sequence=chain, - msa=chain_msa, - modifications=chain_modifications, - ) - - return StructurePredictionInput( - sequences=cleaned_sequences, - distogram_conditioning=input.distogram_conditioning, - covalent_bonds=input.covalent_bonds, - ) - - -class ESMFold2InputBuilder: - def __init__(self, ccd_cache: Path | None = None): - load_ccd(ccd_cache) - - def prepare_input( - self, - input: StructurePredictionInput, - seed: int | None = None, - device: torch.device | str | None = None, - ) -> tuple[dict, list[ChainInfo]]: - """Prepare raw input for the folding model. - - Converts user-provided StructurePredictionInput into batched tensors - ready for model inference. - - Parameters - ---------- - input : StructurePredictionInput - Input specification (sequences, structures, constraints, etc.). - seed : int, optional - Random seed for reproducibility. - device : torch.device or str, optional - Target device for the returned tensors. Defaults to CPU; pass - ``model.device`` to skip a separate ``.to(...)`` step. ``fold()`` - forwards ``model.device`` automatically. - - Returns - ------- - tuple[dict, list[ChainInfo]] - Batched input tensors and chain metadata for output processing. - """ - structure_prediction_input = clean_esmfold2_input(input) - with _seed_context(seed) if seed is not None else nullcontext(): - features, chain_infos = prepare_esmfold2_input( - structure_prediction_input, seed=seed - ) - features = { - k: (v[None].to(device) if device is not None else v[None]) - if isinstance(v, torch.Tensor) - else v - for k, v in features.items() - } - - return features, chain_infos - - def __call__( - self, - input: StructurePredictionInput, - seed: int | None = None, - device: torch.device | str | None = None, - ) -> tuple[dict, list[ChainInfo]]: - return self.prepare_input(input, seed=seed, device=device) - - def decode( - self, - output: dict[str, torch.Tensor], - features: dict[str, torch.Tensor], - chain_infos: list[ChainInfo], - *, - num_diffusion_samples: int = 1, - complex_id: str = "pred", - ) -> MolecularComplexResult | list[MolecularComplexResult]: - """Convert raw model outputs into one MolecularComplexResult per sample. - - Parameters - ---------- - output : dict[str, Tensor] - Output dict returned by ESMFold2Model.forward. - features : dict[str, Tensor] - Feature dict from :meth:`prepare_input` (batched, on the model device). - chain_infos : list[ChainInfo] - Chain metadata returned alongside `features`. - num_diffusion_samples : int - Number of diffusion samples present in the output (Bm = B * num_diffusion_samples). - complex_id : str - Identifier assigned to each MolecularComplex. - - Returns - ------- - MolecularComplexResult or list[MolecularComplexResult] - A single result when num_diffusion_samples == 1, otherwise a list of length Bm. - """ - atom_mask = features["atom_attention_mask"][0] - ref_element = features["ref_element"][0] - ref_atom_name_chars = features["ref_atom_name_chars"][0] - - sample_coords = output["sample_atom_coords"] - plddts = output["plddt"] - Bm = sample_coords.shape[0] - - ptm_t = output.get("ptm") - iptm_t = output.get("iptm") - pae_t = output.get("pae") - distogram_t = output.get("distogram_logits") - pair_chains_t = output.get("pair_chains_iptm") - residue_index_t = output.get("residue_index") - entity_id_t = output.get("entity_id") - - results: list[MolecularComplexResult] = [] - for i in range(Bm): - mc = build_molecular_complex_from_features( - coords=sample_coords[i], - plddt=plddts[i], - atom_mask=atom_mask, - ref_element=ref_element, - ref_atom_name_chars=ref_atom_name_chars, - chain_infos=chain_infos, - complex_id=complex_id, - ) - results.append( - MolecularComplexResult( - complex=mc, - plddt=plddts[i].detach().cpu(), - ptm=float(ptm_t[i].item()) if ptm_t is not None else None, - iptm=float(iptm_t[i].item()) if iptm_t is not None else None, - pae=pae_t[i].detach().cpu() if pae_t is not None else None, - distogram=( - distogram_t[0].detach().cpu() - if distogram_t is not None - else None - ), - pair_chains_iptm=( - pair_chains_t[i].detach().cpu() - if pair_chains_t is not None - else None - ), - residue_index=( - residue_index_t[0].detach().cpu() - if residue_index_t is not None - else None - ), - entity_id=( - entity_id_t[0].detach().cpu() - if entity_id_t is not None - else None - ), - ) - ) - - if num_diffusion_samples == 1 and len(results) == 1: - return results[0] - return results - - def fold( - self, - model: Any, - input: StructurePredictionInput, - *, - num_loops: int = 3, - num_sampling_steps: int = 200, - num_diffusion_samples: int = 1, - seed: int | None = None, - noise_scale: float | None = None, - step_scale: float | None = None, - max_inference_sigma: int | None = None, - early_exit: bool = False, - complex_id: str = "pred", - ) -> MolecularComplexResult | list[MolecularComplexResult]: - """Fold a structure end-to-end: encode → model → decode. - - Parameters - ---------- - model : ESMFold2Model - The folding model. Must already be on the target device and in eval mode. - input : StructurePredictionInput - User-facing input specification. - num_loops, num_sampling_steps, num_diffusion_samples : int - Inference knobs forwarded to the model. - seed : int, optional - Seeds both input prep (SMILES conformer generation) and diffusion sampling. - noise_scale, step_scale, max_inference_sigma, early_exit - Optional sampler overrides forwarded to the model when not None. - complex_id : str - Identifier assigned to the predicted MolecularComplex(es). - - Returns - ------- - MolecularComplexResult or list[MolecularComplexResult] - A single result when num_diffusion_samples == 1, otherwise a list. - """ - features, chain_infos = self.prepare_input( - input, seed=seed, device=model.device - ) - - sampler_kwargs: dict[str, Any] = {} - if noise_scale is not None: - sampler_kwargs["noise_scale"] = noise_scale - if step_scale is not None: - sampler_kwargs["step_scale"] = step_scale - if max_inference_sigma is not None: - sampler_kwargs["max_inference_sigma"] = max_inference_sigma - - with torch.no_grad(): - with _seed_context(seed) if seed is not None else nullcontext(): - output = model( - **features, - num_loops=num_loops, - num_sampling_steps=num_sampling_steps, - num_diffusion_samples=num_diffusion_samples, - early_exit=early_exit, - **sampler_kwargs, - ) - - return self.decode( - output, - features, - chain_infos, - num_diffusion_samples=num_diffusion_samples, - complex_id=complex_id, - ) - - -__all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input"] diff --git a/fastplms/esmfold2/esmfold2_protein_structure.py b/fastplms/esmfold2/esmfold2_protein_structure.py deleted file mode 100644 index 650370c..0000000 --- a/fastplms/esmfold2/esmfold2_protein_structure.py +++ /dev/null @@ -1,306 +0,0 @@ -from __future__ import annotations - -from typing import Tuple, TypeVar - -import numpy as np -import torch -import torch.nn.functional as F -from torch import Tensor -from torch.amp import autocast # type: ignore - -from . import esmfold2_residue_constants as residue_constants -from .esmfold2_misc import unbinpack -from .esmfold2_affine3d import Affine3D - -ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) - - -def index_by_atom_name( - atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2 -) -> ArrayOrTensor: - squeeze = False - if isinstance(atom_names, str): - atom_names = [atom_names] - squeeze = True - indices = [residue_constants.atom_order[atom_name] for atom_name in atom_names] - dim = dim % atom37.ndim - index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim)) - result = atom37[index] # type: ignore - if squeeze: - result = result.squeeze(dim) - return result - - -def infer_cbeta_from_atom37( - atom37: ArrayOrTensor, L: float = 1.522, A: float = 1.927, D: float = -2.143 -): - """ - Inspired by a util in trDesign: - https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L92 - - input: atom37, (L)ength, (A)ngle, and (D)ihedral - output: 4th coord - """ - N = index_by_atom_name(atom37, "N", dim=-2) - CA = index_by_atom_name(atom37, "CA", dim=-2) - C = index_by_atom_name(atom37, "C", dim=-2) - - if isinstance(atom37, np.ndarray): - - def normalize(x: ArrayOrTensor): - return x / np.linalg.norm(x, axis=-1, keepdims=True) - - cross = np.cross - else: - normalize = F.normalize # type: ignore - cross = torch.cross - - with np.errstate(invalid="ignore"): # inf - inf = nan is ok here - vec_nca = N - CA - vec_nc = N - C - nca = normalize(vec_nca) - n = normalize(cross(vec_nc, nca)) # type: ignore - m = [nca, cross(n, nca), n] - d = [L * np.cos(A), L * np.sin(A) * np.cos(D), -L * np.sin(A) * np.sin(D)] - return CA + sum([m * d for m, d in zip(m, d)]) - - -@torch.no_grad() -@autocast("cuda", enabled=False) -def compute_alignment_tensors( - mobile: torch.Tensor, - target: torch.Tensor, - atom_exists_mask: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, -): - """ - Align two batches of structures with support for masking invalid atoms using PyTorch. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N) - - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking. - - Returns: - - centered_mobile (torch.Tensor): Batch of coordinates of structure centered mobile (B, N, 3) - - centroid_mobile (torch.Tensor): Batch of coordinates of mobile centeroid (B, 3) - - centered_target (torch.Tensor): Batch of coordinates of structure centered target (B, N, 3) - - centroid_target (torch.Tensor): Batch of coordinates of target centeroid (B, 3) - - rotation_matrix (torch.Tensor): Batch of coordinates of rotation matrix (B, 3, 3) - - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,) - """ - - # Ensure both batches have the same number of structures, atoms, and dimensions - if sequence_id is not None: - mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan) - target = unbinpack(target, sequence_id, pad_value=torch.nan) - if atom_exists_mask is not None: - atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=0) - else: - atom_exists_mask = torch.isfinite(target).all(-1) - - assert mobile.shape == target.shape, "Batch structure shapes do not match!" - - # Number of structures in the batch - batch_size = mobile.shape[0] - - # if [B, Nres, Natom, 3], resize - if mobile.dim() == 4: - mobile = mobile.view(batch_size, -1, 3) - if target.dim() == 4: - target = target.view(batch_size, -1, 3) - if atom_exists_mask is not None and atom_exists_mask.dim() == 3: - atom_exists_mask = atom_exists_mask.view(batch_size, -1) - - # Number of atoms - num_atoms = mobile.shape[1] - - # Apply masks if provided - if atom_exists_mask is not None: - mobile = mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0) - target = target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0) - else: - atom_exists_mask = torch.ones( - batch_size, num_atoms, dtype=torch.bool, device=mobile.device - ) - - num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True) - # Compute centroids for each batch - centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) - centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) - - # Handle potential division by zero if all atoms are invalid in a structure - centroid_mobile[num_valid_atoms == 0] = 0 - centroid_target[num_valid_atoms == 0] = 0 - - # Center structures by subtracting centroids - centered_mobile = mobile - centroid_mobile - centered_target = target - centroid_target - - centered_mobile = centered_mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0) - centered_target = centered_target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0) - - # Compute covariance matrix for each batch - covariance_matrix = torch.matmul(centered_mobile.transpose(1, 2), centered_target) - - # Singular Value Decomposition for each batch - u, _, v = torch.svd(covariance_matrix) - - # Calculate rotation matrices for each batch - rotation_matrix = torch.matmul(u, v.transpose(1, 2)) - - return ( - centered_mobile, - centroid_mobile, - centered_target, - centroid_target, - rotation_matrix, - num_valid_atoms, - ) - - -@torch.no_grad() -@autocast("cuda", enabled=False) -def compute_rmsd_no_alignment( - aligned: torch.Tensor, - target: torch.Tensor, - num_valid_atoms: torch.Tensor, - reduction: str = "batch", -) -> torch.Tensor: - """ - Compute RMSD between two batches of structures without alignment. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,) - - reduction (str): One of "batch", "per_sample", "per_residue". - - Returns: - - If reduction == "batch": - (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch - If reduction == "per_sample": - (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch - If reduction == "per_residue": - (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch - """ - if reduction not in ("per_residue", "per_sample", "batch"): - raise ValueError("Unrecognized reduction: '{reduction}'") - # Compute RMSD for each batch - diff = aligned - target - if reduction == "per_residue": - mean_squared_error = diff.square().view(diff.size(0), -1, 9).mean(dim=-1) - else: - mean_squared_error = diff.square().sum(dim=(1, 2)) / ( - num_valid_atoms.squeeze(-1) - ) - - rmsd = torch.sqrt(mean_squared_error) - if reduction in ("per_sample", "per_residue"): - return rmsd - elif reduction == "batch": - avg_rmsd = rmsd.masked_fill(num_valid_atoms.squeeze(-1) == 0, 0).sum() / ( - (num_valid_atoms > 0).sum() + 1e-8 - ) - return avg_rmsd - else: - raise ValueError(reduction) - - -@torch.no_grad() -@autocast("cuda", enabled=False) -def compute_affine_and_rmsd( - mobile: torch.Tensor, - target: torch.Tensor, - atom_exists_mask: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, -) -> Tuple[Affine3D, torch.Tensor]: - """ - Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N) - - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking. - - Returns: - - affine (Affine3D): Transformation between mobile and target structure - - avg_rmsd (torch.Tensor): Average Root Mean Square Deviation between the structures for each batch - """ - - ( - centered_mobile, - centroid_mobile, - centered_target, - centroid_target, - rotation_matrix, - num_valid_atoms, - ) = compute_alignment_tensors( - mobile=mobile, - target=target, - atom_exists_mask=atom_exists_mask, - sequence_id=sequence_id, - ) - - # Apply rotation to mobile centroid - translation = torch.matmul(-centroid_mobile, rotation_matrix) + centroid_target - affine = Affine3D.from_tensor_pair( - translation, rotation_matrix.unsqueeze(dim=-3).transpose(-2, -1) - ) - - # Apply transformation to centered structure to compute rmsd - rotated_mobile = torch.matmul(centered_mobile, rotation_matrix) - avg_rmsd = compute_rmsd_no_alignment( - rotated_mobile, centered_target, num_valid_atoms, reduction="batch" - ) - - return affine, avg_rmsd - - -def compute_gdt_ts_no_alignment( - aligned: torch.Tensor, - target: torch.Tensor, - atom_exists_mask: torch.Tensor, - reduction: str = "batch", -) -> torch.Tensor: - """ - Compute GDT_TS between two batches of structures without alignment. - - Args: - - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3) - - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3) - - atom_exists_mask (torch.Tensor): Mask for Whether an atom exists of shape (B, N). noo - - reduction (str): One of "batch", "per_sample". - - Returns: - If reduction == "batch": - (torch.Tensor): 0-dim, GDT_TS between the structures for each batch - If reduction == "per_sample": - (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch - """ - if reduction not in ("per_sample", "batch"): - raise ValueError("Unrecognized reduction: '{reduction}'") - - if atom_exists_mask is None: - atom_exists_mask = torch.isfinite(target).all(dim=-1) - - deviation = torch.linalg.vector_norm(aligned - target, dim=-1) - num_valid_atoms = atom_exists_mask.sum(dim=-1) - - # Compute GDT_TS - score = ( - ((deviation < 1) * atom_exists_mask).sum(dim=-1) / num_valid_atoms - + ((deviation < 2) * atom_exists_mask).sum(dim=-1) / num_valid_atoms - + ((deviation < 4) * atom_exists_mask).sum(dim=-1) / num_valid_atoms - + ((deviation < 8) * atom_exists_mask).sum(dim=-1) / num_valid_atoms - ) * 0.25 - - if reduction == "batch": - return score.mean() - elif reduction == "per_sample": - return score - else: - raise ValueError("Unrecognized reduction: '{reduction}'") diff --git a/fastplms/esmfold2/esmfold2_residue_constants.py b/fastplms/esmfold2/esmfold2_residue_constants.py deleted file mode 100644 index 81b379e..0000000 --- a/fastplms/esmfold2/esmfold2_residue_constants.py +++ /dev/null @@ -1,1223 +0,0 @@ -# Copyright 2025 EvolutionaryScale -# Copyright 2021 AlQuraishi Laboratory -# Copyright 2021 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Constants used in AlphaFold.""" - -import collections -import functools -from pathlib import Path -from typing import List, Mapping, Tuple - -import numpy as np - -# import tree - -# Internal import (35fd). - - -# Distance from one CA to next CA [trans configuration: omega = 180]. -ca_ca = 3.80209737096 - -# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in -# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have -# chi angles so their chi angle lists are empty. -chi_angles_atoms = { - "ALA": [], - # Chi5 in arginine is always 0 +- 5 degrees, so ignore it. - "ARG": [ - ["N", "CA", "CB", "CG"], - ["CA", "CB", "CG", "CD"], - ["CB", "CG", "CD", "NE"], - ["CG", "CD", "NE", "CZ"], - ], - "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], - "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], - "CYS": [["N", "CA", "CB", "SG"]], - "GLN": [ - ["N", "CA", "CB", "CG"], - ["CA", "CB", "CG", "CD"], - ["CB", "CG", "CD", "OE1"], - ], - "GLU": [ - ["N", "CA", "CB", "CG"], - ["CA", "CB", "CG", "CD"], - ["CB", "CG", "CD", "OE1"], - ], - "GLY": [], - "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]], - "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]], - "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], - "LYS": [ - ["N", "CA", "CB", "CG"], - ["CA", "CB", "CG", "CD"], - ["CB", "CG", "CD", "CE"], - ["CG", "CD", "CE", "NZ"], - ], - "MET": [ - ["N", "CA", "CB", "CG"], - ["CA", "CB", "CG", "SD"], - ["CB", "CG", "SD", "CE"], - ], - "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], - "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]], - "SER": [["N", "CA", "CB", "OG"]], - "THR": [["N", "CA", "CB", "OG1"]], - "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], - "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], - "VAL": [["N", "CA", "CB", "CG1"]], - "UNK": [], -} - -# If chi angles given in fixed-length array, this matrix determines how to mask -# them for each AA type. The order is as per restype_order (see below). -chi_angles_mask = [ - [0.0, 0.0, 0.0, 0.0], # ALA - [1.0, 1.0, 1.0, 1.0], # ARG - [1.0, 1.0, 0.0, 0.0], # ASN - [1.0, 1.0, 0.0, 0.0], # ASP - [1.0, 0.0, 0.0, 0.0], # CYS - [1.0, 1.0, 1.0, 0.0], # GLN - [1.0, 1.0, 1.0, 0.0], # GLU - [0.0, 0.0, 0.0, 0.0], # GLY - [1.0, 1.0, 0.0, 0.0], # HIS - [1.0, 1.0, 0.0, 0.0], # ILE - [1.0, 1.0, 0.0, 0.0], # LEU - [1.0, 1.0, 1.0, 1.0], # LYS - [1.0, 1.0, 1.0, 0.0], # MET - [1.0, 1.0, 0.0, 0.0], # PHE - [1.0, 1.0, 0.0, 0.0], # PRO - [1.0, 0.0, 0.0, 0.0], # SER - [1.0, 0.0, 0.0, 0.0], # THR - [1.0, 1.0, 0.0, 0.0], # TRP - [1.0, 1.0, 0.0, 0.0], # TYR - [1.0, 0.0, 0.0, 0.0], # VAL - [0.0, 0.0, 0.0, 0.0], # UNK -] - -# The following chi angles are pi periodic: they can be rotated by a multiple -# of pi without affecting the structure. -chi_pi_periodic = [ - [0.0, 0.0, 0.0, 0.0], # ALA - [0.0, 0.0, 0.0, 0.0], # ARG - [0.0, 0.0, 0.0, 0.0], # ASN - [0.0, 1.0, 0.0, 0.0], # ASP - [0.0, 0.0, 0.0, 0.0], # CYS - [0.0, 0.0, 0.0, 0.0], # GLN - [0.0, 0.0, 1.0, 0.0], # GLU - [0.0, 0.0, 0.0, 0.0], # GLY - [0.0, 0.0, 0.0, 0.0], # HIS - [0.0, 0.0, 0.0, 0.0], # ILE - [0.0, 0.0, 0.0, 0.0], # LEU - [0.0, 0.0, 0.0, 0.0], # LYS - [0.0, 0.0, 0.0, 0.0], # MET - [0.0, 1.0, 0.0, 0.0], # PHE - [0.0, 0.0, 0.0, 0.0], # PRO - [0.0, 0.0, 0.0, 0.0], # SER - [0.0, 0.0, 0.0, 0.0], # THR - [0.0, 0.0, 0.0, 0.0], # TRP - [0.0, 1.0, 0.0, 0.0], # TYR - [0.0, 0.0, 0.0, 0.0], # VAL - [0.0, 0.0, 0.0, 0.0], # UNK -] - -# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi, -# psi and chi angles: -# 0: 'backbone group', -# 1: 'pre-omega-group', (empty) -# 2: 'phi-group', (currently empty, because it defines only hydrogens) -# 3: 'psi-group', -# 4,5,6,7: 'chi1,2,3,4-group' -# The atom positions are relative to the axis-end-atom of the corresponding -# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis -# is defined such that the dihedral-angle-definiting atom (the last entry in -# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate). -# format: [atomname, group_idx, rel_position] -rigid_group_atom_positions = { - "ALA": [ - ["N", 0, (-0.525, 1.363, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, -0.000, -0.000)], - ["CB", 0, (-0.529, -0.774, -1.205)], - ["O", 3, (0.627, 1.062, 0.000)], - ], - "ARG": [ - ["N", 0, (-0.524, 1.362, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, -0.000, -0.000)], - ["CB", 0, (-0.524, -0.778, -1.209)], - ["O", 3, (0.626, 1.062, 0.000)], - ["CG", 4, (0.616, 1.390, -0.000)], - ["CD", 5, (0.564, 1.414, 0.000)], - ["NE", 6, (0.539, 1.357, -0.000)], - ["NH1", 7, (0.206, 2.301, 0.000)], - ["NH2", 7, (2.078, 0.978, -0.000)], - ["CZ", 7, (0.758, 1.093, -0.000)], - ], - "ASN": [ - ["N", 0, (-0.536, 1.357, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, -0.000, -0.000)], - ["CB", 0, (-0.531, -0.787, -1.200)], - ["O", 3, (0.625, 1.062, 0.000)], - ["CG", 4, (0.584, 1.399, 0.000)], - ["ND2", 5, (0.593, -1.188, 0.001)], - ["OD1", 5, (0.633, 1.059, 0.000)], - ], - "ASP": [ - ["N", 0, (-0.525, 1.362, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.527, 0.000, -0.000)], - ["CB", 0, (-0.526, -0.778, -1.208)], - ["O", 3, (0.626, 1.062, -0.000)], - ["CG", 4, (0.593, 1.398, -0.000)], - ["OD1", 5, (0.610, 1.091, 0.000)], - ["OD2", 5, (0.592, -1.101, -0.003)], - ], - "CYS": [ - ["N", 0, (-0.522, 1.362, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.524, 0.000, 0.000)], - ["CB", 0, (-0.519, -0.773, -1.212)], - ["O", 3, (0.625, 1.062, -0.000)], - ["SG", 4, (0.728, 1.653, 0.000)], - ], - "GLN": [ - ["N", 0, (-0.526, 1.361, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, 0.000, 0.000)], - ["CB", 0, (-0.525, -0.779, -1.207)], - ["O", 3, (0.626, 1.062, -0.000)], - ["CG", 4, (0.615, 1.393, 0.000)], - ["CD", 5, (0.587, 1.399, -0.000)], - ["NE2", 6, (0.593, -1.189, -0.001)], - ["OE1", 6, (0.634, 1.060, 0.000)], - ], - "GLU": [ - ["N", 0, (-0.528, 1.361, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, -0.000, -0.000)], - ["CB", 0, (-0.526, -0.781, -1.207)], - ["O", 3, (0.626, 1.062, 0.000)], - ["CG", 4, (0.615, 1.392, 0.000)], - ["CD", 5, (0.600, 1.397, 0.000)], - ["OE1", 6, (0.607, 1.095, -0.000)], - ["OE2", 6, (0.589, -1.104, -0.001)], - ], - "GLY": [ - ["N", 0, (-0.572, 1.337, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.517, -0.000, -0.000)], - ["O", 3, (0.626, 1.062, -0.000)], - ], - "HIS": [ - ["N", 0, (-0.527, 1.360, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, 0.000, 0.000)], - ["CB", 0, (-0.525, -0.778, -1.208)], - ["O", 3, (0.625, 1.063, 0.000)], - ["CG", 4, (0.600, 1.370, -0.000)], - ["CD2", 5, (0.889, -1.021, 0.003)], - ["ND1", 5, (0.744, 1.160, -0.000)], - ["CE1", 5, (2.030, 0.851, 0.002)], - ["NE2", 5, (2.145, -0.466, 0.004)], - ], - "ILE": [ - ["N", 0, (-0.493, 1.373, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.527, -0.000, -0.000)], - ["CB", 0, (-0.536, -0.793, -1.213)], - ["O", 3, (0.627, 1.062, -0.000)], - ["CG1", 4, (0.534, 1.437, -0.000)], - ["CG2", 4, (0.540, -0.785, -1.199)], - ["CD1", 5, (0.619, 1.391, 0.000)], - ], - "LEU": [ - ["N", 0, (-0.520, 1.363, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, -0.000, -0.000)], - ["CB", 0, (-0.522, -0.773, -1.214)], - ["O", 3, (0.625, 1.063, -0.000)], - ["CG", 4, (0.678, 1.371, 0.000)], - ["CD1", 5, (0.530, 1.430, -0.000)], - ["CD2", 5, (0.535, -0.774, 1.200)], - ], - "LYS": [ - ["N", 0, (-0.526, 1.362, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, 0.000, 0.000)], - ["CB", 0, (-0.524, -0.778, -1.208)], - ["O", 3, (0.626, 1.062, -0.000)], - ["CG", 4, (0.619, 1.390, 0.000)], - ["CD", 5, (0.559, 1.417, 0.000)], - ["CE", 6, (0.560, 1.416, 0.000)], - ["NZ", 7, (0.554, 1.387, 0.000)], - ], - "MET": [ - ["N", 0, (-0.521, 1.364, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, 0.000, 0.000)], - ["CB", 0, (-0.523, -0.776, -1.210)], - ["O", 3, (0.625, 1.062, -0.000)], - ["CG", 4, (0.613, 1.391, -0.000)], - ["SD", 5, (0.703, 1.695, 0.000)], - ["CE", 6, (0.320, 1.786, -0.000)], - ], - "PHE": [ - ["N", 0, (-0.518, 1.363, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.524, 0.000, -0.000)], - ["CB", 0, (-0.525, -0.776, -1.212)], - ["O", 3, (0.626, 1.062, -0.000)], - ["CG", 4, (0.607, 1.377, 0.000)], - ["CD1", 5, (0.709, 1.195, -0.000)], - ["CD2", 5, (0.706, -1.196, 0.000)], - ["CE1", 5, (2.102, 1.198, -0.000)], - ["CE2", 5, (2.098, -1.201, -0.000)], - ["CZ", 5, (2.794, -0.003, -0.001)], - ], - "PRO": [ - ["N", 0, (-0.566, 1.351, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.527, -0.000, 0.000)], - ["CB", 0, (-0.546, -0.611, -1.293)], - ["O", 3, (0.621, 1.066, 0.000)], - ["CG", 4, (0.382, 1.445, 0.0)], - # ['CD', 5, (0.427, 1.440, 0.0)], - ["CD", 5, (0.477, 1.424, 0.0)], # manually made angle 2 degrees larger - ], - "SER": [ - ["N", 0, (-0.529, 1.360, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, -0.000, -0.000)], - ["CB", 0, (-0.518, -0.777, -1.211)], - ["O", 3, (0.626, 1.062, -0.000)], - ["OG", 4, (0.503, 1.325, 0.000)], - ], - "THR": [ - ["N", 0, (-0.517, 1.364, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, 0.000, -0.000)], - ["CB", 0, (-0.516, -0.793, -1.215)], - ["O", 3, (0.626, 1.062, 0.000)], - ["CG2", 4, (0.550, -0.718, -1.228)], - ["OG1", 4, (0.472, 1.353, 0.000)], - ], - "TRP": [ - ["N", 0, (-0.521, 1.363, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.525, -0.000, 0.000)], - ["CB", 0, (-0.523, -0.776, -1.212)], - ["O", 3, (0.627, 1.062, 0.000)], - ["CG", 4, (0.609, 1.370, -0.000)], - ["CD1", 5, (0.824, 1.091, 0.000)], - ["CD2", 5, (0.854, -1.148, -0.005)], - ["CE2", 5, (2.186, -0.678, -0.007)], - ["CE3", 5, (0.622, -2.530, -0.007)], - ["NE1", 5, (2.140, 0.690, -0.004)], - ["CH2", 5, (3.028, -2.890, -0.013)], - ["CZ2", 5, (3.283, -1.543, -0.011)], - ["CZ3", 5, (1.715, -3.389, -0.011)], - ], - "TYR": [ - ["N", 0, (-0.522, 1.362, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.524, -0.000, -0.000)], - ["CB", 0, (-0.522, -0.776, -1.213)], - ["O", 3, (0.627, 1.062, -0.000)], - ["CG", 4, (0.607, 1.382, -0.000)], - ["CD1", 5, (0.716, 1.195, -0.000)], - ["CD2", 5, (0.713, -1.194, -0.001)], - ["CE1", 5, (2.107, 1.200, -0.002)], - ["CE2", 5, (2.104, -1.201, -0.003)], - ["OH", 5, (4.168, -0.002, -0.005)], - ["CZ", 5, (2.791, -0.001, -0.003)], - ], - "VAL": [ - ["N", 0, (-0.494, 1.373, -0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.527, -0.000, -0.000)], - ["CB", 0, (-0.533, -0.795, -1.213)], - ["O", 3, (0.627, 1.062, -0.000)], - ["CG1", 4, (0.540, 1.429, -0.000)], - ["CG2", 4, (0.533, -0.776, 1.203)], - ], - # Assume alanine positions for unknown AA - "UNK": [ - ["N", 0, (-0.525, 1.363, 0.000)], - ["CA", 0, (0.000, 0.000, 0.000)], - ["C", 0, (1.526, -0.000, -0.000)], - ], -} - -# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention. -residue_atoms = { - "ALA": ["C", "CA", "CB", "N", "O"], - "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"], - "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"], - "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"], - "CYS": ["C", "CA", "CB", "N", "O", "SG"], - "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"], - "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"], - "GLY": ["C", "CA", "N", "O"], - "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"], - "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"], - "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"], - "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"], - "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"], - "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"], - "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"], - "SER": ["C", "CA", "CB", "N", "O", "OG"], - "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"], - "TRP": [ - "C", - "CA", - "CB", - "CG", - "CD1", - "CD2", - "CE2", - "CE3", - "CZ2", - "CZ3", - "CH2", - "N", - "NE1", - "O", - ], - "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"], - "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"], - "UNK": ["C", "CA", "N"], -} - -# Naming swaps for ambiguous atom names. -# Due to symmetries in the amino acids the naming of atoms is ambiguous in -# 4 of the 20 amino acids. -# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities -# in LEU, VAL and ARG can be resolved by using the 3d constellations of -# the 'ambiguous' atoms and their neighbours) -# TODO: ^ interpret this -residue_atom_renaming_swaps = { - "ASP": {"OD1": "OD2"}, - "GLU": {"OE1": "OE2"}, - "PHE": {"CD1": "CD2", "CE1": "CE2"}, - "TYR": {"CD1": "CD2", "CE1": "CE2"}, -} - -# Van der Waals radii [Angstroem] of the atoms (from Wikipedia) -van_der_waals_radius = {"C": 1.7, "N": 1.55, "O": 1.52, "S": 1.8} - -Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"]) -BondAngle = collections.namedtuple( - "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"] -) - - -@functools.lru_cache(maxsize=None) -def load_stereo_chemical_props() -> ( - Tuple[ - Mapping[str, List[Bond]], - Mapping[str, List[Bond]], - Mapping[str, List[BondAngle]], - ] -): - """Load stereo_chemical_props.txt into a nice structure. - - Load literature values for bond lengths and bond angles and translate - bond angles into the length of the opposite edge of the triangle - ("residue_virtual_bonds"). - - Returns: - residue_bonds: dict that maps resname --> list of Bond tuples - residue_virtual_bonds: dict that maps resname --> list of Bond tuples - residue_bond_angles: dict that maps resname --> list of BondAngle tuples - """ - stereo_chemical_props = Path( - "evolutionaryscale/structure/stereo_chemical_props.txt" - ).read_text() - - lines_iter = iter(stereo_chemical_props.splitlines()) - # Load bond lengths. - residue_bonds = {} - next(lines_iter) # Skip header line. - for line in lines_iter: - if line.strip() == "-": - break - bond, resname, length, stddev = line.split() - atom1, atom2 = bond.split("-") - if resname not in residue_bonds: - residue_bonds[resname] = [] - residue_bonds[resname].append(Bond(atom1, atom2, float(length), float(stddev))) - residue_bonds["UNK"] = [] - - # Load bond angles. - residue_bond_angles = {} - next(lines_iter) # Skip empty line. - next(lines_iter) # Skip header line. - for line in lines_iter: - if line.strip() == "-": - break - bond, resname, angle_degree, stddev_degree = line.split() - atom1, atom2, atom3 = bond.split("-") - if resname not in residue_bond_angles: - residue_bond_angles[resname] = [] - residue_bond_angles[resname].append( - BondAngle( - atom1, - atom2, - atom3, - float(angle_degree) / 180.0 * np.pi, - float(stddev_degree) / 180.0 * np.pi, - ) - ) - residue_bond_angles["UNK"] = [] - - def make_bond_key(atom1_name, atom2_name): - """Unique key to lookup bonds.""" - return "-".join(sorted([atom1_name, atom2_name])) - - # Translate bond angles into distances ("virtual bonds"). - residue_virtual_bonds = {} - for resname, bond_angles in residue_bond_angles.items(): - # Create a fast lookup dict for bond lengths. - bond_cache = {} - for b in residue_bonds[resname]: - bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b - residue_virtual_bonds[resname] = [] - for ba in bond_angles: - bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)] - bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)] - - # Compute distance between atom1 and atom3 using the law of cosines - # c^2 = a^2 + b^2 - 2ab*cos(gamma). - gamma = ba.angle_rad - length = np.sqrt( - bond1.length**2 - + bond2.length**2 - - 2 * bond1.length * bond2.length * np.cos(gamma) - ) - - # Propagation of uncertainty assuming uncorrelated errors. - dl_outer = 0.5 / length - dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer - dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer - dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer - stddev = np.sqrt( - (dl_dgamma * ba.stddev) ** 2 - + (dl_db1 * bond1.stddev) ** 2 - + (dl_db2 * bond2.stddev) ** 2 - ) - residue_virtual_bonds[resname].append( - Bond(ba.atom1_name, ba.atom3name, length, stddev) - ) - - return (residue_bonds, residue_virtual_bonds, residue_bond_angles) - - -# Between-residue bond lengths for general bonds (first element) and for Proline -# (second element). -between_res_bond_length_c_n = [1.329, 1.341] -between_res_bond_length_stddev_c_n = [0.014, 0.016] - -# Between-residue cos_angles. -between_res_cos_angles_c_n_ca = [-0.5203, 0.0353] # degrees: 121.352 +- 2.315 -between_res_cos_angles_ca_c_n = [-0.4473, 0.0311] # degrees: 116.568 +- 1.995 - -# This mapping is used when we need to store atom data in a format that requires -# fixed atom data size for every residue (e.g. a numpy array). -atom_types = [ - "N", - "CA", - "C", - "CB", - "O", - "CG", - "CG1", - "CG2", - "OG", - "OG1", - "SG", - "CD", - "CD1", - "CD2", - "ND1", - "ND2", - "OD1", - "OD2", - "SD", - "CE", - "CE1", - "CE2", - "CE3", - "NE", - "NE1", - "NE2", - "OE1", - "OE2", - "CH2", - "NH1", - "NH2", - "OH", - "CZ", - "CZ2", - "CZ3", - "NZ", - "OXT", -] -atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)} -atom_type_num = len(atom_types) # := 37. - -# A compact atom encoding with 14 columns -# pylint: disable=line-too-long -# pylint: disable=bad-whitespace -restype_name_to_atom14_names = { - "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""], - "ARG": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD", - "NE", - "CZ", - "NH1", - "NH2", - "", - "", - "", - ], - "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""], - "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""], - "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""], - "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""], - "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""], - "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""], - "HIS": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "ND1", - "CD2", - "CE1", - "NE2", - "", - "", - "", - "", - ], - "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""], - "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""], - "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""], - "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""], - "PHE": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD1", - "CD2", - "CE1", - "CE2", - "CZ", - "", - "", - "", - ], - "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""], - "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""], - "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""], - "TRP": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD1", - "CD2", - "NE1", - "CE2", - "CE3", - "CZ2", - "CZ3", - "CH2", - ], - "TYR": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD1", - "CD2", - "CE1", - "CE2", - "CZ", - "OH", - "", - "", - ], - "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""], - "UNK": ["N", "CA", "C", "", "", "", "", "", "", "", "", "", "", ""], -} -# pylint: enable=line-too-long -# pylint: enable=bad-whitespace - - -# This is the standard residue order when coding AA type as a number. -# Reproduce it by taking 3-letter AA codes and sorting them alphabetically. -restypes = [ - "A", - "R", - "N", - "D", - "C", - "Q", - "E", - "G", - "H", - "I", - "L", - "K", - "M", - "F", - "P", - "S", - "T", - "W", - "Y", - "V", -] -restype_order = {restype: i for i, restype in enumerate(restypes)} -restype_num = len(restypes) # := 20. -unk_restype_index = restype_num # Catch-all index for unknown restypes. - -restypes_with_x = restypes + ["X"] -restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)} - -bb_atoms = ["N", "CA", "C", "O"] - -# Hydrophobicity by residue (positive values are hydrophobic). Derived from Black & Mould (1991), normalized by subtracting 0.5. -hydrophobicity = { - "ALA": 0.116, - "ARG": -0.5, - "ASN": -0.264, - "ASP": -0.472, - "CYS": 0.18, - "GLN": -0.249, - "GLU": -0.457, - "GLY": 0.001, - "HIS": -0.335, - "ILE": 0.443, - "LEU": 0.443, - "LYS": -0.217, - "MET": 0.238, - "PHE": 0.5, - "PRO": 0.211, - "SER": -0.141, - "THR": -0.05, - "TRP": 0.378, - "TYR": 0.38, - "VAL": 0.325, -} - -# Side chain max accessible surface area in Ala-X-Ala tripeptide (from Chennamsetty et al. 2010). -side_chain_asa = { - "ALA": 64.7809, - "ARG": 210.02, - "ASN": 113.187, - "ASP": 110.209, - "CYS": 95.2439, - "GLN": 147.855, - "GLU": 143.924, - "GLY": 23.1338, - "HIS": 146.449, - "ILE": 151.242, - "LEU": 139.524, - "LYS": 177.366, - "MET": 164.674, - "PHE": 186.7, - "PRO": 111.533, - "SER": 81.2159, - "THR": 111.597, - "TRP": 229.619, - "TYR": 200.306, - "VAL": 124.237, -} - -# Approximate Volumes of amino acids in cubic angstroms. -# https://www.imgt.org/IMGTeducation/Aide-memoire/_UK/aminoacids/abbreviation.html -amino_acid_volumes = { - "A": 88.6, # Alanine - "R": 173.4, # Arginine - "N": 114.1, # Asparagine - "D": 111.1, # Aspartic acid - "C": 108.5, # Cysteine - "Q": 143.8, # Glutamine - "E": 138.4, # Glutamic acid - "G": 60.1, # Glycine - "H": 153.2, # Histidine - "I": 166.7, # Isoleucine - "L": 166.7, # Leucine - "K": 168.6, # Lysine - "M": 162.9, # Methionine - "F": 189.9, # Phenylalanine - "P": 112.7, # Proline - "S": 89.0, # Serine - "T": 116.1, # Threonine - "W": 227.8, # Tryptophan - "Y": 193.6, # Tyrosine - "V": 140.0, # Valine - "X": 88.6, # Unknown, use Alanine as approximation -} - - -def sequence_to_onehot( - sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False -) -> np.ndarray: - """Maps the given sequence into a one-hot encoded matrix. - - Args: - sequence: An amino acid sequence. - mapping: A dictionary mapping amino acids to integers. - map_unknown_to_x: If True, any amino acid that is not in the mapping will be - mapped to the unknown amino acid 'X'. If the mapping doesn't contain - amino acid 'X', an error will be thrown. If False, any amino acid not in - the mapping will throw an error. - - Returns: - A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of - the sequence. - - Raises: - ValueError: If the mapping doesn't contain values from 0 to - num_unique_aas - 1 without any gaps. - """ - num_entries = max(mapping.values()) + 1 - - if sorted(set(mapping.values())) != list(range(num_entries)): - raise ValueError( - "The mapping must have values from 0 to num_unique_aas-1 " - "without any gaps. Got: %s" % sorted(mapping.values()) - ) - - one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32) - - for aa_index, aa_type in enumerate(sequence): - if map_unknown_to_x: - if aa_type.isalpha() and aa_type.isupper(): - aa_id = mapping.get(aa_type, mapping["X"]) - else: - raise ValueError(f"Invalid character in the sequence: {aa_type}") - else: - aa_id = mapping[aa_type] - one_hot_arr[aa_index, aa_id] = 1 - - return one_hot_arr - - -restype_1to3 = { - "A": "ALA", - "R": "ARG", - "N": "ASN", - "D": "ASP", - "C": "CYS", - "Q": "GLN", - "E": "GLU", - "G": "GLY", - "H": "HIS", - "I": "ILE", - "L": "LEU", - "K": "LYS", - "M": "MET", - "F": "PHE", - "P": "PRO", - "S": "SER", - "T": "THR", - "W": "TRP", - "Y": "TYR", - "V": "VAL", - "X": "UNK", -} - - -# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple -# 1-to-1 mapping of 3 letter names to one letter names. The latter contains -# many more, and less common, three letter names as keys and maps many of these -# to the same one letter name (including 'X' and 'U' which we don't use here). -restype_3to1 = {v: k for k, v in restype_1to3.items()} - -# Define a restype name for all unknown residues. -unk_restype = "UNK" - -resnames = [restype_1to3[r] for r in restypes] + [unk_restype] -resname_to_idx = {resname: i for i, resname in enumerate(resnames)} - -hydrophobic_resnames = {"VAL", "ILE", "LEU", "PHE", "MET", "TRP"} - -# The mapping here uses hhblits convention, so that B is mapped to D, J and O -# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the -# remaining 20 amino acids are kept in alphabetical order. -# There are 2 non-amino acid codes, X (representing any amino acid) and -# "-" representing a missing amino acid in an alignment. The id for these -# codes is put at the end (20 and 21) so that they can easily be ignored if -# desired. -HHBLITS_AA_TO_ID = { - "A": 0, - "B": 2, - "C": 1, - "D": 2, - "E": 3, - "F": 4, - "G": 5, - "H": 6, - "I": 7, - "J": 20, - "K": 8, - "L": 9, - "M": 10, - "N": 11, - "O": 20, - "P": 12, - "Q": 13, - "R": 14, - "S": 15, - "T": 16, - "U": 1, - "V": 17, - "W": 18, - "X": 20, - "Y": 19, - "Z": 3, - "-": 21, -} - -# Partial inversion of HHBLITS_AA_TO_ID. -ID_TO_HHBLITS_AA = { - 0: "A", - 1: "C", # Also U. - 2: "D", # Also B. - 3: "E", # Also Z. - 4: "F", - 5: "G", - 6: "H", - 7: "I", - 8: "K", - 9: "L", - 10: "M", - 11: "N", - 12: "P", - 13: "Q", - 14: "R", - 15: "S", - 16: "T", - 17: "V", - 18: "W", - 19: "Y", - 20: "X", # Includes J and O. - 21: "-", -} - -restypes_with_x_and_gap = restypes + ["X", "-"] -MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple( - restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) - for i in range(len(restypes_with_x_and_gap)) -) - - -def _make_standard_atom_mask() -> np.ndarray: - """Returns [num_res_types, num_atom_types] mask array.""" - # +1 to account for unknown (all 0s). - mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32) - for restype, restype_letter in enumerate(restypes): - restype_name = restype_1to3[restype_letter] - atom_names = residue_atoms[restype_name] - for atom_name in atom_names: - atom_type = atom_order[atom_name] - mask[restype, atom_type] = 1 - return mask - - -STANDARD_ATOM_MASK = _make_standard_atom_mask() - - -# A one hot representation for the first and second atoms defining the axis -# of rotation for each chi-angle in each residue. -def chi_angle_atom(atom_index: int) -> np.ndarray: - """Define chi-angle rigid groups via one-hot representations.""" - chi_angles_index = {} - one_hots = [] - - for k, v in chi_angles_atoms.items(): - indices = [atom_types.index(s[atom_index]) for s in v] - indices.extend([-1] * (4 - len(indices))) - chi_angles_index[k] = indices - - for r in restypes: - res3 = restype_1to3[r] - one_hot = np.eye(atom_type_num)[chi_angles_index[res3]] - one_hots.append(one_hot) - - one_hots.append(np.zeros([4, atom_type_num])) # Add zeros for residue `X`. - one_hot = np.stack(one_hots, axis=0) - one_hot = np.transpose(one_hot, [0, 2, 1]) - - return one_hot - - -chi_atom_1_one_hot = chi_angle_atom(1) -chi_atom_2_one_hot = chi_angle_atom(2) - -# An array like chi_angles_atoms but using indices rather than names. -chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes] -# chi_angles_atom_indices = tree.map_structure( -# lambda atom_name: atom_order[atom_name], chi_angles_atom_indices -# ) -chi_angles_atom_indices = np.array( - [ - chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms))) - for chi_atoms in chi_angles_atom_indices - ] -) - -# Mapping from (res_name, atom_name) pairs to the atom's chi group index -# and atom index within that group. -chi_groups_for_atom = collections.defaultdict(list) -for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items(): - for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res): - for atom_i, atom in enumerate(chi_group): - chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i)) -chi_groups_for_atom = dict(chi_groups_for_atom) - - -def _make_rigid_transformation_4x4(ex, ey, translation): - """Create a rigid 4x4 transformation matrix from two axes and transl.""" - # Normalize ex. - ex_normalized = ex / np.linalg.norm(ex) - - # make ey perpendicular to ex - ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized - ey_normalized /= np.linalg.norm(ey_normalized) - - # compute ez as cross product - eznorm = np.cross(ex_normalized, ey_normalized) - m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose() - m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0) - return m - - -# create an array with (restype, atomtype) --> rigid_group_idx -# and an array with (restype, atomtype, coord) for the atom positions -# and compute affine transformation matrices (4,4) from one rigid group to the -# previous group -restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int) -restype_atom37_mask = np.zeros([21, 37], dtype=np.float32) -restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32) -restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int) -restype_atom14_mask = np.zeros([21, 14], dtype=np.float32) -restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32) -restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32) - - -def _make_rigid_group_constants(): - """Fill the arrays above.""" - for restype, restype_letter in enumerate(restypes_with_x): - resname = restype_1to3[restype_letter] - for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]: - atomtype = atom_order[atomname] - restype_atom37_to_rigid_group[restype, atomtype] = group_idx - restype_atom37_mask[restype, atomtype] = 1 - restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position - - atom14idx = restype_name_to_atom14_names[resname].index(atomname) - restype_atom14_to_rigid_group[restype, atom14idx] = group_idx - restype_atom14_mask[restype, atom14idx] = 1 - restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position - - for restype, restype_letter in enumerate(restypes_with_x): - resname = restype_1to3[restype_letter] - atom_positions = { - name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname] - } - - # backbone to backbone is the identity transform - restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4) - - # pre-omega-frame to backbone (currently dummy identity matrix) - restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4) - - # phi-frame to backbone - mat = _make_rigid_transformation_4x4( - ex=atom_positions["N"] - atom_positions["CA"], - ey=np.array([1.0, 0.0, 0.0]), - translation=atom_positions["N"], - ) - restype_rigid_group_default_frame[restype, 2, :, :] = mat - - # psi-frame to backbone - mat = _make_rigid_transformation_4x4( - ex=atom_positions["C"] - atom_positions["CA"], - ey=atom_positions["CA"] - atom_positions["N"], - translation=atom_positions["C"], - ) - restype_rigid_group_default_frame[restype, 3, :, :] = mat - - # chi1-frame to backbone - if chi_angles_mask[restype][0]: - base_atom_names = chi_angles_atoms[resname][0] - base_atom_positions = [atom_positions[name] for name in base_atom_names] - mat = _make_rigid_transformation_4x4( - ex=base_atom_positions[2] - base_atom_positions[1], - ey=base_atom_positions[0] - base_atom_positions[1], - translation=base_atom_positions[2], - ) - restype_rigid_group_default_frame[restype, 4, :, :] = mat - - # chi2-frame to chi1-frame - # chi3-frame to chi2-frame - # chi4-frame to chi3-frame - # luckily all rotation axes for the next frame start at (0,0,0) of the - # previous frame - for chi_idx in range(1, 4): - if chi_angles_mask[restype][chi_idx]: - axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2] - axis_end_atom_position = atom_positions[axis_end_atom_name] - mat = _make_rigid_transformation_4x4( - ex=axis_end_atom_position, - ey=np.array([-1.0, 0.0, 0.0]), - translation=axis_end_atom_position, - ) - restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat - - -_make_rigid_group_constants() - - -def make_atom14_dists_bounds(overlap_tolerance=1.5, bond_length_tolerance_factor=15.0): - """compute upper and lower bounds for bonds to assess violations.""" - restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32) - restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32) - restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32) - residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props() - for restype, restype_letter in enumerate(restypes): - resname = restype_1to3[restype_letter] - atom_list = restype_name_to_atom14_names[resname] - - # create lower and upper bounds for clashes - for atom1_idx, atom1_name in enumerate(atom_list): - if not atom1_name: - continue - atom1_radius = van_der_waals_radius[atom1_name[0]] - for atom2_idx, atom2_name in enumerate(atom_list): - if (not atom2_name) or atom1_idx == atom2_idx: - continue - atom2_radius = van_der_waals_radius[atom2_name[0]] - lower = atom1_radius + atom2_radius - overlap_tolerance - upper = 1e10 - restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower - restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower - restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper - restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper - - # overwrite lower and upper bounds for bonds and angles - for b in residue_bonds[resname] + residue_virtual_bonds[resname]: - atom1_idx = atom_list.index(b.atom1_name) - atom2_idx = atom_list.index(b.atom2_name) - lower = b.length - bond_length_tolerance_factor * b.stddev - upper = b.length + bond_length_tolerance_factor * b.stddev - restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower - restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower - restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper - restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper - restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev - restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev - return { - "lower_bound": restype_atom14_bond_lower_bound, # shape (21,14,14) - "upper_bound": restype_atom14_bond_upper_bound, # shape (21,14,14) - "stddev": restype_atom14_bond_stddev, # shape (21,14,14) - } - - -restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32) -restype_atom14_ambiguous_atoms_swap_idx = np.tile(np.arange(14, dtype=int), (21, 1)) - - -def _make_atom14_ambiguity_feats(): - for res, pairs in residue_atom_renaming_swaps.items(): - res_idx = restype_order[restype_3to1[res]] - for atom1, atom2 in pairs.items(): - atom1_idx = restype_name_to_atom14_names[res].index(atom1) - atom2_idx = restype_name_to_atom14_names[res].index(atom2) - restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1 - restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1 - restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom1_idx] = atom2_idx - restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom2_idx] = atom1_idx - - -_make_atom14_ambiguity_feats() - - -def aatype_to_str_sequence(aatype): - return "".join([restypes_with_x[aatype[i]] for i in range(len(aatype))]) - - -# NOTE(thayes): These are computed based on the average CA->C and CA->N norm from rigid_group_atom_positions -CA_TO_N_NORM = 1.4591 -CA_TO_C_NORM = 1.5252 - - -def _make_restype_atom37_to_atom14(): - """Map from atom37 to atom14 per residue type.""" - restype_atom37_to_atom14 = [] # mapping (restype, atom37) --> atom14 - for rt in restypes: - atom_names = restype_name_to_atom14_names[restype_1to3[rt]] - atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)} - restype_atom37_to_atom14.append( - [ - (atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0) - for name in atom_types - ] - ) - - restype_atom37_to_atom14.append([0] * 37) - restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32) - return restype_atom37_to_atom14 - - -def _make_restype_atom14_to_atom37(): - """Map from atom14 to atom37 per residue type.""" - restype_atom14_to_atom37 = [] # mapping (restype, atom14) --> atom37 - for rt in restypes: - atom_names = restype_name_to_atom14_names[restype_1to3[rt]] - restype_atom14_to_atom37.append( - [(atom_order[name] if name else 0) for name in atom_names] - ) - # Add dummy mapping for restype 'UNK' - restype_atom14_to_atom37.append([0] * 14) - restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32) - return restype_atom14_to_atom37 - - -RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37() -RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14() -CHAIN_BREAK_TOKEN = "|" diff --git a/fastplms/esmfold2/esmfold2_sequential_dataclass.py b/fastplms/esmfold2/esmfold2_sequential_dataclass.py deleted file mode 100644 index b731437..0000000 --- a/fastplms/esmfold2/esmfold2_sequential_dataclass.py +++ /dev/null @@ -1,157 +0,0 @@ -from abc import ABC, abstractmethod -from dataclasses import dataclass, fields, replace -from typing import TypeVar - -import numpy as np - -from .esmfold2_misc import concat_objects, slice_any_object - -T = TypeVar("T") - - -@dataclass(frozen=True) -class SequentialDataclass(ABC): - """ - This is a builder on a dataclass that allows for automatic slicing and concatenation. - - When representing multimodal data, we often have multiple datatypes which have sequence dimensions that are the same (e.g. the length of the protein). - - When applying a transformation like a crop, we want to apply this to all tensors at the same time (e.g. crop the sequence, structure, and function). - - We also have some fields that are not sequential (like an id, or data source), which we don't want to crop. - - The SequentialDataclass abstracts this cropping away, allowing you to define dataclasses that implement `__len__`, `__getitem__` and `concat` automatically. - - This is done through the `metadata` field, which can take 3 values: - `sequence` (bool): True or False, tells the dataclass whether this field is a sequential type. Default: False. - `sequence_dim` (int): Which dimension is the sequential dimension (e.g. for a list of inverse folded sequences, we want to index each sequence in the list, not the list itself). Default: 0. - `join_token` (Any): What token to use to join when concatenating elements. Default: None. - - - Example: - - @dataclass(frozen=True) - class Foo(SequentialDataclass): - id: str - sequence: str = field(metadata={"sequence": True, "join_token": "|"}) - tensor: torch.Tensor = field(metadata={"sequence": True, "join_token": torch.nan}) - - def __len__(self): - # Must implement the __len__ method - return len(self.sequence) - - >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(5)) - Foo(id='foo', sequence='ABCDE', tensor=tensor([ 0.0252, -0.3335, -0.5143, 0.0251, -1.0717])) - - >>> foo[1:4] - Foo(id='foo', sequence='BCD', tensor=tensor([-0.3335, -0.5143, 0.0251])) - - >>> foo[np.arange(5) < 3] - Foo(id='foo', sequence='ABC', tensor=tensor([ 0.0252, -0.3335, -0.5143])) - - >>> Foo.concat([foo[:2], foo[3:]]) - Foo(id='foo', sequence='AB|DE', tensor=tensor([ 0.0252, -0.3335, nan, 0.0251, -1.0717])) - - # Trying to create a type where the sequence lengths do not match raises an error - >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(6)) - ValueError: Mismatch in sequence length for field: tensor. Expected 5, received 6 - - """ - - def __post_init__(self): - self._check_sequence_lengths_match() - - @abstractmethod - def __len__(self): - raise NotImplementedError - - def __getitem__(self, idx: int | list[int] | slice | np.ndarray): - updated_fields = {} - if isinstance(idx, int): - # make it so that things remain sequential - idx = [idx] - - for fld in fields(self): - if fld.metadata.get("sequence", False): - # this is a sequence, should be the same length as all other sequences - sequence_dim = fld.metadata.get("sequence_dim", 0) - value = getattr(self, fld.name) - if value is None: - continue - match sequence_dim: - case 0: - # sequence is first dimension - value = getattr(self, fld.name) - value = slice_any_object(value, idx) - updated_fields[fld.name] = value - case 1: - new_value = [slice_any_object(item, idx) for item in value] - updated_fields[fld.name] = value.__class__(new_value) - case _: - raise NotImplementedError( - "Arbitrary slicing for different sequence length fields is not implemented" - ) - - return replace(self, **updated_fields) - - def _check_sequence_lengths_match(self): - """Checks if sequence lengths of all "sequence" fields match.""" - for fld in fields(self): - if fld.metadata.get("sequence", False) and fld.name != "complex": - # this is a sequence, should be the same length as all other sequences - sequence_dim = fld.metadata.get("sequence_dim", 0) - value = getattr(self, fld.name) - if value is None: - continue - match sequence_dim: - case 0: - # sequence is first dimension - value = getattr(self, fld.name) - if len(value) != len(self): - raise ValueError( - f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(value)}" - ) - case 1: - for item in value: - if len(item) != len(self): - raise ValueError( - f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(item)}" - ) - case _: - raise NotImplementedError( - "Arbitrary matching for different sequence length fields is not implemented" - ) - - @classmethod - def concat(cls, items: list[T], **kwargs) -> T: - updated_fields = {} - for fld in fields(cls): - if fld.metadata.get("sequence", False): - # this is a sequence, should be the same length as all other sequences - sequence_dim = fld.metadata.get("sequence_dim", 0) - join_value = fld.metadata.get("join_token", None) - if getattr(items[0], fld.name) is None: - continue - values = [getattr(item, fld.name) for item in items] - match sequence_dim: - case 0: - # sequence is first dimension - value = concat_objects(values, join_value) - updated_fields[fld.name] = value - case 1: - new_value = [ - concat_objects(item, join_value) for item in zip(*values) - ] - updated_fields[fld.name] = getattr( - items[0], fld.name - ).__class__(new_value) - case _: - raise NotImplementedError( - "Arbitrary joining for different sequence length fields is not implemented" - ) - updated_fields.update(kwargs) - - return replace( - items[0], # type: ignore - **updated_fields, - ) diff --git a/fastplms/esmfold2/esmfold2_system.py b/fastplms/esmfold2/esmfold2_system.py deleted file mode 100644 index c2800e5..0000000 --- a/fastplms/esmfold2/esmfold2_system.py +++ /dev/null @@ -1,45 +0,0 @@ -import io -import subprocess -import typing as T -from pathlib import Path - -PathLike = T.Union[str, Path] -PathOrBuffer = T.Union[PathLike, io.StringIO] - - -def run_subprocess_with_errorcheck( - *popenargs, - capture_output: bool = False, - quiet: bool = False, - env: dict[str, str] | None = None, - shell: bool = False, - executable: str | None = None, - **kws, -) -> subprocess.CompletedProcess: - """A command similar to subprocess.run, however the errormessage will - contain the stderr when using this function. This makes it significantly - easier to diagnose issues. - """ - try: - if capture_output: - stdout = subprocess.PIPE - elif quiet: - stdout = subprocess.DEVNULL - else: - stdout = None - - p = subprocess.run( - *popenargs, - stderr=subprocess.PIPE, - stdout=stdout, - check=True, - env=env, - shell=shell, - executable=executable, - **kws, - ) - except subprocess.CalledProcessError as e: - raise RuntimeError( - f"Command failed with errorcode {e.returncode}." f"\n\n{e.stderr.decode()}" - ) - return p diff --git a/fastplms/esmfold2/esmfold2_types.py b/fastplms/esmfold2/esmfold2_types.py deleted file mode 100644 index 67bc145..0000000 --- a/fastplms/esmfold2/esmfold2_types.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Re-exports of the canonical SPI dataclasses from input_builder. - -This module exists so the HF processor and downstream code can import the -ESMFold2 input types from a single namespace without picking up internal-only -sibling utilities. The actual definitions live in -``esm.utils.structure.input_builder``. -""" - -from .esmfold2_msa import MSA -from .esmfold2_parsing import FastaEntry -from .esmfold2_input_builder import ( - CovalentBond, - DistogramConditioning, - DNAInput, - LigandInput, - Modification, - ProteinInput, - RNAInput, - StructurePredictionInput, -) - -__all__ = [ - "FastaEntry", - "MSA", - "Modification", - "ProteinInput", - "RNAInput", - "DNAInput", - "LigandInput", - "DistogramConditioning", - "CovalentBond", - "StructurePredictionInput", -] diff --git a/fastplms/esmfold2/esmfold2_utils_types.py b/fastplms/esmfold2/esmfold2_utils_types.py deleted file mode 100644 index 7bf5ffb..0000000 --- a/fastplms/esmfold2/esmfold2_utils_types.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -import io -from dataclasses import dataclass -from pathlib import Path -from typing import Union - -from cloudpathlib import CloudPath - -PathLike = Union[str, Path, CloudPath] -PathOrBuffer = Union[PathLike, io.StringIO] - - -@dataclass -class FunctionAnnotation: - """Represents an annotation of a protein's function over a range of residues. - - Fields: - label (str): An entry in either the function_tokens or residue_annotations tokenizer vocabs - start (int): Start index of this annotation. 1-indexed, inclusive. - end (int): End index of this annotation. 1-indexed, inclusive. - """ - - label: str - start: int - end: int - - def to_tuple(self) -> tuple[str, int, int]: - return self.label, self.start, self.end - - def __len__(self) -> int: - """Length of the annotation.""" - return self.end - self.start + 1 diff --git a/fastplms/esmfold2/get_weights.py b/fastplms/esmfold2/get_weights.py deleted file mode 100644 index c5f31e2..0000000 --- a/fastplms/esmfold2/get_weights.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Copy Biohub ESMFold2 checkpoints into FastPLMs AutoModel repos. - -Usage: - python -m fastplms.esmfold2.get_weights - python -m fastplms.esmfold2.get_weights --skip-weights - python -m fastplms.esmfold2.get_weights --repo_ids Synthyra/ESMFold2-Fast -""" - -import argparse -import os -import tempfile -from pathlib import Path - -import torch -from huggingface_hub import HfApi, login - -from fastplms.esmfold2.configuration_esmfold2 import ESMFold2Config, normalize_esmc_id -from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model -from fastplms.esmfold2.modeling_esmfold2_experimental import ESMFold2ExperimentalModel - -SOURCE_REPOS = { - "Synthyra/ESMFold2": "biohub/ESMFold2", - "Synthyra/ESMFold2-Fast": "biohub/ESMFold2-Fast", - "Synthyra/ESMFold2-Experimental-Fast": "biohub/ESMFold2-Experimental-Fast", - "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025": "biohub/ESMFold2-Experimental-Fast-Cutoff2025", - "Synthyra/ESMFold2-Experimental": "biohub/ESMFold2-Experimental", - "Synthyra/ESMFold2-Experimental-Cutoff2025": "biohub/ESMFold2-Experimental-Cutoff2025", -} -for _size in ("300M", "600M", "6B"): - for _step in ("250", "500", "750", "1000", "1500"): - _name = f"ESMFold2-Experimental-Fast-base{_size}-step{_step}k" - SOURCE_REPOS[f"Synthyra/{_name}"] = f"biohub/{_name}" -SHARD_SIZE = "5GB" -RELEASE_AUTO_MAP = { - "AutoConfig": "configuration_esmfold2.ESMFold2Config", - "AutoModel": "modeling_esmfold2.ESMFold2Model", -} -EXPERIMENTAL_AUTO_MAP = { - "AutoConfig": "configuration_esmfold2.ESMFold2Config", - "AutoModel": "modeling_esmfold2_experimental.ESMFold2ExperimentalModel", -} -IGNORE_PATTERNS = [ - "__pycache__/*", - "*.pyc", - "configuration_esmc.py", - "configuration_esmc_sae.py", - "get_weights.py", - "modeling_esmc.py", - "modeling_esmc_sae.py", -] - - -def _prepare_config(source_repo: str) -> ESMFold2Config: - config = ESMFold2Config.from_pretrained(source_repo) - config.esmc_id = normalize_esmc_id(config.esmc_id) - config.esmc_attn_backend = "flex" - if config.type == "experimental": - config.auto_map = EXPERIMENTAL_AUTO_MAP - config.architectures = ["ESMFold2ExperimentalModel"] - else: - config.auto_map = RELEASE_AUTO_MAP - config.architectures = ["ESMFold2Model"] - return config - - -def _build_esmplusplus_composite() -> str: - from update_HF import build_composite - - composite_code = build_composite( - "fastplms/esm_plusplus/modeling_esm_plusplus.py", - include_embedding_mixin=True, - ) - compile(composite_code, "modeling_esm_plusplus.py", "exec") - return composite_code - - -def _upload_code(api: HfApi, repo_id: str, package_dir: Path) -> None: - api.upload_folder( - folder_path=str(package_dir), - repo_id=repo_id, - repo_type="model", - ignore_patterns=IGNORE_PATTERNS, - ) - composite_code = _build_esmplusplus_composite() - with tempfile.TemporaryDirectory() as tmpdir: - composite_path = Path(tmpdir) / "modeling_esm_plusplus.py" - composite_path.write_text(composite_code, encoding="utf-8") - api.upload_file( - path_or_fileobj=str(composite_path), - path_in_repo="modeling_esm_plusplus.py", - repo_id=repo_id, - repo_type="model", - ) - api.upload_file( - path_or_fileobj=str(package_dir.parent / "test_time_training.py"), - path_in_repo="test_time_training.py", - repo_id=repo_id, - repo_type="model", - ) - - -def convert_and_push( - repo_ids: list[str] | None = None, - hf_token: str | None = None, - dry_run: bool = False, - skip_weights: bool = False, -) -> None: - if hf_token is not None: - login(token=hf_token) - - api = HfApi() - package_dir = Path(__file__).resolve().parent - targets = repo_ids if repo_ids is not None else list(SOURCE_REPOS) - - for target_repo in targets: - assert target_repo in SOURCE_REPOS, ( - f"Unknown repo_id {target_repo}. Expected one of {sorted(SOURCE_REPOS)}." - ) - source_repo = SOURCE_REPOS[target_repo] - config = _prepare_config(source_repo) - - if dry_run: - _build_esmplusplus_composite() - print(f"[dry-run] validated config and code for {target_repo} <- {source_repo}") - continue - - if skip_weights: - config.push_to_hub(target_repo) - _upload_code(api, target_repo, package_dir) - print(f"[skip-weights] uploaded config/code for {target_repo}") - continue - - model_cls = ( - ESMFold2ExperimentalModel - if config.type == "experimental" - else ESMFold2Model - ) - print(f"Loading {source_repo} with FastPLMs ESMFold2 code...") - model = model_cls.from_pretrained( - source_repo, - config=config, - load_esmc=False, - dtype=torch.float32, - ) - print(f"Pushing {target_repo}...") - model.push_to_hub(target_repo, max_shard_size=SHARD_SIZE) - _upload_code(api, target_repo, package_dir) - print(f"Done. Model available at https://huggingface.co/{target_repo}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--repo_ids", - nargs="*", - default=None, - ) - parser.add_argument( - "--hf_token", - type=str, - default=os.environ["HF_TOKEN"] if "HF_TOKEN" in os.environ else None, - ) - parser.add_argument( - "--dry_run", - action="store_true", - ) - parser.add_argument( - "--skip-weights", - action="store_true", - ) - args = parser.parse_args() - - convert_and_push( - repo_ids=args.repo_ids, - hf_token=args.hf_token, - dry_run=args.dry_run, - skip_weights=args.skip_weights, - ) diff --git a/fastplms/esmfold2/modeling_esmc.py b/fastplms/esmfold2/modeling_esmc.py deleted file mode 100644 index 9711102..0000000 --- a/fastplms/esmfold2/modeling_esmc.py +++ /dev/null @@ -1,1674 +0,0 @@ -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""PyTorch ESMC model.""" - -import importlib -import math -import re -from dataclasses import dataclass -from typing import Optional, cast - -import torch -import torch.nn as nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from torch.nn import functional as F - -from transformers.modeling_outputs import ( - MaskedLMOutput, - ModelOutput, - SequenceClassifierOutput, - TokenClassifierOutput, -) -from transformers.modeling_utils import PreTrainedModel -from transformers.utils import ( - auto_docstring, - can_return_tuple, - is_flash_attn_2_available, - logging, -) -from .configuration_esmc import ESMCConfig -from .modeling_esmc_sae import _ESMCSAELayer - -logger = logging.get_logger(__name__) - -_CONFIG_FOR_DOC = "ESMCConfig" - -# Optional accelerated kernels. Pure-PyTorch fallbacks below if absent. -if is_flash_attn_2_available(): - flash_attn_module = importlib.import_module("flash_attn") - flash_bert_padding = importlib.import_module("flash_attn.bert_padding") - flash_attn_varlen_qkvpacked_func = ( - flash_attn_module.flash_attn_varlen_qkvpacked_func - ) - pad_input = flash_bert_padding.pad_input - unpad_input = flash_bert_padding.unpad_input - - _flash_attn_available = True -else: - pad_input = unpad_input = flash_attn_varlen_qkvpacked_func = None - _flash_attn_available = False - -try: - flash_rotary = importlib.import_module("flash_attn.ops.triton.rotary") - apply_triton_rotary = flash_rotary.apply_rotary - - _flash_attn_rotary_available = torch.cuda.is_available() -except ImportError: - apply_triton_rotary = None # type: ignore[assignment] - _flash_attn_rotary_available = False - -# Transformer Engine: fused LayerNorm+Linear / LayerNorm+MLP kernels with -# fp32 reduction inside the LayerNorm. Recommended on GPU for accurate bf16 -# inference; without it the pure-PyTorch fallback drifts ~O(10) in fp32 and -# ~O(100) in bf16 on the unnormalized residual stream (perplexity stays -# within rounding noise). -try: - te = importlib.import_module("transformer_engine.pytorch") - - _te_available = True -except ImportError: - te = None # type: ignore[assignment] - _te_available = False - -# xformers: preferred SDPA implementation on GPU. Provides a fused -# bf16 attention kernel with deterministic reduction order. Flash -# Attention 2 and PyTorch's ``F.scaled_dot_product_attention`` are -# progressively-less-preferred fallbacks. -try: - xops = importlib.import_module("xformers.ops") - - _xformers_available = True -except ImportError: - xops = None # type: ignore[assignment] - _xformers_available = False - -# Flash Attention 2: secondary SDPA fallback. Used when xformers is not -# installed; fp16 / bf16 only. -if _flash_attn_available: - flash_attn_func = flash_attn_module.flash_attn_func -else: - flash_attn_func = None # type: ignore[assignment] - -if not _te_available: - logger.warning( - "ESMC: transformer_engine is not installed; falling back to " - "pure-PyTorch LayerNorm+Linear / LayerNorm+MLP. Outputs will differ " - "numerically — measured on the unnormalized residual stream (before " - "the final LayerNorm), ~O(10) max-diff in fp32 and ~O(100) in bf16; " - "after the final LayerNorm these shrink to a few ULP and perplexity " - "stays within rounding noise. Install with " - "`pip install transformer-engine[pytorch]` to enable fused fp32-" - "reduction LayerNorm." - ) - -if not _xformers_available and not _flash_attn_available: - logger.warning( - "ESMC: neither xformers nor flash-attn is installed; falling back " - "to PyTorch ``F.scaled_dot_product_attention``. The attention " - "reduction order in bf16 differs from a fused kernel by ~1 bf16 " - "ULP per attention block; compounded across the 80-block stack " - "this reaches ~O(100) max-diff on the unnormalized residual stream. " - "Install xformers (preferred) with `pip install xformers` for a " - "fused attention kernel." - ) - -if torch.cuda.is_available() and not _flash_attn_rotary_available: - logger.warning( - "ESMC: flash-attn rotary kernel not installed; falling back to " - "pure-PyTorch RoPE. For faster GPU inference run `pip install flash-attn`." - ) - - -# --------------------------------------------------------------------------- -# Output dataclasses -# --------------------------------------------------------------------------- - - -@dataclass -class ESMCOutput(ModelOutput): - """ - Args: - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`): - Sequence of hidden states at the output of the last layer, after layer normalisation. - hidden_states (`torch.FloatTensor`, *optional*): - Stacked hidden states for all encoder layers. - Shape ``(n_layers, batch_size, sequence_length, d_model)``. - Returned when ``output_hidden_states=True``. - sae_outputs (`dict[str, torch.Tensor]`, *optional*): - SAE feature magnitudes keyed by SAE model name (sparse tensors). - Only populated when SAE models have been registered via - ``add_sae_models`` and ``compute_sae=True``. - attentions (`tuple(torch.FloatTensor)`, *optional*): - Per-layer attention weights of shape - ``(batch_size, num_heads, sequence_length, sequence_length)``. - Returned when ``output_attentions=True``. Not available on the - ``flash_attention_2`` path. - """ - - last_hidden_state: torch.FloatTensor | None = None - hidden_states: torch.FloatTensor | None = None - sae_outputs: dict[str, torch.Tensor] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None - - -@dataclass -class ESMCMaskedLMOutput(MaskedLMOutput): - """ - Args: - loss (`torch.FloatTensor` of shape `(1,)`, *optional*): - Masked language modelling loss. Returned when ``labels`` are provided. - logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, vocab_size)`): - Prediction scores of the language modelling head. - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`): - Final hidden states after layer normalisation. - hidden_states (`torch.FloatTensor`, *optional*): - Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``. - sae_outputs (`dict[str, torch.Tensor]`, *optional*): - SAE feature magnitudes keyed by SAE model name (sparse tensors). - attentions (`tuple(torch.FloatTensor)`, *optional*): - Per-layer attention weights of shape - ``(batch_size, num_heads, sequence_length, sequence_length)``. - Returned when ``output_attentions=True``. - """ - - loss: torch.FloatTensor | None = None - logits: torch.FloatTensor | None = None - last_hidden_state: torch.FloatTensor | None = None - hidden_states: torch.FloatTensor | None = None - sae_outputs: dict[str, torch.Tensor] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None - - -@dataclass -class ESMCTokenClassifierOutput(TokenClassifierOutput): - """ - Args: - loss (`torch.FloatTensor` of shape `(1,)`, *optional*): - Token classification loss. Returned when ``labels`` are provided. - logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_labels)`): - Classification scores (before SoftMax). - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`): - Final hidden states after layer normalisation. - hidden_states (`torch.FloatTensor`, *optional*): - Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``. - sae_outputs (`dict[str, torch.Tensor]`, *optional*): - SAE feature magnitudes keyed by SAE model name (sparse tensors). - attentions (`tuple(torch.FloatTensor)`, *optional*): - Per-layer attention weights of shape - ``(batch_size, num_heads, sequence_length, sequence_length)``. - Returned when ``output_attentions=True``. - """ - - loss: torch.FloatTensor | None = None - logits: torch.FloatTensor | None = None - last_hidden_state: torch.FloatTensor | None = None - hidden_states: torch.FloatTensor | None = None - sae_outputs: dict[str, torch.Tensor] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None - - -@dataclass -class ESMCSequenceClassifierOutput(SequenceClassifierOutput): - """ - Args: - loss (`torch.FloatTensor` of shape `(1,)`, *optional*): - Sequence classification loss. Returned when ``labels`` are provided. - logits (`torch.FloatTensor` of shape `(batch_size, num_labels)`): - Classification scores (before SoftMax). - last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`): - Final hidden states after layer normalisation. - hidden_states (`torch.FloatTensor`, *optional*): - Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``. - sae_outputs (`dict[str, torch.Tensor]`, *optional*): - SAE feature magnitudes keyed by SAE model name (sparse tensors). - attentions (`tuple(torch.FloatTensor)`, *optional*): - Per-layer attention weights of shape - ``(batch_size, num_heads, sequence_length, sequence_length)``. - Returned when ``output_attentions=True``. - """ - - loss: torch.FloatTensor | None = None - logits: torch.FloatTensor | None = None - last_hidden_state: torch.FloatTensor | None = None - hidden_states: torch.FloatTensor | None = None - sae_outputs: dict[str, torch.Tensor] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None - - -# --------------------------------------------------------------------------- -# Rotary position embedding helpers -# --------------------------------------------------------------------------- - - -def _rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - x1, x2 = x[..., ::2], x[..., 1::2] - return torch.stack((-x2, x1), dim=-1).flatten(-2, -1) - - -def _apply_rotary_emb_torch( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False -) -> torch.Tensor: - """Apply rotary position embeddings (pure PyTorch, no Triton dependency). - - Args: - x: ``(batch, seqlen, n_heads, head_dim)`` - cos: ``(seqlen, rotary_dim / 2)`` - sin: ``(seqlen, rotary_dim / 2)`` - """ - ro_dim = cos.shape[-1] * 2 - seqlen = x.size(1) - cos = cos[:seqlen].unsqueeze(1).repeat(1, 1, 2) - sin = sin[:seqlen].unsqueeze(1).repeat(1, 1, 2) - return torch.cat( - [ - x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim], interleaved) * sin, - x[..., ro_dim:], - ], - dim=-1, - ) - - -class RotaryEmbedding(nn.Module): - """Rotary position embeddings (RoPE) as described in `RoFormer`_. - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - - Args: - dim: Size of a single attention head. - base: Frequency base for the sinusoidal positions. - interleaved: If ``True`` rotate adjacent pairs (GPT-J style) instead of - splitting the head dimension in half (GPT-NeoX style). - scaling_factor: Linear scaling factor applied to position indices. - pos_idx_in_fp32: Compute position indices in float32 to avoid bf16 - rounding errors at large sequence lengths. - """ - - def __init__( - self, - dim: int, - base: float = 10000.0, - interleaved: bool = False, - scale_base: float | None = None, - scaling_factor: float = 1.0, - pos_idx_in_fp32: bool = True, - device=None, - ): - super().__init__() - self.dim = dim - self.base = base - self.interleaved = interleaved - self.scale_base = scale_base - self.scaling_factor = scaling_factor - self.pos_idx_in_fp32 = pos_idx_in_fp32 - - self._seq_len_cached = 0 - self._cos_cached: torch.Tensor | None = None - self._sin_cached: torch.Tensor | None = None - self._cos_k_cached: torch.Tensor | None = None - self._sin_k_cached: torch.Tensor | None = None - - self.reset_parameters(device=device) - - def reset_parameters(self, device=None): - inv_freq = self._compute_inv_freq(device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - arange = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) - scale = ( - (arange + 0.4 * self.dim) / (1.4 * self.dim) - if self.scale_base is not None - else None - ) - self.register_buffer("scale", scale, persistent=False) - - def _compute_inv_freq(self, device=None) -> torch.Tensor: - return 1.0 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) - / self.dim - ) - ) - - def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None): - if self.inv_freq.is_meta: - self.reset_parameters(device=device) - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - if self.pos_idx_in_fp32: - t = ( - torch.arange(seqlen, device=device, dtype=torch.float32) - / self.scaling_factor - ) - inv_freq = ( - self.inv_freq.to(torch.float32) - if self.inv_freq.dtype != torch.float32 - else self.inv_freq - ) - else: - t = ( - torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) # type: ignore[call-overload] - / self.scaling_factor - ) - inv_freq = self.inv_freq - freqs = torch.outer(t, inv_freq) # type: ignore[arg-type] - - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - _scale: torch.Tensor = self.scale # type: ignore[assignment] - power = ( - torch.arange(seqlen, dtype=_scale.dtype, device=_scale.device) - - seqlen // 2 - ) / self.scale_base # type: ignore[operator] - scale = _scale.to(device=power.device) ** power.unsqueeze(-1) - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def _apply(self, fn, recurse=True): - if self.inv_freq.is_meta: - self.reset_parameters(device="cpu") - result = super()._apply(fn, recurse=recurse) - # Recompute inv_freq on the new device: CPU vs CUDA ``pow`` differ by - # ~1 fp32 ULP, which compounds across attention layers. Keep this - # buffer fp32 even when the module is cast to bf16/fp16; otherwise the - # rounded RoPE frequencies drift from the internal ESMC path. - new_inv_freq = self._compute_inv_freq(device=self.inv_freq.device) - self.register_buffer("inv_freq", new_inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - return result - - def forward( - self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0 - ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply RoPE to query and key tensors. - - Args: - q: ``(batch, seqlen, n_heads, head_dim)`` - k: ``(batch, seqlen, n_heads, head_dim)`` - seqlen_offset: Offset used in incremental decoding. - - Returns: - Tuple of rotated ``(q, k)`` tensors with the same shape as the inputs. - """ - self._update_cos_sin_cache( - q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype - ) - assert self._cos_cached is not None and self._sin_cached is not None - - if self.scale is not None: - raise NotImplementedError("XPos scaling is not supported in this path.") - - cos = self._cos_cached[seqlen_offset:] - sin = self._sin_cached[seqlen_offset:] - - if _flash_attn_rotary_available and q.device.type == "cuda": - q_rot = apply_triton_rotary(q, cos, sin, interleaved=self.interleaved) # type: ignore[misc] - k_rot = apply_triton_rotary(k, cos, sin, interleaved=self.interleaved) # type: ignore[misc] - else: - q_rot = _apply_rotary_emb_torch(q, cos, sin, self.interleaved) - k_rot = _apply_rotary_emb_torch(k, cos, sin, self.interleaved) - return q_rot, k_rot - - -class _TritonRotaryEmbedding(RotaryEmbedding): - """RoPE variant that delegates to the Flash-Attention Triton kernel. - - Only used inside :class:`_FlashMultiHeadAttention` when Flash Attention 2 - is available. The ``forward`` signature differs from :class:`RotaryEmbedding` - because Flash Attention packs Q, K, V together. - """ - - def forward( - self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int - ) -> torch.Tensor: # type: ignore[override] - """Apply RoPE in-place to a packed ``(N, 3, n_heads, head_dim)`` tensor.""" - self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) - assert self._cos_cached is not None and self._sin_cached is not None - assert apply_triton_rotary is not None - - apply_triton_rotary( - qkv[:, 0], - self._cos_cached, - self._sin_cached, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - inplace=True, - ) - apply_triton_rotary( - qkv[:, 1], - self._cos_cached, - self._sin_cached, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - inplace=True, - ) - return qkv - - -# --------------------------------------------------------------------------- -# Feed-forward network helpers -# --------------------------------------------------------------------------- - - -def _swiglu_hidden_dim(expansion_ratio: float, d_model: int) -> int: - """Round hidden dim to the nearest multiple of 256 after applying expansion_ratio.""" - return int(((expansion_ratio * d_model) + 255) // 256 * 256) - - -class _SwiGLU(nn.Module): - """SwiGLU activation: ``silu(x1) * x2`` where ``x`` is split along the last dim.""" - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return F.silu(x1) * x2 - - -class _PyTorchLayerNormLinear(nn.Module): - """LayerNorm followed by a Linear projection, sharing the parameter - names ``layer_norm_weight``, ``layer_norm_bias`` and ``weight`` so the - state-dict layout matches the accelerated TE module loaded on GPU. - """ - - def __init__(self, d_in: int, d_out: int, eps: float = 1e-5) -> None: - super().__init__() - self.d_in = d_in - self.eps = eps - self.layer_norm_weight = nn.Parameter(torch.ones(d_in)) - self.layer_norm_bias = nn.Parameter(torch.zeros(d_in)) - self.weight = nn.Parameter(torch.empty(d_out, d_in)) - nn.init.normal_(self.weight, std=0.02) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = F.layer_norm( - x, - (self.d_in,), - self.layer_norm_weight.to(dtype=x.dtype), - self.layer_norm_bias.to(dtype=x.dtype), - self.eps, - ) - x = x.to(dtype=self.weight.dtype) - return F.linear(x, self.weight) - - -class _PyTorchLayerNormMLP(nn.Module): - """LayerNorm + SwiGLU MLP, sharing the parameter names - ``layer_norm_weight``, ``layer_norm_bias``, ``fc1_weight``, - ``fc2_weight`` so the state-dict layout matches the accelerated TE - module loaded on GPU. - """ - - def __init__( - self, hidden_size: int, ffn_hidden_size: int, eps: float = 1e-5 - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.eps = eps - self.layer_norm_weight = nn.Parameter(torch.ones(hidden_size)) - self.layer_norm_bias = nn.Parameter(torch.zeros(hidden_size)) - self.fc1_weight = nn.Parameter(torch.empty(2 * ffn_hidden_size, hidden_size)) - self.fc2_weight = nn.Parameter(torch.empty(hidden_size, ffn_hidden_size)) - nn.init.normal_(self.fc1_weight, std=0.02) - nn.init.normal_(self.fc2_weight, std=0.02) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = F.layer_norm( - x, - (self.hidden_size,), - self.layer_norm_weight.to(dtype=x.dtype), - self.layer_norm_bias.to(dtype=x.dtype), - self.eps, - ) - x = x.to(dtype=self.fc1_weight.dtype) - x = F.linear(x, self.fc1_weight) - x1, x2 = x.chunk(2, dim=-1) - x = F.silu(x1) * x2 - x = x.to(dtype=self.fc2_weight.dtype) - return F.linear(x, self.fc2_weight) - - -def _swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module: - """LayerNorm + SwiGLU MLP. Uses Transformer Engine's fused LN+MLP when - available; otherwise returns the pure-PyTorch fallback with matching - state-dict layout.""" - assert not bias, "ESMC was trained with bias=False; bias=True not supported" - hidden = _swiglu_hidden_dim(expansion_ratio, d_model) - if _te_available: - return te.LayerNormMLP( # type: ignore[union-attr] - hidden_size=d_model, - ffn_hidden_size=hidden, - bias=bias, - activation="swiglu", - init_method=None, - output_layer_init_method=None, - ) - return _PyTorchLayerNormMLP(hidden_size=d_model, ffn_hidden_size=hidden) - - -def _make_attn_layernorm_qkv(d_model: int, bias: bool) -> nn.Module: - """LayerNorm + fused QKV projection. Uses Transformer Engine when - available; pure-PyTorch fallback otherwise.""" - assert not bias, "ESMC was trained with bias=False; bias=True not supported" - if _te_available: - return te.LayerNormLinear( # type: ignore[union-attr] - d_model, d_model * 3, bias=bias, init_method=None - ) - return _PyTorchLayerNormLinear(d_model, d_model * 3) - - -def _make_attn_out_proj(d_model: int, bias: bool) -> nn.Module: - """Attention output projection. Uses Transformer Engine when available; - pure-PyTorch ``nn.Linear`` otherwise.""" - if _te_available: - return te.Linear( # type: ignore[union-attr] - d_model, d_model, bias=bias, init_method=None - ) - return nn.Linear(d_model, d_model, bias=bias) - - -def _gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Sequential: - hidden = int(expansion_ratio * d_model) - return nn.Sequential( - nn.LayerNorm(d_model), - nn.Linear(d_model, hidden, bias=bias), - nn.GELU(), - nn.Linear(hidden, d_model, bias=bias), - ) - - -# --------------------------------------------------------------------------- -# Attention -# --------------------------------------------------------------------------- - - -def _scaled_dot_product_attention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - n_heads: int, - d_head: int, - seq_id: torch.Tensor | None, -) -> torch.Tensor: - """Scaled dot-product attention with optional chain-aware mask. - - Dispatches in order of preference: - 1. xformers ``memory_efficient_attention`` — preferred fused kernel, - requires ``xformers``, no chain mask. - 2. Flash Attention 2 (``flash_attn.flash_attn_func``) — secondary - fused kernel, requires ``flash-attn``, no chain mask, fp16 / - bf16 only. - 3. PyTorch's ``F.scaled_dot_product_attention`` — last-resort path; - also handles the chain-aware mask when ``seq_id`` is present - and the fp32 path that Flash Attention 2 does not support. - """ - if seq_id is None and _xformers_available: - b, s, _ = q.shape - q4 = q.view(b, s, n_heads, d_head) - k4 = k.view(b, s, n_heads, d_head) - v4 = v.view(b, s, n_heads, d_head) - context = xops.memory_efficient_attention( # type: ignore[union-attr] - q4, k4, v4, attn_bias=None, scale=d_head**-0.5 - ) - return context.reshape(b, s, n_heads * d_head) - if ( - seq_id is None - and _flash_attn_available - and q.dtype in (torch.float16, torch.bfloat16) - ): - b, s, _ = q.shape - q4 = q.view(b, s, n_heads, d_head) - k4 = k.view(b, s, n_heads, d_head) - v4 = v.view(b, s, n_heads, d_head) - context = flash_attn_func( # type: ignore[misc] - q4, k4, v4, dropout_p=0.0, softmax_scale=d_head**-0.5 - ) - return context.reshape(b, s, n_heads * d_head) # type: ignore[union-attr] - b, s, _ = q.shape - q = q.view(b, s, n_heads, -1).transpose(1, 2) - k = k.view(b, s, n_heads, -1).transpose(1, 2) - v = v.view(b, s, n_heads, -1).transpose(1, 2) - if seq_id is not None: - mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1) - context = F.scaled_dot_product_attention(q, k, v, mask) - else: - context = F.scaled_dot_product_attention(q, k, v) - _, h, _, d_out = context.shape - return context.transpose(1, 2).reshape(b, s, h * d_out) - - -class MultiHeadAttention(nn.Module): - """Multi-head self-attention with QK LayerNorm and RoPE. - - Args: - d_model: Model hidden dimension. - n_heads: Number of attention heads. - bias: Whether to use bias in linear layers. - qk_layernorm: Whether to apply LayerNorm to queries and keys before - computing attention scores. - """ - - def __init__( - self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True - ): - super().__init__() - self.d_model = d_model - self.n_heads = n_heads - self.d_head = d_model // n_heads - - assert not bias, "ESMC was trained with bias=False; bias=True not supported" - self.layernorm_qkv = _make_attn_layernorm_qkv(d_model, bias) - self.out_proj = _make_attn_out_proj(d_model, bias) - - if qk_layernorm: - self.q_ln = nn.LayerNorm(d_model, bias=bias) - self.k_ln = nn.LayerNorm(d_model, bias=bias) - else: - self.q_ln = nn.Identity() - self.k_ln = nn.Identity() - - self.rotary = RotaryEmbedding(d_model // n_heads) - - def _apply_rotary( - self, q: torch.Tensor, k: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - q = q.unflatten(-1, (self.n_heads, self.d_head)) - k = k.unflatten(-1, (self.n_heads, self.d_head)) - q, k = self.rotary(q, k) - q = q.flatten(-2, -1) - k = k.flatten(-2, -1) - return q, k - - def forward( - self, - x: torch.Tensor, - seq_id: torch.Tensor | None, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Return ``(context, attn_weights)``. - - ``attn_weights`` is ``None`` unless ``output_attentions=True`` — the - fused SDPA backends (xformers, flash-attn 2, ``F.scaled_dot_product_attention``) - don't expose attention probabilities, so capturing them forces a - materialized ``softmax(Q @ K.T / sqrt(d)) @ V`` path with shape - ``(B, H, L, L)``. - """ - qkv = self.layernorm_qkv(x) - q, k, v = torch.chunk(qkv, 3, dim=-1) - q = self.q_ln(q).to(q.dtype) - k = self.k_ln(k).to(q.dtype) - q, k = self._apply_rotary(q, k) - - b, s, _ = q.shape - - if output_attentions: - # Manual SDPA so attention probabilities are observable. - q4 = q.view(b, s, self.n_heads, self.d_head).transpose(1, 2) - k4 = k.view(b, s, self.n_heads, self.d_head).transpose(1, 2) - v4 = v.view(b, s, self.n_heads, self.d_head).transpose(1, 2) - scale = self.d_head**-0.5 - attn_scores = (q4 @ k4.transpose(-2, -1)) * scale - if seq_id is not None: - mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1) - attn_scores = attn_scores.masked_fill(~mask, float("-inf")) - attn_weights = torch.softmax(attn_scores, dim=-1) - context = (attn_weights @ v4).transpose(1, 2).reshape(b, s, -1) - return self.out_proj(context), attn_weights - - context = _scaled_dot_product_attention( - q, k, v, n_heads=self.n_heads, d_head=self.d_head, seq_id=seq_id - ) - return self.out_proj(context), None - - -class _FlashMultiHeadAttention(MultiHeadAttention): - """Flash-Attention 2 variant of :class:`MultiHeadAttention`.""" - - def __init__( - self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True - ): - super().__init__( - d_model=d_model, n_heads=n_heads, bias=bias, qk_layernorm=qk_layernorm - ) - self.rotary = _TritonRotaryEmbedding(d_model // n_heads) - - def forward( - self, - x: torch.Tensor, - seq_id: torch.Tensor | None, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - if output_attentions: - raise ValueError( - "output_attentions=True is not supported with " - "attn_implementation='flash_attention_2'. " - "Re-load the model with attn_implementation='sdpa' (or 'eager')." - ) - assert seq_id is not None and seq_id.dtype == torch.bool - - seqlens = seq_id.sum(dim=-1, dtype=torch.int32) - cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) - max_seqlen = int(seqlens.max().item()) - - qkv = self.layernorm_qkv(x) - q, k, v = torch.chunk(qkv, 3, dim=-1) - q = self.q_ln(q).to(q.dtype) - k = self.k_ln(k).to(q.dtype) - - # ``q``/``k``/``v`` are 2D ``(T, D)`` here: the parent ``ESMCModel.forward`` - # calls ``unpad_input`` before the transformer stack to produce the - # varlen-flat layout that ``flash_attn_varlen_qkvpacked_func`` requires. - T = q.shape[0] - qkv_packed = torch.stack([q, k, v], dim=1).view(T, 3, self.n_heads, self.d_head) - qkv_packed = self.rotary(qkv_packed, cu_seqlens, max_seqlen) - - context = flash_attn_varlen_qkvpacked_func( # type: ignore[misc] - qkv_packed, cu_seqlens, max_seqlen, softmax_scale=self.d_head**-0.5 - ) - n_out, h_out, d_out = context.shape # type: ignore[union-attr] - return ( - self.out_proj(context.reshape(n_out, h_out * d_out)), # type: ignore[union-attr] - None, - ) - - -# --------------------------------------------------------------------------- -# Transformer blocks -# --------------------------------------------------------------------------- - - -class UnifiedTransformerBlock(nn.Module): - """Single transformer block: pre-norm attention + pre-norm FFN with residual scaling. - - Args: - d_model: Hidden dimension. - n_heads: Number of attention heads. - use_flash_attn: Use Flash Attention 2 kernel if available. - bias: Whether linear layers include bias terms. - expansion_ratio: Hidden-dim expansion ratio for the FFN. - residue_scaling_factor: Scales residual connections to stabilise deep - networks (``1 / sqrt(n_layers / 36)`` is the ESM3 scheme). - qk_layernorm: Whether to apply QK LayerNorm in attention. - ffn_type: Feed-forward activation: ``"swiglu"`` or ``"gelu"``. - """ - - def __init__( - self, - d_model: int, - n_heads: int, - use_flash_attn: bool = False, - bias: bool = False, - expansion_ratio: float = 4.0, - residue_scaling_factor: float = 1.0, - qk_layernorm: bool = True, - ffn_type: str = "swiglu", - ): - super().__init__() - - attn_cls = _FlashMultiHeadAttention if use_flash_attn else MultiHeadAttention - self.attn = attn_cls(d_model, n_heads, bias=bias, qk_layernorm=qk_layernorm) - - if ffn_type == "swiglu": - self.ffn = _swiglu_ln_ffn(d_model, expansion_ratio, bias) - elif ffn_type == "gelu": - self.ffn = _gelu_ln_ffn(d_model, expansion_ratio, bias) - else: - raise ValueError( - f"Unknown ffn_type: {ffn_type!r}. Choose 'swiglu' or 'gelu'." - ) - - self.scaling_factor = residue_scaling_factor - - def forward( - self, - x: torch.Tensor, - sequence_id: torch.Tensor | None, - output_attentions: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - Args: - x: ``(batch, seq_len, d_model)`` - sequence_id: ``(batch, seq_len)`` chain-ID tensor used to restrict - attention to tokens within the same chain. SDPA blocks accept - an integer tensor (``-1`` marks padding); the flash-attn block - takes a ``bool`` padding mask — the caller selects which. - ``None`` skips chain-aware masking entirely (fast path). - output_attentions: When ``True``, returns the per-head attention - weights for this block alongside the residual output. - - Returns: - ``(output, attn_weights_or_None)``. Shape of ``output`` is - ``(batch, seq_len, d_model)``; ``attn_weights`` shape is - ``(batch, num_heads, seq_len, seq_len)`` or ``None``. - """ - attn_out, attn_weights = self.attn( - x, sequence_id, output_attentions=output_attentions - ) - x = x + attn_out / self.scaling_factor - x = x + self.ffn(x) / self.scaling_factor - return x, attn_weights - - -class TransformerStack(nn.Module): - """Stack of :class:`UnifiedTransformerBlock` layers with a final LayerNorm. - - Args: - d_model: Hidden dimension. - n_heads: Number of attention heads. - n_layers: Number of transformer blocks. - scale_residue: When ``True`` apply ESM3 residue scaling - ``sqrt(n_layers / 36)`` to each block. - bias: Bias flag forwarded to every sub-module. - qk_layernorm: QK LayerNorm flag forwarded to every block. - ffn_type: FFN activation type (``"swiglu"`` or ``"gelu"``). - expansion_ratio: FFN expansion ratio. - use_flash_attn: Use Flash Attention 2 kernel when available. - """ - - def __init__( - self, - d_model: int, - n_heads: int, - n_layers: int, - scale_residue: bool = True, - bias: bool = False, - qk_layernorm: bool = True, - ffn_type: str = "swiglu", - expansion_ratio: float = 8 / 3, - use_flash_attn: bool = False, - ): - super().__init__() - self.blocks = nn.ModuleList( - [ - UnifiedTransformerBlock( - d_model, - n_heads, - use_flash_attn=use_flash_attn, - residue_scaling_factor=math.sqrt(n_layers / 36) - if scale_residue - else 1.0, - expansion_ratio=expansion_ratio, - bias=bias, - qk_layernorm=qk_layernorm, - ffn_type=ffn_type, - ) - for _ in range(n_layers) - ] - ) - self.norm = nn.LayerNorm(d_model, bias=False) - - def forward( - self, - x: torch.Tensor, - sequence_id: torch.Tensor | None = None, - layers_to_collect: list[int] | None = None, - output_attentions: bool = False, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - tuple[torch.Tensor, ...], - tuple[torch.Tensor, ...] | None, - ]: - """Run the full transformer stack. - - Args: - x: ``(batch, seq_len, d_model)`` - sequence_id: Optional chain-id tensor forwarded to each block. - layers_to_collect: Layer indices (0-based pre-block inputs plus - ``n_layers`` for the post-norm output) whose hidden states - should be returned. - output_attentions: When ``True``, collects the per-block attention - weights and returns them as the fourth tuple element. - - Returns: - ``(post_norm, pre_norm, hidden_states, attentions)`` where - ``hidden_states`` is a (possibly empty) tuple of tensors and - ``attentions`` is a tuple of per-block ``(B, H, L, L)`` tensors - or ``None`` when ``output_attentions`` is ``False``. - """ - if layers_to_collect is None: - layers_to_collect = [] - - collected: list[torch.Tensor] = [] - all_attentions: list[torch.Tensor] = [] - for layer_idx, block in enumerate(self.blocks): - if layer_idx in layers_to_collect: - collected.append(x) - x, attn_weights = block(x, sequence_id, output_attentions=output_attentions) - if output_attentions and attn_weights is not None: - all_attentions.append(attn_weights) - - norm_x = self.norm(x) - if len(self.blocks) in layers_to_collect: - collected.append(norm_x) - - attentions = tuple(all_attentions) if output_attentions else None - return norm_x, x, tuple(collected), attentions - - -# --------------------------------------------------------------------------- -# Pre-trained model base class -# --------------------------------------------------------------------------- - - -@auto_docstring -class ESMCPreTrainedModel(PreTrainedModel): - """Base class for ESMC models. - - Handles weight initialisation and declares module-level capabilities. - """ - - config_class = ESMCConfig - base_model_prefix = "esmc" - supports_gradient_checkpointing = False - _supports_sdpa = True - _supports_flash_attn = True - _supports_attention_backend = True - _no_split_modules = ["UnifiedTransformerBlock"] - _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] - - def _init_weights(self, module: nn.Module): - std = self.config.initializer_range - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, RotaryEmbedding): - module.reset_parameters(device=self.device) - - -# --------------------------------------------------------------------------- -# Base encoder model -# --------------------------------------------------------------------------- - - -@auto_docstring -class ESMCModel(ESMCPreTrainedModel): - """The bare ESMC encoder outputting raw hidden states. - - ESMC is a protein language model trained by EvolutionaryScale using a - masked-token objective over amino acid sequences. The architecture is a - standard Transformer encoder with RoPE positional embeddings, QK LayerNorm, - and SwiGLU feed-forward networks. - - Args: - config: An :class:`ESMCConfig` instance. - """ - - def __init__(self, config: ESMCConfig): - super().__init__(config) - self._use_flash_attn = ( - _flash_attn_available and config._attn_implementation == "flash_attention_2" - ) - self.embed = nn.Embedding(config.vocab_size, config.d_model) - self.transformer = TransformerStack( - config.d_model, - config.n_heads, - config.n_layers, - use_flash_attn=self._use_flash_attn, - ) - self._sae_models: nn.ModuleDict = nn.ModuleDict() - self.post_init() - - def get_input_embeddings(self) -> nn.Embedding: - return self.embed - - def set_input_embeddings(self, value: nn.Embedding): - self.embed = value - - def add_sae_models(self, sae_models: list[_ESMCSAELayer]) -> None: - """Register one or more SAEs obtained from an :class:`ESMCSAEModel`. - - Each is keyed by ``f"layer{N}"`` (the backbone-layer index ``N`` the - SAE is trained against, set by - :meth:`ESMCSAEModel.initialize_layers`). Attaching two SAEs for the - same backbone layer raises — only one SAE per layer can be active. - - Example:: - - sae = ESMCSAEModel.from_pretrained( - "biohub/esmc-600m-2024-12-sae-k64-codebook16384" - ) - sae.initialize_layers([27, 33]) - model.add_sae_models([sae.layers["27"], sae.layers["33"]]) - """ - for layer in sae_models: - assert isinstance(layer, _ESMCSAELayer), ( - f"Expected an SAE layer (model.layers['']), got " - f"{type(layer).__name__}." - ) - key = f"layer{int(layer.layer)}" - if key in self._sae_models: - raise ValueError( - f"An SAE is already registered at {key!r}. Only one SAE " - "per backbone layer can be active — pick a different " - "layer on one of them, or attach in a fresh model." - ) - self._sae_models[key] = layer - - _SAE_KEY_RE = re.compile(r"layer(\d+)") - - def _get_sae_layer_num_requested(self, model_name: str) -> int: - """Recover the backbone-layer index from a key written by - :meth:`add_sae_models` (``"layer{N}"`` → ``N``).""" - match = self._SAE_KEY_RE.fullmatch(model_name) - assert ( - match is not None - ), f"Unexpected SAE key {model_name!r}; expected 'layer{{N}}'." - return int(match.group(1)) - - def _validate_sae_inputs(self, input_ids: torch.Tensor) -> None: - assert torch.all(input_ids != self.config.mask_token_id), ( - "SAE inputs must not contain mask tokens. " - "SAEs were trained on unmasked sequences." - ) - - def _get_sae_outputs( - self, - hidden_states: torch.Tensor, - layers_to_collect: list[int], - token_mask: torch.Tensor, - normalize_sae: bool = False, - ) -> dict[str, torch.Tensor]: - """Run all registered SAEs and return their feature magnitudes. - - Args: - hidden_states: Stacked tensor of shape - ``(len(layers_to_collect), batch, seq_len, d_model)``. - layers_to_collect: The ESMC layer indices that were collected, - in the same order as the first dim of ``hidden_states``. - token_mask: Boolean mask ``(batch, seq_len)`` — ``True`` for - real (non-padding) tokens. - normalize_sae: When ``True``, scale features by ``idf / max`` - using the per-feature stats trained alongside each SAE. - """ - layer_to_idx = {layer: idx for idx, layer in enumerate(layers_to_collect)} - sae_outputs: dict[str, torch.Tensor] = {} - - for model_name, sae_module in self._sae_models.items(): - # `nn.ModuleDict` only stores `nn.Module`s at the type level; - # ``add_sae_models`` enforces that each entry is an ``_ESMCSAELayer``. - assert isinstance(sae_module, _ESMCSAELayer) - layer: _ESMCSAELayer = sae_module - requested_layer = self._get_sae_layer_num_requested(model_name) - layer_idx = layer_to_idx[requested_layer] - layer_states = hidden_states[layer_idx].clone().to(self.device) - - sae_out = layer.get_sae_output(layer_states, token_mask) - features = sae_out.feature_magnitudes.detach() - - if normalize_sae: - # ``register_buffer`` is typed as ``Tensor | Module`` on - # ``nn.Module``; narrow here since these are Tensors. - idf = cast(torch.Tensor, layer.idf) - max_val = cast(torch.Tensor, layer.max) - features = (features / max_val) * idf - - sae_outputs[model_name] = features.to_sparse() - - return sae_outputs - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - return_dict: Optional[bool] = None, - compute_sae: bool = True, - normalize_sae: bool = False, - ) -> tuple[torch.Tensor, ...] | ESMCOutput: - r""" - sequence_id (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Integer chain-ID tensor for chain-aware attention masking. Tokens with the same - non-negative integer value can attend to each other; tokens with different values - cannot (cross-chain masking). Padding positions should be set to ``-1``. - When provided, ``attention_mask`` is ignored. The ``flash_attention_2`` backend - only supports single-chain inputs (all non-padding values must be ``0``); pass - multi-chain ``sequence_id`` with ``attn_implementation='sdpa'`` (or ``'eager'``). - output_attentions (`bool`, *optional*): - Whether to return the per-block attention weights of shape - ``(batch_size, num_heads, sequence_length, sequence_length)``. - Forces a manual-SDPA path inside :class:`MultiHeadAttention` so the - attention probabilities are observable; raises on the - ``flash_attention_2`` path. - compute_sae (`bool`, *optional*, defaults to ``True``): - Whether to run any SAE models registered via :meth:`add_sae_models`. - Has no effect when no SAEs are registered. - normalize_sae (`bool`, *optional*, defaults to ``False``): - When ``True``, scale SAE feature magnitudes by ``idf / max`` (only - applied when the SAE's normalization buffers contain non-trivial values). - - Examples: - - ```python - >>> from transformers import AutoTokenizer, ESMCModel - - >>> model = ESMCModel.from_pretrained("Biohub/ESMC-600M-2024-12") - >>> tokenizer = AutoTokenizer.from_pretrained("Biohub/ESMC-600M-2024-12") - >>> inputs = tokenizer(["MLKNVQVQLV"], return_tensors="pt") - >>> outputs = model(**inputs) - >>> outputs.last_hidden_state.shape - torch.Size([1, 12, 960]) - ``` - """ - output_hidden_states = ( - output_hidden_states - if output_hidden_states is not None - else self.config.output_hidden_states - ) - output_attentions = ( - output_attentions - if output_attentions is not None - else self.config.output_attentions - ) - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - output_sae = compute_sae and len(self._sae_models) > 0 - - # Determine which intermediate layers to collect. When SAEs are - # registered we must collect at least the layers they target, even if - # the caller did not ask for all hidden states. - if output_hidden_states: - layers_to_collect: list[int] = list(range(self.config.n_layers + 1)) - elif output_sae: - layers_to_collect = sorted( - {self._get_sae_layer_num_requested(name) for name in self._sae_models} - ) - else: - layers_to_collect = [] - - user_supplied_sequence_id = sequence_id is not None - if sequence_id is not None: - bool_mask = sequence_id >= 0 - else: - if attention_mask is None: - attention_mask = input_ids != self.config.pad_token_id - assert attention_mask is not None - bool_mask = attention_mask.bool() - sequence_id = bool_mask.to(torch.long) - 1 - - x = self.embed(input_ids) - b, l_ = x.shape[:2] - - if self._use_flash_attn: - if user_supplied_sequence_id and (sequence_id > 0).any(): - raise ValueError( - "Multi-chain ``sequence_id`` (any value > 0) is not " - "supported with attn_implementation='flash_attention_2'. " - "Re-load the model with attn_implementation='sdpa' (or " - "'eager') for chain-aware attention masking." - ) - assert unpad_input is not None - x, indices, *_ = unpad_input(x, bool_mask) - else: - indices = None - - if self._use_flash_attn: - trans_seq_id = bool_mask - elif user_supplied_sequence_id: - trans_seq_id = sequence_id - elif bool_mask.all() and not output_attentions: - # Fused SDPA fast path (xformers / flash) is correct only when the - # mask is uniform; output_attentions forces the manual branch. - trans_seq_id = None - else: - trans_seq_id = sequence_id - last_hidden_state, _, collected, attentions = self.transformer( - x, - sequence_id=trans_seq_id, - layers_to_collect=layers_to_collect, - output_attentions=output_attentions, - ) - - if self._use_flash_attn: - assert indices is not None and pad_input is not None - last_hidden_state = pad_input(last_hidden_state, indices, b, l_) - collected = [pad_input(h, indices, b, l_) for h in collected] - - # Stack once; reused for both SAE and hidden-state output. - collected_tensor: torch.Tensor | None = ( - torch.stack(collected, dim=0) if collected else None # type: ignore[arg-type] - ) - - sae_outputs: dict[str, torch.Tensor] | None = None - if output_sae and collected_tensor is not None: - assert input_ids is not None - self._validate_sae_inputs(input_ids) - sae_outputs = self._get_sae_outputs( - collected_tensor, layers_to_collect, bool_mask, normalize_sae - ) - - hidden_states_tensor = collected_tensor if output_hidden_states else None - - if not return_dict: - return tuple( - v - for v in [ - last_hidden_state, - hidden_states_tensor, - sae_outputs, - attentions, - ] - if v is not None - ) - - return ESMCOutput( - last_hidden_state=last_hidden_state, - hidden_states=hidden_states_tensor, - sae_outputs=sae_outputs, - attentions=attentions, - ) - - -# --------------------------------------------------------------------------- -# LM head -# --------------------------------------------------------------------------- - - -def _esmc_lm_head( - d_model: int, output_dim: int, hidden_dim: int | None = None -) -> nn.Sequential: - """Linear → GELU → LayerNorm → Linear projection head for masked LM.""" - hidden_dim = hidden_dim if hidden_dim is not None else d_model - return nn.Sequential( - nn.Linear(d_model, hidden_dim), - nn.GELU(), - nn.LayerNorm(hidden_dim), - nn.Linear(hidden_dim, output_dim), - ) - - -# --------------------------------------------------------------------------- -# Masked language model -# --------------------------------------------------------------------------- - - -@auto_docstring -class ESMCForMaskedLM(ESMCPreTrainedModel): - """ESMC with a masked language modelling head. - - This is the primary pre-training objective of ESMC. The LM head consists - of a single hidden layer with GELU activation followed by LayerNorm and a - linear projection to ``vocab_size``. - """ - - def __init__(self, config: ESMCConfig): - super().__init__(config) - self.esmc = ESMCModel(config) - self.lm_head = _esmc_lm_head(config.d_model, config.vocab_size) - self.post_init() - - def get_output_embeddings(self) -> nn.Linear: - return self.lm_head[-1] # type: ignore[return-value] - - def set_output_embeddings(self, new_embeddings: nn.Linear): - self.lm_head[-1] = new_embeddings - - def add_sae_models(self, sae_models: list[_ESMCSAELayer]) -> None: - """Proxy to :meth:`ESMCModel.add_sae_models`.""" - self.esmc.add_sae_models(sae_models) - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - return_dict: Optional[bool] = None, - labels: Optional[torch.Tensor] = None, - compute_sae: bool = True, - normalize_sae: bool = False, - ) -> tuple[torch.Tensor, ...] | ESMCMaskedLMOutput: - r""" - sequence_id (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Integer chain-ID tensor forwarded to the encoder for chain-aware - attention masking. See :meth:`ESMCModel.forward` for the encoding. - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for masked language modelling loss. Positions with label ``-100`` - are ignored. Other positions must be in ``[0, config.vocab_size)``. - output_attentions (`bool`, *optional*): - Whether to return per-block attention weights. Forwarded to the - backbone; raises on the ``flash_attention_2`` path. - compute_sae (`bool`, *optional*, defaults to ``True``): - Whether to run registered SAE models. Has no effect when none are registered. - normalize_sae (`bool`, *optional*, defaults to ``False``): - When ``True``, scale SAE features by ``idf / max`` normalization buffers. - - Examples: - - ```python - >>> from transformers import AutoTokenizer, ESMCForMaskedLM - >>> import torch - - >>> model = ESMCForMaskedLM.from_pretrained("Biohub/ESMC-600M-2024-12") - >>> tokenizer = AutoTokenizer.from_pretrained("Biohub/ESMC-600M-2024-12") - >>> inputs = tokenizer(["MLKNVQLV"], return_tensors="pt") - >>> outputs = model(**inputs) - >>> outputs.logits.shape - torch.Size([1, 11, 64]) - ``` - """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - encoder_outputs = self.esmc( - input_ids=input_ids, - attention_mask=attention_mask, - sequence_id=sequence_id, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - return_dict=True, - compute_sae=compute_sae, - normalize_sae=normalize_sae, - ) - - logits = self.lm_head(encoder_outputs.last_hidden_state) - - loss: torch.Tensor | None = None - if labels is not None: - loss = CrossEntropyLoss(ignore_index=-100)( - logits.view(-1, self.config.vocab_size), labels.view(-1) - ) - - if not return_dict: - return tuple( - v - for v in [ - loss, - logits, - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - encoder_outputs.sae_outputs, - encoder_outputs.attentions, - ] - if v is not None - ) - - return ESMCMaskedLMOutput( - loss=loss, - logits=logits, - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - sae_outputs=encoder_outputs.sae_outputs, - attentions=encoder_outputs.attentions, - ) - - -# --------------------------------------------------------------------------- -# Classification heads -# --------------------------------------------------------------------------- - - -class _ESMCClassificationHead(nn.Module): - """Dense classification head applied to the ```` token representation.""" - - def __init__(self, config: ESMCConfig): - super().__init__() - self.dense = nn.Linear(config.d_model, config.d_model) - self.dropout = nn.Dropout(config.classifier_dropout) - self.out_proj = nn.Linear(config.d_model, config.num_labels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - x = hidden_states[:, 0, :] # token - x = self.dropout(x) - x = torch.tanh(self.dense(x)) - x = self.dropout(x) - return self.out_proj(x) - - -# --------------------------------------------------------------------------- -# Sequence classification -# --------------------------------------------------------------------------- - - -@auto_docstring -class ESMCForSequenceClassification(ESMCPreTrainedModel): - """ESMC with a sequence-level classification head. - - A linear layer is applied to the ```` token representation. - Supports regression (``num_labels == 1``), single-label classification, - and multi-label classification. - """ - - def __init__(self, config: ESMCConfig): - super().__init__(config) - self.num_labels = config.num_labels - self.esmc = ESMCModel(config) - self.classifier = _ESMCClassificationHead(config) - self.post_init() - - def add_sae_models(self, sae_models: list[_ESMCSAELayer]) -> None: - """Proxy to :meth:`ESMCModel.add_sae_models`.""" - self.esmc.add_sae_models(sae_models) - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - return_dict: Optional[bool] = None, - labels: Optional[torch.Tensor] = None, - compute_sae: bool = True, - normalize_sae: bool = False, - ) -> tuple[torch.Tensor, ...] | ESMCSequenceClassifierOutput: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for sequence classification loss. Indices must be in - ``[0, config.num_labels - 1]``. For regression pass a float - tensor of shape ``(batch_size,)``. - output_attentions (`bool`, *optional*): - Whether to return per-block attention weights. Forwarded to the - backbone; raises on the ``flash_attention_2`` path. - compute_sae (`bool`, *optional*, defaults to ``True``): - Whether to run registered SAE models. Has no effect when none are registered. - normalize_sae (`bool`, *optional*, defaults to ``False``): - When ``True``, scale SAE features by ``idf / max`` normalization buffers. - """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - encoder_outputs = self.esmc( - input_ids, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - return_dict=True, - compute_sae=compute_sae, - normalize_sae=normalize_sae, - ) - logits = self.classifier(encoder_outputs.last_hidden_state) - - loss: torch.Tensor | None = None - if labels is not None: - labels = labels.to(logits.device) - - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - loss = loss_fct( - logits.squeeze() if self.num_labels == 1 else logits, - labels.squeeze() if self.num_labels == 1 else labels, - ) - elif self.config.problem_type == "single_label_classification": - loss = CrossEntropyLoss()( - logits.view(-1, self.num_labels), labels.view(-1) - ) - elif self.config.problem_type == "multi_label_classification": - loss = BCEWithLogitsLoss()(logits, labels) - - if not return_dict: - return tuple( - v - for v in [ - loss, - logits, - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - encoder_outputs.sae_outputs, - encoder_outputs.attentions, - ] - if v is not None - ) - - return ESMCSequenceClassifierOutput( - loss=loss, - logits=logits, - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - sae_outputs=encoder_outputs.sae_outputs, - attentions=encoder_outputs.attentions, - ) - - -# --------------------------------------------------------------------------- -# Token classification -# --------------------------------------------------------------------------- - - -@auto_docstring -class ESMCForTokenClassification(ESMCPreTrainedModel): - """ESMC with a per-token classification head. - - Useful for tasks such as secondary structure prediction, contact-map - prediction, or per-residue labelling. - """ - - def __init__(self, config: ESMCConfig): - super().__init__(config) - self.num_labels = config.num_labels - self.esmc = ESMCModel(config) - self.dropout = nn.Dropout(config.classifier_dropout) - self.classifier = nn.Linear(config.d_model, config.num_labels) - self.post_init() - - def add_sae_models(self, sae_models: list[_ESMCSAELayer]) -> None: - """Proxy to :meth:`ESMCModel.add_sae_models`.""" - self.esmc.add_sae_models(sae_models) - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - output_hidden_states: Optional[bool] = None, - output_attentions: Optional[bool] = None, - return_dict: Optional[bool] = None, - labels: Optional[torch.Tensor] = None, - compute_sae: bool = True, - normalize_sae: bool = False, - ) -> tuple[torch.Tensor, ...] | ESMCTokenClassifierOutput: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Per-token labels. Indices must be in ``[0, config.num_labels - 1]``. - Positions with index ``-100`` are ignored in the loss. - output_attentions (`bool`, *optional*): - Whether to return per-block attention weights. Forwarded to the - backbone; raises on the ``flash_attention_2`` path. - compute_sae (`bool`, *optional*, defaults to ``True``): - Whether to run registered SAE models. Has no effect when none are registered. - normalize_sae (`bool`, *optional*, defaults to ``False``): - When ``True``, scale SAE features by ``idf / max`` normalization buffers. - """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - encoder_outputs = self.esmc( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - return_dict=True, - compute_sae=compute_sae, - normalize_sae=normalize_sae, - ) - - sequence_output = self.dropout(encoder_outputs.last_hidden_state) - logits = self.classifier(sequence_output) - - loss: torch.Tensor | None = None - if labels is not None: - loss = CrossEntropyLoss(ignore_index=-100)( - logits.view(-1, self.num_labels), labels.to(logits.device).view(-1) - ) - - if not return_dict: - return tuple( - v - for v in [ - loss, - logits, - encoder_outputs.last_hidden_state, - encoder_outputs.hidden_states, - encoder_outputs.sae_outputs, - encoder_outputs.attentions, - ] - if v is not None - ) - - return ESMCTokenClassifierOutput( - loss=loss, - logits=logits, - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - sae_outputs=encoder_outputs.sae_outputs, - attentions=encoder_outputs.attentions, - ) - - -__all__ = [ - "ESMCModel", - "ESMCForMaskedLM", - "ESMCForSequenceClassification", - "ESMCForTokenClassification", - "ESMCPreTrainedModel", -] diff --git a/fastplms/esmfold2/modeling_esmc_sae.py b/fastplms/esmfold2/modeling_esmc_sae.py deleted file mode 100644 index 69f9363..0000000 --- a/fastplms/esmfold2/modeling_esmc_sae.py +++ /dev/null @@ -1,363 +0,0 @@ -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""PyTorch ESMC SAE (Sparse Autoencoder) model. - -* :class:`ESMCSAEModel` — the published HF container, one repo per - ``(backbone, codebook_dim, k)`` group. Each backbone layer ships as a - ``layer_{i}.safetensors`` shard; ``from_pretrained`` downloads the whole - snapshot but loads no weights — callers materialize the layers they need - via :meth:`initialize_layers`. Single-layer repos auto-load so bare - ``forward(x)`` works. -* :class:`_ESMCSAELayer` — internal ``nn.Module`` that holds the weights for - one ``(backbone, codebook_dim, k, layer)`` SAE. Not a published HF artifact; - obtained only via ``model.layers[""]``. -""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from pathlib import Path -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F -from safetensors.torch import load_file, save_file - -from transformers.modeling_outputs import ModelOutput -from transformers.modeling_utils import PreTrainedModel -from transformers.utils import auto_docstring -from .configuration_esmc_sae import ESMCSAEConfig, ESMCSAEParams - - -@dataclass -@auto_docstring( - custom_intro=""" - Output type of [`ESMCSAEModel`]. - """ -) -class ESMCSAEOutput(ModelOutput): - feature_magnitudes: torch.Tensor - reconstruction_loss: Optional[torch.Tensor] = None - - def to_sparse(self) -> None: - self.feature_magnitudes = self.feature_magnitudes.to_sparse() - - -class _ESMCSAELayer(nn.Module): - """One backbone layer's SAE — internal building block of :class:`ESMCSAEModel`. - - Not exposed via ``AutoModel`` and not loadable on its own. Obtain one - via ``model.layers[""]`` after calling ``initialize_layers``. - """ - - def __init__(self, params: ESMCSAEParams): - super().__init__() - self.params = params - - self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim)) - self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model)) - self.b_dec = nn.Parameter(torch.zeros(params.d_model)) - # Per-feature normalization stats. Trained alongside the SAE for some - # variants; for variants that don't ship them, leaving these as ones - # makes ``_get_sae_outputs``'s ``features / max * idf`` a no-op. - self.register_buffer("idf", torch.ones(params.codebook_dim)) - self.register_buffer("max", torch.ones(params.codebook_dim)) - - @property - def layer(self) -> int: - """Backbone-layer index this SAE is trained against.""" - return self.params.layer - - def forward(self, x: torch.Tensor, **_kwargs: object) -> ESMCSAEOutput: - del _kwargs - x = self._zscore_normalize_representation(x) - - x_with_pre_encoder_bias = x - self.b_dec - preactivations = F.relu(x_with_pre_encoder_bias @ self.W_enc) - - topk = torch.topk(preactivations, self.params.k, dim=-1) - feature_magnitudes = torch.zeros_like(preactivations).scatter( - -1, topk.indices, topk.values - ) - - reconstructed = feature_magnitudes @ self.W_dec + self.b_dec - - reconstruction_loss = (reconstructed - x).pow(2).mean(dim=-1) - - return ESMCSAEOutput( - feature_magnitudes=feature_magnitudes, - reconstruction_loss=reconstruction_loss, - ) - - def get_sae_output( - self, layer_states: torch.Tensor, token_mask: torch.Tensor - ) -> ESMCSAEOutput: - _, _, v_len = layer_states.shape - nonpad_states = layer_states[token_mask].view(-1, v_len) - return self(nonpad_states) - - def _zscore_normalize_representation(self, x: torch.Tensor) -> torch.Tensor: - x_mean = x.mean(dim=-1, keepdim=True) - x = x - x_mean - x_std = x.std(dim=-1, keepdim=True) - return x / (x_std + 1e-5) - - -@auto_docstring -class ESMCSAEPreTrainedModel(PreTrainedModel): - config_class = ESMCSAEConfig - base_model_prefix = "esmc_sae" - - -@auto_docstring( - custom_intro=""" - HF container holding one SAE per backbone layer, all sharing the same - ``(d_model, codebook_dim, k)``. - - ``from_pretrained`` downloads the entire repo (every ``layer_{i}.safetensors``) - into the local HF cache but does **not** load any weights into memory. - Callers materialize the layers they actually need by calling - :meth:`initialize_layers`. The full set is available on disk after the - first call, so subsequent layer switches read from the local cache without - re-downloading. - - Examples:: - - model = ESMCSAEModel.from_pretrained( - "biohub/esmc-6b-2024-12-sae-k64-codebook16384" - ) - model.initialize_layers([60]) # ~2.5 GB into memory - out = model(layer_states, layer=60) # forward through layer 60 - model.initialize_layers([45]) # add layer 45 (cached locally) - model.release_layer(60) # free layer 60 - """ -) -class ESMCSAEModel(ESMCSAEPreTrainedModel): - def __init__(self, config: ESMCSAEConfig): - super().__init__(config) - # Layers are populated lazily by ``initialize_layers``; the container - # starts empty so ``from_pretrained`` doesn't materialize hundreds of - # GB of unused parameters. - self.layers = nn.ModuleDict() - # Zero-element buffer that rides along with ``.to(device/dtype)``. - # ``initialize_layers`` reads its current device/dtype so SAEs added - # after ``model.to("cuda")`` land on CUDA without re-passing ``device=``. - self.register_buffer("_device_marker", torch.empty(0), persistent=False) - self._snapshot_dir: Optional[str] = None - self.post_init() - - @classmethod - def from_pretrained( # type: ignore[override] - cls, pretrained_model_name_or_path: str | os.PathLike, *model_args, **kwargs - ) -> "ESMCSAEModel": - """Download (or reuse cached) the full repo and return the model. - - By default no weights are read into memory and the caller must invoke - :meth:`initialize_layers` before running :meth:`forward`. The single - exception is when the repo ships exactly one layer: that layer is - auto-loaded (honoring ``torch_dtype`` / ``device`` if passed) so the - bare ``forward(x)`` call just works. - - Honored kwargs: ``revision``, ``cache_dir``, ``token``, - ``allow_patterns``, ``local_files_only``, ``force_download`` (forwarded - to ``snapshot_download``); ``torch_dtype`` and ``device`` (used by the - single-layer auto-load path; otherwise pass them to - :meth:`initialize_layers`). Behavioral kwargs that imply work we do - not perform (``device_map``, ``low_cpu_mem_usage``, - ``quantization_config``, ``attn_implementation``) raise so the user - isn't silently misled. Other HF housekeeping kwargs (``config``, - ``trust_remote_code``, ``adapter_kwargs``, …) are accepted and - ignored — they only matter for the standard loader, which we bypass. - """ - del model_args - torch_dtype = kwargs.pop("torch_dtype", None) - device = kwargs.pop("device", None) - local_dir = _resolve_snapshot_dir(pretrained_model_name_or_path, kwargs) - unsupported = { - "device_map", - "low_cpu_mem_usage", - "quantization_config", - "attn_implementation", - "max_memory", - "offload_folder", - "offload_state_dict", - } & kwargs.keys() - if unsupported: - raise TypeError( - f"Unsupported kwargs to ESMCSAEModel.from_pretrained: " - f"{sorted(unsupported)}. The standard HF loader is bypassed —" - " call initialize_layers(..., device=, dtype=) instead." - ) - config = ESMCSAEConfig.from_pretrained(local_dir) - model = cls(config) - model._snapshot_dir = str(local_dir) - if device is not None: - model.to(device) - if torch_dtype is not None: - model.to(torch_dtype) - if len(config.available_layers) == 1: - model.initialize_layers(list(config.available_layers)) - return model - - def initialize_layers( - self, - layers: list[int], - *, - device: torch.device | str | None = None, - dtype: torch.dtype | None = None, - ) -> None: - """Load the requested layers from the local snapshot into memory. - - Layers already present in :attr:`self.layers` are skipped — calling - ``initialize_layers([23])`` twice is idempotent. ``device`` / ``dtype`` - default to wherever the model itself lives (via the ``_device_marker`` - buffer that moves with ``.to(...)``), so the common pattern of - ``model.to("cuda"); model.initialize_layers([7])`` Just Works. - """ - assert self._snapshot_dir is not None, ( - "ESMCSAEModel has no snapshot directory — call " - "from_pretrained first, or set _snapshot_dir manually." - ) - if device is None: - device = self._device_marker.device - if dtype is None: - dtype = self._device_marker.dtype - snapshot_dir = Path(self._snapshot_dir) - available = set(self.config.available_layers) - for layer_idx in layers: - key = str(layer_idx) - if key in self.layers: - continue - if layer_idx not in available: - raise KeyError( - f"Layer {layer_idx} is not in this repo. " - f"available_layers={sorted(available)}" - ) - shard = snapshot_dir / f"layer_{layer_idx}.safetensors" - if not shard.exists(): - raise FileNotFoundError( - f"Missing layer file {shard} — config lists layer " - f"{layer_idx} as available but the shard is not on disk." - ) - params = ESMCSAEParams( - d_model=self.config.d_model, - codebook_dim=self.config.codebook_dim, - k=self.config.k, - layer=layer_idx, - ) - # Build on the meta device so we don't allocate weights that - # ``load_state_dict`` would immediately overwrite. - with torch.device("meta"): - layer = _ESMCSAELayer(params) - layer.to_empty(device=device) - layer.load_state_dict(load_file(str(shard))) - layer.to(dtype=dtype) - self.layers[key] = layer - - def release_layer(self, layer: int) -> None: - """Drop the named layer from memory. No-op if not loaded.""" - key = str(layer) - if key in self.layers: - del self.layers[key] - - def loaded_layers(self) -> list[int]: - """Sorted list of layer indices currently materialized in memory.""" - return sorted(int(k) for k in self.layers.keys()) - - def forward( - self, x: torch.Tensor, layer: int | None = None, **kwargs: object - ) -> ESMCSAEOutput: - if layer is None: - if len(self.layers) == 1: - # Unambiguous: exactly one layer loaded → use it. - ((_only_key, only_layer),) = self.layers.items() - return only_layer(x, **kwargs) - if len(self.layers) == 0: - raise RuntimeError( - "No layers loaded — call " - f"initialize_layers([...]) first. " - f"available_layers={self.config.available_layers}" - ) - raise RuntimeError( - "Multiple layers are loaded — please select one via " - f"forward(x, layer=). Loaded layers: {self.loaded_layers()}" - ) - key = str(layer) - if key not in self.layers: - raise KeyError( - f"Layer {layer} is not loaded. Call " - f"initialize_layers([{layer}]) first. Loaded layers: " - f"{self.loaded_layers()}" - ) - return self.layers[key](x, **kwargs) - - def save_pretrained( # type: ignore[override] - self, save_directory: str | os.PathLike, *args, **kwargs - ) -> None: - """Write ``config.json`` plus one ``layer_{i}.safetensors`` per loaded layer. - - Only layers currently in :attr:`self.layers` are written. - ``available_layers`` in the saved config is synced to what's actually - on disk so a ``release_layer`` + ``save_pretrained`` round-trip never - advertises a layer whose shard is missing. - """ - del args, kwargs - save_directory = Path(save_directory) - save_directory.mkdir(parents=True, exist_ok=True) - # Sync available_layers to what we're about to write — never advertise - # a layer that isn't on disk in this repo. - self.config.available_layers = self.loaded_layers() - self.config.save_pretrained(str(save_directory)) - for key, layer in self.layers.items(): - shard = save_directory / f"layer_{key}.safetensors" - save_file( - { - k: v.detach().cpu().contiguous() - for k, v in layer.state_dict().items() - }, - str(shard), - ) - - -def _resolve_snapshot_dir( - pretrained_model_name_or_path: str | os.PathLike, kwargs: dict -) -> str: - """Local dir → return as-is; hub id → ``snapshot_download`` it. - - A directory only counts as "local" if it actually contains ``config.json``, - so a stale subdir named like a hub id (``./biohub/esmc-...``) - doesn't accidentally shadow the hub fetch. - - Pops the standard ``snapshot_download`` keyword args from ``kwargs`` so - callers can forward them via ``from_pretrained``. - """ - path = Path(pretrained_model_name_or_path) - if path.is_dir() and (path / "config.json").exists(): - return str(path) - from huggingface_hub import snapshot_download - - return snapshot_download( - repo_id=str(pretrained_model_name_or_path), - revision=kwargs.pop("revision", None), - cache_dir=kwargs.pop("cache_dir", None), - token=kwargs.pop("token", None), - allow_patterns=kwargs.pop("allow_patterns", None), - local_files_only=kwargs.pop("local_files_only", False), - force_download=kwargs.pop("force_download", False), - ) - - -__all__ = ["ESMCSAEModel", "ESMCSAEOutput", "ESMCSAEPreTrainedModel"] diff --git a/fastplms/esmfold2/protein_utils.py b/fastplms/esmfold2/protein_utils.py deleted file mode 100644 index 75b785d..0000000 --- a/fastplms/esmfold2/protein_utils.py +++ /dev/null @@ -1,488 +0,0 @@ -# coding=utf-8 -# Copyright 2026 Biohub. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Self-contained protein featurization for ESMFold2 inference. - -Lets ``ESMFold2ExperimentalModel.infer_protein_as_pdb`` fold a protein sequence -ESMFold-style without the ``esm`` companion package. The featurization -mirrors ``ESMFold2InputBuilder.prepare_input`` for the protein-only path — -``test_prepare_protein_features.py`` enforces tensor-exact parity. -""" - -from __future__ import annotations - -import math - -import torch -from torch import Tensor - -MOL_TYPE_PROTEIN = 0 -PROTEIN_UNK_RES_TYPE = 22 -MSA_GAP_TOKEN_ID = 1 - -PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = { - "ALA": 2, - "ARG": 3, - "ASN": 4, - "ASP": 5, - "CYS": 6, - "GLN": 7, - "GLU": 8, - "GLY": 9, - "HIS": 10, - "ILE": 11, - "LEU": 12, - "LYS": 13, - "MET": 14, - "PHE": 15, - "PRO": 16, - "SER": 17, - "THR": 18, - "TRP": 19, - "TYR": 20, - "VAL": 21, -} - -PROTEIN_1TO3: dict[str, str] = { - "A": "ALA", - "R": "ARG", - "N": "ASN", - "D": "ASP", - "C": "CYS", - "Q": "GLN", - "E": "GLU", - "G": "GLY", - "H": "HIS", - "I": "ILE", - "L": "LEU", - "K": "LYS", - "M": "MET", - "F": "PHE", - "P": "PRO", - "S": "SER", - "T": "THR", - "W": "TRP", - "Y": "TYR", - "V": "VAL", - "X": "UNK", -} - -ESM_PROTEIN_VOCAB: dict[str, int] = { - "L": 4, - "A": 5, - "G": 6, - "V": 7, - "S": 8, - "E": 9, - "R": 10, - "T": 11, - "I": 12, - "D": 13, - "P": 14, - "K": 15, - "Q": 16, - "N": 17, - "F": 18, - "Y": 19, - "M": 20, - "H": 21, - "W": 22, - "C": 23, - "X": 3, -} - -# Heavy atoms per canonical residue, in training-time order. -PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = { - "ALA": ["N", "CA", "C", "O", "CB"], - "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"], - "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"], - "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"], - "CYS": ["N", "CA", "C", "O", "CB", "SG"], - "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"], - "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"], - "GLY": ["N", "CA", "C", "O"], - "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"], - "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"], - "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"], - "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"], - "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"], - "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"], - "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"], - "SER": ["N", "CA", "C", "O", "CB", "OG"], - "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"], - "TRP": [ - "N", - "CA", - "C", - "O", - "CB", - "CG", - "CD1", - "CD2", - "NE1", - "CE2", - "CE3", - "CZ2", - "CZ3", - "CH2", - ], - "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"], - "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"], - "UNK": ["N", "CA", "C", "O"], -} - -PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = { - "ALA": { - "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452), - "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918), - "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088), - "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582), - "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896), - }, - "ARG": { - "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555), - "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113), - "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917), - "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399), - "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575), - "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137), - "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481), - "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593), - "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172), - "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516), - "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767), - }, - "ASN": { - "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759), - "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145), - "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467), - "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498), - "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526), - "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335), - "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128), - "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051), - }, - "ASP": { - "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178), - "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976), - "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495), - "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676), - "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351), - "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941), - "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447), - "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614), - }, - "CYS": { - "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378), - "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164), - "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977), - "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664), - "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387), - "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649), - }, - "GLN": { - "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079), - "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778), - "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646), - "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664), - "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865), - "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648), - "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318), - "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527), - "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538), - }, - "GLU": { - "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136), - "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053), - "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486), - "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864), - "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338), - "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481), - "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136), - "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355), - "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173), - }, - "GLY": { - "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463), - "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014), - "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134), - "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525), - }, - "HIS": { - "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072), - "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936), - "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935), - "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316), - "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062), - "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286), - "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016), - "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453), - "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805), - "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204), - }, - "ILE": { - "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667), - "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499), - "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528), - "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855), - "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409), - "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797), - "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326), - "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125), - }, - "LEU": { - "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443), - "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385), - "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608), - "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547), - "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259), - "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605), - "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489), - "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017), - }, - "LYS": { - "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242), - "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756), - "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961), - "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573), - "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517), - "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443), - "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504), - "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164), - "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088), - }, - "MET": { - "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634), - "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499), - "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433), - "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872), - "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532), - "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378), - "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546), - "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728), - }, - "PHE": { - "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116), - "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363), - "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713), - "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241), - "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157), - "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873), - "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768), - "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914), - "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812), - "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617), - "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575), - }, - "PRO": { - "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674), - "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277), - "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273), - "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332), - "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146), - "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489), - "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918), - }, - "SER": { - "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464), - "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794), - "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995), - "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707), - "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838), - "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003), - }, - "THR": { - "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494), - "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415), - "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276), - "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509), - "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815), - "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944), - "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266), - }, - "TRP": { - "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244), - "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307), - "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397), - "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029), - "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117), - "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055), - "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062), - "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285), - "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188), - "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866), - "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853), - "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229), - "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332), - "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151), - }, - "TYR": { - "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505), - "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751), - "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458), - "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212), - "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438), - "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647), - "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119), - "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414), - "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867), - "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026), - "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173), - "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002), - }, - "VAL": { - "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747), - "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094), - "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547), - "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899), - "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792), - "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943), - "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565), - }, - "UNK": { - "N": (0.0, 0.0, 0.0), - "CA": (0.0, 0.0, 0.0), - "C": (0.0, 0.0, 0.0), - "O": (0.0, 0.0, 0.0), - }, -} - -# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the -# opensource constants for the protein subset). -PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = { - ("LYS", "NZ"): 1, - ("ARG", "NH2"): 1, - ("HIS", "ND1"): 1, -} - -# Only the elements that appear in canonical protein heavy atoms. -_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16} - - -def _encode_atom_name(name: str) -> list[int]: - padded = name.ljust(4)[:4] - return [ord(c) - 32 if c != " " else 0 for c in padded] - - -def prepare_protein_features(sequence: str) -> dict[str, Tensor]: - """Featurize a single protein sequence for ESMFold2ExperimentalModel.forward. - - Returns the same keys with the same dtypes/shapes as - ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))`` - restricted to a single-chain protein with no MSA, modifications, - distogram conditioning, or covalent bonds. All tensors have a - leading batch dim of 1; the caller is responsible for moving them - to the model device. - """ - if not sequence: - raise ValueError("sequence must be non-empty") - - res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence] - L = len(sequence) - - token_atom_starts: list[int] = [] - atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = [] - res_type_vals: list[int] = [] - input_id_vals: list[int] = [] - distogram_rep_atom_idx: list[int] = [] - - atom_cursor = 0 - for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)): - atom_names = PROTEIN_HEAVY_ATOMS[res_3] - res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE) - input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"]) - - token_atom_starts.append(atom_cursor) - for name in atom_names: - charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0) - element = name[0] # protein heavy atoms are all single-letter C/N/O/S - ref_pos = PROTEIN_REF_POS[res_3][name] - atom_records.append((t_idx, name, element, charge, ref_pos)) - atom_cursor += 1 - - rep_name = "CB" if "CB" in atom_names else "CA" - distogram_rep_atom_idx.append( - token_atom_starts[t_idx] + atom_names.index(rep_name) - ) - - res_type_vals.append(res_type) - input_id_vals.append(input_id) - - n_real_atoms = len(atom_records) - n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32 - - ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32) - ref_element = torch.zeros(n_atoms, dtype=torch.int64) - ref_charge = torch.zeros(n_atoms, dtype=torch.int8) - ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64) - ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64) - atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool) - atom_to_token = torch.zeros(n_atoms, dtype=torch.int64) - - for i, (t_idx, name, element, charge, pos) in enumerate(atom_records): - ref_pos[i] = torch.tensor(pos, dtype=torch.float32) - ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element] - ref_charge[i] = charge - ref_atom_name_chars[i] = torch.tensor( - _encode_atom_name(name), dtype=torch.int64 - ) - ref_space_uid[i] = t_idx - atom_attention_mask[i] = True - atom_to_token[i] = t_idx - - token_index = torch.arange(L, dtype=torch.int64) - residue_index = torch.arange(L, dtype=torch.int64) - asym_id = torch.zeros(L, dtype=torch.int64) - sym_id = torch.zeros(L, dtype=torch.int64) - entity_id = torch.ones(L, dtype=torch.int64) - mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64) - res_type = torch.tensor(res_type_vals, dtype=torch.int64) - input_ids = torch.tensor(input_id_vals, dtype=torch.int64) - token_bonds = torch.zeros(L, L, 1, dtype=torch.float32) - token_attention_mask = torch.ones(L, dtype=torch.bool) - distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64) - - # Single-sequence MSA: depth 1, row 0 is the sequence itself. - msa = res_type.unsqueeze(0) - msa_attention_mask = torch.ones(1, L, dtype=torch.bool) - has_deletion = torch.zeros(1, L, dtype=torch.bool) - deletion_value = torch.zeros(1, L, dtype=torch.float32) - deletion_mean = torch.zeros(L, dtype=torch.float32) - - features = { - "token_index": token_index, - "residue_index": residue_index, - "asym_id": asym_id, - "sym_id": sym_id, - "entity_id": entity_id, - "mol_type": mol_type, - "res_type": res_type, - "input_ids": input_ids, - "token_bonds": token_bonds, - "token_attention_mask": token_attention_mask, - "ref_pos": ref_pos, - "ref_element": ref_element, - "ref_charge": ref_charge, - "ref_atom_name_chars": ref_atom_name_chars, - "ref_space_uid": ref_space_uid, - "atom_attention_mask": atom_attention_mask, - "atom_to_token": atom_to_token, - "distogram_atom_idx": distogram_atom_idx, - "msa": msa, - "msa_attention_mask": msa_attention_mask, - "has_deletion": has_deletion, - "deletion_value": deletion_value, - "deletion_mean": deletion_mean, - } - return {k: v.unsqueeze(0) for k, v in features.items()} diff --git a/fastplms/fine_tuning_example.py b/fastplms/fine_tuning_example.py deleted file mode 100644 index 3356320..0000000 --- a/fastplms/fine_tuning_example.py +++ /dev/null @@ -1,602 +0,0 @@ -#! /usr/bin/env python3 -""" -This script is a simple example of how to fine-tune a Synthyra FastPLM model for a protein sequence regression or classification task. -For regression we look at the binding affinity of two proteins (pkd) -For classification we look at the solubility of a protein (membrane bound or not) -""" -import torch -import numpy as np -import matplotlib.pyplot as plt -import seaborn as sns -from datasets import load_dataset -from torch.utils.data import Dataset as TorchDataset -from typing import List, Tuple, Dict, Union, Any -from transformers import ( - AutoModelForSequenceClassification, - Trainer, - TrainingArguments, - EarlyStoppingCallback, - EvalPrediction -) -from peft import LoraConfig, get_peft_model -from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay -from scipy.stats import spearmanr - - -# Shared arguments for the trainer -BASE_TRAINER_KWARGS = { - "warmup_steps": 500, - "weight_decay": 0.01, - "logging_steps": 100, - "eval_strategy": "steps", - "eval_steps": 500, - "save_strategy": "steps", - "save_steps": 500, - "load_best_model_at_end": True, - "metric_for_best_model": "eval_loss", - "greater_is_better": False, - "report_to": "none", - "label_names": ["labels"] -} - - -# Dataset classes -class PairDatasetHF(TorchDataset): - """ - Dataset class for protein pair data (e.g., protein-protein interactions). - - Args: - data: The dataset containing protein sequences and labels - col_a: Column name for the first protein sequence - col_b: Column name for the second protein sequence - label_col: Column name for the labels - max_length: Maximum sequence length to consider - """ - def __init__(self, dataset: Any, col_a: str, col_b: str, label_col: str, max_length: int = 2048): - self.seqs_a = dataset[col_a] - self.seqs_b = dataset[col_b] - self.labels = dataset[label_col] - self.max_length = max_length - - def __len__(self) -> int: - return len(self.seqs_a) - - def __getitem__(self, idx: int) -> Tuple[str, str, Union[float, int]]: - seq_a = self.seqs_a[idx][:self.max_length] - seq_b = self.seqs_b[idx][:self.max_length] - label = self.labels[idx] - return seq_a, seq_b, label - - -class SequenceDatasetHF(TorchDataset): - """ - Dataset class for single protein sequence data. - - Args: - dataset: The dataset containing protein sequences and labels - col_name: Column name for the protein sequences - label_col: Column name for the labels - max_length: Maximum sequence length to consider - """ - def __init__(self, dataset: Any, col_name: str = 'seqs', label_col: str = 'labels', max_length: int = 2048): - self.seqs = dataset[col_name] - self.labels = dataset[label_col] - self.max_length = max_length - - def __len__(self) -> int: - return len(self.seqs) - - def __getitem__(self, idx: int) -> Tuple[str, Union[float, int]]: - seq = self.seqs[idx][:self.max_length] - label = self.labels[idx] - return seq, label - - -class PairCollator: - """ - Collator for protein pair data that handles tokenization and tensor conversion. - - Args: - tokenizer: The tokenizer to use for encoding sequences - regression: Whether this is a regression task (True) or classification (False) - """ - def __init__(self, tokenizer: Any, regression: bool = False): - self.tokenizer = tokenizer - self.regression = regression - - def __call__(self, batch: List[Tuple[str, str, Union[float, int]]]) -> Dict[str, torch.Tensor]: - seqs_a, seqs_b, labels = zip(*batch) - labels = torch.tensor(labels) - if self.regression: - labels = labels.float() - else: - labels = labels.long() - tokenized = self.tokenizer( - seqs_a, seqs_b, - padding='longest', - pad_to_multiple_of=8, - return_tensors='pt' - ) - return { - 'input_ids': tokenized['input_ids'], - 'attention_mask': tokenized['attention_mask'], - 'labels': labels - } - - -class SequenceCollator: - """ - Collator for single protein sequence data that handles tokenization and tensor conversion. - - Args: - tokenizer: The tokenizer to use for encoding sequences - regression: Whether this is a regression task (True) or classification (False) - """ - def __init__(self, tokenizer: Any, regression: bool = False): - self.tokenizer = tokenizer - self.regression = regression - - def __call__(self, batch: List[Tuple[str, Union[float, int]]]) -> Dict[str, torch.Tensor]: - seqs, labels = zip(*batch) - labels = torch.tensor(labels) - if self.regression: - labels = labels.float() - else: - labels = labels.long() - tokenized = self.tokenizer( - seqs, - padding='longest', - pad_to_multiple_of=8, - return_tensors='pt' - ) - return { - 'input_ids': tokenized['input_ids'], - 'attention_mask': tokenized['attention_mask'], - 'labels': labels - } - - -# Get the model ready, with or without LoRA -def initialize_model(model_name: str, num_labels: int, use_lora: bool = True, lora_config: Any = None): - """ - Initialize a model with optional LoRA support - - Args: - model_name: Name or path of the pretrained model - num_labels: Number of labels for the task (1 for regression) - use_lora: Whether to use LoRA for fine-tuning - lora_config: Custom LoRA configuration (optional) - - Returns: - model: The initialized model - tokenizer: The model's tokenizer - """ - print(f"Loading model {model_name} with {num_labels} labels...") - - # Load base model - model = AutoModelForSequenceClassification.from_pretrained( - model_name, - trust_remote_code=True, - num_labels=num_labels - ) - tokenizer = model.tokenizer - - # Apply LoRA if requested - if use_lora: - # Default LoRA configuration if none provided - if lora_config is None: - # Target modules for ESM++ or ESM2 models - target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"] - - lora_config = LoraConfig( - r=8, - lora_alpha=16, - lora_dropout=0.01, - bias="none", - target_modules=target_modules, - ) - - # Apply LoRA to the model - model = get_peft_model(model, lora_config) - - # Unfreeze the classifier head - for param in model.classifier.parameters(): - param.requires_grad = True - - # Print parameter statistics - total_params = sum(p.numel() for p in model.parameters()) - trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - non_trainable_params = total_params - trainable_params - print(f"Total parameters: {total_params}") - print(f"Trainable parameters: {trainable_params}") - print(f"Non-trainable parameters: {non_trainable_params}") - print(f"Percentage of parameters being trained: {100 * trainable_params / total_params:.2f}%") - - return model, tokenizer - - -# For computing performance metrics, it's fairly straightforward to add more metrics here -def compute_metrics_regression(p: EvalPrediction) -> Dict[str, float]: - """Compute Spearman correlation for regression tasks""" - predictions, labels = p.predictions, p.label_ids - predictions = predictions[0] if isinstance(predictions, tuple) else predictions - - # Calculate Spearman correlation - correlation, p_value = spearmanr(predictions.flatten(), labels.flatten()) - - return { - "spearman_correlation": correlation, - "p_value": p_value - } - - -def compute_metrics_classification(p: EvalPrediction) -> Dict[str, float]: - """Compute accuracy for classification tasks""" - predictions, labels = p.predictions, p.label_ids - predictions = predictions[0] if isinstance(predictions, tuple) else predictions - predictions = np.argmax(predictions, axis=-1) - - accuracy = (predictions.flatten() == labels.flatten()).mean() - - return { - "accuracy": accuracy - } - - -# For plotting the results, it's fairly straightforward to add more plots here -def plot_regression_results(preds: np.ndarray, labels: np.ndarray, task_name: str = "Regression") -> float: - """ - Plot regression results with Spearman correlation - - Args: - preds: Predicted values - labels: True values - task_name: Name of the task for plot title and filename - - Returns: - correlation: Spearman correlation coefficient - """ - # Calculate Spearman correlation - correlation, p_value = spearmanr(preds, labels) - - # Create scatter plot - plt.figure(figsize=(10, 8)) - sns.scatterplot(x=labels, y=preds, alpha=0.6) - - # Add regression line - sns.regplot(x=labels, y=preds, scatter=False, color='red') - - plt.title(f'{task_name} - Spearman Correlation: {correlation:.3f} (p={p_value:.3e})') - plt.xlabel('True Values') - plt.ylabel('Predicted Values') - - # Add correlation text - plt.annotate(f'ρ = {correlation:.3f}', xy=(0.05, 0.95), xycoords='axes fraction', - fontsize=12, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8)) - - plt.tight_layout() - plt.savefig(f'{task_name.lower().replace(" ", "_")}_results.png') - plt.show() - return correlation - - -def plot_classification_results(trainer: Trainer, test_dataset: Any, task_name: str = "Classification") -> float: - """ - Plot classification results with confusion matrix - - Args: - trainer: The trained model trainer - test_dataset: Dataset to evaluate on - task_name: Name of the task for plot title and filename - - Returns: - accuracy: Classification accuracy - """ - # Get predictions - predictions, labels, _ = trainer.predict(test_dataset) - preds = predictions[0] if isinstance(predictions, tuple) else predictions - pred_values = np.argmax(preds, axis=1) - - # Calculate accuracy - accuracy = (pred_values == labels).mean() - - # Create confusion matrix - cm = confusion_matrix(labels, pred_values) - - # Plot confusion matrix - plt.figure(figsize=(10, 8)) - disp = ConfusionMatrixDisplay(confusion_matrix=cm) - disp.plot(cmap=plt.cm.Blues) - - plt.title(f'{task_name} - Accuracy: {accuracy:.3f}') - plt.tight_layout() - plt.savefig(f'{task_name.lower().replace(" ", "_")}_results.png') - plt.show() - - return accuracy - - -# Training functions -def train_regression_model( - model_name: str = 'Synthyra/ESMplusplus_small', - use_lora: bool = True, - custom_lora_config: Any = None, - batch_size: int = 8, - learning_rate: float = 5e-5, - num_epochs: int = 10, - max_length: int = 1024, - gradient_accumulation_steps: int = 1, - patience: int = 3 - ) -> Tuple[Trainer, Any]: - """ - Train a regression model for protein-protein affinity prediction - - Args: - model_name: Name or path of the pretrained model - use_lora: Whether to use LoRA for fine-tuning - custom_lora_config: Custom LoRA configuration (optional) - batch_size: Batch size for training - learning_rate: Learning rate for training - num_epochs: Number of epochs for training - max_length: Maximum sequence length to consider - gradient_accumulation_steps: Number of gradient accumulation steps - patience: Number of evaluation calls with no improvement after which training will be stopped - - Returns: - trainer: The trained model trainer - test_dataset: The test dataset used for evaluation - """ - print("Loading datasets for regression task...") - - # Filter sequences that exceed max_length - def _filter_pair_by_length(example: Any) -> bool: - return len(example['SeqA']) + len(example['SeqB']) <= max_length - - # Load datasets - train_data = load_dataset('Synthyra/ProteinProteinAffinity', split='train').filter(_filter_pair_by_length) - valid_data = load_dataset('Synthyra/AffinityBenchmarkv5.5', split='train').filter(_filter_pair_by_length) - test_data = load_dataset('Synthyra/haddock_benchmark', split='train').filter(_filter_pair_by_length) - - # Create datasets - train_dataset = PairDatasetHF(train_data, 'SeqA', 'SeqB', 'labels', max_length=max_length) - valid_dataset = PairDatasetHF(valid_data, 'SeqA', 'SeqB', 'labels', max_length=max_length) - test_dataset = PairDatasetHF(test_data, 'SeqA', 'SeqB', 'labels', max_length=max_length) - - # Initialize model with modular function - model, tokenizer = initialize_model( - model_name=model_name, - num_labels=1, # Regression task - use_lora=use_lora, - lora_config=custom_lora_config - ) - - # Create data collator - data_collator = PairCollator(tokenizer, regression=True) - - # Define training arguments - output_dir = "./results_regression_lora" if use_lora else "./results_regression" - logging_dir = "./logs_regression_lora" if use_lora else "./logs_regression" - - training_args = TrainingArguments( - output_dir=output_dir, - num_train_epochs=num_epochs, - gradient_accumulation_steps=gradient_accumulation_steps, - per_device_train_batch_size=batch_size, - per_device_eval_batch_size=batch_size, - logging_dir=logging_dir, - learning_rate=learning_rate, - **BASE_TRAINER_KWARGS - ) - - # Create trainer - trainer = Trainer( - model=model, - args=training_args, - train_dataset=train_dataset, - eval_dataset=valid_dataset, - data_collator=data_collator, - compute_metrics=compute_metrics_regression, - callbacks=[EarlyStoppingCallback(early_stopping_patience=patience)] - ) - - metrics = trainer.evaluate(test_dataset) - print(f"Initial metrics: {metrics}") - print("Training regression model...") - trainer.train() - - # Evaluate and visualize results - print("Evaluating and visualizing results...") - predictions, labels, metrics = trainer.predict(test_dataset) - preds = predictions[0] if isinstance(predictions, tuple) else predictions - correlation = plot_regression_results(preds.flatten(), labels.flatten(), "Protein-Protein Affinity") - print(f"Final Spearman correlation on test set: {correlation:.3f}") - return trainer, test_dataset - - -def train_classification_model( - model_name: str = 'Synthyra/ESMplusplus_small', - use_lora: bool = True, - custom_lora_config: Any = None, - batch_size: int = 8, - learning_rate: float = 5e-5, - num_epochs: int = 10, - max_length: int = 512, - gradient_accumulation_steps: int = 1, - patience: int = 3 - ) -> Tuple[Trainer, Any]: - """ - Train a classification model for protein solubility prediction - - Args: - model_name: Name or path of the pretrained model - use_lora: Whether to use LoRA for fine-tuning - custom_lora_config: Custom LoRA configuration (optional) - batch_size: Batch size for training - learning_rate: Learning rate for training - num_epochs: Number of epochs for training - max_length: Maximum sequence length to consider - gradient_accumulation_steps: Number of gradient accumulation steps - patience: Number of evaluation calls with no improvement after which training will be stopped - - Returns: - trainer: The trained model trainer - """ - print("Loading datasets for classification task...") - - # Filter sequences that exceed max_length - def _filter_by_length(example: Any) -> bool: - return len(example['seqs']) <= max_length - - # Load datasets - data = load_dataset('GleghornLab/DL2_reg') - train_data = data['train'].filter(_filter_by_length) - valid_data = data['valid'].filter(_filter_by_length) - test_data = data['test'].filter(_filter_by_length) - - # Create datasets - train_dataset = SequenceDatasetHF(train_data, 'seqs', 'labels', max_length=max_length) - valid_dataset = SequenceDatasetHF(valid_data, 'seqs', 'labels', max_length=max_length) - test_dataset = SequenceDatasetHF(test_data, 'seqs', 'labels', max_length=max_length) - - # Get number of labels - num_labels = len(set(train_data['labels'])) - - # Initialize model with modular function - model, tokenizer = initialize_model( - model_name=model_name, - num_labels=num_labels, - use_lora=use_lora, - lora_config=custom_lora_config - ) - - # Create data collator - data_collator = SequenceCollator(tokenizer, regression=False) - - # Define training arguments - output_dir = "./results_classification_lora" if use_lora else "./results_classification" - logging_dir = "./logs_classification_lora" if use_lora else "./logs_classification" - - training_args = TrainingArguments( - output_dir=output_dir, - num_train_epochs=num_epochs, - gradient_accumulation_steps=gradient_accumulation_steps, - per_device_train_batch_size=batch_size, - per_device_eval_batch_size=batch_size, - logging_dir=logging_dir, - learning_rate=learning_rate, - **BASE_TRAINER_KWARGS - ) - - # Create trainer - trainer = Trainer( - model=model, - args=training_args, - train_dataset=train_dataset, - eval_dataset=valid_dataset, - data_collator=data_collator, - compute_metrics=compute_metrics_classification, - callbacks=[EarlyStoppingCallback(early_stopping_patience=patience)] - ) - - metrics = trainer.evaluate(test_dataset) - print(f"Initial metrics: {metrics}") - print("Training classification model...") - trainer.train() - - # Evaluate and visualize results - print("Evaluating and visualizing results...") - accuracy = plot_classification_results(trainer, test_dataset, "Protein Solubility") - print(f"Final accuracy on test set: {accuracy:.3f}") - - return trainer - - -# Main function -if __name__ == "__main__": - """ - With default arguments on 4070 laptop GPU - py -m fine_tuning_example --task classification --batch_size 8 --epochs 2 - Runs in 80 seconds with test accuracy of ~89% - py -m fine_tuning_example --task regression --batch_size 2 --max_length 1024 --grad_accum 4 --epochs 2 - Runs in 7 minutes with test Spearman correlation of ~0.72 - """ - import argparse - - # Examples of PLMs with efficient implementations offered by Synthyra - MODEL_LIST = [ - 'Synthyra/ESMplusplus_small', - 'Synthyra/ESMplusplus_large', - 'Synthyra/ESM2-8M', - 'Synthyra/ESM2-35M', - 'Synthyra/ESM2-150M', - 'Synthyra/ESM2-650M', - ] - - parser = argparse.ArgumentParser(description="Train models for protein tasks") - parser.add_argument("--task", type=str, choices=["regression", "classification", "both"], - default="both", help="Task to train model for") - parser.add_argument("--model_path", type=str, default="Synthyra/ESM2-8M", - help="Path to the model to train") - parser.add_argument("--use_lora", action="store_true", default=True, - help="Whether to use LoRA for fine-tuning") - parser.add_argument("--batch_size", type=int, default=2, - help="Batch size for training") - parser.add_argument("--lr", type=float, default=5e-5, - help="Learning rate for training") - parser.add_argument("--epochs", type=float, default=1.0, - help="Number of epochs for training") - parser.add_argument("--max_length", type=int, default=512, - help="Maximum length of input sequences") - parser.add_argument("--grad_accum", type=int, default=1, - help="Number of gradient accumulation steps") - parser.add_argument("--patience", type=int, default=3, - help="Early stopping patience - number of evaluation calls with no improvement after which training will be stopped") - args = parser.parse_args() - - # Print training configuration - print("\n" + "="*50) - print("TRAINING CONFIGURATION") - print("="*50) - print(f"Task: {args.task}") - print(f"Using LoRA: {args.use_lora}") - print(f"Batch size: {args.batch_size}") - print(f"Learning rate: {args.lr}") - print(f"Number of epochs: {args.epochs}") - print(f"Max sequence length: {args.max_length}") - print(f"Gradient Accumulation Steps: {args.grad_accum}") - print(f"Early stopping patience: {args.patience}") - print("="*50 + "\n") - - # Train regression model if required - if args.task in ["regression", "both"]: - print("\n" + "="*50) - print("TRAINING REGRESSION MODEL") - print("="*50) - regression_trainer, test_dataset = train_regression_model( - model_name=args.model_path, - use_lora=args.use_lora, - batch_size=args.batch_size, - learning_rate=args.lr, - num_epochs=args.epochs, - max_length=args.max_length, - gradient_accumulation_steps=args.grad_accum, - patience=args.patience - ) - - # Train classification model if required - if args.task in ["classification", "both"]: - print("\n" + "="*50) - print("TRAINING CLASSIFICATION MODEL") - print("="*50) - classification_trainer = train_classification_model( - model_name=args.model_path, - use_lora=args.use_lora, - batch_size=args.batch_size, - learning_rate=args.lr, - num_epochs=args.epochs, - max_length=args.max_length, - gradient_accumulation_steps=args.grad_accum, - patience=args.patience - ) - - print("\nTraining completed!") diff --git a/fastplms/test_time_training.py b/fastplms/test_time_training.py deleted file mode 100644 index bcae5db..0000000 --- a/fastplms/test_time_training.py +++ /dev/null @@ -1,539 +0,0 @@ -from __future__ import annotations - -import typing as T -from dataclasses import dataclass, fields - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -@dataclass -class TTTConfig: - lr: float = 4e-4 - steps: int = 30 - ags: int = 16 - batch_size: int = 2 - mask_ratio: float = 0.15 - crop_size: int = 1024 - bert_leave_prob: float = 0.1 - bert_replace_prob: float = 0.1 - optimizer: str = "sgd" - momentum: float = 0.0 - weight_decay: float = 0.0 - seed: int | None = 0 - lora_rank: int = 8 - lora_alpha: float = 32.0 - lora_target_replace_module: str | None = None - lora_target_modules: tuple[str, ...] | None = None - initial_state_reset: bool = True - automatic_best_state_reset: bool = False - eval_each_step: bool = False - gradient_clip: bool = False - gradient_clip_max_norm: float = 1.0 - - @classmethod - def from_kwargs(cls, **kwargs: T.Any) -> "TTTConfig": - valid_names = {field.name for field in fields(cls)} - unknown_names = set(kwargs) - valid_names - assert len(unknown_names) == 0, f"Unknown TTTConfig fields: {sorted(unknown_names)}" - return cls(**kwargs) - - def merged(self, overrides: T.Mapping[str, T.Any] | "TTTConfig" | None) -> "TTTConfig": - if overrides is None: - return self - if isinstance(overrides, TTTConfig): - return overrides - values = {field.name: self.__dict__[field.name] for field in fields(self)} - for name, value in overrides.items(): - assert name in values, f"Unknown TTTConfig field: {name}" - values[name] = value - return TTTConfig(**values) - - def verify(self) -> None: - assert self.lr > 0.0, "TTT learning rate must be positive." - assert self.steps >= 1, "TTT steps must be >= 1." - assert self.ags >= 1, "TTT gradient accumulation steps must be >= 1." - assert self.batch_size >= 1, "TTT batch_size must be >= 1." - assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]." - assert self.crop_size >= 1, "TTT crop_size must be >= 1." - assert self.lora_rank >= 1, "TTT v1 is LoRA-only, so lora_rank must be >= 1." - assert self.lora_alpha > 0.0, "TTT lora_alpha must be positive." - assert self.optimizer in {"adamw", "sgd"}, "TTT optimizer must be 'adamw' or 'sgd'." - assert 0.0 <= self.bert_leave_prob <= 1.0, "bert_leave_prob must be in [0, 1]." - assert 0.0 <= self.bert_replace_prob <= 1.0, "bert_replace_prob must be in [0, 1]." - assert self.bert_leave_prob + self.bert_replace_prob <= 1.0, ( - "bert_leave_prob + bert_replace_prob must be <= 1." - ) - if self.gradient_clip: - assert self.gradient_clip_max_norm > 0.0, "gradient_clip_max_norm must be positive." - - -class LoraInjectedLinear(nn.Module): - def __init__(self, linear: nn.Module, rank: int, alpha: float) -> None: - super().__init__() - weight = linear._parameters["weight"] - assert weight.ndim == 2, "LoRA can only wrap 2D linear weights." - self.linear = linear - self.linear.requires_grad_(False) - self.rank = rank - self.scale = alpha - in_features = weight.shape[1] - out_features = weight.shape[0] - self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32) - self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32) - self.lora_down.to(device=weight.device) - self.lora_up.to(device=weight.device) - nn.init.normal_(self.lora_down.weight, std=1.0 / rank) - nn.init.zeros_(self.lora_up.weight) - - @property - def weight(self) -> torch.Tensor: - return self.linear._parameters["weight"] - - @property - def bias(self) -> torch.Tensor | None: - return self.linear._parameters["bias"] - - def forward(self, x: torch.Tensor) -> torch.Tensor: - base = self.linear(x) - delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale - return base + delta.to(dtype=base.dtype) - - -class FastPLMTestTimeTrainingMixin: - def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None: - base_config = TTTConfig() - self._ttt_cfg = base_config.merged(ttt_config) - self._ttt_cfg.verify() - self._ttt_initialized = False - self._ttt_initial_state: list[dict[str, torch.Tensor]] | None = None - - @property - def ttt_config(self) -> TTTConfig: - if "_ttt_cfg" not in self.__dict__: - self.init_ttt() - return self._ttt_cfg - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self] - - def _ttt_get_frozen_modules(self) -> list[nn.Module]: - return [] - - def _ttt_tokenize( - self, - seq: str | list[str] | None = None, - input_ids: torch.Tensor | None = None, - **kwargs: T.Any, - ) -> torch.Tensor | dict[str, torch.Tensor]: - del kwargs - if input_ids is not None: - return input_ids - assert seq is not None, "Pass either seq or input_ids for TTT." - tokenized = self.tokenizer(seq, return_tensors="pt", padding=True) - return tokenized["input_ids"] - - def _ttt_mask_token(self) -> int: - return int(self.tokenizer.mask_token_id) - - def _ttt_padding_token(self) -> int: - return int(self.tokenizer.pad_token_id) - - def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: - tokenizer = self.tokenizer - special_ids = set(tokenizer.all_special_ids) - vocab_size = int(self.config.vocab_size) - ids = [idx for idx in range(vocab_size) if idx not in special_ids] - assert len(ids) > 0, "TTT replacement token set is empty." - return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) - - def _ttt_predict_logits( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - **kwargs: T.Any, - ) -> torch.Tensor: - del kwargs - if isinstance(batch, dict): - output = self(**batch) - return output.logits - attention_mask = batch.ne(self._ttt_padding_token()) - output = self(input_ids=batch, attention_mask=attention_mask) - return output.logits - - def _ttt_eval_step( - self, - step: int, - loss: float, - seq: str | list[str] | None = None, - input_ids: torch.Tensor | None = None, - **kwargs: T.Any, - ) -> tuple[dict[str, T.Any], float | None]: - del step, loss, seq, input_ids, kwargs - return {}, None - - def _ttt_is_lora_target( - self, - name: str, - full_name: str, - module: nn.Module, - active: bool, - target_modules: tuple[str, ...] | None, - ) -> bool: - if not active: - return False - if isinstance(module, LoraInjectedLinear): - return False - if ( - target_modules is not None - and name not in target_modules - and full_name not in target_modules - ): - return False - if isinstance(module, nn.Linear): - return True - if "weight" not in module._parameters: - return False - weight = module._parameters["weight"] - if weight is None or weight.ndim != 2: - return False - return "Linear" in module.__class__.__name__ - - def _ttt_inject_lora(self) -> int: - cfg = self.ttt_config - cfg.verify() - target_class = cfg.lora_target_replace_module - target_modules = cfg.lora_target_modules - wrapped = 0 - - def inject(module: nn.Module, prefix: str, active: bool) -> None: - nonlocal wrapped - for name, child in list(module.named_children()): - full_name = f"{prefix}.{name}" if prefix else name - child_active = active - if target_class is not None: - child_active = active or child.__class__.__name__ == target_class - if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules): - setattr( - module, - name, - LoraInjectedLinear(child, rank=cfg.lora_rank, alpha=cfg.lora_alpha), - ) - wrapped += 1 - continue - inject(child, full_name, child_active) - - for trainable_module in self._ttt_get_trainable_modules(): - inject(trainable_module, "", target_class is None) - assert wrapped > 0, "TTT LoRA injection did not find any target modules." - return wrapped - - def _ttt_lora_modules(self) -> list[LoraInjectedLinear]: - return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)] - - def _ttt_lora_parameters(self) -> list[nn.Parameter]: - params: list[nn.Parameter] = [] - for module in self._ttt_lora_modules(): - params.extend(module.lora_down.parameters()) - params.extend(module.lora_up.parameters()) - assert len(params) > 0, "TTT has no LoRA parameters." - return params - - def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]: - snapshot = [] - for module in self._ttt_lora_modules(): - snapshot.append( - { - "lora_down.weight": module.lora_down.weight.detach().clone(), - "lora_up.weight": module.lora_up.weight.detach().clone(), - } - ) - assert len(snapshot) > 0, "TTT has no LoRA state to snapshot." - return snapshot - - def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None: - modules = self._ttt_lora_modules() - assert len(modules) == len(state), "TTT LoRA state/module count mismatch." - with torch.no_grad(): - for module, module_state in zip(modules, state): - module.lora_down.weight.copy_(module_state["lora_down.weight"]) - module.lora_up.weight.copy_(module_state["lora_up.weight"]) - - def _ttt_ensure_initialized(self) -> None: - if "_ttt_cfg" not in self.__dict__: - self.init_ttt() - if self._ttt_initialized: - return - self._ttt_inject_lora() - self._ttt_initial_state = self._ttt_snapshot_lora_state() - self._ttt_initialized = True - - def ttt_reset(self) -> None: - self._ttt_ensure_initialized() - assert self._ttt_initial_state is not None, "TTT initial state is not available." - self._ttt_restore_lora_state(self._ttt_initial_state) - - def _ttt_make_optimizer(self) -> torch.optim.Optimizer: - cfg = self.ttt_config - params = self._ttt_lora_parameters() - if cfg.optimizer == "sgd": - return torch.optim.SGD( - params, - lr=cfg.lr, - momentum=cfg.momentum, - weight_decay=cfg.weight_decay, - ) - return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay) - - def _ttt_to_device( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - device: torch.device, - ) -> torch.Tensor | dict[str, torch.Tensor]: - if isinstance(batch, dict): - return {name: tensor.to(device) for name, tensor in batch.items()} - return batch.to(device) - - def _ttt_input_ids_from_batch( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - ) -> torch.Tensor: - if isinstance(batch, dict): - return batch["input_ids"] - return batch - - def _ttt_set_input_ids( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - input_ids: torch.Tensor, - ) -> torch.Tensor | dict[str, torch.Tensor]: - if isinstance(batch, dict): - updated = dict(batch) - updated["input_ids"] = input_ids - return updated - return input_ids - - def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: - pad_token = self._ttt_padding_token() - mask = input_ids.ne(pad_token) - special_ids = set(self.tokenizer.all_special_ids) - for special_id in special_ids: - mask = mask & input_ids.ne(int(special_id)) - return mask - - def _ttt_sample_crop( - self, - batch: torch.Tensor | dict[str, torch.Tensor], - generator: torch.Generator, - ) -> torch.Tensor | dict[str, torch.Tensor]: - input_ids = self._ttt_input_ids_from_batch(batch) - cfg = self.ttt_config - if input_ids.shape[1] <= cfg.crop_size: - return batch - high = input_ids.shape[1] - cfg.crop_size + 1 - start = int( - torch.randint( - high, - (1,), - generator=generator, - device=input_ids.device, - ).item() - ) - end = start + cfg.crop_size - if isinstance(batch, dict): - cropped = {} - for name, tensor in batch.items(): - if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]: - cropped[name] = tensor[:, start:end] - else: - cropped[name] = tensor - return cropped - return input_ids[:, start:end] - - def _ttt_sample_batch( - self, - tokenized: torch.Tensor | dict[str, torch.Tensor], - generator: torch.Generator, - ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]: - cfg = self.ttt_config - batch = self._ttt_sample_crop(tokenized, generator) - input_ids = self._ttt_input_ids_from_batch(batch) - rows = torch.randint( - input_ids.shape[0], - (cfg.batch_size,), - generator=generator, - device=input_ids.device, - ) - if isinstance(batch, dict): - sampled: torch.Tensor | dict[str, torch.Tensor] = {} - for name, tensor in batch.items(): - if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]: - sampled[name] = tensor.index_select(0, rows) - else: - sampled[name] = tensor - else: - sampled = input_ids.index_select(0, rows) - - sampled_ids = self._ttt_input_ids_from_batch(sampled) - labels = sampled_ids.clone() - non_special = self._ttt_non_special_mask(sampled_ids) - label_mask = torch.zeros_like(non_special) - for row_idx in range(sampled_ids.shape[0]): - candidate_positions = torch.where(non_special[row_idx])[0] - if candidate_positions.numel() == 0: - continue - num_mask = max(1, int(round(candidate_positions.numel() * cfg.mask_ratio))) - order = torch.randperm( - candidate_positions.numel(), - generator=generator, - device=sampled_ids.device, - ) - chosen = candidate_positions[order[:num_mask]] - label_mask[row_idx, chosen] = True - labels = labels.masked_fill(~label_mask, -100) - - masked_ids = sampled_ids.clone() - chosen_positions = torch.where(label_mask) - if chosen_positions[0].numel() > 0: - random_values = torch.rand( - chosen_positions[0].shape, - generator=generator, - device=sampled_ids.device, - ) - leave = random_values < cfg.bert_leave_prob - replace = (random_values >= cfg.bert_leave_prob) & ( - random_values < cfg.bert_leave_prob + cfg.bert_replace_prob - ) - mask = ~(leave | replace) - if mask.any(): - masked_ids[ - chosen_positions[0][mask], - chosen_positions[1][mask], - ] = self._ttt_mask_token() - if replace.any(): - replacement_tokens = self._ttt_replacement_tokens(sampled_ids) - replacement_idx = torch.randint( - replacement_tokens.shape[0], - (int(replace.sum().item()),), - generator=generator, - device=sampled_ids.device, - ) - masked_ids[ - chosen_positions[0][replace], - chosen_positions[1][replace], - ] = replacement_tokens[replacement_idx] - - return self._ttt_set_input_ids(sampled, masked_ids), labels - - def ttt( - self, - seq: str | list[str] | None = None, - input_ids: torch.Tensor | None = None, - ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None, - **kwargs: T.Any, - ) -> dict[str, T.Any]: - if ttt_config is not None: - if "_ttt_initialized" in self.__dict__ and self._ttt_initialized: - next_cfg = self.ttt_config.merged(ttt_config) - assert next_cfg.lora_rank == self.ttt_config.lora_rank, ( - "Changing lora_rank after TTT initialization is not supported." - ) - assert next_cfg.lora_alpha == self.ttt_config.lora_alpha, ( - "Changing lora_alpha after TTT initialization is not supported." - ) - assert ( - next_cfg.lora_target_replace_module - == self.ttt_config.lora_target_replace_module - ), "Changing LoRA target class after TTT initialization is not supported." - assert next_cfg.lora_target_modules == self.ttt_config.lora_target_modules, ( - "Changing LoRA target modules after TTT initialization is not supported." - ) - self._ttt_cfg = next_cfg - else: - self.init_ttt(ttt_config) - - self._ttt_ensure_initialized() - cfg = self.ttt_config - if cfg.initial_state_reset: - self.ttt_reset() - - device = next(self.parameters()).device - tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs) - tokenized = self._ttt_to_device(tokenized, device) - generator_device = device if device.type == "cuda" else torch.device("cpu") - generator = torch.Generator(device=generator_device) - if cfg.seed is not None: - generator.manual_seed(cfg.seed) - - module_modes = {module: module.training for module in self.modules()} - requires_grad = {param: param.requires_grad for param in self.parameters()} - losses: list[float] = [] - step_metrics: list[dict[str, T.Any]] = [] - best_state: list[dict[str, torch.Tensor]] | None = None - best_metric: float | None = None - best_step = 0 - - try: - self.train() - for param in self.parameters(): - param.requires_grad_(False) - for param in self._ttt_lora_parameters(): - param.requires_grad_(True) - - optimizer = self._ttt_make_optimizer() - optimizer.zero_grad(set_to_none=True) - total_micro_steps = cfg.steps * cfg.ags - for micro_step in range(total_micro_steps): - batch, labels = self._ttt_sample_batch(tokenized, generator) - logits = self._ttt_predict_logits(batch, **kwargs) - labels = labels.to(device=logits.device) - loss = F.cross_entropy( - logits.reshape(-1, logits.shape[-1]), - labels.reshape(-1), - ignore_index=-100, - ) - (loss / cfg.ags).backward() - if (micro_step + 1) % cfg.ags != 0: - continue - - if cfg.gradient_clip: - torch.nn.utils.clip_grad_norm_( - self._ttt_lora_parameters(), - cfg.gradient_clip_max_norm, - ) - optimizer.step() - optimizer.zero_grad(set_to_none=True) - step = (micro_step + 1) // cfg.ags - loss_value = float(loss.detach().item()) - losses.append(loss_value) - if cfg.eval_each_step: - metrics, metric = self._ttt_eval_step( - step=step, - loss=loss_value, - seq=seq, - input_ids=input_ids, - **kwargs, - ) - if len(metrics) > 0: - step_metrics.append(metrics) - if metric is not None and ( - best_metric is None or metric > best_metric - ): - best_metric = metric - best_step = step - best_state = self._ttt_snapshot_lora_state() - - if cfg.automatic_best_state_reset and best_state is not None: - self._ttt_restore_lora_state(best_state) - finally: - for param, value in requires_grad.items(): - param.requires_grad_(value) - for module, training in module_modes.items(): - module.train(training) - - return { - "losses": losses, - "step_metrics": step_metrics, - "best_step": best_step, - "best_metric": best_metric, - } diff --git a/fastplms/weight_parity_utils.py b/fastplms/weight_parity_utils.py deleted file mode 100644 index 06bc56f..0000000 --- a/fastplms/weight_parity_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import torch -from torch.nn.functional import mse_loss -from typing import Dict, List, Mapping - - -def assert_model_parameters_fp32(model: torch.nn.Module, model_name: str) -> None: - non_fp32: List[Dict[str, str]] = [] - parameter_count = 0 - for name, parameter in model.named_parameters(): - parameter_count += 1 - if parameter.dtype != torch.float32: - non_fp32.append({"name": name, "dtype": str(parameter.dtype)}) - - assert parameter_count > 0, f"{model_name} has no parameters." - assert len(non_fp32) == 0, ( - f"{model_name} parameters must all be torch.float32. " - f"non_fp32_count={len(non_fp32)} sample={non_fp32[:5]}" - ) - - -def assert_state_dict_floating_tensors_fp32( - state_dict: Mapping[str, torch.Tensor], - state_dict_name: str, -) -> None: - non_fp32: List[Dict[str, str]] = [] - for tensor_name in sorted(state_dict.keys()): - tensor = state_dict[tensor_name] - assert torch.is_tensor(tensor), ( - f"{state_dict_name} state_dict entry must be a tensor. " - f"name={tensor_name} type={type(tensor)}" - ) - if tensor.is_floating_point() and tensor.dtype != torch.float32: - non_fp32.append({"name": tensor_name, "dtype": str(tensor.dtype)}) - - assert len(non_fp32) == 0, ( - f"{state_dict_name} floating tensors must be torch.float32. " - f"non_fp32_count={len(non_fp32)} sample={non_fp32[:5]}" - ) - - -def assert_state_dict_equal( - reference_state_dict: Mapping[str, torch.Tensor], - candidate_state_dict: Mapping[str, torch.Tensor], - context: str, - max_report: int = 10, -) -> None: - error_msgs = [] - for (ref_name, ref_tensor), (cand_name, cand_tensor) in zip(reference_state_dict.items(), candidate_state_dict.items()): - if ref_name != cand_name: - msg = f"Name mismatch: {ref_name} != {cand_name}" - print(msg) - error_msgs.append(msg) - else: - diff = mse_loss(ref_tensor, cand_tensor).item() - if diff > 0.0: - msg = f"{ref_name}: {diff}" - print(msg) - error_msgs.append(msg) - assert not error_msgs, ( - f"{context} state_dict parity failed:{' | '.join(error_msgs[:max_report])}" - ) - - -def assert_models_fp32_and_equal( - reference_model: torch.nn.Module, - candidate_model: torch.nn.Module, - context: str, - max_report: int = 5, -) -> None: - assert_model_parameters_fp32(model=reference_model, model_name=f"{context} reference model") - assert_model_parameters_fp32(model=candidate_model, model_name=f"{context} candidate model") - assert_state_dict_equal( - reference_state_dict=reference_model.state_dict(), - candidate_state_dict=candidate_model.state_dict(), - context=context, - max_report=max_report, - ) diff --git a/kernels.lock b/kernels.lock new file mode 100644 index 0000000..1fb0d88 --- /dev/null +++ b/kernels.lock @@ -0,0 +1,142 @@ +[ + { + "repo_id": "kernels-community/flash-attn2", + "sha": "db6b51744f0cd7061386442c09df890fc6d9f47e", + "variants": { + "torch211-cxx11-cpu-x86_64-linux": { + "hash": "sha256-1e8f5c640ee6d45aae98e0337b1eac6682084d86da7883c9e043613d60da953a", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu126-aarch64-linux": { + "hash": "sha256-eff0d2520681d106364765c317b4d497f73d2838a32f521c9ec419a097d7e4d8", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu126-x86_64-linux": { + "hash": "sha256-89f8ecab7835608292fee924d527d26542381f9659d4ba1c9994b9ccdec8fb0a", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu128-aarch64-linux": { + "hash": "sha256-432851c98b56edde900b138ffeb1195b6a85640e8a4216b7fd02f36133058161", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu128-x86_64-linux": { + "hash": "sha256-e3ade9cf9cf1ea3d0c0e5dff398d04eb9ad9fa781786142c5a8a475462b33685", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu130-aarch64-linux": { + "hash": "sha256-2ab7710c2d8191a1e82aac1f7858232fdc07c68ba1a0775985cbb73f0efa15ce", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu130-x86_64-linux": { + "hash": "sha256-3e8f86d87b649fd2ac8aeccc5872654d07b3157240498db75a3f8baf1cc3706d", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-xpu20253-x86_64-linux": { + "hash": "sha256-44ade6d45eb6ec4d12018201e2373593312890a733197ed35d51915e50a2228c", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cpu-x86_64-linux": { + "hash": "sha256-ee019e1930c47a94595469751dfd3822e88eb39f533ac60105ae92d2d8f9e5f8", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu126-aarch64-linux": { + "hash": "sha256-d8b01b0baeab9bdf73b04beebde13df36c6abda16eaf7a576d291ad19e79704e", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu126-x86_64-linux": { + "hash": "sha256-61d7cac70063b44574e3593fbaf275686673c273826c1eb6632c2a6d9f2ad563", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu130-aarch64-linux": { + "hash": "sha256-35ff5caf9744fa19e95c21fa8267668150312ac44b988c2d688a737484d69faa", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu130-x86_64-linux": { + "hash": "sha256-fef69620fff871f8535e00aaf9ab1e940f529f690c511e2843e89ffccdc44f49", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu132-aarch64-linux": { + "hash": "sha256-404a8ceb6cc75cf58e3d6bce8962bb5a667af6e07b25f0c43ac3ebbd1404ff42", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu132-x86_64-linux": { + "hash": "sha256-87bc38d4b142f771f40475351b06024601d9b91634585a8a99a2919cbcdedaef", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-xpu20253-x86_64-linux": { + "hash": "sha256-614f0afb8dbcfa7a57687c32f82228f62ddfddb992b802e6073d04caf535d510", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cpu-x86_64-linux": { + "hash": "sha256-06221756fa4f85fb4c683fd399df989e6faa9ebddc1243d894e8bf3c96a0b222", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu126-aarch64-linux": { + "hash": "sha256-516d899591dbc8bba1155c9b6c4e2c1628be58dde68b7f24c09f82532751f623", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu126-x86_64-linux": { + "hash": "sha256-570f2215ceeb777750dbc123206814797eda1657033c0fcbf69fd2f3c580c5bd", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu130-aarch64-linux": { + "hash": "sha256-ef1028a133882222ba5d33654248acf9c01d851d0b0eed8c71753cd242a77d27", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu130-x86_64-linux": { + "hash": "sha256-238cdad1945962331ad685a07119bb9e893ed976f11ecbf257e03d36682f95e4", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu132-aarch64-linux": { + "hash": "sha256-4f26ede2a43e29301478405ba3c7c50fadb800324d1bc03fbbc32e9f449818b3", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-cu132-x86_64-linux": { + "hash": "sha256-a48c5ff3def3cbff8474980badaf2d8e0d4a0b6752e6c376db81166ed54e0e6e", + "hash_type": "git_lfs_concat" + }, + "torch213-cxx11-xpu20260-x86_64-linux": { + "hash": "sha256-c500992b32efbccad8add1478de2399596c4645e58dbb58fd5db392005953851", + "hash_type": "git_lfs_concat" + } + } + }, + { + "repo_id": "kernels-community/flash-attn3", + "sha": "43f0bd269777115d94ff826e0d113ce9c1c9087b", + "variants": { + "torch-stable-abi29-cu126-x86_64-linux": { + "hash": "sha256-f130ad586e33e2810a6a731e7a88f9183fbe98f95d8e085ef7db6f9153c64672", + "hash_type": "git_lfs_concat" + }, + "torch-stable-abi29-cu128-x86_64-linux": { + "hash": "sha256-1651bf516e1deb4f1db318e497cc081606a8193899bea95650f80b456ef78904", + "hash_type": "git_lfs_concat" + }, + "torch-stable-abi29-cu130-x86_64-linux": { + "hash": "sha256-8dc3c4645b8ed2c5ce27873f8c6deb4ecf60060b5f08f538389b3d79e8842a2f", + "hash_type": "git_lfs_concat" + }, + "torch210-cxx11-cu128-x86_64-linux": { + "hash": "sha256-43044b0121f22902a363c475e4d5edc46120af7abc126ad5db9782ef174dc284", + "hash_type": "git_lfs_concat" + }, + "torch210-cxx11-cu130-x86_64-linux": { + "hash": "sha256-3467b35235769ebec82256e08ec299856ac34733ed1ecc2e854890a4fcd0fdee", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu128-x86_64-linux": { + "hash": "sha256-67f2af03de8af4c76cce249662d210be955a0022c4b56ec9fd4cb200f0c86cdd", + "hash_type": "git_lfs_concat" + }, + "torch211-cxx11-cu130-x86_64-linux": { + "hash": "sha256-5283c2e8951a09c27b853f43fd757d8859fe560fb945b4ed0a68c64a618e236d", + "hash_type": "git_lfs_concat" + }, + "torch212-cxx11-cu130-x86_64-linux": { + "hash": "sha256-cbee545d181d496af2349ffca6770f9d0941b30ffa25ddd1784abde66576de5f", + "hash_type": "git_lfs_concat" + } + } + } +] \ No newline at end of file diff --git a/model_cards/README.md b/model_cards/README.md new file mode 100644 index 0000000..03aea29 --- /dev/null +++ b/model_cards/README.md @@ -0,0 +1,32 @@ +# Generated model cards + +Checkpoint cards in this directory are generated from +`src/fastplms/models.toml`: + +```bash +PYTHONPATH=src python -m tools.artifacts.generate_docs +``` + +Use `--check` to reject stale cards without writing files. Edit the typed +manifest or the renderer, not an individual generated card. Keep model cards +in this directory rather than beside runtime modules under +`src/fastplms/models/`. + +Each card combines: + +- installation and platform requirements before the Hub quick start; +- direct Hub and offline artifact loading; +- family-appropriate preparation, inference, embedding, generation, or folding + examples; +- declared AutoClasses, backends, precision, and generation behavior; +- immutable checkpoint and upstream source records; +- validation boundaries, limitations, and checkpoint terms. +- explicit AutoClass weight status and weight-publication policy; +- ESMC backend diagnostic tables whose missing frozen-head measurements remain + labeled pending rather than inferred. + +Examples must follow the current public API. Do not restore legacy backend +names, pickle embedding output, unsupported TTT paths, or capabilities from +another family. Release tests compile the examples, validate local links and +Hub license metadata, and check representative family-specific sections. They +do not execute every model example or establish a scientific result. diff --git a/model_cards/ankh2_large.md b/model_cards/ankh2_large.md new file mode 100644 index 0000000..8a4ac34 --- /dev/null +++ b/model_cards/ankh2_large.md @@ -0,0 +1,318 @@ +--- +library_name: transformers +license: "cc-by-nc-sa-4.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ANKH2_large + +This checkpoint packages the FastPLMs `ANKH` implementation. + +Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to- +sequence use. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, +`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: encoder or explicitly prepared decoder states | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ANKH2_large/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ANKH2_large" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ANKH2_large` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`. An unavailable requested backend raises +instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +`Synthyra/ANKH2_large` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH2_large", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ANKH2_large" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ANKH2_large", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Encoder and sequence-to-sequence use + +`Synthyra/ANKH2_large` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "Synthyra/ANKH2_large" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +## Notes and limitations + +ANKH parity covers the official encoder and sequence-to-sequence heads. +AutoModelForMaskedLM exposes the separately named FastPLMs synthesized +masked-LM extension and is not an official ANKH head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ANKH2_large` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `ElnaggarLab/ankh2-ext2` +- Artifact source: `official` +- State transform: `ankh_t5_to_fastplms_v1` +- Pinned upstreams: `ankh` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is +`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/ankh3_large.md b/model_cards/ankh3_large.md new file mode 100644 index 0000000..5650d9b --- /dev/null +++ b/model_cards/ankh3_large.md @@ -0,0 +1,318 @@ +--- +library_name: transformers +license: "cc-by-nc-sa-4.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ANKH3_large + +This checkpoint packages the FastPLMs `ANKH` implementation. + +Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to- +sequence use. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, +`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: encoder or explicitly prepared decoder states | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ANKH3_large/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ANKH3_large" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ANKH3_large` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`. An unavailable requested backend raises +instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +`Synthyra/ANKH3_large` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH3_large", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ANKH3_large" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ANKH3_large", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Encoder and sequence-to-sequence use + +`Synthyra/ANKH3_large` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "Synthyra/ANKH3_large" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +## Notes and limitations + +ANKH parity covers the official encoder and sequence-to-sequence heads. +AutoModelForMaskedLM exposes the separately named FastPLMs synthesized +masked-LM extension and is not an official ANKH head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ANKH3_large` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `60acb7ef86e85dc0c51fc1edf4c8e69a0480049723b6b2c95e6e9faa720c112a` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `ElnaggarLab/ankh3-large` +- Artifact source: `official` +- State transform: `ankh_t5_to_fastplms_v1` +- Pinned upstreams: `ankh` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is +`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/ankh3_xl.md b/model_cards/ankh3_xl.md new file mode 100644 index 0000000..19cafe7 --- /dev/null +++ b/model_cards/ankh3_xl.md @@ -0,0 +1,320 @@ +--- +library_name: transformers +license: "cc-by-nc-sa-4.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ANKH3_xl + +This checkpoint packages the FastPLMs `ANKH` implementation. + +Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to- +sequence use. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, +`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: encoder or explicitly prepared decoder states | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ANKH3_xl/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ANKH3_xl" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ANKH3_xl` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`. An unavailable requested backend raises +instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +`Synthyra/ANKH3_xl` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH3_xl", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ANKH3_xl" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ANKH3_xl", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Encoder and sequence-to-sequence use + +`Synthyra/ANKH3_xl` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "Synthyra/ANKH3_xl" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +## Notes and limitations + +ANKH parity covers the official encoder and sequence-to-sequence heads. +AutoModelForMaskedLM exposes the separately named FastPLMs synthesized +masked-LM extension and is not an official ANKH head. The official PyTorch +shard index is deliberately excluded: the builder verifies every declared +source shard directly and writes a new canonical safetensors index. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ANKH3_xl` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `dd2188e0d2ca65232135714eef6de394239734d843ddae4928c7398685d858e7` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `ElnaggarLab/ankh3-xl` +- Artifact source: `official` +- State transform: `ankh_t5_to_fastplms_v1` +- Pinned upstreams: `ankh` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is +`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/ankh_base.md b/model_cards/ankh_base.md new file mode 100644 index 0000000..bd5720b --- /dev/null +++ b/model_cards/ankh_base.md @@ -0,0 +1,318 @@ +--- +library_name: transformers +license: "cc-by-nc-sa-4.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ANKH_base + +This checkpoint packages the FastPLMs `ANKH` implementation. + +Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to- +sequence use. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, +`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: encoder or explicitly prepared decoder states | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ANKH_base/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ANKH_base" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ANKH_base` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`. An unavailable requested backend raises +instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +`Synthyra/ANKH_base` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH_base", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ANKH_base" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ANKH_base", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Encoder and sequence-to-sequence use + +`Synthyra/ANKH_base` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "Synthyra/ANKH_base" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +## Notes and limitations + +ANKH parity covers the official encoder and sequence-to-sequence heads. +AutoModelForMaskedLM exposes the separately named FastPLMs synthesized +masked-LM extension and is not an official ANKH head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ANKH_base` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `cdd8d30d88e5bf41f44e1eef4470d8e46607aba5f7c7c805b06c035b89c8c16f` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `ElnaggarLab/ankh-base` +- Artifact source: `official` +- State transform: `ankh_t5_to_fastplms_v1` +- Pinned upstreams: `ankh` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is +`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/ankh_large.md b/model_cards/ankh_large.md new file mode 100644 index 0000000..d98ba60 --- /dev/null +++ b/model_cards/ankh_large.md @@ -0,0 +1,318 @@ +--- +library_name: transformers +license: "cc-by-nc-sa-4.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ANKH_large + +This checkpoint packages the FastPLMs `ANKH` implementation. + +Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to- +sequence use. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, +`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: encoder or explicitly prepared decoder states | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ANKH_large/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ANKH_large" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ANKH_large` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`. An unavailable requested backend raises +instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +`Synthyra/ANKH_large` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "Synthyra/ANKH_large", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ANKH_large" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ANKH_large", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Encoder and sequence-to-sequence use + +`Synthyra/ANKH_large` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "Synthyra/ANKH_large" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +## Notes and limitations + +ANKH parity covers the official encoder and sequence-to-sequence heads. +AutoModelForMaskedLM exposes the separately named FastPLMs synthesized +masked-LM extension and is not an official ANKH head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ANKH_large` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `e498a2e9aea76ef784cbe3e596c6b3f5e9a40e209ad837f7e3207099e4d74483` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `ElnaggarLab/ankh-large` +- Artifact source: `official` +- State transform: `ankh_t5_to_fastplms_v1` +- Pinned upstreams: `ankh` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is +`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/boltz2.md b/model_cards/boltz2.md new file mode 100644 index 0000000..dcc3768 --- /dev/null +++ b/model_cards/boltz2.md @@ -0,0 +1,183 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/Boltz2 + +This checkpoint packages the FastPLMs `Boltz2` implementation. + +Accepted inputs are raw amino-acid sequences through the convenience API, or +prepared model features. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Unavailable for this structure-only checkpoint | +| Test-time training | Unavailable for this inference-only checkpoint | +| Attention variants | Supported: `eager` | +| Compliance | Unavailable: this provisional family has no compliance tier | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/Boltz2/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/Boltz2" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="eager", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/Boltz2` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `eager` explicitly. Declared variants are `eager`. An unavailable requested backend raises instead +of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family does not declare the `compliance` tier. Boltz2 remains provisional +and its structure checks must not be broadened into parity claims. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Protein structure prediction + +The high-level helper prepares a protein-only input, runs the declared Boltz2 +inference core, and returns coordinates and confidence fields: + +```python +import torch + +model = model.cuda().eval() +output = model.predict_structure( + amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG", + recycling_steps=3, + num_sampling_steps=50, + diffusion_samples=1, + seed=7, +) +model.save_as_cif(output, "prediction.cif") + +print(output.sample_atom_coords.shape) +print(output.plddt, output.ptm, output.iptm) +``` + +The validation boundary below describes the currently supported inference +subset and its provisional status. The helper scopes and restores Python, +NumPy, CPU Torch, and CUDA RNG state. Parameters and prepared features remain +FP32; supported CUDA inference executes inside BF16 autocast. + +## Notes and limitations + +Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared +inference-core state, feature preparation, and seeded execution remain tested, +but native-environment BF16 end-to-end inference currently exceeds the fixed +numerical-equivalence limits. FastPLMs therefore does not claim official +inference equivalence for this checkpoint yet. Work on that numerical gap +continues independently of the ESM++ and ESMFold2 release gates. + +## Runtime contract + +- Public input: Raw amino-acid sequences through the convenience API, or prepared model features +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/Boltz2` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `boltz-community/boltz-2` +- Artifact source: `fast` +- State transform: `boltz2_inference_core_v1` +- Pinned upstreams: `boltz` +- Release tiers: `structure`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm2_150m.md b/model_cards/dplm2_150m.md new file mode 100644 index 0000000..3207561 --- /dev/null +++ b/model_cards/dplm2_150m.md @@ -0,0 +1,288 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM2-150M + +This checkpoint packages the FastPLMs `DPLM2` implementation. + +Accepted inputs are tokenized amino-acid and structure tracks with explicit +modality boundaries. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM2-150M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM2-150M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM2-150M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`. An unavailable requested backend raises instead +of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM2-150M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM2-150M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Amino-acid and structure co-generation + +DPLM2 uses separate structure and amino-acid tracks with modality-specific +boundary and mask tokens: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM2-150M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +vocab = tokenizer.get_vocab() +l = 64 +structure = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +amino_acids = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +input_ids = torch.tensor([structure + amino_acids], device="cuda") + +with torch.inference_mode(): + generated = generator.generate(input_ids, max_iter=100)["output_tokens"] +print(generated.shape) +``` + +Generic `cls_token`, `eos_token`, `mask_token`, and `unk_token` aliases are +intentionally unset. Callers constructing multimodal tensors must choose the +amino-acid or structure token explicitly. Raw amino-acid sequences remain +supported by `model.embed_dataset(...)`. + +Plain `AutoModel` omits the optional ESM pooler because this co-generation +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +The checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) and [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) document the license +basis for the pretrained DPLM1 and DPLM2 weights. Complete publication remains +subject to all artifact, legal, parity, and atomic-publication preflights. + +## Runtime contract + +- Public input: Tokenized amino-acid and structure tracks with explicit modality boundaries +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM2-150M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm2_150m` +- Artifact source: `official` +- State transform: `dplm2_to_fastplms_v1` +- Tokenizer class: `fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm2_3b.md b/model_cards/dplm2_3b.md new file mode 100644 index 0000000..52992cd --- /dev/null +++ b/model_cards/dplm2_3b.md @@ -0,0 +1,294 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM2-3B + +This checkpoint packages the FastPLMs `DPLM2` implementation. + +Accepted inputs are tokenized amino-acid and structure tracks with explicit +modality boundaries. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM2-3B/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM2-3B" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM2-3B` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`. An unavailable requested backend raises instead +of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM2-3B" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM2-3B", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Amino-acid and structure co-generation + +DPLM2 uses separate structure and amino-acid tracks with modality-specific +boundary and mask tokens: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM2-3B" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +vocab = tokenizer.get_vocab() +l = 64 +structure = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +amino_acids = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +input_ids = torch.tensor([structure + amino_acids], device="cuda") + +with torch.inference_mode(): + generated = generator.generate(input_ids, max_iter=100)["output_tokens"] +print(generated.shape) +``` + +Generic `cls_token`, `eos_token`, `mask_token`, and `unk_token` aliases are +intentionally unset. Callers constructing multimodal tensors must choose the +amino-acid or structure token explicitly. Raw amino-acid sequences remain +supported by `model.embed_dataset(...)`. + +Plain `AutoModel` omits the optional ESM pooler because this co-generation +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +The checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) and [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) document the license +basis for the pretrained DPLM1 and DPLM2 weights. Complete publication remains +subject to all artifact, legal, parity, and atomic-publication preflights. + +## Notes and limitations + +The pinned official DPLM2-3B sampler fails before generation, so live +generation equivalence cannot be established for this checkpoint. State, +tokenizer, and inference parity remain required. + +## Runtime contract + +- Public input: Tokenized amino-acid and structure tracks with explicit modality boundaries +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `official_unavailable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM2-3B` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `8c46ec09115dbe6cbfb91d94ab5e906369d57e27fe620a7741c6f8cb1b6ca890` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm2_3b` +- Artifact source: `official` +- State transform: `dplm2_to_fastplms_v1` +- Tokenizer class: `fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm2_650m.md b/model_cards/dplm2_650m.md new file mode 100644 index 0000000..832c6fa --- /dev/null +++ b/model_cards/dplm2_650m.md @@ -0,0 +1,288 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM2-650M + +This checkpoint packages the FastPLMs `DPLM2` implementation. + +Accepted inputs are tokenized amino-acid and structure tracks with explicit +modality boundaries. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM2-650M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM2-650M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM2-650M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`. An unavailable requested backend raises instead +of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM2-650M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM2-650M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Amino-acid and structure co-generation + +DPLM2 uses separate structure and amino-acid tracks with modality-specific +boundary and mask tokens: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM2-650M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +vocab = tokenizer.get_vocab() +l = 64 +structure = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +amino_acids = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +input_ids = torch.tensor([structure + amino_acids], device="cuda") + +with torch.inference_mode(): + generated = generator.generate(input_ids, max_iter=100)["output_tokens"] +print(generated.shape) +``` + +Generic `cls_token`, `eos_token`, `mask_token`, and `unk_token` aliases are +intentionally unset. Callers constructing multimodal tensors must choose the +amino-acid or structure token explicitly. Raw amino-acid sequences remain +supported by `model.embed_dataset(...)`. + +Plain `AutoModel` omits the optional ESM pooler because this co-generation +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +The checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) and [README](https://github.com/bytedance/dplm/blob/main/README.md#overview) document the license +basis for the pretrained DPLM1 and DPLM2 weights. Complete publication remains +subject to all artifact, legal, parity, and atomic-publication preflights. + +## Runtime contract + +- Public input: Tokenized amino-acid and structure tracks with explicit modality boundaries +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM2-650M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Canonical transformed state SHA-256: `cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4` +- Conversion equality attestation: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm2_650m` +- Artifact source: `official` +- State transform: `dplm2_to_fastplms_v1` +- Tokenizer class: `fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm_150m.md b/model_cards/dplm_150m.md new file mode 100644 index 0000000..81df9b1 --- /dev/null +++ b/model_cards/dplm_150m.md @@ -0,0 +1,306 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM-150M + +This checkpoint packages the FastPLMs `DPLM` implementation. + +Accepted inputs are amino-acid sequences tokenized to masked or partially +masked residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM-150M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM-150M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM-150M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_3`. +An unavailable requested backend raises instead of silently switching +implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/DPLM-150M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM-150M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM-150M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Diffusion sequence generation + +DPLM defines the requested length from biological positions in a tokenized +input, masks those positions, and iteratively retains confident predictions: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM-150M" +tokenizer = AutoTokenizer.from_pretrained(model_id) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +input_ids = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() + +with torch.inference_mode(): + generated_ids = generator.generate(input_ids, max_iter=100) + +sequence = tokenizer.decode( + generated_ids[0], + skip_special_tokens=True, +).replace(" ", "") +print(sequence) +``` + +Omitting `max_iter` uses the official 500-step schedule. A shorter schedule +changes the sampling process rather than providing an equivalent faster mode. + +Plain `AutoModel` omits the optional ESM pooler because this diffusion +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) is Apache-2.0 and the +[README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +explicitly scopes the repository release to the pretrained DPLM1 and DPLM2 +weights. FastPLMs artifacts record `weights_license_status="resolved"` and +`redistributable=true`; complete publication is permitted only after all +artifact, legal, parity, and atomic-publication preflights pass. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to masked or partially masked residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM-150M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm_150m` +- Artifact source: `fast` +- State transform: `dplm_to_fastplms_v1` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm_3b.md b/model_cards/dplm_3b.md new file mode 100644 index 0000000..d0e02e7 --- /dev/null +++ b/model_cards/dplm_3b.md @@ -0,0 +1,306 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM-3B + +This checkpoint packages the FastPLMs `DPLM` implementation. + +Accepted inputs are amino-acid sequences tokenized to masked or partially +masked residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM-3B/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM-3B" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM-3B` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_3`. +An unavailable requested backend raises instead of silently switching +implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/DPLM-3B" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM-3B" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM-3B", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Diffusion sequence generation + +DPLM defines the requested length from biological positions in a tokenized +input, masks those positions, and iteratively retains confident predictions: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM-3B" +tokenizer = AutoTokenizer.from_pretrained(model_id) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +input_ids = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() + +with torch.inference_mode(): + generated_ids = generator.generate(input_ids, max_iter=100) + +sequence = tokenizer.decode( + generated_ids[0], + skip_special_tokens=True, +).replace(" ", "") +print(sequence) +``` + +Omitting `max_iter` uses the official 500-step schedule. A shorter schedule +changes the sampling process rather than providing an equivalent faster mode. + +Plain `AutoModel` omits the optional ESM pooler because this diffusion +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) is Apache-2.0 and the +[README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +explicitly scopes the repository release to the pretrained DPLM1 and DPLM2 +weights. FastPLMs artifacts record `weights_license_status="resolved"` and +`redistributable=true`; complete publication is permitted only after all +artifact, legal, parity, and atomic-publication preflights pass. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to masked or partially masked residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM-3B` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm_3b` +- Artifact source: `fast` +- State transform: `dplm_to_fastplms_v1` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/dplm_650m.md b/model_cards/dplm_650m.md new file mode 100644 index 0000000..f154c38 --- /dev/null +++ b/model_cards/dplm_650m.md @@ -0,0 +1,306 @@ +--- +library_name: transformers +license: "apache-2.0" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/DPLM-650M + +This checkpoint packages the FastPLMs `DPLM` implementation. + +Accepted inputs are amino-acid sequences tokenized to masked or partially +masked residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/DPLM-650M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/DPLM-650M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/DPLM-650M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_3`. +An unavailable requested backend raises instead of silently switching +implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/DPLM-650M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/DPLM-650M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/DPLM-650M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Diffusion sequence generation + +DPLM defines the requested length from biological positions in a tokenized +input, masks those positions, and iteratively retains confident predictions: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/DPLM-650M" +tokenizer = AutoTokenizer.from_pretrained(model_id) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +input_ids = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() + +with torch.inference_mode(): + generated_ids = generator.generate(input_ids, max_iter=100) + +sequence = tokenizer.decode( + generated_ids[0], + skip_special_tokens=True, +).replace(" ", "") +print(sequence) +``` + +Omitting `max_iter` uses the official 500-step schedule. A shorter schedule +changes the sampling process rather than providing an equivalent faster mode. + +Plain `AutoModel` omits the optional ESM pooler because this diffusion +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE](https://github.com/bytedance/dplm/blob/main/LICENSE) is Apache-2.0 and the +[README](https://github.com/bytedance/dplm/blob/main/README.md#overview) +explicitly scopes the repository release to the pretrained DPLM1 and DPLM2 +weights. FastPLMs artifacts record `weights_license_status="resolved"` and +`redistributable=true`; complete publication is permitted only after all +artifact, legal, parity, and atomic-publication preflights pass. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to masked or partially masked residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `required` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/DPLM-650M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `airkingbd/dplm_650m` +- Artifact source: `fast` +- State transform: `dplm_to_fastplms_v1` +- Pinned upstreams: `dplm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: Apache-2.0. The Hub model-card identifier is +`apache-2.0`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/e1_150m.md b/model_cards/e1_150m.md new file mode 100644 index 0000000..df13722 --- /dev/null +++ b/model_cards/e1_150m.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "other" +license_name: "profluent-e1-clickthrough-license-agreement" +license_link: "https://github.com/Profluent-AI/E1/blob/main/LICENSE" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/Profluent-E1-150M + +This checkpoint packages the FastPLMs `E1` implementation. + +Accepted inputs are raw amino-acid sequences prepared by the native E1 adapter. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: tokenizer-free raw-sequence preparation | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/Profluent-E1-150M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/Profluent-E1-150M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/Profluent-E1-150M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`, `flex_attention`. An unavailable requested +backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/Profluent-E1-150M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = sequence_model.prep_tokens.get_batch_kwargs( + sequences, + device=sequence_model.device, +) +biological = batch["sequence_ids"].ne(-1) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/Profluent-E1-150M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Tokenizer-free E1 input + +E1 has no tokenizer. The model retains native raw-sequence preparation, +boundary tokens, sequence positions, and retrieval-augmented context behavior. +The ordinary representation path accepts sequences directly: + +```python +result = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean",), +) +print(result[0].tensor.shape) +``` + +Lower-level masked-language-model calls must use the E1 batch preparer rather +than an `AutoTokenizer`. E1 launch messages and distributed legal files retain +the attribution required by the upstream agreement. + +## Runtime contract + +- Public input: Raw amino-acid sequences prepared by the native E1 adapter +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa`, `flex_attention` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/Profluent-E1-150M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `Profluent-Bio/E1-150m` +- Artifact source: `fast` +- State transform: `e1_to_fastplms_v1` +- Pinned upstreams: `e1` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: [Profluent-E1 Clickthrough License Agreement](https://github.com/Profluent-AI/E1/blob/main/LICENSE). The Hub model-card identifier is +`other`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/e1_300m.md b/model_cards/e1_300m.md new file mode 100644 index 0000000..26988b7 --- /dev/null +++ b/model_cards/e1_300m.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "other" +license_name: "profluent-e1-clickthrough-license-agreement" +license_link: "https://github.com/Profluent-AI/E1/blob/main/LICENSE" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/Profluent-E1-300M + +This checkpoint packages the FastPLMs `E1` implementation. + +Accepted inputs are raw amino-acid sequences prepared by the native E1 adapter. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: tokenizer-free raw-sequence preparation | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/Profluent-E1-300M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/Profluent-E1-300M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/Profluent-E1-300M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`, `flex_attention`. An unavailable requested +backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/Profluent-E1-300M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = sequence_model.prep_tokens.get_batch_kwargs( + sequences, + device=sequence_model.device, +) +biological = batch["sequence_ids"].ne(-1) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/Profluent-E1-300M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Tokenizer-free E1 input + +E1 has no tokenizer. The model retains native raw-sequence preparation, +boundary tokens, sequence positions, and retrieval-augmented context behavior. +The ordinary representation path accepts sequences directly: + +```python +result = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean",), +) +print(result[0].tensor.shape) +``` + +Lower-level masked-language-model calls must use the E1 batch preparer rather +than an `AutoTokenizer`. E1 launch messages and distributed legal files retain +the attribution required by the upstream agreement. + +## Runtime contract + +- Public input: Raw amino-acid sequences prepared by the native E1 adapter +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa`, `flex_attention` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/Profluent-E1-300M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `Profluent-Bio/E1-300m` +- Artifact source: `fast` +- State transform: `e1_to_fastplms_v1` +- Pinned upstreams: `e1` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: [Profluent-E1 Clickthrough License Agreement](https://github.com/Profluent-AI/E1/blob/main/LICENSE). The Hub model-card identifier is +`other`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/e1_600m.md b/model_cards/e1_600m.md new file mode 100644 index 0000000..3c1d0bc --- /dev/null +++ b/model_cards/e1_600m.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "other" +license_name: "profluent-e1-clickthrough-license-agreement" +license_link: "https://github.com/Profluent-AI/E1/blob/main/LICENSE" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/Profluent-E1-600M + +This checkpoint packages the FastPLMs `E1` implementation. + +Accepted inputs are raw amino-acid sequences prepared by the native E1 adapter. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Special: tokenizer-free raw-sequence preparation | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/Profluent-E1-600M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/Profluent-E1-600M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/Profluent-E1-600M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `sdpa`, `flex_attention`. An unavailable requested +backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/Profluent-E1-600M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = sequence_model.prep_tokens.get_batch_kwargs( + sequences, + device=sequence_model.device, +) +biological = batch["sequence_ids"].ne(-1) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/Profluent-E1-600M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Tokenizer-free E1 input + +E1 has no tokenizer. The model retains native raw-sequence preparation, +boundary tokens, sequence positions, and retrieval-augmented context behavior. +The ordinary representation path accepts sequences directly: + +```python +result = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean",), +) +print(result[0].tensor.shape) +``` + +Lower-level masked-language-model calls must use the E1 batch preparer rather +than an `AutoTokenizer`. E1 launch messages and distributed legal files retain +the attribution required by the upstream agreement. + +## Runtime contract + +- Public input: Raw amino-acid sequences prepared by the native E1 adapter +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `sdpa`, `flex_attention` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/Profluent-E1-600M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `Profluent-Bio/E1-600m` +- Artifact source: `fast` +- State transform: `e1_to_fastplms_v1` +- Pinned upstreams: `e1` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: [Profluent-E1 Clickthrough License Agreement](https://github.com/Profluent-AI/E1/blob/main/LICENSE). The Hub model-card identifier is +`other`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm2_150m.md b/model_cards/esm2_150m.md new file mode 100644 index 0000000..d60d617 --- /dev/null +++ b/model_cards/esm2_150m.md @@ -0,0 +1,297 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM2-150M + +This checkpoint packages the FastPLMs `ESM2` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM2-150M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM2-150M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM2-150M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESM2-150M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESM2-150M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-150M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/ESM2-150M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM2-150M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esm2_t30_150M_UR50D` +- Artifact source: `fast` +- State transform: `esm2_hf_to_fastplms_v1` +- Pinned upstreams: `fair-esm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm2_35m.md b/model_cards/esm2_35m.md new file mode 100644 index 0000000..672c285 --- /dev/null +++ b/model_cards/esm2_35m.md @@ -0,0 +1,297 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM2-35M + +This checkpoint packages the FastPLMs `ESM2` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM2-35M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM2-35M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM2-35M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESM2-35M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESM2-35M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-35M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/ESM2-35M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM2-35M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esm2_t12_35M_UR50D` +- Artifact source: `fast` +- State transform: `esm2_hf_to_fastplms_v1` +- Pinned upstreams: `fair-esm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm2_3b.md b/model_cards/esm2_3b.md new file mode 100644 index 0000000..708ad68 --- /dev/null +++ b/model_cards/esm2_3b.md @@ -0,0 +1,305 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM2-3B + +This checkpoint packages the FastPLMs `ESM2` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM2-3B/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM2-3B" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM2-3B` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESM2-3B" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESM2-3B" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-3B", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/ESM2-3B" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +## Notes and limitations + +The pinned default SDPA BF16 path uses a checkpoint-specific numeric +calibration: relative L2 target/hard limit 0.06/0.07, relative Q99.9 0.15/0.18, +first-percentile residue cosine 0.994/0.992, and pooled cosine 0.998/0.997. +Exact state identity and the global logits-distribution contract remain +required. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM2-3B` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esm2_t36_3B_UR50D` +- Artifact source: `fast` +- State transform: `esm2_hf_to_fastplms_v1` +- Pinned upstreams: `fair-esm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm2_650m.md b/model_cards/esm2_650m.md new file mode 100644 index 0000000..649955c --- /dev/null +++ b/model_cards/esm2_650m.md @@ -0,0 +1,297 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM2-650M + +This checkpoint packages the FastPLMs `ESM2` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM2-650M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM2-650M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM2-650M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESM2-650M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESM2-650M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-650M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/ESM2-650M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM2-650M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esm2_t33_650M_UR50D` +- Artifact source: `fast` +- State transform: `esm2_hf_to_fastplms_v1` +- Pinned upstreams: `fair-esm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm2_8m.md b/model_cards/esm2_8m.md new file mode 100644 index 0000000..f524840 --- /dev/null +++ b/model_cards/esm2_8m.md @@ -0,0 +1,297 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM2-8M + +This checkpoint packages the FastPLMs `ESM2` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, +`AutoModelForTokenClassification`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Supported: base weights with an untrained task head | +| Token classification | Supported: base weights with an untrained task head | +| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM2-8M/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM2-8M" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM2-8M` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESM2-8M" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +from transformers import AutoTokenizer +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESM2-8M" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head, so the separately trained +`classifier` is saved with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESM2-8M", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "Synthyra/ESM2-8M" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM2-8M` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esm2_t6_8M_UR50D` +- Artifact source: `fast` +- State transform: `esm2_hf_to_fastplms_v1` +- Pinned upstreams: `fair-esm` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esm3_small.md b/model_cards/esm3_small.md new file mode 100644 index 0000000..d335124 --- /dev/null +++ b/model_cards/esm3_small.md @@ -0,0 +1,238 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESM3_small + +This checkpoint packages the FastPLMs `ESM3` implementation. + +Accepted inputs are sequence, structure, and function tracks prepared through +the multimodal helpers. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESM3_small/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESM3_small" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESM3_small` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModel + +ttt_model = AutoModel.from_pretrained( + "Synthyra/ESM3_small", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## Sequence inference and masked-sequence generation + +ESM3 owns its sequence preparation. This example exercises the sequence track; +the public input contract also supports structure and function tracks through +the multimodal helpers: + +```python +import torch + +batch = model.tokenize_sequences( + ["MKTAYIAKQ", "GGGG"], + device=model.device, +) +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +print(output.logits.shape) +print(output.structure_logits.shape) +print(output.function_logits.shape) +``` + +When `return_dict=False`, ESM3 follows the standard base-model tuple prefix: +`last_hidden_state`, then requested `hidden_states` and `attentions`. Multimodal +logits and extensions follow that prefix. Prefer named fields for individual +tracks. + +Generate masked sequence positions with an explicit seed: + +```python +from fastplms.models.esm3.modeling_esm3 import FastESM3GenerationConfig + +config = FastESM3GenerationConfig( + num_steps=8, + temperature=1.0, + seed=7, +) +generated = model.generate("MK____A", config) +print(generated) +``` + +Underscores mark positions to generate. Model outputs are predictions over +tracks, not experimental measurements of structure or function. + +## Runtime contract + +- Public input: Sequence, structure, and function tracks prepared through the multimodal helpers +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESM3_small` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/esm3-sm-open-v1` +- Artifact source: `fast` +- State transform: `esm3_to_fastplms_v1` +- Pinned upstreams: `biohub-esm`, `biohub-transformers` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmc_6b.md b/model_cards/esmc_6b.md new file mode 100644 index 0000000..6c84dcc --- /dev/null +++ b/model_cards/esmc_6b.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMplusplus_6B + +This checkpoint packages the FastPLMs `ESMC` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMplusplus_6B/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMplusplus_6B" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMplusplus_6B` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESMplusplus_6B" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESMplusplus_6B", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## ESMC behavior + +This artifact exposes the Biohub ESMC sequence encoder and masked-language-model +head through Transformers. It is also the language-model family used by +ESMFold2. SDPA is the default and the recommended choice for highest numerical +fidelity. Flex Attention and FlashAttention 3 are supported, non-experimental +backends, but their BF16 arithmetic may be numerically divergent from SDPA. +Those deviations produce diagnostic warnings rather than strict parity +failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic +biological disagreement remain hard gates. + +The current GH200/aarch64 release environment validates eager, SDPA, and Flex. +Flash requests fail closed because compatible locked kernels are unavailable +on this platform. + +When `sequence_id` is supplied, it is authoritative for ESMC attention grouping +and padding, and `attention_mask` is ignored. Values greater than or equal to +zero are valid sequence-group IDs; `-1` denotes padding. Omit `sequence_id` to +use `attention_mask` as the padding contract. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flash_attention_2` | Supported | Unavailable on current GH200/aarch64 lock | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | +| `flash_attention_3` | Supported, numerically divergent | Unavailable on current GH200/aarch64 lock | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMplusplus_6B` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMC-6B` +- Artifact source: `fast` +- State transform: `esmc_to_fastplms_v1` +- Pinned upstreams: `biohub-esm`, `biohub-transformers` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmc_large.md b/model_cards/esmc_large.md new file mode 100644 index 0000000..8eb1fd7 --- /dev/null +++ b/model_cards/esmc_large.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMplusplus_large + +This checkpoint packages the FastPLMs `ESMC` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMplusplus_large/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMplusplus_large" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMplusplus_large` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESMplusplus_large" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESMplusplus_large", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## ESMC behavior + +This artifact exposes the Biohub ESMC sequence encoder and masked-language-model +head through Transformers. It is also the language-model family used by +ESMFold2. SDPA is the default and the recommended choice for highest numerical +fidelity. Flex Attention and FlashAttention 3 are supported, non-experimental +backends, but their BF16 arithmetic may be numerically divergent from SDPA. +Those deviations produce diagnostic warnings rather than strict parity +failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic +biological disagreement remain hard gates. + +The current GH200/aarch64 release environment validates eager, SDPA, and Flex. +Flash requests fail closed because compatible locked kernels are unavailable +on this platform. + +When `sequence_id` is supplied, it is authoritative for ESMC attention grouping +and padding, and `attention_mask` is ignored. Values greater than or equal to +zero are valid sequence-group IDs; `-1` denotes padding. Omit `sequence_id` to +use `attention_mask` as the padding contract. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flash_attention_2` | Supported | Unavailable on current GH200/aarch64 lock | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | +| `flash_attention_3` | Supported, numerically divergent | Unavailable on current GH200/aarch64 lock | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMplusplus_large` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMC-600M` +- Artifact source: `fast` +- State transform: `esmc_to_fastplms_v1` +- Pinned upstreams: `biohub-esm`, `biohub-transformers` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmc_small.md b/model_cards/esmc_small.md new file mode 100644 index 0000000..29866a0 --- /dev/null +++ b/model_cards/esmc_small.md @@ -0,0 +1,256 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMplusplus_small + +This checkpoint packages the FastPLMs `ESMC` implementation. + +Accepted inputs are amino-acid sequences tokenized to residue IDs. +Supported Transformers entry points are `AutoConfig`, `AutoModel`, +`AutoModelForMaskedLM`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Supported: shared ordered embedding API | +| Test-time training | Supported: low-rank masked-residue adaptation | +| Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMplusplus_small/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMplusplus_small" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMplusplus_small` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, +`flash_attention_3`. An unavailable requested backend raises instead of +silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "Synthyra/ESMplusplus_small" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import AutoModelForMaskedLM + +ttt_model = AutoModelForMaskedLM.from_pretrained( + "Synthyra/ESMplusplus_small", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +## ESMC behavior + +This artifact exposes the Biohub ESMC sequence encoder and masked-language-model +head through Transformers. It is also the language-model family used by +ESMFold2. SDPA is the default and the recommended choice for highest numerical +fidelity. Flex Attention and FlashAttention 3 are supported, non-experimental +backends, but their BF16 arithmetic may be numerically divergent from SDPA. +Those deviations produce diagnostic warnings rather than strict parity +failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic +biological disagreement remain hard gates. + +The current GH200/aarch64 release environment validates eager, SDPA, and Flex. +Flash requests fail closed because compatible locked kernels are unavailable +on this platform. + +When `sequence_id` is supplied, it is authoritative for ESMC attention grouping +and padding, and `attention_mask` is ignored. Values greater than or equal to +zero are valid sequence-group IDs; `-1` denotes padding. Omit `sequence_id` to +use `attention_mask` as the padding contract. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flash_attention_2` | Supported | Unavailable on current GH200/aarch64 lock | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | +| `flash_attention_3` | Supported, numerically divergent | Unavailable on current GH200/aarch64 lock | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Runtime contract + +- Public input: Amino-acid sequences tokenized to residue IDs +- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3` +- Precision policies: `default` +- BF16 execution: `static_parameters` +- Generation contract: `not_applicable` +- Artifact dependency set: `core` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMplusplus_small` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMC-300M` +- Artifact source: `fast` +- State transform: `esmc_to_fastplms_v1` +- Pinned upstreams: `biohub-esm`, `biohub-transformers` +- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmfold.md b/model_cards/esmfold.md new file mode 100644 index 0000000..6a4d2e1 --- /dev/null +++ b/model_cards/esmfold.md @@ -0,0 +1,176 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/FastESMFold + +This checkpoint packages the FastPLMs `ESMFold` implementation. + +Accepted inputs are raw amino-acid sequences through folding helpers, or +prepared residue tensors. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Unavailable for this structure-only checkpoint | +| Test-time training | Unavailable: the checkpoint has no trained MLM head | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/FastESMFold/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/FastESMFold" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/FastESMFold` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Protein structure prediction + +ESMFold accepts a raw sequence and returns structure tensors and confidence: + +```python +import torch + +model = model.cuda().eval() +with torch.inference_mode(): + output = model.infer( + "MKTLLILAVVAAALA", + num_recycles=4, + ) + +print(output["mean_plddt"]) + +summary = model.fold_protein( + "MKTLLILAVVAAALA", + return_pdb_string=True, +) +with open("prediction.pdb", "w", encoding="utf-8") as handle: + handle.write(summary["pdb_string"]) +print(summary["plddt"], summary["ptm"]) +``` + +FastPLMs does not expose ProteinTTT for ESMFold. The pinned folding checkpoint +does not contain a trained masked-language-model head for that objective, so +`ttt()` and TTT folding requests raise explicitly. + +## Runtime contract + +- Public input: Raw amino-acid sequences through folding helpers, or prepared residue tensors +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `default` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/FastESMFold` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `facebook/esmfold_v1` +- Artifact source: `fast` +- State transform: `esmfold_meta_to_fastplms_v1` +- Pinned upstreams: `fair-esm`, `openfold` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmfold2.md b/model_cards/esmfold2.md new file mode 100644 index 0000000..4bf6aaa --- /dev/null +++ b/model_cards/esmfold2.md @@ -0,0 +1,288 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMFold2 + +This checkpoint packages the FastPLMs `ESMFold2` implementation. + +Accepted inputs are raw amino-acid sequences or typed molecular-complex +specifications; low-level forward accepts prepared feature tensors. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Special: ESMC state mixture to 256-wide residue embeddings | +| Test-time training | Special: opt-in folding TTT on the ESMC backbone | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMFold2/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMFold2" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMFold2` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Alignment-conditioning contract + +This is a full 48-block ESMFold2 checkpoint. It supports both +single-sequence inference and optional MSA-conditioned inference. Typed +multichain and multimolecule inputs may attach an MSA to each applicable +protein chain. + + +## Protein folding + +The single-protein helper returns typed structure and confidence outputs: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = model.result_to_pdb(result) +cif_text = model.result_to_cif(result) +print(result.ptm, result.plddt.mean().item()) +``` + +No target structure is required. For complexes, construct the input from the +types exposed by the loaded artifact: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +The typed interface also supports RNA, protein MSAs, modifications, covalent +bonds, and distogram conditioning. The public schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning +instead of silently ignoring it. Prepared `ref_pos` values are component +reference geometries created during featurization, not target coordinates. +Predicted coordinates and confidence scores are outputs and do not establish +biochemical activity. + +## Learned representation and ESMC precision + +ESMFold2 applies its learned state mixture and projection as +`H: (b, l, 81, 2560) -> Z: (b, l, 256)`. Retrieve `Z` through the public +embedding API: + +```python +representations = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) +``` + +`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue +tensor per single-chain input. It rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti` in the embedding path. + +Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading. +`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only, +and strict: + +```python +model.reload_esmc(precision="fp8", device="cuda:0") +print(model.esmc_precision_status) +``` + +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Canonical BF16 weights are retained, and transient quantization state is never +serialized. + +The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex +Attention is supported and non-experimental but can be numerically divergent; +ESMFold2 does not advertise FlashAttention for the folding interface. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Hash-pinned CCD runtime asset + +Structure preparation requires `ccd.pkl` from +`biohub/ESMFold2`. The manifest pins +its 417,306,584-byte size and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` +under MIT terms. This is a trusted-deserialization boundary: FastPLMs only +allows the exact manifest repository/revision snapshot link to resolve within +that repository's contained blob directory; user-supplied asset and `cache_dir` +symlinks are rejected. The loader creates a private temporary snapshot, verifies +its size and SHA-256, and unpickles only that loader-owned snapshot, closing +path-replacement and in-place source-write races. Offline execution requires the +exact cache object and never downloads a replacement. + +## Optional folding TTT + +The standard and Fast checkpoints expose opt-in folding TTT on their ESMC +backbone: + +```python +adapted = model.fold_protein_ttt( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=50, + seed=7, + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +print(adapted.ttt_metrics) +``` + +Entering a gradient-enabled path reloads canonical BF16 ESMC weights. TTT adds +latency and memory, can worsen a prediction, and does not calibrate confidence +or establish biological validity. Folding TTT is result-scoped: its transient +ESMC adapter modules are excluded from checkpoint state, so it is not a generic +`save_pretrained` adapter-persistence path. + +## Runtime contract + +- Public input: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental) +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMFold2` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMFold2` +- Artifact source: `fast` +- State transform: `identity` +- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmfold2_experimental_cutoff2025.md b/model_cards/esmfold2_experimental_cutoff2025.md new file mode 100644 index 0000000..a2c0cd5 --- /dev/null +++ b/model_cards/esmfold2_experimental_cutoff2025.md @@ -0,0 +1,293 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMFold2-Experimental-Cutoff2025 + +This checkpoint packages the FastPLMs `ESMFold2` implementation. + +Accepted inputs are raw amino-acid sequences or typed molecular-complex +specifications; low-level forward accepts prepared feature tensors. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Special: ESMC state mixture to 256-wide residue embeddings | +| Test-time training | Unavailable for this experimental checkpoint | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMFold2-Experimental-Cutoff2025/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMFold2-Experimental-Cutoff2025" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMFold2-Experimental-Cutoff2025` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Alignment-conditioning contract + +This is a full 48-block ESMFold2 checkpoint. It supports both +single-sequence inference and optional MSA-conditioned inference. Typed +multichain and multimolecule inputs may attach an MSA to each applicable +protein chain. + + +## Protein folding + +The single-protein helper returns typed structure and confidence outputs: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = model.result_to_pdb(result) +cif_text = model.result_to_cif(result) +print(result.ptm, result.plddt.mean().item()) +``` + +No target structure is required. For complexes, construct the input from the +types exposed by the loaded artifact: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +The typed interface also supports RNA, protein MSAs, modifications, covalent +bonds, and distogram conditioning. The public schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning +instead of silently ignoring it. Prepared `ref_pos` values are component +reference geometries created during featurization, not target coordinates. +Predicted coordinates and confidence scores are outputs and do not establish +biochemical activity. + +## Learned representation and ESMC precision + +ESMFold2 applies its learned state mixture and projection as +`H: (b, l, 81, 2560) -> Z: (b, l, 256)`. Retrieve `Z` through the public +embedding API: + +```python +representations = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) +``` + +`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue +tensor per single-chain input. It rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti` in the embedding path. + +Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading. +`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only, +and strict: + +```python +model.reload_esmc(precision="fp8", device="cuda:0") +print(model.esmc_precision_status) +``` + +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Canonical BF16 weights are retained, and transient quantization state is never +serialized. + +The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex +Attention is supported and non-experimental but can be numerically divergent; +ESMFold2 does not advertise FlashAttention for the folding interface. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Hash-pinned CCD runtime asset + +Structure preparation requires `ccd.pkl` from +`biohub/ESMFold2`. The manifest pins +its 417,306,584-byte size and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` +under MIT terms. This is a trusted-deserialization boundary: FastPLMs only +allows the exact manifest repository/revision snapshot link to resolve within +that repository's contained blob directory; user-supplied asset and `cache_dir` +symlinks are rejected. The loader creates a private temporary snapshot, verifies +its size and SHA-256, and unpickles only that loader-owned snapshot, closing +path-replacement and in-place source-write races. Offline execution requires the +exact cache object and never downloads a replacement. + +## Test-time training + +This experimental checkpoint does not expose folding TTT. Use the corresponding +standard or Fast checkpoint when opt-in ESMC-backbone adaptation is required. + +## Binder-design research example + +The FastPLMs binder-design workflow uses the experimental Fast Cutoff2025 +checkpoint for differentiable inversion, both experimental Cutoff2025 +checkpoints as critics, and ESM++ as the sequence prior: + +![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png) + +```bash +python examples/binder_design_fastplms.py \ + --target-name pd-l1 \ + --binder-name minibinder \ + --batch-size 4 \ + --steps 150 \ + --output-dir artifacts/binder-design +``` + +The workflow ranks candidates by mean iPTM across the approved critics after +the minibinder isoelectric-point filter. These are model-based prioritization +signals, not experimental evidence of affinity or specificity. See the +[complete workflow](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md). + +## Runtime contract + +- Public input: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental) +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMFold2-Experimental-Cutoff2025` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMFold2-Experimental-Cutoff2025` +- Artifact source: `fast` +- State transform: `identity` +- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmfold2_experimental_fast_cutoff2025.md b/model_cards/esmfold2_experimental_fast_cutoff2025.md new file mode 100644 index 0000000..2301ee9 --- /dev/null +++ b/model_cards/esmfold2_experimental_fast_cutoff2025.md @@ -0,0 +1,297 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMFold2-Experimental-Fast-Cutoff2025 + +This checkpoint packages the FastPLMs `ESMFold2` implementation. + +Accepted inputs are raw amino-acid sequences or typed molecular-complex +specifications; low-level forward accepts prepared feature tensors. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Special: ESMC state mixture to 256-wide residue embeddings | +| Test-time training | Unavailable for this experimental checkpoint | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMFold2-Experimental-Fast-Cutoff2025/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMFold2-Experimental-Fast-Cutoff2025` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Alignment-conditioning contract + +This 24-block Fast checkpoint is inference-optimized for single-sequence +conditioning and was trained without MSA conditioning. It is not +MSA-conditioned and rejects `ProteinInput.msa` and low-level MSA-derived +features. Typed multichain and multimolecule inputs remain supported when +every protein chain uses `msa=None`. Use the corresponding full ESMFold2 +checkpoint for MSA-conditioned inference. This follows the official Biohub +architecture description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). + + +## Protein folding + +The single-protein helper returns typed structure and confidence outputs: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = model.result_to_pdb(result) +cif_text = model.result_to_cif(result) +print(result.ptm, result.plddt.mean().item()) +``` + +No target structure is required. For complexes, construct the input from the +types exposed by the loaded artifact: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +The typed interface also supports RNA, modifications, covalent bonds, and +distogram conditioning. Protein MSA inputs are not supported by this Fast +checkpoint; every protein chain must use `msa=None`. The public schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning +instead of silently ignoring it. Prepared `ref_pos` values are component +reference geometries created during featurization, not target coordinates. +Predicted coordinates and confidence scores are outputs and do not establish +biochemical activity. + +## Learned representation and ESMC precision + +ESMFold2 applies its learned state mixture and projection as +`H: (b, l, 81, 2560) -> Z: (b, l, 256)`. Retrieve `Z` through the public +embedding API: + +```python +representations = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) +``` + +`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue +tensor per single-chain input. It rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti` in the embedding path. + +Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading. +`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only, +and strict: + +```python +model.reload_esmc(precision="fp8", device="cuda:0") +print(model.esmc_precision_status) +``` + +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Canonical BF16 weights are retained, and transient quantization state is never +serialized. + +The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex +Attention is supported and non-experimental but can be numerically divergent; +ESMFold2 does not advertise FlashAttention for the folding interface. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Hash-pinned CCD runtime asset + +Structure preparation requires `ccd.pkl` from +`biohub/ESMFold2`. The manifest pins +its 417,306,584-byte size and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` +under MIT terms. This is a trusted-deserialization boundary: FastPLMs only +allows the exact manifest repository/revision snapshot link to resolve within +that repository's contained blob directory; user-supplied asset and `cache_dir` +symlinks are rejected. The loader creates a private temporary snapshot, verifies +its size and SHA-256, and unpickles only that loader-owned snapshot, closing +path-replacement and in-place source-write races. Offline execution requires the +exact cache object and never downloads a replacement. + +## Test-time training + +This experimental checkpoint does not expose folding TTT. Use the corresponding +standard or Fast checkpoint when opt-in ESMC-backbone adaptation is required. + +## Binder-design research example + +The FastPLMs binder-design workflow uses the experimental Fast Cutoff2025 +checkpoint for differentiable inversion, both experimental Cutoff2025 +checkpoints as critics, and ESM++ as the sequence prior: + +![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png) + +```bash +python examples/binder_design_fastplms.py \ + --target-name pd-l1 \ + --binder-name minibinder \ + --batch-size 4 \ + --steps 150 \ + --output-dir artifacts/binder-design +``` + +The workflow ranks candidates by mean iPTM across the approved critics after +the minibinder isoelectric-point filter. These are model-based prioritization +signals, not experimental evidence of affinity or specificity. See the +[complete workflow](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md). + +## Runtime contract + +- Public input: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental) +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMFold2-Experimental-Fast-Cutoff2025` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMFold2-Experimental-Fast-Cutoff2025` +- Artifact source: `fast` +- State transform: `identity` +- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/model_cards/esmfold2_fast.md b/model_cards/esmfold2_fast.md new file mode 100644 index 0000000..1668e2f --- /dev/null +++ b/model_cards/esmfold2_fast.md @@ -0,0 +1,292 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# Synthyra/ESMFold2-Fast + +This checkpoint packages the FastPLMs `ESMFold2` implementation. + +Accepted inputs are raw amino-acid sequences or typed molecular-complex +specifications; low-level forward accepts prepared feature tensors. +Supported Transformers entry points are `AutoConfig`, `AutoModel`. + +## Capabilities + +| Feature | Status | +| --- | --- | +| Sequence classification | Unavailable: no advertised AutoClass | +| Token classification | Unavailable: no advertised AutoClass | +| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model | +| Embeddings | Special: ESMC state mixture to 256-wide residue embeddings | +| Test-time training | Special: opt-in folding TTT on the ESMC backbone | +| Attention variants | Supported: `eager`, `sdpa`, `flex_attention` | +| Compliance | Declared: exact release evidence is required | + +A supported interface is not a pretrained downstream predictor. Classification +heads start untrained, and declared compliance metadata is not a claim that an +arbitrary local build passed its release gate. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMFold2-Fast/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMFold2-Fast" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMFold2-Fast` path and pass `local_files_only=True`. + +## Attention and compliance + +The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable +requested backend raises instead of silently switching implementations. +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +This family declares the `compliance` tier. Release evidence binds the exact +checkpoint, backend, dtype, hardware, inputs, and reference revision. + +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, get_peft_model + +peft_model = get_peft_model( + model, + LoraConfig( + r=8, + lora_alpha=16, + target_modules="all-linear", + ), +) +``` + +This checkpoint has no advertised classifier. Supply the task-specific +objective and preserve any new head through `modules_to_save`. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Alignment-conditioning contract + +This 24-block Fast checkpoint is inference-optimized for single-sequence +conditioning and was trained without MSA conditioning. It is not +MSA-conditioned and rejects `ProteinInput.msa` and low-level MSA-derived +features. Typed multichain and multimolecule inputs remain supported when +every protein chain uses `msa=None`. Use the corresponding full ESMFold2 +checkpoint for MSA-conditioned inference. This follows the official Biohub +architecture description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). + + +## Protein folding + +The single-protein helper returns typed structure and confidence outputs: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = model.result_to_pdb(result) +cif_text = model.result_to_cif(result) +print(result.ptm, result.plddt.mean().item()) +``` + +No target structure is required. For complexes, construct the input from the +types exposed by the loaded artifact: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +The typed interface also supports RNA, modifications, covalent bonds, and +distogram conditioning. Protein MSA inputs are not supported by this Fast +checkpoint; every protein chain must use `msa=None`. The public schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning +instead of silently ignoring it. Prepared `ref_pos` values are component +reference geometries created during featurization, not target coordinates. +Predicted coordinates and confidence scores are outputs and do not establish +biochemical activity. + +## Learned representation and ESMC precision + +ESMFold2 applies its learned state mixture and projection as +`H: (b, l, 81, 2560) -> Z: (b, l, 256)`. Retrieve `Z` through the public +embedding API: + +```python +representations = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) +``` + +`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue +tensor per single-chain input. It rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti` in the embedding path. + +Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading. +`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only, +and strict: + +```python +model.reload_esmc(precision="fp8", device="cuda:0") +print(model.esmc_precision_status) +``` + +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Canonical BF16 weights are retained, and transient quantization state is never +serialized. + +The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex +Attention is supported and non-experimental but can be numerically divergent; +ESMFold2 does not advertise FlashAttention for the folding interface. + +| Backend | Support | Measurement status | +| --- | --- | --- | +| `sdpa` | Recommended fidelity path | Pending release measurement | +| `eager` | Supported | Pending release measurement | +| `flex_attention` | Supported, numerically divergent | Pending release measurement | + +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). + + +## Hash-pinned CCD runtime asset + +Structure preparation requires `ccd.pkl` from +`biohub/ESMFold2`. The manifest pins +its 417,306,584-byte size and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` +under MIT terms. This is a trusted-deserialization boundary: FastPLMs only +allows the exact manifest repository/revision snapshot link to resolve within +that repository's contained blob directory; user-supplied asset and `cache_dir` +symlinks are rejected. The loader creates a private temporary snapshot, verifies +its size and SHA-256, and unpickles only that loader-owned snapshot, closing +path-replacement and in-place source-write races. Offline execution requires the +exact cache object and never downloads a replacement. + +## Optional folding TTT + +The standard and Fast checkpoints expose opt-in folding TTT on their ESMC +backbone: + +```python +adapted = model.fold_protein_ttt( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=50, + seed=7, + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +print(adapted.ttt_metrics) +``` + +Entering a gradient-enabled path reloads canonical BF16 ESMC weights. TTT adds +latency and memory, can worsen a prediction, and does not calibrate confidence +or establish biological validity. Folding TTT is result-scoped: its transient +ESMC adapter modules are excluded from checkpoint state, so it is not a generic +`save_pretrained` adapter-persistence path. + +## Runtime contract + +- Public input: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors +- Advertised AutoClasses: `AutoConfig`, `AutoModel` +- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained` +- Attention implementations: `eager`, `sdpa`, `flex_attention` +- Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental) +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Artifact dependency set: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Release record + +- FastPLMs weights: `Synthyra/ESMFold2-Fast` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +- Official checkpoint: `biohub/ESMFold2-Fast` +- Artifact source: `fast` +- State transform: `identity` +- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..5a94abe --- /dev/null +++ b/mypy.ini @@ -0,0 +1,4 @@ +[mypy] +python_version = 3.11 +strict = true +warn_unreachable = true diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a4e09eb --- /dev/null +++ b/pytest.ini @@ -0,0 +1,20 @@ +[pytest] +minversion = 9.0 +testpaths = tests +pythonpath = + src + . +addopts = -ra --strict-markers --strict-config +markers = + cpu_contract: hermetic, deterministic CPU contract required before merge + gpu: requires a CUDA accelerator + slow: loads large checkpoints or performs long numerical comparisons + large: requires at least 24 GiB of accelerator memory + network: performs external network access + checkpoint: loads a non-micro checkpoint + reference: executes a pinned official implementation + structure: exercises protein structure models + compliance: compares against a pinned official upstream implementation + feature: exercises generation, multimodal, TTT, binder, and adapter workflows + artifact: validates a built remote-code artifact in an isolated environment + benchmark: short benchmark-harness smoke test; performance runs live outside pytest diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 8280b4d..0000000 --- a/requirements.txt +++ /dev/null @@ -1,47 +0,0 @@ -# Local GPU installs should install torch==2.11.0 and torchvision==0.26.0 -# from the cu128 PyTorch index before this file. Dockerfiles enforce this pin. -matplotlib==3.10.9 -tqdm==4.67.3 -# testing -pytest==9.0.3 -# standard -numpy==1.26.4 -einops==0.8.2 -scikit-learn==1.8.0 -scipy==1.17.1 -seaborn==0.13.2 -attrs==26.1.0 -biopython==1.87 -pandas==3.0.3 -# huggingface -tensorflow==2.20.0 -transformers==4.57.6 -accelerate==1.12.0 -datasets==4.5.0 -hf_transfer==0.1.9 -hf-xet==1.5.0 -peft==0.19.1 -sentencepiece==0.2.1 -# embedding mixin -networkx==3.6.1 -# for boltz -omegaconf==2.3.0 -rich==15.0.0 -trifast==0.1.13 -# ESM++ / ESM3 / ESMFold2 -abnumber==0.4.4 -biotite==1.6.0 -brotli==1.2.0 -cloudpathlib==0.24.0 -dna_features_viewer==3.1.5 -httpx==0.28.1 -msgpack==1.1.2 -msgpack-numpy==0.4.8 -pydssp==0.9.1 -pygtrie==2.5.0 -py3dmol==2.5.5 -rdkit==2026.3.2 -tenacity==9.1.4 -zstd==1.5.7.3 -# E1 -kernels==0.12.3 diff --git a/requirements/README.md b/requirements/README.md new file mode 100644 index 0000000..9a393ca --- /dev/null +++ b/requirements/README.md @@ -0,0 +1,40 @@ +# Dependency profiles + +FastPLMs is loaded from Hugging Face model repositories with +`trust_remote_code=True`. This repository is a source, test, and dependency +workspace, not an installable Python distribution. + +`core.in` and `features/*.in` are the direct dependency declarations. +`profiles/*.in` compose those declarations for the environments exercised by +the repository. Validation commands constrain Torch and Transformers with +`constraints/validation.txt`. + +Create a local validation environment with: + +```bash +uv venv +uv pip install \ + -r requirements/profiles/cpu-validation.in \ + -c requirements/constraints/validation.txt \ + --torch-backend cpu +``` + +CUDA container profiles use the same command without `--torch-backend cpu`. +The FP8 profiles additionally pass +`--overrides requirements/overrides/cuda.txt`. + +When a fully transitive lock is needed for a maintained environment, compile +the corresponding profile rather than locking every mutually incompatible +feature together: + +```bash +uv pip compile \ + requirements/profiles/candidate.in \ + -c requirements/constraints/validation.txt \ + --universal \ + --generate-hashes \ + -o requirements/locks/candidate.txt +``` + +Generated locks are environment-specific test inputs. Public Hugging Face +artifacts should expose only their direct runtime dependencies. diff --git a/requirements/constraints/validation.txt b/requirements/constraints/validation.txt new file mode 100644 index 0000000..d06f1de --- /dev/null +++ b/requirements/constraints/validation.txt @@ -0,0 +1,2 @@ +torch==2.13.0 +transformers==5.13.0 diff --git a/requirements/core.in b/requirements/core.in new file mode 100644 index 0000000..f191840 --- /dev/null +++ b/requirements/core.in @@ -0,0 +1,8 @@ +torch>=2.13,<2.14 +transformers>=5.13,<5.14 +huggingface-hub>=0.34,<2 +tokenizers>=0.22,<0.23 +safetensors>=0.5,<1 +numpy>=1.26,<3 +einops>=0.8,<1 +tqdm>=4.67,<5 diff --git a/requirements/features/binder.in b/requirements/features/binder.in new file mode 100644 index 0000000..a5787a7 --- /dev/null +++ b/requirements/features/binder.in @@ -0,0 +1,4 @@ +abnumber==0.4.4 +anarcii==2.0.8 +pandas>=3.0,<3.1 +pyarrow>=25,<26 diff --git a/requirements/features/cpu.in b/requirements/features/cpu.in new file mode 100644 index 0000000..b4969a0 --- /dev/null +++ b/requirements/features/cpu.in @@ -0,0 +1 @@ +torch==2.13.0 diff --git a/requirements/features/cueq.in b/requirements/features/cueq.in new file mode 100644 index 0000000..c1d48cc --- /dev/null +++ b/requirements/features/cueq.in @@ -0,0 +1,3 @@ +cuequivariance==0.10.0; platform_system == "Linux" +cuequivariance-torch==0.10.0; platform_system == "Linux" +cuequivariance-ops-torch-cu13==0.10.0; platform_system == "Linux" diff --git a/requirements/features/dev.in b/requirements/features/dev.in new file mode 100644 index 0000000..3b1e040 --- /dev/null +++ b/requirements/features/dev.in @@ -0,0 +1,5 @@ +accelerate>=1.10,<2 +mypy>=1.18,<2 +pytest>=9,<10 +pytest-xdist>=3.8,<4 +ruff>=0.14,<1 diff --git a/requirements/features/flash.in b/requirements/features/flash.in new file mode 100644 index 0000000..0716453 --- /dev/null +++ b/requirements/features/flash.in @@ -0,0 +1 @@ +kernels>=0.15,<0.16 diff --git a/requirements/features/fp8.in b/requirements/features/fp8.in new file mode 100644 index 0000000..4fbcbb6 --- /dev/null +++ b/requirements/features/fp8.in @@ -0,0 +1,4 @@ +accelerate>=1.10,<2 +transformer-engine==2.12.0; platform_system == "Linux" +transformer-engine-cu13==2.12.0; platform_system == "Linux" +transformer-engine-torch==2.12.0; platform_system == "Linux" diff --git a/requirements/features/reporting.in b/requirements/features/reporting.in new file mode 100644 index 0000000..6b704d5 --- /dev/null +++ b/requirements/features/reporting.in @@ -0,0 +1,4 @@ +matplotlib>=3.10,<4 +scikit-learn>=1.7,<2 +scipy>=1.15,<2 +seaborn>=0.13,<1 diff --git a/requirements/features/structure.in b/requirements/features/structure.in new file mode 100644 index 0000000..ff3d520 --- /dev/null +++ b/requirements/features/structure.in @@ -0,0 +1,10 @@ +accelerate>=1.10,<2 +biopython>=1.85,<2 +biotite>=1.4,<2 +brotli>=1.1,<2 +msgpack>=1.1,<2 +msgpack-numpy>=0.4.8,<1 +omegaconf>=2.3,<3 +rdkit>=2025.9,<2027 +scipy>=1.15,<2 +zstandard>=0.23,<1 diff --git a/requirements/features/train.in b/requirements/features/train.in new file mode 100644 index 0000000..4856c22 --- /dev/null +++ b/requirements/features/train.in @@ -0,0 +1,3 @@ +accelerate>=1.10,<2 +datasets>=4,<5 +peft>=0.18,<1 diff --git a/requirements/overrides/cuda.txt b/requirements/overrides/cuda.txt new file mode 100644 index 0000000..8507680 --- /dev/null +++ b/requirements/overrides/cuda.txt @@ -0,0 +1 @@ +transformer-engine-cu12; sys_platform == "never" diff --git a/requirements/profiles/artifact.in b/requirements/profiles/artifact.in new file mode 100644 index 0000000..e69a49f --- /dev/null +++ b/requirements/profiles/artifact.in @@ -0,0 +1,4 @@ +-r ../core.in +-r ../features/dev.in +-r ../features/flash.in +-r ../features/structure.in diff --git a/requirements/profiles/binder.in b/requirements/profiles/binder.in new file mode 100644 index 0000000..a87e829 --- /dev/null +++ b/requirements/profiles/binder.in @@ -0,0 +1,3 @@ +-r ../core.in +-r ../features/structure.in +-r ../features/binder.in diff --git a/requirements/profiles/candidate-fp8.in b/requirements/profiles/candidate-fp8.in new file mode 100644 index 0000000..0688fd3 --- /dev/null +++ b/requirements/profiles/candidate-fp8.in @@ -0,0 +1,2 @@ +-r candidate-structure.in +-r ../features/fp8.in diff --git a/requirements/profiles/candidate-structure.in b/requirements/profiles/candidate-structure.in new file mode 100644 index 0000000..254ebdb --- /dev/null +++ b/requirements/profiles/candidate-structure.in @@ -0,0 +1,6 @@ +-r ../core.in +-r ../features/dev.in +-r ../features/flash.in +-r ../features/structure.in +-r ../features/cueq.in +-r ../features/train.in diff --git a/requirements/profiles/candidate.in b/requirements/profiles/candidate.in new file mode 100644 index 0000000..1b369a8 --- /dev/null +++ b/requirements/profiles/candidate.in @@ -0,0 +1,4 @@ +-r ../core.in +-r ../features/dev.in +-r ../features/flash.in +-r ../features/train.in diff --git a/requirements/profiles/cpu-validation.in b/requirements/profiles/cpu-validation.in new file mode 100644 index 0000000..b2c1693 --- /dev/null +++ b/requirements/profiles/cpu-validation.in @@ -0,0 +1,5 @@ +-r ../core.in +-r ../features/cpu.in +-r ../features/dev.in +-r ../features/structure.in +-r ../features/train.in diff --git a/requirements/profiles/reporting.in b/requirements/profiles/reporting.in new file mode 100644 index 0000000..4fa2a10 --- /dev/null +++ b/requirements/profiles/reporting.in @@ -0,0 +1,3 @@ +-r ../core.in +-r ../features/train.in +-r ../features/reporting.in diff --git a/requirements/profiles/runtime-fp8.in b/requirements/profiles/runtime-fp8.in new file mode 100644 index 0000000..3036e7a --- /dev/null +++ b/requirements/profiles/runtime-fp8.in @@ -0,0 +1,3 @@ +-r ../core.in +-r ../features/structure.in +-r ../features/fp8.in diff --git a/requirements/profiles/runtime.in b/requirements/profiles/runtime.in new file mode 100644 index 0000000..1662d87 --- /dev/null +++ b/requirements/profiles/runtime.in @@ -0,0 +1 @@ +-r ../core.in diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..8a0c270 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,6 @@ +line-length = 100 +target-version = "py311" +exclude = ["vendor/upstream"] + +[lint] +select = ["E", "F", "UP", "B", "SIM", "RUF"] diff --git a/src/fastplms/__init__.py b/src/fastplms/__init__.py new file mode 100644 index 0000000..c5eebf7 --- /dev/null +++ b/src/fastplms/__init__.py @@ -0,0 +1,49 @@ +"""FastPLMs public package interface. + +The module uses lazy exports so importing :mod:`fastplms` does not initialize +Torch, download checkpoints, construct tokenizers, or compile kernels. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +__version__ = "1.0.0" + +_LAZY_EXPORTS = { + "CheckpointSource": ("fastplms.registry", "CheckpointSource"), + "EmbeddingInput": ("fastplms.embeddings", "EmbeddingInput"), + "EmbeddingRecord": ("fastplms.embeddings", "EmbeddingRecord"), + "EmbeddingResult": ("fastplms.embeddings", "EmbeddingResult"), + "FileDigest": ("fastplms.registry", "FileDigest"), + "ModelFamily": ("fastplms.registry", "ModelFamily"), + "ModelRegistry": ("fastplms.registry", "ModelRegistry"), + "ModelSpec": ("fastplms.registry", "ModelSpec"), + "OracleAsset": ("fastplms.registry", "OracleAsset"), + "RegistryError": ("fastplms.registry", "RegistryError"), + "RuntimeProfile": ("fastplms.runtime", "RuntimeProfile"), + "UpstreamSource": ("fastplms.registry", "UpstreamSource"), + "embed_dataset": ("fastplms.embeddings", "embed_dataset"), + "get_model_registry": ("fastplms.registry", "get_model_registry"), + "get_model_spec": ("fastplms.registry", "get_model_spec"), + "load_model_registry": ("fastplms.registry", "load_model_registry"), + "runtime_profile": ("fastplms.runtime", "runtime_profile"), +} + +__all__ = ["__version__", *_LAZY_EXPORTS] + + +def __getattr__(name: str) -> Any: + try: + module_name, attribute_name = _LAZY_EXPORTS[name] + except KeyError as error: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from error + value = getattr(import_module(module_name), attribute_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()).union(__all__)) diff --git a/src/fastplms/attention/__init__.py b/src/fastplms/attention/__init__.py new file mode 100644 index 0000000..b0ec8ab --- /dev/null +++ b/src/fastplms/attention/__init__.py @@ -0,0 +1,64 @@ +"""Shared attention backends, masks, and optional optimized kernels.""" + +from ._core import ( + VALID_ATTENTION_BACKENDS, + AttentionBackend, + BlockMask, + _ensure_flash_kernels_loaded, + _get_flex_attention_fn, + _get_flex_block_mask, + _kernels_flash_forward, + _kernels_flash_varlen_forward, + _unpad_input, + bool_to_additive_mask, + clear_flex_attention_caches, + create_block_mask, + flex_attention, + get_attention_mask, + get_attn_implementation, + index_first_axis, + index_put_first_axis, + kernels_flash_attention_func, + pad_input, + resolve_attention_backend, + resolve_attention_backend_for_call, + set_config_attn_implementation, + warn_attention_backend_fallback, +) +from .interfaces import ( + FASTPLMS_ATTENTION_FUNCTIONS, + FASTPLMS_ATTENTION_MASKS, + FastPLMsAttentionMixin, + validate_transformers_attention_interfaces, +) + + +__all__ = [ + "FASTPLMS_ATTENTION_FUNCTIONS", + "FASTPLMS_ATTENTION_MASKS", + "VALID_ATTENTION_BACKENDS", + "AttentionBackend", + "BlockMask", + "FastPLMsAttentionMixin", + "_ensure_flash_kernels_loaded", + "_get_flex_attention_fn", + "_get_flex_block_mask", + "_kernels_flash_forward", + "_kernels_flash_varlen_forward", + "_unpad_input", + "bool_to_additive_mask", + "clear_flex_attention_caches", + "create_block_mask", + "flex_attention", + "get_attention_mask", + "get_attn_implementation", + "index_first_axis", + "index_put_first_axis", + "kernels_flash_attention_func", + "pad_input", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "set_config_attn_implementation", + "validate_transformers_attention_interfaces", + "warn_attention_backend_fallback", +] diff --git a/src/fastplms/attention/_core.py b/src/fastplms/attention/_core.py new file mode 100644 index 0000000..14cb220 --- /dev/null +++ b/src/fastplms/attention/_core.py @@ -0,0 +1,800 @@ +"""Low-level attention kernels and mask construction. + +The public backend contract lives in :mod:`fastplms.attention`. Optional +kernels are resolved only after a caller explicitly requests them, so importing +FastPLMs never downloads or compiles code. +""" + +from __future__ import annotations + +import warnings +import torch +from collections import OrderedDict +from collections.abc import Callable +from enum import Enum +from threading import RLock +from einops import rearrange +from torch.nn import functional as F + +from ._kernel_lock import load_locked_kernel + + +try: + from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention +except ImportError: + create_block_mask = None + flex_attention = None + BlockMask = None + +_MAX_FLEX_CACHE_ENTRIES = 128 +_compiled_flex_attention: OrderedDict[tuple, object] = OrderedDict() +_flex_block_masks: OrderedDict[tuple, BlockMask] = OrderedDict() +_flex_cache_lock = RLock() + + +def _remember(cache: OrderedDict, key: tuple, value): + """Insert an item into a bounded least-recently-used cache.""" + cache[key] = value + cache.move_to_end(key) + while len(cache) > _MAX_FLEX_CACHE_ENTRIES: + cache.popitem(last=False) + return value + + +def clear_flex_attention_caches() -> None: + """Drop FastPLMs-owned compiled Flex callables and block masks. + + This deliberately does not call :func:`torch.compiler.reset`, which would + clear process-global Torch compilation state owned by unrelated models. + Active forwards retain their local references and can complete safely. + """ + + with _flex_cache_lock: + _compiled_flex_attention.clear() + _flex_block_masks.clear() + + +def _get_flex_attention_fn( + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + shape: tuple[int, ...] | None = None, + sequence_lengths: tuple[int, ...] | None = None, + mask_semantics: str = "padding", +): + """Return a compiled Flex callable for an explicit execution signature. + + Compilation depends on execution shape, device, dtype, and mask semantics. + Per-example padding lengths are represented by the ``BlockMask`` argument + and must not create a new compiled graph for every batch composition. + """ + if flex_attention is None: + return None + # Retain the keyword for compatibility with remote-code artifacts while + # deliberately excluding data-dependent lengths from the compile key. + del sequence_lengths + flex_mod = torch.nn.attention.flex_attention + if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False): + return flex_attention + key = ( + None if device is None else str(device), + None if dtype is None else str(dtype), + shape, + mask_semantics, + ) + with _flex_cache_lock: + compiled = _compiled_flex_attention.get(key) + if compiled is None: + compiled = torch.compile(flex_attention, dynamic=False) + _remember(_compiled_flex_attention, key, compiled) + else: + _compiled_flex_attention.move_to_end(key) + return compiled + + +def _get_flex_block_mask( + *, + mask_pattern: torch.Tensor, + batch_size: int, + query_length: int, + key_value_length: int, + device: torch.device, + dtype: torch.dtype | None, + mask_semantics: str, + mask_mod: Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + torch.Tensor, + ], +) -> BlockMask: + """Return a bounded, exact-pattern cached Flex ``BlockMask``. + + The complete pattern is transferred to the host once to avoid a CUDA + synchronization per batch row. Execution dtype remains part of the key + because compiled Flex plans can specialize on it even though the pattern + tensor itself is boolean or integer. + """ + if create_block_mask is None: + raise RuntimeError( + "'flex_attention' was requested, but torch.create_block_mask is unavailable." + ) + pattern = mask_pattern.detach().to(device=device).contiguous() # mask_pattern.shape + # One device-to-host transfer is required for an exact cache identity. Use + # the contiguous buffer directly instead of materializing one Python int + # per byte, which is prohibitively expensive for long batched sequences. + host_pattern = pattern.to(device="cpu").contiguous() # mask_pattern.shape + pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C") # bytes + cache_key = ( + str(device), + None if dtype is None else str(dtype), + (batch_size, query_length, key_value_length), + str(pattern.dtype), + pattern_bytes, + mask_semantics, + ) + with _flex_cache_lock: + flex_block_mask = _flex_block_masks.get(cache_key) + if flex_block_mask is None: + flex_block_mask = create_block_mask( + mask_mod, + batch_size, + 1, + query_length, + key_value_length, + device=device, + ) + _remember(_flex_block_masks, cache_key, flex_block_mask) + else: + _flex_block_masks.move_to_end(cache_key) + return flex_block_mask + + +# Hugging Face `kernels` exposes slightly different APIs for FlashAttention 2 +# and 3. Detect the loaded variant once so every caller uses the same dispatch. +def _infer_kernels_flash_variant(kernel) -> str | None: + if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"): + return "flash_attn2" + if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"): + return "flash_attn3" + return None + + +def _load_kernels_flash(implementation: str) -> tuple[object, str]: + """Load exactly the requested FlashAttention kernel. + + Loading is deferred until backend selection. A FlashAttention-2 request + never falls through to FlashAttention-3, or vice versa. + """ + from fastplms.registry import get_model_registry + + kernel_spec = get_model_registry().attention_kernels[implementation] + repository = kernel_spec.repository + try: + flash_kernel = load_locked_kernel(repository, kernel_spec.revision) + except Exception as error: + raise RuntimeError( + f"Unable to load the manifest-pinned kernel " + f"{repository}@{kernel_spec.revision} for {implementation!r}." + ) from error + flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel) + if flash_kernel_variant != kernel_spec.expected_variant: + raise RuntimeError( + f"{repository}@{kernel_spec.revision} exposed {flash_kernel_variant!r}; " + f"expected {kernel_spec.expected_variant!r}." + ) + if not all( + callable(getattr(flash_kernel, name, None)) + for name in ("flash_attn_func", "flash_attn_varlen_func") + ): + raise RuntimeError( + f"{repository}@{kernel_spec.revision} does not expose the " + "autograd-enabled flash_attn_func and flash_attn_varlen_func APIs." + ) + return flash_kernel, flash_kernel_variant + + +_FLASH_KERNELS: dict[str, tuple[object, str]] = {} + + +def _validate_kernels_flash_dtype( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + implementation: str, +) -> torch.dtype: + """Reject dtypes outside the immutable kernel manifest before dispatch.""" + + # query_states, key_states, value_states: (b, l, h, d) or (t, h, d) + tensor_dtypes = {query_states.dtype, key_states.dtype, value_states.dtype} + if len(tensor_dtypes) != 1: + observed = ", ".join(sorted(str(dtype) for dtype in tensor_dtypes)) + raise RuntimeError( + f"{implementation!r} requires Q, K, and V to share one dtype; received {observed}." + ) + runtime_dtype = query_states.dtype + if ( + runtime_dtype == torch.float32 + and query_states.is_cuda + and torch.is_autocast_enabled("cuda") + ): + runtime_dtype = torch.get_autocast_dtype("cuda") + dtype_names = { + torch.float32: "float32", + torch.bfloat16: "bfloat16", + torch.float16: "float16", + } + runtime_dtype_name = dtype_names.get(runtime_dtype, str(runtime_dtype)) + from fastplms.registry import get_model_registry + + supported = get_model_registry().attention_kernels[implementation].dtypes + if runtime_dtype_name not in supported: + expected = ", ".join(supported) + raise RuntimeError( + f"{implementation!r} supports only manifest-declared dtype(s) {expected}; " + f"received {runtime_dtype_name}. Use CUDA BF16 autocast for FP32-resident " + "models." + ) + return runtime_dtype + + +def _validate_kernels_flash_device( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + implementation: str, +) -> torch.device: + """Require Q, K, and V on one CUDA device before loading a kernel.""" + + # query_states, key_states, value_states: (b, l, h, d) or (t, h, d) + devices = (query_states.device, key_states.device, value_states.device) + if len(set(devices)) != 1: + observed = ", ".join(str(device) for device in devices) + raise RuntimeError( + f"{implementation!r} requires Q, K, and V on one device; received {observed}." + ) + device = devices[0] + if device.type != "cuda" or not all( + tensor.is_cuda for tensor in (query_states, key_states, value_states) + ): + raise RuntimeError( + f"{implementation!r} requires CUDA Q, K, and V; received device {device}." + ) + return device + + +def _ensure_flash_kernels_loaded(implementation: str) -> tuple[object, str]: + cached = _FLASH_KERNELS.get(implementation) + if cached is not None: + return cached + loaded = _load_kernels_flash(implementation) + _FLASH_KERNELS[implementation] = loaded + return loaded + + +def _kernels_flash_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + causal: bool = False, + softmax_scale: float | None = None, + implementation: str = "flash_attention_3", +) -> torch.Tensor: + """Flash-attention forward, optionally overriding the softmax scale. + + When `softmax_scale is None`, the flash kernel applies its default + `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already + pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold). + Failing to override when Q is pre-scaled applies the scale twice and breaks + parity with eager attention and SDPA. + """ + # query_states, key_states, value_states: (b, l, h, d) + flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation) + if flash_kernel_variant == "flash_attn2": + output = flash_kernel.flash_attn_func( # (b, l, h, d) or tuple with that first + q=query_states, + k=key_states, + v=value_states, + dropout_p=0.0, + softmax_scale=softmax_scale, + causal=causal, + ) + return output[0] if isinstance(output, tuple) else output # (b, l, h, d) + if flash_kernel_variant == "flash_attn3": + output = flash_kernel.flash_attn_func( # (b, l, h, d) or tuple with that first + q=query_states, + k=key_states, + v=value_states, + softmax_scale=softmax_scale, + causal=causal, + ) + if isinstance(output, tuple): + return output[0] # (b, l, h, d) + return output # (b, l, h, d) + raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}") + + +def _kernels_flash_varlen_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_in_batch_q: int, + max_seqlen_in_batch_k: int, + causal: bool = False, + softmax_scale: float | None = None, + implementation: str = "flash_attention_3", +) -> torch.Tensor: + """Varlen flash-attention forward, optionally overriding the softmax scale. + + See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be + passed when Q has been pre-scaled by the caller. + """ + # query_states, key_states, value_states: (t, h, d) + # cu_seqlens_q, cu_seqlens_k: (b + 1,) + flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation) + if flash_kernel_variant == "flash_attn2": + output = flash_kernel.flash_attn_varlen_func( # (t, h, d) or tuple with that first + q=query_states, + k=key_states, + v=value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=0.0, + softmax_scale=softmax_scale, + causal=causal, + ) + return output[0] if isinstance(output, tuple) else output # (t, h, d) + if flash_kernel_variant == "flash_attn3": + output = flash_kernel.flash_attn_varlen_func( # (t, h, d) or tuple with that first + q=query_states, + k=key_states, + v=value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + softmax_scale=softmax_scale, + causal=causal, + ) + if isinstance(output, tuple): + return output[0] # (t, h, d) + return output # (t, h, d) + raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}") + + +# Varlen flash attention runs only on real tokens. These helpers remove padding +# before the kernel call and restore the original padded batch shape afterward. +class IndexFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, input, indices) -> torch.Tensor: + # input: (n, ...); indices: (m,) + ctx.save_for_backward(indices) + if input.ndim < 2: + raise ValueError( + "index_first_axis input must have at least two dimensions; " + f"received shape {tuple(input.shape)}." + ) + if indices.ndim != 1: + raise ValueError( + "index_first_axis indices must be one-dimensional; " + f"received shape {tuple(indices.shape)}." + ) + ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] + second_dim = other_shape.numel() + return torch.gather( # (m, ...) + rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim) + ).reshape(-1, *other_shape) + + @staticmethod + def backward(ctx, grad_output) -> tuple[torch.Tensor, None]: + # grad_output: (m, ...) + (indices,) = ctx.saved_tensors + if grad_output.ndim < 2: + raise RuntimeError( + "index_first_axis received an invalid gradient with fewer than " + "two dimensions." + ) + other_shape = grad_output.shape[1:] + grad_output = rearrange(grad_output, "b ... -> b (...)") # (m, product(...)) + grad_input = torch.zeros( # (n, product(...)) + [ctx.first_axis_dim, grad_output.shape[1]], + device=grad_output.device, + dtype=grad_output.dtype, + ) + grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output) + return grad_input.reshape(ctx.first_axis_dim, *other_shape), None # (n, ...), None + + +class IndexPutFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor: + # values: (m, ...); indices: (m,) + ctx.save_for_backward(indices) + if indices.ndim != 1: + raise ValueError( + "index_put_first_axis indices must be one-dimensional; " + f"received shape {tuple(indices.shape)}." + ) + if values.ndim < 2: + raise ValueError( + "index_put_first_axis values must have at least two dimensions; " + f"received shape {tuple(values.shape)}." + ) + output = torch.zeros( # (n, ...) + first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype + ) + output[indices] = values + return output # (n, ...) + + @staticmethod + def backward(ctx, grad_output) -> tuple[torch.Tensor, None, None]: + # grad_output: (n, ...) + (indices,) = ctx.saved_tensors + return grad_output[indices], None, None # (m, ...), None, None + + +index_first_axis = IndexFirstAxis.apply +index_put_first_axis = IndexPutFirstAxis.apply + + +def pad_input( + hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int +) -> torch.Tensor: + # hidden_states: (t, ...); indices: (t,) + output = index_put_first_axis(hidden_states, indices, batch * seqlen) # (b * l, ...) + return rearrange(output, "(b s) ... -> b s ...", b=batch) # (b, l, ...) + + +def _unpad_input( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + attention_mask_2d: torch.Tensor, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, torch.Tensor], + tuple[int, int], +]: + # query_layer, key_layer, value_layer: (b, l, h, d); attention_mask_2d: (b, l) + batch_size, seq_len, num_heads, head_dim = query_layer.shape + seqlens = attention_mask_2d.sum(dim=1).int() # (b,) + cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0)) # (b + 1,) + max_seqlen = int(seqlens.max().item()) + indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten() # (t,) + query_layer = index_first_axis( # (t, h, d) + query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices + ) + key_layer = index_first_axis( # (t, h, d) + key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices + ) + value_layer = index_first_axis( # (t, h, d) + value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices + ) + return ( + query_layer, + key_layer, + value_layer, + indices, + (cu_seqlens, cu_seqlens), + (max_seqlen, max_seqlen), + ) + + +def _validate_flash_padding_mask( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask_2d: torch.Tensor, +) -> torch.Tensor: + """Validate the self-attention padding mask used by the varlen kernels.""" + + # query_states, key_states, value_states: (b, l, h, d); attention_mask_2d: (b, l) + if attention_mask_2d.ndim != 2: + raise ValueError("FlashAttention padding masks must have shape (batch, sequence_length).") + expected_shape = query_states.shape[:2] + if tuple(attention_mask_2d.shape) != tuple(expected_shape): + raise ValueError( + "FlashAttention padding mask shape must match the query batch and " + f"sequence dimensions; expected {tuple(expected_shape)}, received " + f"{tuple(attention_mask_2d.shape)}." + ) + if key_states.shape[:2] != expected_shape or value_states.shape[:2] != expected_shape: + raise ValueError( + "Masked FlashAttention requires Q, K, and V to share batch and sequence dimensions." + ) + if attention_mask_2d.device != query_states.device: + raise ValueError("FlashAttention padding mask and Q, K, and V must be on the same device.") + return attention_mask_2d.to(dtype=torch.bool) # (b, l) + + +def kernels_flash_attention_func( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + causal: bool = False, + softmax_scale: float | None = None, + implementation: str = "flash_attention_3", +) -> torch.Tensor: + """Public flash-attention entry point with optional padding handling. + + `softmax_scale`: + None -> kernel applies its default `1 / sqrt(head_dim)`. + float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled + by the caller). + + Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)` + before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass + `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale + again, yielding an effective `1/head_dim` scale that drifts across layers. + """ + # query_states, key_states, value_states: (b, l, h, d) + # attention_mask_2d: (b, l) or None + _validate_kernels_flash_device( + query_states, + key_states, + value_states, + implementation, + ) + runtime_dtype = _validate_kernels_flash_dtype( + query_states, + key_states, + value_states, + implementation, + ) + if query_states.dtype != runtime_dtype: + query_states = query_states.to(dtype=runtime_dtype) # (b, l, h, d) + key_states = key_states.to(dtype=runtime_dtype) # (b, l, h, d) + value_states = value_states.to(dtype=runtime_dtype) # (b, l, h, d) + if attention_mask_2d is not None: + attention_mask_2d = _validate_flash_padding_mask( # (b, l) + query_states, + key_states, + value_states, + attention_mask_2d, + ) + _ensure_flash_kernels_loaded(implementation) + if attention_mask_2d is not None: + batch_size, q_len = query_states.shape[:2] + ( + query_states, + key_states, + value_states, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_q, max_seqlen_k), + ) = _unpad_input( # (t, h, d), (t, h, d), (t, h, d), (t,), (b + 1,), scalars + query_states, + key_states, + value_states, + attention_mask_2d, + ) + attn_output_unpad = _kernels_flash_varlen_forward( # (t, h, d) + query_states=query_states, + key_states=key_states, + value_states=value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_in_batch_q=max_seqlen_q, + max_seqlen_in_batch_k=max_seqlen_k, + causal=causal, + softmax_scale=softmax_scale, + implementation=implementation, + ) + output = pad_input(attn_output_unpad, indices_q, batch_size, q_len) # (b, l, h, d) + return output.masked_fill(~attention_mask_2d[:, :, None, None], 0) # (b, l, h, d) + else: + return _kernels_flash_forward( # (b, l, h, d) + query_states=query_states, + key_states=key_states, + value_states=value_states, + causal=causal, + softmax_scale=softmax_scale, + implementation=implementation, + ) + + +# User-facing backend strings follow the Transformers attention interface. +# Keep ``str`` plus ``Enum`` so stringification stays compatible with existing +# configuration serialization rather than adopting ``StrEnum.__str__``. +class AttentionBackend(str, Enum): # noqa: UP042 + EAGER = "eager" + SDPA = "sdpa" + FLEX_ATTENTION = "flex_attention" + FLASH_ATTENTION_2 = "flash_attention_2" + FLASH_ATTENTION_3 = "flash_attention_3" + + # Internal spelling retained to keep attention modules concise. It is an + # enum alias, not an accepted public backend string. + FLEX = FLEX_ATTENTION + + @property + def is_flash(self) -> bool: + return self in { + AttentionBackend.FLASH_ATTENTION_2, + AttentionBackend.FLASH_ATTENTION_3, + } + + +VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend) + + +def warn_attention_backend_fallback( + requested_backend: str | AttentionBackend, + *, + effective_backend: str | AttentionBackend, + reason: str, +) -> None: + """Warn when one forward call cannot honor the configured backend.""" + + requested = resolve_attention_backend(requested_backend).value + effective = resolve_attention_backend(effective_backend).value + if requested == effective: + return + warnings.warn( + f"{reason} The requested {requested!r} attention implementation cannot " + f"satisfy this call, so FastPLMs is using {effective!r} attention for this " + "call only. This can change performance and memory use; the configured " + "backend remains unchanged for subsequent calls.", + RuntimeWarning, + stacklevel=3, + ) + + +def resolve_attention_backend_for_call( + requested_backend: str | AttentionBackend, + *, + output_attentions: bool, +) -> AttentionBackend: + """Resolve the effective backend for one call and report substitutions once.""" + + requested = resolve_attention_backend(requested_backend) + if not output_attentions or requested == AttentionBackend.EAGER: + return requested + warn_attention_backend_fallback( + requested, + effective_backend=AttentionBackend.EAGER, + reason=( + "output_attentions=True requires the full materialized attention probability " + "matrix, which optimized PyTorch attention APIs do not return." + ), + ) + return AttentionBackend.EAGER + + +def resolve_attention_backend( + requested_backend: str | AttentionBackend | None, +) -> AttentionBackend: + """Validate a backend without silently substituting another implementation.""" + if requested_backend is None: + requested_backend = AttentionBackend.SDPA.value + if isinstance(requested_backend, AttentionBackend): + resolved = requested_backend + else: + try: + resolved = AttentionBackend(requested_backend) + except ValueError as error: + raise ValueError( + f"Unsupported attention implementation {requested_backend!r}; " + f"expected one of {VALID_ATTENTION_BACKENDS}." + ) from error + if resolved == AttentionBackend.FLEX_ATTENTION and flex_attention is None: + raise RuntimeError( + "'flex_attention' was requested, but this PyTorch build does not provide it." + ) + return resolved + + +def get_attn_implementation(config) -> str: + """Read the Transformers attention setting, defaulting to SDPA.""" + requested = getattr(config, "_attn_implementation", None) + if requested is None: + requested = getattr(config, "attn_backend", None) + return resolve_attention_backend(requested).value + + +def set_config_attn_implementation(config, implementation: str) -> str: + """Set both the Transformers field and the internal dispatch field.""" + resolved = resolve_attention_backend(implementation).value + if hasattr(config, "_attn_implementation_internal"): + config._attn_implementation_internal = resolved + else: + config._attn_implementation = resolved + # Existing checkpoint configs contain this field. Keeping it synchronized + # preserves their state schema while the public API uses attn_implementation. + config.attn_backend = resolved + return resolved + + +@torch.compiler.disable +def get_attention_mask( + effective_backend: AttentionBackend, + batch_size: int, + seq_len: int, + device: torch.device, + attention_mask: torch.Tensor | None = None, + dtype: torch.dtype | None = None, + mask_semantics: str = "padding", +) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]: + """Build padding masks once for all encoder layers. + + Returns (attention_mask_2d, attention_mask_4d, flex_block_mask). + """ + # attention_mask: (b, l) or None + if attention_mask is None: + return None, None, None + + if attention_mask.ndim != 2: + raise ValueError( + "attention_mask must have shape (batch, sequence_length); " + f"received rank {attention_mask.ndim} with shape {tuple(attention_mask.shape)}." + ) + expected_shape = (batch_size, seq_len) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + "attention_mask shape must match the input batch and sequence dimensions; " + f"expected {expected_shape}, received {tuple(attention_mask.shape)}." + ) + attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool) # (b, l) + if not bool(attention_mask_2d.any(dim=1).all()): + raise ValueError("attention_mask must keep at least one valid key per batch row.") + + effective_backend = resolve_attention_backend(effective_backend) + + if effective_backend.is_flash: + return attention_mask_2d, None, None # (b, l), None, None + + if effective_backend == AttentionBackend.FLEX_ATTENTION: + if create_block_mask is None: + raise RuntimeError( + "'flex_attention' was requested, but torch.create_block_mask is unavailable." + ) + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + del head_idx, q_idx + # Match eager and SDPA: padding masks suppress invalid keys only. + # Invalid queries still attend to real keys and therefore remain + # finite; downstream residue masks exclude their outputs. + return attention_mask_2d[batch_idx, kv_idx] + + flex_block_mask = _get_flex_block_mask( + mask_pattern=attention_mask_2d, + batch_size=batch_size, + query_length=seq_len, + key_value_length=seq_len, + device=device, + dtype=dtype, + mask_semantics=mask_semantics, + mask_mod=mask_mod, + ) + return attention_mask_2d, None, flex_block_mask # (b, l), None, BlockMask + + # SDPA/manual masks only keys. Padding queries still attend to real keys, so + # their outputs stay finite instead of softmaxing over all -inf scores. + attention_mask_4d = attention_mask_2d[:, None, None, :] # (b, 1, 1, l) + return attention_mask_2d, attention_mask_4d, None # (b, l), (b, 1, 1, l), None + + +def bool_to_additive_mask( + bool_mask: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid). + + Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))` + directly on a bool tensor returns a bool tensor because `-inf` casts to `True`. + That silently drops the mask. Always allocate a float tensor first, then fill it. + This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask. + """ + # bool_mask: (...) + if bool_mask.dtype != torch.bool: + raise TypeError( + f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}" + ) + additive = torch.zeros_like(bool_mask, dtype=dtype) # (...) + additive.masked_fill_(bool_mask.logical_not(), float("-inf")) + return additive # (...) diff --git a/src/fastplms/attention/_kernel_lock.py b/src/fastplms/attention/_kernel_lock.py new file mode 100644 index 0000000..a859f57 --- /dev/null +++ b/src/fastplms/attention/_kernel_lock.py @@ -0,0 +1,175 @@ +"""Resolve and validate Hugging Face kernels before importing their binaries.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + + +def require_kernels_package() -> None: + """Fail early when the precompiled-kernel runtime is not installed.""" + try: + import kernels # noqa: F401 + except ImportError as error: + raise RuntimeError( + "Precompiled FlashAttention requires requirements/features/flash.in." + ) from error + + +def _kernel_lock_path() -> Path: + """Return the kernel lock from a Hub artifact or source checkout.""" + source_path = Path(__file__).resolve() + for candidate in ( + source_path.parents[1] / "kernels.lock", + source_path.parents[3] / "kernels.lock", + ): + if candidate.is_file(): + return candidate + + raise RuntimeError( + "kernels.lock is missing from the Hugging Face artifact or source checkout." + ) + + +def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]: + try: + lock_entries = json.loads(lock_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"Unable to read the kernel lock: {lock_path}") from error + if not isinstance(lock_entries, list): + raise RuntimeError("kernels.lock must contain a JSON list.") + if any(not isinstance(entry, dict) for entry in lock_entries): + raise RuntimeError("Every kernels.lock entry must be a JSON object.") + matches = [entry for entry in lock_entries if entry.get("repo_id") == repository] + if len(matches) != 1: + raise RuntimeError( + f"kernels.lock must contain exactly one entry for {repository!r}; found {len(matches)}." + ) + return matches[0] + + +def _offline_mode() -> bool: + """Return whether Hub access was explicitly disabled for this process.""" + + enabled_values = {"1", "on", "true", "yes"} + return any( + os.environ.get(name, "").strip().lower() in enabled_values + for name in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ) + + +def _offline_snapshot_path(repository: str, revision: str) -> Path: + """Locate one exact, possibly sparse, kernel snapshot without using Hub APIs.""" + + try: + from huggingface_hub import constants + from huggingface_hub.file_download import repo_folder_name + except ImportError as error: + raise RuntimeError("Offline kernel loading requires huggingface-hub.") from error + + cache_root = Path(os.environ.get("KERNELS_CACHE") or constants.HF_HUB_CACHE).resolve() + repository_root = ( + cache_root / repo_folder_name(repo_id=repository, repo_type="kernel") + ).resolve() + snapshot = repository_root / "snapshots" / revision + if not snapshot.is_dir(): + raise RuntimeError( + f"The exact offline kernel snapshot {repository}@{revision} is not cached under " + f"{cache_root}. Run `kernels download` before enabling offline mode." + ) + if repository_root not in snapshot.resolve().parents: + raise RuntimeError(f"Refusing kernel snapshot outside its cache repository: {snapshot}") + return snapshot + + +def _load_offline_locked_kernel( + repository: str, + revision: str, + variant_locks: dict[str, object], +) -> object: + """Validate and import the one compatible variant from a sparse Hub snapshot.""" + snapshot = _offline_snapshot_path(repository, revision) + build_root = snapshot / "build" + if not build_root.is_dir(): + raise RuntimeError(f"The cached kernel snapshot has no build directory: {snapshot}") + + cached_names = sorted(entry.name for entry in build_root.iterdir() if entry.is_dir()) + unexpected = sorted(set(cached_names).difference(variant_locks)) + if unexpected: + raise RuntimeError( + f"The cached {repository}@{revision} snapshot contains unlocked variants: " + f"{', '.join(unexpected)}" + ) + + try: + from kernels import get_local_kernel + from kernels.utils import validate_kernel + from kernels.variants import get_variants_local, resolve_variants + except ImportError as error: + raise RuntimeError( + "Precompiled FlashAttention requires requirements/features/flash.in." + ) from error + + cached_variants = get_variants_local(build_root) + parsed_names = {variant.variant_str for variant in cached_variants} + invalid = sorted(set(cached_names).difference(parsed_names)) + if invalid: + raise RuntimeError( + f"The cached {repository}@{revision} snapshot contains invalid variants: " + f"{', '.join(invalid)}" + ) + + compatible, _ = resolve_variants(cached_variants) + if len(compatible) != 1: + names = ", ".join(variant.variant_str for variant in compatible) or "none" + raise RuntimeError( + f"Expected exactly one compatible cached variant for {repository}@{revision}; " + f"found {names}." + ) + variant_name = compatible[0].variant_str + variant_lock = variant_locks.get(variant_name) + expected_hash = getattr(variant_lock, "hash", None) + if not isinstance(expected_hash, str) or not expected_hash.startswith("sha256-"): + raise RuntimeError(f"The kernel lock for {variant_name} has no valid SHA-256 digest.") + + # Hash validation deliberately happens before import. This operates on the + # sparse snapshot produced by `kernels download` and avoids Hub 1.23's + # full-snapshot completeness check in offline mode. + validate_kernel(repo_path=snapshot, variant=variant_name, hash=expected_hash) + return get_local_kernel(build_root / variant_name) + + +def load_locked_kernel(repository: str, revision: str) -> object: + """Download, hash-validate, then import one immutable precompiled kernel.""" + require_kernels_package() + try: + from kernels import get_local_kernel, install_kernel + from kernels.lockfile import KernelLock + except ImportError as error: + raise RuntimeError( + "Precompiled FlashAttention requires requirements/features/flash.in." + ) from error + + lock_path = _kernel_lock_path() + kernel_lock = KernelLock.from_json(_locked_entry(lock_path, repository)) + if kernel_lock.sha != revision: + raise RuntimeError( + f"The typed manifest pins {repository}@{revision}, but kernels.lock pins " + f"{kernel_lock.sha}." + ) + + if _offline_mode(): + return _load_offline_locked_kernel(repository, revision, kernel_lock.variants) + + # `install_kernel` downloads data without importing it and validates the + # selected build against the tracked variant hash. Only then is the exact + # validated path imported directly. Offline mode uses the sparse-cache + # resolver above because Hub 1.23 rejects partial snapshots as incomplete. + validated_path = install_kernel( + repository, + revision=kernel_lock.sha, + variant_locks=kernel_lock.variants, + ) + return get_local_kernel(validated_path) diff --git a/src/fastplms/attention/interfaces.py b/src/fastplms/attention/interfaces.py new file mode 100644 index 0000000..19aedde --- /dev/null +++ b/src/fastplms/attention/interfaces.py @@ -0,0 +1,242 @@ +"""Transformers-compatible attention selection for FastPLMs models.""" + +from __future__ import annotations + +import torch +from collections.abc import Mapping +from functools import partial +from typing import Any +from transformers import AttentionInterface, AttentionMaskInterface + +from ._core import ( + AttentionBackend, + get_attn_implementation, + kernels_flash_attention_func, + resolve_attention_backend, + set_config_attn_implementation, +) +from ._kernel_lock import require_kernels_package + + +def _kernels_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + *, + implementation: str, + **kwargs: Any, +) -> tuple[torch.Tensor, None]: + """Run one canonical FlashAttention backend through Hugging Face kernels. + + Transformers attention functions receive Q, K, and V with shape + (b, h, l, d) and return an output with shape (b, l, h, d). The shared + FastPLMs kernel adapter uses the latter layout internally. + """ + + # query, key, value: (b, h, l, d); attention_mask: (b, l) or None + dropout = float(kwargs.get("dropout", 0.0) or 0.0) + if module.training and dropout: + raise RuntimeError( + "Hugging Face kernels FlashAttention is inference-only when attention dropout " + "is nonzero. Use SDPA for this training configuration." + ) + causal = bool(kwargs.get("is_causal", getattr(module, "is_causal", False))) + softmax_scale = kwargs.get("scaling") + output = kernels_flash_attention_func( + query_states=query.transpose(1, 2).contiguous(), # (b, l, h, d) + key_states=key.transpose(1, 2).contiguous(), # (b, l, h, d) + value_states=value.transpose(1, 2).contiguous(), # (b, l, h, d) + attention_mask_2d=attention_mask, + causal=causal, + softmax_scale=softmax_scale, + implementation=implementation, + ) # (b, l, h, d) + return output, None # (b, l, h, d), None + + +# Keep FastPLMs' kernels-only adapters local to this registry instance. +# ``GeneralInterface.register`` updates Transformers' class-wide mapping, so +# using it here would replace the canonical FlashAttention handlers for every +# model in the process, including models unrelated to FastPLMs. +FASTPLMS_ATTENTION_FUNCTIONS = AttentionInterface() +FASTPLMS_ATTENTION_MASKS = AttentionMaskInterface() +FASTPLMS_ATTENTION_FUNCTIONS["flash_attention_2"] = partial( + _kernels_attention_forward, + implementation="flash_attention_2", +) +FASTPLMS_ATTENTION_FUNCTIONS["flash_attention_3"] = partial( + _kernels_attention_forward, + implementation="flash_attention_3", +) +for _flash_name in ("flash_attention_2", "flash_attention_3"): + FASTPLMS_ATTENTION_MASKS[_flash_name] = FASTPLMS_ATTENTION_MASKS[_flash_name] + + +class FastPLMsAttentionMixin: + """Synchronize Transformers attention selection with custom model layers. + + Model families retain their checkpoint parameter names. Only runtime + attributes are updated when ``set_attn_implementation`` is called. + """ + + _supports_sdpa = True + _supports_flex_attn = True + # Transformers 5.13 uses the singular flag during model construction. A + # family opts in only when its manifest entry advertises at least one of + # the two FastPLMs kernels-only FlashAttention implementations. + _supports_flash_attn = False + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + ) + + def _validate_attention_name(self, implementation: str) -> None: + if implementation not in self._fastplms_attention_implementations: + raise ValueError( + f"{type(self).__name__} does not support {implementation!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + + def _check_and_adjust_attn_implementation( + self, + attn_implementation: str | None, + is_init_check: bool = False, + allow_all_kernels: bool = False, + ) -> str: + """Resolve attention without invoking Transformers' source-Flash probe. + + The standard ``flash_attention_2`` and ``flash_attention_3`` names are + retained for the Transformers API, but FastPLMs resolves them only + through the exact Hugging Face ``kernels`` artifacts pinned by + ``models.toml``. Repository-qualified or otherwise external kernels + are never accepted through this model hook. + """ + + if allow_all_kernels: + raise ValueError("FastPLMs does not load external attention kernels.") + if attn_implementation is None: + return super()._check_and_adjust_attn_implementation( + None, + is_init_check=is_init_check, + allow_all_kernels=False, + ) + + self._validate_attention_name(attn_implementation) + if attn_implementation in {"flash_attention_2", "flash_attention_3"}: + if not self._supports_flash_attn: + raise ValueError( + f"{type(self).__name__} does not advertise kernels-only FlashAttention." + ) + # Validate the lightweight Python dependency here, but defer binary + # download and import until Q, K, and V have passed the CUDA gate. + require_kernels_package() + return attn_implementation + + return super()._check_and_adjust_attn_implementation( + attn_implementation, + is_init_check=is_init_check, + allow_all_kernels=False, + ) + + def __init__(self, config, *args: Any, **kwargs: Any) -> None: + sentinel = object() + internal = getattr(config, "_attn_implementation_internal", sentinel) + canonical = ( + getattr(config, "_attn_implementation", None) if internal is sentinel else internal + ) + legacy = getattr(config, "attn_backend", None) + requested = canonical if canonical is not None else legacy + if requested is not None: + if not isinstance(requested, str): + raise TypeError( + "The configured attention implementation must be a string or None; " + f"received {type(requested).__name__}." + ) + self._validate_attention_name(requested) + # ``PreTrainedModel.__init__`` resolves a missing Transformers + # implementation to the family default. Legacy FastPLMs configs + # persist their explicit choice in ``attn_backend``, so forward it + # into the canonical Transformers field before the base class can + # replace it with SDPA. A non-None canonical value still wins, + # including an explicit ``attn_implementation=...`` load override. + if canonical is None and legacy is not None: + set_config_attn_implementation(config, legacy) + super().__init__(config, *args, **kwargs) + # Transformers resolves an unspecified implementation during the base + # model initialization. Synchronize that choice before family layers + # are constructed. + resolved = get_attn_implementation(config) + self._validate_attention_name(resolved) + set_config_attn_implementation(config, resolved) + + def set_attn_implementation( + self, + attn_implementation: str | Mapping[str, str], + allow_all_kernels: bool = False, + ) -> None: + """Select an advertised backend and update every instantiated layer.""" + if isinstance(attn_implementation, Mapping): + if set(attn_implementation) == {""}: + attn_implementation = attn_implementation[""] + else: + raise ValueError( + "FastPLMs models have one attention backbone; pass a string or {'': name}." + ) + resolved_name = self._check_and_adjust_attn_implementation( + attn_implementation, + is_init_check=False, + allow_all_kernels=allow_all_kernels, + ) + set_config_attn_implementation(self.config, resolved_name) + resolved = resolve_attention_backend(resolved_name) + for module in self.modules(): + if module is self: + continue + for attribute in ("attn_backend", "attention_backend", "_attn_backend"): + if attribute not in module.__dict__: + continue + current = module.__dict__[attribute] + module.__dict__[attribute] = ( + resolved if isinstance(current, AttentionBackend) else resolved_name + ) + + +def validate_transformers_attention_interfaces() -> None: + """Verify that Transformers exposes functions and masks for every backend. + + Transformers 5.13 registers these canonical names. The FastPLMs function + overrides remain instance-local and do not replace process-global handlers. + """ + function_registry = FASTPLMS_ATTENTION_FUNCTIONS + mask_registry = FASTPLMS_ATTENTION_MASKS + missing_functions = [ + name + for name in ( + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + if name not in function_registry + ] + missing_masks = [ + name + for name in ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + if name not in mask_registry + ] + if missing_functions or missing_masks: + raise RuntimeError( + "Transformers attention registry is incomplete: " + f"functions={missing_functions}, masks={missing_masks}." + ) diff --git a/src/fastplms/embeddings/__init__.py b/src/fastplms/embeddings/__init__.py new file mode 100644 index 0000000..dd1ddd5 --- /dev/null +++ b/src/fastplms/embeddings/__init__.py @@ -0,0 +1,66 @@ +"""Ordered, residue-aware protein embedding utilities.""" + +from .pooling import POOLING_NAMES, Pooler, pagerank_weights +from .runner import ( + EmbeddingMixin, + embed_dataset, + iter_fasta, + parse_fasta, + select_hidden_state_embeddings, +) +from .storage import ( + DEFAULT_SHARD_SIZE, + append_sqlite_records, + convert_legacy_sqlite, + garbage_collect_safetensors_generations, + initialize_sqlite_run, + load_legacy_pth, + load_result, + load_safetensors_result, + load_sqlite_result, + save_result, + save_safetensors_result, + save_sqlite_result, + tensor_sha256, + update_sqlite_run_metadata, +) +from .types import ( + EmbeddingBatch, + EmbeddingInput, + EmbeddingRecord, + EmbeddingResult, + LazyTensorReference, + TensorValue, +) + + +__all__ = [ + "DEFAULT_SHARD_SIZE", + "POOLING_NAMES", + "EmbeddingBatch", + "EmbeddingInput", + "EmbeddingMixin", + "EmbeddingRecord", + "EmbeddingResult", + "LazyTensorReference", + "Pooler", + "TensorValue", + "append_sqlite_records", + "convert_legacy_sqlite", + "embed_dataset", + "garbage_collect_safetensors_generations", + "initialize_sqlite_run", + "iter_fasta", + "load_legacy_pth", + "load_result", + "load_safetensors_result", + "load_sqlite_result", + "pagerank_weights", + "parse_fasta", + "save_result", + "save_safetensors_result", + "save_sqlite_result", + "select_hidden_state_embeddings", + "tensor_sha256", + "update_sqlite_run_metadata", +] diff --git a/src/fastplms/embeddings/pooling.py b/src/fastplms/embeddings/pooling.py new file mode 100644 index 0000000..08d577c --- /dev/null +++ b/src/fastplms/embeddings/pooling.py @@ -0,0 +1,218 @@ +"""Residue-aware pooling implemented entirely with PyTorch.""" + +from __future__ import annotations + +import math +import torch +from collections.abc import Sequence +from torch import Tensor + + +POOLING_NAMES = frozenset({"mean", "max", "norm", "median", "std", "var", "cls", "parti"}) + + +def _validate_inputs(X: Tensor, M: Tensor) -> Tensor: + # X: (b, l, d); M: (b, l) + if not isinstance(X, Tensor) or not isinstance(M, Tensor): + raise TypeError("X and M must be tensors.") + if X.ndim != 3: + raise ValueError(f"X must have shape (b, l, d), got {tuple(X.shape)}.") + if not X.is_floating_point(): + raise TypeError("X must use a floating-point embedding dtype.") + if M.shape != X.shape[:2]: + raise ValueError(f"M must have shape (b, l)={tuple(X.shape[:2])}, got {tuple(M.shape)}.") + if M.is_complex(): + raise TypeError("M must be a boolean or binary numeric residue mask.") + if not bool(torch.isfinite(M).all()) or not bool(((M == 0) | (M == 1)).all()): + raise ValueError("M must contain only finite binary mask values.") + M = M.to(device=X.device, dtype=torch.bool) # (b, l) + if not bool(M.any(dim=1).all()): + raise ValueError("Every sample must contain at least one biological residue.") + if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()): + raise ValueError("Biological residue embeddings produced non-finite output.") + return M # (b, l) + + +def _pooled_attention(attentions: Tensor | Sequence[Tensor], *, batch_size: int) -> Tensor: + """Max-pool layer/head attention A to shape ``(b, l, l)``. + + ``parti`` historically keeps the strongest directed edge across the + available attention maps before PageRank. Replacing NetworkX with Torch + must not change that reduction. + """ + + if isinstance(attentions, Sequence): + if not attentions: + raise ValueError("parti received an empty attention sequence.") + # Each A_i: (b, h, l, l). + A = torch.stack(tuple(attentions), dim=1) # (b, n, h, l, l) + else: + A = attentions # (b, ..., l, l) + + if A.ndim == 5: + if A.shape[0] != batch_size and A.shape[1] == batch_size: + A = A.transpose(0, 1) # (b, n, h, l, l) + if A.shape[0] != batch_size: + raise ValueError("Five-dimensional attentions must use (b, n, h, l, l).") + A = A.flatten(1, 2).amax(dim=1) # (b, l, l) + elif A.ndim == 4: + if A.shape[0] != batch_size: + raise ValueError("Four-dimensional attentions must use (b, h, l, l).") + A = A.amax(dim=1) # (b, l, l) + elif A.ndim == 3: + if A.shape[0] != batch_size: + raise ValueError("Three-dimensional attentions must use (b, l, l).") + else: + raise ValueError("Attentions must have shape (b, l, l), (b, h, l, l), or (b, n, h, l, l).") + return A # (b, l, l) + + +def pagerank_weights( + A: Tensor, + *, + damping: float = 0.85, + tolerance: float = 1e-6, + max_iterations: int = 100, +) -> Tensor: + """Compute PageRank weights for a non-negative attention matrix A. + + A has shape ``(l, l)``. Rows are normalized into transition + probabilities; dangling rows transition uniformly. + """ + + # A: (l, l) + if not isinstance(A, Tensor): + raise TypeError("A must be a tensor.") + if A.ndim != 2 or A.shape[0] != A.shape[1]: + raise ValueError(f"A must be square, got shape {tuple(A.shape)}.") + if not A.is_floating_point(): + raise TypeError("A must use a floating-point attention dtype.") + if not isinstance(damping, (int, float)) or isinstance(damping, bool): + raise TypeError("damping must be a finite float in [0, 1).") + if not math.isfinite(float(damping)) or not 0 <= damping < 1: + raise ValueError("damping must be a finite float in [0, 1).") + if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool): + raise TypeError("tolerance must be a positive finite float.") + if not math.isfinite(float(tolerance)) or tolerance <= 0: + raise ValueError("tolerance must be a positive finite float.") + if not isinstance(max_iterations, int) or isinstance(max_iterations, bool): + raise TypeError("max_iterations must be a positive integer.") + if max_iterations <= 0: + raise ValueError("max_iterations must be a positive integer.") + length = A.shape[0] + if length == 0: + raise ValueError("PageRank requires at least one residue.") + if not bool(torch.isfinite(A).all()): + raise ValueError("A must contain only finite attention values.") + work_dtype = torch.float64 if A.dtype == torch.float64 else torch.float32 + P = A.detach().to(dtype=work_dtype).clamp_min(0) # (l, l) + row_sum = P.sum(dim=-1, keepdim=True) # (l, 1) + uniform = torch.full_like(P, 1.0 / length) # (l, l) + P = torch.where( # (l, l) + row_sum > 0, + P / row_sum.clamp_min(torch.finfo(work_dtype).tiny), + uniform, + ) + p = torch.full((length,), 1.0 / length, device=P.device, dtype=work_dtype) # (l,) + teleport = (1.0 - damping) / length + for _ in range(max_iterations): + p_next = teleport + damping * (P.transpose(0, 1) @ p) # (l,) + if torch.linalg.vector_norm(p_next - p, ord=1) <= tolerance: + p = p_next # (l,) + break + p = p_next # (l,) + return p / p.sum() # (l,) + + +class Pooler: + """Apply one or more pooling operations to biological residue rows.""" + + def __init__(self, pooling: str | Sequence[str] = ("mean",)) -> None: + pooling_value: object = pooling + if isinstance(pooling_value, (bytes, bytearray)) or not isinstance( + pooling_value, (str, Sequence) + ): + raise TypeError("pooling must be a name or a sequence of names.") + names = (pooling_value,) if isinstance(pooling_value, str) else tuple(pooling_value) + if not all(isinstance(name, str) for name in names): + raise TypeError("pooling names must be strings.") + if not names: + raise ValueError("At least one pooling operation is required.") + unknown = set(names) - POOLING_NAMES + if unknown: + raise ValueError(f"Unknown pooling operations: {sorted(unknown)}.") + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + raise ValueError(f"Duplicate pooling operations are not supported: {duplicates}.") + self.names = names + + def output_slices(self, d: int) -> dict[str, tuple[int, int]]: + """Return the output interval assigned to each pooler.""" + + if not isinstance(d, int) or isinstance(d, bool): + raise TypeError("d must be a positive integer.") + if d <= 0: + raise ValueError("d must be a positive integer.") + return {name: (i * d, (i + 1) * d) for i, name in enumerate(self.names)} + + def __call__( + self, + X: Tensor, + residue_mask: Tensor, + *, + attentions: Tensor | Sequence[Tensor] | None = None, + attention_backend: str | None = None, + ) -> Tensor: + # X: (b, l, d); residue_mask: (b, l) + M = _validate_inputs(X, residue_mask) # (b, l) + M_expanded = M.unsqueeze(-1) # (b, l, 1) + count = M_expanded.sum(dim=1).clamp_min(1) # (b, 1) + X_residues = X.masked_fill(~M_expanded, 0) # (b, l, d) + outputs: list[Tensor] = [] + + for name in self.names: + if name == "mean": + Y = X_residues.sum(dim=1) / count # (b, d) + elif name == "max": + Y = X.masked_fill(~M_expanded, -torch.inf).max(dim=1).values # (b, d) + elif name == "norm": + Y = torch.linalg.vector_norm(X_residues, ord=2, dim=1) # (b, d) + elif name == "median": + Y = X.masked_fill(~M_expanded, torch.nan).nanmedian(dim=1).values # (b, d) + elif name in {"var", "std"}: + mean = X_residues.sum(dim=1, keepdim=True) / count.unsqueeze(1) # (b, 1, d) + centered = (X - mean).masked_fill(~M_expanded, 0) # (b, l, d) + variance = (centered**2).sum(dim=1) / count # (b, d) + Y = variance.sqrt() if name == "std" else variance # (b, d) + elif name == "cls": + Y = X[:, 0] # (b, d) + else: + if attention_backend != "eager": + raise ValueError( + "parti requires attn_implementation='eager' so full " + "attention matrices are available." + ) + if attentions is None: + raise ValueError("parti requires model attention matrices.") + if int(M.sum(dim=1).max().item()) > 2048: + raise ValueError("parti supports at most 2,048 biological residues.") + A = _pooled_attention(attentions, batch_size=X.shape[0]).to(X.device) # (b, l, l) + pooled: list[Tensor] = [] + for X_i, M_i, A_i in zip(X, M, A, strict=True): + # X_i: (l, d); M_i: (l,); A_i: (l, l) + indices = M_i.nonzero(as_tuple=True)[0] # (r,) + A_residue = A_i.index_select(0, indices).index_select(1, indices) # (r, r) + w = pagerank_weights(A_residue).to(dtype=X.dtype) # (r,) + pooled.append(w @ X_i.index_select(0, indices)) # (d,) + Y = torch.stack(pooled) # (b, d) + if not bool(torch.isfinite(Y).all()): + raise ValueError( + f"Pooling operation {name!r} produced non-finite output from " + "biological residue embeddings." + ) + outputs.append(Y) + + return torch.cat(outputs, dim=-1) # (b, len(self.names) * d) + + +__all__ = ["POOLING_NAMES", "Pooler", "pagerank_weights"] diff --git a/src/fastplms/embeddings/runner.py b/src/fastplms/embeddings/runner.py new file mode 100644 index 0000000..200082f --- /dev/null +++ b/src/fastplms/embeddings/runner.py @@ -0,0 +1,1583 @@ +"""Model-independent dataset embedding orchestration.""" + +from __future__ import annotations + +import hashlib +import json +import platform +import sqlite3 +import tempfile +import torch +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Any, overload +from torch import Tensor + +from .pooling import Pooler +from .storage import ( + SafetensorsStreamWriter, + append_sqlite_records, + initialize_sqlite_run, + load_result, + load_sqlite_result, + safetensors_result_exists, + save_result, + tensor_sha256, + update_sqlite_run_metadata, +) +from .types import ( + EmbeddingBatch, + EmbeddingInput, + EmbeddingRecord, + EmbeddingResult, + LazyTensorReference, +) + + +_MAX_PARTI_RESIDUES = 2_048 +_RUN_FINGERPRINT_SCHEMA_VERSION = 3 +_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2 +_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16 +_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"}) + + +def _validate_parti_length(M: Tensor) -> None: + """Reject an oversized attention graph before model inference.""" + + # M: (b, l) + n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item()) + if n_residues > _MAX_PARTI_RESIDUES: + raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.") + + +def select_hidden_state_embeddings( + last_hidden_state: Tensor, + hidden_states: tuple[Tensor, ...] | None, + *, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, +) -> Tensor: + """Select one hidden state or stack every state without changing values.""" + # last_hidden_state and each hidden_states entry: (b, l, d) + if store_all_hidden_states: + if not hidden_states: + raise ValueError("store_all_hidden_states requires model hidden states.") + # H has shape (b, n, l, d), where n follows the model's output order. + return torch.stack(hidden_states, dim=1) # (b, n, l, d) + if hidden_state_index == -1: + return last_hidden_state # (b, l, d) + if not hidden_states: + raise ValueError("hidden_state_index requires model hidden states.") + return hidden_states[hidden_state_index] # (b, l, d) + + +def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]: + """Yield FASTA records in source order without reading the file into memory.""" + + identifier: str | None = None + sequence_parts: list[str] = [] + found_record = False + with Path(path).open("r", encoding="utf-8") as handle: + for line_number, raw_line in enumerate(handle, start=1): + line = raw_line.strip() + if not line: + continue + if line.startswith(">"): + if identifier is not None: + found_record = True + yield EmbeddingInput(identifier, "".join(sequence_parts)) + identifier = line[1:].strip().split(maxsplit=1)[0] + if not identifier: + raise ValueError(f"Missing FASTA identifier on line {line_number}.") + sequence_parts = [] + else: + if identifier is None: + raise ValueError( + f"Sequence data precedes the first FASTA header on line {line_number}." + ) + sequence_parts.append("".join(line.split())) + if identifier is not None: + found_record = True + yield EmbeddingInput(identifier, "".join(sequence_parts)) + if not found_record: + raise ValueError(f"No FASTA records found in {path}.") + + +def parse_fasta(path: str | Path) -> list[EmbeddingInput]: + """Parse FASTA records while preserving identifiers, order, and duplicates.""" + + return list(iter_fasta(path)) + + +def _normalize_input_item( + position: int, + item: str | EmbeddingInput | tuple[str, str], +) -> EmbeddingInput: + if isinstance(item, EmbeddingInput): + return item + if isinstance(item, str): + return EmbeddingInput(str(position), item) + if isinstance(item, tuple) and len(item) == 2: + return EmbeddingInput(str(item[0]), str(item[1])) + raise TypeError( + "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples." + ) + + +class _InputSpool(Sequence[EmbeddingInput]): + """Immutable disk-backed normalized inputs with an incremental digest.""" + + def __init__( + self, + values: Iterable[str | EmbeddingInput | tuple[str, str]], + ) -> None: + self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory( + prefix="fastplms-inputs-" + ) + self.path = Path(self._temporary.name) / "inputs.sqlite" + self._connection: sqlite3.Connection | None = sqlite3.connect(self.path) + self._connection.execute( + "CREATE TABLE inputs (" + "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)" + ) + digest = hashlib.sha256() + count = 0 + pending: list[tuple[int, str, str]] = [] + try: + for position, item in enumerate(values): + record = _normalize_input_item(position, item) + for value in (record.id, record.sequence): + encoded = value.encode("utf-8") + digest.update(len(encoded).to_bytes(8, "big")) + digest.update(encoded) + pending.append((position, record.id, record.sequence)) + count += 1 + if len(pending) == 1_024: + self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending) + pending.clear() + if pending: + self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending) + if count == 0: + raise ValueError("inputs must contain at least one sequence.") + self._connection.commit() + self._connection.close() + self._connection = sqlite3.connect( + f"{self.path.resolve().as_uri()}?mode=ro", + uri=True, + ) + except BaseException: + self.close() + raise + digest.update(count.to_bytes(8, "big")) + self.input_fingerprint = digest.hexdigest() + self._count = count + + def _require_connection(self) -> sqlite3.Connection: + if self._connection is None: + raise RuntimeError("Input spool is closed.") + return self._connection + + def __len__(self) -> int: + return self._count + + def __iter__(self) -> Iterator[EmbeddingInput]: + cursor = self._require_connection().execute( + "SELECT input_id, sequence FROM inputs ORDER BY position" + ) + while rows := cursor.fetchmany(1_024): + for input_id, sequence in rows: + yield EmbeddingInput(input_id, sequence) + + @overload + def __getitem__(self, index: int, /) -> EmbeddingInput: ... + + @overload + def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ... + + def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]: + connection = self._require_connection() + + if isinstance(index, slice): + start, stop, step = index.indices(self._count) + if step != 1: + return [self[position] for position in range(start, stop, step)] + rows = connection.execute( + "SELECT input_id, sequence FROM inputs " + "WHERE position >= ? AND position < ? ORDER BY position", + (start, stop), + ).fetchall() + return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows] + position = index + self._count if index < 0 else index + if position < 0 or position >= self._count: + raise IndexError(index) + row = connection.execute( + "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,) + ).fetchone() + if row is None: + raise IndexError(index) + return EmbeddingInput(row[0], row[1]) + + def close(self) -> None: + connection = getattr(self, "_connection", None) + if connection is not None: + connection.close() + self._connection = None + temporary = getattr(self, "_temporary", None) + if temporary is not None: + temporary.cleanup() + self._temporary = None + + def __del__(self) -> None: + self.close() + + +def _normalize_inputs( + inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path), + *, + disk_backed: bool, +) -> Sequence[EmbeddingInput]: + is_fasta_path = isinstance(inputs, Path) + if isinstance(inputs, str): + try: + is_fasta_path = Path(inputs).is_file() + except OSError: + is_fasta_path = False + should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping)) + values: Iterable[str | EmbeddingInput | tuple[str, str]] + if isinstance(inputs, Path): + values = iter_fasta(inputs) + elif isinstance(inputs, str): + values = iter_fasta(inputs) if is_fasta_path else [inputs] + elif isinstance(inputs, Mapping): + values = inputs.items() + else: + values = inputs + if should_spool: + return _InputSpool(values) + records: list[EmbeddingInput] = [] + for position, item in enumerate(values): + records.append(_normalize_input_item(position, item)) + if not records: + raise ValueError("inputs must contain at least one sequence.") + return records + + +def _validate_untruncated_lengths( + records: Sequence[EmbeddingInput], + *, + max_length: int | None, + truncate: bool, +) -> None: + """Fail before inference when a biological-residue limit would be exceeded.""" + + if max_length is None or truncate: + return + for position, record in enumerate(records): + residue_count = len(record.sequence) + if residue_count > max_length: + raise ValueError( + f"Input at position {position} with id {record.id!r} has " + f"{residue_count} biological residues, exceeding max_length={max_length} " + "while truncate=False." + ) + + +def _model_device(model: Any) -> torch.device: + try: + return torch.device(next(model.parameters()).device) + except (AttributeError, StopIteration): + return torch.device("cpu") + + +def _attention_backend(model: Any) -> str | None: + config = getattr(model, "config", None) + for name in ("_attn_implementation", "attn_implementation", "attn_backend"): + value = getattr(config, name, None) + if value: + return str(value) + return None + + +def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None: + if backend not in {"flash_attention_2", "flash_attention_3"}: + return None + from fastplms.registry import get_model_registry + + spec = get_model_registry().attention_kernels[backend] + return { + "repository": spec.repository, + "revision": spec.revision, + "version": spec.version, + "expected_variant": spec.expected_variant, + "dtypes": list(spec.dtypes), + } + + +def _fingerprint_jsonable(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _fingerprint_jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_fingerprint_jsonable(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted((_fingerprint_jsonable(item) for item in value), key=repr) + if isinstance(value, Path): + return str(value) + if isinstance(value, Tensor): + return { + "dtype": str(value.dtype).removeprefix("torch."), + "shape": list(value.shape), + "sha256": tensor_sha256(value), + } + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, torch.device): + return str(value) + if value is None or isinstance(value, (str, int, float, bool)): + return value + return { + "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}", + "value": str(value), + } + + +def _tokenizer_content_sha256(tokenizer: Any) -> str: + content: dict[str, Any] = { + "init_kwargs": getattr(tokenizer, "init_kwargs", None), + "special_tokens_map": getattr(tokenizer, "special_tokens_map", None), + "model_max_length": getattr(tokenizer, "model_max_length", None), + "padding_side": getattr(tokenizer, "padding_side", None), + "truncation_side": getattr(tokenizer, "truncation_side", None), + } + get_vocab = getattr(tokenizer, "get_vocab", None) + if callable(get_vocab): + content["vocabulary"] = get_vocab() + get_added_vocab = getattr(tokenizer, "get_added_vocab", None) + if callable(get_added_vocab): + content["added_vocabulary"] = get_added_vocab() + backend = getattr(tokenizer, "backend_tokenizer", None) + backend_to_str = getattr(backend, "to_str", None) + if callable(backend_to_str): + content["backend"] = backend_to_str() + serialized = json.dumps( + _fingerprint_jsonable(content), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode() + return hashlib.sha256(serialized).hexdigest() + + +def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]: + resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None) + if resolved is None: + # Raw-sequence families such as E1 retain their loader context on the + # model/encoder rather than exposing a Transformers tokenizer. Bind the + # non-secret source policy to resume identity without serializing a Hub + # token or forcing lazy tokenizer initialization. + for candidate in (model, getattr(model, "model", None)): + settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs") + if isinstance(settings, Mapping): + token_value = settings.get("token") + return { + "mode": "native-sequence", + "source": ( + str(settings.get("tokenizer_source")) + if settings.get("tokenizer_source") is not None + else None + ), + "revision": settings.get("revision"), + "cache_dir": ( + str(settings.get("cache_dir")) + if settings.get("cache_dir") is not None + else None + ), + "local_files_only": bool(settings.get("local_files_only", False)), + "token_policy": ( + "disabled" + if token_value is False + else "provided" + if token_value is not None + else "default" + ), + } + return {"mode": "native-sequence"} + return { + "mode": "tokenizer", + "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}", + "name_or_path": getattr(resolved, "name_or_path", None), + "vocab_size": getattr(resolved, "vocab_size", None), + "special_token_ids": list(getattr(resolved, "all_special_ids", ())), + "content_sha256": _tokenizer_content_sha256(resolved), + } + + +@contextmanager +def _temporary_eval(model: Any) -> Iterator[None]: + was_training = getattr(model, "training", None) + eval_method = getattr(model, "eval", None) + train_method = getattr(model, "train", None) + if ( + not isinstance(was_training, bool) + or not callable(eval_method) + or not callable(train_method) + ): + yield + return + eval_method() + try: + yield + finally: + train_method(was_training) + + +def _software_versions() -> dict[str, str | None]: + try: + import fastplms + + fastplms_version = fastplms.__version__ + except (AttributeError, ImportError): + fastplms_version = None + try: + import safetensors + + safetensors_version = safetensors.__version__ + except ImportError: + safetensors_version = None + try: + import transformers + + transformers_version = transformers.__version__ + except ImportError: + transformers_version = None + return { + "fastplms": fastplms_version, + "python": platform.python_version(), + "safetensors": safetensors_version, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "transformers": transformers_version, + } + + +def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None: + """Return deterministic PEFT/adapter identity without tensor payloads.""" + + peft_config = getattr(model, "peft_config", None) + if not isinstance(peft_config, Mapping) or not peft_config: + return None + configurations: dict[str, Any] = {} + for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])): + to_dict = getattr(config, "to_dict", None) + if callable(to_dict): + value = to_dict() + else: + try: + value = vars(config) + except TypeError: + value = config + configurations[str(name)] = _fingerprint_jsonable(value) + active_adapters = getattr(model, "active_adapters", None) + if callable(active_adapters): + active_adapters = active_adapters() + return { + "active": _fingerprint_jsonable(active_adapters), + "configurations": configurations, + } + + +def _execution_identity_metadata(model: Any) -> dict[str, Any]: + """Capture runtime policy that can change persisted numerical results.""" + + parameter_dtypes = sorted( + { + str(parameter.dtype).removeprefix("torch.") + for parameter in getattr(model, "parameters", lambda: ())() + } + ) + return { + "device": _model_device(model).type, + "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)), + "parameter_dtypes": parameter_dtypes, + "software": _software_versions(), + } + + +def _biological_residue_mask( + input_ids: Tensor, + attention_mask: Tensor, + tokenizer: Any, +) -> Tensor: + """Remove padding and tokenizer-declared special tokens from M.""" + + # input_ids, attention_mask: (b, l) + M = attention_mask.to(dtype=torch.bool) # (b, l) + special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ())) + if special_ids: + specials = torch.tensor( # (n_special,) + special_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + M = M & ~torch.isin(input_ids, specials) # (b, l) + return M # (b, l) + + +def _generic_embedding_batch( + model: Any, + sequences: list[str], + *, + tokenizer: Any | None, + max_length: int | None, + truncate: bool, + need_attentions: bool, + model_kwargs: dict[str, Any], +) -> EmbeddingBatch: + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "")).lower() + if tokenizer is None: + tokenizer = getattr(model, "tokenizer", None) + + if tokenizer is None and model_type == "e1": + output = model._embed(sequences, return_attention_mask=True, **model_kwargs) + if not isinstance(output, tuple) or len(output) != 2: + raise TypeError("E1 _embed must return (X, residue_mask).") + X, M = output # (b, l, d), (b, l) + preparer = getattr(model, "prep_tokens", None) + if preparer is not None and hasattr(preparer, "get_batch_kwargs"): + prepared = preparer.get_batch_kwargs(sequences, device=X.device) + input_ids = prepared["input_ids"] # (b, l) + boundary_ids = preparer.boundary_token_ids.to( # (n_boundary,) + device=input_ids.device, dtype=input_ids.dtype + ) + # E1 wraps each raw sequence in BOS, context-label, terminal-label, + # and EOS tokens. Only amino-acid rows are biological residues. + M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids) # (b, l) + if need_attentions: + raise ValueError("parti is not available for tokenizer-free E1 embedding.") + return EmbeddingBatch( # X: (b, l, d); residue_mask: (b, l) + X=X, + residue_mask=M.to(dtype=torch.bool), + ) + if tokenizer is None: + raise ValueError("A tokenizer is required for this model's embedding path.") + + tokenize_kwargs: dict[str, Any] = { + "return_tensors": "pt", + "padding": True, + "truncation": truncate, + } + if max_length is not None and truncate: + # ``max_length`` is a biological-residue limit. Tokenizer limits include + # boundary tokens, so reserve their declared width instead of dropping + # residues at the exact boundary. + special_token_count = 0 + num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None) + if callable(num_special_tokens_to_add): + special_token_count = int(num_special_tokens_to_add(pair=False)) + tokenize_kwargs["max_length"] = max_length + special_token_count + sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None) + if callable(sequence_tokenizer): + encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs) + else: + encoded = tokenizer(sequences, **tokenize_kwargs) + device = _model_device(model) + input_ids = encoded["input_ids"].to(device) # (b, l) + attention_mask = encoded.get( # (b, l) + "attention_mask", + input_ids.new_ones(input_ids.shape), + ).to(device) + M = _biological_residue_mask(input_ids, attention_mask, tokenizer) # (b, l) + if need_attentions: + # Validate l before either the backbone or its quadratic attention graph + # is materialized. M has shape (b, l). + _validate_parti_length(M) + X = model._embed(input_ids, attention_mask, **model_kwargs) # (b, l, d) + attentions = None + if need_attentions: + output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + return_dict=True, + ) + attentions = getattr(output, "attentions", None) # each: (b, h, l, l) + if attentions is None: + raise ValueError("The model did not return attentions required by parti.") + return EmbeddingBatch( # X: (b, l, d); M: (b, l) + X=X, + residue_mask=M, + attentions=attentions, + ) + + +def _first_metadata_value(*values: Any) -> Any: + for value in values: + if isinstance(value, str): + if value.strip(): + return value + elif value is not None: + return value + return None + + +def _model_identity_metadata(model: Any) -> dict[str, Any]: + """Resolve model and checkpoint identity, including local artifact fallbacks.""" + + config = getattr(model, "config", None) + checkpoint_revision = _first_metadata_value( + getattr(config, "fastplms_checkpoint_revision", None), + getattr(config, "_commit_hash", None), + ) + return { + "model_id": _first_metadata_value( + getattr(config, "fastplms_model_id", None), + getattr(config, "_name_or_path", None), + ), + "model_revision": _first_metadata_value( + getattr(config, "_commit_hash", None), + checkpoint_revision, + ), + "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None), + "checkpoint_revision": checkpoint_revision, + "checkpoint_hash": _first_metadata_value( + getattr(model, "checkpoint_hash", None), + getattr(config, "checkpoint_hash", None), + getattr(config, "fastplms_checkpoint_hash", None), + ), + "weights_revision": getattr(config, "fastplms_weights_revision", None), + "runtime_revision": getattr(config, "fastplms_runtime_revision", None), + "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None), + "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None), + } + + +def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]: + """Yield X in logical row-major order without materializing a full copy.""" + + # X: (...) + if X.numel() == 0: + return + if X.ndim == 0: + yield X + return + trailing_elements = 1 + for size in X.shape[1:]: + trailing_elements *= int(size) + if trailing_elements <= max_elements: + rows_per_chunk = max(1, max_elements // trailing_elements) + for start in range(0, X.shape[0], rows_per_chunk): + yield X[start : start + rows_per_chunk] # (chunk_rows, ...) + return + for row in X: + yield from _bounded_tensor_chunks(row, max_elements) + + +def _model_state_sha256(model: Any) -> str: + """Hash named parameters and persistent buffers using bounded CPU copies.""" + + # Never cache this digest from tensor identity or ``Tensor._version``. + # ``Parameter.data`` and independent tensor aliases can mutate shared storage + # without changing either signal, while persisted resume identity must bind + # the authoritative bytes visible at the start of this run. + state = model.state_dict(keep_vars=True) + digest = hashlib.sha256() + for name, value in sorted(state.items()): + if not isinstance(value, Tensor): + raise TypeError(f"Model state entry {name!r} is not a tensor.") + if value.is_meta: + raise ValueError( + f"Cannot fingerprint meta-device model state entry {name!r}; pass " + "model_state_fingerprint with a caller-owned state identity." + ) + header = json.dumps( + { + "name": name, + "dtype": str(value.dtype).removeprefix("torch."), + "shape": list(value.shape), + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + digest.update(len(header).to_bytes(8, "big")) + digest.update(header) + max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size()) + for chunk in _bounded_tensor_chunks(value.detach(), max_elements): + cpu_chunk = chunk.to(device="cpu").contiguous() # chunk.shape + digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def _input_sha256(records: Iterable[EmbeddingInput]) -> str: + """Hash an ordered input stream without constructing a duplicate JSON payload.""" + + precomputed = getattr(records, "input_fingerprint", None) + if isinstance(precomputed, str): + return precomputed + digest = hashlib.sha256() + count = 0 + for record in records: + count += 1 + for value in (record.id, record.sequence): + encoded = value.encode("utf-8") + digest.update(len(encoded).to_bytes(8, "big")) + digest.update(encoded) + digest.update(count.to_bytes(8, "big")) + return digest.hexdigest() + + +def _run_fingerprint( + model: Any, + records: Sequence[EmbeddingInput], + *, + pooling: Sequence[str], + full_embeddings: bool, + max_length: int | None, + truncate: bool, + dtype: torch.dtype | None, + model_kwargs: dict[str, Any], + tokenizer_metadata: dict[str, Any], + model_state_fingerprint: str | None, + persist_output: bool, + embedding_context: Mapping[str, Any], + batch_size: int, + batch_window_size: int, + max_tokens_per_batch: int | None, +) -> tuple[str, str, str | None, str]: + input_fingerprint = _input_sha256(records) + attention_backend = _attention_backend(model) + model_identity = _model_identity_metadata(model) + if model_state_fingerprint is None and persist_output: + resolved_model_state_fingerprint = _model_state_sha256(model) + model_state_fingerprint_source = "computed" + elif model_state_fingerprint is not None: + resolved_model_state_fingerprint = model_state_fingerprint.strip() + if not resolved_model_state_fingerprint: + raise ValueError("model_state_fingerprint must not be empty.") + model_state_fingerprint_source = "caller" + else: + resolved_model_state_fingerprint = None + model_state_fingerprint_source = "not-computed" + payload = { + "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION, + "input_fingerprint": input_fingerprint, + "model_state_fingerprint": resolved_model_state_fingerprint, + "model_state_fingerprint_source": model_state_fingerprint_source, + "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}", + **model_identity, + "attention_backend": attention_backend, + "attention_kernel": _attention_kernel_metadata(attention_backend), + "layer": repr( + getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1)) + ), + "projection": getattr(model, "embedding_projection", None), + "esmc_source": getattr(model, "_esmc_source", None), + "esmc_revision": getattr(model, "_esmc_source_revision", None), + "esmc_files": getattr(model, "_esmc_source_files", None), + "token_policy": getattr(model, "embedding_token_policy", None), + "tokenizer": tokenizer_metadata, + "adapter": _adapter_identity_metadata(model), + "execution": _execution_identity_metadata(model), + "embedding_context": _fingerprint_jsonable(embedding_context), + "pooling": list(pooling), + "full_embeddings": full_embeddings, + "max_length": max_length, + "truncate": truncate, + "dtype": str(dtype) if dtype is not None else None, + "batching": { + "batch_size": batch_size, + "batch_window_size": batch_window_size, + "max_tokens_per_batch": max_tokens_per_batch, + "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"), + }, + "model_kwargs": { + key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items()) + }, + "residue_mask_policy": "attention-mask-minus-special-tokens", + } + run_fingerprint = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + return ( + input_fingerprint, + run_fingerprint, + resolved_model_state_fingerprint, + model_state_fingerprint_source, + ) + + +def _output_exists(path: str | Path, format: str) -> bool: + path = Path(path) + if format == "sqlite": + return path.is_file() + return safetensors_result_exists(path) + + +def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]: + tensor = record.tensor + if isinstance(tensor, LazyTensorReference): + dtype = tensor.dtype + shape = tensor.shape + digest = tensor.sha256 + else: + dtype = str(tensor.dtype).removeprefix("torch.") + shape = tuple(tensor.shape) + digest = tensor_sha256(tensor) + return { + "position": position, + "id": record.id, + "dtype": dtype, + "shape": shape, + "sha256": digest, + } + + +def _ordered_string_sha256(values: Sequence[str]) -> str: + digest = hashlib.sha256() + for value in values: + encoded = value.encode("utf-8") + digest.update(len(encoded).to_bytes(8, "big")) + digest.update(encoded) + digest.update(len(values).to_bytes(8, "big")) + return digest.hexdigest() + + +def _embedding_context( + model: Any, + records: Sequence[EmbeddingInput], + *, + hidden_state_source: str, + decoder_inputs: Sequence[str] | None, + decoder_input_ids: Tensor | None, + decoder_attention_mask: Tensor | None, + model_kwargs: Mapping[str, Any], +) -> tuple[dict[str, Any], tuple[str, ...] | None]: + if hidden_state_source not in {"encoder", "decoder"}: + raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.") + hidden_state_index = model_kwargs.get("hidden_state_index", -1) + if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool): + raise TypeError("hidden_state_index must be an integer.") + store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False) + if not isinstance(store_all_hidden_states, bool): + raise TypeError("store_all_hidden_states must be a boolean.") + normalized_decoder_inputs: tuple[str, ...] | None = None + has_decoder_inputs = decoder_inputs is not None + has_decoder_ids = decoder_input_ids is not None + if hidden_state_source == "encoder": + if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None: + raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.") + else: + if has_decoder_inputs == has_decoder_ids: + raise ValueError( + "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids." + ) + decoder_input_fingerprint: str | None = None + if decoder_inputs is not None: + if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence): + raise TypeError("decoder_inputs must be an aligned sequence of strings.") + normalized_decoder_inputs = tuple(decoder_inputs) + if not all(isinstance(value, str) and value for value in normalized_decoder_inputs): + raise ValueError("decoder_inputs must contain non-empty strings.") + if len(normalized_decoder_inputs) != len(records): + raise ValueError("decoder_inputs must align one-to-one with embedding inputs.") + decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs) + if decoder_attention_mask is not None: + raise ValueError("decoder_attention_mask requires decoder_input_ids.") + if decoder_input_ids is not None: + if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2: + raise ValueError("decoder_input_ids must have shape (batch, sequence).") + if decoder_input_ids.shape[0] != len(records): + raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.") + if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point(): + raise TypeError("decoder_input_ids must use an integer token dtype.") + decoder_input_fingerprint = tensor_sha256(decoder_input_ids) + decoder_mask_fingerprint: str | None = None + if decoder_attention_mask is not None: + if not isinstance(decoder_attention_mask, Tensor): + raise TypeError("decoder_attention_mask must be a tensor.") + if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape: + raise ValueError("decoder_attention_mask must match decoder_input_ids shape.") + decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask) + + context: dict[str, Any] = { + "hidden_state_source": hidden_state_source, + "hidden_state_index": hidden_state_index, + "store_all_hidden_states": store_all_hidden_states, + "decoder_input_fingerprint": decoder_input_fingerprint, + "decoder_attention_mask_fingerprint": decoder_mask_fingerprint, + "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None, + } + metadata_hook = getattr(model, "_embedding_metadata", None) + model_metadata: Mapping[str, Any] | None = None + if callable(metadata_hook): + model_metadata = metadata_hook(**context) + if not isinstance(model_metadata, Mapping): + raise TypeError("_embedding_metadata must return a mapping.") + context["model_embedding"] = _fingerprint_jsonable(model_metadata) + if hidden_state_source == "decoder": + has_decoder_batch = callable(getattr(model, "_embedding_batch", None)) + declares_decoder_stack = ( + model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder" + ) + if not has_decoder_batch or not declares_decoder_stack: + raise ValueError( + f"{model.__class__.__name__} does not declare decoder embedding support." + ) + return context, normalized_decoder_inputs + + +def _planned_batches( + records: Sequence[EmbeddingInput], + positions: range, + *, + batch_size: int, + max_tokens_per_batch: int | None, + max_length: int | None, + truncate: bool, +) -> Iterator[list[int]]: + """Length-bucket one bounded window while retaining stable output positions.""" + + def effective_length(position: int) -> int: + length = len(records[position].sequence) + return min(length, max_length) if truncate and max_length is not None else length + + ordered = sorted(positions, key=lambda position: (-effective_length(position), position)) + batch: list[int] = [] + longest = 0 + for position in ordered: + length = effective_length(position) + if max_tokens_per_batch is not None and length > max_tokens_per_batch: + raise ValueError( + f"Input at position {position} has {length} residues, exceeding " + f"max_tokens_per_batch={max_tokens_per_batch}." + ) + candidate_longest = max(longest, length) + exceeds_tokens = ( + max_tokens_per_batch is not None + and candidate_longest * (len(batch) + 1) > max_tokens_per_batch + ) + if batch and (len(batch) >= batch_size or exceeds_tokens): + yield batch + batch = [] + longest = 0 + batch.append(position) + longest = max(longest, length) + if batch: + yield batch + + +def embed_dataset( + model: Any, + inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path), + *, + batch_size: int = 2, + pooling: str | Sequence[str] | None = None, + full_embeddings: bool = False, + output: str | Path | None = None, + format: str = "safetensors", + resume: bool = True, + tokenizer: Any | None = None, + max_length: int | None = None, + truncate: bool = True, + dtype: torch.dtype | None = torch.float32, + shard_size: int = 2 * 1024**3, + model_state_fingerprint: str | None = None, + batch_window_size: int | None = None, + max_tokens_per_batch: int | None = None, + hidden_state_source: str = "encoder", + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: Tensor | None = None, + decoder_attention_mask: Tensor | None = None, + _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None, + _embedding_batch_identity: Mapping[str, Any] | None = None, + _allowed_unsupported_pooling: Sequence[str] = (), + **model_kwargs: Any, +) -> EmbeddingResult: + """Embed protein sequences with stable ordering and residue-only pooling.""" + + for name, value in ( + ("batch_size", batch_size), + ("shard_size", shard_size), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be a positive integer.") + if value <= 0: + raise ValueError(f"{name} must be a positive integer.") + for optional_name, optional_value in ( + ("max_length", max_length), + ("max_tokens_per_batch", max_tokens_per_batch), + ("batch_window_size", batch_window_size), + ): + if optional_value is not None and ( + not isinstance(optional_value, int) or isinstance(optional_value, bool) + ): + raise TypeError(f"{optional_name} must be a positive integer when provided.") + if optional_value is not None and optional_value <= 0: + raise ValueError(f"{optional_name} must be a positive integer when provided.") + for name, value in ( + ("full_embeddings", full_embeddings), + ("resume", resume), + ("truncate", truncate), + ): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a boolean.") + if not isinstance(format, str): + raise TypeError("format must be a string.") + if output is not None and not isinstance(output, (str, Path)): + raise TypeError("output must be a path or None.") + if model_state_fingerprint is not None and ( + not isinstance(model_state_fingerprint, str) or not model_state_fingerprint + ): + raise ValueError("model_state_fingerprint must be a non-empty string when provided.") + if hidden_state_source not in {"encoder", "decoder"}: + raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.") + hidden_state_index = model_kwargs.get("hidden_state_index", -1) + if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool): + raise TypeError("hidden_state_index must be an integer.") + store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False) + if not isinstance(store_all_hidden_states, bool): + raise TypeError("store_all_hidden_states must be a boolean.") + if decoder_input_ids is not None: + if not isinstance(decoder_input_ids, Tensor): + raise TypeError("decoder_input_ids must be a tensor.") + if decoder_input_ids.is_meta: + raise ValueError("decoder_input_ids cannot be a meta tensor.") + if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0: + raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).") + if decoder_input_ids.dtype not in {torch.int32, torch.int64}: + raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.") + if decoder_attention_mask is not None: + if not isinstance(decoder_attention_mask, Tensor): + raise TypeError("decoder_attention_mask must be a tensor.") + if decoder_attention_mask.is_meta: + raise ValueError("decoder_attention_mask cannot be a meta tensor.") + if decoder_attention_mask.is_complex() or not bool( + torch.isfinite(decoder_attention_mask).all() + ): + raise ValueError("decoder_attention_mask must contain finite binary values.") + if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()): + raise ValueError("decoder_attention_mask must contain finite binary values.") + pooling_names = ( + (("mean",) if not full_embeddings else ()) + if pooling is None + else ((pooling,) if isinstance(pooling, str) else tuple(pooling)) + ) + if full_embeddings and pooling is not None: + raise ValueError("full_embeddings=True cannot be combined with pooling.") + if not full_embeddings and not pooling_names: + raise ValueError("pooling is required unless full_embeddings=True.") + pooler = Pooler(pooling_names) if pooling_names else None + + if batch_size <= 0: + raise ValueError("batch_size must be positive.") + if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"): + raise ValueError("Writing pickle-based .pth embeddings is not supported.") + if format not in _SUPPORTED_STORAGE_FORMATS: + raise ValueError("format must be 'safetensors' or 'sqlite'.") + if max_length is not None and max_length <= 0: + raise ValueError("max_length must be positive when provided.") + if max_tokens_per_batch is not None and max_tokens_per_batch <= 0: + raise ValueError("max_tokens_per_batch must be positive when provided.") + if not isinstance(dtype, (torch.dtype, type(None))): + raise TypeError("dtype must be a torch.dtype or None.") + if batch_window_size is not None and batch_window_size <= 0: + raise ValueError("batch_window_size must be positive when provided.") + if _embedding_batch_fn is not None and not callable(_embedding_batch_fn): + raise TypeError("_embedding_batch_fn must be callable when provided.") + if _embedding_batch_fn is not None and _embedding_batch_identity is None: + raise ValueError( + "_embedding_batch_identity is required with _embedding_batch_fn so persisted " + "runs bind the family-specific embedding behavior." + ) + if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping): + raise TypeError("_embedding_batch_identity must be a mapping when provided.") + if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance( + _allowed_unsupported_pooling, Sequence + ): + raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.") + if not all(isinstance(name, str) for name in _allowed_unsupported_pooling): + raise TypeError("_allowed_unsupported_pooling must contain only strings.") + allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling) + if allowed_unsupported_pooling and _embedding_batch_fn is None: + raise ValueError( + "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn." + ) + resolved_batch_window_size = ( + batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER + if batch_window_size is None + else batch_window_size + ) + if resolved_batch_window_size < batch_size: + raise ValueError("batch_window_size must be at least batch_size.") + records = _normalize_inputs(inputs, disk_backed=output is not None) + _validate_untruncated_lengths( + records, + max_length=max_length, + truncate=truncate, + ) + pooling_names = ( + (("mean",) if not full_embeddings else ()) + if pooling is None + else ((pooling,) if isinstance(pooling, str) else tuple(pooling)) + ) + if full_embeddings: + if pooling is not None: + raise ValueError("full_embeddings=True cannot be combined with pooling.") + elif not pooling_names: + raise ValueError("pooling is required unless full_embeddings=True.") + store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False)) + if store_all_hidden_states and not full_embeddings: + raise ValueError("store_all_hidden_states=True requires full_embeddings=True.") + + unsupported = set(getattr(model, "embedding_unsupported_pooling", ())) + unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported) + if unknown_pooling_overrides: + raise ValueError( + "_allowed_unsupported_pooling may only override poolers declared unsupported " + f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}." + ) + unsupported.difference_update(allowed_unsupported_pooling) + requested_unsupported = unsupported.intersection(pooling_names) + if requested_unsupported: + raise ValueError( + f"{model.__class__.__name__} does not support pooling operations " + f"{sorted(requested_unsupported)}." + ) + + # Constructing the pooler validates names and duplicate operations before + # any checkpoint hashing, tokenization, or inference occurs. + pooler = Pooler(pooling_names) if pooling_names else None + embedding_context, normalized_decoder_inputs = _embedding_context( + model, + records, + hidden_state_source=hidden_state_source, + decoder_inputs=decoder_inputs, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + model_kwargs=model_kwargs, + ) + if _embedding_batch_identity is not None: + embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity) + if allowed_unsupported_pooling: + embedding_context["family_adapter_pooling_override"] = sorted( + allowed_unsupported_pooling + ) + + tokenizer_metadata = _tokenizer_metadata(model, tokenizer) + ( + input_fingerprint, + run_fingerprint, + resolved_model_state_fingerprint, + model_state_fingerprint_source, + ) = _run_fingerprint( + model, + records, + pooling=pooling_names, + full_embeddings=full_embeddings, + max_length=max_length, + truncate=truncate, + dtype=dtype, + model_kwargs=model_kwargs, + tokenizer_metadata=tokenizer_metadata, + model_state_fingerprint=model_state_fingerprint, + persist_output=output is not None, + embedding_context=embedding_context, + batch_size=batch_size, + batch_window_size=resolved_batch_window_size, + max_tokens_per_batch=max_tokens_per_batch, + ) + output_already_exists = output is not None and _output_exists(output, format) + existing: EmbeddingResult | None = None + start_position = 0 + if output is not None and resume and output_already_exists: + if format == "sqlite": + try: + existing = load_sqlite_result(output, run_id=run_fingerprint) + except KeyError: + existing = load_result(output, format=format) + else: + existing = load_result(output, format=format) + if existing.metadata.get("fingerprint_schema_version") != (_RUN_FINGERPRINT_SCHEMA_VERSION): + raise ValueError( + "Existing embeddings use an incompatible run fingerprint schema; " + "choose another output or set resume=False." + ) + if existing.metadata.get("run_fingerprint") != run_fingerprint: + raise ValueError( + "Existing embeddings were produced by a different run fingerprint; " + "choose another output or set resume=False." + ) + if len(existing) > len(records): + raise ValueError( + "Existing embeddings are not an ordered prefix of the requested inputs." + ) + prefix_matches = all( + (observed.id, observed.sequence) == (expected.id, expected.sequence) + for expected, observed in zip(records, existing, strict=False) + ) + if not prefix_matches: + raise ValueError( + "Existing embeddings are not an ordered prefix of the requested inputs." + ) + if len(existing) == len(records) and existing.metadata.get("complete", True): + return existing + start_position = len(existing) + + sqlite_run_id: str | None = None + sqlite_replace_on_first_commit = False + sqlite_initial_metadata: dict[str, Any] | None = None + if output is not None and format == "sqlite": + sqlite_initial_metadata = { + "format_version": 1, + "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION, + "run_fingerprint": run_fingerprint, + "input_fingerprint": input_fingerprint, + "model_state_fingerprint": resolved_model_state_fingerprint, + "model_state_fingerprint_source": model_state_fingerprint_source, + "complete": False, + } + sqlite_run_id = run_fingerprint + if not resume and output_already_exists: + try: + load_sqlite_result(output, run_id=run_fingerprint) + except KeyError: + pass + else: + # Keep an exact prior run readable until replacement inference + # has produced the first complete commit window. + sqlite_replace_on_first_commit = True + if not sqlite_replace_on_first_commit: + initialize_sqlite_run( + output, + sqlite_initial_metadata, + resume=resume, + ) + + stream_safetensors = output is not None and format == "safetensors" + attention_backend = _attention_backend(model) + output_records: list[EmbeddingRecord] = ( + [] if sqlite_run_id is not None or stream_safetensors else list(existing or ()) + ) + output_descriptors: list[dict[str, Any]] | None = [] if output is None else None + pool_slices: dict[str, tuple[int, int]] = {} + if existing and pooler is not None: + pooled_width = existing[0].load_tensor().shape[-1] + if pooled_width % len(pooling_names) != 0: + raise ValueError("Stored pooled width is inconsistent with pooling metadata.") + pool_slices = pooler.output_slices(pooled_width // len(pooling_names)) + + safetensors_writer: SafetensorsStreamWriter | None = None + if stream_safetensors: + if output is None: + raise RuntimeError("Safetensors streaming was enabled without an output destination.") + transactional_overwrite = output_already_exists and not resume + safetensors_writer = SafetensorsStreamWriter( + output, + { + "format_version": 1, + "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION, + "run_fingerprint": run_fingerprint, + "input_fingerprint": input_fingerprint, + "model_state_fingerprint": resolved_model_state_fingerprint, + "model_state_fingerprint_source": model_state_fingerprint_source, + "complete": False, + }, + shard_size=shard_size, + existing=existing or (), + reuse_existing=bool(resume and existing is not None), + publish_initial=not transactional_overwrite, + publish_incremental=not transactional_overwrite, + ) + need_attentions = "parti" in pooling_names + + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "")).lower() + resolved_tokenizer = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None) + with _temporary_eval(model), torch.inference_mode(): + for window_start in range(start_position, len(records), resolved_batch_window_size): + window_stop = min(window_start + resolved_batch_window_size, len(records)) + window_records = records[window_start:window_stop] + if not isinstance(window_records, Sequence): + raise RuntimeError("The immutable embedding spool returned a non-sequence window.") + window_results: dict[int, EmbeddingRecord] = {} + for local_positions in _planned_batches( + window_records, + range(len(window_records)), + batch_size=batch_size, + max_tokens_per_batch=max_tokens_per_batch, + max_length=max_length, + truncate=truncate, + ): + batch_positions = [window_start + position for position in local_positions] + batch_records = [window_records[position] for position in local_positions] + sequences = [ + record.sequence[:max_length] + if truncate and max_length is not None + else record.sequence + for record in batch_records + ] + batch_model_kwargs = dict(model_kwargs) + if model_type == "fast_ankh" or hidden_state_source == "decoder": + batch_model_kwargs["hidden_state_source"] = hidden_state_source + if normalized_decoder_inputs is not None: + batch_model_kwargs["decoder_inputs"] = [ + normalized_decoder_inputs[position] for position in batch_positions + ] + if decoder_input_ids is not None: + # decoder_input_ids: (n_records, l_decoder) + indices = torch.tensor( # (b,) + batch_positions, + device=decoder_input_ids.device, + dtype=torch.long, + ) + batch_model_kwargs["decoder_input_ids"] = ( # (b, l_decoder) + decoder_input_ids.index_select(0, indices) + ) + if decoder_attention_mask is not None: + # decoder_attention_mask: (n_records, l_decoder) + indices = torch.tensor( # (b,) + batch_positions, + device=decoder_attention_mask.device, + dtype=torch.long, + ) + batch_model_kwargs["decoder_attention_mask"] = ( + decoder_attention_mask.index_select(0, indices) # (b, l_decoder) + ) + custom_batch = _embedding_batch_fn or getattr(model, "_embedding_batch", None) + if custom_batch is not None: + if model_type == "fast_ankh": + batch = custom_batch( + sequences, + tokenizer=resolved_tokenizer, + max_length=max_length, + truncate=truncate, + need_attentions=need_attentions, + **batch_model_kwargs, + ) + else: + batch = custom_batch(sequences, **batch_model_kwargs) + if not isinstance(batch, EmbeddingBatch): + raise TypeError("_embedding_batch must return EmbeddingBatch.") + else: + batch = _generic_embedding_batch( + model, + sequences, + tokenizer=tokenizer, + max_length=max_length, + truncate=truncate, + need_attentions=need_attentions, + model_kwargs=batch_model_kwargs, + ) + X = batch.X # (b, l, d) or (b, n_states, l, d) + raw_mask = batch.residue_mask # (b, l) + if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor): + raise TypeError("Embedding batches must provide Tensor X and residue_mask.") + if X.is_meta or raw_mask.is_meta: + raise ValueError("Embedding batches cannot contain meta tensors.") + if not X.is_floating_point(): + raise TypeError("Embedding batches must use a floating-point X dtype.") + if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()): + raise ValueError("Embedding residue_mask must contain finite binary values.") + if not bool(((raw_mask == 0) | (raw_mask == 1)).all()): + raise ValueError("Embedding residue_mask must contain finite binary values.") + M = raw_mask.to(device=X.device, dtype=torch.bool) # (b, l) + valid_X_shape = ( + X.ndim == 3 + and X.shape[0] == len(batch_records) + and X.shape[-1] > 0 + and M.shape == X.shape[:2] + ) + valid_all_states_shape = ( + X.ndim == 4 + and store_all_hidden_states + and full_embeddings + and X.shape[0] == len(batch_records) + and X.shape[1] > 0 + and X.shape[-1] > 0 + and M.shape == (X.shape[0], X.shape[2]) + ) + if not (valid_X_shape or valid_all_states_shape): + raise ValueError( + "Embedding batches must provide X with shape (b, l, d), or " + "(b, states, l, d) when storing all hidden states, and " + "residue_mask with shape (b, l)." + ) + if not bool(M.any(dim=1).all()): + raise ValueError("Every embedding sample must contain a biological residue.") + finite_selected = ( # X.shape + torch.isfinite(X) | ~M.unsqueeze(-1) + if X.ndim == 3 + else torch.isfinite(X) | ~M[:, None, :, None] + ) + if not bool(finite_selected.all()): + raise ValueError("Biological residue embeddings produced non-finite output.") + if need_attentions: + # Validate the biological graph only after mask integrity is established. + _validate_parti_length(M) + if dtype is not None: + X = X.to(dtype=dtype) # unchanged shape + + if full_embeddings: + if X.ndim == 4: + values = [ + X_i[:, M_i, :].detach().cpu() # (n_states, r_i, d) + for X_i, M_i in zip(X, M, strict=True) + ] + else: + values = [ + X_i[M_i].detach().cpu() # (r_i, d) + for X_i, M_i in zip(X, M, strict=True) + ] + else: + if pooler is None: + raise RuntimeError( + "Pooled embedding output was requested without an initialized pooler." + ) + Y = pooler( # (b, n_poolers * d) + X, + M, + attentions=batch.attentions, + attention_backend=attention_backend, + ) + pool_slices = pooler.output_slices(X.shape[-1]) + values = list(Y.detach().cpu().unbind(0)) # each: (n_poolers * d,) + for position, record, value in zip( + batch_positions, batch_records, values, strict=True + ): + window_results[position] = EmbeddingRecord(record.id, record.sequence, value) + + new_records = [ + window_results[position] for position in range(window_start, window_stop) + ] + if output_descriptors is not None: + output_descriptors.extend( + _output_descriptor(window_start + offset, record) + for offset, record in enumerate(new_records) + ) + if output is not None and sqlite_run_id is not None: + append_sqlite_records( + output, + sqlite_run_id, + window_start, + new_records, + replace_metadata=( + sqlite_initial_metadata if sqlite_replace_on_first_commit else None + ), + ) + sqlite_replace_on_first_commit = False + elif safetensors_writer is not None: + safetensors_writer.append(new_records) + else: + output_records.extend(new_records) + + software_versions = _software_versions() + projection = getattr(model, "embedding_projection", None) + resolved_layer = getattr( + model, + "embedding_layer", + model_kwargs.get("hidden_state_index", -1), + ) + token_policy = getattr( + model, + "embedding_token_policy", + { + "unit": "residue", + "include": ["biological residues"], + "exclude": [ + "BOS", + "EOS", + "padding", + "chain delimiters", + "non-protein tokens", + ], + }, + ) + model_identity = _model_identity_metadata(model) + metadata: dict[str, Any] = { + "format_version": 1, + "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION, + "run_fingerprint": run_fingerprint, + "input_fingerprint": input_fingerprint, + "model_state_fingerprint": resolved_model_state_fingerprint, + "model_state_fingerprint_source": model_state_fingerprint_source, + "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}", + **model_identity, + "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model", + "attention_backend": attention_backend, + "attention_kernel": _attention_kernel_metadata(attention_backend), + "layer": resolved_layer, + "projection": projection, + "esmc_source": getattr(model, "_esmc_source", None), + "esmc_revision": getattr(model, "_esmc_source_revision", None), + "esmc_files": getattr(model, "_esmc_source_files", None), + "token_policy": token_policy, + "tokenizer": tokenizer_metadata, + **embedding_context, + "pooling": list(pooling_names), + "pool_slices": pool_slices, + "full_embeddings": full_embeddings, + "max_length": max_length, + "truncate": truncate, + "truncation": {"enabled": truncate, "max_length": max_length}, + "batching": { + "batch_size": batch_size, + "batch_window_size": resolved_batch_window_size, + "max_tokens_per_batch": max_tokens_per_batch, + "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"), + "ordering": "bounded-length-bucketed-stable-output", + "resume_commit_granularity": ( + "not-applicable" + if output is None + else "batch-window" + if format == "sqlite" + else "shard-flush" + ), + }, + "residue_mask_policy": "biological-residues-only", + "record_count": len(records), + "descriptor_index": ( + "memory-metadata" + if output is None + else "sqlite-records" + if format == "sqlite" + else "safetensors-generation-index" + ), + "storage_format": format if output is not None else "memory", + "software": software_versions, + "execution": _execution_identity_metadata(model), + "adapter": _adapter_identity_metadata(model), + "torch_version": software_versions["torch"], + "transformers_version": software_versions["transformers"], + "complete": True, + } + if output_descriptors is not None: + metadata["outputs"] = output_descriptors + metadata["tensor_hashes"] = [item["sha256"] for item in output_descriptors] + status = getattr(model, "esmc_precision_status", None) + if status is not None: + metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status + if output is not None and sqlite_run_id is not None: + update_sqlite_run_metadata(output, sqlite_run_id, metadata) + return load_sqlite_result(output, run_id=sqlite_run_id) + if safetensors_writer is not None: + return safetensors_writer.publish(complete=True, metadata=metadata) + result = EmbeddingResult(output_records, metadata) + if output is not None: + return save_result(result, output, format=format, shard_size=shard_size) + return result + + +class EmbeddingMixin: + """Small delegation mixin shared by FastPLMs model classes.""" + + def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult: + return embed_dataset(self, inputs, **kwargs) + + +__all__ = [ + "EmbeddingMixin", + "embed_dataset", + "iter_fasta", + "parse_fasta", + "select_hidden_state_embeddings", +] diff --git a/src/fastplms/embeddings/storage.py b/src/fastplms/embeddings/storage.py new file mode 100644 index 0000000..7fc4d1d --- /dev/null +++ b/src/fastplms/embeddings/storage.py @@ -0,0 +1,1601 @@ +"""Lossless, reproducible storage for :mod:`fastplms.embeddings`.""" + +from __future__ import annotations + +import hashlib +import io +import json +import sqlite3 +import struct +import numpy as np +import torch +from bisect import bisect_right +from collections.abc import Iterable, Iterator, Sequence +from pathlib import Path +from typing import Any, cast, overload +from uuid import uuid4 +from torch import Tensor + +from .types import ( + EmbeddingRecord, + EmbeddingResult, + LazyTensorReference, +) + + +_DTYPE_NAMES: dict[torch.dtype, str] = { + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.float32: "float32", + torch.float64: "float64", + torch.int64: "int64", + torch.int32: "int32", + torch.int16: "int16", + torch.int8: "int8", + torch.uint8: "uint8", + torch.bool: "bool", +} +_NAME_DTYPES = {name: dtype for dtype, name in _DTYPE_NAMES.items()} +DEFAULT_SHARD_SIZE = 2 * 1024**3 +_MAX_RECORDS_PER_DESCRIPTOR_SHARD = 1_024 +_TENSOR_HASH_CHUNK_BYTES = 16 * 1024**2 + + +def _jsonable(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, Path): + return str(value) + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, torch.device): + return str(value) + if value is None or isinstance(value, (str, int, float, bool)): + return value + return repr(value) + + +def _persistent_metadata( + metadata: dict[str, Any], + *, + descriptor_index: str, + record_count: int | None = None, +) -> dict[str, Any]: + """Remove per-record copies from metadata and identify the authoritative index.""" + + cleaned_value = _jsonable(metadata) + if not isinstance(cleaned_value, dict): + raise TypeError("Embedding metadata must serialize to a JSON object.") + cleaned: dict[str, Any] = cleaned_value + cleaned.pop("outputs", None) + cleaned.pop("tensor_hashes", None) + cleaned["descriptor_index"] = descriptor_index + if record_count is not None: + cleaned["record_count"] = record_count + return cleaned + + +def _tensor_bytes(X: Tensor) -> bytes: + """Return the exact contiguous byte representation of X.""" + + # X: (...) + X = X.detach().cpu().contiguous() # (...) + return X.view(torch.uint8).numpy().tobytes() + + +def _bounded_tensor_chunks(X: Tensor, max_bytes: int) -> Iterator[Tensor]: + """Yield row-major CPU chunks without materializing one full byte string.""" + + # X: (...) + flattened = X.detach().to(device="cpu").reshape(-1) # (n,) + if flattened.numel() == 0: + return + chunk_elements = max(1, max_bytes // flattened.element_size()) + for start in range(0, flattened.numel(), chunk_elements): + chunk = flattened[start : start + chunk_elements] # (n_chunk,) + if chunk.stride(0) != 1: + chunk = chunk.clone(memory_format=torch.contiguous_format) # (n_chunk,) + yield chunk # (n_chunk,) + + +def _tensor_hash_chunks(X: Tensor) -> Iterator[bytes]: + for chunk in _bounded_tensor_chunks(X, _TENSOR_HASH_CHUNK_BYTES): + yield chunk.view(torch.uint8).numpy().tobytes() + + +def tensor_sha256(X: Tensor) -> str: + """Hash dtype, shape, and exact tensor bytes.""" + + if not isinstance(X, Tensor): + raise TypeError("X must be a tensor.") + if X.dtype not in _DTYPE_NAMES: + raise TypeError(f"Unsupported tensor dtype {X.dtype}.") + if X.is_meta: + raise ValueError("Cannot hash a meta tensor without storage.") + if X.layout != torch.strided: + raise TypeError("Only strided tensors can be hashed.") + digest = hashlib.sha256() + digest.update(_DTYPE_NAMES[X.dtype].encode()) + digest.update(json.dumps(tuple(X.shape)).encode()) + for chunk in _tensor_hash_chunks(X): + digest.update(chunk) + return digest.hexdigest() + + +def _encode_tensor(X: Tensor) -> tuple[str, str, bytes]: + if X.dtype not in _DTYPE_NAMES: + raise TypeError(f"Unsupported tensor dtype {X.dtype}.") + shape = json.dumps(tuple(X.shape), separators=(",", ":")) + return _DTYPE_NAMES[X.dtype], shape, _tensor_bytes(X) + + +def _decode_tensor(dtype_name: str, shape_json: str, data: bytes) -> Tensor: + try: + dtype = _NAME_DTYPES[dtype_name] + except KeyError as error: + raise ValueError(f"Unsupported stored dtype {dtype_name!r}.") from error + shape = tuple(json.loads(shape_json)) + # uint8 is used only as a byte-level carrier, preserving BF16 bits exactly. + byte_array = np.frombuffer(data, dtype=np.uint8).copy() # (n_bytes,) + X = torch.from_numpy(byte_array).view(dtype) # (n_elements,) + return X.reshape(shape).clone() # shape + + +def _index_path(path: str | Path) -> Path: + path = Path(path) + if path.suffix == ".json": + return path + if path.suffix == ".safetensors": + return path.with_suffix(".json") + return path / "index.json" + + +def _run_manifest_path(path: str | Path) -> Path: + path = Path(path) + if path.name == "index.json": + return path.with_name("run.json") + if path.suffix == ".json": + return path.with_name(f"{path.stem}.run.json") + if path.suffix == ".safetensors": + return path.with_suffix(".run.json") + return path / "run.json" + + +def _resolve_index_child(root: Path, relative: str, *, label: str) -> Path: + relative_path = Path(relative) + candidate = (root / relative_path).resolve() + if relative_path.is_absolute() or candidate.parent != root.resolve(): + raise ValueError(f"Safetensors {label} references a file outside its output directory.") + return candidate + + +def _canonical_json_bytes(payload: dict[str, Any]) -> bytes: + return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8") + + +def _load_authoritative_index( + path: str | Path, +) -> tuple[dict[str, Any], Path, dict[str, Any]]: + """Load the index selected by the atomic run-manifest commit record.""" + + stable_index_path = _index_path(path) + run_manifest_path = _run_manifest_path(path) + if not run_manifest_path.is_file(): + raise ValueError(f"Missing safetensors run manifest: {run_manifest_path}.") + run_manifest = json.loads(run_manifest_path.read_text(encoding="utf-8")) + if not isinstance(run_manifest, dict): + raise ValueError("Safetensors run manifest must contain a JSON object.") + if run_manifest.get("format") != "fastplms-embedding-run": + raise ValueError(f"Not a FastPLMs embedding run manifest: {run_manifest_path}.") + version = run_manifest.get("version") + index_reference = run_manifest.get("index") + if not isinstance(index_reference, dict): + raise ValueError("Safetensors run manifest contains an invalid index reference.") + if version == 1: + snapshot = run_manifest.get("index_payload") + if isinstance(snapshot, dict): + payload = snapshot + index_bytes = _canonical_json_bytes(payload) + elif snapshot is None: + index_bytes = stable_index_path.read_bytes() + payload = json.loads(index_bytes.decode("utf-8")) + if not isinstance(payload, dict): + raise ValueError("Safetensors index must contain a JSON object.") + else: + raise ValueError("Safetensors run manifest contains an invalid index snapshot.") + expected = { + "file": stable_index_path.name, + "sha256": hashlib.sha256(index_bytes).hexdigest(), + } + index_path = stable_index_path + elif version == 2: + relative = index_reference.get("file") + if not isinstance(relative, str): + raise ValueError("Safetensors run manifest index file is invalid.") + index_path = _resolve_index_child(stable_index_path.parent, relative, label="run manifest") + index_bytes = index_path.read_bytes() + payload = json.loads(index_bytes.decode("utf-8")) + if not isinstance(payload, dict): + raise ValueError("Safetensors generation index must contain a JSON object.") + if payload.get("version") != 2: + raise ValueError("Safetensors v2 run manifest must reference a v2 generation index.") + expected = { + "file": relative, + "sha256": hashlib.sha256(index_bytes).hexdigest(), + } + else: + raise ValueError(f"Unsupported safetensors run manifest version {version!r}.") + if index_reference != expected: + raise ValueError("Safetensors run manifest does not match its index.") + if payload.get("format") != "fastplms-embedding-safetensors": + raise ValueError(f"Not a FastPLMs embedding index: {index_path}.") + record_count = payload.get("record_count") + if record_count is None: + legacy_records = payload.get("records", ()) + if not isinstance(legacy_records, list): + raise ValueError("Safetensors index contains invalid records.") + record_count = len(legacy_records) + if not isinstance(record_count, int) or isinstance(record_count, bool) or record_count < 0: + raise ValueError("Safetensors record count must be a non-negative integer.") + if run_manifest.get("record_count") != record_count: + raise ValueError("Safetensors run manifest record count does not match its index.") + metadata = payload.get("metadata", {}) + if not isinstance(metadata, dict): + raise ValueError("Safetensors index metadata must contain a JSON object.") + if metadata.get("record_count", record_count) != record_count: + raise ValueError("Safetensors metadata record count does not match its index.") + if version == 1 and run_manifest.get("metadata") != payload.get("metadata"): + raise ValueError("Safetensors run manifest metadata does not match its index.") + return payload, index_path, run_manifest + + +def safetensors_result_exists(path: str | Path) -> bool: + """Return whether an authoritative committed safetensors run exists.""" + + try: + _load_authoritative_index(path) + except (OSError, ValueError, json.JSONDecodeError): + return False + return True + + +def _load_safetensor(path: Path, key: str) -> Tensor: + try: + from safetensors import safe_open + except ImportError as error: + raise ImportError("Loading embeddings requires the 'safetensors' package.") from error + with safe_open(path, framework="pt", device="cpu") as handle: + return cast(Tensor, handle.get_tensor(key)) + + +def _safetensors_shard_prefix(path: str | Path) -> str: + requested_path = Path(path) + if requested_path.suffix in {".json", ".safetensors"}: + return f"{requested_path.stem}-embeddings" + return "embeddings" + + +def _authoritative_index_payload(path: str | Path) -> dict[str, Any] | None: + """Return the last atomically committed generation index when available.""" + + try: + payload, _, _ = _load_authoritative_index(path) + except (OSError, ValueError, json.JSONDecodeError): + return None + return payload + + +def _referenced_shards( + index_path: Path, + payload: dict[str, Any] | None = None, +) -> set[Path]: + if payload is None: + payload = _authoritative_index_payload(index_path) + if payload is None: + return set() + shards: set[Path] = set() + for descriptor_shard in payload.get("descriptor_shards", ()): + tensor_file = descriptor_shard.get("tensor_file") + if isinstance(tensor_file, str): + candidate = _resolve_index_child( + index_path.parent, tensor_file, label="descriptor index" + ) + shards.add(candidate) + for item in payload.get("records", ()): + relative = item.get("tensor", {}).get("file") + if not isinstance(relative, str): + continue + candidate = (index_path.parent / relative).resolve() + if candidate.parent == index_path.parent.resolve(): + shards.add(candidate) + return shards + + +def _validate_tensor_descriptor( + tensor: dict[str, Any], +) -> tuple[str, str, tuple[int, ...], str]: + key = tensor.get("key") + if not isinstance(key, str) or not key: + raise ValueError("Safetensors descriptor tensor key is invalid.") + dtype = tensor.get("dtype") + if not isinstance(dtype, str) or dtype not in _NAME_DTYPES: + raise ValueError("Safetensors descriptor tensor dtype is invalid.") + raw_shape = tensor.get("shape") + if not isinstance(raw_shape, (list, tuple)) or not all( + isinstance(dimension, int) and not isinstance(dimension, bool) and dimension >= 0 + for dimension in raw_shape + ): + raise ValueError("Safetensors descriptor tensor shape is invalid.") + sha256 = tensor.get("sha256") + if ( + not isinstance(sha256, str) + or len(sha256) != 64 + or sha256 != sha256.lower() + or any(character not in "0123456789abcdef" for character in sha256) + ): + raise ValueError("Safetensors descriptor tensor SHA-256 is invalid.") + return key, dtype, tuple(raw_shape), sha256 + + +def _record_from_safetensors_descriptor(root: Path, item: dict[str, Any]) -> EmbeddingRecord: + if not isinstance(item, dict): + raise ValueError("Safetensors record descriptor must contain a JSON object.") + record_id = item.get("id") + sequence = item.get("sequence") + if not isinstance(record_id, str) or not record_id: + raise ValueError("Safetensors descriptor record ID is invalid.") + if not isinstance(sequence, str) or not sequence: + raise ValueError("Safetensors descriptor sequence is invalid.") + tensor = item.get("tensor") + if not isinstance(tensor, dict): + raise ValueError("Safetensors descriptor is missing tensor metadata.") + relative = tensor.get("file") + if not isinstance(relative, str) or not relative: + raise ValueError("Safetensors descriptor tensor file is invalid.") + key, dtype, shape, sha256 = _validate_tensor_descriptor(tensor) + tensor_path = _resolve_index_child(root, relative, label="descriptor") + if not tensor_path.is_file(): + raise ValueError(f"Safetensors tensor shard is missing: {relative}.") + + def load_tensor() -> Tensor: + return _load_safetensor(tensor_path, key) + + reference = LazyTensorReference( + source=str(tensor_path), + key=key, + dtype=dtype, + shape=shape, + sha256=sha256, + _loader=load_tensor, + ) + return EmbeddingRecord(record_id, sequence, reference) + + +class _SafetensorsRecordSequence(Sequence[EmbeddingRecord]): + """Lazy immutable view over bounded descriptor JSONL shards.""" + + _fastplms_immutable_sequence = True + + def __init__(self, root: Path, descriptor_shards: Sequence[dict[str, Any]]) -> None: + if not isinstance(descriptor_shards, (list, tuple)): + raise ValueError("Safetensors generation index has invalid descriptor shards.") + self.root = root + self.shards = tuple(descriptor_shards) + cumulative: list[int] = [] + total = 0 + for shard in self.shards: + if not isinstance(shard, dict): + raise ValueError("Safetensors descriptor shard entry is invalid.") + relative = shard.get("file") + declared_count = shard.get("count") + if ( + not isinstance(declared_count, int) + or isinstance(declared_count, bool) + or declared_count < 0 + ): + raise ValueError("Safetensors descriptor shard count is invalid.") + declared_sha256 = shard.get("sha256") + if not isinstance(declared_sha256, str) or len(declared_sha256) != 64: + raise ValueError("Safetensors descriptor shard SHA-256 is invalid.") + if not isinstance(relative, str): + raise ValueError("Safetensors descriptor index file is invalid.") + descriptor_path = _resolve_index_child(root, relative, label="index") + tensor_file = shard.get("tensor_file") + if not isinstance(tensor_file, str): + raise ValueError("Safetensors descriptor tensor file is invalid.") + tensor_path = _resolve_index_child(root, tensor_file, label="index") + if not tensor_path.is_file(): + raise ValueError(f"Safetensors tensor shard is missing: {tensor_file}.") + digest = hashlib.sha256() + count = 0 + with descriptor_path.open("rb") as handle: + for line in handle: + digest.update(line) + if line.strip(): + item = json.loads(line) + if not isinstance(item, dict): + raise ValueError("Safetensors record descriptor must be a JSON object.") + item_tensor = item.get("tensor") + if not isinstance(item_tensor, dict): + raise ValueError("Safetensors descriptor is missing tensor metadata.") + item_tensor_file = item_tensor.get("file") + if not isinstance(item_tensor_file, str): + raise ValueError("Safetensors descriptor tensor file is invalid.") + _resolve_index_child(root, item_tensor_file, label="descriptor") + if item_tensor_file != tensor_file: + raise ValueError( + "Safetensors descriptor tensor file does not match its shard." + ) + count += 1 + _validate_tensor_descriptor(item_tensor) + if digest.hexdigest() != declared_sha256 or count != declared_count: + raise ValueError( + f"Safetensors descriptor shard failed integrity validation: {relative}." + ) + total += count + cumulative.append(total) + self._cumulative = tuple(cumulative) + self._count = total + + def __len__(self) -> int: + return self._count + + def _iter_shard(self, shard_index: int) -> Iterator[EmbeddingRecord]: + descriptor_path = _resolve_index_child( + self.root, str(self.shards[shard_index]["file"]), label="index" + ) + with descriptor_path.open("r", encoding="utf-8") as handle: + for line in handle: + if line.strip(): + yield _record_from_safetensors_descriptor(self.root, json.loads(line)) + + def __iter__(self) -> Iterator[EmbeddingRecord]: + for shard_index in range(len(self.shards)): + yield from self._iter_shard(shard_index) + + @overload + def __getitem__(self, index: int, /) -> EmbeddingRecord: ... + + @overload + def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ... + + def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]: + if isinstance(index, slice): + start, stop, step = index.indices(self._count) + return [self[position] for position in range(start, stop, step)] + position = index + self._count if index < 0 else index + if position < 0 or position >= self._count: + raise IndexError(index) + shard_index = bisect_right(self._cumulative, position) + previous = self._cumulative[shard_index - 1] if shard_index else 0 + local_position = position - previous + for offset, record in enumerate(self._iter_shard(shard_index)): + if offset == local_position: + return record + raise IndexError(index) + + +class SafetensorsStreamWriter: + """Bounded-memory, resumable publisher with immutable retained generations.""" + + def __init__( + self, + path: str | Path, + metadata: dict[str, Any], + *, + shard_size: int = DEFAULT_SHARD_SIZE, + existing: Iterable[EmbeddingRecord] = (), + reuse_existing: bool = False, + publish_initial: bool = True, + publish_incremental: bool = True, + ) -> None: + try: + from safetensors.torch import save_file + except ImportError as error: + raise ImportError("Saving embeddings requires the 'safetensors' package.") from error + if shard_size <= 0: + raise ValueError("shard_size must be positive.") + + self.path = Path(path) + self.index_path = _index_path(path) + self.run_manifest_path = _run_manifest_path(path) + self.index_path.parent.mkdir(parents=True, exist_ok=True) + self.metadata = _persistent_metadata( + metadata, + descriptor_index="safetensors-generation-index", + record_count=0, + ) + self.shard_size = shard_size + self.publish_incremental = publish_incremental + self._save_file = save_file + authoritative_payload = _authoritative_index_payload(path) + prefix = _safetensors_shard_prefix(path) + # A random generation identity prevents a new writer from reusing a + # previously published or interrupted generation name. Published files + # are immutable and remain available to lazy readers until explicit GC. + self._generation = uuid4().hex + self._prefix = prefix + self._shard_index = 0 + self._seed_index = 0 + self._commit_index = 0 + self._descriptor_shards: list[dict[str, Any]] = [] + self._record_count = 0 + self._current: dict[str, Tensor] = {} + self._pending: list[tuple[EmbeddingRecord, str, str, tuple[int, ...], str]] = [] + self._current_size = 0 + if reuse_existing: + if authoritative_payload is None: + raise ValueError("Cannot resume without an authoritative safetensors index.") + authoritative_metadata = authoritative_payload.get("metadata") + if not isinstance(authoritative_metadata, dict) or authoritative_metadata.get( + "run_fingerprint" + ) != self.metadata.get("run_fingerprint"): + raise ValueError("Cannot resume a safetensors run with a different fingerprint.") + expected_prefix_length = ( + len(existing) if isinstance(existing, Sequence) else sum(1 for _ in existing) + ) + if authoritative_payload.get("version") == 2: + self._descriptor_shards = list(authoritative_payload.get("descriptor_shards", ())) + self._record_count = int(authoritative_payload.get("record_count", 0)) + else: + legacy_records = list(authoritative_payload.get("records", ())) + self._record_count = len(legacy_records) + if legacy_records: + self._descriptor_shards.extend(self._write_descriptor_seed(legacy_records)) + if expected_prefix_length != self._record_count: + raise ValueError( + "The resumable safetensors prefix does not match the validated " + "embedding records." + ) + + if publish_initial: + self._publish_metadata(complete=False) + + def _write_descriptor_file( + self, + name: str, + descriptors: Sequence[dict[str, Any]], + *, + tensor_file: str, + ) -> dict[str, Any]: + temporary = self.index_path.parent / f".{name}.tmp" + destination = self.index_path.parent / name + if temporary.exists() or destination.exists(): + raise FileExistsError( + f"Refusing to reuse immutable safetensors generation path {destination}." + ) + digest = hashlib.sha256() + with temporary.open("wb") as handle: + for item in descriptors: + encoded = ( + json.dumps(item, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n" + ) + handle.write(encoded) + digest.update(encoded) + temporary.replace(destination) + return { + "file": name, + "sha256": digest.hexdigest(), + "count": len(descriptors), + "tensor_file": tensor_file, + } + + def _write_descriptor_seed(self, records: Sequence[dict[str, Any]]) -> list[dict[str, Any]]: + groups: list[tuple[str, list[dict[str, Any]]]] = [] + for record in records: + tensor_file = str(record["tensor"]["file"]) + if ( + not groups + or groups[-1][0] != tensor_file + or len(groups[-1][1]) == _MAX_RECORDS_PER_DESCRIPTOR_SHARD + ): + groups.append((tensor_file, [])) + groups[-1][1].append(record) + descriptor_shards: list[dict[str, Any]] = [] + for tensor_file, descriptors in groups: + self._seed_index += 1 + name = ( + f"{self._prefix}-records-run-{self._generation}-seed-{self._seed_index:05d}.jsonl" + ) + descriptor_shards.append( + self._write_descriptor_file(name, descriptors, tensor_file=tensor_file) + ) + return descriptor_shards + + def _write_shard(self) -> None: + if not self._current: + return + self._shard_index += 1 + name = f"{self._prefix}-run-{self._generation}-{self._shard_index:05d}.safetensors" + temporary = self.index_path.parent / f".{name}.tmp" + destination = self.index_path.parent / name + if temporary.exists() or destination.exists(): + raise FileExistsError( + f"Refusing to reuse immutable safetensors generation path {destination}." + ) + self._save_file(self._current, temporary) + temporary.replace(destination) + descriptors: list[dict[str, Any]] = [] + for record, key, dtype_name, shape, digest in self._pending: + descriptors.append( + { + "id": record.id, + "sequence": record.sequence, + "tensor": { + "file": name, + "key": key, + "dtype": dtype_name, + "shape": list(shape), + "sha256": digest, + }, + } + ) + descriptor_name = ( + f"{self._prefix}-records-run-{self._generation}-{self._shard_index:05d}.jsonl" + ) + self._descriptor_shards.append( + self._write_descriptor_file(descriptor_name, descriptors, tensor_file=name) + ) + self._record_count += len(descriptors) + self._current = {} + self._pending = [] + self._current_size = 0 + + def append( + self, + records: Iterable[EmbeddingRecord], + *, + publish: bool | None = None, + ) -> None: + """Persist records while retaining at most one shard of tensors.""" + + for record in records: + position = self._record_count + len(self._pending) + tensor = record.load_tensor().detach().cpu().contiguous() # (...) + if tensor.dtype not in _DTYPE_NAMES: + raise TypeError(f"Unsupported tensor dtype {tensor.dtype}.") + nbytes = tensor.numel() * tensor.element_size() + if nbytes > self.shard_size: + raise ValueError( + f"Embedding {position} requires {nbytes} bytes and cannot fit in a " + f"{self.shard_size}-byte safetensors shard." + ) + if self._current and ( + self._current_size + nbytes > self.shard_size + or len(self._pending) == _MAX_RECORDS_PER_DESCRIPTOR_SHARD + ): + self._write_shard() + if self.publish_incremental: + self._publish_metadata(complete=False) + position = self._record_count + key = f"embedding_{position:08d}" + self._current[key] = tensor + self._current_size += nbytes + self._pending.append( + ( + record, + key, + _DTYPE_NAMES[tensor.dtype], + tuple(tensor.shape), + tensor_sha256(tensor), + ) + ) + if publish: + self.publish(complete=False) + + def _publish_metadata( + self, + *, + complete: bool, + metadata: dict[str, Any] | None = None, + ) -> EmbeddingResult: + """Atomically expose one self-consistent metadata generation.""" + + if metadata is not None: + self.metadata = _persistent_metadata( + metadata, + descriptor_index="safetensors-generation-index", + ) + self.metadata["complete"] = complete + self.metadata["record_count"] = self._record_count + self._commit_index += 1 + payload = { + "version": 2, + "format": "fastplms-embedding-safetensors", + "metadata": self.metadata, + "record_count": self._record_count, + "descriptor_shards": self._descriptor_shards, + } + generation_index_name = ( + f"{self._prefix}-index-run-{self._generation}-{self._commit_index:05d}.json" + ) + generation_index_path = self.index_path.parent / generation_index_name + temporary_generation_index = generation_index_path.with_name( + f".{generation_index_path.name}.tmp" + ) + if temporary_generation_index.exists() or generation_index_path.exists(): + raise FileExistsError( + f"Refusing to reuse immutable safetensors generation index {generation_index_path}." + ) + encoded_index = _canonical_json_bytes(payload) + temporary_generation_index.write_bytes(encoded_index) + temporary_generation_index.replace(generation_index_path) + + index_sha256 = hashlib.sha256(encoded_index).hexdigest() + index_reference = { + "file": generation_index_name, + "sha256": index_sha256, + } + run_manifest = { + "version": 2, + "format": "fastplms-embedding-run", + "index": index_reference, + "record_count": self._record_count, + } + pointer_identity = f"{self._generation}-{self._commit_index:05d}" + temporary_manifest = self.run_manifest_path.with_name( + f".{self.run_manifest_path.name}.{pointer_identity}.tmp" + ) + temporary_manifest.write_bytes(_canonical_json_bytes(run_manifest)) + temporary_manifest.replace(self.run_manifest_path) + + # ``index.json`` is a non-authoritative convenience pointer. The run + # manifest is committed first, so interruption here cannot invalidate + # the newly committed generation. + stable_pointer = { + "version": 2, + "format": "fastplms-embedding-index-pointer", + "index": index_reference, + } + temporary_index = self.index_path.with_name( + f".{self.index_path.name}.{pointer_identity}.tmp" + ) + temporary_index.write_bytes(_canonical_json_bytes(stable_pointer)) + temporary_index.replace(self.index_path) + + return load_safetensors_result(self.index_path) + + def publish( + self, + *, + complete: bool, + metadata: dict[str, Any] | None = None, + ) -> EmbeddingResult: + """Flush the current shard and atomically expose a consistent generation.""" + + self._write_shard() + return self._publish_metadata(complete=complete, metadata=metadata) + + +def save_safetensors_result( + result: EmbeddingResult, + path: str | Path, + *, + shard_size: int = DEFAULT_SHARD_SIZE, +) -> EmbeddingResult: + """Write sharded safetensors without materializing the full result.""" + + writer = SafetensorsStreamWriter( + path, + result.metadata, + shard_size=shard_size, + publish_initial=False, + publish_incremental=False, + ) + writer.append(result, publish=False) + return writer.publish(complete=bool(result.metadata.get("complete", True))) + + +def load_safetensors_result(path: str | Path) -> EmbeddingResult: + """Load an indexed safetensors result without loading tensor payloads.""" + + payload, index_path, _ = _load_authoritative_index(path) + if payload.get("version") == 2: + lazy_records = _SafetensorsRecordSequence( + index_path.parent, payload.get("descriptor_shards", ()) + ) + if len(lazy_records) != payload.get("record_count"): + raise ValueError("Safetensors descriptor count does not match its generation index.") + return EmbeddingResult(lazy_records, payload.get("metadata", {})) + + records: list[EmbeddingRecord] = [] + for item in payload["records"]: + records.append(_record_from_safetensors_descriptor(index_path.parent, item)) + return EmbeddingResult(records, payload.get("metadata", {})) + + +def garbage_collect_safetensors_generations( + path: str | Path, + *, + dry_run: bool = True, + confirm_no_active_readers_or_writers: bool = False, +) -> tuple[Path, ...]: + """Remove non-authoritative generations after an explicit exclusivity check. + + Safetensors results retain immutable historical generations because an + already-open :class:`EmbeddingResult` resolves tensors through those exact + descriptor and shard paths. Destructive collection is therefore safe only + when the caller guarantees that no reader or writer for ``path`` remains + active. ``dry_run=True`` is the default and returns the paths that would be + removed without changing the output directory. + """ + + if not isinstance(dry_run, bool): + raise TypeError("dry_run must be a bool.") + if not isinstance(confirm_no_active_readers_or_writers, bool): + raise TypeError("confirm_no_active_readers_or_writers must be a bool.") + if not dry_run and not confirm_no_active_readers_or_writers: + raise ValueError( + "Destructive safetensors generation collection requires " + "confirm_no_active_readers_or_writers=True." + ) + + # Validate the full descriptor graph before identifying anything as stale. + load_safetensors_result(path) + payload, authoritative_index_path, _ = _load_authoritative_index(path) + stable_index_path = _index_path(path) + run_manifest_path = _run_manifest_path(path) + root = stable_index_path.parent + prefix = _safetensors_shard_prefix(path) + protected = { + stable_index_path.resolve(), + run_manifest_path.resolve(), + authoritative_index_path.resolve(), + *_referenced_shards(stable_index_path, payload), + } + for descriptor_shard in payload.get("descriptor_shards", ()): + relative = descriptor_shard.get("file") + if isinstance(relative, str): + protected.add(_resolve_index_child(root, relative, label="index").resolve()) + + candidates: set[Path] = set() + for pattern in ( + f"{prefix}-run-*-*.safetensors", + f"{prefix}-records-run-*.jsonl", + f"{prefix}-index-run-*.json", + f".{prefix}-*.tmp", + ): + candidates.update(root.glob(pattern)) + candidates.update(root.glob(f".{stable_index_path.name}.*.tmp")) + candidates.update(root.glob(f".{run_manifest_path.name}.*.tmp")) + + stale = tuple( + sorted( + (candidate for candidate in candidates if candidate.resolve() not in protected), + key=lambda candidate: candidate.name, + ) + ) + if not dry_run: + for candidate in stale: + candidate.unlink(missing_ok=True) + return stale + + +def _ensure_sqlite_schema(connection: sqlite3.Connection) -> None: + connection.executescript( + """ + PRAGMA foreign_keys = ON; + CREATE TABLE IF NOT EXISTS runs ( + run_id TEXT PRIMARY KEY, + metadata_json TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, + published_order INTEGER + ); + CREATE TABLE IF NOT EXISTS tensors ( + run_id TEXT NOT NULL, + position INTEGER NOT NULL, + dtype TEXT NOT NULL, + shape_json TEXT NOT NULL, + data BLOB NOT NULL, + sha256 TEXT NOT NULL, + PRIMARY KEY (run_id, position), + FOREIGN KEY (run_id) REFERENCES runs(run_id) ON DELETE CASCADE + ); + CREATE TABLE IF NOT EXISTS records ( + run_id TEXT NOT NULL, + position INTEGER NOT NULL, + record_id TEXT NOT NULL, + sequence TEXT NOT NULL, + PRIMARY KEY (run_id, position), + FOREIGN KEY (run_id, position) REFERENCES tensors(run_id, position) + ON DELETE CASCADE + ); + """ + ) + run_columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(runs)").fetchall()} + if "published_order" not in run_columns: + connection.execute("ALTER TABLE runs ADD COLUMN published_order INTEGER") + # Databases created before staged publication exposed every stored run. + # Preserve that view for historical runs containing committed records. + connection.execute( + "UPDATE runs SET published_order = rowid " + "WHERE published_order IS NULL AND EXISTS (" + "SELECT 1 FROM records WHERE records.run_id = runs.run_id)" + ) + connection.execute( + "CREATE INDEX IF NOT EXISTS runs_published_order_idx ON runs(published_order)" + ) + if "published_order" not in run_columns: + # Schema upgrades run before callers open their data transaction. + # End the migration transaction explicitly so BEGIN IMMEDIATE below + # remains valid on existing databases. + connection.commit() + + +def save_sqlite_result(result: EmbeddingResult, path: str | Path) -> EmbeddingResult: + """Transactionally store an ordered result in normalized SQLite tables.""" + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + run_id = str(result.metadata.get("run_fingerprint", "")) + if not run_id: + raise ValueError("SQLite results require metadata['run_fingerprint'].") + metadata_json = json.dumps( + _persistent_metadata( + result.metadata, + descriptor_index="sqlite-records", + record_count=len(result), + ), + sort_keys=True, + ) + with sqlite3.connect(path, timeout=30) as connection: + _ensure_sqlite_schema(connection) + connection.execute("PRAGMA journal_mode = WAL") + connection.execute("BEGIN IMMEDIATE") + connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,)) + connection.execute( + "INSERT INTO runs(run_id, metadata_json, published_order) " + "SELECT ?, ?, COALESCE(MAX(published_order), 0) + 1 FROM runs", + (run_id, metadata_json), + ) + for position, record in enumerate(result): + X = record.load_tensor().detach().cpu().contiguous() # (...) + dtype_name, shape_json, data = _encode_tensor(X) + digest = tensor_sha256(X) + connection.execute( + "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)", + (run_id, position, dtype_name, shape_json, data, digest), + ) + connection.execute( + "INSERT INTO records VALUES (?, ?, ?, ?)", + (run_id, position, record.id, record.sequence), + ) + connection.commit() + return load_sqlite_result(path, run_id=run_id) + + +def initialize_sqlite_run( + path: str | Path, + metadata: dict[str, Any], + *, + resume: bool, +) -> str: + """Create a resumable SQLite run without buffering tensor results.""" + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + run_id = str(metadata.get("run_fingerprint", "")) + if not run_id: + raise ValueError("SQLite runs require metadata['run_fingerprint'].") + with sqlite3.connect(path, timeout=30) as connection: + _ensure_sqlite_schema(connection) + connection.execute("PRAGMA journal_mode = WAL") + connection.execute("BEGIN IMMEDIATE") + exists = connection.execute("SELECT 1 FROM runs WHERE run_id = ?", (run_id,)).fetchone() + if exists and not resume: + connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,)) + exists = None + if exists is None: + initial_metadata = _persistent_metadata( + metadata, + descriptor_index="sqlite-records", + record_count=0, + ) + connection.execute( + "INSERT INTO runs(run_id, metadata_json) VALUES (?, ?)", + (run_id, json.dumps(initial_metadata, sort_keys=True)), + ) + connection.commit() + return run_id + + +def append_sqlite_records( + path: str | Path, + run_id: str, + start_position: int, + records: list[EmbeddingRecord], + *, + replace_metadata: dict[str, Any] | None = None, +) -> None: + """Commit one ordered embedding batch so an interrupted run can resume.""" + + if not isinstance(run_id, str) or not run_id: + raise ValueError("run_id must be a non-empty string.") + if not isinstance(start_position, int) or isinstance(start_position, bool): + raise TypeError("start_position must be a non-negative integer.") + if start_position < 0: + raise ValueError("start_position must be a non-negative integer.") + if not isinstance(records, list) or not all( + isinstance(record, EmbeddingRecord) for record in records + ): + raise TypeError("records must be a list of EmbeddingRecord values.") + + with sqlite3.connect(Path(path), timeout=30) as connection: + _ensure_sqlite_schema(connection) + connection.execute("PRAGMA journal_mode = WAL") + connection.execute("BEGIN IMMEDIATE") + if replace_metadata is not None: + replacement_run_id = str(replace_metadata.get("run_fingerprint", "")) + if replacement_run_id != run_id: + raise ValueError("Replacement metadata must match the SQLite run ID.") + initial_metadata = _persistent_metadata( + replace_metadata, + descriptor_index="sqlite-records", + record_count=0, + ) + connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,)) + connection.execute( + "INSERT INTO runs(run_id, metadata_json) VALUES (?, ?)", + (run_id, json.dumps(initial_metadata, sort_keys=True)), + ) + if connection.execute("SELECT 1 FROM runs WHERE run_id = ?", (run_id,)).fetchone() is None: + raise KeyError(f"Missing SQLite embedding run {run_id}.") + current_count, minimum_position, maximum_position = connection.execute( + "SELECT COUNT(*), MIN(position), MAX(position) FROM records WHERE run_id = ?", + (run_id,), + ).fetchone() + if current_count and (minimum_position != 0 or maximum_position != current_count - 1): + raise ValueError("SQLite embedding run has a non-contiguous record prefix.") + if start_position != current_count: + raise ValueError( + f"start_position={start_position} does not match the contiguous " + f"SQLite prefix length {current_count}." + ) + for offset, record in enumerate(records): + position = start_position + offset + X = record.load_tensor().detach().cpu().contiguous() # (...) + dtype_name, shape_json, data = _encode_tensor(X) + digest = tensor_sha256(X) + connection.execute( + "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)", + (run_id, position, dtype_name, shape_json, data, digest), + ) + connection.execute( + "INSERT INTO records VALUES (?, ?, ?, ?)", + (run_id, position, record.id, record.sequence), + ) + row = connection.execute( + "SELECT metadata_json FROM runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None: + raise KeyError(f"Missing SQLite embedding run {run_id}.") + metadata = json.loads(row[0]) + if not isinstance(metadata, dict): + raise ValueError("SQLite run metadata must contain a JSON object.") + metadata["record_count"] = start_position + len(records) + metadata["descriptor_index"] = "sqlite-records" + connection.execute( + "UPDATE runs SET metadata_json = ? WHERE run_id = ?", + (json.dumps(metadata, sort_keys=True), run_id), + ) + if records: + connection.execute( + "UPDATE runs SET published_order = (" + "SELECT COALESCE(MAX(published_order), 0) + 1 FROM runs" + ") WHERE run_id = ? AND published_order IS NULL", + (run_id,), + ) + connection.commit() + + +def update_sqlite_run_metadata(path: str | Path, run_id: str, metadata: dict[str, Any]) -> None: + """Finalize reproducibility metadata after the last streamed batch.""" + + with sqlite3.connect(Path(path), timeout=30) as connection: + row = connection.execute( + "SELECT COUNT(*) FROM records WHERE run_id = ?", (run_id,) + ).fetchone() + record_count = int(row[0]) if row is not None else 0 + cleaned_metadata = _persistent_metadata( + metadata, + descriptor_index="sqlite-records", + record_count=record_count, + ) + updated = connection.execute( + "UPDATE runs SET metadata_json = ? WHERE run_id = ?", + (json.dumps(cleaned_metadata, sort_keys=True), run_id), + ).rowcount + if updated != 1: + raise KeyError(f"Missing SQLite embedding run {run_id}.") + connection.commit() + + +def _connect_sqlite_read_only(path: Path) -> sqlite3.Connection: + if not path.is_file(): + raise FileNotFoundError(path) + return sqlite3.connect(f"{path.resolve().as_uri()}?mode=ro", uri=True, timeout=30) + + +def _validate_sqlite_result_schema(connection: sqlite3.Connection, path: Path) -> None: + tables = { + str(row[0]) + for row in connection.execute( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ).fetchall() + } + required = {"runs", "records", "tensors"} + if not required.issubset(tables): + raise ValueError( + f"Not a FastPLMs embedding SQLite database: {path}. " + "Use convert_legacy_sqlite() for a legacy embeddings table." + ) + + +def _load_sqlite_tensor(path: Path, run_id: str, position: int) -> Tensor: + with _connect_sqlite_read_only(path) as connection: + row = connection.execute( + "SELECT dtype, shape_json, data FROM tensors WHERE run_id = ? AND position = ?", + (run_id, position), + ).fetchone() + if row is None: + raise KeyError(f"Missing SQLite tensor {run_id}:{position}.") + return _decode_tensor(*row) + + +def _validate_sqlite_descriptor_row( + row: Sequence[Any], +) -> tuple[int, str, str, str, str, str]: + if len(row) != 6: + raise ValueError("SQLite embedding descriptor has an invalid column count.") + position, record_id, sequence, dtype_name, shape_json, digest = row + if not isinstance(position, int) or isinstance(position, bool) or position < 0: + raise ValueError("SQLite embedding position is invalid.") + if not isinstance(record_id, str) or not record_id: + raise ValueError("SQLite embedding record ID is invalid.") + if not isinstance(sequence, str) or not sequence: + raise ValueError("SQLite embedding sequence is invalid.") + if not isinstance(shape_json, str): + raise ValueError("SQLite embedding tensor shape is invalid.") + try: + shape = json.loads(shape_json) + except json.JSONDecodeError as error: + raise ValueError("SQLite embedding tensor shape is invalid.") from error + _validate_tensor_descriptor( + { + "key": f"embedding_{position}", + "dtype": dtype_name, + "shape": shape, + "sha256": digest, + } + ) + return position, record_id, sequence, dtype_name, shape_json, digest + + +def _sqlite_record_from_row(path: Path, run_id: str, row: Sequence[Any]) -> EmbeddingRecord: + position, record_id, sequence, dtype_name, shape_json, digest = _validate_sqlite_descriptor_row( + row + ) + + def load_tensor() -> Tensor: + return _load_sqlite_tensor(path, run_id, position) + + reference = LazyTensorReference( + source=str(path), + key=f"{run_id}:{position}", + dtype=dtype_name, + shape=tuple(json.loads(shape_json)), + sha256=digest, + _loader=load_tensor, + ) + return EmbeddingRecord(record_id, sequence, reference) + + +class _SQLiteRecordSequence(Sequence[EmbeddingRecord]): + """Lazy immutable descriptor view over one SQLite embedding run.""" + + _fastplms_immutable_sequence = True + + def __init__(self, path: Path, run_id: str, count: int) -> None: + self.path = path + self.run_id = run_id + self._count = count + + @staticmethod + def _row_query() -> str: + return ( + "SELECT r.position, r.record_id, r.sequence, t.dtype, t.shape_json, t.sha256 " + "FROM records r JOIN tensors t USING (run_id, position) " + "WHERE r.run_id = ?" + ) + + def __len__(self) -> int: + return self._count + + def __iter__(self) -> Iterator[EmbeddingRecord]: + with _connect_sqlite_read_only(self.path) as connection: + cursor = connection.execute(f"{self._row_query()} ORDER BY r.position", (self.run_id,)) + while rows := cursor.fetchmany(1_024): + for row in rows: + yield _sqlite_record_from_row(self.path, self.run_id, row) + + @overload + def __getitem__(self, index: int, /) -> EmbeddingRecord: ... + + @overload + def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ... + + def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]: + if isinstance(index, slice): + start, stop, step = index.indices(self._count) + return [self[position] for position in range(start, stop, step)] + position = index + self._count if index < 0 else index + if position < 0 or position >= self._count: + raise IndexError(index) + with _connect_sqlite_read_only(self.path) as connection: + row = connection.execute( + f"{self._row_query()} AND r.position = ?", + (self.run_id, position), + ).fetchone() + if row is None: + raise IndexError(index) + return _sqlite_record_from_row(self.path, self.run_id, row) + + +def load_sqlite_result( + path: str | Path, + *, + run_id: str | None = None, + positions: Iterable[int] | None = None, + record_ids: Iterable[str] | None = None, + sequences: Iterable[str] | None = None, +) -> EmbeddingResult: + """Load one SQLite run read-only, optionally in explicit selector order. + + Exactly one selector may be supplied. Repeated selectors are retained. An + ID or sequence selector that matches multiple stored rows returns those + rows in their original order for every occurrence of that selector. + """ + + path = Path(path).resolve() + supplied_selectors = sum( + selector is not None for selector in (positions, record_ids, sequences) + ) + if supplied_selectors > 1: + raise ValueError("Choose at most one of positions, record_ids, or sequences.") + normalized_positions = tuple(positions) if positions is not None else None + normalized_ids = tuple(record_ids) if record_ids is not None else None + normalized_sequences = tuple(sequences) if sequences is not None else None + if normalized_positions is not None and not all( + isinstance(position, int) and not isinstance(position, bool) and position >= 0 + for position in normalized_positions + ): + raise ValueError("positions must contain non-negative integers.") + for name, values in ( + ("record_ids", normalized_ids), + ("sequences", normalized_sequences), + ): + if values is not None and not all(isinstance(value, str) for value in values): + raise TypeError(f"{name} must contain strings.") + + with _connect_sqlite_read_only(path) as connection: + _validate_sqlite_result_schema(connection, path) + if run_id is None: + run_columns = { + str(info[1]) for info in connection.execute("PRAGMA table_info(runs)").fetchall() + } + if "published_order" in run_columns: + row = connection.execute( + "SELECT run_id, metadata_json FROM runs " + "WHERE published_order IS NOT NULL " + "ORDER BY published_order DESC, rowid DESC LIMIT 1" + ).fetchone() + else: + row = connection.execute( + "SELECT run_id, metadata_json FROM runs " + "ORDER BY created_at DESC, rowid DESC LIMIT 1" + ).fetchone() + else: + row = connection.execute( + "SELECT run_id, metadata_json FROM runs WHERE run_id = ?", (run_id,) + ).fetchone() + if row is None: + raise KeyError(f"No embedding run found in {path}.") + selected_run, metadata_json = row + metadata = json.loads(metadata_json) + if not isinstance(metadata, dict): + raise ValueError("SQLite run metadata must contain a JSON object.") + row_prefix = ( + "SELECT r.position, r.record_id, r.sequence, t.dtype, t.shape_json, t.sha256 " + "FROM records r JOIN tensors t USING (run_id, position) " + "WHERE r.run_id = ?" + ) + record_count, minimum_position, maximum_position = connection.execute( + "SELECT COUNT(*), MIN(position), MAX(position) FROM records WHERE run_id = ?", + (selected_run,), + ).fetchone() + (tensor_count,) = connection.execute( + "SELECT COUNT(*) FROM tensors WHERE run_id = ?", (selected_run,) + ).fetchone() + (joined_count,) = connection.execute( + "SELECT COUNT(*) FROM records r JOIN tensors t USING (run_id, position) " + "WHERE r.run_id = ?", + (selected_run,), + ).fetchone() + if ( + tensor_count != record_count + or joined_count != record_count + or (record_count and (minimum_position != 0 or maximum_position != record_count - 1)) + ): + raise ValueError("SQLite embedding run has inconsistent or non-contiguous records.") + metadata_count = metadata.get("record_count") + if ( + not isinstance(metadata_count, int) + or isinstance(metadata_count, bool) + or metadata_count != record_count + ): + raise ValueError("SQLite metadata record count does not match stored records.") + descriptor_cursor = connection.execute(f"{row_prefix} ORDER BY r.position", (selected_run,)) + while descriptor_rows := descriptor_cursor.fetchmany(1_024): + for descriptor_row in descriptor_rows: + _validate_sqlite_descriptor_row(descriptor_row) + if supplied_selectors == 0: + rows: list[tuple[Any, ...]] | None = None + else: + selector_values: tuple[Any, ...] + selector_column: str + if normalized_positions is not None: + selector_values = normalized_positions + selector_column = "r.position" + elif normalized_ids is not None: + selector_values = normalized_ids + selector_column = "r.record_id" + else: + if normalized_sequences is None: + raise RuntimeError("Filtered SQLite retrieval resolved no selector values.") + selector_values = normalized_sequences + selector_column = "r.sequence" + fetched: list[tuple[Any, ...]] = [] + unique_values = tuple(dict.fromkeys(selector_values)) + for start in range(0, len(unique_values), 900): + chunk = unique_values[start : start + 900] + placeholders = ",".join("?" for _ in chunk) + fetched.extend( + connection.execute( + f"{row_prefix} AND {selector_column} IN ({placeholders}) " + "ORDER BY r.position", + (selected_run, *chunk), + ).fetchall() + ) + value_index = ( + 0 if normalized_positions is not None else (1 if normalized_ids is not None else 2) + ) + matched: dict[Any, list[tuple[Any, ...]]] = {} + for fetched_row in sorted(fetched, key=lambda item: int(item[0])): + matched.setdefault(fetched_row[value_index], []).append(fetched_row) + missing = [value for value in selector_values if value not in matched] + if missing: + raise KeyError(f"SQLite embedding selectors were not found: {missing!r}.") + rows = [ + fetched_row for value in selector_values for fetched_row in matched.get(value, ()) + ] + + if rows is None: + return EmbeddingResult( + _SQLiteRecordSequence(path, selected_run, int(record_count)), + metadata, + ) + records = [_sqlite_record_from_row(path, selected_run, selected_row) for selected_row in rows] + if supplied_selectors: + metadata = dict(metadata) + metadata["selection"] = { + "kind": ( + "positions" + if normalized_positions is not None + else "record_ids" + if normalized_ids is not None + else "sequences" + ), + "count": len(rows), + "duplicate_policy": "preserve-request-order", + } + return EmbeddingResult(records, metadata) + + +def load_legacy_pth(path: str | Path, *, allow_unsafe_pickle: bool = False) -> EmbeddingResult: + """Import a legacy mapping-only ``.pth`` file after explicit opt-in.""" + + if not allow_unsafe_pickle: + raise ValueError( + "Legacy .pth loading can execute pickle payloads. Pass " + "allow_unsafe_pickle=True only for a trusted file." + ) + payload = torch.load(Path(path), map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError("A legacy .pth embedding file must contain a mapping.") + records: list[EmbeddingRecord] = [] + for position, (sequence, X) in enumerate(payload.items()): + # X: (...) + if not isinstance(sequence, str) or not isinstance(X, Tensor): + raise ValueError("Legacy embedding mappings must use str keys and Tensor values.") + records.append(EmbeddingRecord(str(position), sequence, X.detach().cpu())) + return EmbeddingResult(records, {"format": "legacy-pth", "unsafe_pickle": True}) + + +_LEGACY_COMPACT_VERSION = 0x01 +_LEGACY_CODE_DTYPES: dict[int, tuple[np.dtype[Any], torch.dtype]] = { + 0: (np.dtype(np.float16), torch.float16), + # Legacy BF16 blobs stored FP16 payload bytes and converted back to BF16. + 1: (np.dtype(np.float16), torch.bfloat16), + 2: (np.dtype(np.float32), torch.float32), +} + + +def _decode_legacy_sqlite_blob( + data: bytes, + *, + fallback_shape: tuple[int, ...] | None, + allow_unsafe_pickle: bool, +) -> Tensor: + if len(data) >= 6 and data[0] == _LEGACY_COMPACT_VERSION: + dtype_code = int(data[1]) + if dtype_code not in _LEGACY_CODE_DTYPES: + raise ValueError(f"Unsupported legacy compact dtype code {dtype_code}.") + (ndim,) = struct.unpack_from(" 16 or len(data) < 6 + 4 * ndim: + raise ValueError("Malformed legacy compact embedding header.") + shape = tuple(int(value) for value in struct.unpack_from(f"<{ndim}i", data, 6)) + if any(size < 0 for size in shape): + raise ValueError("Malformed negative legacy embedding dimension.") + numpy_dtype, target_dtype = _LEGACY_CODE_DTYPES[dtype_code] + offset = 6 + 4 * ndim + expected = int(np.prod(shape, dtype=np.int64)) * numpy_dtype.itemsize + if len(data) - offset != expected: + raise ValueError("Legacy compact embedding payload length does not match shape.") + array = ( # shape + np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape) + ) + return torch.from_numpy(array).to(dtype=target_dtype) # shape + + try: + loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True) + except Exception as safe_error: + if allow_unsafe_pickle: + loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=False) + elif fallback_shape is None: + raise ValueError( + "Legacy embedding blob is neither compact nor safely loadable. " + "Provide fallback_shape for raw FP32 bytes, or set " + "allow_unsafe_pickle=True only for a trusted database." + ) from safe_error + else: + expected = int(np.prod(fallback_shape, dtype=np.int64)) * 4 + if len(data) != expected: + raise ValueError( + "Legacy raw FP32 payload length does not match fallback_shape." + ) from safe_error + array = np.frombuffer(data, dtype=np.float32).copy().reshape( # fallback_shape + fallback_shape + ) + return torch.from_numpy(array) # fallback_shape + if not isinstance(loaded, Tensor): + raise ValueError("Legacy serialized embedding payload must contain one tensor.") + return loaded.detach().cpu() # (...) + + +def convert_legacy_sqlite( + source: str | Path, + output: str | Path, + *, + fallback_shape: tuple[int, ...] | None = None, + allow_unsafe_pickle: bool = False, + metadata: dict[str, Any] | None = None, +) -> EmbeddingResult: + """Convert the v0 ``embeddings(sequence, embedding)`` database safely. + + The source is opened read-only. Compact blobs and ``weights_only`` Torch + tensors are accepted by default. Unsafe general pickle deserialization + remains an explicit opt-in. + """ + + source_path = Path(source) + output_path = Path(output) + if source_path.resolve() == output_path.resolve(): + raise ValueError("Legacy SQLite conversion requires a different output path.") + if fallback_shape is not None and ( + not fallback_shape or any(not isinstance(size, int) or size < 0 for size in fallback_shape) + ): + raise ValueError("fallback_shape must contain non-negative integer dimensions.") + with _connect_sqlite_read_only(source_path) as connection: + columns = { + str(row[1]) for row in connection.execute("PRAGMA table_info(embeddings)").fetchall() + } + if not {"sequence", "embedding"}.issubset(columns): + raise ValueError("Legacy SQLite database must contain embeddings(sequence, embedding).") + rows = connection.execute( + "SELECT sequence, embedding FROM embeddings ORDER BY rowid" + ).fetchall() + if not rows: + raise ValueError("Legacy SQLite database contains no embeddings.") + + records: list[EmbeddingRecord] = [] + content_digest = hashlib.sha256() + for position, (sequence, data) in enumerate(rows): + if not isinstance(sequence, str) or not sequence: + raise ValueError("Legacy embedding sequences must be non-empty strings.") + if not isinstance(data, bytes): + data = bytes(data) + tensor = _decode_legacy_sqlite_blob( + data, + fallback_shape=fallback_shape, + allow_unsafe_pickle=allow_unsafe_pickle, + ) + tensor_digest = tensor_sha256(tensor) + for value in (sequence.encode("utf-8"), tensor_digest.encode("ascii")): + content_digest.update(len(value).to_bytes(8, "big")) + content_digest.update(value) + records.append(EmbeddingRecord(str(position), sequence, tensor)) + + content_sha256 = content_digest.hexdigest() + run_fingerprint = hashlib.sha256( + f"fastplms-legacy-sqlite-v1:{content_sha256}".encode("ascii") + ).hexdigest() + converted_metadata: dict[str, Any] = { + "format_version": 1, + "run_fingerprint": run_fingerprint, + "source_format": "legacy-fastplms-sqlite-v0", + "source_content_sha256": content_sha256, + "unsafe_pickle": allow_unsafe_pickle, + "complete": True, + } + if metadata: + converted_metadata["conversion_metadata"] = _jsonable(metadata) + return save_sqlite_result( + EmbeddingResult(records, converted_metadata), + output_path, + ) + + +def save_result( + result: EmbeddingResult, + path: str | Path, + *, + format: str = "safetensors", + shard_size: int = DEFAULT_SHARD_SIZE, +) -> EmbeddingResult: + if format == "safetensors": + return save_safetensors_result(result, path, shard_size=shard_size) + if format == "sqlite": + return save_sqlite_result(result, path) + if format == "pth": + raise ValueError("Writing pickle-based .pth embeddings is not supported.") + raise ValueError("format must be 'safetensors' or 'sqlite'.") + + +def load_result(path: str | Path, *, format: str = "safetensors") -> EmbeddingResult: + if format == "safetensors": + return load_safetensors_result(path) + if format == "sqlite": + return load_sqlite_result(path) + raise ValueError("format must be 'safetensors' or 'sqlite'.") + + +__all__ = [ + "DEFAULT_SHARD_SIZE", + "SafetensorsStreamWriter", + "append_sqlite_records", + "convert_legacy_sqlite", + "garbage_collect_safetensors_generations", + "initialize_sqlite_run", + "load_legacy_pth", + "load_result", + "load_safetensors_result", + "load_sqlite_result", + "safetensors_result_exists", + "save_result", + "save_safetensors_result", + "save_sqlite_result", + "tensor_sha256", + "update_sqlite_run_metadata", +] diff --git a/src/fastplms/embeddings/types.py b/src/fastplms/embeddings/types.py new file mode 100644 index 0000000..5166ebd --- /dev/null +++ b/src/fastplms/embeddings/types.py @@ -0,0 +1,186 @@ +"""Public value types for dataset embedding.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, overload +from torch import Tensor + + +@dataclass(frozen=True, slots=True) +class EmbeddingInput: + """One named protein sequence supplied to :func:`embed_dataset`.""" + + id: str + sequence: str + + def __post_init__(self) -> None: + if not isinstance(self.id, str) or not self.id: + raise ValueError("EmbeddingInput.id must be a non-empty string.") + if not isinstance(self.sequence, str) or not self.sequence: + raise ValueError("EmbeddingInput.sequence must be a non-empty string.") + + +@dataclass(frozen=True, slots=True) +class LazyTensorReference: + """A tensor stored outside memory and loaded only when requested.""" + + source: str + key: str + dtype: str + shape: tuple[int, ...] + sha256: str + _loader: Callable[[], Tensor] = field(repr=False, compare=False) + + def load(self, *, verify: bool = True) -> Tensor: + """Load X and optionally verify its content digest.""" + + if not isinstance(verify, bool): + raise TypeError("verify must be a boolean.") + X = self._loader() # self.shape + if not isinstance(X, Tensor): + raise TypeError(f"Stored tensor loader for {self.key!r} must return a Tensor.") + if tuple(X.shape) != self.shape: + raise ValueError( + f"Stored tensor {self.key!r} has shape {tuple(X.shape)}, expected {self.shape}." + ) + dtype = str(X.dtype).removeprefix("torch.") + if dtype != self.dtype: + raise ValueError( + f"Stored tensor {self.key!r} has dtype {dtype!r}, expected {self.dtype!r}." + ) + if verify: + from .storage import tensor_sha256 + + digest = tensor_sha256(X) + if digest != self.sha256: + raise ValueError(f"Stored tensor {self.key!r} failed SHA-256 verification.") + return X # self.shape + + +TensorValue = Tensor | LazyTensorReference + + +@dataclass(frozen=True, slots=True) +class EmbeddingRecord: + """One ordered embedding result.""" + + id: str + sequence: str + tensor: TensorValue + + def __post_init__(self) -> None: + if not isinstance(self.id, str) or not self.id: + raise ValueError("EmbeddingRecord.id must be a non-empty string.") + if not isinstance(self.sequence, str) or not self.sequence: + raise ValueError("EmbeddingRecord.sequence must be a non-empty string.") + if not isinstance(self.tensor, (Tensor, LazyTensorReference)): + raise TypeError("EmbeddingRecord.tensor must be a Tensor or LazyTensorReference.") + + def load_tensor(self, *, verify: bool = True) -> Tensor: + """Return X regardless of whether this record is memory-backed or lazy.""" + + if not isinstance(verify, bool): + raise TypeError("verify must be a boolean.") + if isinstance(self.tensor, LazyTensorReference): + return self.tensor.load(verify=verify) # (...) + return self.tensor # (...) + + +class EmbeddingResult(Sequence[EmbeddingRecord]): + """Ordered embedding records and the metadata needed to reproduce them.""" + + def __init__( + self, + records: Sequence[EmbeddingRecord], + metadata: Mapping[str, Any] | None = None, + ) -> None: + self.records: Sequence[EmbeddingRecord] = ( + records if getattr(records, "_fastplms_immutable_sequence", False) else tuple(records) + ) + self.metadata = dict(metadata or {}) + + def __len__(self) -> int: + return len(self.records) + + def __iter__(self) -> Iterator[EmbeddingRecord]: + return iter(self.records) + + @overload + def __getitem__(self, index: int, /) -> EmbeddingRecord: ... + + @overload + def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ... + + def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]: + return self.records[index] + + def as_dict( + self, + *, + key: Literal["id", "sequence"] = "id", + duplicates: Literal["error", "first", "last"] = "error", + materialize: bool = True, + ) -> dict[str, TensorValue]: + """Convert records to a mapping under an explicit duplicate policy.""" + + if key not in {"id", "sequence"}: + raise ValueError("key must be 'id' or 'sequence'.") + if duplicates not in {"error", "first", "last"}: + raise ValueError("duplicates must be 'error', 'first', or 'last'.") + if not isinstance(materialize, bool): + raise TypeError("materialize must be a boolean.") + output: dict[str, TensorValue] = {} + for record in self.records: + record_key = getattr(record, key) + if record_key in output: + if duplicates == "error": + raise ValueError( + f"Duplicate {key} {record_key!r}; choose duplicates='first' " + "or duplicates='last' explicitly." + ) + if duplicates == "first": + continue + output[record_key] = record.load_tensor() if materialize else record.tensor + return output + + def materialize(self, *, verify: bool = True) -> EmbeddingResult: + """Return an equivalent result with every X loaded into CPU memory.""" + + if not isinstance(verify, bool): + raise TypeError("verify must be a boolean.") + return EmbeddingResult( + [ + EmbeddingRecord( + id=record.id, + sequence=record.sequence, + tensor=record.load_tensor(verify=verify), + ) + for record in self.records + ], + self.metadata, + ) + + +@dataclass(frozen=True, slots=True) +class EmbeddingBatch: + """Internal model-to-runner contract. + + ``X`` has shape ``(b, l, d)`` and ``residue_mask`` has shape ``(b, l)``. + ``attentions`` may contain layer/head attention matrices for ``parti``. + """ + + X: Tensor + residue_mask: Tensor + attentions: Tensor | tuple[Tensor, ...] | None = None + + +__all__ = [ + "EmbeddingBatch", + "EmbeddingInput", + "EmbeddingRecord", + "EmbeddingResult", + "LazyTensorReference", + "TensorValue", +] diff --git a/src/fastplms/models.toml b/src/fastplms/models.toml new file mode 100644 index 0000000..b39acbb --- /dev/null +++ b/src/fastplms/models.toml @@ -0,0 +1,1238 @@ +schema_version = 1 +legal_files = [ + "LICENSE=sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a", + "THIRD_PARTY_NOTICES.md=sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa", +] + +[[attention_kernels]] +implementation = "flash_attention_2" +repository = "kernels-community/flash-attn2" +revision = "db6b51744f0cd7061386442c09df890fc6d9f47e" +version = 2 +expected_variant = "flash_attn2" +dtypes = ["bfloat16"] + +[[attention_kernels]] +implementation = "flash_attention_3" +repository = "kernels-community/flash-attn3" +revision = "43f0bd269777115d94ff826e0d113ce9c1c9087b" +version = 1 +expected_variant = "flash_attn3" +dtypes = ["bfloat16"] + +[[runtime_assets]] +id = "esmfold2_ccd" +repository = "biohub/ESMFold2" +revision = "1ebf0e3481a5184eb6171d40615c79e384b48796" +path = "ccd.pkl" +sha256 = "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5" +size = 417306584 +consumer_family = "esmfold2" +trust_kind = "hash_pinned_pickle" +license = "MIT" +offline_behavior = "requires_cached_verified_file" + +[[upstreams]] +id = "ankh" +path = "vendor/upstream/ankh" +url = "https://github.com/agemagician/Ankh.git" +revision = "02b4e25ce5389b9e771c9df6e546c62af1216f8e" +license = "CC-BY-NC-SA-4.0" +license_files = ["LICENSE.md"] +license_digests = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"] +distribution_files = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"] + +[[upstreams]] +id = "biohub-esm" +path = "vendor/upstream/biohub-esm" +url = "https://github.com/Biohub/esm.git" +revision = "82ee35553d39169d678f784c8d3f8712ffd7d2c4" +license = "MIT" +license_files = ["LICENSE.md", "THIRD_PARTY_NOTICE.md"] +license_digests = [ + "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a", + "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23", +] +distribution_files = [ + "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a", + "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23", +] + +[[upstreams]] +id = "biohub-transformers" +path = "vendor/upstream/biohub-transformers" +url = "https://github.com/Biohub/transformers.git" +revision = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf" +license = "Apache-2.0" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"] +distribution_files = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"] + +[[upstreams]] +id = "boltz" +path = "vendor/upstream/boltz" +url = "https://github.com/jwohlwend/boltz.git" +revision = "b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc" +license = "MIT" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"] +distribution_files = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"] + +[[upstreams]] +id = "dplm" +path = "vendor/upstream/dplm" +url = "https://github.com/bytedance/dplm.git" +revision = "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d" +license = "Apache-2.0" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"] +distribution_files = [ + "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30", + "PROVENANCE.md=sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e", +] + +[[upstreams]] +id = "e1" +path = "vendor/upstream/e1" +url = "https://github.com/Profluent-AI/E1.git" +revision = "bfd2620a602248499f3d2583d85a7ecddf0b6e02" +license = "Apache-2.0 AND Profluent-E1-Agreement" +license_files = ["LICENSE", "ATTRIBUTION", "NOTICE"] +license_digests = [ + "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81", + "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367", + "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3", +] +distribution_files = [ + "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81", + "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367", + "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3", + "Apache-2.0.txt=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30", + "BSD-3-Clause.txt=sha256:36e1987f2f17db7f8ad36cd7a37dbb7aeaaf0ab68b97ab4b9d3556f3a7a76ae8", + "MODIFICATIONS.md=sha256:2506f47c0f5475af8e8ff2cff13eb8b79e8e25a08a054cdd617bf336536750ca", +] + +[[upstreams]] +id = "fair-esm" +path = "vendor/upstream/fair-esm" +url = "https://github.com/facebookresearch/esm.git" +revision = "2b369911bb5b4b0dda914521b9475cad1656b2ac" +license = "MIT" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93"] +distribution_files = [ + "LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93", + "PROVENANCE.md=sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7", +] + +[[upstreams]] +id = "openfold" +path = "vendor/upstream/openfold" +url = "https://github.com/aqlaboratory/openfold.git" +revision = "4b41059694619831a7db195b7e0988fc4ff3a307" +license = "Apache-2.0" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"] +distribution_files = [ + "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30", + "MODIFICATIONS.md=sha256:fd6f0aa1086a0c996cf967b326d18e965660cda0ad5c7f36a3474a8490720da3", + "PROVENANCE.md=sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081", +] + +[[upstreams]] +id = "protein-ttt" +path = "vendor/upstream/protein-ttt" +url = "https://github.com/anton-bushuiev/ProteinTTT.git" +revision = "fde2817cd84b936167cc76ccabf31e5c0fe49962" +license = "MIT" +license_files = ["LICENSE"] +license_digests = ["LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df"] +distribution_files = [ + "LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df", + "PROVENANCE.md=sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5", +] + +[families.esm2] +architecture = "ESM2" +upstreams = ["fair-esm"] +tokenizer_mode = "tokenizer" +public_input = "Amino-acid sequences tokenized to residue IDs" +extra = "core" +reference_container = "reference-esm2" +reference_adapter = "tests.parity.support.reference_adapters.esm2" +attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "sequence" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "esm2_hf_to_fastplms_v1" +conversion_provenance = "Input: the pinned official ESM2 state dictionary. Transformation: apply the deterministic esm2_hf_to_fastplms_v1 key map while preserving tensor values and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra FastPLMs checkpoint. Validation: release parity compares exact keys and values after the declared non-aliasing transform, tokenizer behavior, and inference. Limitation: any numerical rewrite requires a new transform identifier and exact conversion test." +representative = "esm2_8m" +documentation = "docs/models.md#esm2" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esm2", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.esm2.modeling_fastesm.FastEsmConfig", AutoModel = "fastplms.models.esm2.modeling_fastesm.FastEsmModel", AutoModelForMaskedLM = "fastplms.models.esm2.modeling_fastesm.FastEsmForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForTokenClassification" } + +[families.esm_plusplus] +architecture = "ESMC" +upstreams = ["biohub-esm", "biohub-transformers"] +tokenizer_mode = "tokenizer" +public_input = "Amino-acid sequences tokenized to residue IDs" +extra = "core" +reference_container = "reference-biohub-esm" +reference_adapter = "tests.parity.support.reference_adapters.esm_plusplus" +attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "static_parameters" +precisions = ["default"] +vram_tier = "sequence" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "esmc_to_fastplms_v1" +conversion_provenance = "Input: the pinned Biohub ESMC checkpoint. Transformation: apply the deterministic esmc_to_fastplms_v1 parameter map into the FastPLMs ESMC modules. Output: the pinned Synthyra ESMplusplus checkpoint. Validation: release parity compares keys, shapes, dtypes, values, aliases, and live inference. Limitation: runtime attention and precision selection are not serialized weight transforms." +representative = "esmc_small" +documentation = "docs/models.md#esm-and-esmc" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm_plusplus", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM" } + +[families.esm3] +architecture = "ESM3" +upstreams = ["biohub-esm", "biohub-transformers"] +tokenizer_mode = "tokenizer" +public_input = "Sequence, structure, and function tracks prepared through the multimodal helpers" +extra = "core" +reference_container = "reference-biohub-esm" +reference_adapter = "tests.parity.support.reference_adapters.esm3" +attention = ["eager", "sdpa", "flex_attention"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "large-sequence" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "esm3_to_fastplms_v1" +conversion_provenance = "Input: the pinned Biohub ESM3 checkpoint. Transformation: apply the deterministic esm3_to_fastplms_v1 parameter map for the supported sequence and multimodal modules and expand BF16 checkpoint tensors to FP32 storage. Output: the pinned Synthyra ESM3 checkpoint. Validation: release parity compares exact state identity after the declared map and live feature behavior. Limitation: unsupported upstream modalities may not be inferred from this record." +representative = "esm3_small" +documentation = "docs/models.md#esm3" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm3", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model" } + +[families.e1] +architecture = "E1" +upstreams = ["e1"] +tokenizer_mode = "sequence" +public_input = "Raw amino-acid sequences prepared by the native E1 adapter" +extra = "core" +reference_container = "reference-e1" +reference_adapter = "tests.parity.support.reference_adapters.e1" +attention = ["sdpa", "flex_attention"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "static_parameters" +precisions = ["default"] +vram_tier = "sequence" +checkpoint_license = "Profluent-E1-Agreement" +hub_license = "other" +hub_license_name = "Profluent-E1 Clickthrough License Agreement" +hub_license_link = "https://github.com/Profluent-AI/E1/blob/main/LICENSE" +weights_publication_allowed = true +state_transform = "e1_to_fastplms_v1" +conversion_provenance = "Input: the pinned Profluent-E1 checkpoint and tokenizer-free sequence contract. Transformation: apply e1_to_fastplms_v1 to the FastPLMs encoder and official task heads, storing floating tensors in BF16. Output: the pinned Synthyra Profluent-E1 checkpoint. Validation: release parity covers state identity after the declared cast, sequence and RAG preparation, aliases, and inference. Limitation: the FastPLMs scoring extension is not represented as an official E1 head." +representative = "e1_150m" +documentation = "docs/models.md#e1" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/e1", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.e1.modeling_e1.E1Config", AutoModel = "fastplms.models.e1.modeling_e1.E1Model", AutoModelForMaskedLM = "fastplms.models.e1.modeling_e1.E1ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.e1.modeling_e1.E1ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.e1.modeling_e1.E1ForTokenClassification" } + +[families.dplm] +architecture = "DPLM" +upstreams = ["dplm"] +tokenizer_mode = "tokenizer" +public_input = "Amino-acid sequences tokenized to masked or partially masked residue IDs" +extra = "core" +reference_container = "reference-dplm" +reference_adapter = "tests.parity.support.reference_adapters.dplm" +attention = ["eager", "sdpa", "flex_attention", "flash_attention_3"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "sequence" +checkpoint_license = "Apache-2.0" +hub_license = "apache-2.0" +weights_publication_allowed = true +state_transform = "dplm_to_fastplms_v1" +conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transformation: apply dplm_to_fastplms_v1, omitting the unused absolute-position table for rotary checkpoints and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra DPLM checkpoint. Validation: release parity compares exact state identity after the declared transform, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned source record; no broader rights are inferred." +representative = "dplm_150m" +documentation = "docs/models.md#dplm" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.dplm.modeling_dplm.DPLMConfig", AutoModel = "fastplms.models.dplm.modeling_dplm.DPLMModel", AutoModelForMaskedLM = "fastplms.models.dplm.modeling_dplm.DPLMForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm.modeling_dplm.DPLMForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm.modeling_dplm.DPLMForTokenClassification" } + +[families.dplm2] +architecture = "DPLM2" +upstreams = ["dplm"] +tokenizer_mode = "tokenizer" +public_input = "Tokenized amino-acid and structure tracks with explicit modality boundaries" +extra = "core" +reference_container = "reference-dplm" +reference_adapter = "tests.parity.support.reference_adapters.dplm2" +attention = ["sdpa"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "sequence" +checkpoint_license = "Apache-2.0" +hub_license = "apache-2.0" +weights_publication_allowed = true +state_transform = "dplm2_to_fastplms_v1" +conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transformation: apply dplm2_to_fastplms_v1, retaining the independent language-model head and trained encoder contact head while omitting the unused absolute-position table for rotary checkpoints. Output: the pinned Synthyra DPLM2 checkpoint. Validation: release parity compares exact keys and values after the declared omission, non-aliasing, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: no head exception is permitted by this source record, and redistribution remains subject to Apache-2.0." +representative = "dplm2_150m" +documentation = "docs/models.md#dplm2" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.dplm2.modeling_dplm2.DPLM2Config", AutoModel = "fastplms.models.dplm2.modeling_dplm2.DPLM2Model", AutoModelForMaskedLM = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForTokenClassification" } +tokenizer_class = "fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer" + +[families.ankh] +architecture = "ANKH" +upstreams = ["ankh"] +tokenizer_mode = "tokenizer" +public_input = "Amino-acid sequences tokenized for encoder or sequence-to-sequence use" +extra = "core" +reference_container = "reference-ankh" +reference_adapter = "tests.parity.support.reference_adapters.ankh" +attention = ["eager", "sdpa"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "static_parameters" +precisions = ["default"] +vram_tier = "large-sequence" +checkpoint_license = "CC-BY-NC-SA-4.0" +hub_license = "cc-by-nc-sa-4.0" +weights_publication_allowed = true +state_transform = "ankh_t5_to_fastplms_v1" +conversion_provenance = "Input: the pinned official ANKH T5 checkpoint. Transformation: apply ankh_t5_to_fastplms_v1 to the official encoder and sequence-to-sequence heads. Output: the pinned Synthyra ANKH checkpoint. Validation: release parity compares exact mapped state, tokenizer behavior, official heads, and inference. Limitation: the separately named FastPLMs masked-language-model extension is not an official ANKH head." +representative = "ankh_base" +documentation = "docs/models.md#ankh" +test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"] +requires_complete_weight_publication = false +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/ankh", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.ankh.modeling_ankh.FastAnkhConfig", AutoModel = "fastplms.models.ankh.modeling_ankh.FastAnkhModel", AutoModelForMaskedLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForMaskedLMExtension", AutoModelForSeq2SeqLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForConditionalGeneration", AutoModelForSequenceClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForTokenClassification" } + +[families.boltz2] +architecture = "Boltz2" +upstreams = ["boltz"] +tokenizer_mode = "structure" +public_input = "Raw amino-acid sequences through the convenience API, or prepared model features" +extra = "structure" +reference_container = "reference-boltz2" +reference_adapter = "tests.parity.support.reference_adapters.boltz" +attention = ["eager"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "structure" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "boltz2_inference_core_v1" +conversion_provenance = "Input: the pinned official Boltz2 checkpoint. Transformation: select and map the supported Boltz2 inference-core parameters with boltz2_inference_core_v1. Output: the pinned Synthyra Boltz2 checkpoint. Validation: release parity covers state identity for the declared subset, feature preparation, seeded inference, and structure outputs. Limitation: this record does not claim support for undeclared upstream training components." +representative = "boltz2" +documentation = "docs/models.md#boltz2" +test_tiers = ["structure", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "models/boltz"] +auto_map = { AutoConfig = "fastplms.models.boltz.modeling_boltz2.Boltz2Config", AutoModel = "fastplms.models.boltz.modeling_boltz2.Boltz2Model" } + +[families.esmfold] +architecture = "ESMFold" +upstreams = ["fair-esm", "openfold"] +tokenizer_mode = "structure" +public_input = "Raw amino-acid sequences through folding helpers, or prepared residue tensors" +extra = "structure" +reference_container = "reference-esmfold" +reference_adapter = "tests.parity.support.reference_adapters.esmfold" +attention = ["eager", "sdpa", "flex_attention"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["default"] +vram_tier = "structure" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "esmfold_meta_to_fastplms_v1" +conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus its pinned ESM2 backbone. Transformation: apply esmfold_meta_to_fastplms_v1 to map native ESM2 names into the structure-only FastPLMs backbone, retain folding tensors, omit five deterministically reconstructed geometry buffers, omit the folding-unused ESM2 masked-LM and contact-regression heads, and remove the obsolete random FastPLMs TTT head from earlier mirrors. Output: canonical FP32 FastPLMs ESMFold state with an explicit CUDA BF16-autocast execution path. Validation: release parity compares exact mapped keys, shapes, dtypes, values, aliases, semantic configuration, FP32 and BF16-compute seeded inference, and structure metrics with pLDDT normalized to (0, 1). Limitation: ESMFold TTT is rejected because the official checkpoint contains no trained masked-language-model head." +representative = "esmfold" +documentation = "docs/models.md#esmfold" +test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esmfold"] +auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding" } + +[families.esmfold2] +architecture = "ESMFold2" +upstreams = ["biohub-esm", "biohub-transformers", "protein-ttt"] +backbone_model = "esmc_6b" +tokenizer_mode = "structure" +public_input = "Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors" +extra = "structure" +reference_container = "reference-esmfold2" +reference_adapter = "tests.parity.support.reference_adapters.esmfold2" +attention = ["eager", "sdpa", "flex_attention"] +dtypes = ["float32", "bfloat16"] +bf16_execution = "fp32_parameters_autocast" +precisions = ["auto", "fp32", "bf16", "fp8"] +experimental_precisions = ["fp8"] +vram_tier = "structure-6b" +checkpoint_license = "MIT" +hub_license = "mit" +weights_publication_allowed = true +state_transform = "identity" +conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its separately pinned ESMC checkpoint. Transformation: apply identity to preserve the folding checkpoint exactly, load its parameters in FP32 for CUDA BF16-autocast execution, retain canonical BF16 ESMC weights, and optionally rebuild exactly 80 ESMC attention output projections as transient Transformer Engine linears. Output: the corresponding pinned Synthyra ESMFold2 checkpoint plus its declared ESMC precision policy. Validation: release parity covers exact canonical state, learned projection, prepared features, and seeded BF16 folding; experimental FP8 validation covers strict unavailable-device behavior, all four variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the four manifest-listed ESMFold2 variants are supported; FP8 is experimental, applies only to inference-time ESMC execution, and requires direct CUDA loading with Transformer Engine availability." +representative = "esmfold2" +documentation = "docs/esmfold2.md" +test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"] +runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model" } + +[[models]] +id = "esm2_8m" +family = "esm2" +size_category = "small" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:6975e86d1d8f27488bf2a676551feaa48cc19254c9d24b6acb09198122745609", tensors = "tests/goldens/esm2_8m.safetensors=sha256:b40217566c33c71988d28869de353be54a3b3ebfc21fdfd29056e88cf7e99f4c" } +fast_repo = "Synthyra/ESM2-8M" +fast_revision = "185ecbd45665d050a8dae326d91886d330c5f9d0" +fast_files = [ + "config.json=git-sha1:46d0a7b517f59123c6ebc6d1011585731cbab259", + "model.safetensors=sha256:c824e6ded5fb71c72bc5ac05300699947819023cb26cdaf6897665e6b2645e1b", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "facebook/esm2_t6_8M_UR50D" +official_revision = "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" +official_files = [ + "config.json=git-sha1:c2c6e65a87d9d20d47699ae236d605b80c741dd3", + "model.safetensors=sha256:24c5fa474c48f3b754b86efe752d5f189d2bcd88190fa2270fc92b2ef3034189", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esm2_t6_8M_UR50D.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt" +sha256 = "46f002a9870c9bdecd0ea887acb1f9a38a6b561e8f8bf8a6990b679b9d31b928" +size = 30099493 + +[[models.oracle_assets]] +role = "contact_regression" +path = "regression/esm2_t6_8M_UR50D-contact-regression.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t6_8M_UR50D-contact-regression.pt" +sha256 = "8f7a4557d57713b97ba0e484303007efb7230d25299c0ac47a0a1b12a87bbb9d" +size = 1511 + +[[models]] +id = "esm2_35m" +family = "esm2" +size_category = "small" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:e919d3ce6d20b6a942d27d92323814ae7594a0129dc9c4de27c5053e96675bcd", tensors = "tests/goldens/esm2_35m.safetensors=sha256:c9b8bb616cf884fb7744521a2fcc6eed23586342d11241e6c9ef16454ec31e17" } +fast_repo = "Synthyra/ESM2-35M" +fast_revision = "37ab9f56b41e365b3bd9e25d6fefe9150fd910f0" +fast_files = [ + "config.json=git-sha1:4d428c9934572f39e2a00db162249971f37c88e4", + "model.safetensors=sha256:21d95ab6bb9aa91bfec87eff11da61a657b732f2df279cbddbae6a7f1f0bba9c", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "facebook/esm2_t12_35M_UR50D" +official_revision = "6fbf070e65b0b7291e7bbcd451118c216cff79d8" +official_files = [ + "config.json=git-sha1:3f64131bb610ed1ce482c4b5421fc358c785278f", + "model.safetensors=sha256:e35647818e0e064351d4531ed480d225a002567b4b2b93ad3a9246d753150fc0", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esm2_t12_35M_UR50D.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt" +sha256 = "7f21e80e61d16a71735163ef555d3009afb0c98da74c48e29df08606973cc55e" +size = 134095705 + +[[models.oracle_assets]] +role = "contact_regression" +path = "regression/esm2_t12_35M_UR50D-contact-regression.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t12_35M_UR50D-contact-regression.pt" +sha256 = "16641e05d830d0ce863dd152dbb8c2f3ddfa3c3ec2a66080152c8abad01d8585" +size = 1959 + +[[models]] +id = "esm2_150m" +family = "esm2" +size_category = "medium" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:c04c93486024ba0fa1c81fbfbe92ee79d1d4c7f1cfcc2c9886728522f752feab", tensors = "tests/goldens/esm2_150m.safetensors=sha256:c03fe9916dba137b452a6bbe944c7dc414db4019a6f0921e87b92d4bb6a8a42f" } +fast_repo = "Synthyra/ESM2-150M" +fast_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a" +fast_files = [ + "config.json=git-sha1:efeae2af182b7d34dc35740a45f157661e7acdf4", + "model.safetensors=sha256:d1f7c60f98c31af328381519a750972b6a31b13b97aa7cca2e71b5ae1b3f8f53", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "facebook/esm2_t30_150M_UR50D" +official_revision = "a695f6045e2e32885fa60af20c13cb35398ce30c" +official_files = [ + "config.json=git-sha1:52e04179e6fbad6663a94ea5cc44f09d764c5cd4", + "model.safetensors=sha256:c3f1da8aea53bddd32c246c86168c23b9fd72341fb9db9a94436f855f5053566", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esm2_t30_150M_UR50D.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt" +sha256 = "881c7176cf198ef8dec26a3c375d40eb58d0c33df95c22562ca6cc6d3f812c62" +size = 592774773 + +[[models.oracle_assets]] +role = "contact_regression" +path = "regression/esm2_t30_150M_UR50D-contact-regression.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t30_150M_UR50D-contact-regression.pt" +sha256 = "6a604b96722ed052eef8a094ad90b275ba2e987d406315dbed0bdc6b3c4238a7" +size = 3431 + +[[models]] +id = "esm2_650m" +family = "esm2" +size_category = "large" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:f18332172fcb3abf5dd2485fd55f5b0d193ad3b93a44cc744e0d02817c927477", tensors = "tests/goldens/esm2_650m.safetensors=sha256:c3a66b75add03628e62e238cb63da6a9e4d321f8160e84bdf2a131c096977f86" } +fast_repo = "Synthyra/ESM2-650M" +fast_revision = "ca0718a5d52b80d5c60dd76860e55e061a95fb0a" +fast_files = [ + "config.json=git-sha1:88f6bd240680b29c3244df8292246048401f5caf", + "model.safetensors=sha256:a15142e94ecf36f0edde9b37796f591e609ebe1694ca411e93640f0ee384994a", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "facebook/esm2_t33_650M_UR50D" +official_revision = "08e4846e537177426273712802403f7ba8261b6c" +official_files = [ + "config.json=git-sha1:a956a25d277f30bd870d3760b9a116f19ead885e", + "model.safetensors=sha256:a08adabb949fa67ad3c14b509d04fd60368b35007b0095e3358f81200c4f4db0", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esm2_t33_650M_UR50D.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt" +sha256 = "ea9d0522b335a8778dea6535a65301f10208dece28cd5865482b0b1fc446168c" +size = 2604537549 + +[[models.oracle_assets]] +role = "contact_regression" +path = "regression/esm2_t33_650M_UR50D-contact-regression.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt" +sha256 = "8ffe6edbd4173dc8d45c2cd5cb27d43aad77ec26b4c768200c58ae1f96693575" +size = 3687 + +[[models]] +id = "esm2_3b" +family = "esm2" +size_category = "xlarge" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:5043b2333c57a34d54fac53916722d1acb4b6fd50395b9abafa805435b184a48", tensors = "tests/goldens/esm2_3b.safetensors=sha256:dfd5a8cb05d3e814a080185c4808c8e7ec2277f070f395562fcfbe4376789e4e" } +notes = "The pinned default SDPA BF16 path uses a checkpoint-specific numeric calibration: relative L2 target/hard limit 0.06/0.07, relative Q99.9 0.15/0.18, first-percentile residue cosine 0.994/0.992, and pooled cosine 0.998/0.997. Exact state identity and the global logits-distribution contract remain required." +fast_repo = "Synthyra/ESM2-3B" +fast_revision = "ff89d0180f414ab9c677219a25da79bf09185456" +fast_files = [ + "config.json=git-sha1:94944ad6cabaa40a3ce1cbe6699cf464fdc1b2c0", + "model-00001-of-00003.safetensors=sha256:04b57854545c23779b562ee2ae22f10021ba0f4d586ba0ad482ee6eda187d562", + "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102", + "model-00003-of-00003.safetensors=sha256:a6b3a55b9e3b2e1778de34c665c3dd17bdfdf6da9d6d5c97730c57168709ccae", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "facebook/esm2_t36_3B_UR50D" +official_revision = "476b639933c8baad5ad09a60ac1a87f987b656fc" +official_files = [ + "config.json=git-sha1:69e7563923f87d2d7439bfb83e5a19b44b46d71b", + "pytorch_model-00001-of-00002.bin=sha256:0f971f11c449d21422aa982b791619c10351972992c735f4c3cd43fe09790412", + "pytorch_model-00002-of-00002.bin=sha256:7560b46fc383c691fb74b915b7d4bcef40d3df181447f16ba4b298845e308d0c", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esm2_t36_3B_UR50D.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt" +sha256 = "7de8b4082ba15891959ab368b77ce3886697af1efb16d3c9e9e7b0c5d3f07500" +size = 5678116398 + +[[models.oracle_assets]] +role = "contact_regression" +path = "regression/esm2_t36_3B_UR50D-contact-regression.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt" +sha256 = "4da500eab246481dc9c8c95bc7b1d02f2803d761c380b0e95186d4a07d0fc84e" +size = 6759 + +[[models]] +id = "esmc_small" +family = "esm_plusplus" +size_category = "medium" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:bb02652cf3cc484756b98ffa4ba55ed4c55870d2cea3342adb1d920ba9dfe10a", tensors = "tests/goldens/esmc_small.safetensors=sha256:03378d0f0fdd8161178ebb2c1f0da1b9776a726c8e8d3a10c009808a24de5654" } +notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration." +fast_repo = "Synthyra/ESMplusplus_small" +fast_revision = "46c5f7d562e47d4c14165b424c71ab7db008e6fb" +fast_files = [ + "config.json=git-sha1:df2f44187157b0cc371c48c887b77b1783679201", + "model.safetensors=sha256:d099223765bc4f1ae8d6c7e18561ce41df1d54073fdc5327ef0a229235a8f52a", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71", + "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756", +] +official_repo = "biohub/ESMC-300M" +official_revision = "a59b831785f907e96e6a246b1d142bfb76df31ee" +official_files = [ + "config.json=git-sha1:9a49eacf4e65c39f74381f0f0d240e3b89ef43d7", + "model.safetensors=sha256:0772d8fe64bb25e14fe6f23b80e3c9a7d215d0da3c6cba5bd356d7c0e0bb22cc", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61", +] + +[[models]] +id = "esmc_large" +family = "esm_plusplus" +size_category = "large" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:7a4d614f67b6fde417f3fd89f61e7ec442ae284769734b2b73e14945a816a8fd", tensors = "tests/goldens/esmc_large.safetensors=sha256:e13302df4cf7e8381552f1043a8fd0f31f3e0d50b2ab6009fb86b7940ae8ff79" } +notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration." +fast_repo = "Synthyra/ESMplusplus_large" +fast_revision = "f813401638b3fddab09748aec1ad2bf537aa4208" +fast_files = [ + "config.json=git-sha1:5736371902fe5d04e2859be30ac7dbd31b271b25", + "model.safetensors=sha256:4aff3f8c5de68c4d3e3824eb2c478e4a47355d3f849f3c745e5c8a5ee6cff851", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71", + "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756", +] +official_repo = "biohub/ESMC-600M" +official_revision = "a7e82012c83126b9eedb055fea9fa84b6c02f094" +official_files = [ + "config.json=git-sha1:71c8241dc28a5fb636248267a0927c0242b264c1", + "model.safetensors=sha256:e4232c30fd35fe2f57051ec88a703996ac94520580b4b836894207a3d45d9ff8", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61", +] + +[[models]] +id = "esmc_6b" +family = "esm_plusplus" +size_category = "xlarge" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:e229d938719782f280fab22dfc4c43e86109fdb0cc523631168c5a491afaace3", tensors = "tests/goldens/esmc_6b.safetensors=sha256:a948945e985c7deaca7be8b7eed09c0a9521a2af3f2b10fc2ec7a7d2a0f99ada" } +notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration." +fast_repo = "Synthyra/ESMplusplus_6B" +fast_revision = "0d579cce3b0f09efa6b3baddf6cc3fd8c9b616c8" +fast_files = [ + "config.json=git-sha1:e740cbcf211f2511c70c25a1ff6017a757ba7a69", + "model-00001-of-00006.safetensors=sha256:d30d18703453019f2d2d050866309888720c28eebc9a10307d1ddf3799e85a65", + "model-00002-of-00006.safetensors=sha256:b3d85378ab5023f4160a96e9c8cbd4cc6f78a771a83c856e88d48112f555bc13", + "model-00003-of-00006.safetensors=sha256:52595519b59349c5c6e373e6f5ca4a3d48ea6dde345f7e61e24766df5fab0e5b", + "model-00004-of-00006.safetensors=sha256:e46c6113c89c6f3e9b072c1bef02d763a625c37bcd8f9da2ed9363891c9a0758", + "model-00005-of-00006.safetensors=sha256:6d92cb2bf9791de644de2ae86f8523d802ac3b4aaabfff0716ab6c2b97f6fb14", + "model-00006-of-00006.safetensors=sha256:5fc1a8632490bb34162823c35d0d591337b9e4195b22cc0560741397a6e9d0b3", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71", + "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756", +] +official_repo = "biohub/ESMC-6B" +official_revision = "45b0fa5d7fb06faefbd5e3b89bdcef35d564e79a" +official_files = [ + "config.json=git-sha1:19f5fb09e4f630fb5b748a497183c22a87ec5102", + "model-00001-of-00006.safetensors=sha256:bd90149ff223e6ac1a0cac6147a5ae0df20d3a21df4f65356a1f19cd14f4aa8a", + "model-00002-of-00006.safetensors=sha256:f75e2144d8269fe2eb4b3e0823fb089b94f176d8024153e85b8fb573a42294fa", + "model-00003-of-00006.safetensors=sha256:f699f01ecc9691d9c6470492765fe54b8b5d2e9f277c139e89427433ffdfe0b2", + "model-00004-of-00006.safetensors=sha256:46add1b7be098bbfdc3073884851ba3057f1b33ea23a158b650a37007dabd13d", + "model-00005-of-00006.safetensors=sha256:1e1cb62f060a34e18f54a31a76683ef888b8cec59e73315f5b31d25d45a1f88c", + "model-00006-of-00006.safetensors=sha256:56c73e13ae96e777ce65eee99364056069ef93b646470f352f83c5f1037b1b18", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61", +] + +[[models]] +id = "esm3_small" +family = "esm3" +tokenizer_source = "esmc_small" +size_category = "large" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:5470e8596cbba0e2882647eccbc53c36d8b48b0f3947d1fe0bcea68da1078c32", tensors = "tests/goldens/esm3_small.safetensors=sha256:d957922f810c9ab4c557d80d5aaaf6a3aab79a5a45e4638012a634a4134803b1" } +fast_repo = "Synthyra/ESM3_small" +fast_revision = "7ddb5a740f9e5f93933eb6410c0ee8684bc63ec1" +fast_files = [ + "config.json=git-sha1:60526e2fdd8af9d4fba17f323775458ef5a1a1f9", + "model-00001-of-00002.safetensors=sha256:a4c9b736c4c59d51180e966005a164859b47d5cd36e1f8ecdea619fbd34a0e92", + "model-00002-of-00002.safetensors=sha256:bea60e4e91b03bb00b6cedd29b07606b8543f0869fb74454af7b26e216d80d2b", + "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71", + "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756", +] +official_repo = "biohub/esm3-sm-open-v1" +official_revision = "47f0545b2b6daf26a93439a3cd610f4f7f3d5478" +official_files = [ + "config.json=git-sha1:0967ef424bce6791893e9a57bb952f80fd536e93", + "data/weights/esm3_function_decoder_v0.pth=sha256:f76d074efcaccfe21365a4fa96f212dadd66798e1e49d809ab7ffbe025d227c9", + "data/weights/esm3_sm_open_v1.pth=sha256:5ead5a135c658068db6a4f1b933e72d6110992c4668822e1c0e2dcc53e38acd9", + "data/weights/esm3_structure_decoder_v0.pth=sha256:3b726258a44274792b40ce7ea307e10c5da09936368a4ffa2970264d909da65b", + "data/weights/esm3_structure_encoder_v0.pth=sha256:467acbaee703ba3ccde6e75241a912a316952e5ff071355f85c1d33c68704f40", +] + +[[models]] +id = "e1_150m" +family = "e1" +size_category = "small" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:701a64a6ab1a2fec5a427555b6af96232526c15cb3d5b4dc7fb253ac8f20b922", tensors = "tests/goldens/e1_150m.safetensors=sha256:6558bc8f1a7b20629eaaaa6f72601d0c2cdb859a5dc13595549b1773b6e2de41" } +fast_repo = "Synthyra/Profluent-E1-150M" +fast_revision = "7c5f3bbf697226a2e0900db7a100f9201774a907" +fast_files = [ + "config.json=git-sha1:562ef21e722ca708064fc3d54d25b731d4ac8171", + "model.safetensors=sha256:d779ed3a4e23799aafc932dc09c9963428d10aa7075999b5f8851b39c76b67f6", +] +official_repo = "Profluent-Bio/E1-150m" +official_revision = "c4dbfe827e4aa6ed7f95eaef50dc1e084f4d77dc" +official_files = [ + "config.json=git-sha1:485e649199b46fe6ee7456bebf7aae9b3d4baeab", + "model.safetensors=sha256:ba2656339005e6598642836acfdafde480fecc7e145ce0058eb54adf572c3484", +] + +[[models]] +id = "e1_300m" +family = "e1" +size_category = "medium" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d3478f3f5957a0e0377864074dde0107de890019f96cb63548ee17ffb8f3ec3a", tensors = "tests/goldens/e1_300m.safetensors=sha256:92778b9ef95a803ddc84b3e3ca764c59e045872a94bcff0eb0cd47647732c188" } +fast_repo = "Synthyra/Profluent-E1-300M" +fast_revision = "5ef52c0ad2ae2578f40622696b763523810e8e26" +fast_files = [ + "config.json=git-sha1:f5c91498b76a3e3282a0d716d87738abb1a1b6c1", + "model.safetensors=sha256:9271c4176a8a2e0905a0bb769570ba1c2978fb999a87da92db4cf2b041224864", +] +official_repo = "Profluent-Bio/E1-300m" +official_revision = "5a2871c587eadbcc9237bc686ea45e5b4d28dfb3" +official_files = [ + "config.json=git-sha1:918cb09e6e96d4719ed85951f38c693360f9cdb8", + "model.safetensors=sha256:31e09a2542f45b04e6ce4adafb3b657f21e2d56d12bf68fd2266b1576a80bc9b", +] + +[[models]] +id = "e1_600m" +family = "e1" +size_category = "large" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:914be191c28141c1f84535cdb69ead0588a2057bb19d46c5bc7f3891a3d6739e", tensors = "tests/goldens/e1_600m.safetensors=sha256:22ed8417a4651ded255099f6d15c63c2c40552e700d2b0470d1adfde3a39c513" } +fast_repo = "Synthyra/Profluent-E1-600M" +fast_revision = "6c8bf0ec83b0e0178677c528b101efffd0677742" +fast_files = [ + "config.json=git-sha1:1d35c0b35b473259875fd29ee80167487a0d6afe", + "model.safetensors=sha256:793483b1b3411eab73fe5214b94d1424ca0545992dfac6889cfc0186af472363", +] +official_repo = "Profluent-Bio/E1-600m" +official_revision = "52d959fb87a609d15cf223a485127b29ed5c382a" +official_files = [ + "config.json=git-sha1:8a0a439ed4201462bc01189c9f8b43523b257b5c", + "model.safetensors=sha256:cfc108d4b98baaa62932331b40be265eae39dc382595bc3cde4a5ab55db1bf7a", +] + +[[models]] +id = "dplm_150m" +family = "dplm" +size_category = "small" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:3228551fe3bed951db9ec97347143ec4462ce7c221ac240b7ce7730948c1dc1f", tensors = "tests/goldens/dplm_150m.safetensors=sha256:392992235195beed97ab8359b90a2e11e52f4326606f99a471447bed81d146bd" } +fast_repo = "Synthyra/DPLM-150M" +fast_revision = "90ba742754151a774f3b7ed580170d0a76b3e69d" +fast_files = [ + "config.json=git-sha1:117ac2c1222152ef378abaad1f605e18c4a18ab0", + "model.safetensors=sha256:8bac5ac767ceb8deb511b272d32883f811768d56cb25e920cea94ba9b979ca14", + "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15", + "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "airkingbd/dplm_150m" +official_revision = "49b7125a5d28c6418fcc2f3c4fe799352ac1488b" +official_files = [ + "config.json=git-sha1:4910cb02f1840e9ac577026f601829604af58c74", + "pytorch_model.bin=sha256:ea4eaa99536b60ed76f945f71a1a5e604f08447ec3def5104a93ca6001a59961", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models]] +id = "dplm_650m" +family = "dplm" +size_category = "large" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:bf58d0ce73aaac7e6fb1923ef3d9adad67122df2a3dd414c3229488ef9587a6d", tensors = "tests/goldens/dplm_650m.safetensors=sha256:073f0a6abea7e48f28c2d921ff8329a28e22627f01979277cb324908a01b3378" } +fast_repo = "Synthyra/DPLM-650M" +fast_revision = "05dc16d97c5c028aed924c9ed681cee4ab609760" +fast_files = [ + "config.json=git-sha1:3537150eb87b213a676d5840548625e220b60e8b", + "model.safetensors=sha256:e27a47b8ec1c078b3fccb36542210e20f0380c88828db2ca9acf3d8a25048bd8", + "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15", + "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "airkingbd/dplm_650m" +official_revision = "7a7e651baa667d094aba05e9dc1cf52a3332110a" +official_files = [ + "config.json=git-sha1:625574d625a4178ca6966e9545fee56026c0b634", + "pytorch_model.bin=sha256:db4e54343a89e7600f41c3aacbc593db1b0caee82ec28cab25ff2ae090eba39c", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models]] +id = "dplm_3b" +family = "dplm" +size_category = "xlarge" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:a5b6df8b9c7b371976892ec1d6c45581a32ad3a6325c6c0a0b3267012848c8ed", tensors = "tests/goldens/dplm_3b.safetensors=sha256:75b0a0854fc391133920b0feaaeb8f69ab7568a88b3759627aca1556c4338c1e" } +fast_repo = "Synthyra/DPLM-3B" +fast_revision = "7d764dd3d70ecf1ac0e64693de64a0064aacac65" +fast_files = [ + "config.json=git-sha1:7f5baf9426be06760c86882948b0f4af2e681e22", + "model-00001-of-00003.safetensors=sha256:37b54855d087ef3e7d883464ae9d5ea3127ec15a16c6323d91ad16a6b98305c9", + "model-00002-of-00003.safetensors=sha256:042604fefb05ea8c360a48416ce7ba662a4f90b176b4baf646c5c1814c35e6e8", + "model-00003-of-00003.safetensors=sha256:b9ae04012665163c3fc9781dd04fcd69738ac20c07e615e98fc4483fd2c4de45", + "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15", + "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] +official_repo = "airkingbd/dplm_3b" +official_revision = "53849d4a7fe944ae0b9cf2bbc0d2cc0054795b51" +official_files = [ + "config.json=git-sha1:f6206456e8c2f22ebe1d37fce3b5d50fd8073e68", + "pytorch_model-00001-of-00004.bin=sha256:0bcb86a115fe744ed686756db143f78851304e855e2f83cec58681c6080ced5f", + "pytorch_model-00002-of-00004.bin=sha256:daf3324f3be949e7dd1c3c84b28da7fec5151b1890cb0904e73427266856a06f", + "pytorch_model-00003-of-00004.bin=sha256:dbbeb7924a21059854f994931e23590b054aa000b10370a71c052c4aa36e9246", + "pytorch_model-00004-of-00004.bin=sha256:21c01740d091487db43446489d8a893dea1fcc6f2e1c1991ece13945f7ab4e07", + "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2", +] + +[[models]] +id = "dplm2_150m" +family = "dplm2" +size_category = "small" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:d269de779ea1503de72c77e7b2e6224afc9797bd945b40c571ff6faec782e4aa", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:17fc26600938ba5364b8ecb96750786d33e9f92bcd4ea4df3e12a389340748eb" } +artifact_source = "official" +canonical_state_sha256 = "82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001" +fast_repo = "Synthyra/DPLM2-150M" +fast_revision = "182745b8dc5661f898481a4fa60a7af9d53385c4" +fast_files = [ + "config.json=git-sha1:07905a2e4327d27d073cd0390f140aec2976125a", + "model.safetensors=sha256:0a7751b3113027b1d9c966a5bda2d6ab831855de7aaa047b911731665a7c3cc6", + "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b", + "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259", + "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2", +] +official_repo = "airkingbd/dplm2_150m" +official_revision = "3451d984d06497f835ed49634bd68c9dfb54d730" +official_files = [ + "config.json=git-sha1:20f1e55c64fdc4d1d30f7b1df64b6167fa23dc7c", + "pytorch_model.bin=sha256:be7f5cf9e421f59fcc437e63ce1c7391099a314a4e9a4f10b8688785fa581238", + "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37", +] + +[[models]] +id = "dplm2_650m" +family = "dplm2" +size_category = "large" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:d9a7548f9af657a72d441ca70f27379863724fcce8ddd3da4f672104b7bfb772", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c4e0e467c252c3ac813363d2d4b17a5e3bd99e75fad315e76d97689b4655ddac" } +artifact_source = "official" +canonical_state_sha256 = "cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4" +fast_repo = "Synthyra/DPLM2-650M" +fast_revision = "b9d8527a9473a54954fa2764f590b9ea1b435bb2" +fast_files = [ + "config.json=git-sha1:3e079579b214d48a09db57f2c60be6a1acea5baf", + "model.safetensors=sha256:92db08c7dbfd6c5e03fbfeaea3f36b09640ee794dcf5ea8d550527869a9f1d63", + "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b", + "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259", + "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2", +] +official_repo = "airkingbd/dplm2_650m" +official_revision = "0bc69b644976c6680ab7e26669854d1979e8876e" +official_files = [ + "config.json=git-sha1:4cce8d9dc212cdace0e20e89169790bcf199c158", + "pytorch_model.bin=sha256:8d6e08cc05e4858064a714013c74cc88c9caa2cc8b12c34605a3c24bcd877cfb", + "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37", +] + +[[models]] +id = "dplm2_3b" +family = "dplm2" +size_category = "xlarge" +# The pinned public sampler fails before generation because cls_token_id is None. +# State, tokenizer, and inference parity remain required for this checkpoint. +generation_contract = "official_unavailable" +official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:d6e0e02af53b13cb129192f06e264758aa21c9ebf4ee82411cf67037082d2329", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:838b11824d08f83bcb0c0b3268e579f3a87dbfb965370cfe5c3f8793b96b1964" } +notes = "The pinned official DPLM2-3B sampler fails before generation, so live generation equivalence cannot be established for this checkpoint. State, tokenizer, and inference parity remain required." +artifact_source = "official" +canonical_state_sha256 = "8c46ec09115dbe6cbfb91d94ab5e906369d57e27fe620a7741c6f8cb1b6ca890" +fast_repo = "Synthyra/DPLM2-3B" +fast_revision = "2a63babe8848abf5233d31bd55891dff8285fc50" +fast_files = [ + "config.json=git-sha1:5932b1d501fed28b84614e0d2c1ecc4e89f10d6e", + "model-00001-of-00003.safetensors=sha256:2ff393f6e8df1568ce075d50de69ff4e5e9d9886e5ec47e43d6c24df23459be3", + "model-00002-of-00003.safetensors=sha256:feb3cea852c2aa849cc30783a984a97f0d076990ade6606cda5e38bf2a5a9621", + "model-00003-of-00003.safetensors=sha256:9be363ddb98436af20901981ffbed2f1097377424987f6c1baad27d512b62e71", + "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b", + "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259", + "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2", +] +official_repo = "airkingbd/dplm2_3b" +official_revision = "9e77567926f98d1b997ea9131a8eeb035b9bf827" +official_files = [ + "config.json=git-sha1:22d51ce44cd6da8d819e0d00566987bb51d74753", + "pytorch_model-00001-of-00004.bin=sha256:d8c641eae6bf891581ec64d543169891b093e296f5679ac75c695bcf596b4211", + "pytorch_model-00002-of-00004.bin=sha256:6478ad86ec5fef3d1d26580493af2d8666009d3ff884f3f88548080c8bbf94b5", + "pytorch_model-00003-of-00004.bin=sha256:dde8f88dac4a6355488c2fb433ee12cd69f1169950566624fba43684d4d99dc6", + "pytorch_model-00004-of-00004.bin=sha256:17ec0145152bc10e4dd3b4c2edff337979f6b99ee7c7bfd6cf4e6dbd7262d079", + "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37", +] + +[[models]] +id = "ankh_base" +family = "ankh" +size_category = "medium" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:ebce8d7de821827ee995789c9b38d79252d3b2f76888130b0a8a7eedafaefe2b", tensors = "tests/goldens/ankh_base.safetensors=sha256:f0e78aa15d11749e0c64ff57f9e88c51cec6538a0adf8951f839df70cc708b65" } +notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head." +artifact_source = "official" +canonical_state_sha256 = "cdd8d30d88e5bf41f44e1eef4470d8e46607aba5f7c7c805b06c035b89c8c16f" +fast_repo = "Synthyra/ANKH_base" +fast_revision = "a3afa1db21c876dff57b3540fa7241e138fb1ed6" +fast_files = [ + "config.json=git-sha1:d1b81bb97129bc75dea04daef1ea2af373018e6b", + "model-00001-of-00001.safetensors=sha256:c943d25cacdafd2c8e3518a74450b5f90f715becf30ceb24c327f1c5a0bc8b5d", + "model.safetensors.index.json=git-sha1:ca251ab9277c06081b33e027f68f5bdc0808b443", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0", +] +official_repo = "ElnaggarLab/ankh-base" +official_revision = "d99cb6b966530dfc2ae96bc69d9255c2a07308b0" +official_files = [ + "config.json=git-sha1:abd44a36b5469e9a7cb019e4059b5ac1392d8422", + "pytorch_model.bin=sha256:9b2a886374f0ff4a893f4e7a989deed76bb2458c8998bd5202ea8e97d92ddcc3", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0", +] + +[[models]] +id = "ankh_large" +family = "ankh" +size_category = "large" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:59492518b021de5cfaea87d672c9448c8558e99a3443ba2cc7ab544963196ecb", tensors = "tests/goldens/ankh_large.safetensors=sha256:3fb8d3ac27716d15a9ea92aeef6acf2b977bcc887d9b535000539e523673459b" } +notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head." +artifact_source = "official" +canonical_state_sha256 = "e498a2e9aea76ef784cbe3e596c6b3f5e9a40e209ad837f7e3207099e4d74483" +fast_repo = "Synthyra/ANKH_large" +fast_revision = "92d2403bbe3c32acaa944fbb8dc2beb5f571f008" +fast_files = [ + "config.json=git-sha1:46eef0fff286107820f8ffc523127fa981435aeb", + "model-00001-of-00002.safetensors=sha256:79301f0b6a4fcbfd3b8bd10ca892846d79b1aad6ad06976da7380249e36f5158", + "model-00002-of-00002.safetensors=sha256:20062a5049fcde509030024527665a75062a95d64966558dfcaa9245b441cbec", + "model.safetensors.index.json=git-sha1:6b707ca3ce7255d241a52feeca68c0cbbe2a383f", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09", +] +official_repo = "ElnaggarLab/ankh-large" +official_revision = "74b371dbfa3ee0a05d32ae74df0c2e0b82d6b9a6" +official_files = [ + "config.json=git-sha1:1abf33e52ee3d6be67d780ec57d32ac2b27b5306", + "pytorch_model.bin=sha256:517b6e8b279dedcb477af240b35c46bd6eb3307723eb281e60d4b2c8a87b889b", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09", +] + +[[models]] +id = "ankh2_large" +family = "ankh" +size_category = "large" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:e8df38994ca1a1e0c598ace34a0b257b264937e4fdbb01bc41544985116b02a4", tensors = "tests/goldens/ankh2_large.safetensors=sha256:25fe1569f55c635fab8fa49c1d62a889a35a2a738bad921f5764a85b58fd4b5d" } +notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head." +artifact_source = "official" +canonical_state_sha256 = "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3" +fast_repo = "Synthyra/ANKH2_large" +fast_revision = "729167c1980316ae61691338838447491926033f" +fast_files = [ + "config.json=git-sha1:dd5d59e6b74bc8afa9fd4a5bda13526c235dabb8", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "model-00001-of-00002.safetensors=sha256:7c0c297f60bcf81c732cdfeae6e99e140272807eb52afd70356fc6fdfa94e5a8", + "model-00002-of-00002.safetensors=sha256:f3d425d3e8741ccbdd925446559a9bf317c2c91e328f2eee44924423b56e3a3d", + "model.safetensors.index.json=git-sha1:6b707ca3ce7255d241a52feeca68c0cbbe2a383f", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25", +] +official_repo = "ElnaggarLab/ankh2-ext2" +official_revision = "aa9b9fa72288c47d9f618ce80c011e24b54e17a8" +official_files = [ + "config.json=git-sha1:9286bed4ecbc4f7113024919d16ec9719b0c0748", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "pytorch_model.bin=sha256:2df583f28f111276ee22a7b76007f4297e9a69766d60bccd9c8d7169c06ac606", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25", +] + +[[models]] +id = "ankh3_large" +family = "ankh" +size_category = "large" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:2e5bb05b3baa5baa78f61fef7d2a2c669b0da5dbfaf6b50b12abd3e17253a961", tensors = "tests/goldens/ankh3_large.safetensors=sha256:e5c494ac418e0a2fe7bdad1376676d48960d58ec9e044d19bfffccb8c3288513" } +notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head." +artifact_source = "official" +canonical_state_sha256 = "60acb7ef86e85dc0c51fc1edf4c8e69a0480049723b6b2c95e6e9faa720c112a" +fast_repo = "Synthyra/ANKH3_large" +fast_revision = "c6d16ca2a1b3b27a27bcf3875e816a059029d264" +fast_files = [ + "config.json=git-sha1:813ffc6c319549a2c1f3503e1309c36110202d65", + "generation_config.json=git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2", + "model-00001-of-00002.safetensors=sha256:7f1f5c5dcff4b6bc6b8464fe9a7eebdd99b0789ee8da895f42a41bdb04191654", + "model-00002-of-00002.safetensors=sha256:c1a67cef9b76202362ff00c9d2b2dc4b5fc7acd1f22d30c8b3f2e3d2597d0f22", + "model.safetensors.index.json=git-sha1:a20cfc1f8517ef47d12d08604dc93c064f1e6736", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] +official_repo = "ElnaggarLab/ankh3-large" +official_revision = "2be091622e8a393f0ef21735070084123c874b6e" +official_files = [ + "config.json=git-sha1:f5278f77d158cdd8a173df888e3ed365e84a80a3", + "generation_config.json=git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2", + "pytorch_model.bin=sha256:26321a345e07a25b21c6c41b651c4db91b420892e52c0dcbc55bd7a8f510f95b", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] + +[[models]] +id = "ankh3_xl" +family = "ankh" +size_category = "xlarge" +generation_contract = "required" +official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:66bb12e033e4163be225d636108a479393228a4f5061015c8af114e766c3c486", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:72d34567d0228cb6f1ee701c578ed4039fead4346e3f161a52e0e74df28dc8ae" } +notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head. The official PyTorch shard index is deliberately excluded: the builder verifies every declared source shard directly and writes a new canonical safetensors index." +artifact_source = "official" +canonical_state_sha256 = "dd2188e0d2ca65232135714eef6de394239734d843ddae4928c7398685d858e7" +fast_repo = "Synthyra/ANKH3_xl" +fast_revision = "d2856892e7535af2f55c2c4de043b1b272a29ed8" +fast_files = [ + "config.json=git-sha1:791460a5c0d6c03bebbac1d7eec7e35805eaf7b7", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "model-00001-of-00005.safetensors=sha256:f6b841f6b800e436b08e362d04f8442fd044839b1e32ba6ac01ecb30a9d2bae5", + "model-00002-of-00005.safetensors=sha256:ea556d511d4747ada49b9d0c24ef503774410093e53c0d003ffb9407efc2be31", + "model-00003-of-00005.safetensors=sha256:125884f5dcb5b44435e3b76330582f31f74547b67c9b6cadf9a4d7cf38748eb7", + "model-00004-of-00005.safetensors=sha256:e776ef6c5d6a50d4b3fcf0bbb2431de7ce813eddd2903290e54809c801ddb241", + "model-00005-of-00005.safetensors=sha256:58ee3b065cfcccd179fdbecef9827dfb91feb33e0d8385b692e6309d56ec530e", + "model.safetensors.index.json=git-sha1:74d149f64234c3f43eb85971f40b4a1c6d05a407", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] +official_repo = "ElnaggarLab/ankh3-xl" +official_revision = "e00113df5c95ef71df7ea3f5a73d56bd00e473a4" +official_files = [ + "config.json=git-sha1:f8997040e8913df75fd2eebe71a2a8eb750ed0d0", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "pytorch_model-00001-of-00003.bin=sha256:2c9793cbee16697cd4149debe07d3a27143e280f6e970fa46042aae820fea981", + "pytorch_model-00002-of-00003.bin=sha256:31c5a860e414513c829ae52affb0970d7cef2c0545df2d6e1338b6806ab7174b", + "pytorch_model-00003-of-00003.bin=sha256:055a853bdd3623db95a637935aa299427e837cd8ea69fc04708b0262508bec75", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] + +[[models]] +id = "boltz2" +family = "boltz2" +size_category = "structure" +generation_contract = "not_applicable" +notes = "Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared inference-core state, feature preparation, and seeded execution remain tested, but native-environment BF16 end-to-end inference currently exceeds the fixed numerical-equivalence limits. FastPLMs therefore does not claim official inference equivalence for this checkpoint yet. Work on that numerical gap continues independently of the ESM++ and ESMFold2 release gates." +fast_repo = "Synthyra/Boltz2" +fast_revision = "3b148fc5efea109c065ec82ba8683d024de7134e" +fast_files = [ + "config.json=git-sha1:8682ccb12e177e73bc7a351ff7e3af484bfb6fac", + "model.safetensors=sha256:5c863fd200a1613a0e311071e2ad73ab350635e3fd336e6822cf45c52cb960e5", +] +official_repo = "boltz-community/boltz-2" +official_revision = "6fdef46d763fee7fbb83ca5501ccceff43b85607" +official_files = [ + "boltz2_conf.ckpt=sha256:090e82ac8c92f5e943fa1b39e7410a44027bea7243c0bbb3caa67a77fc1428e1", + "mols.tar=sha256:39e076d96dbec6b4e86982bbda16f3a53a2a60c9bdc17828d88f6f9a0c7d1fd7", +] + +[[models]] +id = "esmfold" +family = "esmfold" +size_category = "structure" +generation_contract = "not_applicable" +official_golden = { metadata = "tests/goldens/esmfold.json=sha256:380b9a96168410717d1f698feaabb826b1606444cbdeec86c2ea06d9ffe8f186", tensors = "tests/goldens/esmfold.safetensors=sha256:873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb" } +fast_repo = "Synthyra/FastESMFold" +fast_revision = "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823" +fast_files = [ + "config.json=git-sha1:18e0091dcbf6140bf68924d53c4c8917b9cd90b1", + "model-00001-of-00003.safetensors=sha256:36fab9e5c96d409b2a34a8b4f1273acac8c07f119c32c4fcfa7d47bbcd55b83c", + "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102", + "model-00003-of-00003.safetensors=sha256:2f1178cda0e6cff3b1e158e1acc59c83e3f4fc46e246388a5127bc56b8d9c4f2", + "special_tokens_map.json=git-sha1:53cd95604a28eb7e23da763c8da23f5006ab2179", + "tokenizer_config.json=git-sha1:10213f69b51b4b38876a29271b8f908e853a5800", + "vocab.txt=git-sha1:eee0a1fc93c82568f78f086550fbd7c591cf423a", +] +official_repo = "facebook/esmfold_v1" +official_revision = "75a3841ee059df2bf4d56688166c8fb459ddd97a" +official_files = [ + "config.json=git-sha1:1232d0aee4be551021d8e70e66ed2b062df917bf", + "pytorch_model.bin=sha256:2ee07356b125d1e3e57503c204111fd7323347fc4735d41d3caac57c2a78e116", + "special_tokens_map.json=git-sha1:121c8d54f8ea66cdf678f48b3cb37c05b4de5c0d", + "tokenizer_config.json=git-sha1:aad24fba9f1bad2d74ed79d414ddcd60e6b0f812", + "vocab.txt=git-sha1:9abfdf5472c0ed970648b683b86ab131256b3e42", +] + +[[models.oracle_assets]] +role = "weights" +path = "models/esmfold_3B_v1.pt" +url = "https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt" +sha256 = "e9a52579027e77d2d2e0a18218e755821f395730e86624cab9413dc117f5ca62" +size = 2771653574 + +[[models]] +id = "esmfold2" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = true +official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:f6e0ed1ec400b9a0fcc817db51774be968dc454b7a32645a07c479e42423ab20", tensors = "tests/goldens/esmfold2.safetensors=sha256:e4d6be4344c528e26b13f79a9303549e3de7e582da195c0078db3ce957fad420" } +fast_repo = "Synthyra/ESMFold2" +fast_revision = "cd5a0927cec585a778d983b99a8db23d2e9b281e" +fast_files = [ + "config.json=git-sha1:67e81ff571f393f0b630cd5a22398bd84979c030", + "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e", +] +official_repo = "biohub/ESMFold2" +official_revision = "1ebf0e3481a5184eb6171d40615c79e384b48796" +official_files = [ + "config.json=git-sha1:0300c084b990b2bd600efd9f538aa5de27109fea", + "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e", +] + +[[models]] +id = "esmfold2_fast" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = false +official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:091b004c0b330217b59c12acd6da3d6edaf91e48d95f6d5f40fc20399cef9478", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:6e2e1cd07401538b4d9df994f82abe7a5b38a01e8d1ee26681e1216d44a81990" } +fast_repo = "Synthyra/ESMFold2-Fast" +fast_revision = "407875bfcaa42552bfcb25acd67ee1888b790170" +fast_files = [ + "config.json=git-sha1:62ccca15a416a5dcbd02cd6ce161f432c7b4de58", + "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a", +] +official_repo = "biohub/ESMFold2-Fast" +official_revision = "b28d8ace5e05e61e5bec1e6820cfd3e221819d12" +official_files = [ + "config.json=git-sha1:c0ca526090fa7f8342ee4666d56e7fe3a4b8cbb2", + "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a", +] + +[[models]] +id = "esmfold2_experimental_cutoff2025" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = true +official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:cfd0e35b2bc468a0dc4f614d3acfa2fce004f96e9ae2433256ed095b829d55cc", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:9347466bbe803b6f5dc82e3356ca6cbbf2c2edd8765f9fd273385bda255019f6" } +fast_repo = "Synthyra/ESMFold2-Experimental-Cutoff2025" +fast_revision = "632ff4a9e68f1de78ee956a613267bdcdb5b354d" +fast_files = [ + "config.json=git-sha1:41119745d38bc5503a0212ad923e75211dec565f", + "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3", +] +official_repo = "biohub/ESMFold2-Experimental-Cutoff2025" +official_revision = "56f94f5c1069ecde17512c96928850518340d287" +official_files = [ + "config.json=git-sha1:79ed0dc0f867b8f09bfa004d6f77397c2ab9b38d", + "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3", +] +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" } + +[[models]] +id = "esmfold2_experimental_fast_cutoff2025" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = false +official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:1d0b2da4f1579243f37ae04bd4b834b747005cd8e8e7665e00d088123c43afd9", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:516e216d05d7e6bee59e77126d3e595e2bb7821929433f00c259c5d5241964bb" } +fast_repo = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025" +fast_revision = "8f022c2514a6c32692aaca078a8391d6bc6c4bac" +fast_files = [ + "config.json=git-sha1:b9d39e941050179ca51faaed58cbbd77778c1143", + "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", +] +official_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025" +official_revision = "74b88548bf19688b8727432db0d698cb2e1d8783" +official_files = [ + "config.json=git-sha1:0333d68ddb12ed2f066741dcb801142f466c0a2c", + "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", +] +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" } diff --git a/src/fastplms/models/__init__.py b/src/fastplms/models/__init__.py new file mode 100644 index 0000000..340b1a1 --- /dev/null +++ b/src/fastplms/models/__init__.py @@ -0,0 +1,11 @@ +"""Lazy model-family namespace for FastPLMs. + +Model classes are resolved through Transformers AutoClasses and the typed +registry. Importing this package therefore does not load checkpoints, create +tokenizers, compile kernels, or initialize an accelerator runtime. +""" + +from __future__ import annotations + + +__all__: tuple[str, ...] = () diff --git a/src/fastplms/models/_diffusion_generation.py b/src/fastplms/models/_diffusion_generation.py new file mode 100644 index 0000000..bcc6a93 --- /dev/null +++ b/src/fastplms/models/_diffusion_generation.py @@ -0,0 +1,530 @@ +"""Discrete diffusion generation shared by DPLM and DPLM2. + +The implementation keeps model-specific vocabulary rules at the public entry +points and shares only the categorical sampling and confidence-based remasking +mechanism. It has no dependency on the pinned upstream checkout. +""" + +from __future__ import annotations + +import math +import torch +from collections.abc import Iterable, Iterator, Mapping +from contextlib import contextmanager +from typing import Any, Protocol +from tqdm.auto import tqdm + + +class _MaskedLanguageModel(Protocol): + """Structural type used by the two generation entry points.""" + + config: Any + + def eval(self) -> Any: ... + + def modules(self) -> Iterable[torch.nn.Module]: ... + + def __call__(self, **kwargs: Any) -> Any: ... + + +_DPLM2_AA_BOUNDARY = 33 +_DPLM2_AA_BOS = 0 +_DPLM2_PAD = 1 +_DPLM2_AA_EOS = 2 +_DPLM2_AA_UNK = 3 +_DPLM2_AA_X = 24 +_DPLM2_AA_B = 25 +_DPLM2_AA_U = 26 +_DPLM2_AA_Z = 27 +_DPLM2_AA_O = 28 +_DPLM2_AA_MASK = 32 +_DPLM2_STRUCT_BOS = 33 +_DPLM2_STRUCT_EOS = 34 +_DPLM2_STRUCT_UNK = 35 + + +@contextmanager +def _temporary_eval(model: _MaskedLanguageModel) -> Iterator[None]: + """Run one generation forward in eval mode and restore every module flag.""" + training_states = tuple((module, module.training) for module in model.modules()) + model.eval() + try: + yield + finally: + for module, training in training_states: + module.training = training + + +def _resolve_max_iter(model: _MaskedLanguageModel, max_iter: int | None) -> int: + if max_iter is None: + max_iter = int(getattr(model.config, "num_diffusion_timesteps", 500)) + if isinstance(max_iter, bool) or not isinstance(max_iter, int) or max_iter <= 0: + raise ValueError("max_iter must be a positive integer") + return max_iter + + +def _validate_inputs( + input_tokens: torch.Tensor, + partial_masks: torch.Tensor | None, +) -> torch.Tensor | None: + # input_tokens: (b, l); partial_masks: (b, l) or None + if input_tokens.ndim != 2 or input_tokens.dtype not in { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + }: + raise ValueError("input_tokens must be an integer tensor with shape (b, l)") + if input_tokens.shape[-1] == 0: + raise ValueError("input_tokens must contain at least one token") + if partial_masks is None: + return None + if partial_masks.shape != input_tokens.shape or partial_masks.dtype != torch.bool: + raise ValueError("partial_masks must be boolean with the same shape as input_tokens") + if partial_masks.device != input_tokens.device: + raise ValueError("partial_masks and input_tokens must be on the same device") + return partial_masks # (b, l) + + +def _validate_temperature(temperature: float | None, *, default: float = 1.0) -> float: + if temperature is None: + temperature = default + temperature = float(temperature) + if not math.isfinite(temperature) or temperature < 0: + raise ValueError("temperature must be finite and non-negative") + return temperature + + +def _steps(max_iter: int, *, show_progress: bool) -> Iterable[int]: + return tqdm(range(max_iter), desc="Decoding", disable=not show_progress) + + +def _categorical( + logits: torch.Tensor, + *, + temperature: float, +) -> tuple[torch.Tensor, torch.Tensor]: + # logits: (..., c) + if temperature == 0: + scores, tokens = logits.log_softmax(dim=-1).max(dim=-1) # (...), (...) + return tokens, scores # (...), (...) + distribution = torch.distributions.Categorical(logits=logits.div(temperature)) + tokens = distribution.sample() # (...) + return tokens, distribution.log_prob(tokens) # (...), (...) + + +def _gumbel_argmax( + logits: torch.Tensor, + *, + noise_scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + # logits: (..., c) + uniform = torch.rand_like(logits) # (..., c) + noise = -torch.log(-torch.log(uniform + 1e-8) + 1e-8) # (..., c) + return _categorical(logits + noise_scale * noise, temperature=0.0) # (...), (...) + + +def _top_p(logits: torch.Tensor, probability: float = 0.95) -> torch.Tensor: + """Apply the nucleus filter used by the official DPLM samplers.""" + + # logits: (..., c) + original_shape = logits.shape + flattened = logits.reshape(-1, original_shape[-1]) # (n, c) + sorted_logits, sorted_indices = flattened.sort(dim=-1, descending=True) # (n, c), (n, c) + cumulative = sorted_logits.softmax(dim=-1).cumsum(dim=-1) # (n, c) + remove = cumulative > probability # (n, c) + remove[..., 1:] = remove[..., :-1].clone() + remove[..., 0] = False + sorted_logits.masked_fill_(remove, -math.inf) + return sorted_logits.gather(1, sorted_indices.argsort(dim=-1)).reshape( # (..., c) + original_shape + ) + + +def _lowest_confidence_mask( + scores: torch.Tensor, + eligible: torch.Tensor, + *, + rate: float, + stochastic_temperature: float | None = None, +) -> torch.Tensor: + # scores, eligible: (b, l) + selection_scores = scores.masked_fill(~eligible, 1000.0) # (b, l) + if stochastic_temperature is not None: + uniform = torch.rand_like(selection_scores) # (b, l) + noise = -torch.log(-torch.log(uniform + 1e-8) + 1e-8) # (b, l) + selection_scores = selection_scores + stochastic_temperature * rate * noise # (b, l) + cutoff_index = ( # (b, 1) + eligible.sum(dim=-1, keepdim=True).to(scores.dtype) * rate + ).long() + cutoff_index.clamp_(min=0, max=scores.shape[-1] - 1) + sorted_scores = selection_scores.sort(dim=-1).values # (b, l) + cutoff = sorted_scores.gather(dim=-1, index=cutoff_index) # (b, 1) + return (selection_scores < cutoff) & eligible # (b, l) + + +def _reparameterize( + output_tokens: torch.Tensor, + output_scores: torch.Tensor, + candidate_tokens: torch.Tensor, + candidate_scores: torch.Tensor, + active_mask: torch.Tensor, + eligible: torch.Tensor, + *, + mask_token_id: int, + rate: float, + stochastic_temperature: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # All tensor inputs except vocabulary-bearing candidate logits: (b, l). + remask = _lowest_confidence_mask( # (b, l) + candidate_scores, + eligible, + rate=rate, + stochastic_temperature=stochastic_temperature, + ) + output_tokens.masked_fill_(remask, mask_token_id) + output_scores.masked_fill_(remask, -math.inf) + accept = active_mask & eligible & ~remask # (b, l) + output_tokens.masked_scatter_(accept, candidate_tokens[accept]) + output_scores.masked_scatter_(accept, candidate_scores[accept]) + return remask, output_tokens, output_scores # (b, l), (b, l), (b, l) + + +def _logits(output: object) -> torch.Tensor: + value = output.get("logits") if isinstance(output, Mapping) else getattr(output, "logits", None) + if not torch.is_tensor(value): + raise RuntimeError("The masked-language model did not return logits") + return value # (b, l, c) + + +def _suppress_token_ids(logits: torch.Tensor, token_ids: Iterable[int]) -> None: + vocabulary_size = logits.shape[-1] + for token_id in token_ids: + if 0 <= token_id < vocabulary_size: + logits[..., token_id] = -math.inf + + +def _dplm_special_id( + model: _MaskedLanguageModel, + tokenizer: object | None, + name: str, + default: int, +) -> int: + value = getattr(model.config, name, None) + if value is None: + if tokenizer is None: + tokenizer = getattr(model, "tokenizer", None) + value = getattr(tokenizer, name, None) + return default if value is None else int(value) + + +def _dplm_resample_repeats( + model: _MaskedLanguageModel, + candidate_tokens: torch.Tensor, + candidate_scores: torch.Tensor, + *, + invalid_token_ids: tuple[int, ...], + mask_token_id: int, + ratio: float, +) -> None: + # candidate_tokens, candidate_scores: (b, l) + selected_rows: list[int] = [] + resample_tokens: list[torch.Tensor] = [] + resample_scores: list[torch.Tensor] = [] + resample_masks: list[torch.Tensor] = [] + for row_index, row in enumerate(candidate_tokens): + positions: dict[int, list[int]] = {} + for position, token in enumerate(row.tolist()): + positions.setdefault(int(token), []).append(position) + repeated = [indices for indices in positions.values() if len(indices) > row.numel() * ratio] + if not repeated: + continue + M = torch.zeros_like(row, dtype=torch.bool) # (l,) + for indices in repeated: + M[indices] = True + selected_rows.append(row_index) + resample_masks.append(M) # (l,) + resample_tokens.append(row.masked_fill(M, mask_token_id)) # (l,) + resample_scores.append(candidate_scores[row_index]) # (l,) + + if not selected_rows: + return + X = torch.stack(resample_tokens) # (r, l) + S = torch.stack(resample_scores) # (r, l) + M = torch.stack(resample_masks) # (r, l) + with _temporary_eval(model), torch.no_grad(): + logits = _logits(model(input_ids=X, return_dict=True)) # (r, l, c) + if logits.dtype != S.dtype: + logits = logits.to(S.dtype) # (r, l, c) + _suppress_token_ids(logits, invalid_token_ids) + logits = _top_p(logits) # (r, l, c) + sampled_tokens, sampled_scores = _gumbel_argmax(logits, noise_scale=1.0) # (r, l), (r, l) + X.masked_scatter_(M, sampled_tokens[M]) + S.masked_scatter_(M, sampled_scores[M]) + candidate_tokens[selected_rows] = X + candidate_scores[selected_rows] = S + + +def generate_dplm( + model: _MaskedLanguageModel, + input_tokens: torch.Tensor, + *, + tokenizer: object | None = None, + max_iter: int | None = None, + temperature: float | None = None, + partial_masks: torch.Tensor | None = None, + sampling_strategy: str = "gumbel_argmax", + disable_resample: bool = False, + resample_ratio: float = 0.25, + show_progress: bool = False, +) -> torch.Tensor: + """Generate DPLM sequences with the official iterative unmasking process. + + ``input_tokens`` is X with shape (b, l). ``partial_masks=True`` marks fixed + positions. The return value is the generated token tensor X with shape + (b, l), matching the official DPLM public API. + """ + + partial_masks = _validate_inputs(input_tokens, partial_masks) + max_iter = _resolve_max_iter(model, max_iter) + # Upstream treats ``None`` as the falsey, zero-temperature branch for + # vanilla categorical sampling. Gumbel and argmax strategies ignore it. + temperature = _validate_temperature(temperature, default=0.0) + if sampling_strategy not in {"vanilla", "argmax", "gumbel_argmax"}: + raise ValueError(f"Unsupported DPLM sampling strategy: {sampling_strategy!r}") + if not 0 < float(resample_ratio) <= 1: + raise ValueError("resample_ratio must be in (0, 1]") + + pad_id = _dplm_special_id(model, tokenizer, "pad_token_id", 1) + bos_id = _dplm_special_id(model, tokenizer, "bos_token_id", 0) + eos_id = _dplm_special_id(model, tokenizer, "eos_token_id", 2) + mask_id = _dplm_special_id(model, tokenizer, "mask_token_id", 32) + x_id = 24 + X = input_tokens.clone() # (b, l) + mutable = X.ne(pad_id) & X.ne(bos_id) & X.ne(eos_id) # (b, l) + if partial_masks is not None: + mutable &= ~partial_masks + X.masked_fill_(mutable, mask_id) + S = torch.zeros_like(X, dtype=torch.float32) # (b, l) + active = mutable.clone() # (b, l) + invalid_ids = (mask_id, x_id, pad_id, bos_id, eos_id) + for step in _steps(max_iter, show_progress=show_progress): + with _temporary_eval(model), torch.no_grad(): + logits = _logits(model(input_ids=X, return_dict=True)) # (b, l, c) + if logits.dtype != S.dtype: + logits = logits.to(S.dtype) # (b, l, c) + _suppress_token_ids(logits, invalid_ids) + if sampling_strategy == "vanilla": + candidate_tokens, candidate_scores = _categorical( # (b, l), (b, l) + logits, + temperature=temperature, + ) + elif sampling_strategy == "argmax": + candidate_scores, candidate_tokens = logits.max(dim=-1) # (b, l), (b, l) + else: + candidate_tokens, candidate_scores = _gumbel_argmax( # (b, l), (b, l) + logits, + noise_scale=1.0, + ) + if not disable_resample: + _dplm_resample_repeats( + model, + candidate_tokens, + candidate_scores, + invalid_token_ids=invalid_ids, + mask_token_id=mask_id, + ratio=float(resample_ratio), + ) + + eligible = X.ne(pad_id) & X.ne(bos_id) & X.ne(eos_id) # (b, l) + if partial_masks is not None: + eligible &= ~partial_masks + rate = 1.0 - (step + 1) / max_iter + active, X, S = _reparameterize( # (b, l), (b, l), (b, l) + X.clone(), + S.clone(), + candidate_tokens, + candidate_scores, + active, + eligible, + mask_token_id=mask_id, + rate=rate, + ) + return X # (b, l) + + +def _normalize_dplm2_special_ids(X: torch.Tensor, vocabulary_size: int) -> torch.Tensor: + # X: (b, l) + normalized = X.clone() # (b, l) + replacements = { + vocabulary_size: _DPLM2_AA_EOS, + vocabulary_size + 1: _DPLM2_AA_UNK, + vocabulary_size + 2: _DPLM2_AA_BOS, + vocabulary_size + 3: _DPLM2_AA_MASK, + } + for generic_id, native_id in replacements.items(): + normalized.masked_fill_(X.eq(generic_id), native_id) + return normalized # (b, l) + + +def _dplm2_types(X: torch.Tensor) -> torch.Tensor: + # X: (b, l) + valid = X.ne(_DPLM2_PAD) # (b, l) + types = ((X < _DPLM2_AA_BOUNDARY) & valid).to(torch.int64) # (b, l) + types.masked_fill_(~valid, 2) + return types # (b, l) + + +def _dplm2_mutable(X: torch.Tensor, partial_masks: torch.Tensor | None) -> torch.Tensor: + # X, partial_masks: (b, l) + mutable = ( # (b, l) + X.ne(_DPLM2_PAD) + & X.ne(_DPLM2_AA_BOS) + & X.ne(_DPLM2_AA_EOS) + & X.ne(_DPLM2_STRUCT_BOS) + & X.ne(_DPLM2_STRUCT_EOS) + ) + if partial_masks is not None: + mutable &= ~partial_masks + return mutable # (b, l) + + +def _dplm2_unmasking_temperature(strategy: str) -> float | None: + if strategy == "deterministic": + return None + if strategy.startswith("stochastic"): + suffix = strategy.removeprefix("stochastic") + value = 1.0 if not suffix else float(suffix) + if not math.isfinite(value) or value < 0: + raise ValueError("The stochastic unmasking temperature must be non-negative") + return value + raise ValueError(f"Unsupported DPLM2 unmasking strategy: {strategy!r}") + + +def _annealing_temperature(strategy: str, step: int, max_iter: int) -> float | None: + if not strategy.startswith("annealing"): + return None + try: + maximum, minimum = map(float, strategy.split("@", maxsplit=1)[1].split(":")) + except (IndexError, ValueError) as error: + raise ValueError("Annealing must use the form 'annealing@maximum:minimum'") from error + if not all(math.isfinite(value) and value >= 0 for value in (maximum, minimum)): + raise ValueError("Annealing temperatures must be finite and non-negative") + rate = 1.0 - step / max_iter + return minimum + (maximum - minimum) * rate + + +def generate_dplm2( + model: _MaskedLanguageModel, + input_tokens: torch.Tensor, + *, + max_iter: int | None = None, + temperature: float = 1.0, + partial_masks: torch.Tensor | None = None, + unmasking_strategy: str = "stochastic1.0", + sampling_strategy: str = "annealing@2.0:0.1", + show_progress: bool = False, +) -> dict[str, torch.Tensor]: + """Generate packed DPLM2 sequence and structure tracks. + + ``input_tokens`` is X with shape (b, l). A packed co-generation input has + two equal-length modality tracks. ``partial_masks=True`` marks fixed + positions. The output mapping matches the official DPLM2 public API. + """ + + partial_masks = _validate_inputs(input_tokens, partial_masks) + max_iter = _resolve_max_iter(model, max_iter) + temperature = _validate_temperature(temperature) + unmasking_temperature = _dplm2_unmasking_temperature(unmasking_strategy) + if sampling_strategy.startswith("annealing"): + _annealing_temperature(sampling_strategy, 0, max_iter) + elif sampling_strategy not in {"argmax", "gumbel_argmax"}: + raise ValueError(f"Unsupported DPLM2 sampling strategy: {sampling_strategy!r}") + vocabulary_size = int(model.config.vocab_size) + if vocabulary_size <= _DPLM2_STRUCT_UNK + 1: + raise ValueError("DPLM2 generation requires the multimodal vocabulary") + struct_mask_id = vocabulary_size - 1 + + X = _normalize_dplm2_special_ids(input_tokens, vocabulary_size) # (b, l) + if X.numel() and (X.min() < 0 or X.max() >= vocabulary_size): + raise ValueError("input_tokens contains an ID outside the DPLM2 vocabulary") + mutable = _dplm2_mutable(X, partial_masks) # (b, l) + types = _dplm2_types(X) # (b, l) + X.masked_fill_(mutable & types.eq(1), _DPLM2_AA_MASK) + X.masked_fill_(mutable & types.eq(0), struct_mask_id) + S = torch.zeros_like(X, dtype=torch.float32) # (b, l) + active = mutable.clone() # (b, l) + invalid_ids = ( + _DPLM2_AA_BOS, + _DPLM2_AA_EOS, + _DPLM2_AA_MASK, + _DPLM2_STRUCT_BOS, + _DPLM2_STRUCT_EOS, + struct_mask_id, + _DPLM2_PAD, + _DPLM2_AA_UNK, + _DPLM2_STRUCT_UNK, + _DPLM2_AA_X, + _DPLM2_AA_B, + _DPLM2_AA_U, + _DPLM2_AA_Z, + _DPLM2_AA_O, + ) + for step in _steps(max_iter, show_progress=show_progress): + eligible = _dplm2_mutable(X, partial_masks) # (b, l) + types = _dplm2_types(X) # (b, l) + with _temporary_eval(model), torch.no_grad(): + logits = _logits( # (b, l, c) + model(input_ids=X, return_dict=True) + ).log_softmax(dim=-1) + if logits.dtype != S.dtype: + logits = logits.to(S.dtype) # (b, l, c) + aa_rows, aa_columns = torch.where(types.eq(1) & eligible) # (n_aa,), (n_aa,) + struct_rows, struct_columns = torch.where( # (n_struct,), (n_struct,) + types.eq(0) & eligible + ) + logits[aa_rows, aa_columns, _DPLM2_AA_BOUNDARY:] = -math.inf + logits[struct_rows, struct_columns, :_DPLM2_AA_BOUNDARY] = -math.inf + _suppress_token_ids(logits, invalid_ids) + logits = _top_p(logits) # (b, l, c) + + if sampling_strategy == "argmax": + candidate_scores, candidate_tokens = logits.max(dim=-1) # (b, l), (b, l) + elif sampling_strategy == "gumbel_argmax": + candidate_tokens, candidate_scores = _gumbel_argmax( # (b, l), (b, l) + logits, + noise_scale=temperature, + ) + candidate_tokens.masked_scatter_(~eligible, X[~eligible]) + else: + annealed = _annealing_temperature(sampling_strategy, step, max_iter) + sample_temperature = temperature if annealed is None else annealed + candidate_tokens, candidate_scores = _categorical( # (b, l), (b, l) + logits, + temperature=sample_temperature, + ) + + rate = 1.0 - (step + 1) / max_iter + new_active = torch.zeros_like(active) # (b, l) + for modality, mask_id in ((1, _DPLM2_AA_MASK), (0, struct_mask_id)): + modality_positions = types.eq(modality) & eligible # (b, l) + if not bool(modality_positions.any()): + continue + modality_active, X, S = _reparameterize( # (b, l), (b, l), (b, l) + X, + S, + candidate_tokens, + candidate_scores, + active, + modality_positions, + mask_token_id=mask_id, + rate=rate, + stochastic_temperature=unmasking_temperature, + ) + new_active |= modality_active + active = new_active # (b, l) + return {"output_tokens": X} # (b, l) + + +__all__ = ["generate_dplm", "generate_dplm2"] diff --git a/src/fastplms/models/_esm_rotary.py b/src/fastplms/models/_esm_rotary.py new file mode 100644 index 0000000..57b3c5c --- /dev/null +++ b/src/fastplms/models/_esm_rotary.py @@ -0,0 +1,92 @@ +"""Stable ESM rotary embeddings independent of Transformers internals. + +Transformers 5 changed both the name and call contract of its private ESM +rotary helper. FastPLMs checkpoints use the earlier two-tensor contract, so +the small mathematical primitive lives here instead of importing a private +Transformers implementation. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +def _rotate_half(tensor: torch.Tensor) -> torch.Tensor: + """Rotate the final dimension of X by 90 degrees in paired subspaces.""" + + # tensor: (..., d) + first, second = tensor.chunk(2, dim=-1) # (..., d / 2), (..., d / 2) + return torch.cat((-second, first), dim=-1) # (..., d) + + +def apply_rotary_pos_emb( + tensor: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + """Apply cached rotary factors to X with shape ``(b, h, l, d)``.""" + + # tensor: (b, h, l, d); cos, sin: (1, 1, l_cache, d) + cos = cos[:, :, : tensor.shape[-2], :] # (1, 1, l, d) + sin = sin[:, :, : tensor.shape[-2], :] # (1, 1, l, d) + return tensor * cos + _rotate_half(tensor) * sin # (b, h, l, d) + + +class RotaryEmbedding(nn.Module): + """Apply rotary position embeddings to query and key tensors.""" + + inv_freq: torch.Tensor + + def __init__(self, dim: int) -> None: + super().__init__() + frequencies = 1.0 / ( # (d / 2,) + 10_000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim) + ) + # Keep this persistent to preserve the historical checkpoint schema. + self.register_buffer("inv_freq", frequencies) + self._seq_len_cached: int | None = None + self._cos_cached: torch.Tensor | None = None + self._sin_cached: torch.Tensor | None = None + + def _update_cos_sin_tables( + self, + tensor: torch.Tensor, + seq_dimension: int = 2, + ) -> tuple[torch.Tensor, torch.Tensor]: + # tensor: (..., l, d) + seq_len = tensor.shape[seq_dimension] + cache_stale = ( + self._cos_cached is None + or self._sin_cached is None + or self._seq_len_cached != seq_len + or self._cos_cached.device != tensor.device + ) + if cache_stale: + self._seq_len_cached = seq_len + positions = torch.arange(seq_len, device=tensor.device).type_as( # (l,) + self.inv_freq + ) + angles = torch.outer(positions, self.inv_freq) # (l, d / 2) + angles = torch.cat((angles, angles), dim=-1).to(tensor.device) # (l, d) + self._cos_cached = angles.cos()[None, None, :, :] # (1, 1, l, d) + self._sin_cached = angles.sin()[None, None, :, :] # (1, 1, l, d) + + assert self._cos_cached is not None + assert self._sin_cached is not None + return self._cos_cached, self._sin_cached # (1, 1, l, d), (1, 1, l, d) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # query, key: (b, h, l, d) + cos, sin = self._update_cos_sin_tables( # (1, 1, l, d), (1, 1, l, d) + key, + seq_dimension=-2, + ) + return ( + apply_rotary_pos_emb(query, cos, sin).to(dtype=query.dtype), # (b, h, l, d) + apply_rotary_pos_emb(key, cos, sin).to(dtype=key.dtype), # (b, h, l, d) + ) diff --git a/__init__.py b/src/fastplms/models/ankh/__init__.py similarity index 100% rename from __init__.py rename to src/fastplms/models/ankh/__init__.py diff --git a/src/fastplms/models/ankh/modeling_ankh.py b/src/fastplms/models/ankh/modeling_ankh.py new file mode 100644 index 0000000..86677a0 --- /dev/null +++ b/src/fastplms/models/ankh/modeling_ankh.py @@ -0,0 +1,1645 @@ +from __future__ import annotations + +import math +import torch +import torch.nn as nn +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from numbers import Real +from typing import Any, ClassVar +from tokenizers import pre_tokenizers +from torch.nn import functional as F +from transformers import ( + AutoTokenizer, + PretrainedConfig, + PreTrainedModel, + T5ForConditionalGeneration, +) +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + + +try: + from fastplms.attention import ( + AttentionBackend, + FastPLMsAttentionMixin, + bool_to_additive_mask, + get_attention_mask, + resolve_attention_backend, + resolve_attention_backend_for_call, + set_config_attn_implementation, + ) + from fastplms.embeddings import ( + EmbeddingBatch, + EmbeddingMixin, + select_hidden_state_embeddings, + ) + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "EmbeddingBatch", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "bool_to_additive_mask", + "get_attention_mask", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + "set_config_attn_implementation", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +# --------------------------------------------------------------------------- +# Output dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class AnkhEncoderOutput(ModelOutput): + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +@dataclass +class AnkhMaskedLMOutput(ModelOutput): + loss: torch.Tensor | None = None + logits: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +class FastAnkhConfig(PretrainedConfig): + model_type = "fast_ankh" + attribute_map: ClassVar[dict[str, str]] = { + "head_dim": "d_kv", + "hidden_size": "d_model", + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + } + + def __init__( + self, + vocab_size: int = 144, + d_model: int = 768, + d_kv: int = 64, + d_ff: int = 3072, + num_heads: int = 12, + num_layers: int = 48, + num_decoder_layers: int | None = None, + relative_attention_num_buckets: int = 64, + relative_attention_max_distance: int = 128, + dense_act_fn: str = "gelu_new", + feed_forward_proj: str | None = None, + dropout_rate: float = 0.0, + layer_norm_epsilon: float = 1e-6, + initializer_factor: float = 1.0, + pad_token_id: int = 0, + eos_token_id: int = 1, + decoder_start_token_id: int | None = None, + use_cache: bool = True, + tie_word_embeddings: bool = True, + attn_backend: str | None = None, + **kwargs, + ): + if feed_forward_proj is None: + feed_forward_proj = ( + "gated-gelu" if dense_act_fn == "gelu_new" else f"gated-{dense_act_fn}" + ) + if decoder_start_token_id is None: + decoder_start_token_id = pad_token_id + if isinstance(dropout_rate, bool) or not isinstance(dropout_rate, Real): + raise TypeError("dropout_rate must be a real number in [0, 1).") + dropout_rate = float(dropout_rate) + if not 0.0 <= dropout_rate < 1.0: + raise ValueError("dropout_rate must be in [0, 1).") + serialized_encoder_decoder = kwargs.pop("is_encoder_decoder", True) + if serialized_encoder_decoder is not True: + raise ValueError( + "FastAnkhConfig requires is_encoder_decoder=true to match the official " + "T5 configuration." + ) + super().__init__( + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + decoder_start_token_id=decoder_start_token_id, + is_encoder_decoder=True, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + self.vocab_size = vocab_size + self.d_model = d_model + self.d_kv = d_kv + self.d_ff = d_ff + self.num_heads = num_heads + self.num_layers = num_layers + self.num_decoder_layers = num_layers if num_decoder_layers is None else num_decoder_layers + self.relative_attention_num_buckets = relative_attention_num_buckets + self.relative_attention_max_distance = relative_attention_max_distance + self.dense_act_fn = dense_act_fn + self.feed_forward_proj = feed_forward_proj + self.is_gated_act = feed_forward_proj.startswith("gated-") + self.dropout_rate = dropout_rate + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_factor = initializer_factor + self.use_cache = use_cache + self.scale_decoder_outputs = tie_word_embeddings + self.tie_word_embeddings = tie_word_embeddings + self.attn_backend = attn_backend + + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + return output + + +_TOKENIZER_LOAD_CONTEXT_KEYS = ( + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "subfolder", + "token", + "trust_remote_code", +) + + +def configure_ankh_tokenizer(tokenizer: Any) -> Any: + """Apply ANKH's residue-aware pre-tokenizer to a tokenizer instance. + + The tokenizer files published by the official checkpoints use a leading + metaspace convention intended for natural-language text. Protein inputs + are already residue-delimited, so retaining that convention emits a + leading ```` token. FastPLMs configures the fast tokenizer to split + raw residue strings and tight sentinel prompts without manufacturing a + whitespace token. + """ + + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is None: + if getattr(tokenizer, "is_fast", None) is False: + raise TypeError( + "ANKH requires a fast tokenizer so its residue-aware pre-tokenizer " + "can be configured." + ) + # Lightweight tokenizer doubles used by offline CPU contracts need not + # expose a Rust tokenizer backend. + return tokenizer + backend.pre_tokenizer = pre_tokenizers.Metaspace( + replacement="\u2581", + prepend_scheme="never", + split=True, + ) + return tokenizer + + +def normalize_ankh_sequence(sequence: str) -> str: + """Return one ANKH protein sequence in canonical raw-residue form.""" + + if not isinstance(sequence, str): + raise TypeError("ANKH protein sequences must be strings.") + normalized = "".join(sequence.split()) + if not normalized: + raise ValueError("ANKH protein sequences must not be empty or whitespace-only.") + return normalized + + +def normalize_ankh_decoder_prompt(prompt: str) -> str: + """Return a decoder prompt with residues and sentinels directly adjacent.""" + + if not isinstance(prompt, str): + raise TypeError("ANKH decoder prompts must be strings.") + normalized = "".join(prompt.split()) + if not normalized: + raise ValueError("ANKH decoder prompts must not be empty or whitespace-only.") + return normalized + + +def _normalize_ankh_text_batch( + values: str | Sequence[str], + *, + field: str, +) -> str | list[str]: + normalizer = normalize_ankh_sequence if field == "sequence" else normalize_ankh_decoder_prompt + if isinstance(values, str): + return normalizer(values) + if isinstance(values, bytes) or not isinstance(values, Sequence): + raise TypeError(f"ANKH {field} inputs must be a string or a sequence of strings.") + return [normalizer(value) for value in values] + + +def tokenize_ankh_sequences( + tokenizer: Any, + sequences: str | Sequence[str], + **tokenizer_kwargs: Any, +) -> Any: + """Tokenize raw ANKH protein sequences with one model-wide contract.""" + + configured = configure_ankh_tokenizer(tokenizer) + normalized = _normalize_ankh_text_batch(sequences, field="sequence") + return configured(normalized, **tokenizer_kwargs) + + +def tokenize_ankh_decoder_prompts( + tokenizer: Any, + prompts: str | Sequence[str], + **tokenizer_kwargs: Any, +) -> Any: + """Tokenize explicit ANKH decoder prompts without whitespace ```` tokens.""" + + configured = configure_ankh_tokenizer(tokenizer) + normalized = _normalize_ankh_text_batch(prompts, field="decoder prompt") + return configured(normalized, **tokenizer_kwargs) + + +def _load_ankh_tokenizer( + config: FastAnkhConfig, + load_context: Mapping[str, Any] | None = None, +): + """Load the tokenizer from the same immutable checkpoint as the model.""" + name_or_path = str(getattr(config, "_name_or_path", "")).strip() + if not name_or_path: + raise RuntimeError( + "ANKH tokenizer loading requires a model loaded with from_pretrained " + "so checkpoint provenance is available." + ) + tokenizer_kwargs = { + key: value + for key, value in dict(load_context or {}).items() + if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None + } + # The resolved commit is authoritative. In particular, do not reload the + # tokenizer from a moving branch when Transformers resolved model weights to + # an immutable Hub commit. + revision = getattr(config, "_commit_hash", None) + if revision: + tokenizer_kwargs["revision"] = revision + tokenizer = AutoTokenizer.from_pretrained(name_or_path, **tokenizer_kwargs) + return configure_ankh_tokenizer(tokenizer) + + +class _AnkhTokenizerLoadMixin: + """Keep tokenizer loading scoped to the model instance and weight request.""" + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs} + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + model = loaded[0] if isinstance(loaded, tuple) else loaded + model.__dict__["_fastplms_tokenizer_load_context"] = load_context + model.__dict__["_fastplms_tokenizer"] = None + return loaded + + @property + def tokenizer(self): + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + tokenizer = _load_ankh_tokenizer( + self.config, + self.__dict__.get("_fastplms_tokenizer_load_context"), + ) + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value) -> None: + self.__dict__["_fastplms_tokenizer"] = configure_ankh_tokenizer(value) + + def _tokenize_sequence_batch( + self, + sequences: Sequence[str], + *, + tokenizer: Any | None = None, + **tokenizer_kwargs: Any, + ) -> Any: + resolved_tokenizer = tokenizer if tokenizer is not None else self.tokenizer + return tokenize_ankh_sequences( + resolved_tokenizer, + sequences, + **tokenizer_kwargs, + ) + + def embed_dataset(self, inputs: Any, **kwargs: Any) -> Any: + explicit_tokenizer = kwargs.get("tokenizer") + if explicit_tokenizer is not None: + kwargs["tokenizer"] = configure_ankh_tokenizer(explicit_tokenizer) + decoder_inputs = kwargs.get("decoder_inputs") + if ( + decoder_inputs is not None + and not isinstance(decoder_inputs, (str, bytes)) + and isinstance(decoder_inputs, Sequence) + ): + kwargs["decoder_inputs"] = [ + normalize_ankh_decoder_prompt(value) for value in decoder_inputs + ] + return EmbeddingMixin.embed_dataset(self, inputs, **kwargs) + + +def _validate_hidden_state_source(hidden_state_source: str) -> str: + if hidden_state_source not in {"encoder", "decoder"}: + raise ValueError( + "hidden_state_source must be either 'encoder' or 'decoder'; " + f"received {hidden_state_source!r}." + ) + return hidden_state_source + + +def _require_encoder_embedding_source( + hidden_state_source: str, + *, + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, +) -> None: + source = _validate_hidden_state_source(hidden_state_source) + if source == "decoder": + raise ValueError( + "Decoder hidden states require FastAnkhForConditionalGeneration loaded " + "through AutoModelForSeq2SeqLM; the encoder-only ANKH view does not " + "allocate a decoder." + ) + decoder_values = (decoder_inputs, decoder_input_ids, decoder_attention_mask) + if any(value is not None for value in decoder_values): + raise ValueError( + "decoder_inputs, decoder_input_ids, and decoder_attention_mask are only " + "valid when hidden_state_source='decoder'." + ) + + +def _biological_token_mask( + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + tokenizer: Any, +) -> torch.Tensor: + mask = attention_mask.to(device=input_ids.device, dtype=torch.bool) + special_ids = tuple(int(value) for value in getattr(tokenizer, "all_special_ids", ())) + if special_ids: + mask = mask & ~torch.isin( + input_ids, + torch.tensor(special_ids, device=input_ids.device, dtype=input_ids.dtype), + ) + return mask + + +# --------------------------------------------------------------------------- +# Submodules +# --------------------------------------------------------------------------- + + +class AnkhRMSNorm(nn.Module): + """T5-style RMS layer norm: scales without mean subtraction or bias.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # hidden_states: (..., d) + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) # (..., 1) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(self.weight.dtype) + + +def _gelu_new(x: torch.Tensor) -> torch.Tensor: + return ( + 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + ) + + +class AnkhGatedFFN(nn.Module): + """T5-style gated feed-forward: activation(wi_0(x)) * wi_1(x) -> wo.""" + + def __init__(self, config: FastAnkhConfig) -> None: + super().__init__() + self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.act = F.silu if config.dense_act_fn == "silu" else _gelu_new + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # hidden_states: (b, l, d) + hidden_states = self.act(self.wi_0(hidden_states)) * self.wi_1(hidden_states) + return self.wo(self.dropout(hidden_states)) + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class AnkhSelfAttention(nn.Module): + """T5-style self-attention with relative position bias and multi-backend dispatch. + + Only layer 0 has ``has_relative_attention_bias=True`` and owns the + ``nn.Embedding`` that produces the position bias. All other layers + receive the precomputed bias through the forward call. + """ + + def __init__( + self, + config: FastAnkhConfig, + has_relative_attention_bias: bool = False, + ) -> None: + super().__init__() + self.num_heads = config.num_heads + self.d_kv = config.d_kv + self.inner_dim = self.num_heads * self.d_kv + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + + self.q = nn.Linear(config.d_model, self.inner_dim, bias=False) + self.k = nn.Linear(config.d_model, self.inner_dim, bias=False) + self.v = nn.Linear(config.d_model, self.inner_dim, bias=False) + self.o = nn.Linear(self.inner_dim, config.d_model, bias=False) + # T5/ANKH attention is unscaled: scores = Q K^T (no 1/sqrt(d_kv)). + # The learned relative position bias absorbs any temperature. + self.scale = 1.0 + self.dropout_prob = float(config.dropout_rate) + + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding( + config.relative_attention_num_buckets, config.num_heads + ) + + self.attn_backend: AttentionBackend = AttentionBackend.SDPA # set by encoder + + # ---- T5 relative position bucketing ---- + + @staticmethod + def _relative_position_bucket( + relative_position: torch.Tensor, + num_buckets: int = 32, + max_distance: int = 128, + ) -> torch.Tensor: + """Bidirectional log-bucketed relative position mapping (T5 style).""" + # Bidirectional: half buckets for negative, half for positive + num_buckets //= 2 + relative_buckets = (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.clamp(relative_position_if_large, max=num_buckets - 1) + + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) + return relative_buckets + + def compute_bias( + self, query_length: int, key_length: int, device: torch.device + ) -> torch.Tensor: + """Compute the position-bias tensor A with shape (1, h, q, k).""" + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] + relative_position = memory_position - context_position + buckets = self._relative_position_bucket( + relative_position, + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias(buckets) # (q, k, h) + return values.permute(2, 0, 1).unsqueeze(0) # (1, h, q, k) + + # ---- Forward ---- + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + output_attentions: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Returns (attn_output, attn_weights_or_none, position_bias).""" + # hidden_states: (b, l, d) + batch_size, seq_length = hidden_states.shape[:2] + hidden_shape = (batch_size, seq_length, self.num_heads, self.d_kv) + + query_heads = self.q(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + key_heads = self.k(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + value_heads = self.v(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + + # The first layer computes the bias once; later layers reuse it. + if position_bias is None and self.has_relative_attention_bias: + position_bias = self.compute_bias(seq_length, seq_length, hidden_states.device) + # Fold padding mask into position bias so layers don't need separate mask. + if attention_mask_4d is not None: + position_bias = position_bias + bool_to_additive_mask( + attention_mask_4d, position_bias.dtype + ) + + if effective_backend is None: + effective_backend = resolve_attention_backend_for_call( + self.attn_backend, + output_attentions=output_attentions, + ) + if output_attentions: + attn_output, attn_weights = self._manual_attn( + query_heads, key_heads, value_heads, position_bias + ) + return self.o(attn_output), attn_weights, position_bias + + if effective_backend == AttentionBackend.EAGER: + attn_output, _ = self._manual_attn(query_heads, key_heads, value_heads, position_bias) + elif effective_backend == AttentionBackend.SDPA: + attn_output = self._sdpa_attn(query_heads, key_heads, value_heads, position_bias) + else: + raise AssertionError(f"Unsupported backend for ANKH: {effective_backend}") + + return self.o(attn_output), None, position_bias + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + position_bias: torch.Tensor | None, + ) -> torch.Tensor: + # position_bias: (1, h, l, l), including padding + # Never mutate torch.backends.cuda process-global reduction policy from + # a model forward. Concurrent model requests must not change each + # other's numerical behavior or restore a stale process setting. + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=position_bias, + dropout_p=self.dropout_prob if self.training else 0.0, + scale=self.scale, + ) # (b, h, l, d_h) + return ( + context_heads.transpose(1, 2) + .contiguous() + .view(query_heads.shape[0], -1, self.inner_dim) + ) + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + position_bias: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = ( + torch.matmul(query_heads, key_heads.transpose(-1, -2)) * self.scale + ) # (b, h, l, l) + if position_bias is not None: + attn_weights = attn_weights + position_bias + attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights) + if self.dropout_prob > 0 and self.training: + attn_weights = F.dropout( + attn_weights, + p=self.dropout_prob, + training=self.training, + ) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = ( + context_heads.transpose(1, 2) + .contiguous() + .view(query_heads.shape[0], -1, self.inner_dim) + ) + return attn_output, attn_weights + + +# --------------------------------------------------------------------------- +# Encoder block & stack (T5-compatible key naming) +# --------------------------------------------------------------------------- + + +class AnkhSelfAttentionLayer(nn.Module): + """Wraps AnkhSelfAttention + layer_norm to match T5Block.layer[0] key naming.""" + + def __init__( + self, + config: FastAnkhConfig, + has_relative_attention_bias: bool = False, + ) -> None: + super().__init__() + self.SelfAttention = AnkhSelfAttention(config, has_relative_attention_bias) + self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + output_attentions: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + normed = self.layer_norm(hidden_states) + attn_output, attn_weights, position_bias = self.SelfAttention( + normed, + attention_mask_4d=attention_mask_4d, + position_bias=position_bias, + output_attentions=output_attentions, + effective_backend=effective_backend, + ) + hidden_states = hidden_states + self.dropout(attn_output) + return hidden_states, attn_weights, position_bias + + +class AnkhFFLayer(nn.Module): + """Wraps AnkhGatedFFN + layer_norm to match T5Block.layer[1] key naming.""" + + def __init__(self, config: FastAnkhConfig) -> None: + super().__init__() + self.DenseReluDense = AnkhGatedFFN(config) + self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normed = self.layer_norm(hidden_states) + hidden_states = hidden_states + self.dropout(self.DenseReluDense(normed)) + return hidden_states + + +class AnkhBlock(nn.Module): + """Single transformer block with T5-compatible .layer ModuleList naming.""" + + def __init__( + self, + config: FastAnkhConfig, + has_relative_attention_bias: bool = False, + ) -> None: + super().__init__() + self.layer = nn.ModuleList( + [ + AnkhSelfAttentionLayer(config, has_relative_attention_bias), + AnkhFFLayer(config), + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + output_attentions: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + hidden_states, attn_weights, position_bias = self.layer[0]( + hidden_states, + attention_mask_4d=attention_mask_4d, + position_bias=position_bias, + output_attentions=output_attentions, + effective_backend=effective_backend, + ) + hidden_states = self.layer[1](hidden_states) + return hidden_states, attn_weights, position_bias + + +# --------------------------------------------------------------------------- +# PreTrainedModel base +# --------------------------------------------------------------------------- + + +class AnkhPreTrainedModel( + _AnkhTokenizerLoadMixin, + FastPLMsAttentionMixin, + PreTrainedModel, +): + config_class = FastAnkhConfig + base_model_prefix = "encoder" + supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["AnkhBlock"] + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _supports_flex_attn = False + _fastplms_attention_implementations = ("eager", "sdpa") + embedding_unsupported_pooling = ("cls",) + + def __init__(self, config: FastAnkhConfig, *args, **kwargs) -> None: + super().__init__(config, *args, **kwargs) + self.__dict__["_fastplms_tokenizer"] = None + self.__dict__["_fastplms_tokenizer_load_context"] = {} + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + factor = self.config.initializer_factor + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=factor * (self.config.d_model**-0.5)) + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=factor * 1.0) + elif isinstance(module, AnkhRMSNorm): + module.weight.data.fill_(1.0) + + def post_init(self) -> None: + super().post_init() + + def get_output_embeddings(self): + return None + + def _embedding_metadata(self, **context: Any) -> Mapping[str, Any]: + source = _validate_hidden_state_source(context.get("hidden_state_source", "encoder")) + if source != "encoder": + raise ValueError( + "Decoder hidden states require FastAnkhForConditionalGeneration loaded " + "through AutoModelForSeq2SeqLM." + ) + return { + "architecture": "ANKH-T5", + "hidden_state_stack": "encoder", + "layer_order": "embedding-plus-transformer-blocks", + } + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"{type(self).__name__} does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.config.attn_backend = backend + resolved = resolve_attention_backend(backend) + for module in self.modules(): + if isinstance(module, FAST_ANKH_ENCODER): + module.attention_backend = resolved + elif isinstance(module, AnkhSelfAttention): + module.attn_backend = resolved + + +# --------------------------------------------------------------------------- +# FAST_ANKH_ENCODER (mirrors T5Stack key naming) +# --------------------------------------------------------------------------- + + +class FAST_ANKH_ENCODER(AnkhPreTrainedModel, EmbeddingMixin): + """Inner encoder that mirrors T5Stack attribute naming for weight compliance. + + State dict keys: embed_tokens.*, block.{i}.layer.0.SelfAttention.*, + block.{i}.layer.1.DenseReluDense.*, final_layer_norm.*. + """ + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + AnkhPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + + resolved = resolve_attention_backend(config.attn_backend) + if resolved.is_flash: + raise ValueError( + "ANKH does not support FlashAttention because it requires relative position bias." + ) + self.attention_backend = resolved + + self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model) + self.block = nn.ModuleList( + [ + AnkhBlock(config, has_relative_attention_bias=(i == 0)) + for i in range(config.num_layers) + ] + ) + for blk in self.block: + blk.layer[0].SelfAttention.attn_backend = self.attention_backend + + self.final_layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + self.gradient_checkpointing = False + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + hidden_state_source: str = "encoder", + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + _require_encoder_embedding_source( + hidden_state_source, + decoder_inputs=decoder_inputs, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + ) + hidden_states = self.embed_tokens(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + encoder_output = self._run_encoder( + hidden_states, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + ) + return select_hidden_state_embeddings( + encoder_output.last_hidden_state, + encoder_output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def _run_encoder( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool = False, + output_attentions: bool = False, + ) -> AnkhEncoderOutput: + # T5Stack applies this module both to the input embeddings and after + # final normalization. Keeping those as separate calls preserves the + # official training-time stochastic path without affecting eval mode. + hidden_states = self.dropout(hidden_states) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + batch_size, seq_len = hidden_states.shape[:2] + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + _, attention_mask_4d, _ = get_attention_mask( + effective_backend=effective_backend, + batch_size=batch_size, + seq_len=seq_len, + device=hidden_states.device, + attention_mask=attention_mask, + dtype=hidden_states.dtype, + mask_semantics="padding", + ) + + position_bias = None + + for layer_module in self.block: + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + if self.gradient_checkpointing and self.training: + hidden_states, attn_weights, position_bias = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + attention_mask_4d, + position_bias, + output_attentions, + effective_backend, + ) + else: + hidden_states, attn_weights, position_bias = layer_module( + hidden_states, + attention_mask_4d=attention_mask_4d, + position_bias=position_bias, + output_attentions=output_attentions, + effective_backend=effective_backend, + ) + + if all_attentions is not None: + all_attentions = (*all_attentions, attn_weights) + + hidden_states = self.dropout(self.final_layer_norm(hidden_states)) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + return AnkhEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + ) -> AnkhEncoderOutput | tuple[torch.Tensor, ...]: + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + hidden_states = self.embed_tokens(input_ids) + elif inputs_embeds is not None: + hidden_states = inputs_embeds + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + outputs = self._run_encoder( + hidden_states, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states or False, + output_attentions=output_attentions or False, + ) + return outputs if return_dict else outputs.to_tuple() + + +# --------------------------------------------------------------------------- +# Model classes +# --------------------------------------------------------------------------- + + +class FastAnkhModel(AnkhPreTrainedModel, EmbeddingMixin): + """ANKH encoder model for embedding extraction.""" + + _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"} + # The published ANKH checkpoint is the complete official T5 state. AutoModel + # intentionally exposes only its encoder view without allocating a decoder. + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [ + r"^decoder\.", + r"^lm_head\.", + ] + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + AnkhPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.encoder = FAST_ANKH_ENCODER(config) + self.encoder.embed_tokens = self.shared + self.post_init() + + def get_input_embeddings(self): + return self.encoder.embed_tokens + + def set_input_embeddings(self, value): + self.shared = value + self.encoder.embed_tokens = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **embedding_kwargs, + ) -> torch.Tensor: + return self.encoder._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + **embedding_kwargs, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + ) -> AnkhEncoderOutput | tuple[torch.Tensor, ...]: + return self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=return_dict, + ) + + +class FastAnkhForMaskedLMExtension( + FastPLMTestTimeTrainingMixin, AnkhPreTrainedModel, EmbeddingMixin +): + """ANKH encoder with LM head for masked language modeling. + + NOTE: The LM head is initialized from the shared embedding weights but is NOT + tied. The original ANKH models were trained with T5's span corruption objective + using an encoder-decoder architecture. This encoder-only MaskedLM variant is + not pre-trained for standard MLM and requires additional fine-tuning. + """ + + _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"} + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"^decoder\."] + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + # The historical Synthyra extension stores an independent output head. + config.tie_word_embeddings = False + AnkhPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.encoder = FAST_ANKH_ENCODER(config) + self.encoder.embed_tokens = self.shared + self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + self.init_ttt({"lora_target_replace_module": "AnkhSelfAttention"}) + + def get_input_embeddings(self): + return self.encoder.embed_tokens + + def set_input_embeddings(self, value): + self.shared = value + self.encoder.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **embedding_kwargs, + ) -> torch.Tensor: + return self.encoder._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + **embedding_kwargs, + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.encoder] + + def _ttt_tokenize( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + del kwargs + if input_ids is not None: + return input_ids + if seq is None: + raise ValueError("Pass either seq or input_ids for ANKH TTT.") + sequences = [seq] if isinstance(seq, str) else seq + tokenized = tokenize_ankh_sequences( + self.tokenizer, + sequences, + return_tensors="pt", + padding=True, + ) + return tokenized["input_ids"] + + def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + amino_acids = "ACDEFGHIKLMNPQRSTVWY" + ids = [self.tokenizer.convert_tokens_to_ids(aa) for aa in amino_acids] + return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + ) -> MaskedLMOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state + logits = self.lm_head(sequence_output) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (logits, *outputs.to_tuple()[1:]) + return (loss, *output) if loss is not None else output + + return MaskedLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class FastAnkhForConditionalGeneration( + _AnkhTokenizerLoadMixin, + T5ForConditionalGeneration, + EmbeddingMixin, +): + """Official ANKH sequence-to-sequence architecture with exact T5 state keys. + + ANKH generation checkpoints are ordinary T5 conditional-generation models. + This class intentionally delegates their decoder, cross-attention, language + model head, caching, generation, and tied-weight behavior to Transformers. + The optimized encoder-only implementation remains available through + :class:`FastAnkhModel`. + """ + + config_class = FastAnkhConfig + embedding_unsupported_pooling = ("cls",) + _fastplms_attention_implementations = ("eager",) + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + requested_backend = getattr(config, "_attn_implementation", None) or config.attn_backend + if requested_backend not in (None, "eager"): + raise ValueError( + "ANKH sequence-to-sequence checkpoints support only eager attention; " + f"received {requested_backend!r}. Use FastAnkhModel for optimized " + "encoder embeddings." + ) + set_config_attn_implementation(config, "eager") + super().__init__(config, **kwargs) + self.__dict__["_fastplms_tokenizer"] = None + self.__dict__["_fastplms_tokenizer_load_context"] = {} + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + encoder_outputs: Any | None = None, + past_key_values: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> Any: + """Run official T5 seq2seq behavior with a fail-closed public signature. + + Transformers' T5 forward currently accepts and silently ignores arbitrary + keyword arguments. FastPLMs keeps the supported T5 arguments explicit so a + misspelled generation, cache, or conditioning argument cannot appear to + have taken effect. + """ + + return super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + def _prepare_decoder_embedding_inputs( + self, + *, + batch_size: int, + decoder_inputs: Sequence[str] | None, + decoder_input_ids: torch.Tensor | None, + decoder_attention_mask: torch.Tensor | None, + tokenizer: Any | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, Any]: + if (decoder_inputs is None) == (decoder_input_ids is None): + raise ValueError( + "hidden_state_source='decoder' requires exactly one of " + "decoder_inputs or decoder_input_ids. Decoder inputs are task-specific " + "and FastPLMs will not synthesize shifted encoder tokens." + ) + resolved_tokenizer = configure_ankh_tokenizer( + tokenizer if tokenizer is not None else self.tokenizer + ) + if decoder_inputs is not None: + if decoder_attention_mask is not None: + raise ValueError( + "decoder_attention_mask may only accompany decoder_input_ids; " + "decoder_inputs are tokenized with their own attention mask." + ) + values = [decoder_inputs] if isinstance(decoder_inputs, str) else list(decoder_inputs) + if len(values) != batch_size: + raise ValueError( + "decoder_inputs must align one-to-one with encoder inputs; " + f"expected {batch_size}, received {len(values)}." + ) + encoded = tokenize_ankh_decoder_prompts( + resolved_tokenizer, + values, + return_tensors="pt", + padding=True, + truncation=False, + ) + decoder_input_ids = encoded["input_ids"] + decoder_attention_mask = encoded.get("attention_mask") + if decoder_input_ids is None: + raise RuntimeError( + "Decoder input resolution completed without decoder_input_ids." + ) + if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[0] != batch_size: + raise ValueError( + "decoder_input_ids must have shape (batch, decoder_sequence_length); " + f"expected batch {batch_size}, received {tuple(decoder_input_ids.shape)}." + ) + if decoder_attention_mask is None: + pad_token_id = self.config.pad_token_id + decoder_attention_mask = ( + torch.ones_like(decoder_input_ids, dtype=torch.bool) + if pad_token_id is None + else decoder_input_ids.ne(pad_token_id) + ) + decoder_start_token_id = self.config.decoder_start_token_id + if ( + decoder_input_ids.shape[1] > 0 + and decoder_start_token_id is not None + and decoder_start_token_id == pad_token_id + ): + decoder_attention_mask[:, 0] |= decoder_input_ids[:, 0].eq(decoder_start_token_id) + if tuple(decoder_attention_mask.shape) != tuple(decoder_input_ids.shape): + raise ValueError( + "decoder_attention_mask must have the same shape as decoder_input_ids; " + f"received {tuple(decoder_attention_mask.shape)} and " + f"{tuple(decoder_input_ids.shape)}." + ) + device = self.shared.weight.device + return ( + decoder_input_ids.to(device=device), + decoder_attention_mask.to(device=device), + resolved_tokenizer, + ) + + def _extract_embedding_stack( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + *, + hidden_state_source: str, + hidden_state_index: int, + store_all_hidden_states: bool, + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + tokenizer: Any | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...] | None]: + source = _validate_hidden_state_source(hidden_state_source) + device = self.shared.weight.device + input_ids = input_ids.to(device=device) + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + attention_mask = attention_mask.to(device=device) + need_hidden_states = store_all_hidden_states or hidden_state_index != -1 + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=need_hidden_states, + output_attentions=output_attentions and source == "encoder", + return_dict=True, + ) + if source == "encoder": + if any( + value is not None + for value in (decoder_inputs, decoder_input_ids, decoder_attention_mask) + ): + raise ValueError( + "Decoder inputs are only valid when hidden_state_source='decoder'." + ) + X = select_hidden_state_embeddings( + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + return X, input_ids, attention_mask, encoder_outputs.attentions + + decoder_input_ids, decoder_attention_mask, _ = self._prepare_decoder_embedding_inputs( + batch_size=input_ids.shape[0], + decoder_inputs=decoder_inputs, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + tokenizer=tokenizer, + ) + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=attention_mask, + use_cache=False, + output_hidden_states=need_hidden_states, + output_attentions=output_attentions, + return_dict=True, + ) + X = select_hidden_state_embeddings( + decoder_outputs.last_hidden_state, + decoder_outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + return X, decoder_input_ids, decoder_attention_mask, decoder_outputs.attentions + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + hidden_state_source: str = "encoder", + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + tokenizer: Any | None = None, + ) -> torch.Tensor: + X, _, _, _ = self._extract_embedding_stack( + input_ids, + attention_mask, + hidden_state_source=hidden_state_source, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + decoder_inputs=decoder_inputs, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + tokenizer=tokenizer, + ) + return X + + def _embedding_batch( + self, + sequences: Sequence[str], + *, + tokenizer: Any | None = None, + max_length: int | None = None, + truncate: bool = True, + need_attentions: bool = False, + hidden_state_source: str = "encoder", + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + decoder_inputs: Sequence[str] | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + ) -> EmbeddingBatch: + del max_length, truncate # The shared runner already crops biological residues. + resolved_tokenizer = configure_ankh_tokenizer( + tokenizer if tokenizer is not None else self.tokenizer + ) + encoded = tokenize_ankh_sequences( + resolved_tokenizer, + list(sequences), + return_tensors="pt", + padding=True, + truncation=False, + ) + device = self.shared.weight.device + input_ids = encoded["input_ids"].to(device=device) + attention_mask = encoded.get("attention_mask", input_ids.new_ones(input_ids.shape)).to( + device=device + ) + X, selected_ids, selected_attention_mask, attentions = self._extract_embedding_stack( + input_ids, + attention_mask, + hidden_state_source=hidden_state_source, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + decoder_inputs=decoder_inputs, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + tokenizer=resolved_tokenizer, + output_attentions=need_attentions, + ) + residue_mask = _biological_token_mask( + selected_ids, + selected_attention_mask, + resolved_tokenizer, + ) + return EmbeddingBatch(X=X, residue_mask=residue_mask, attentions=attentions) + + def _embedding_metadata(self, **context: Any) -> Mapping[str, Any]: + source = _validate_hidden_state_source(context.get("hidden_state_source", "encoder")) + return { + "architecture": "ANKH-T5", + "hidden_state_stack": source, + "layer_order": "embedding-plus-transformer-blocks", + "decoder_inputs": "explicit-task-inputs-required" if source == "decoder" else None, + "decoder_residue_mask": ( + "attention-mask-minus-tokenizer-specials" if source == "decoder" else None + ), + } + + +class FastAnkhForSequenceClassification(AnkhPreTrainedModel, EmbeddingMixin): + _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"} + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [ + r"^decoder\.", + r"^lm_head\.", + ] + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + AnkhPreTrainedModel.__init__(self, config, **kwargs) + self.num_labels = config.num_labels + self.config = config + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.encoder = FAST_ANKH_ENCODER(config) + self.encoder.embed_tokens = self.shared + self.classifier = nn.Linear(config.d_model, config.num_labels) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.post_init() + + def get_input_embeddings(self): + return self.encoder.embed_tokens + + def set_input_embeddings(self, value): + self.shared = value + self.encoder.embed_tokens = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **embedding_kwargs, + ) -> torch.Tensor: + return self.encoder._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + **embedding_kwargs, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + ) -> SequenceClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + ) + # Pool: mean over non-padding tokens + sequence_output = outputs.last_hidden_state + if attention_mask is not None: + mask = attention_mask.unsqueeze(-1).to(sequence_output.dtype) + pooled = (sequence_output * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1) + else: + pooled = sequence_output.mean(dim=1) + logits = self.classifier(pooled) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss = ( + self.mse(logits.squeeze(), labels.squeeze()) + if self.num_labels == 1 + else self.mse(logits, labels) + ) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + if not return_dict: + output = (logits, *outputs.to_tuple()[1:]) + return (loss, *output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class FastAnkhForTokenClassification(AnkhPreTrainedModel, EmbeddingMixin): + _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"} + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [ + r"^decoder\.", + r"^lm_head\.", + ] + + def __init__(self, config: FastAnkhConfig, **kwargs) -> None: + AnkhPreTrainedModel.__init__(self, config, **kwargs) + self.num_labels = config.num_labels + self.shared = nn.Embedding(config.vocab_size, config.d_model) + self.encoder = FAST_ANKH_ENCODER(config) + self.encoder.embed_tokens = self.shared + self.classifier = nn.Linear(config.d_model, config.num_labels) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + + def get_input_embeddings(self): + return self.encoder.embed_tokens + + def set_input_embeddings(self, value): + self.shared = value + self.encoder.embed_tokens = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **embedding_kwargs, + ) -> torch.Tensor: + return self.encoder._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + **embedding_kwargs, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + ) -> TokenClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits, *outputs.to_tuple()[1:]) + return (loss, *output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fastplms/boltz/__init__.py b/src/fastplms/models/boltz/__init__.py similarity index 59% rename from fastplms/boltz/__init__.py rename to src/fastplms/models/boltz/__init__.py index a91ef52..d7dd650 100644 --- a/fastplms/boltz/__init__.py +++ b/src/fastplms/models/boltz/__init__.py @@ -1,11 +1,14 @@ -from fastplms.boltz.modeling_boltz2 import ( +from fastplms.models.boltz.modeling_boltz2 import ( Boltz2Config, Boltz2Model, + Boltz2ModelOutput, Boltz2StructureOutput, ) + __all__ = [ "Boltz2Config", "Boltz2Model", + "Boltz2ModelOutput", "Boltz2StructureOutput", ] diff --git a/src/fastplms/models/boltz/_pair_attention.py b/src/fastplms/models/boltz/_pair_attention.py new file mode 100644 index 0000000..9afea91 --- /dev/null +++ b/src/fastplms/models/boltz/_pair_attention.py @@ -0,0 +1,50 @@ +"""Shared tensor operations for Boltz2 pair-biased attention layers.""" + +from __future__ import annotations + +import math +import torch +from torch import Tensor + + +def reshape_heads(states: Tensor, num_heads: int) -> Tensor: + """Reshape X from ``(b, l, d)`` to ``(b, l, h, d_head)``.""" + + # states: (b, l, d); num_heads: h + batch_size, sequence_length, width = states.shape # b, l, d + if width % num_heads: + raise ValueError(f"width {width} is not divisible by {num_heads} heads") + # states.view(...): (b, l, h, d_h) + return states.view(batch_size, sequence_length, num_heads, width // num_heads) + + +def pair_biased_attention( + query_states: Tensor, + key_states: Tensor, + value_states: Tensor, + pair_bias: Tensor, + key_mask: Tensor, + mask_value: float, +) -> Tensor: + """Return pair-biased attention values with shape ``(b, l_q, h, d_head)``.""" + + # query_states: (b, l_q, h, d_h); key_states, value_states: (b, l_k, h, d_h) + # pair_bias: (b, h, l_q, l_k); key_mask: (b, l_k) + head_dim = query_states.shape[-1] # d_h + with torch.autocast("cuda", enabled=False): + scores = torch.einsum( + "bihd,bjhd->bhij", + query_states.float(), + key_states.float(), + ) # (b, h, l_q, l_k) + scores = scores / math.sqrt(head_dim) + pair_bias.float() # (b, h, l_q, l_k) + scores = ( + scores + (1 - key_mask[:, None, None].float()) * -mask_value + ) # (b, h, l_q, l_k) + probabilities = scores.softmax(dim=-1) # (b, h, l_q, l_k) + output = torch.einsum( + "bhij,bjhd->bihd", + probabilities, + value_states.float(), + ) # (b, l_q, h, d_h) + return output.to(value_states.dtype) # (b, l_q, h, d_h) diff --git a/src/fastplms/models/boltz/cif_writer.py b/src/fastplms/models/boltz/cif_writer.py new file mode 100644 index 0000000..b7f4ebe --- /dev/null +++ b/src/fastplms/models/boltz/cif_writer.py @@ -0,0 +1,147 @@ +import numpy as np +import torch +from pathlib import Path + +from .minimal_structures import ProteinStructureTemplate + + +def _confidence_per_atom( + plddt: torch.Tensor | None, + atom_to_residue: list[int], + num_atoms: int, + sample_index: int, +) -> np.ndarray: + # plddt: (n_i,) or (n_s, n_i); atom_to_residue: length n_a; num_atoms: n_a + if plddt is None: + return np.ones((num_atoms,), dtype=np.float32) * 100.0 # (n_a,) + + values = plddt.detach().cpu() # (n_i,) or (n_s, n_i) + if values.ndim == 1: + values = values.unsqueeze(0) # (1, n_i) + if values.ndim != 2: + raise ValueError("Expected pLDDT tensor S with shape (n_samples, n_items).") + # values: (n_s, n_i) beyond this point + if not 0 <= sample_index < values.shape[0]: + raise IndexError("sample_index out of range for pLDDT.") + + selected = values[sample_index] # (n_i,) + if selected.shape[0] == num_atoms: + return (selected.numpy() * 100.0).astype(np.float32) # (n_a,) + + num_residues = max(atom_to_residue) + 1 # n_r + if selected.shape[0] == num_residues: + expanded = np.zeros((num_atoms,), dtype=np.float32) # (n_a,) + selected_np = selected.numpy() # (n_r,) + for atom_idx, residue_idx in enumerate(atom_to_residue): + expanded[atom_idx] = selected_np[residue_idx] * 100.0 # () -> () + return expanded # (n_a,) + + raise ValueError( + "pLDDT item count must match either atoms or residues: " + f"received {selected.shape[0]}, expected {num_atoms} atoms or " + f"{num_residues} residues." + ) + + +def write_cif( + structure_template: ProteinStructureTemplate, + atom_coords: torch.Tensor, + atom_mask: torch.Tensor, + output_path: str, + plddt: torch.Tensor | None = None, + sample_index: int = 0, +) -> str: + # atom_coords: (n_a, 3) or (n_s, n_a, 3) + # atom_mask: (n_a,) or (n_m, n_a); plddt: (n_i,) or (n_s, n_i) + coords = atom_coords.detach().cpu() # (n_a, 3) or (n_s, n_a, 3) + if coords.ndim == 2: + coords = coords.unsqueeze(0) # (1, n_a, 3) + if coords.ndim != 3 or coords.shape[-1] != 3: + raise ValueError( + "Expected coordinate tensor X with shape (n_samples, n_atoms, 3)." + ) + # coords: (n_s, n_a, 3) beyond this point + if not 0 <= sample_index < coords.shape[0]: + raise IndexError("sample_index out of range.") + selected_coords_tensor = coords[sample_index] # (n_a, 3) + all_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor)) # (n_a, 3) + if torch.any(all_non_finite): + raise ValueError( + "CIF export received non-finite coordinates. " + f"Non-finite count: {int(all_non_finite.sum().item())}" + ) + selected_coords = selected_coords_tensor.numpy() # (n_a, 3) + + mask = atom_mask.detach().cpu() # (n_a,) or (n_m, n_a) + if mask.ndim == 2: + mask = mask[0] # (n_a,) + if mask.ndim != 1: + raise ValueError("Expected atom mask M with shape (n_atoms,).") + # mask: (n_a,) beyond this point + if mask.shape[0] != selected_coords.shape[0]: + raise ValueError("Atom mask/coord size mismatch.") + if not torch.any(mask > 0): + raise ValueError("Atom mask has no valid atoms for CIF export.") + valid_non_finite = torch.logical_not(torch.isfinite(selected_coords_tensor[mask > 0])) + # valid_non_finite: (n_v, 3), where n_v is the valid-atom count + if torch.any(valid_non_finite): + raise ValueError( + "CIF export has non-finite coordinates in unmasked atoms. " + f"Non-finite count: {int(valid_non_finite.sum().item())}" + ) + + b_iso = _confidence_per_atom( + plddt=plddt, + atom_to_residue=structure_template.atom_residue_index, + num_atoms=structure_template.num_atoms, + sample_index=sample_index, + ) # (n_a,) + if b_iso.shape[0] != structure_template.num_atoms: + raise RuntimeError("CIF confidence values do not match the structure atom count.") + + lines = [ + "data_boltz2_prediction", + "#", + "loop_", + "_atom_site.group_PDB", + "_atom_site.id", + "_atom_site.type_symbol", + "_atom_site.label_atom_id", + "_atom_site.label_comp_id", + "_atom_site.label_asym_id", + "_atom_site.label_seq_id", + "_atom_site.Cartn_x", + "_atom_site.Cartn_y", + "_atom_site.Cartn_z", + "_atom_site.occupancy", + "_atom_site.B_iso_or_equiv", + "_atom_site.pdbx_PDB_model_num", + ] + + atom_id = 1 + for idx in range(structure_template.num_atoms): + if mask[idx] <= 0: + continue + + residue_idx = structure_template.atom_residue_index[idx] + residue_name = structure_template.residue_names[residue_idx] + atom_name = structure_template.atom_names[idx] + element = structure_template.atom_elements[idx] + chain_id = structure_template.atom_chain_id[idx] + x_val, y_val, z_val = selected_coords[idx].tolist() # (3,) -> three scalars + b_factor = float(b_iso[idx]) # () -> Python float + + line = ( + f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} " + f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1" + ) + lines.append(line) + atom_id += 1 + + lines.append("#") + text = "\n".join(lines) + "\n" + + out_path = Path(output_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(text, encoding="utf-8") + return str(out_path) diff --git a/src/fastplms/models/boltz/minimal_featurizer.py b/src/fastplms/models/boltz/minimal_featurizer.py new file mode 100644 index 0000000..7c87bcd --- /dev/null +++ b/src/fastplms/models/boltz/minimal_featurizer.py @@ -0,0 +1,850 @@ +import math +import numpy as np +import torch +from torch.nn.functional import one_hot + +from . import vb_const as const +from .minimal_structures import ProteinStructureTemplate + + +_ELEMENT_TO_Z = { + "H": 1, + "C": 6, + "N": 7, + "O": 8, + "P": 15, + "S": 16, +} + +# Canonical formal charges and tetrahedral chirality values extracted from the +# pinned Boltz2 molecule archive. Values use Boltz2's public feature encoding. +_FORMAL_CHARGES = { + ("ARG", "NH1"): 1.0, + ("HIS", "ND1"): 1.0, + ("LYS", "NZ"): 1.0, +} +_CHIRAL_ATOMS = { + "ALA": frozenset({"CA"}), + "ARG": frozenset({"CA"}), + "ASN": frozenset({"CA"}), + "ASP": frozenset({"CA"}), + "CYS": frozenset({"CA"}), + "GLN": frozenset({"CA"}), + "GLU": frozenset({"CA"}), + "HIS": frozenset({"CA"}), + "ILE": frozenset({"CA", "CB"}), + "LEU": frozenset({"CA"}), + "LYS": frozenset({"CA"}), + "MET": frozenset({"CA"}), + "PHE": frozenset({"CA"}), + "PRO": frozenset({"CA"}), + "SER": frozenset({"CA"}), + "THR": frozenset({"CA", "CB"}), + "TRP": frozenset({"CA"}), + "TYR": frozenset({"CA"}), + "UNK": frozenset({"CA"}), + "VAL": frozenset({"CA"}), +} + + +def _normalize_sequence(sequence: str) -> str: + if not isinstance(sequence, str): + raise TypeError("Amino acid sequence must be a string.") + seq = sequence.strip().upper() + if not seq: + raise ValueError("Amino acid sequence must be non-empty.") + for aa in seq: + if aa not in const.prot_letter_to_token: + raise ValueError(f"Unsupported residue code '{aa}'.") + return seq + + +def _atom_name_to_element(atom_name: str) -> str: + name = atom_name.strip().upper() + if len(name) == 0: + return "C" + if name[0].isdigit(): + name = name[1:] + if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"): + return name[0] + return name[0] + + +def _atom_name_to_codes(atom_name: str) -> torch.Tensor: + clipped = atom_name.strip()[:4] + vals = [ord(ch) - 32 for ch in clipped] + while len(vals) < 4: + vals.append(0) + out = torch.tensor(vals, dtype=torch.long) # (4,) + if not (torch.all(out >= 0) and torch.all(out < 64)): + raise ValueError(f"Invalid atom-name encoding for '{atom_name}'.") + return out + + +# Raw first-conformer coordinates extracted at full float32 precision from the +# hash-pinned Boltz2 molecule archive. The feature builder applies the official +# seeded centering, rotation, and translation policy below. +_RDKIT_CONFORMERS: dict[str, dict[str, list[float]]] = { + "ALA": { + "N": [-0.9241582155227661, 1.1821246147155762, 0.712748110294342], + "CA": [-0.2663755416870117, -0.08827890455722809, 0.4008508622646332], + "C": [1.1188693046569824, 0.1387452930212021, -0.14366498589515686], + "O": [1.2882297039031982, 0.8058497309684753, -1.2001261711120605], + "CB": [-1.1134333610534668, -0.8915302753448486, -0.5876986384391785], + }, + "ARG": { + "N": [3.253326416015625, -1.699564814567566, -0.9852627515792847], + "CA": [2.191511392593384, -0.692007839679718, -0.9506462812423706], + "C": [2.80021333694458, 0.672307550907135, -1.1160222291946411], + "O": [2.465440034866333, 1.3987600803375244, -2.0899009704589844], + "CB": [1.3811039924621582, -0.7949594855308533, 0.36316436529159546], + "CG": [0.20344306528568268, 0.19287171959877014, 0.457282155752182], + "CD": [-0.902410626411438, -0.09953747689723969, -0.5671625137329102], + "NE": [-2.0867979526519775, 0.7146661281585693, -0.2911669909954071], + "CZ": [-3.018158197402954, 0.4804283082485199, 0.7831193804740906], + "NH1": [-4.080137729644775, 1.4154621362686157, 0.9938872456550598], + "NH2": [-2.9226150512695312, -0.5637124180793762, 1.5533623695373535], + }, + "ASN": { + "N": [-1.5767980813980103, -1.7227835655212402, 0.2427234649658203], + "CA": [-0.7041534781455994, -0.5622008442878723, 0.42767149209976196], + "C": [-1.246688961982727, 0.6072975993156433, -0.3483567237854004], + "O": [-1.554918885231018, 0.47644123435020447, -1.5638835430145264], + "CB": [0.7263866662979126, -0.9032993316650391, -0.028588544577360153], + "CG": [1.6871562004089355, 0.20810098946094513, 0.28412675857543945], + "OD1": [2.19687557220459, 0.28916797041893005, 1.433493971824646], + "ND2": [1.996717095375061, 1.196706771850586, -0.6981719136238098], + }, + "ASP": { + "N": [-0.11774874478578568, -1.6310220956802368, 0.374360054731369], + "CA": [-0.3847021758556366, -0.19290897250175476, 0.28121331334114075], + "C": [-1.806337594985962, 0.059408850967884064, -0.1472136378288269], + "O": [-2.2237648963928223, -0.36817988753318787, -1.257502555847168], + "CB": [0.5981687307357788, 0.4785478413105011, -0.6920451521873474], + "CG": [2.0088446140289307, 0.36018264293670654, -0.20304732024669647], + "OD1": [2.747267961502075, -0.5711832046508789, -0.6216948628425598], + "OD2": [2.49090313911438, 1.2577818632125854, 0.7443997859954834], + }, + "CYS": { + "N": [-0.058702241629362106, 1.771048665046692, 0.2434082329273224], + "CA": [-0.06546340137720108, 0.47663936018943787, -0.4456723928451538], + "C": [-1.273919701576233, -0.3504071533679962, -0.08361759781837463], + "O": [-1.7324607372283936, -0.34313708543777466, 1.0909277200698853], + "CB": [1.2432398796081543, -0.28917407989501953, -0.19254730641841888], + "SG": [1.484184980392456, -0.7191365361213684, 1.5649693012237549], + }, + "GLN": { + "N": [-1.8543237447738647, -1.0024770498275757, -1.6278940439224243], + "CA": [-1.292184591293335, -0.6786512732505798, -0.3153715133666992], + "C": [-2.2264492511749268, 0.25171151757240295, 0.4067918062210083], + "O": [-2.7255098819732666, -0.08834805339574814, 1.5129321813583374], + "CB": [0.11662524193525314, -0.06213820353150368, -0.45663803815841675], + "CG": [0.8140791058540344, 0.10461423546075821, 0.9002283811569214], + "CD": [2.195122003555298, 0.6569428443908691, 0.7154991030693054], + "OE1": [2.3870291709899902, 1.9003654718399048, 0.7825236916542053], + "NE2": [3.287371873855591, -0.21307705342769623, 0.41959717869758606], + }, + "GLU": { + "N": [-1.3492857217788696, -1.114271640777588, -1.3739068508148193], + "CA": [-1.2765676975250244, -0.506334662437439, -0.04287439212203026], + "C": [-1.9195666313171387, 0.8547031879425049, -0.04214814677834511], + "O": [-1.9091347455978394, 1.5701956748962402, -1.080504298210144], + "CB": [0.18272103369235992, -0.43209096789360046, 0.45213374495506287], + "CG": [1.1058456897735596, 0.3659400939941406, -0.481784850358963], + "CD": [2.5090205669403076, 0.35135790705680847, 0.03646872192621231], + "OE1": [2.896414041519165, 1.2445733547210693, 0.8365193605422974], + "OE2": [3.3727896213531494, -0.6754258275032043, -0.33141785860061646], + }, + "GLY": { + "N": [-1.291549801826477, 0.6080796122550964, -0.4228580892086029], + "CA": [-0.4895951449871063, -0.2882237136363983, 0.40191715955734253], + "C": [0.9350062608718872, -0.2543502449989319, -0.04786944016814232], + "O": [1.3473576307296753, -1.0836502313613892, -0.9022219777107239], + }, + "HIS": { + "N": [1.0371263027191162, -1.5621215105056763, 0.4178937077522278], + "CA": [1.1980191469192505, -0.40241551399230957, -0.46312251687049866], + "C": [2.6561269760131836, -0.07123222947120667, -0.6407948136329651], + "O": [3.1795547008514404, -0.12371637672185898, -1.7859554290771484], + "CB": [0.44051393866539, 0.8116919994354248, 0.09850303083658218], + "CG": [-1.0356801748275757, 0.5519806742668152, 0.15254715085029602], + "ND1": [-1.9089158773422241, 0.6490958333015442, -0.968245267868042], + "CD2": [-1.7095766067504883, 0.1138235554099083, 1.2061808109283447], + "CE1": [-3.0751633644104004, 0.28129827976226807, -0.5511555075645447], + "NE2": [-3.0650453567504883, -0.07846951484680176, 0.8257263898849487], + }, + "ILE": { + "N": [-1.2378733158111572, -1.8146690130233765, -0.16437499225139618], + "CA": [-1.2780015468597412, -0.3994033634662628, 0.22437156736850739], + "C": [-1.9879958629608154, 0.4497596025466919, -0.80372154712677], + "O": [-2.0291709899902344, 0.10288607329130173, -2.015406608581543], + "CB": [0.14225825667381287, 0.146738663315773, 0.5415176153182983], + "CG1": [1.1062071323394775, 0.04519597440958023, -0.6694174408912659], + "CG2": [0.7215384244918823, -0.5685611963272095, 1.7750537395477295], + "CD1": [2.371805191040039, 0.8852734565734863, -0.49045464396476746], + }, + "LEU": { + "N": [1.6452248096466064, -1.017622470855713, -1.0626215934753418], + "CA": [1.3129304647445679, 0.08074887096881866, -0.15053576231002808], + "C": [2.431154727935791, 0.28310757875442505, 0.8357383608818054], + "O": [2.9231173992156982, -0.7013152837753296, 1.451209306716919], + "CB": [0.006050171796232462, -0.2141101360321045, 0.6198461651802063], + "CG": [-1.2534235715866089, -0.3772832751274109, -0.26637399196624756], + "CD1": [-2.455688714981079, -0.7604416012763977, 0.6070646643638611], + "CD2": [-1.5742294788360596, 0.8985711932182312, -1.059889554977417], + }, + "LYS": { + "N": [-2.3938536643981934, -1.4751482009887695, -0.9336304068565369], + "CA": [-2.2094509601593018, -0.6203567981719971, 0.2421264797449112], + "C": [-3.4513463973999023, 0.1961250752210617, 0.4710204303264618], + "O": [-3.9688456058502197, 0.2503468096256256, 1.6188287734985352], + "CB": [-1.008726716041565, 0.3269270658493042, 0.0479682981967926], + "CG": [0.336457222700119, -0.41559070348739624, 0.047045741230249405], + "CD": [1.5088860988616943, 0.5716065168380737, -0.00632853340357542], + "CE": [2.853410005569458, -0.1678532361984253, -0.020898228511214256], + "NZ": [3.9690465927124023, 0.7765668630599976, -0.07007527351379395], + }, + "MET": { + "N": [-1.743452548980713, -0.44039490818977356, 1.851780652999878], + "CA": [-1.2168668508529663, 0.4041755199432373, 0.7779996991157532], + "C": [-2.3118340969085693, 0.683431088924408, -0.2134547382593155], + "O": [-2.6660380363464355, 1.8705639839172363, -0.44476020336151123], + "CB": [-0.0005482881097123027, -0.262584388256073, 0.09933014959096909], + "CG": [0.6952047944068909, 0.6768104434013367, -0.8930339813232422], + "SD": [2.136805772781372, -0.14913176000118256, -1.6563159227371216], + "CE": [3.3404407501220703, 0.348065584897995, -0.38104408979415894], + }, + "PHE": { + "N": [2.7976438999176025, -1.6209964752197266, -0.4779351055622101], + "CA": [1.5799994468688965, -0.8289614915847778, -0.6608061194419861], + "C": [1.9535551071166992, 0.5758142471313477, -1.0445665121078491], + "O": [1.5211102962493896, 1.072593331336975, -2.1190123558044434], + "CB": [0.7137221097946167, -0.8474348187446594, 0.6163292527198792], + "CG": [-0.6229045987129211, -0.18666480481624603, 0.3912639319896698], + "CD1": [-0.8259273767471313, 1.0978055000305176, 0.7381877303123474], + "CD2": [-1.7160141468048096, -0.9425657391548157, -0.2700555920600891], + "CE1": [-2.131312131881714, 1.741182565689087, 0.4800238013267517], + "CE2": [-2.8960046768188477, -0.3551671802997589, -0.505269467830658], + "CZ": [-3.11456298828125, 1.0511738061904907, -0.11009909957647324], + }, + "PRO": { + "N": [-0.577339768409729, -0.5123927593231201, -1.150407314300537], + "CA": [0.5119893550872803, 0.23565764725208282, -0.5195383429527283], + "C": [1.8393173217773438, -0.46712079644203186, -0.6550371646881104], + "O": [2.0836586952209473, -1.166333556175232, -1.6751221418380737], + "CB": [0.1029478907585144, 0.45346882939338684, 0.9297074675559998], + "CG": [-1.4098080396652222, 0.5215330123901367, 0.8787456750869751], + "CD": [-1.7971025705337524, -0.06041542813181877, -0.47876954078674316], + }, + "SER": { + "N": [0.8959873914718628, -1.423298716545105, -0.2692987024784088], + "CA": [-0.018787339329719543, -0.28751760721206665, -0.3863537311553955], + "C": [0.767615795135498, 0.9929561018943787, -0.42147380113601685], + "O": [0.5327092409133911, 1.8535151481628418, -1.3116248846054077], + "CB": [-0.9973848462104797, -0.2769756019115448, 0.7988247871398926], + "OG": [-1.8999865055084229, 0.7915746569633484, 0.691534698009491], + }, + "THR": { + "N": [0.005045319441705942, 1.7655576467514038, 0.0639338344335556], + "CA": [0.4231603145599365, 0.3755212128162384, 0.2734847664833069], + "C": [1.7184523344039917, 0.11012311279773712, -0.4465082287788391], + "O": [1.9252793788909912, 0.5967074036598206, -1.5912789106369019], + "CB": [-0.6634359359741211, -0.615561306476593, -0.21173778176307678], + "OG1": [-1.0762907266616821, -0.29800114035606384, -1.5171915292739868], + "CG2": [-1.8869022130966187, -0.6176311373710632, 0.7084044218063354], + }, + "TRP": { + "N": [-3.0439648628234863, 1.0628471374511719, 0.015692168846726418], + "CA": [-2.463679552078247, 0.024657616391777992, -0.8414031863212585], + "C": [-2.186711072921753, -1.2228965759277344, -0.04691409692168236], + "O": [-2.3364906311035156, -2.3535354137420654, -0.5830773115158081], + "CB": [-1.2062920331954956, 0.5195542573928833, -1.592403531074524], + "CG": [-0.06450498104095459, 0.884079098701477, -0.6835198402404785], + "CD1": [0.16109998524188995, 2.085163116455078, -0.14801783859729767], + "CD2": [0.9781660437583923, -0.005544683896005154, -0.17733941972255707], + "NE1": [1.3106837272644043, 2.0506503582000732, 0.698072075843811], + "CE2": [1.7528777122497559, 0.7044392228126526, 0.6196392774581909], + "CE3": [1.236093521118164, -1.4355946779251099, -0.4235689342021942], + "CZ2": [2.9070138931274414, 0.10196790844202042, 1.3070299625396729], + "CZ3": [2.2851555347442627, -2.007185697555542, 0.1949695348739624], + "CH2": [3.1585631370544434, -1.2034809589385986, 1.1002938747406006], + }, + "TYR": { + "N": [-1.7950150966644287, 0.49119046330451965, -1.3951144218444824], + "CA": [-1.8436503410339355, -0.2694968581199646, -0.14338290691375732], + "C": [-3.240288019180298, -0.28215137124061584, 0.41974759101867676], + "O": [-3.8150112628936768, 0.7977103590965271, 0.7253130674362183], + "CB": [-0.8541494011878967, 0.30990728735923767, 0.8830214738845825], + "CG": [0.5672639012336731, 0.2125871181488037, 0.3916495442390442], + "CD1": [1.2694358825683594, -0.924761176109314, 0.5431669354438782], + "CD2": [1.1872504949569702, 1.3605931997299194, -0.31554359197616577], + "CE1": [2.6519992351531982, -1.0213464498519897, 0.029911600053310394], + "CE2": [2.4384164810180664, 1.2693941593170166, -0.7831515073776245], + "CZ": [3.211331367492676, 0.022552935406565666, -0.6021429300308228], + "OH": [4.513507843017578, -0.05260590463876724, -1.0929937362670898], + }, + "VAL": { + "N": [0.9408224821090698, -1.2608877420425415, 0.652370810508728], + "CA": [0.7287879586219788, -0.3937721848487854, -0.5120788216590881], + "C": [1.7670493125915527, 0.6979415416717529, -0.5508439540863037], + "O": [2.2449657917022705, 1.1683598756790161, 0.5171348452568054], + "CB": [-0.7015937566757202, 0.21480272710323334, -0.5226467847824097], + "CG1": [-1.7662651538848877, -0.845522403717041, -0.8411717414855957], + "CG2": [-1.0500152111053467, 0.9374979138374329, 0.7892321944236755], + }, + "UNK": { + "N": [1.6134154796600342, -1.304404377937317, -0.35241976380348206], + "CA": [0.6100443601608276, -0.23816797137260437, -0.3852081298828125], + "C": [1.0861350297927856, 0.9232651591300964, 0.4423142373561859], + "O": [1.374688982963562, 0.7629808187484741, 1.6591525077819824], + "CB": [-0.7607048153877258, -0.7552146315574646, 0.10266945511102676], + }, +} + + +def _get_atom_position(res_name: str, atom_name: str, atom_idx: int) -> np.ndarray: + """Get the canonical RDKit conformer position for an atom. + + Uses pre-extracted positions from official Boltz2 mol files. Falls back + to a simple geometric placement for unknown residue/atom combinations. + """ + if res_name in _RDKIT_CONFORMERS and atom_name in _RDKIT_CONFORMERS[res_name]: + return np.array(_RDKIT_CONFORMERS[res_name][atom_name], dtype=np.float32) # (3,) + # Fallback for unknown atoms (should not happen for canonical AAs) + angle = (atom_idx + 1) * 0.7 + radius = 1.4 + 0.03 * atom_idx + return np.array( + [ + radius * math.cos(angle), + radius * math.sin(angle), + 0.1 * ((atom_idx % 5) - 2), + ], + dtype=np.float32, + ) # (3,) + + +def _build_template( + sequence: str, +) -> tuple[ + ProteinStructureTemplate, + list[str], + list[int], + list[int], + list[int], + list[np.ndarray], + list[int], +]: + residue_names: list[str] = [] + residue_token_ids: list[int] = [] + atom_names: list[str] = [] + atom_elements: list[str] = [] + atom_residue_index: list[int] = [] + atom_chain_id: list[str] = [] + atom_positions: list[np.ndarray] = [] # each: (3,) + residue_center_atom_idx: list[int] = [] + residue_disto_atom_idx: list[int] = [] + residue_frame_atom_idx: list[int] = [] + + global_atom_idx = 0 + for res_idx, aa in enumerate(sequence): + token_name = const.prot_letter_to_token[aa] + residue_names.append(token_name) + residue_token_ids.append(const.token_ids[token_name]) + + residue_atoms = const.ref_atoms[token_name] + if not residue_atoms: + raise RuntimeError(f"No reference atoms for residue {token_name}.") + center_atom_name = const.res_to_center_atom[token_name] + disto_atom_name = const.res_to_disto_atom[token_name] + + center_idx = -1 + disto_idx = -1 + n_idx = -1 + ca_idx = -1 + c_idx = -1 + + for local_idx, atom_name in enumerate(residue_atoms): + atom_names.append(atom_name) + element = _atom_name_to_element(atom_name) + atom_elements.append(element) + atom_residue_index.append(res_idx) + atom_chain_id.append("A") + + atom_pos = _get_atom_position(token_name, atom_name, local_idx) # (3,) + atom_positions.append(atom_pos) + + if atom_name == center_atom_name: + center_idx = global_atom_idx + if atom_name == disto_atom_name: + disto_idx = global_atom_idx + if atom_name == "N": + n_idx = global_atom_idx + if atom_name == "CA": + ca_idx = global_atom_idx + if atom_name == "C": + c_idx = global_atom_idx + global_atom_idx += 1 + + if center_idx == -1: + center_idx = global_atom_idx - len(residue_atoms) + if disto_idx == -1: + disto_idx = center_idx + if n_idx == -1: + n_idx = center_idx + if ca_idx == -1: + ca_idx = center_idx + if c_idx == -1: + c_idx = center_idx + + residue_center_atom_idx.append(center_idx) + residue_disto_atom_idx.append(disto_idx) + residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx]) + + template = ProteinStructureTemplate( + sequence=sequence, + residue_names=residue_names, + atom_names=atom_names, + atom_elements=atom_elements, + atom_residue_index=atom_residue_index, + atom_chain_id=atom_chain_id, + ) + + return ( + template, + residue_names, + residue_token_ids, + residue_center_atom_idx, + residue_disto_atom_idx, + atom_positions, + residue_frame_atom_idx, + ) + + +def _random_rotation_matrix() -> torch.Tensor: + """Sample a uniform random 3x3 rotation matrix (Algorithm 19 from AF2/Boltz).""" + quaternion = torch.randn((1, 4), dtype=torch.float32) # (1, 4) + squared_norm = (quaternion * quaternion).sum(dim=1) # (1,) + norm = torch.sqrt(squared_norm) # (1,) + signed_norm = torch.where(quaternion[:, 0] < 0, -norm, norm) # (1,) + quaternion = quaternion / signed_norm[:, None] # (1, 4) + + real, i, j, k = torch.unbind(quaternion, dim=-1) # each: (1,) + two_s = 2.0 / (quaternion * quaternion).sum(dim=-1) # (1,) + rotation = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * real), + two_s * (i * k + j * real), + two_s * (i * j + k * real), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * real), + two_s * (i * k - j * real), + two_s * (j * k + i * real), + 1 - two_s * (i * i + j * j), + ), + dim=-1, + ) # (1, 9) + return rotation.reshape(1, 3, 3)[0] # (3, 3) + + +def _center_and_augment_atoms_per_residue( + atom_positions: torch.Tensor, + atom_residue_index: list[int], + num_residues: int, +) -> torch.Tensor: + """Apply Boltz2's seeded conformer augmentation per residue. + + The pinned implementation intentionally excludes the final residue because + it iterates to the maximum reference-space identifier rather than through it. + """ + # atom_positions: (a, 3), where a is the unpadded atom count. + result = atom_positions.clone() # (a, 3) + residue_index_tensor = torch.tensor(atom_residue_index, dtype=torch.long) # (a,) + for residue_idx in range(max(num_residues - 1, 0)): + residue_mask = residue_index_tensor == residue_idx # (a,) + if not torch.any(residue_mask): + raise RuntimeError(f"Residue index {residue_idx} has no atoms.") + residue_coords = result[residue_mask][None] # (1, a_r, 3) + resolved_mask = torch.ones( + residue_coords.shape[:2], dtype=torch.bool, device=residue_coords.device + ) # (1, a_r) + residue_center = torch.sum( + residue_coords * resolved_mask[:, :, None], dim=1, keepdim=True + ) / torch.sum(resolved_mask[:, :, None], dim=1, keepdim=True) # (1, 1, 3) + residue_coords = residue_coords - residue_center # (1, a_r, 3) + rotation = _random_rotation_matrix()[None] # (1, 3, 3) + residue_coords = torch.einsum( + "bmd,bds->bms", residue_coords, rotation + ) # (1, a_r, 3) + residue_coords = ( + residue_coords + torch.randn_like(residue_coords[:, 0:1, :]) + ) # (1, a_r, 3) + result[residue_mask] = residue_coords[0] # (a_r, 3) + return result # (a, 3) + + +def build_boltz2_features( + amino_acid_sequence: str, + num_bins: int = 64, + atoms_per_window_queries: int = 32, +) -> tuple[dict[str, torch.Tensor], ProteinStructureTemplate]: + sequence = _normalize_sequence(amino_acid_sequence) + ( + template, + residue_names, + residue_token_ids, + residue_center_atom_idx, + residue_disto_atom_idx, + atom_positions_np, + residue_frame_atom_idx_flat, + ) = _build_template(sequence) + + num_tokens = len(residue_names) + num_atoms = len(atom_positions_np) + if num_tokens <= 0 or num_atoms <= 0: + raise RuntimeError("Boltz2 feature construction produced an empty protein template.") + + # t is the token count, a is the unpadded atom count, and a_p is the padded atom count. + atom_positions = torch.tensor( + np.asarray(atom_positions_np), dtype=torch.float32 + ) # (a, 3) + atom_positions = _center_and_augment_atoms_per_residue( + atom_positions=atom_positions, + atom_residue_index=template.atom_residue_index, + num_residues=num_tokens, + ) # (a, 3) + + token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0) # (1, t) + residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0) # (1, t) + asym_id = torch.zeros((1, num_tokens), dtype=torch.long) # (1, t) + entity_id = torch.zeros((1, num_tokens), dtype=torch.long) # (1, t) + sym_id = torch.zeros((1, num_tokens), dtype=torch.long) # (1, t) + mol_type = torch.full( + (1, num_tokens), + fill_value=const.chain_type_ids["PROTEIN"], + dtype=torch.long, + ) # (1, t) + + res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long) # (t,) + res_type = one_hot( + res_type_ids, num_classes=const.num_tokens + ).unsqueeze(0) # (1, t, n_res_type) + + # token_bonds encodes explicit covalent cross-links from structure bonds, + # NOT backbone peptide bonds (those are implicit via residue_index + asym_id). + # For a standard single-chain protein without cross-links, this is all zeros. + # This matches the official Boltz2 featurizer (featurizerv2.py lines 696-705). + token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32) # (t, t) + type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long) # (t, t) + token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1) # (1, t, t, 1) + type_bonds = type_bonds.unsqueeze(0) # (1, t, t) + + token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32) # (1, t) + token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32) # (1, t) + token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32) # (1, t) + + num_contact_classes = len(const.contact_conditioning_info) + unspecified_id = const.contact_conditioning_info["UNSPECIFIED"] + contact_ids = torch.full( + (num_tokens, num_tokens), + fill_value=unspecified_id, + dtype=torch.long, + ) # (t, t) + contact_conditioning = one_hot( + contact_ids, + num_classes=num_contact_classes, + ).unsqueeze(0) # (1, t, t, n_contact) + contact_threshold = torch.zeros( + (1, num_tokens, num_tokens), dtype=torch.float32 + ) # (1, t, t) + + if "x-ray diffraction" not in const.method_types_ids: + raise RuntimeError("Boltz2 method metadata omits x-ray diffraction.") + method_feature = torch.full( + (1, num_tokens), + fill_value=const.method_types_ids["x-ray diffraction"], + dtype=torch.long, + ) # (1, t) + modified = torch.zeros((1, num_tokens), dtype=torch.long) # (1, t) + cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32) # (1, t) + affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32) # (1, t) + + ref_pos = atom_positions.unsqueeze(0) # (1, a, 3) + atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32) # (1, a) + atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.bool) # (1, a) + + atom_name_codes = torch.stack( + [_atom_name_to_codes(atom_name) for atom_name in template.atom_names], + dim=0, + ) # (a, 4) + ref_atom_name_chars = one_hot( + atom_name_codes, num_classes=64 + ).unsqueeze(0) # (1, a, 4, 64) + + atomic_numbers = [] + for element in template.atom_elements: + z_value = _ELEMENT_TO_Z[element] if element in _ELEMENT_TO_Z else _ELEMENT_TO_Z["C"] + if z_value >= const.num_elements: + raise RuntimeError( + f"Atomic number {z_value} exceeds the Boltz2 element vocabulary." + ) + atomic_numbers.append(z_value) + ref_element = one_hot( + torch.tensor(atomic_numbers, dtype=torch.long), + num_classes=const.num_elements, + ).unsqueeze(0) # (1, a, n_element) + + ref_charge_values = [] + ref_chirality_values = [] + for atom_name, residue_idx in zip( + template.atom_names, + template.atom_residue_index, + strict=True, + ): + residue_name = residue_names[residue_idx] + ref_charge_values.append(_FORMAL_CHARGES.get((residue_name, atom_name), 0.0)) + ref_chirality_values.append( + 2 if atom_name in _CHIRAL_ATOMS.get(residue_name, frozenset()) else 0 + ) + ref_charge = torch.tensor(ref_charge_values, dtype=torch.float32).unsqueeze(0) # (1, a) + ref_chirality = torch.tensor( + ref_chirality_values, dtype=torch.long + ).unsqueeze(0) # (1, a) + ref_space_uid = torch.tensor( + template.atom_residue_index, dtype=torch.long + ).unsqueeze(0) # (1, a) + + atom_to_token = one_hot( + torch.tensor(template.atom_residue_index, dtype=torch.long), + num_classes=num_tokens, + ).unsqueeze(0) # (1, a, t) + token_to_rep_atom = one_hot( + torch.tensor(residue_disto_atom_idx, dtype=torch.long), + num_classes=num_atoms, + ).unsqueeze(0) # (1, t, a) + token_to_center_atom = one_hot( + torch.tensor(residue_center_atom_idx, dtype=torch.long), + num_classes=num_atoms, + ).unsqueeze(0) # (1, t, a) + r_set_to_rep_atom = token_to_center_atom.clone() # (1, t, a) + + num_backbone_classes = ( + 1 + len(const.protein_backbone_atom_index) + len(const.nucleic_backbone_atom_index) + ) + backbone_ids = [] + for atom_name in template.atom_names: + if atom_name in const.protein_backbone_atom_index: + backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1) + else: + backbone_ids.append(0) + atom_backbone_feat = one_hot( + torch.tensor(backbone_ids, dtype=torch.long), + num_classes=num_backbone_classes, + ).unsqueeze(0) # (1, a, n_backbone) + + # X contains no observed coordinates for sequence-only inference. + coords = torch.zeros((1, 1, num_atoms, 3), dtype=torch.float32) # (1, 1, a, 3) + disto_coords_ensemble = torch.zeros( + (1, 1, num_tokens, 3), + dtype=torch.float32, + ) # (1, 1, t, 3) + + bfactor = torch.zeros((1, num_atoms), dtype=torch.float32) # (1, a) + atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32) # (1, a) + + if atoms_per_window_queries <= 0: + raise ValueError("atoms_per_window_queries must be positive.") + pad_atoms = ( + (num_atoms - 1) // atoms_per_window_queries + 1 + ) * atoms_per_window_queries - num_atoms + if pad_atoms > 0: + ref_pos = torch.nn.functional.pad( + ref_pos, (0, 0, 0, pad_atoms), value=0.0 + ) # (1, a_p, 3) + atom_pad_mask = torch.nn.functional.pad( + atom_pad_mask, (0, pad_atoms), value=0.0 + ) # (1, a_p) + atom_resolved_mask = torch.nn.functional.pad( + atom_resolved_mask, + (0, pad_atoms), + value=0.0, + ) # (1, a_p) + ref_atom_name_chars = torch.nn.functional.pad( + ref_atom_name_chars, + (0, 0, 0, 0, 0, pad_atoms), + value=0.0, + ) # (1, a_p, 4, 64) + ref_element = torch.nn.functional.pad( + ref_element, (0, 0, 0, pad_atoms), value=0.0 + ) # (1, a_p, n_element) + ref_charge = torch.nn.functional.pad( + ref_charge, (0, pad_atoms), value=0.0 + ) # (1, a_p) + ref_chirality = torch.nn.functional.pad( + ref_chirality, (0, pad_atoms), value=0 + ) # (1, a_p) + atom_backbone_feat = torch.nn.functional.pad( + atom_backbone_feat, + (0, 0, 0, pad_atoms), + value=0.0, + ) # (1, a_p, n_backbone) + ref_space_uid = torch.nn.functional.pad( + ref_space_uid, (0, pad_atoms), value=0 + ) # (1, a_p) + coords = torch.nn.functional.pad( + coords, (0, 0, 0, pad_atoms), value=0.0 + ) # (1, 1, a_p, 3) + atom_to_token = torch.nn.functional.pad( + atom_to_token, (0, 0, 0, pad_atoms), value=0.0 + ) # (1, a_p, t) + token_to_rep_atom = torch.nn.functional.pad( + token_to_rep_atom, + (0, pad_atoms), + value=0.0, + ) # (1, t, a_p) + token_to_center_atom = torch.nn.functional.pad( + token_to_center_atom, + (0, pad_atoms), + value=0.0, + ) # (1, t, a_p) + r_set_to_rep_atom = torch.nn.functional.pad( + r_set_to_rep_atom, + (0, pad_atoms), + value=0.0, + ) # (1, t, a_p) + bfactor = torch.nn.functional.pad( + bfactor, (0, pad_atoms), value=0.0 + ) # (1, a_p) + atom_plddt = torch.nn.functional.pad( + atom_plddt, (0, pad_atoms), value=0.0 + ) # (1, a_p) + + frames_idx = torch.tensor( + residue_frame_atom_idx_flat, + dtype=torch.long, + ).reshape(num_tokens, 3) # (t, 3) + frames_idx = frames_idx.unsqueeze(0).unsqueeze(1) # (1, 1, t, 3) + frame_resolved_mask = torch.zeros( + (1, 1, num_tokens), dtype=torch.bool + ) # (1, 1, t) + + msa = ( + torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0) + ) # (1, 1, t) + msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32) # (1, 1, t) + deletion_value = torch.zeros( + (1, 1, num_tokens), dtype=torch.float32 + ) # (1, 1, t) + has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.bool) # (1, 1, t) + msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.long) # (1, 1, t) + deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32) # (1, t) + profile = ( + one_hot( + torch.tensor(residue_token_ids, dtype=torch.long), + num_classes=const.num_tokens, + ) + .float() + .unsqueeze(0) + ) # (1, t, n_res_type) + + template_restype = one_hot( + torch.zeros((1, 1, num_tokens), dtype=torch.long), + num_classes=const.num_tokens, + ) # (1, 1, t, n_res_type) + template_frame_rot = torch.zeros( + (1, 1, num_tokens, 3, 3), dtype=torch.float32 + ) # (1, 1, t, 3, 3) + template_frame_t = torch.zeros( + (1, 1, num_tokens, 3), dtype=torch.float32 + ) # (1, 1, t, 3) + template_cb = torch.zeros( + (1, 1, num_tokens, 3), dtype=torch.float32 + ) # (1, 1, t, 3) + template_ca = torch.zeros( + (1, 1, num_tokens, 3), dtype=torch.float32 + ) # (1, 1, t, 3) + template_mask_cb = torch.zeros( + (1, 1, num_tokens), dtype=torch.float32 + ) # (1, 1, t) + template_mask_frame = torch.zeros( + (1, 1, num_tokens), dtype=torch.float32 + ) # (1, 1, t) + template_mask = torch.zeros( + (1, 1, num_tokens), dtype=torch.float32 + ) # (1, 1, t) + query_to_template = torch.zeros( + (1, 1, num_tokens), dtype=torch.long + ) # (1, 1, t) + visibility_ids = torch.zeros( + (1, 1, num_tokens), dtype=torch.float32 + ) # (1, 1, t) + + disto_target = torch.zeros( + (1, num_tokens, num_tokens, 1, num_bins), + dtype=torch.float32, + ) # (1, t, t, 1, n_bin) + disto_target[..., 0] = 1.0 # (1, t, t, 1) + disto_center = torch.zeros( + (1, num_tokens, 3), dtype=torch.float32 + ) # (1, t, 3) + + features: dict[str, torch.Tensor] = { + "token_index": token_index, + "residue_index": residue_index, + "asym_id": asym_id, + "entity_id": entity_id, + "sym_id": sym_id, + "mol_type": mol_type, + "res_type": res_type, + "disto_center": disto_center, + "token_bonds": token_bonds, + "type_bonds": type_bonds, + "token_pad_mask": token_pad_mask, + "token_resolved_mask": token_resolved_mask, + "token_disto_mask": token_disto_mask, + "contact_conditioning": contact_conditioning, + "contact_threshold": contact_threshold, + "method_feature": method_feature, + "modified": modified, + "cyclic_period": cyclic_period, + "affinity_token_mask": affinity_token_mask, + "ref_pos": ref_pos, + "atom_resolved_mask": atom_resolved_mask, + "ref_atom_name_chars": ref_atom_name_chars, + "ref_element": ref_element, + "ref_charge": ref_charge, + "ref_chirality": ref_chirality, + "atom_backbone_feat": atom_backbone_feat, + "ref_space_uid": ref_space_uid, + "coords": coords, + "atom_pad_mask": atom_pad_mask, + "atom_to_token": atom_to_token, + "token_to_rep_atom": token_to_rep_atom, + "r_set_to_rep_atom": r_set_to_rep_atom, + "token_to_center_atom": token_to_center_atom, + "disto_target": disto_target, + "disto_coords_ensemble": disto_coords_ensemble, + "bfactor": bfactor, + "plddt": atom_plddt, + "frames_idx": frames_idx, + "frame_resolved_mask": frame_resolved_mask, + "msa": msa, + "msa_paired": msa_paired, + "deletion_value": deletion_value, + "has_deletion": has_deletion, + "deletion_mean": deletion_mean, + "profile": profile, + "msa_mask": msa_mask, + "template_restype": template_restype, + "template_frame_rot": template_frame_rot, + "template_frame_t": template_frame_t, + "template_cb": template_cb, + "template_ca": template_ca, + "template_mask_cb": template_mask_cb, + "template_mask_frame": template_mask_frame, + "template_mask": template_mask, + "query_to_template": query_to_template, + "visibility_ids": visibility_ids, + } + + return features, template # each feature shape is traced at its binding above diff --git a/fastplms/boltz/minimal_structures.py b/src/fastplms/models/boltz/minimal_structures.py similarity index 61% rename from fastplms/boltz/minimal_structures.py rename to src/fastplms/models/boltz/minimal_structures.py index 5246fbe..666972a 100644 --- a/fastplms/boltz/minimal_structures.py +++ b/src/fastplms/models/boltz/minimal_structures.py @@ -1,15 +1,14 @@ from dataclasses import dataclass -from typing import List @dataclass class ProteinStructureTemplate: sequence: str - residue_names: List[str] - atom_names: List[str] - atom_elements: List[str] - atom_residue_index: List[int] - atom_chain_id: List[str] + residue_names: list[str] + atom_names: list[str] + atom_elements: list[str] + atom_residue_index: list[int] + atom_chain_id: list[str] @property def num_atoms(self) -> int: diff --git a/fastplms/boltz/modeling_boltz2.py b/src/fastplms/models/boltz/modeling_boltz2.py similarity index 51% rename from fastplms/boltz/modeling_boltz2.py rename to src/fastplms/models/boltz/modeling_boltz2.py index 83f4d75..bf1829a 100644 --- a/fastplms/boltz/modeling_boltz2.py +++ b/src/fastplms/models/boltz/modeling_boltz2.py @@ -1,54 +1,70 @@ -import entrypoint_setup import copy import inspect -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union - +import random +import numpy as np import torch -import torch._dynamo import torch.nn as nn +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from typing import Any, ClassVar from torch import Tensor -from transformers import PreTrainedModel, PretrainedConfig +from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import ModelOutput +from . import vb_const as const +from . import vb_layers_initialize as init from .cif_writer import write_cif from .minimal_featurizer import build_boltz2_features from .minimal_structures import ProteinStructureTemplate from .vb_const import bond_types as _vb_const_bond_types # noqa: F401 from .vb_layers_attention import AttentionPairBias as _vb_layers_attention_marker # noqa: F401 from .vb_layers_attentionv2 import AttentionPairBias as _vb_layers_attentionv2_marker # noqa: F401 -from .vb_layers_confidence_utils import compute_ptms as _vb_layers_confidence_utils_marker # noqa: F401 +from .vb_layers_confidence_utils import ( + compute_ptms as _vb_layers_confidence_utils_marker, # noqa: F401 +) from .vb_layers_dropout import get_dropout_mask as _vb_layers_dropout_marker # noqa: F401 from .vb_layers_initialize import gating_init_ as _vb_layers_initialize_marker # noqa: F401 -from .vb_layers_outer_product_mean import OuterProductMean as _vb_layers_outer_product_mean_marker # noqa: F401 -from .vb_layers_pair_averaging import PairWeightedAveraging as _vb_layers_pair_averaging_marker # noqa: F401 +from .vb_layers_outer_product_mean import ( + OuterProductMean as _vb_layers_outer_product_mean_marker, # noqa: F401 +) +from .vb_layers_pair_averaging import ( + PairWeightedAveraging as _vb_layers_pair_averaging_marker, # noqa: F401 +) +from .vb_layers_pairformer import PairformerModule from .vb_layers_transition import Transition as _vb_layers_transition_marker # noqa: F401 -from .vb_layers_triangular_mult import TriangleMultiplicationIncoming as _vb_layers_triangular_mult_marker # noqa: F401 +from .vb_layers_triangular_mult import ( + TriangleMultiplicationIncoming as _vb_layers_triangular_mult_marker, # noqa: F401 +) from .vb_loss_diffusionv2 import weighted_rigid_align as _vb_loss_diffusionv2_marker # noqa: F401 -from .vb_modules_transformersv2 import DiffusionTransformer as _vb_modules_transformersv2_marker # noqa: F401 -from .vb_modules_utils import LinearNoBias as _vb_modules_utils_marker # noqa: F401 -from .vb_potentials_potentials import get_potentials as _vb_potentials_potentials_marker # noqa: F401 -from .vb_potentials_schedules import ParameterSchedule as _vb_potentials_schedules_marker # noqa: F401 -from .vb_tri_attn_attention import TriangleAttentionStartingNode as _vb_tri_attn_attention_marker # noqa: F401 -from .vb_tri_attn_primitives import Attention as _vb_tri_attn_primitives_marker # noqa: F401 -from .vb_tri_attn_utils import permute_final_dims as _vb_tri_attn_utils_marker # noqa: F401 -from . import vb_const as const -from . import vb_layers_initialize as init -from .vb_layers_pairformer import PairformerModule from .vb_modules_confidencev2 import ConfidenceModule from .vb_modules_diffusion_conditioning import DiffusionConditioning from .vb_modules_diffusionv2 import AtomDiffusion, DiffusionModule from .vb_modules_encodersv2 import RelativePositionEncoder +from .vb_modules_transformersv2 import ( + DiffusionTransformer as _vb_modules_transformersv2_marker, # noqa: F401 +) from .vb_modules_trunkv2 import ( ContactConditioning, DistogramModule, InputEmbedder, MSAModule, ) +from .vb_modules_utils import LinearNoBias as _vb_modules_utils_marker # noqa: F401 +from .vb_potentials_potentials import ( + get_potentials as _vb_potentials_potentials_marker, # noqa: F401 +) +from .vb_potentials_schedules import ( + ParameterSchedule as _vb_potentials_schedules_marker, # noqa: F401 +) +from .vb_tri_attn_attention import ( + TriangleAttentionStartingNode as _vb_tri_attn_attention_marker, # noqa: F401 +) +from .vb_tri_attn_primitives import Attention as _vb_tri_attn_primitives_marker # noqa: F401 +from .vb_tri_attn_utils import permute_final_dims as _vb_tri_attn_utils_marker # noqa: F401 -def _default_steering_args() -> Dict[str, Any]: +def _default_steering_args() -> dict[str, Any]: return { "fk_steering": False, "num_particles": 3, @@ -60,7 +76,7 @@ def _default_steering_args() -> Dict[str, Any]: } -def _boltz2_reference_diffusion_overrides() -> Dict[str, Any]: +def _boltz2_reference_diffusion_overrides() -> dict[str, Any]: # Match Boltz2 CLI inference defaults from boltz.main/Boltz2DiffusionParams. return { "gamma_0": 0.8, @@ -79,24 +95,25 @@ def _boltz2_reference_diffusion_overrides() -> Dict[str, Any]: } -def _enforce_pairformer_v2(pairformer_args: Mapping[str, Any], context: str) -> Dict[str, Any]: - assert isinstance(pairformer_args, Mapping), ( - f"Expected {context} pairformer_args to be a dictionary." - ) +def _enforce_pairformer_v2(pairformer_args: Mapping[str, Any], context: str) -> dict[str, Any]: + if not isinstance(pairformer_args, Mapping): + raise TypeError(f"Expected {context} pairformer_args to be a dictionary.") out = _to_plain_python(copy.deepcopy(pairformer_args)) - if "v2" in out: - assert out["v2"], f"{context} pairformer_args['v2'] must be True for Boltz2." + if "v2" in out and not out["v2"]: + raise ValueError(f"{context} pairformer_args['v2'] must be True for Boltz2.") out["v2"] = True return out -def _require_key(mapping: Dict[str, Any], key: str) -> Any: - assert key in mapping, f"Missing required key '{key}' in checkpoint hyperparameters." +def _require_key(mapping: dict[str, Any], key: str) -> Any: + if key not in mapping: + raise KeyError(f"Missing required key '{key}' in checkpoint hyperparameters.") return mapping[key] -def _state_dict_without_wrappers(state_dict: Dict[str, Tensor]) -> Dict[str, Tensor]: - cleaned: Dict[str, Tensor] = {} +def _state_dict_without_wrappers(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + # Every state tensor has its parameter-defined shape (...). + cleaned: dict[str, Tensor] = {} for key, value in state_dict.items(): if key.startswith("ema."): continue @@ -105,15 +122,15 @@ def _state_dict_without_wrappers(state_dict: Dict[str, Tensor]) -> Dict[str, Ten new_key = new_key[len("model.") :] if new_key.startswith("module."): new_key = new_key[len("module.") :] - cleaned[new_key] = value + cleaned[new_key] = value # (...) return cleaned def _to_cpu_detached(value: Any) -> Any: if torch.is_tensor(value): - return value.detach().cpu() + return value.detach().cpu() # same shape (...) if isinstance(value, dict): - out: Dict[Any, Any] = {} + out: dict[Any, Any] = {} for key, nested_value in value.items(): out[key] = _to_cpu_detached(nested_value) return out @@ -124,9 +141,50 @@ def _to_cpu_detached(value: Any) -> Any: return value +@dataclass(frozen=True) +class _RandomState: + python: object + numpy: tuple[Any, ...] + torch_cpu: Tensor # (n_cpu_rng_state,) + torch_cuda: list[Tensor] | None # each: (n_cuda_rng_state,) + + +@contextmanager +def _seed_context(seed: int | None) -> Iterator[None]: + """Temporarily seed every RNG used by Boltz2 and restore caller state.""" + + if seed is None: + yield + return + if isinstance(seed, bool) or not isinstance(seed, int): + raise TypeError( + "Boltz2 seed must be an int or None; bool and coerced values are not accepted." + ) + state = _RandomState( + python=random.getstate(), + numpy=np.random.get_state(), + torch_cpu=torch.random.get_rng_state(), + torch_cuda=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, + ) + normalized_seed = seed % (2**32) + random.seed(normalized_seed) + np.random.seed(normalized_seed) + torch.manual_seed(normalized_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(normalized_seed) + try: + yield + finally: + random.setstate(state.python) + np.random.set_state(state.numpy) + torch.random.set_rng_state(state.torch_cpu) + if state.torch_cuda is not None: + torch.cuda.set_rng_state_all(state.torch_cuda) + + def _to_plain_python(value: Any) -> Any: if isinstance(value, Mapping): - out: Dict[Any, Any] = {} + out: dict[Any, Any] = {} for key, nested_value in value.items(): out[key] = _to_plain_python(nested_value) return out @@ -139,29 +197,62 @@ def _to_plain_python(value: Any) -> Any: return value -def _filtered_kwargs(target: Any, kwargs: Dict[str, Any]) -> Dict[str, Any]: +def _filtered_kwargs(target: Any, kwargs: dict[str, Any]) -> dict[str, Any]: signature = inspect.signature(target.__init__) allowed = set(signature.parameters.keys()) allowed.discard("self") - filtered: Dict[str, Any] = {} + filtered: dict[str, Any] = {} for key, value in kwargs.items(): if key in allowed: filtered[key] = value return filtered +@dataclass +class Boltz2ModelOutput(ModelOutput): + """Raw Boltz2 inference output with standard AutoModel controls.""" + + last_hidden_state: torch.Tensor | None = None # (b, t, d_s) or None + hidden_states: tuple[torch.Tensor, ...] | None = None # each shape traced at source + attentions: tuple[torch.Tensor, ...] | None = None # unsupported; always None + pdistogram: torch.Tensor | None = None # (b, t, t, n_distogram, n_bin) or None + s: torch.Tensor | None = None # (b, t, d_s) or None + z: torch.Tensor | None = None # (b, t, t, d_z) or None + sample_atom_coords: torch.Tensor | None = None # (b * m, a, 3) or None + diff_token_repr: torch.Tensor | None = None # (...) or None + s_conf: torch.Tensor | None = None # (...) or None + z_conf: torch.Tensor | None = None # (...) or None + pde_logits: torch.Tensor | None = None # (...) or None + plddt_logits: torch.Tensor | None = None # (...) or None + resolved_logits: torch.Tensor | None = None # (...) or None + pde: torch.Tensor | None = None # (...) or None + plddt: torch.Tensor | None = None # (...) or None + complex_plddt: torch.Tensor | None = None # (...) or None + complex_iplddt: torch.Tensor | None = None # (...) or None + complex_pde: torch.Tensor | None = None # (...) or None + complex_ipde: torch.Tensor | None = None # (...) or None + pae_logits: torch.Tensor | None = None # (...) or None + pae: torch.Tensor | None = None # (...) or None + ptm: torch.Tensor | None = None # (...) or None + iptm: torch.Tensor | None = None # (...) or None + ligand_iptm: torch.Tensor | None = None # (...) or None + protein_iptm: torch.Tensor | None = None # (...) or None + pair_chains_iptm: Any = None + + @dataclass class Boltz2StructureOutput(ModelOutput): - sample_atom_coords: Optional[torch.Tensor] = None - atom_pad_mask: Optional[torch.Tensor] = None - plddt: Optional[torch.Tensor] = None - confidence_score: Optional[torch.Tensor] = None - complex_plddt: Optional[torch.Tensor] = None - iptm: Optional[torch.Tensor] = None - ptm: Optional[torch.Tensor] = None - sequence: Optional[str] = None - structure_template: Optional[ProteinStructureTemplate] = None - raw_output: Optional[Dict[str, torch.Tensor]] = None + sample_atom_coords: torch.Tensor | None = None # (m, a_p, 3) or None + atom_pad_mask: torch.Tensor | None = None # (a_p,) or None + plddt: torch.Tensor | None = None # (m, a_p) or None + confidence_score: torch.Tensor | None = None # (m,) or None + complex_plddt: torch.Tensor | None = None # (m,) or None + iptm: torch.Tensor | None = None # (m,) or None + ptm: torch.Tensor | None = None # (m,) or None + sequence: str | None = None + structure_template: ProteinStructureTemplate | None = None + raw_output: dict[str, torch.Tensor] | None = None # values keep model-output shapes + seed: int | None = None class Boltz2Config(PretrainedConfig): @@ -169,7 +260,7 @@ class Boltz2Config(PretrainedConfig): def __init__( self, - core_kwargs: Optional[Dict[str, Any]] = None, + core_kwargs: dict[str, Any] | None = None, num_bins: int = 64, default_recycling_steps: int = 3, default_sampling_steps: int = 200, @@ -188,13 +279,14 @@ def __init__( @classmethod def from_hyperparameters( cls, - hparams: Dict[str, Any], + hparams: dict[str, Any], use_kernels: bool = False, - default_recycling_steps: Optional[int] = None, - default_sampling_steps: Optional[int] = None, - default_diffusion_samples: Optional[int] = None, + default_recycling_steps: int | None = None, + default_sampling_steps: int | None = None, + default_diffusion_samples: int | None = None, ) -> "Boltz2Config": - assert isinstance(hparams, dict), "Expected checkpoint hyperparameters as a dictionary." + if not isinstance(hparams, dict): + raise TypeError("Expected checkpoint hyperparameters as a dictionary.") required = [ "atom_s", "atom_z", @@ -214,14 +306,12 @@ def from_hyperparameters( hparams["pairformer_args"], context="checkpoint", ) - diffusion_process_args = _to_plain_python( - copy.deepcopy(hparams["diffusion_process_args"]) - ) + diffusion_process_args = _to_plain_python(copy.deepcopy(hparams["diffusion_process_args"])) diffusion_overrides = _boltz2_reference_diffusion_overrides() for key in diffusion_overrides: diffusion_process_args[key] = diffusion_overrides[key] - core_kwargs: Dict[str, Any] = { + core_kwargs: dict[str, Any] = { "atom_s": hparams["atom_s"], "atom_z": hparams["atom_z"], "token_s": hparams["token_s"], @@ -248,98 +338,34 @@ def from_hyperparameters( else: core_kwargs["confidence_model_args"] = None - if "confidence_prediction" in hparams: - core_kwargs["confidence_prediction"] = hparams["confidence_prediction"] - else: - core_kwargs["confidence_prediction"] = True - - if "token_level_confidence" in hparams: - core_kwargs["token_level_confidence"] = hparams["token_level_confidence"] - else: - core_kwargs["token_level_confidence"] = True - - if "alpha_pae" in hparams: - core_kwargs["alpha_pae"] = hparams["alpha_pae"] - else: - core_kwargs["alpha_pae"] = 0.0 - - if "atoms_per_window_queries" in hparams: - core_kwargs["atoms_per_window_queries"] = hparams["atoms_per_window_queries"] - else: - core_kwargs["atoms_per_window_queries"] = 32 - - if "atoms_per_window_keys" in hparams: - core_kwargs["atoms_per_window_keys"] = hparams["atoms_per_window_keys"] - else: - core_kwargs["atoms_per_window_keys"] = 128 - - if "atom_feature_dim" in hparams: - core_kwargs["atom_feature_dim"] = hparams["atom_feature_dim"] - else: - core_kwargs["atom_feature_dim"] = 128 - - if "bond_type_feature" in hparams: - core_kwargs["bond_type_feature"] = hparams["bond_type_feature"] - else: - core_kwargs["bond_type_feature"] = False - - if "run_trunk_and_structure" in hparams: - core_kwargs["run_trunk_and_structure"] = hparams["run_trunk_and_structure"] - else: - core_kwargs["run_trunk_and_structure"] = True - - if "skip_run_structure" in hparams: - core_kwargs["skip_run_structure"] = hparams["skip_run_structure"] - else: - core_kwargs["skip_run_structure"] = False - - if "fix_sym_check" in hparams: - core_kwargs["fix_sym_check"] = hparams["fix_sym_check"] - else: - core_kwargs["fix_sym_check"] = False - - if "cyclic_pos_enc" in hparams: - core_kwargs["cyclic_pos_enc"] = hparams["cyclic_pos_enc"] - else: - core_kwargs["cyclic_pos_enc"] = False - - if "use_no_atom_char" in hparams: - core_kwargs["use_no_atom_char"] = hparams["use_no_atom_char"] - else: - core_kwargs["use_no_atom_char"] = False - - if "use_atom_backbone_feat" in hparams: - core_kwargs["use_atom_backbone_feat"] = hparams["use_atom_backbone_feat"] - else: - core_kwargs["use_atom_backbone_feat"] = False - - if "use_residue_feats_atoms" in hparams: - core_kwargs["use_residue_feats_atoms"] = hparams["use_residue_feats_atoms"] - else: - core_kwargs["use_residue_feats_atoms"] = False - - if "conditioning_cutoff_min" in hparams: - core_kwargs["conditioning_cutoff_min"] = hparams["conditioning_cutoff_min"] - else: - core_kwargs["conditioning_cutoff_min"] = 4.0 - - if "conditioning_cutoff_max" in hparams: - core_kwargs["conditioning_cutoff_max"] = hparams["conditioning_cutoff_max"] - else: - core_kwargs["conditioning_cutoff_max"] = 20.0 + core_kwargs["confidence_prediction"] = hparams.get("confidence_prediction", True) + core_kwargs["token_level_confidence"] = hparams.get("token_level_confidence", True) + core_kwargs["alpha_pae"] = hparams.get("alpha_pae", 0.0) + core_kwargs["atoms_per_window_queries"] = hparams.get("atoms_per_window_queries", 32) + core_kwargs["atoms_per_window_keys"] = hparams.get("atoms_per_window_keys", 128) + core_kwargs["atom_feature_dim"] = hparams.get("atom_feature_dim", 128) + core_kwargs["bond_type_feature"] = hparams.get("bond_type_feature", False) + core_kwargs["run_trunk_and_structure"] = hparams.get("run_trunk_and_structure", True) + core_kwargs["skip_run_structure"] = hparams.get("skip_run_structure", False) + core_kwargs["fix_sym_check"] = hparams.get("fix_sym_check", False) + core_kwargs["cyclic_pos_enc"] = hparams.get("cyclic_pos_enc", False) + core_kwargs["use_no_atom_char"] = hparams.get("use_no_atom_char", False) + core_kwargs["use_atom_backbone_feat"] = hparams.get("use_atom_backbone_feat", False) + core_kwargs["use_residue_feats_atoms"] = hparams.get("use_residue_feats_atoms", False) + core_kwargs["conditioning_cutoff_min"] = hparams.get("conditioning_cutoff_min", 4.0) + core_kwargs["conditioning_cutoff_max"] = hparams.get("conditioning_cutoff_max", 20.0) if "steering_args" in hparams and hparams["steering_args"] is not None: - core_kwargs["steering_args"] = _to_plain_python( - copy.deepcopy(hparams["steering_args"]) - ) + core_kwargs["steering_args"] = _to_plain_python(copy.deepcopy(hparams["steering_args"])) else: core_kwargs["steering_args"] = _default_steering_args() if "validation_args" in hparams: validation_args = hparams["validation_args"] - assert isinstance(validation_args, Mapping), ( - "Expected 'validation_args' in checkpoint hyperparameters to be a mapping." - ) + if not isinstance(validation_args, Mapping): + raise TypeError( + "Expected 'validation_args' in checkpoint hyperparameters to be a mapping." + ) if default_recycling_steps is None and "recycling_steps" in validation_args: default_recycling_steps = validation_args["recycling_steps"] if default_sampling_steps is None and "sampling_steps" in validation_args: @@ -371,12 +397,12 @@ def __init__( token_s: int, token_z: int, num_bins: int, - embedder_args: Dict[str, Any], - msa_args: Dict[str, Any], - pairformer_args: Dict[str, Any], - score_model_args: Dict[str, Any], - diffusion_process_args: Dict[str, Any], - confidence_model_args: Optional[Dict[str, Any]] = None, + embedder_args: dict[str, Any], + msa_args: dict[str, Any], + pairformer_args: dict[str, Any], + score_model_args: dict[str, Any], + diffusion_process_args: dict[str, Any], + confidence_model_args: dict[str, Any] | None = None, atom_feature_dim: int = 128, confidence_prediction: bool = True, token_level_confidence: bool = True, @@ -394,7 +420,7 @@ def __init__( conditioning_cutoff_min: float = 4.0, conditioning_cutoff_max: float = 20.0, use_kernels: bool = False, - steering_args: Optional[Dict[str, Any]] = None, + steering_args: dict[str, Any] | None = None, ) -> None: super().__init__() self.use_kernels = use_kernels @@ -404,9 +430,13 @@ def __init__( self.run_trunk_and_structure = run_trunk_and_structure self.skip_run_structure = skip_run_structure self.bond_type_feature = bond_type_feature - self.steering_args = steering_args if steering_args is not None else _default_steering_args() - assert "v2" in pairformer_args, "Boltz2 requires pairformer_args['v2']." - assert pairformer_args["v2"], "Boltz2 requires pairformer_args['v2']=True." + self.steering_args = ( + steering_args if steering_args is not None else _default_steering_args() + ) + if "v2" not in pairformer_args: + raise ValueError("Boltz2 requires pairformer_args['v2'].") + if not pairformer_args["v2"]: + raise ValueError("Boltz2 requires pairformer_args['v2']=True.") full_embedder_args = { "atom_s": atom_s, @@ -449,17 +479,18 @@ def __init__( init.gating_init_(self.s_recycle.weight) init.gating_init_(self.z_recycle.weight) - torch._dynamo.config.cache_size_limit = 512 # noqa: SLF001 - torch._dynamo.config.accumulated_cache_size_limit = 512 # noqa: SLF001 - - msa_kwargs = _filtered_kwargs(MSAModule, {"token_z": token_z, "token_s": token_s, **msa_args}) + msa_kwargs = _filtered_kwargs( + MSAModule, + {"token_z": token_z, "token_s": token_s, **msa_args}, + ) self.msa_module = MSAModule(**msa_kwargs) pairformer_kwargs = _filtered_kwargs( PairformerModule, {"token_s": token_s, "token_z": token_z, **pairformer_args}, ) - assert "token_s" in pairformer_kwargs and "token_z" in pairformer_kwargs + if "token_s" not in pairformer_kwargs or "token_z" not in pairformer_kwargs: + raise RuntimeError("Boltz2 pairformer construction lost its token dimensions.") pairformer_token_s = pairformer_kwargs.pop("token_s") pairformer_token_z = pairformer_kwargs.pop("token_z") self.pairformer_module = PairformerModule( @@ -514,9 +545,10 @@ def __init__( self.distogram_module = DistogramModule(token_z, num_bins) if self.confidence_prediction: - assert confidence_model_args is not None, ( - "confidence_prediction=True requires confidence_model_args in config." - ) + if confidence_model_args is None: + raise ValueError( + "confidence_prediction=True requires confidence_model_args in config." + ) confidence_kwargs = { "token_s": token_s, "token_z": token_z, @@ -533,50 +565,58 @@ def __init__( def forward( self, - feats: Dict[str, Tensor], + feats: dict[str, Tensor], recycling_steps: int = 3, - num_sampling_steps: Optional[int] = None, + num_sampling_steps: int | None = None, diffusion_samples: int = 1, - max_parallel_samples: Optional[int] = None, + max_parallel_samples: int | None = None, run_confidence_sequentially: bool = True, detach_confidence: bool = True, - ) -> Dict[str, Tensor]: - s_inputs = self.input_embedder(feats) - s_init = self.s_init(s_inputs) - - z_init = self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :] - relative_position_encoding = self.rel_pos(feats) - z_init = z_init + relative_position_encoding - z_init = z_init + self.token_bonds(feats["token_bonds"].float()) + ) -> dict[str, Tensor]: + # b is the batch size, t the token count, a the padded atom count, + # d_s the token width, d_z the pair width, and m the diffusion multiplicity. + s_inputs = self.input_embedder(feats) # (b, t, d_s) + s_init = self.s_init(s_inputs) # (b, t, d_s) + + z_init = ( + self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :] + ) # (b, t, t, d_z) + relative_position_encoding = self.rel_pos(feats) # (b, t, t, d_z) + z_init = z_init + relative_position_encoding # (b, t, t, d_z) + z_init = z_init + self.token_bonds( + feats["token_bonds"].float() + ) # (b, t, t, d_z) if self.bond_type_feature: - z_init = z_init + self.token_bonds_type(feats["type_bonds"].long()) - z_init = z_init + self.contact_conditioning(feats) + z_init = z_init + self.token_bonds_type( + feats["type_bonds"].long() + ) # (b, t, t, d_z) + z_init = z_init + self.contact_conditioning(feats) # (b, t, t, d_z) - s = torch.zeros_like(s_init) - z = torch.zeros_like(z_init) - mask = feats["token_pad_mask"].float() - pair_mask = mask[:, :, None] * mask[:, None, :] + s = torch.zeros_like(s_init) # (b, t, d_s) + z = torch.zeros_like(z_init) # (b, t, t, d_z) + mask = feats["token_pad_mask"].float() # (b, t) + pair_mask = mask[:, :, None] * mask[:, None, :] # (b, t, t) if self.run_trunk_and_structure: for _ in range(recycling_steps + 1): - s = s_init + self.s_recycle(self.s_norm(s)) - z = z_init + self.z_recycle(self.z_norm(z)) + s = s_init + self.s_recycle(self.s_norm(s)) # (b, t, d_s) + z = z_init + self.z_recycle(self.z_norm(z)) # (b, t, t, d_z) z = z + self.msa_module( z, s_inputs, feats, use_kernels=self.use_kernels, - ) + ) # (b, t, t, d_z) s, z = self.pairformer_module( s, z, mask=mask, pair_mask=pair_mask, use_kernels=self.use_kernels, - ) + ) # (b, t, d_s), (b, t, t, d_z) - pdistogram = self.distogram_module(z) - output: Dict[str, Tensor] = { + pdistogram = self.distogram_module(z) # (b, t, t, n_distogram, n_bin) + output: dict[str, Tensor] = { "pdistogram": pdistogram, "s": s, "z": z, @@ -590,7 +630,7 @@ def forward( relative_position_encoding=relative_position_encoding, feats=feats, ) - ) + ) # q/c: (b, a, d_a); biases retain their attention-specific axes diffusion_conditioning = { "q": q, "c": c, @@ -610,30 +650,31 @@ def forward( max_parallel_samples=max_parallel_samples, steering_args=self.steering_args, diffusion_conditioning=diffusion_conditioning, - ) + ) # tensor values include sample_atom_coords: (b * m, a, 3) output.update(struct_out) if self.confidence_prediction: if self.skip_run_structure: - x_pred = feats["coords"].repeat_interleave(diffusion_samples, 0) + x_pred = feats["coords"].repeat_interleave( + diffusion_samples, 0 + ) # (b * m, ..., a, 3) else: - assert "sample_atom_coords" in output, ( - "Structure sampling did not produce sample_atom_coords." - ) - x_pred = output["sample_atom_coords"] + if "sample_atom_coords" not in output: + raise RuntimeError("Structure sampling did not produce sample_atom_coords.") + x_pred = output["sample_atom_coords"] # (b * m, a, 3) if detach_confidence: - s_inputs_c = s_inputs.detach() - s_c = s.detach() - z_c = z.detach() - x_pred_c = x_pred.detach() - pdist_c = output["pdistogram"][:, :, :, 0].detach() + s_inputs_c = s_inputs.detach() # (b, t, d_s) + s_c = s.detach() # (b, t, d_s) + z_c = z.detach() # (b, t, t, d_z) + x_pred_c = x_pred.detach() # (b * m, ..., a, 3) + pdist_c = output["pdistogram"][:, :, :, 0].detach() # (b, t, t, n_bin) else: - s_inputs_c = s_inputs - s_c = s - z_c = z - x_pred_c = x_pred - pdist_c = output["pdistogram"][:, :, :, 0] + s_inputs_c = s_inputs # (b, t, d_s) + s_c = s # (b, t, d_s) + z_c = z # (b, t, t, d_z) + x_pred_c = x_pred # (b * m, ..., a, 3) + pdist_c = output["pdistogram"][:, :, :, 0] # (b, t, t, n_bin) output.update( self.confidence_module( @@ -649,41 +690,43 @@ def forward( ) ) - return output + return output # named tensors retain the shapes traced above class Boltz2Model(PreTrainedModel): config_class = Boltz2Config base_model_prefix = "core" - all_tied_weights_keys = {} + all_tied_weights_keys: ClassVar[dict[str, str]] = {} def __init__(self, config: Boltz2Config) -> None: super().__init__(config) - assert isinstance(config.core_kwargs, dict), "config.core_kwargs must be a dictionary." + if not isinstance(config.core_kwargs, dict): + raise TypeError("config.core_kwargs must be a dictionary.") self.core = Boltz2InferenceCore(**config.core_kwargs) - def _init_weights(self, module: nn.Module) -> None: # noqa: ARG002 + def _init_weights(self, module: nn.Module) -> None: return - def _detied_state_dict(self) -> Dict[str, Tensor]: + def _detied_state_dict(self) -> dict[str, Tensor]: + # Every state tensor has its parameter-defined shape (...). raw_state = self.state_dict() - seen_ptrs: Dict[int, str] = {} - out: Dict[str, Tensor] = {} + seen_ptrs: dict[int, str] = {} + out: dict[str, Tensor] = {} for key, tensor in raw_state.items(): if torch.is_tensor(tensor): ptr = tensor.untyped_storage().data_ptr() if ptr in seen_ptrs: - out[key] = tensor.clone() + out[key] = tensor.clone() # (...) else: seen_ptrs[ptr] = key - out[key] = tensor + out[key] = tensor # (...) else: - out[key] = tensor - return out + out[key] = tensor # (...) + return out # each tensor: (...) def save_pretrained(self, save_directory: str, **kwargs: Any) -> None: if "safe_serialization" not in kwargs: - kwargs["safe_serialization"] = False + kwargs["safe_serialization"] = True if "state_dict" not in kwargs: kwargs["state_dict"] = self._detied_state_dict() super().save_pretrained(save_directory, **kwargs) @@ -696,26 +739,47 @@ def device(self) -> torch.device: def from_boltz_checkpoint( cls, checkpoint_path: str, - map_location: Union[str, torch.device] = "cpu", + map_location: str | torch.device = "cpu", use_kernels: bool = False, - default_recycling_steps: Optional[int] = None, - default_sampling_steps: Optional[int] = None, - default_diffusion_samples: Optional[int] = None, + default_recycling_steps: int | None = None, + default_sampling_steps: int | None = None, + default_diffusion_samples: int | None = None, + *, + allow_unsafe_pickle: bool = False, ) -> "Boltz2Model": - # Boltz Lightning checkpoints include OmegaConf objects and require full unpickling. + """Import an official Boltz Lightning checkpoint into the inference model. + + Lightning checkpoints require Python pickle deserialization. The caller + must opt in with the literal value ``True`` after independently verifying + that the checkpoint comes from a trusted source and matches its expected + cryptographic hash. + """ + + if allow_unsafe_pickle is not True: + raise ValueError( + "Boltz Lightning checkpoints contain Python pickle and cannot be loaded " + "safely. Pass allow_unsafe_pickle=True only for a trusted, hash-verified " + "checkpoint." + ) + + # The official Lightning checkpoints contain OmegaConf objects, so this explicit + # trusted-checkpoint boundary requires full Python unpickling. checkpoint = torch.load( checkpoint_path, map_location=map_location, weights_only=False, - ) - assert isinstance(checkpoint, dict), "Checkpoint must deserialize to a dictionary." + ) # mapping values include state tensors with parameter-defined shapes (...) + if not isinstance(checkpoint, dict): + raise TypeError("Checkpoint must deserialize to a dictionary.") _require_key(checkpoint, "hyper_parameters") _require_key(checkpoint, "state_dict") hparams = checkpoint["hyper_parameters"] - assert isinstance(hparams, dict), "Checkpoint hyper_parameters must be a dictionary." - state_dict = checkpoint["state_dict"] - assert isinstance(state_dict, dict), "Checkpoint state_dict must be a dictionary." + if not isinstance(hparams, dict): + raise TypeError("Checkpoint hyper_parameters must be a dictionary.") + state_dict = checkpoint["state_dict"] # each tensor: (...) + if not isinstance(state_dict, dict): + raise TypeError("Checkpoint state_dict must be a dictionary.") config = Boltz2Config.from_hyperparameters( hparams, @@ -725,51 +789,76 @@ def from_boltz_checkpoint( default_diffusion_samples=default_diffusion_samples, ) model = cls(config) - cleaned = _state_dict_without_wrappers(state_dict) + cleaned = _state_dict_without_wrappers(state_dict) # each tensor: (...) target_keys = set(model.core.state_dict().keys()) for key in target_keys: - assert ".attention.norm_s." not in key, ( - "Boltz2 inference core unexpectedly uses v1 attention parameters. " - "Expected pairformer v2 architecture." - ) - filtered: Dict[str, Tensor] = {} + if ".attention.norm_s." in key: + raise RuntimeError( + "Boltz2 inference core unexpectedly uses v1 attention parameters. " + "Expected pairformer v2 architecture." + ) + filtered: dict[str, Tensor] = {} for key, value in cleaned.items(): if key in target_keys: - filtered[key] = value + filtered[key] = value # (...) missing = sorted(target_keys.difference(filtered.keys())) - assert len(missing) == 0, ( - "Checkpoint is missing required parameters for Boltz2 inference core. " - f"Missing keys (first 20): {missing[:20]}" - ) + if missing: + raise RuntimeError( + "Checkpoint is missing required parameters for Boltz2 inference core. " + f"Missing keys (first 20): {missing[:20]}" + ) load_result = model.core.load_state_dict(filtered, strict=False) loaded_missing = sorted(load_result.missing_keys) - assert len(loaded_missing) == 0, ( - "Model has unexpected missing keys after load_state_dict. " - f"Missing keys (first 20): {loaded_missing[:20]}" - ) - assert len(load_result.unexpected_keys) == 0 + if loaded_missing: + raise RuntimeError( + "Model has unexpected missing keys after load_state_dict. " + f"Missing keys (first 20): {loaded_missing[:20]}" + ) + if load_result.unexpected_keys: + raise RuntimeError( + "Model has unexpected checkpoint keys after load_state_dict: " + f"{sorted(load_result.unexpected_keys)[:20]}" + ) model.eval() return model def forward( self, - feats: Dict[str, Tensor], - recycling_steps: Optional[int] = None, - num_sampling_steps: Optional[int] = None, - diffusion_samples: Optional[int] = None, - max_parallel_samples: Optional[int] = None, + feats: dict[str, Tensor], + recycling_steps: int | None = None, + num_sampling_steps: int | None = None, + diffusion_samples: int | None = None, + max_parallel_samples: int | None = None, run_confidence_sequentially: bool = True, detach_confidence: bool = True, - ) -> Dict[str, Tensor]: + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> Boltz2ModelOutput | tuple[Any, ...]: + # Feature shapes follow build_boltz2_features; b and t are batch and token counts. + output_attentions = ( + self.config.output_attentions if output_attentions is None else output_attentions + ) + if output_attentions: + raise NotImplementedError( + "Boltz2 does not expose normalized attention tensors from its structure " + "modules. output_attentions=True is unsupported." + ) + output_hidden_states = ( + self.config.output_hidden_states + if output_hidden_states is None + else output_hidden_states + ) + return_dict = self.config.use_return_dict if return_dict is None else return_dict if recycling_steps is None: recycling_steps = self.config.default_recycling_steps if num_sampling_steps is None: num_sampling_steps = self.config.default_sampling_steps if diffusion_samples is None: diffusion_samples = self.config.default_diffusion_samples - return self.core( + raw_output = self.core( feats=feats, recycling_steps=recycling_steps, num_sampling_steps=num_sampling_steps, @@ -777,72 +866,123 @@ def forward( max_parallel_samples=max_parallel_samples, run_confidence_sequentially=run_confidence_sequentially, detach_confidence=detach_confidence, + ) # named tensors use the Boltz2InferenceCore.forward shapes + token_state = raw_output.get("s") # (b, t, d_s) or None + pair_state = raw_output.get("z") # (b, t, t, d_z) or None + model_output = Boltz2ModelOutput( + last_hidden_state=token_state, + hidden_states=(token_state, pair_state) + if output_hidden_states and token_state is not None and pair_state is not None + else None, + **raw_output, ) + return model_output if return_dict else model_output.to_tuple() def _to_model_device( self, - feats: Dict[str, Tensor], - float_dtype: torch.dtype, - ) -> Dict[str, Tensor]: - moved: Dict[str, Tensor] = {} + feats: dict[str, Tensor], + float_dtype: torch.dtype = torch.float32, + ) -> dict[str, Tensor]: + # Each feature tensor keeps its existing shape (...). + moved: dict[str, Tensor] = {} for key, value in feats.items(): if torch.is_tensor(value): if value.is_floating_point(): - moved[key] = value.to(device=self.device, dtype=float_dtype) + moved[key] = value.to( + device=self.device, dtype=float_dtype + ) # same shape (...) else: - moved[key] = value.to(device=self.device) + moved[key] = value.to(device=self.device) # same shape (...) else: - moved[key] = value - return moved + moved[key] = value # same shape (...) + return moved # each tensor: same shape (...) def predict_structure( self, amino_acid_sequence: str, - recycling_steps: Optional[int] = None, - num_sampling_steps: Optional[int] = None, - diffusion_samples: Optional[int] = None, - max_parallel_samples: Optional[int] = None, + recycling_steps: int | None = None, + num_sampling_steps: int | None = None, + diffusion_samples: int | None = None, + max_parallel_samples: int | None = None, run_confidence_sequentially: bool = True, - float_dtype: Optional[torch.dtype] = None, + float_dtype: torch.dtype = torch.float32, + seed: int | None = None, ) -> Boltz2StructureOutput: - if float_dtype is None: - float_dtype = torch.float32 - - feats, template = build_boltz2_features( - amino_acid_sequence=amino_acid_sequence, - num_bins=self.config.num_bins, - atoms_per_window_queries=self.core.input_embedder.atom_encoder.atoms_per_window_queries, - ) - feats = self._to_model_device(feats, float_dtype=float_dtype) - - with torch.no_grad(): - output = self.forward( - feats=feats, - recycling_steps=recycling_steps, - num_sampling_steps=num_sampling_steps, - diffusion_samples=diffusion_samples, - max_parallel_samples=max_parallel_samples, - run_confidence_sequentially=run_confidence_sequentially, + if float_dtype != torch.float32: + raise ValueError( + "Boltz2 predict_structure requires FP32 features and parameter storage; " + "CUDA compute runs under the declared BF16 autocast policy." + ) + parameter_dtypes = { + parameter.dtype for parameter in self.parameters() if parameter.is_floating_point() + } + if parameter_dtypes and parameter_dtypes != {torch.float32}: + raise ValueError( + "Boltz2 predict_structure requires FP32 parameter storage before CUDA " + f"BF16 autocast; found {sorted(map(str, parameter_dtypes))}." ) - sample_atom_coords = output["sample_atom_coords"].detach().cpu() - non_finite_mask = torch.logical_not(torch.isfinite(sample_atom_coords)) - assert not torch.any(non_finite_mask), ( - "sample_atom_coords contains non-finite values. " - f"Non-finite count: {int(non_finite_mask.sum().item())}" - ) - atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu() - plddt = output["plddt"].detach().cpu() if "plddt" in output else None - complex_plddt = output["complex_plddt"].detach().cpu() if "complex_plddt" in output else None - iptm = output["iptm"].detach().cpu() if "iptm" in output else None - ptm = output["ptm"].detach().cpu() if "ptm" in output else None - - confidence_score = None + with _seed_context(seed): + feats, template = build_boltz2_features( + amino_acid_sequence=amino_acid_sequence, + num_bins=self.config.num_bins, + atoms_per_window_queries=( + self.core.input_embedder.atom_encoder.atoms_per_window_queries + ), + ) # feature shapes are traced in build_boltz2_features + feats = self._to_model_device( + feats, float_dtype=torch.float32 + ) # unchanged feature shapes + autocast_context = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if self.device.type == "cuda" + else nullcontext() + ) + with torch.no_grad(), autocast_context: + output = self.forward( + feats=feats, + recycling_steps=recycling_steps, + num_sampling_steps=num_sampling_steps, + diffusion_samples=diffusion_samples, + max_parallel_samples=max_parallel_samples, + run_confidence_sequentially=run_confidence_sequentially, + return_dict=True, + ) # named tensors use Boltz2InferenceCore.forward shapes + + sample_atom_coords_value = output.get( + "sample_atom_coords" + ) # (m, a_p, 3) or absent + if not isinstance(sample_atom_coords_value, torch.Tensor): + raise RuntimeError("Boltz2 structure sampling did not return coordinate tensors.") + sample_atom_coords = sample_atom_coords_value.detach().cpu() # (m, a_p, 3) + non_finite_mask = torch.logical_not( + torch.isfinite(sample_atom_coords) + ) # (m, a_p, 3) + if torch.any(non_finite_mask): + raise RuntimeError( + "sample_atom_coords contains non-finite values. " + f"Non-finite count: {int(non_finite_mask.sum().item())}" + ) + atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu() # (a_p,) + plddt = ( + output["plddt"].detach().cpu() if "plddt" in output else None + ) # (m, a_p) or None + complex_plddt = ( + output["complex_plddt"].detach().cpu() if "complex_plddt" in output else None + ) # (m,) or None + iptm = ( + output["iptm"].detach().cpu() if "iptm" in output else None + ) # (m,) or None + ptm = ( + output["ptm"].detach().cpu() if "ptm" in output else None + ) # (m,) or None + + confidence_score = None # (m,) or None if (complex_plddt is not None) and (iptm is not None) and (ptm is not None): if torch.allclose(iptm, torch.zeros_like(iptm)): - confidence_score = (4 * complex_plddt + ptm) / 5 + confidence_score = (4 * complex_plddt + ptm) / 5 # (m,) else: - confidence_score = (4 * complex_plddt + iptm) / 5 + confidence_score = (4 * complex_plddt + iptm) / 5 # (m,) return Boltz2StructureOutput( sample_atom_coords=sample_atom_coords, @@ -855,6 +995,7 @@ def predict_structure( sequence=template.sequence, structure_template=template, raw_output={key: _to_cpu_detached(val) for key, val in output.items()}, + seed=seed, ) def save_as_cif( @@ -863,15 +1004,12 @@ def save_as_cif( output_path: str, sample_index: int = 0, ) -> str: - assert structure_output.structure_template is not None, ( - "structure_output.structure_template is required for CIF export." - ) - assert structure_output.sample_atom_coords is not None, ( - "structure_output.sample_atom_coords is required for CIF export." - ) - assert structure_output.atom_pad_mask is not None, ( - "structure_output.atom_pad_mask is required for CIF export." - ) + if structure_output.structure_template is None: + raise ValueError("structure_output.structure_template is required for CIF export.") + if structure_output.sample_atom_coords is None: + raise ValueError("structure_output.sample_atom_coords is required for CIF export.") + if structure_output.atom_pad_mask is None: + raise ValueError("structure_output.atom_pad_mask is required for CIF export.") return write_cif( structure_template=structure_output.structure_template, atom_coords=structure_output.sample_atom_coords, diff --git a/src/fastplms/models/boltz/vb_const.py b/src/fastplms/models/boltz/vb_const.py new file mode 100644 index 0000000..066bbef --- /dev/null +++ b/src/fastplms/models/boltz/vb_const.py @@ -0,0 +1,353 @@ +"""Minimal biological constants required by the local Boltz2 runtime. + +The upstream project contains additional training-time curation tables. They +are intentionally excluded here because FastPLMs neither trains Boltz2 nor +uses upstream data-processing code. Keeping only runtime inputs makes this +module auditable and prevents parity-oracle data from becoming a package +dependency. +""" + +from __future__ import annotations + + +def _index(values: list[str]) -> dict[str, int]: + return {value: index for index, value in enumerate(values)} + + +chain_types = ["PROTEIN", "DNA", "RNA", "NONPOLYMER"] +chain_type_ids = _index(chain_types) + + +canonical_tokens = [ + "ALA", + "ARG", + "ASN", + "ASP", + "CYS", + "GLN", + "GLU", + "GLY", + "HIS", + "ILE", + "LEU", + "LYS", + "MET", + "PHE", + "PRO", + "SER", + "THR", + "TRP", + "TYR", + "VAL", + "UNK", +] +tokens = [ + "", + "-", + *canonical_tokens, + "A", + "G", + "C", + "U", + "N", + "DA", + "DG", + "DC", + "DT", + "DN", +] +token_ids = _index(tokens) +num_tokens = len(tokens) + +prot_letter_to_token = dict( + zip( + "ARNDCEQGHILKMFPSTWYV", + [ + "ALA", + "ARG", + "ASN", + "ASP", + "CYS", + "GLU", + "GLN", + "GLY", + "HIS", + "ILE", + "LEU", + "LYS", + "MET", + "PHE", + "PRO", + "SER", + "THR", + "TRP", + "TYR", + "VAL", + ], + strict=True, + ) +) +prot_letter_to_token.update({letter: "UNK" for letter in "XJBZOU"} | {"-": "-"}) + + +def _parse_atom_rows(rows: str) -> dict[str, list[str]]: + table: dict[str, list[str]] = {} + for row in rows.strip().splitlines(): + residue, *atom_names = row.split() + table[residue] = atom_names + return table + + +ref_atoms = {"PAD": [], "-": []} +ref_atoms.update( + _parse_atom_rows( + """ +UNK N CA C O CB +ALA N CA C O CB +ARG N CA C O CB CG CD NE CZ NH1 NH2 +ASN N CA C O CB CG OD1 ND2 +ASP N CA C O CB CG OD1 OD2 +CYS N CA C O CB SG +GLN N CA C O CB CG CD OE1 NE2 +GLU N CA C O CB CG CD OE1 OE2 +GLY N CA C O +HIS N CA C O CB CG ND1 CD2 CE1 NE2 +ILE N CA C O CB CG1 CG2 CD1 +LEU N CA C O CB CG CD1 CD2 +LYS N CA C O CB CG CD CE NZ +MET N CA C O CB CG SD CE +PHE N CA C O CB CG CD1 CD2 CE1 CE2 CZ +PRO N CA C O CB CG CD +SER N CA C O CB OG +THR N CA C O CB OG1 CG2 +TRP N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2 +TYR N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH +VAL N CA C O CB CG1 CG2 +""" + ) +) + +protein_backbone_atom_names = ["N", "CA", "C", "O"] +nucleic_backbone_atom_names = [ + "P", + "OP1", + "OP2", + "O5'", + "C5'", + "C4'", + "O4'", + "C3'", + "O3'", + "C2'", + "O2'", + "C1'", +] +protein_backbone_atom_index = _index(protein_backbone_atom_names) +nucleic_backbone_atom_index = _index(nucleic_backbone_atom_names) + +_rna_backbone = nucleic_backbone_atom_names +_dna_backbone = [atom for atom in _rna_backbone if atom != "O2'"] +ref_atoms.update( + { + "A": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], + "G": [*_rna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], + "C": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], + "U": [*_rna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"], + "N": list(_rna_backbone), + "DA": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"], + "DG": [*_dna_backbone, "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"], + "DC": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"], + "DT": [*_dna_backbone, "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"], + "DN": list(_dna_backbone), + } +) + +_protein_tokens = ["UNK", *canonical_tokens[:-1]] +_nucleic_tokens = ["A", "G", "C", "U", "N", "DA", "DG", "DC", "DT", "DN"] +res_to_center_atom = { + **dict.fromkeys(_protein_tokens, "CA"), + **dict.fromkeys(_nucleic_tokens, "C1'"), +} +res_to_disto_atom = { + **dict.fromkeys(_protein_tokens, "CB"), + "GLY": "CA", + "A": "C4", + "G": "C4", + "C": "C2", + "U": "C2", + "N": "C1'", + "DA": "C4", + "DG": "C4", + "DC": "C2", + "DT": "C2", + "DN": "C1'", +} + + +num_elements = 128 +bond_types = ["OTHER", "SINGLE", "DOUBLE", "TRIPLE", "AROMATIC", "COVALENT"] +contact_conditioning_info = { + "UNSPECIFIED": 0, + "UNSELECTED": 1, + "POCKET>BINDER": 2, + "BINDER>POCKET": 3, + "CONTACT": 4, +} +chunk_size_threshold = 384 + +_method_groups = { + 0: ("MD",), + 1: ("X-RAY DIFFRACTION",), + 2: ("ELECTRON MICROSCOPY",), + 3: ("SOLUTION NMR",), + 4: ( + "SOLID-STATE NMR", + "NEUTRON DIFFRACTION", + "ELECTRON CRYSTALLOGRAPHY", + "FIBER DIFFRACTION", + "POWDER DIFFRACTION", + "INFRARED SPECTROSCOPY", + "FLUORESCENCE TRANSFER", + "EPR", + "THEORETICAL MODEL", + "SOLUTION SCATTERING", + "OTHER", + ), + 5: ("AFDB",), + 6: ("BOLTZ-1",), + 7: ("FUTURE1",), + 8: ("FUTURE2",), + 9: ("FUTURE3",), + 10: ("FUTURE4",), + 11: ("FUTURE5",), +} +method_types_ids = { + method.lower(): identifier + for identifier, methods in _method_groups.items() + for method in methods +} +num_method_types = len(_method_groups) + +vdw_radii = [ + float(radius) + for radius in [ + "1.2", + "1.4", + "2.2", + "1.9", + "1.8", + "1.7", + "1.6", + "1.55", + "1.5", + "1.54", + "2.4", + "2.2", + "2.1", + "2.1", + "1.95", + "1.8", + "1.8", + "1.88", + "2.8", + "2.4", + "2.3", + "2.15", + "2.05", + "2.05", + "2.05", + "2.05", + "2.0", + "2.0", + "2.0", + "2.1", + "2.1", + "2.1", + "2.05", + "1.9", + "1.9", + "2.02", + "2.9", + "2.55", + "2.4", + "2.3", + "2.15", + "2.1", + "2.05", + "2.05", + "2.0", + "2.05", + "2.1", + "2.2", + "2.2", + "2.25", + "2.2", + "2.1", + "2.1", + "2.16", + "3.0", + "2.7", + "2.5", + "2.48", + "2.47", + "2.45", + "2.43", + "2.42", + "2.4", + "2.38", + "2.37", + "2.35", + "2.33", + "2.32", + "2.3", + "2.28", + "2.27", + "2.25", + "2.2", + "2.1", + "2.05", + "2.0", + "2.0", + "2.05", + "2.1", + "2.05", + "2.2", + "2.3", + "2.3", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.4", + "2.0", + "2.3", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + "2.0", + ] +] diff --git a/src/fastplms/models/boltz/vb_layers_attention.py b/src/fastplms/models/boltz/vb_layers_attention.py new file mode 100644 index 0000000..8649d70 --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_attention.py @@ -0,0 +1,103 @@ +"""Pair-biased self-attention used in the Boltz2 diffusion stack.""" + +from __future__ import annotations + +from collections.abc import Callable, MutableMapping +from einops.layers.torch import Rearrange +from torch import Tensor, nn + +from . import vb_layers_initialize as init +from ._pair_attention import pair_biased_attention, reshape_heads + + +class AttentionPairBias(nn.Module): + """Attend over sequence states while adding a learned pair representation.""" + + def __init__( + self, + c_s: int, + c_z: int, + num_heads: int, + inf: float = 1e6, + initial_norm: bool = True, + ) -> None: + super().__init__() + if c_s % num_heads: + raise ValueError(f"c_s={c_s} must be divisible by num_heads={num_heads}") + + self.c_s = c_s + self.num_heads = num_heads + self.head_dim = c_s // num_heads + self.inf = inf + self.initial_norm = initial_norm + if initial_norm: + self.norm_s = nn.LayerNorm(c_s) + + self.proj_q = nn.Linear(c_s, c_s) + self.proj_k = nn.Linear(c_s, c_s, bias=False) + self.proj_v = nn.Linear(c_s, c_s, bias=False) + self.proj_g = nn.Linear(c_s, c_s, bias=False) + self.proj_z = nn.Sequential( + nn.LayerNorm(c_z), + nn.Linear(c_z, num_heads, bias=False), + Rearrange("b ... h -> b h ..."), + ) + self.proj_o = nn.Linear(c_s, c_s, bias=False) + init.final_init_(self.proj_o.weight) + + def _resolve_pair_bias( + self, + pair_states: Tensor, + model_cache: MutableMapping[str, Tensor] | None, + ) -> Tensor: + # pair_states: (b, l_q, l_k, d_z) + if model_cache is not None and "z" in model_cache: + return model_cache["z"] # (b, h, l_q, l_k) + pair_bias = self.proj_z(pair_states) # (b, h, l_q, l_k) + if model_cache is not None: + model_cache["z"] = pair_bias + return pair_bias # (b, h, l_q, l_k) + + def forward( + self, + s: Tensor, + z: Tensor, + mask: Tensor, + k_in: Tensor | None = None, + multiplicity: int = 1, + to_keys: Callable[[Tensor], Tensor] | None = None, + model_cache: MutableMapping[str, Tensor] | None = None, + ) -> Tensor: + """Transform S with shapes ``S: (b, l_q, d)`` and ``Z: (b, l, l, d_z)``.""" + + # s: (b, l_q, d); z: (b_z, l_q, l_k, d_z); mask: (b_k, l_k). + sequence_states = self.norm_s(s) if self.initial_norm else s # (b, l_q, d) + if to_keys is not None: + key_states = to_keys(sequence_states) # (b, l_k, d) + key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1) # (b_k, l_k) + else: + key_states = sequence_states if k_in is None else k_in # (b, l_k, d) + key_mask = mask # (b_k, l_k) + + query = reshape_heads( + self.proj_q(sequence_states), self.num_heads + ) # (b, l_q, h, d_h) + key = reshape_heads(self.proj_k(key_states), self.num_heads) # (b, l_k, h, d_h) + value = reshape_heads( + self.proj_v(key_states), self.num_heads + ) # (b, l_k, h, d_h) + pair_bias = self._resolve_pair_bias(z, model_cache).repeat_interleave( + multiplicity, + dim=0, + ) # (b, h, l_q, l_k) + attended = pair_biased_attention( + query, + key, + value, + pair_bias, + key_mask, + self.inf, + ) # (b, l_q, h, d_h) + attended = attended.reshape(s.shape[0], -1, self.c_s) # (b, l_q, d) + gate = self.proj_g(sequence_states).sigmoid() # (b, l_q, d) + return self.proj_o(gate * attended) # (b, l_q, d) diff --git a/src/fastplms/models/boltz/vb_layers_attentionv2.py b/src/fastplms/models/boltz/vb_layers_attentionv2.py new file mode 100644 index 0000000..2166de9 --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_attentionv2.py @@ -0,0 +1,76 @@ +"""Cross-attention variant of Boltz2 pair-biased attention.""" + +from __future__ import annotations + +from einops.layers.torch import Rearrange +from torch import Tensor, nn + +from . import vb_layers_initialize as init +from ._pair_attention import pair_biased_attention, reshape_heads + + +class AttentionPairBias(nn.Module): + """Attend from S to K while adding either learned or precomputed pair bias.""" + + def __init__( + self, + c_s: int, + c_z: int | None = None, + num_heads: int | None = None, + inf: float = 1e6, + compute_pair_bias: bool = True, + ) -> None: + super().__init__() + if num_heads is None or c_s % num_heads: + raise ValueError("num_heads must divide c_s") + if compute_pair_bias and c_z is None: + raise ValueError("c_z is required when pair bias is learned") + + self.c_s = c_s + self.num_heads = num_heads + self.head_dim = c_s // num_heads + self.inf = inf + self.proj_q = nn.Linear(c_s, c_s) + self.proj_k = nn.Linear(c_s, c_s, bias=False) + self.proj_v = nn.Linear(c_s, c_s, bias=False) + self.proj_g = nn.Linear(c_s, c_s, bias=False) + self.compute_pair_bias = compute_pair_bias + if compute_pair_bias: + self.proj_z = nn.Sequential( + nn.LayerNorm(c_z), + nn.Linear(c_z, num_heads, bias=False), + Rearrange("b ... h -> b h ..."), + ) + else: + self.proj_z = Rearrange("b ... h -> b h ...") + self.proj_o = nn.Linear(c_s, c_s, bias=False) + init.final_init_(self.proj_o.weight) + + def forward( + self, + s: Tensor, + z: Tensor, + mask: Tensor, + k_in: Tensor, + multiplicity: int = 1, + ) -> Tensor: + """Transform S with K and pair input Z, returning shape ``(b, l_q, d)``.""" + + # s: (b, l_q, d); k_in: (b, l_k, d); z: (b_z, l_q, l_k, d_z or h). + query = reshape_heads(self.proj_q(s), self.num_heads) # (b, l_q, h, d_h) + key = reshape_heads(self.proj_k(k_in), self.num_heads) # (b, l_k, h, d_h) + value = reshape_heads(self.proj_v(k_in), self.num_heads) # (b, l_k, h, d_h) + pair_bias = self.proj_z(z).repeat_interleave( + multiplicity, dim=0 + ) # (b, h, l_q, l_k) + attended = pair_biased_attention( + query, + key, + value, + pair_bias, + mask, + self.inf, + ) # (b, l_q, h, d_h) + attended = attended.reshape(s.shape[0], -1, self.c_s) # (b, l_q, d) + gate = self.proj_g(s).sigmoid() # (b, l_q, d) + return self.proj_o(gate * attended) # (b, l_q, d) diff --git a/src/fastplms/models/boltz/vb_layers_confidence_utils.py b/src/fastplms/models/boltz/vb_layers_confidence_utils.py new file mode 100644 index 0000000..dc48c34 --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_confidence_utils.py @@ -0,0 +1,272 @@ +"""Confidence geometry and aggregate metrics for Boltz2 outputs.""" + +from __future__ import annotations + +import torch +from torch import nn + +from . import vb_const as const + + +def compute_collinear_mask(v1: torch.Tensor, v2: torch.Tensor) -> torch.Tensor: + """Mark nondegenerate vector pairs whose angle defines a stable frame.""" + + # v1/v2: (n, 3) + norm1 = torch.norm(v1, dim=1, keepdim=True) # (n, 1) + norm2 = torch.norm(v2, dim=1, keepdim=True) # (n, 1) + unit1 = v1 / (norm1 + 1e-6) # (n, 3) + unit2 = v2 / (norm2 + 1e-6) # (n, 3) + separated = torch.abs(torch.sum(unit1 * unit2, dim=1)) < 0.9063 # (n,) + return ( + separated & (norm1.reshape(-1) > 1e-2) & (norm2.reshape(-1) > 1e-2) + ) # (n,) + + +def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + return torch.bmm( + feats["atom_to_token"].float(), + feats["asym_id"].unsqueeze(-1).float(), + ).squeeze(-1) # (b, a) + + +def _replace_nonpolymer_frames( + coordinates: torch.Tensor, + frame_indices: torch.Tensor, + feats: dict[str, torch.Tensor], + atom_chain_ids: torch.Tensor, + resolved_mask: torch.Tensor | None, + *, + inference: bool, +) -> None: + # coordinates: (b, m, a, 3); frame_indices: (b, m, t_or_a, 3). + token_chain_ids = feats["asym_id"] # (b, t) + for batch_index, sample_coordinates in enumerate(coordinates): + # sample_coordinates: (m, a, 3) + token_offset = 0 + atom_offset = 0 + for chain_id in torch.unique(token_chain_ids[batch_index]): + token_mask = (token_chain_ids[batch_index] == chain_id) * feats[ + "token_pad_mask" + ][batch_index] # (t,) + atom_mask = (atom_chain_ids[batch_index] == chain_id) * feats[ + "atom_pad_mask" + ][batch_index] # (a,) + token_count = int(token_mask.sum().item()) + atom_count = int(atom_mask.sum().item()) + is_nonpolymer = ( + feats["mol_type"][batch_index, token_offset] == const.chain_type_ids["NONPOLYMER"] + ) + if is_nonpolymer and atom_count >= 3: + chain_resolved = ( + feats["atom_pad_mask"][batch_index] + if inference + else ( + feats["atom_resolved_mask"][batch_index] + if resolved_mask is None + else resolved_mask[batch_index] + ) + ) + chain_atom_mask = atom_mask.bool() # (a,) + chain_coordinates = sample_coordinates[:, chain_atom_mask] # (m, a_c, 3) + differences = ( + chain_coordinates[:, None, :, :] - chain_coordinates[:, :, None, :] + ) # (m, a_c, a_c, 3) + distances = differences.square().sum(dim=-1) ** 0.5 # (m, a_c, a_c) + valid = chain_resolved[chain_atom_mask] # (a_c,) + invalid_pairs = 1 - (valid[None, :] * valid[:, None]).to( + torch.float32 + ) # (a_c, a_c) + invalid_pairs[invalid_pairs == 1] = torch.inf # (a_c, a_c) + nearest = torch.sort( + distances + invalid_pairs, dim=2 + ).indices # (m, a_c, a_c) + frames = ( + torch.cat( + (nearest[:, :, 1:2], nearest[:, :, 0:1], nearest[:, :, 2:3]), + dim=2, + ) + + atom_offset + ) # (m, a_c, 3) + frame_indices[ + batch_index, + :, + token_offset : token_offset + atom_count, + :, + ] = frames + token_offset += token_count + atom_offset += atom_count + + +def compute_frame_pred( + pred_atom_coords: torch.Tensor, + frames_idx_true: torch.Tensor, + feats: dict[str, torch.Tensor], + multiplicity: int, + resolved_mask: torch.Tensor | None = None, + inference: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Construct predicted local frames and their validity mask. + + ``pred_atom_coords`` has shape ``(b * multiplicity, a, 3)``. Polymer + frames retain the supplied atom indices; nonpolymer frames use each + atom's three nearest resolved neighbors. + """ + + atom_chain_ids = _atom_chain_ids(feats) # (b, a) + expanded_batch, _, _ = pred_atom_coords.shape + base_batch = expanded_batch // multiplicity + coordinates = pred_atom_coords.reshape( + base_batch, multiplicity, -1, 3 + ) # (b, m, a, 3) + frame_indices = frames_idx_true.clone().repeat_interleave( + multiplicity, 0 + ) # (b * m, 1, t, 3) or (b * m, t, 3) + frame_indices = frame_indices.reshape( + base_batch, multiplicity, -1, 3 + ) # (b, m, t, 3) + _replace_nonpolymer_frames( + coordinates, + frame_indices, + feats, + atom_chain_ids, + resolved_mask, + inference=inference, + ) + + batch_indices = torch.arange( + base_batch, device=frame_indices.device + )[:, None, None, None] # (b, 1, 1, 1) + sample_indices = torch.arange( + multiplicity, device=frame_indices.device + )[None, :, None, None] # (1, m, 1, 1) + frame_coordinates = coordinates[ + batch_indices, + sample_indices, + frame_indices, + ].reshape(-1, 3, 3) # (b * m * t, 3, 3) + valid_frames = compute_collinear_mask( + frame_coordinates[:, 1] - frame_coordinates[:, 0], + frame_coordinates[:, 1] - frame_coordinates[:, 2], + ).reshape(base_batch, multiplicity, -1) # (b, m, t) + return ( + frame_indices, + valid_frames * feats["token_pad_mask"][:, None, :], + ) # (b, m, t, 3), (b, m, t) + + +def compute_aggregated_metric(logits: torch.Tensor, end: float = 1.0) -> torch.Tensor: + """Convert categorical confidence logits to bin-center expectations.""" + + num_bins = logits.shape[-1] + bin_width = end / num_bins + centers = torch.arange( + 0.5 * bin_width, + end, + bin_width, + device=logits.device, + ) # (n_bin,) + # logits: (..., n_bin) + probabilities = nn.functional.softmax(logits, dim=-1) # (..., n_bin) + broadcast_shape = (1,) * (probabilities.ndim - 1) + centers.shape + return torch.sum( + probabilities * centers.view(broadcast_shape), dim=-1 + ) # (...) + + +def tm_function(d: torch.Tensor, n_res: torch.Tensor) -> torch.Tensor: + """Evaluate the TM-score distance kernel.""" + + # d and n_res broadcast to a shared shape (...). + distance_scale = 1.24 * (torch.clip(n_res, min=19) - 15) ** (1 / 3) - 1.8 # (...) + return 1 / (1 + (d / distance_scale) ** 2) # (...) + + +def _maximum_masked_tm( + expected_tm: torch.Tensor, + pair_mask: torch.Tensor, +) -> torch.Tensor: + # expected_tm/pair_mask: (b * m, t, t). + per_anchor = torch.sum(expected_tm * pair_mask, dim=-1) / ( + torch.sum(pair_mask, dim=-1) + 1e-5 + ) # (b * m, t) + return torch.max(per_anchor, dim=1).values # (b * m,) + + +def compute_ptms( + logits: torch.Tensor, + x_preds: torch.Tensor, + feats: dict[str, torch.Tensor], + multiplicity: int, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + dict[int, dict[int, torch.Tensor]], +]: + """Compute pTM, ipTM, interface-type, and chain-pair confidence scores.""" + + # logits: (b * m, t, t, n_bin); x_preds: (b * m, a, 3). + _, frame_mask = compute_frame_pred( + x_preds, + feats["frames_idx"], + feats, + multiplicity, + inference=True, + ) # frame_mask: (b, m, t) + token_mask = feats["token_pad_mask"].repeat_interleave( + multiplicity, 0 + ) # (b * m, t) + valid_anchor = frame_mask.reshape(-1, frame_mask.shape[-1]) # (b * m, t) + base_pair_mask = ( + valid_anchor[:, :, None] * token_mask[:, None, :] * token_mask[:, :, None] + ) # (b * m, t, t) + asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0) # (b * m, t) + interface_mask = base_pair_mask * ( + asym_id[:, None, :] != asym_id[:, :, None] + ) # (b * m, t, t) + + num_bins = logits.shape[-1] + pae_centers = torch.arange( + 0.5 * (32.0 / num_bins), + 32.0, + 32.0 / num_bins, + device=logits.device, + ).unsqueeze(0) # (1, n_bin) + n_res = token_mask.sum(dim=-1, keepdim=True) # (b * m, 1) + tm_values = tm_function(pae_centers, n_res).unsqueeze(1).unsqueeze( + 2 + ) # (b * m, 1, 1, n_bin) + expected_tm = torch.sum( + nn.functional.softmax(logits, dim=-1) * tm_values, dim=-1 + ) # (b * m, t, t) + + ptm = _maximum_masked_tm(expected_tm, base_pair_mask) # (b * m,) + iptm = _maximum_masked_tm(expected_tm, interface_mask) # (b * m,) + + token_type = feats["mol_type"].repeat_interleave(multiplicity, 0) # (b * m, t) + ligand = (token_type == const.chain_type_ids["NONPOLYMER"]).float() # (b * m, t) + protein = (token_type == const.chain_type_ids["PROTEIN"]).float() # (b * m, t) + ligand_protein = ( + ligand[:, :, None] * protein[:, None, :] + protein[:, :, None] * ligand[:, None, :] + ) # (b * m, t, t) + protein_protein = protein[:, :, None] * protein[:, None, :] # (b * m, t, t) + ligand_iptm = _maximum_masked_tm( + expected_tm, interface_mask * ligand_protein + ) # (b * m,) + protein_iptm = _maximum_masked_tm( + expected_tm, interface_mask * protein_protein + ) # (b * m,) + + chain_pair_iptm: dict[int, dict[int, torch.Tensor]] = {} + for first_chain in torch.unique(asym_id).tolist(): + scores: dict[int, torch.Tensor] = {} + for second_chain in torch.unique(asym_id).tolist(): + chain_mask = base_pair_mask + chain_mask = chain_mask * (asym_id[:, None, :] == first_chain) + chain_mask = chain_mask * (asym_id[:, :, None] == second_chain) + scores[second_chain] = _maximum_masked_tm(expected_tm, chain_mask) + chain_pair_iptm[first_chain] = scores + + return ptm, iptm, ligand_iptm, protein_iptm, chain_pair_iptm diff --git a/src/fastplms/models/boltz/vb_layers_dropout.py b/src/fastplms/models/boltz/vb_layers_dropout.py new file mode 100644 index 0000000..ac8c9ff --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_dropout.py @@ -0,0 +1,42 @@ +"""Broadcast dropout masks used by the Boltz pair stack. + +The mask is sampled in ``float32`` even when the pair tensor uses a lower +precision dtype. This matches the checkpoint implementation while keeping +the broadcast axis explicit. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + + +def _broadcast_mask_shape(Z: Tensor, *, columnwise: bool) -> tuple[int, ...]: + """Return the row-wise or column-wise pair-mask shape for ``Z``.""" + + if Z.ndim != 4: + raise ValueError(f"pair tensor must have four dimensions, got {Z.ndim}") + b, rows, columns, _ = Z.shape + return (b, 1, columns, 1) if columnwise else (b, rows, 1, 1) + + +def get_dropout_mask( + dropout: float, + z: Tensor, + training: bool, + columnwise: bool = False, +) -> Tensor: + """Sample an inverted-dropout mask that broadcasts over pair channels. + + ``Z`` is the pair tensor with shape ``(b, n, n, d)``. Row-wise masks + vary along the first residue axis; column-wise masks vary along the + second. Evaluation returns an all-one mask while retaining the same RNG + call pattern as training. + """ + + probability = float(dropout) if training else 0.0 + sample_shape = _broadcast_mask_shape(z, columnwise=columnwise) + keep = ( + torch.rand(sample_shape, dtype=torch.float32, device=z.device) >= probability + ) # (b, 1, n, 1) or (b, n, 1, 1) + return keep * (1.0 / (1.0 - probability)) # same broadcast-mask shape diff --git a/src/fastplms/models/boltz/vb_layers_initialize.py b/src/fastplms/models/boltz/vb_layers_initialize.py new file mode 100644 index 0000000..37c2b7b --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_initialize.py @@ -0,0 +1,109 @@ +"""Torch-only parameter initialization for the Boltz2 runtime.""" + +from __future__ import annotations + +import math +import torch +from typing import Literal +from torch import Tensor + + +FanMode = Literal["fan_in", "fan_out", "fan_avg"] + + +def _calculate_fan(shape: torch.Size | tuple[int, ...], fan: FanMode = "fan_in") -> float: + """Resolve the selected fan for a two-dimensional linear weight tensor.""" + + if len(shape) != 2: + raise ValueError(f"linear weights must be two-dimensional, received {tuple(shape)}") + fan_out, fan_in = shape + if fan == "fan_in": + return float(fan_in) + if fan == "fan_out": + return float(fan_out) + if fan == "fan_avg": + return (fan_in + fan_out) / 2 + raise ValueError(f"invalid fan mode: {fan!r}") + + +def trunc_normal_init_( + weights: Tensor, + scale: float = 1.0, + fan: FanMode = "fan_in", +) -> None: + """Fill W from a normal distribution truncated at two standard deviations.""" + + # weights: (d_out, d_in); initialization preserves this shape in place. + variance = scale / max(1.0, _calculate_fan(weights.shape, fan)) + std = math.sqrt(variance) + with torch.no_grad(): + torch.nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-2 * std, b=2 * std) + + +def lecun_normal_init_(weights: Tensor) -> None: + """Initialize W using fan-in-scaled truncated normal values.""" + + # weights: (d_out, d_in), mutated in place. + trunc_normal_init_(weights) + + +def he_normal_init_(weights: Tensor) -> None: + """Initialize W using twice the fan-in variance.""" + + trunc_normal_init_(weights, scale=2.0) + + +def glorot_uniform_init_(weights: Tensor) -> None: + """Initialize W with Xavier uniform values.""" + + torch.nn.init.xavier_uniform_(weights, gain=1.0) + + +def _fill_(tensor: Tensor, value: float) -> None: + # tensor: (...), mutated in place without changing shape. + with torch.no_grad(): + tensor.fill_(value) + + +def final_init_(weights: Tensor) -> None: + """Zero the final projection W.""" + + _fill_(weights, 0.0) + + +def gating_init_(weights: Tensor) -> None: + """Zero the gating projection W.""" + + _fill_(weights, 0.0) + + +def bias_init_zero_(bias: Tensor) -> None: + """Set the bias tensor to zero.""" + + _fill_(bias, 0.0) + + +def bias_init_one_(bias: Tensor) -> None: + """Set the bias tensor to one.""" + + _fill_(bias, 1.0) + + +def normal_init_(weights: Tensor) -> None: + """Initialize W with linear Kaiming-normal values.""" + + torch.nn.init.kaiming_normal_(weights, nonlinearity="linear") + + +def ipa_point_weights_init_(weights: Tensor) -> None: + """Initialize W so applying softplus yields one.""" + + _fill_(weights, 0.541324854612918) + # weights: (d_out, d_in), mutated in place. + # weights: (d_out, d_in), mutated in place. + # weights: (...), mutated in place. + # weights: (...), mutated in place. + # bias: (...), mutated in place. + # bias: (...), mutated in place. + # weights: (d_out, d_in), mutated in place. + # weights: (...), mutated in place. diff --git a/src/fastplms/models/boltz/vb_layers_outer_product_mean.py b/src/fastplms/models/boltz/vb_layers_outer_product_mean.py new file mode 100644 index 0000000..c2c5ef7 --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_outer_product_mean.py @@ -0,0 +1,103 @@ +"""Outer-product aggregation from MSA states to pair states.""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + +from . import vb_layers_initialize as init + + +class OuterProductMean(nn.Module): + """Aggregate M with shape ``(b, s, l, d_m)`` into pair tensor Z.""" + + def __init__(self, c_in: int, c_hidden: int, c_out: int) -> None: + super().__init__() + self.c_hidden = c_hidden + self.norm = nn.LayerNorm(c_in) + self.proj_a = nn.Linear(c_in, c_hidden, bias=False) + self.proj_b = nn.Linear(c_in, c_hidden, bias=False) + self.proj_o = nn.Linear(c_hidden * c_hidden, c_out) + init.final_init_(self.proj_o.weight) + init.final_init_(self.proj_o.bias) + + @staticmethod + def _pair_counts(mask: Tensor) -> Tensor: + """Count valid MSA observations for each residue pair.""" + + # mask: (b, s, l, 1) + counts: Tensor | None = None + for start in range(0, mask.shape[1], 64): + mask_slice = mask[:, start : start + 64] # (b, s_chunk, l, 1) + contribution = ( + mask_slice[:, :, None, :] * mask_slice[:, :, :, None] + ).sum(dim=1) # (b, l, l, 1) + counts = ( + contribution if counts is None else counts + contribution + ) # (b, l, l, 1) + if counts is None: + raise ValueError("MSA depth must be positive") + return counts.clamp(min=1) # (b, l, l, 1) + + @staticmethod + def _outer_product(left: Tensor, right: Tensor, counts: Tensor) -> Tensor: + # left/right: (b, s, l, d_h); counts: (b, l, l, 1). + outer = torch.einsum("bsic,bsjd->bijcd", left, right) # (b, l, l, d_h, d_h) + return outer.reshape(*outer.shape[:3], -1) / counts # (b, l, l, d_h**2) + + def _chunked_projection( + self, + left: Tensor, + right: Tensor, + counts: Tensor, + chunk_size: int, + target: Tensor, + ) -> Tensor: + # left/right: (b, s, l, d_h); counts: (b, l, l, 1). + output: Tensor | None = None + for start in range(0, self.c_hidden, chunk_size): + stop = min(start + chunk_size, self.c_hidden) + outer = self._outer_product( + left[..., start:stop], right, counts + ) # (b, l, l, d_chunk * d_h) + weight = self.proj_o.weight[ + :, + start * self.c_hidden : stop * self.c_hidden, + ] # (d_z, d_chunk * d_h) + contribution = outer.to(target) @ weight.T # (b, l, l, d_z) + output = ( + contribution if output is None else output + contribution + ) # (b, l, l, d_z) + if output is None: + raise ValueError("c_hidden must be positive") + return output + self.proj_o.bias # (b, l, l, d_z) + + def forward(self, m: Tensor, mask: Tensor, chunk_size: int | None = None) -> Tensor: + """Return pair tensor Z with shape ``(b, l, l, d_z)``.""" + + # m: (b, s, l, d_m); mask: (b, s, l). + expanded_mask = mask.unsqueeze(-1).to(m) # (b, s, l, 1) + normalized = self.norm(m) # (b, s, l, d_m) + left = self.proj_a(normalized) * expanded_mask # (b, s, l, d_h) + right = self.proj_b(normalized) * expanded_mask # (b, s, l, d_h) + + if chunk_size is not None and not self.training: + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + counts = self._pair_counts(expanded_mask) # (b, l, l, 1) + return self._chunked_projection( + left, + right, + counts, + chunk_size, + normalized, + ) + + pair_mask = ( + expanded_mask[:, :, None, :] * expanded_mask[:, :, :, None] + ) # (b, s, l, l, 1) + counts = pair_mask.sum(dim=1).clamp(min=1) # (b, l, l, 1) + outer = self._outer_product( + left.float(), right.float(), counts + ) # (b, l, l, d_h**2) + return self.proj_o(outer.to(normalized)) # (b, l, l, d_z) diff --git a/src/fastplms/models/boltz/vb_layers_pair_averaging.py b/src/fastplms/models/boltz/vb_layers_pair_averaging.py new file mode 100644 index 0000000..d3dc31e --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_pair_averaging.py @@ -0,0 +1,111 @@ +"""Pair-weighted aggregation of Boltz2 MSA states.""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + +from . import vb_layers_initialize as init + + +class PairWeightedAveraging(nn.Module): + """Update M with pair-derived attention weights while preserving its shape.""" + + def __init__( + self, + c_m: int, + c_z: int, + c_h: int, + num_heads: int, + inf: float = 1e6, + ) -> None: + super().__init__() + self.c_m = c_m + self.c_z = c_z + self.c_h = c_h + self.num_heads = num_heads + self.inf = inf + self.norm_m = nn.LayerNorm(c_m) + self.norm_z = nn.LayerNorm(c_z) + self.proj_m = nn.Linear(c_m, c_h * num_heads, bias=False) + self.proj_g = nn.Linear(c_m, c_h * num_heads, bias=False) + self.proj_z = nn.Linear(c_z, num_heads, bias=False) + self.proj_o = nn.Linear(c_h * num_heads, c_m, bias=False) + init.final_init_(self.proj_o.weight) + + def _attention_weights(self, pair_states: Tensor, mask: Tensor) -> Tensor: + # pair_states: (b, l, l, d_z); mask: (b, l). + logits = self.proj_z(pair_states).permute(0, 3, 1, 2) # (b, h, l, l) + logits = logits + (1 - mask[:, None]) * -self.inf # (b, h, l, l) + return torch.softmax(logits, dim=-1) # (b, h, l, l) + + def _all_heads(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor) -> Tensor: + # msa_states: (b, s, l, d_m); pair_states: (b, l, l, d_z). + values = self.proj_m(msa_states).reshape( + *msa_states.shape[:3], + self.num_heads, + self.c_h, + ) # (b, s, l, h, d_h) + values = values.permute(0, 3, 1, 2, 4) # (b, h, s, l, d_h) + weights = self._attention_weights(pair_states, mask) # (b, h, l, l) + gate = self.proj_g(msa_states).sigmoid() # (b, s, l, h * d_h) + attended = torch.einsum( + "bhij,bhsjd->bhsid", weights, values + ) # (b, h, s, l, d_h) + attended = attended.permute(0, 2, 3, 1, 4) # (b, s, l, h, d_h) + attended = attended.reshape( + *attended.shape[:3], self.num_heads * self.c_h + ) # (b, s, l, h * d_h) + return self.proj_o(gate * attended) # (b, s, l, d_m) + + def _one_head(self, msa_states: Tensor, pair_states: Tensor, mask: Tensor, head: int) -> Tensor: + # msa_states: (b, s, l, d_m); pair_states: (b, l, l, d_z). + start = head * self.c_h + stop = start + self.c_h + value = msa_states @ self.proj_m.weight[start:stop].T # (b, s, l, d_h) + value = value.reshape(*value.shape[:3], 1, self.c_h).permute( + 0, 3, 1, 2, 4 + ) # (b, 1, s, l, d_h) + logits = pair_states @ self.proj_z.weight[head : head + 1].T # (b, l, l, 1) + logits = logits.permute(0, 3, 1, 2) # (b, 1, l, l) + weights = torch.softmax( + logits + (1 - mask[:, None]) * -self.inf, dim=-1 + ) # (b, 1, l, l) + gate = ( + msa_states @ self.proj_g.weight[start:stop].T + ).sigmoid() # (b, s, l, d_h) + attended = torch.einsum( + "bhij,bhsjd->bhsid", weights, value + ) # (b, 1, s, l, d_h) + attended = attended.permute(0, 2, 3, 1, 4).reshape( + *msa_states.shape[:3], + self.c_h, + ) # (b, s, l, d_h) + return (gate * attended) @ self.proj_o.weight[:, start:stop].T # (b, s, l, d_m) + + def forward( + self, + m: Tensor, + z: Tensor, + mask: Tensor, + chunk_heads: bool = True, + ) -> Tensor: + """Return updated M with shape ``(b, s, l, d_m)``.""" + + # m: (b, s, l, d_m); z: (b, l, l, d_z); mask: (b, l). + msa_states = self.norm_m(m) # (b, s, l, d_m) + pair_states = self.norm_z(z) # (b, l, l, d_z) + if not chunk_heads or self.training: + return self._all_heads(msa_states, pair_states, mask) + + output: Tensor | None = None + for head in range(self.num_heads): + contribution = self._one_head( + msa_states, pair_states, mask, head + ) # (b, s, l, d_m) + output = ( + contribution if output is None else output + contribution + ) # (b, s, l, d_m) + if output is None: + raise ValueError("num_heads must be positive") + return output # (b, s, l, d_m) diff --git a/src/fastplms/models/boltz/vb_layers_pairformer.py b/src/fastplms/models/boltz/vb_layers_pairformer.py new file mode 100644 index 0000000..6e323dd --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_pairformer.py @@ -0,0 +1,349 @@ +"""Pairformer layers for joint or pair-only Boltz2 representations.""" + +from __future__ import annotations + +import torch +from typing import Any, cast +from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint + +from . import vb_const as const +from .vb_layers_attention import AttentionPairBias +from .vb_layers_attentionv2 import AttentionPairBias as AttentionPairBiasV2 +from .vb_layers_dropout import get_dropout_mask +from .vb_layers_transition import Transition +from .vb_layers_triangular_mult import ( + TriangleMultiplicationIncoming, + TriangleMultiplicationOutgoing, +) +from .vb_tri_attn_attention import ( + TriangleAttentionEndingNode, + TriangleAttentionStartingNode, +) + + +def _pair_update( + modules: Any, + pair_states: Tensor, + pair_mask: Tensor, + chunk_size: int | None, + use_kernels: bool, + use_cuequiv_mul: bool, + use_cuequiv_attn: bool, +) -> Tensor: + """Apply multiplicative, attention, and transition updates to pair tensor Z.""" + + # pair_states: (b, l, l, d_z); pair_mask: (b, l, l). + multiplication_kernel = use_kernels or use_cuequiv_mul + attention_kernel = use_kernels or use_cuequiv_attn + output = pair_states # (b, l, l, d_z) + for update in (modules.tri_mul_out, modules.tri_mul_in): + dropout = get_dropout_mask( + modules.dropout, output, modules.training + ) # broadcastable to (b, l, l, d_z) + output = output + dropout * update( + output, + mask=pair_mask, + use_kernels=multiplication_kernel, + ) # (b, l, l, d_z) + + dropout = get_dropout_mask( + modules.dropout, output, modules.training + ) # broadcastable to (b, l, l, d_z) + output = output + dropout * modules.tri_att_start( + output, + mask=pair_mask, + chunk_size=chunk_size, + use_kernels=attention_kernel, + ) # (b, l, l, d_z) + dropout = get_dropout_mask( + modules.dropout, + output, + modules.training, + columnwise=True, + ) # broadcastable to (b, l, l, d_z) + output = output + dropout * modules.tri_att_end( + output, + mask=pair_mask, + chunk_size=chunk_size, + use_kernels=attention_kernel, + ) # (b, l, l, d_z) + return cast(Tensor, output + modules.transition_z(output)) # (b, l, l, d_z) + + +def _triangle_chunk_size(pair_states: Tensor, training: bool) -> int | None: + if training: + return None + return 128 if pair_states.shape[1] > const.chunk_size_threshold else 512 + + +class PairformerLayer(nn.Module): + """Update sequence tensor S and pair tensor Z once.""" + + def __init__( + self, + token_s: int, + token_z: int, + num_heads: int = 16, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + v2: bool = False, + ) -> None: + super().__init__() + self.token_z = token_z + self.dropout = dropout + self.num_heads = num_heads + self.post_layer_norm = post_layer_norm + self.pre_norm_s = nn.LayerNorm(token_s) + attention_class = AttentionPairBiasV2 if v2 else AttentionPairBias + self.attention = attention_class(token_s, token_z, num_heads) + self.tri_mul_out = TriangleMultiplicationOutgoing(token_z) + self.tri_mul_in = TriangleMultiplicationIncoming(token_z) + self.tri_att_start = TriangleAttentionStartingNode( + token_z, + pairwise_head_width, + pairwise_num_heads, + inf=1e9, + ) + self.tri_att_end = TriangleAttentionEndingNode( + token_z, + pairwise_head_width, + pairwise_num_heads, + inf=1e9, + ) + self.transition_s = Transition(token_s, token_s * 4) + self.transition_z = Transition(token_z, token_z * 4) + self.s_post_norm = nn.LayerNorm(token_s) if post_layer_norm else nn.Identity() + + def forward( + self, + s: Tensor, + z: Tensor, + mask: Tensor, + pair_mask: Tensor, + chunk_size_tri_attn: int | None = None, + use_kernels: bool = False, + use_cuequiv_mul: bool = False, + use_cuequiv_attn: bool = False, + ) -> tuple[Tensor, Tensor]: + """Return updated ``S: (b, l, d_s)`` and ``Z: (b, l, l, d_z)``.""" + + pair_output = _pair_update( + self, + z, + pair_mask, + chunk_size_tri_attn, + use_kernels, + use_cuequiv_mul, + use_cuequiv_attn, + ) # (b, l, l, d_z) + with torch.autocast("cuda", enabled=False): + normalized = self.pre_norm_s(s.float()) # (b, l, d_s) + sequence_output = s.float() + self.attention( + s=normalized, + z=pair_output.float(), + mask=mask.float(), + k_in=normalized, + ) # (b, l, d_s) + sequence_output = ( + sequence_output + self.transition_s(sequence_output) + ) # (b, l, d_s) + sequence_output = cast( + Tensor, self.s_post_norm(sequence_output) + ) # (b, l, d_s) + return sequence_output, pair_output # (b, l, d_s), (b, l, l, d_z) + + +class PairformerModule(nn.Module): + """Run a stack of joint sequence and pair updates.""" + + def __init__( + self, + token_s: int, + token_z: int, + num_blocks: int, + num_heads: int = 16, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + activation_checkpointing: bool = False, + v2: bool = False, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__() + self.token_z = token_z + self.num_blocks = num_blocks + self.dropout = dropout + self.num_heads = num_heads + self.post_layer_norm = post_layer_norm + self.activation_checkpointing = activation_checkpointing + self.layers = nn.ModuleList( + [ + PairformerLayer( + token_s, + token_z, + num_heads, + dropout, + pairwise_head_width, + pairwise_num_heads, + post_layer_norm, + v2, + ) + for _ in range(num_blocks) + ] + ) + + def forward( + self, + s: Tensor, + z: Tensor, + mask: Tensor, + pair_mask: Tensor, + use_kernels: bool = False, + ) -> tuple[Tensor, Tensor]: + """Return S and Z after every configured pairformer block.""" + + # s: (b, l, d_s); z: (b, l, l, d_z). + chunk_size = _triangle_chunk_size(z, self.training) + sequence_output, pair_output = s, z # (b, l, d_s), (b, l, l, d_z) + for layer in self.layers: + if self.activation_checkpointing: + sequence_output, pair_output = checkpoint( + layer, + sequence_output, + pair_output, + mask, + pair_mask, + chunk_size, + use_kernels, + use_reentrant=False, + ) # (b, l, d_s), (b, l, l, d_z) + else: + sequence_output, pair_output = layer( + sequence_output, + pair_output, + mask, + pair_mask, + chunk_size, + use_kernels, + ) # (b, l, d_s), (b, l, l, d_z) + return sequence_output, pair_output # (b, l, d_s), (b, l, l, d_z) + + +class PairformerNoSeqLayer(nn.Module): + """Update pair tensor Z without a sequence track.""" + + def __init__( + self, + token_z: int, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + ) -> None: + super().__init__() + self.token_z = token_z + self.dropout = dropout + self.post_layer_norm = post_layer_norm + self.tri_mul_out = TriangleMultiplicationOutgoing(token_z) + self.tri_mul_in = TriangleMultiplicationIncoming(token_z) + self.tri_att_start = TriangleAttentionStartingNode( + token_z, + pairwise_head_width, + pairwise_num_heads, + inf=1e9, + ) + self.tri_att_end = TriangleAttentionEndingNode( + token_z, + pairwise_head_width, + pairwise_num_heads, + inf=1e9, + ) + self.transition_z = Transition(token_z, token_z * 4) + + def forward( + self, + z: Tensor, + pair_mask: Tensor, + chunk_size_tri_attn: int | None = None, + use_kernels: bool = False, + use_cuequiv_mul: bool = False, + use_cuequiv_attn: bool = False, + ) -> Tensor: + """Return updated pair tensor Z with shape ``(b, l, l, d_z)``.""" + + return _pair_update( + self, + z, + pair_mask, + chunk_size_tri_attn, + use_kernels, + use_cuequiv_mul, + use_cuequiv_attn, + ) + + +class PairformerNoSeqModule(nn.Module): + """Run a stack of pair-only updates.""" + + def __init__( + self, + token_z: int, + num_blocks: int, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + activation_checkpointing: bool = False, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__() + self.token_z = token_z + self.num_blocks = num_blocks + self.dropout = dropout + self.post_layer_norm = post_layer_norm + self.activation_checkpointing = activation_checkpointing + self.layers = nn.ModuleList( + [ + PairformerNoSeqLayer( + token_z, + dropout, + pairwise_head_width, + pairwise_num_heads, + post_layer_norm, + ) + for _ in range(num_blocks) + ] + ) + + def forward( + self, + z: Tensor, + pair_mask: Tensor, + use_kernels: bool = False, + ) -> Tensor: + """Return Z after every configured pair-only block.""" + + # z: (b, l, l, d_z); pair_mask: (b, l, l). + chunk_size = _triangle_chunk_size(z, self.training) + output = z # (b, l, l, d_z) + for layer in self.layers: + if self.activation_checkpointing: + output = checkpoint( + layer, + output, + pair_mask, + chunk_size, + use_kernels, + use_reentrant=False, + ) # (b, l, l, d_z) + else: + output = layer( + output, pair_mask, chunk_size, use_kernels + ) # (b, l, l, d_z) + return output # (b, l, l, d_z) diff --git a/src/fastplms/models/boltz/vb_layers_transition.py b/src/fastplms/models/boltz/vb_layers_transition.py new file mode 100644 index 0000000..7e35f3c --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_transition.py @@ -0,0 +1,78 @@ +"""Gated feed-forward transition used by Boltz2 pair representations.""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + +from . import vb_layers_initialize as init + + +class Transition(nn.Module): + """Apply a normalized SwiGLU-style transition. + + The checkpoint-facing module names remain ``norm`` and ``fc1`` through + ``fc3``. For an input tensor X with shape ``(..., d)``, the module returns + a tensor with shape ``(..., d_out)``. + """ + + def __init__( + self, + dim: int = 128, + hidden: int = 512, + out_dim: int | None = None, + ) -> None: + super().__init__() + output_dim = dim if out_dim is None else out_dim + self.norm = nn.LayerNorm(dim, eps=1e-5) + self.fc1 = nn.Linear(dim, hidden, bias=False) + self.fc2 = nn.Linear(dim, hidden, bias=False) + self.fc3 = nn.Linear(hidden, output_dim, bias=False) + self.silu = nn.SiLU() + self.hidden = hidden + + init.bias_init_one_(self.norm.weight) + init.bias_init_zero_(self.norm.bias) + init.lecun_normal_init_(self.fc1.weight) + init.lecun_normal_init_(self.fc2.weight) + init.final_init_(self.fc3.weight) + + def _project_hidden_slice(self, normalized: Tensor, start: int, stop: int) -> Tensor: + """Return one hidden-width contribution to the output projection.""" + + # normalized: (..., d); the selected hidden width is d_chunk. + gate = self.silu( + torch.matmul(normalized, self.fc1.weight[start:stop].T) + ) # (..., d_chunk) + value = torch.matmul( + normalized, self.fc2.weight[start:stop].T + ) # (..., d_chunk) + return torch.matmul( + gate * value, self.fc3.weight[:, start:stop].T + ) # (..., d_out) + + def forward(self, x: Tensor, chunk_size: int | None = None) -> Tensor: + """Transform X, optionally accumulating the hidden dimension in chunks.""" + + # X is the normalized input tensor with shape (..., d). + normalized = self.norm(x) # (..., d) + if chunk_size is None or self.training: + # H is the gated hidden tensor with shape (..., d_hidden). + hidden_states = self.silu(self.fc1(normalized)) * self.fc2(normalized) + return self.fc3(hidden_states) # (..., d_out) + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + + output: Tensor | None = None + for start in range(0, self.hidden, chunk_size): + contribution = self._project_hidden_slice( + normalized, + start, + min(start + chunk_size, self.hidden), + ) # (..., d_out) + output = ( + contribution if output is None else output + contribution + ) # (..., d_out) + if output is None: # hidden is normally positive; keep malformed configs explicit. + raise ValueError("hidden must be positive") + return output # (..., d_out) diff --git a/src/fastplms/models/boltz/vb_layers_triangular_mult.py b/src/fastplms/models/boltz/vb_layers_triangular_mult.py new file mode 100644 index 0000000..44779aa --- /dev/null +++ b/src/fastplms/models/boltz/vb_layers_triangular_mult.py @@ -0,0 +1,131 @@ +"""Incoming and outgoing triangular multiplicative pair updates.""" + +from __future__ import annotations + +import importlib +import torch +from importlib.util import find_spec +from typing import Literal +from torch import Tensor, nn + +from . import vb_layers_initialize as init + + +TriangleDirection = Literal["incoming", "outgoing"] + + +@torch.compiler.disable +def kernel_triangular_mult( + x: Tensor, + direction: TriangleDirection, + mask: Tensor, + norm_in_weight: Tensor, + norm_in_bias: Tensor, + p_in_weight: Tensor, + g_in_weight: Tensor, + norm_out_weight: Tensor, + norm_out_bias: Tensor, + p_out_weight: Tensor, + g_out_weight: Tensor, + eps: float, +) -> Tensor: + """Dispatch the optional cuEquivariance triangle primitive lazily.""" + + # x: (b, l, l, d); mask: (b, l, l); returned tensor: (b, l, l, d). + if ( + find_spec("cuequivariance_torch") is None + or find_spec("cuequivariance_ops_torch") is None + ): + raise RuntimeError( + "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 " + "cuequivariance_ops_torch runtime from the 'structure,cueq' extras." + ) + cueq = importlib.import_module("cuequivariance_torch") + return cueq.triangle_multiplicative_update( + x, + direction=direction, + mask=mask, + norm_in_weight=norm_in_weight, + norm_in_bias=norm_in_bias, + p_in_weight=p_in_weight, + g_in_weight=g_in_weight, + norm_out_weight=norm_out_weight, + norm_out_bias=norm_out_bias, + p_out_weight=p_out_weight, + g_out_weight=g_out_weight, + eps=eps, + ) + + +class _TriangleMultiplication(nn.Module): + direction: TriangleDirection + equation: str + + def __init__(self, dim: int, direction: TriangleDirection, equation: str) -> None: + super().__init__() + self.direction = direction + self.equation = equation + self.norm_in = nn.LayerNorm(dim, eps=1e-5) + self.p_in = nn.Linear(dim, 2 * dim, bias=False) + self.g_in = nn.Linear(dim, 2 * dim, bias=False) + self.norm_out = nn.LayerNorm(dim) + self.p_out = nn.Linear(dim, dim, bias=False) + self.g_out = nn.Linear(dim, dim, bias=False) + + init.bias_init_one_(self.norm_in.weight) + init.bias_init_zero_(self.norm_in.bias) + init.lecun_normal_init_(self.p_in.weight) + init.gating_init_(self.g_in.weight) + init.bias_init_one_(self.norm_out.weight) + init.bias_init_zero_(self.norm_out.bias) + init.final_init_(self.p_out.weight) + init.gating_init_(self.g_out.weight) + + def _kernel_forward(self, pair_states: Tensor, mask: Tensor) -> Tensor: + # pair_states: (b, l, l, d); mask: (b, l, l). + return kernel_triangular_mult( + pair_states, + direction=self.direction, + mask=mask, + norm_in_weight=self.norm_in.weight, + norm_in_bias=self.norm_in.bias, + p_in_weight=self.p_in.weight, + g_in_weight=self.g_in.weight, + norm_out_weight=self.norm_out.weight, + norm_out_bias=self.norm_out.bias, + p_out_weight=self.p_out.weight, + g_out_weight=self.g_out.weight, + eps=1e-5, + ) # (b, l, l, d) + + def forward(self, x: Tensor, mask: Tensor, use_kernels: bool = False) -> Tensor: + """Transform pair tensor X with shape ``(b, l, l, d)``.""" + + if use_kernels: + return self._kernel_forward(x, mask) + + # X_norm is the normalized pair tensor used by the output gate. + normalized = self.norm_in(x) # (b, l, l, d) + projected = ( + self.p_in(normalized) * self.g_in(normalized).sigmoid() + ) # (b, l, l, 2 * d) + projected = projected * mask.unsqueeze(-1) # (b, l, l, 2 * d) + left, right = torch.chunk(projected.float(), 2, dim=-1) # each: (b, l, l, d) + combined = torch.einsum(self.equation, left, right) # (b, l, l, d) + return ( + self.p_out(self.norm_out(combined)) * self.g_out(normalized).sigmoid() + ) # (b, l, l, d) + + +class TriangleMultiplicationOutgoing(_TriangleMultiplication): + """Aggregate pair paths that share their destination index.""" + + def __init__(self, dim: int = 128) -> None: + super().__init__(dim, direction="outgoing", equation="bikd,bjkd->bijd") + + +class TriangleMultiplicationIncoming(_TriangleMultiplication): + """Aggregate pair paths that share their source index.""" + + def __init__(self, dim: int = 128) -> None: + super().__init__(dim, direction="incoming", equation="bkid,bkjd->bijd") diff --git a/src/fastplms/models/boltz/vb_loss_diffusionv2.py b/src/fastplms/models/boltz/vb_loss_diffusionv2.py new file mode 100644 index 0000000..7788cd9 --- /dev/null +++ b/src/fastplms/models/boltz/vb_loss_diffusionv2.py @@ -0,0 +1,185 @@ +"""Geometry objectives shared by Boltz diffusion and steering code. + +The rigid-alignment mechanism is based on the Kabsch formulation used by +AlphaFold 3 implementations. The implementation is maintained locally and +does not import an upstream runtime package. +""" + +from __future__ import annotations + +import warnings +import torch +import torch.nn.functional as functional +from einops import einsum + + +def _weighted_centroid( + coordinates: torch.Tensor, + weights: torch.Tensor, +) -> torch.Tensor: + # coordinates: (..., n, 3); weights: (..., n, 1). + return (coordinates * weights).sum(dim=-2, keepdim=True) / weights.sum( + dim=-2, + keepdim=True, + ) # (..., 1, 3) + + +def _warn_if_alignment_is_ambiguous( + mask: torch.Tensor, + singular_values: torch.Tensor, + *, + num_points: int, + coordinate_dim: int, +) -> None: + if torch.any(mask.sum(dim=-1) < coordinate_dim + 1): + warnings.warn( + "The size of one of the point clouds is <= dim+1. " + "`WeightedRigidAlign` cannot return a unique rotation.", + RuntimeWarning, + stacklevel=3, + ) + if (singular_values.abs() <= 1e-15).any() and num_points >= coordinate_dim + 1: + warnings.warn( + "Excessively low rank of cross-correlation between aligned " + "point clouds. `WeightedRigidAlign` cannot return a unique rotation.", + RuntimeWarning, + stacklevel=3, + ) + + +def weighted_rigid_align( + true_coords: torch.Tensor, + pred_coords: torch.Tensor, + weights: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + """Align true coordinates to predicted coordinates with weighted Kabsch. + + ``true_coords`` and ``pred_coords`` have shape ``(..., n, 3)``. The + returned tensor is detached because alignment defines a fixed target for + the diffusion loss. + """ + + output_shape = torch.broadcast_shapes(true_coords.shape, pred_coords.shape) + *batch_shape, num_points, coordinate_dim = output_shape + point_weights = (mask * weights).unsqueeze(-1) # (..., n, 1) + + true_centroid = _weighted_centroid(true_coords, point_weights) # (..., 1, 3) + pred_centroid = _weighted_centroid(pred_coords, point_weights) # (..., 1, 3) + true_centered = true_coords - true_centroid # (..., n, 3) + pred_centered = pred_coords - pred_centroid # (..., n, 3) + + covariance = einsum( + point_weights * pred_centered, + true_centered, + "... n i, ... n j -> ... i j", + ) # (..., 3, 3) + original_dtype = covariance.dtype + covariance_fp32 = covariance.to(torch.float32) # (..., 3, 3) + left_vectors, singular_values, right_vectors_h = torch.linalg.svd( + covariance_fp32, + driver="gesvd" if covariance_fp32.is_cuda else None, + ) # left/right: (..., 3, 3); singular_values: (..., 3) + right_vectors = right_vectors_h.mH # (..., 3, 3) + _warn_if_alignment_is_ambiguous( + mask, + singular_values, + num_points=num_points, + coordinate_dim=coordinate_dim, + ) + + preliminary_rotation = torch.einsum( + "... i j, ... k j -> ... i k", + left_vectors, + right_vectors, + ).to(torch.float32) # (..., 3, 3) + orientation = torch.eye( + coordinate_dim, + dtype=covariance_fp32.dtype, + device=covariance.device, + )[None].repeat(*batch_shape, 1, 1) # (..., 3, 3) + orientation[..., -1, -1] = torch.det(preliminary_rotation) # (...) + rotation = einsum( + left_vectors, + orientation, + right_vectors, + "... i j, ... j k, ... l k -> ... i l", + ).to(original_dtype) # (..., 3, 3) + + aligned = ( + einsum(true_centered, rotation, "... n i, ... j i -> ... n j") + pred_centroid + ) # (..., n, 3) + aligned.detach_() + return aligned # (..., n, 3) + + +def _smooth_lddt_for_example( + pred_coords: torch.Tensor, + true_coords: torch.Tensor, + is_nucleotide: torch.Tensor, + coords_mask: torch.Tensor, + *, + nucleic_acid_cutoff: float, + other_cutoff: float, +) -> torch.Tensor: + # pred_coords/true_coords: (n, 3); is_nucleotide/coords_mask: (n,). + true_distances = torch.cdist(true_coords, true_coords) # (n, n) + nucleotide_rows = is_nucleotide.bool().unsqueeze(-1).expand_as( + true_distances + ) # (n, n) + pair_mask = torch.where( + nucleotide_rows, + true_distances < nucleic_acid_cutoff, + true_distances < other_cutoff, + ) # (n, n) + pair_mask &= ~torch.eye( + pred_coords.shape[0], + dtype=torch.bool, + device=pred_coords.device, + ) # (n, n) + coordinate_rows = coords_mask.bool() # (n,) + pair_mask &= coordinate_rows.unsqueeze(-1) # (n, n) + pair_mask &= coordinate_rows.unsqueeze(-2) # (n, n) + + pair_indices = pair_mask.nonzero() # (n_pair, 2) + true_pair_distances = true_distances[ + pair_indices[:, 0], pair_indices[:, 1] + ] # (n_pair,) + pred_pair_distances = functional.pairwise_distance( + pred_coords[pair_indices[:, 0]], + pred_coords[pair_indices[:, 1]], + ) # (n_pair,) + distance_error = torch.abs(true_pair_distances - pred_pair_distances) # (n_pair,) + smooth_agreement = ( + sum(torch.sigmoid(threshold - distance_error) for threshold in (0.5, 1.0, 2.0, 4.0)) / 4.0 + ) # (n_pair,) + return smooth_agreement.sum() / (pair_indices.shape[0] + 1e-5) # () + + +def smooth_lddt_loss( + pred_coords: torch.Tensor, + true_coords: torch.Tensor, + is_nucleotide: torch.Tensor, + coords_mask: torch.Tensor, + nucleic_acid_cutoff: float = 30.0, + other_cutoff: float = 15.0, + multiplicity: int = 1, +) -> torch.Tensor: + """Return one minus the smooth local-distance agreement. + + Coordinate tensors have shape ``(b, n, 3)``. Sequence-level masks may + be shared across repeated diffusion samples through ``multiplicity``. + """ + + agreements = [ + _smooth_lddt_for_example( + pred_coords[index], + true_coords[index], + is_nucleotide[index // multiplicity], + coords_mask[index // multiplicity], + nucleic_acid_cutoff=nucleic_acid_cutoff, + other_cutoff=other_cutoff, + ) + for index in range(true_coords.shape[0]) + ] # each: () + return 1.0 - torch.stack(agreements).mean(dim=0) # () diff --git a/fastplms/boltz/vb_modules_confidencev2.py b/src/fastplms/models/boltz/vb_modules_confidencev2.py similarity index 67% rename from fastplms/boltz/vb_modules_confidencev2.py rename to src/fastplms/models/boltz/vb_modules_confidencev2.py index da415e8..de693a7 100644 --- a/fastplms/boltz/vb_modules_confidencev2.py +++ b/src/fastplms/models/boltz/vb_modules_confidencev2.py @@ -1,6 +1,5 @@ import torch from torch import nn -from torch.nn.functional import pad from . import vb_const as const from . import vb_layers_initialize as init @@ -16,6 +15,75 @@ from .vb_modules_utils import LinearNoBias +def _token_slot_logits_to_atom_logits( + token_logits: torch.Tensor, + atom_to_token: torch.Tensor, + atom_pad_mask: torch.Tensor, + *, + multiplicity: int, +) -> torch.Tensor: + """Gather per-token atom-slot logits onto each example's atom table. + + ``token_logits`` is ordered as ``(batch * multiplicity, token, slot, + channel)``. ``atom_to_token`` determines both the owning token and the + within-token slot of every atom, so heterogeneous atom counts do not leak + across batch rows. + """ + + if token_logits.ndim != 4: + raise ValueError( + "token_logits must have shape (batch * multiplicity, token, slot, channel), " + f"got {tuple(token_logits.shape)}." + ) + if atom_to_token.ndim != 3: + raise ValueError( + "atom_to_token must have shape (batch, atom, token), " + f"got {tuple(atom_to_token.shape)}." + ) + if atom_pad_mask.shape != atom_to_token.shape[:2]: + raise ValueError( + "atom_pad_mask must match the batch and atom axes of atom_to_token; " + f"got {tuple(atom_pad_mask.shape)} and {tuple(atom_to_token.shape)}." + ) + if multiplicity < 1: + raise ValueError(f"multiplicity must be positive, got {multiplicity}.") + + batch_size, atom_count, token_count = atom_to_token.shape + if token_logits.shape[0] != batch_size * multiplicity: + raise ValueError( + "token_logits batch axis must equal batch * multiplicity; " + f"got {token_logits.shape[0]} and {batch_size} * {multiplicity}." + ) + if token_logits.shape[1] != token_count: + raise ValueError( + "token_logits and atom_to_token disagree on token count; " + f"got {token_logits.shape[1]} and {token_count}." + ) + + valid_atoms = atom_pad_mask.bool() + assignments = atom_to_token.bool() & valid_atoms.unsqueeze(-1) + token_index = assignments.to(dtype=torch.int64).argmax(dim=-1) + # Cumulative one-hot counts give each atom its ordinal within its owning + # token without assuming that atoms from different tokens are contiguous. + cumulative_slots = assignments.to(dtype=torch.int64).cumsum(dim=1) - 1 + slot_index = (cumulative_slots * assignments).sum(dim=-1) + + slots_per_token = token_logits.shape[2] + flattened_index = token_index * slots_per_token + slot_index + flattened_index = flattened_index.masked_fill(~valid_atoms, 0) + flattened_index = flattened_index.repeat_interleave(multiplicity, dim=0) + expanded_atom_mask = valid_atoms.repeat_interleave(multiplicity, dim=0) + + flattened_logits = token_logits.flatten(1, 2) + gather_index = flattened_index.unsqueeze(-1).expand( + -1, + atom_count, + flattened_logits.shape[-1], + ) + atom_logits = torch.gather(flattened_logits, dim=1, index=gather_index) + return atom_logits * expanded_atom_mask.unsqueeze(-1).to(dtype=atom_logits.dtype) + + class ConfidenceModule(nn.Module): """Algorithm 31""" @@ -32,7 +100,7 @@ def __init__( add_z_input_to_z=False, maximum_bond_distance=0, bond_type_feature=False, - confidence_args: dict = None, + confidence_args: dict | None = None, compile_pairformer=False, fix_sym_check=False, cyclic_pos_enc=False, @@ -43,10 +111,7 @@ def __init__( ): super().__init__() self.max_num_atoms_per_token = 23 - if "no_update_s" in pairformer_args: - self.no_update_s = pairformer_args["no_update_s"] - else: - self.no_update_s = False + self.no_update_s = pairformer_args.get("no_update_s", False) boundaries = torch.linspace(2, max_dist, num_dist_bins - 1) self.register_buffer("boundaries", boundaries) self.dist_bin_pairwise_embed = nn.Embedding(num_dist_bins, token_z) @@ -122,15 +187,30 @@ def forward( use_kernels: bool = False, ): if run_sequentially and multiplicity > 1: - assert z.shape[0] == 1, "Not supported with batch size > 1" + batch_size = z.shape[0] + expected_shape = (batch_size, multiplicity) + if x_pred.ndim >= 4 and x_pred.shape[:2] == expected_shape: + sample_coordinates = x_pred + elif x_pred.shape[0] == batch_size * multiplicity: + sample_coordinates = x_pred.reshape( + batch_size, + multiplicity, + *x_pred.shape[1:], + ) + else: + raise ValueError( + "Sequential confidence expected coordinates with leading shape " + f"{expected_shape} or {batch_size * multiplicity}, got " + f"{tuple(x_pred.shape)}." + ) out_dicts = [] for sample_idx in range(multiplicity): - out_dicts.append( # noqa: PERF401 + out_dicts.append( self.forward( s_inputs, s, z, - x_pred[sample_idx : sample_idx + 1], + sample_coordinates[:, sample_idx], feats, pred_distogram_logits, multiplicity=1, @@ -142,16 +222,21 @@ def forward( out_dict = {} for key in out_dicts[0]: if key != "pair_chains_iptm": - out_dict[key] = torch.cat([out[key] for out in out_dicts], dim=0) + values = [out[key] for out in out_dicts] + out_dict[key] = torch.stack(values, dim=1).flatten(0, 1) else: pair_chains_iptm = {} for chain_idx1 in out_dicts[0][key]: chains_iptm = {} for chain_idx2 in out_dicts[0][key][chain_idx1]: - chains_iptm[chain_idx2] = torch.cat( - [out[key][chain_idx1][chain_idx2] for out in out_dicts], - dim=0, - ) + values = [ + out[key][chain_idx1][chain_idx2] + for out in out_dicts + ] + chains_iptm[chain_idx2] = torch.stack( + values, + dim=1, + ).flatten(0, 1) pair_chains_iptm[chain_idx1] = chains_iptm out_dict[key] = pair_chains_iptm return out_dict @@ -192,10 +277,10 @@ def forward( token_to_rep_atom = feats["token_to_rep_atom"] token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0) if len(x_pred.shape) == 4: - B, mult, N, _ = x_pred.shape - x_pred = x_pred.reshape(B * mult, N, -1) + b, multiplicity, n, _ = x_pred.shape + x_pred = x_pred.reshape(b * multiplicity, n, -1) else: - BM, N, _ = x_pred.shape + _, n, _ = x_pred.shape x_pred_repr = torch.bmm(token_to_rep_atom.float(), x_pred) d = torch.cdist(x_pred_repr, x_pred_repr) distogram = (d.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() @@ -270,9 +355,7 @@ def __init__( self.to_plddt_logits = LinearNoBias( token_s, num_plddt_bins * self.max_num_atoms_per_token ) - self.to_resolved_logits = LinearNoBias( - token_s, 2 * self.max_num_atoms_per_token - ) + self.to_resolved_logits = LinearNoBias(token_s, 2 * self.max_num_atoms_per_token) def forward( self, @@ -294,9 +377,7 @@ def forward( pae_intra_logits = pae_intra_logits * is_same_chain.float().unsqueeze(-1) pae_inter_logits = self.to_pae_inter_logits(z) - pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze( - -1 - ) + pae_inter_logits = pae_inter_logits * is_different_chain.float().unsqueeze(-1) pae_logits = pae_inter_logits + pae_intra_logits else: @@ -307,9 +388,7 @@ def forward( pde_intra_logits = pde_intra_logits * is_same_chain.float().unsqueeze(-1) pde_inter_logits = self.to_pde_inter_logits(z + z.transpose(1, 2)) - pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze( - -1 - ) + pde_inter_logits = pde_inter_logits * is_different_chain.float().unsqueeze(-1) pde_logits = pde_inter_logits + pde_intra_logits else: @@ -328,9 +407,7 @@ def forward( if self.token_level_confidence: plddt = compute_aggregated_metric(plddt_logits) token_pad_mask = feats["token_pad_mask"].repeat_interleave(multiplicity, 0) - complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum( - dim=-1 - ) + complex_plddt = (plddt * token_pad_mask).sum(dim=-1) / token_pad_mask.sum(dim=-1) is_contact = (d < 8).float() is_different_chain = ( @@ -341,78 +418,53 @@ def forward( is_contact * is_different_chain * (1 - is_ligand_token).unsqueeze(-1), dim=-1, ).values - token_non_interface_mask = (1 - token_interface_mask) * ( - 1 - is_ligand_token - ) + token_non_interface_mask = (1 - token_interface_mask) * (1 - is_ligand_token) iplddt_weight = ( is_ligand_token * ligand_weight + token_interface_mask * interface_weight + token_non_interface_mask * non_interface_weight ) - complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum( - dim=-1 - ) / torch.sum(token_pad_mask * iplddt_weight, dim=-1) + complex_iplddt = (plddt * token_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum( + token_pad_mask * iplddt_weight, dim=-1 + ) else: # token to atom conversion for resolved logits - B, N, _ = resolved_logits.shape - resolved_logits = resolved_logits.reshape( - B, N, self.max_num_atoms_per_token, 2 - ) - - arange_max_num_atoms = ( - torch.arange(self.max_num_atoms_per_token) - .reshape(1, 1, -1) - .to(resolved_logits.device) - ) - max_num_atoms_mask = ( - feats["atom_to_token"].sum(1).unsqueeze(-1) > arange_max_num_atoms - ) - resolved_logits = resolved_logits[:, max_num_atoms_mask.squeeze(0)] - resolved_logits = pad( + b, n, _ = resolved_logits.shape + resolved_logits = resolved_logits.reshape(b, n, self.max_num_atoms_per_token, 2) + resolved_logits = _token_slot_logits_to_atom_logits( resolved_logits, - ( - 0, - 0, - 0, - int( - feats["atom_pad_mask"].shape[1] - - feats["atom_pad_mask"].sum().item() - ), - ), - value=0, + feats["atom_to_token"], + feats["atom_pad_mask"], + multiplicity=multiplicity, ) - plddt_logits = plddt_logits.reshape(B, N, self.max_num_atoms_per_token, -1) - plddt_logits = plddt_logits[:, max_num_atoms_mask.squeeze(0)] - plddt_logits = pad( + plddt_logits = plddt_logits.reshape(b, n, self.max_num_atoms_per_token, -1) + plddt_logits = _token_slot_logits_to_atom_logits( plddt_logits, - ( - 0, - 0, - 0, - int( - feats["atom_pad_mask"].shape[1] - - feats["atom_pad_mask"].sum().item() - ), - ), - value=0, + feats["atom_to_token"], + feats["atom_pad_mask"], + multiplicity=multiplicity, ) atom_pad_mask = feats["atom_pad_mask"].repeat_interleave(multiplicity, 0) plddt = compute_aggregated_metric(plddt_logits) - complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum( - dim=-1 + complex_plddt = (plddt * atom_pad_mask).sum(dim=-1) / atom_pad_mask.sum(dim=-1) + atom_to_token = feats["atom_to_token"].float().repeat_interleave( + multiplicity, + 0, + ) + chain_id_token = feats["asym_id"].float().repeat_interleave( + multiplicity, + 0, ) - token_type = feats["mol_type"].float() - atom_to_token = feats["atom_to_token"].float() - chain_id_token = feats["asym_id"].float() - atom_type = torch.bmm(atom_to_token, token_type.unsqueeze(-1)).squeeze(-1) + atom_type = torch.bmm( + atom_to_token, + token_type.float().unsqueeze(-1), + ).squeeze(-1) is_ligand_atom = (atom_type == const.chain_type_ids["NONPOLYMER"]).float() d_atom = torch.cdist(x_pred, x_pred) is_contact = (d_atom < 8).float() - chain_id_atom = torch.bmm( - atom_to_token, chain_id_token.unsqueeze(-1) - ).squeeze(-1) + chain_id_atom = torch.bmm(atom_to_token, chain_id_token.unsqueeze(-1)).squeeze(-1) is_different_chain = ( chain_id_atom.unsqueeze(-1) != chain_id_atom.unsqueeze(-2) ).float() @@ -428,9 +480,10 @@ def forward( + atom_non_interface_mask * non_interface_weight ) - complex_iplddt = (plddt * feats["atom_pad_mask"] * iplddt_weight).sum( - dim=-1 - ) / torch.sum(feats["atom_pad_mask"] * iplddt_weight, dim=-1) + complex_iplddt = (plddt * atom_pad_mask * iplddt_weight).sum(dim=-1) / torch.sum( + atom_pad_mask * iplddt_weight, + dim=-1, + ) # Compute the gPDE and giPDE pde = compute_aggregated_metric(pde_logits, end=32) @@ -446,16 +499,15 @@ def forward( token_pad_pair_mask = ( token_pad_mask.unsqueeze(-1) * token_pad_mask.unsqueeze(-2) - * ( - 1 - - torch.eye( - token_pad_mask.shape[1], device=token_pad_mask.device - ).unsqueeze(0) - ) + * (1 - torch.eye(token_pad_mask.shape[1], device=token_pad_mask.device).unsqueeze(0)) ) token_pair_mask = token_pad_pair_mask * prob_contact - complex_pde = (pde * token_pair_mask).sum(dim=(1, 2)) / token_pair_mask.sum( - dim=(1, 2) + complex_pde_numerator = (pde * token_pair_mask).sum(dim=(1, 2)) + complex_pde_denominator = token_pair_mask.sum(dim=(1, 2)) + complex_pde = complex_pde_numerator / torch.where( + complex_pde_denominator > 0, + complex_pde_denominator, + torch.ones_like(complex_pde_denominator), ) asym_id = feats["asym_id"].repeat_interleave(multiplicity, 0) token_interface_pair_mask = token_pair_mask * ( @@ -478,21 +530,13 @@ def forward( out_dict["pae_logits"] = pae_logits out_dict["pae"] = compute_aggregated_metric(pae_logits, end=32) - try: - ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms( - pae_logits, x_pred, feats, multiplicity - ) - out_dict["ptm"] = ptm - out_dict["iptm"] = iptm - out_dict["ligand_iptm"] = ligand_iptm - out_dict["protein_iptm"] = protein_iptm - out_dict["pair_chains_iptm"] = pair_chains_iptm - except Exception as e: - print(f"Error in compute_ptms: {e}") - out_dict["ptm"] = torch.zeros_like(complex_plddt) - out_dict["iptm"] = torch.zeros_like(complex_plddt) - out_dict["ligand_iptm"] = torch.zeros_like(complex_plddt) - out_dict["protein_iptm"] = torch.zeros_like(complex_plddt) - out_dict["pair_chains_iptm"] = torch.zeros_like(complex_plddt) + ptm, iptm, ligand_iptm, protein_iptm, pair_chains_iptm = compute_ptms( + pae_logits, x_pred, feats, multiplicity + ) + out_dict["ptm"] = ptm + out_dict["iptm"] = iptm + out_dict["ligand_iptm"] = ligand_iptm + out_dict["protein_iptm"] = protein_iptm + out_dict["pair_chains_iptm"] = pair_chains_iptm return out_dict diff --git a/src/fastplms/models/boltz/vb_modules_diffusion_conditioning.py b/src/fastplms/models/boltz/vb_modules_diffusion_conditioning.py new file mode 100644 index 0000000..a9fd0df --- /dev/null +++ b/src/fastplms/models/boltz/vb_modules_diffusion_conditioning.py @@ -0,0 +1,145 @@ +"""Precompute pair and atom biases used by the Boltz2 diffusion stack.""" + +from __future__ import annotations + +import torch +from torch import nn + +from .vb_modules_encodersv2 import AtomEncoder, PairwiseConditioning + + +def _bias_projections( + depth: int, + input_dim: int, + num_heads: int, +) -> nn.ModuleList: + """Build one normalized, bias-free projection per transformer block.""" + + return nn.ModuleList( + [ + nn.Sequential( + nn.LayerNorm(input_dim), + nn.Linear(input_dim, num_heads, bias=False), + ) + for _ in range(depth) + ] + ) + + +def _concatenate_biases( + projections: nn.ModuleList, + pair_features: torch.Tensor, +) -> torch.Tensor: + # pair_features: (..., d_pair); each projection: (..., h). + return torch.cat( + [projection(pair_features) for projection in projections], dim=-1 + ) # (..., depth * h) + + +class DiffusionConditioning(nn.Module): + """Prepare conditioned atom features and per-layer attention biases.""" + + def __init__( + self, + token_s: int, + token_z: int, + atom_s: int, + atom_z: int, + atoms_per_window_queries: int = 32, + atoms_per_window_keys: int = 128, + atom_encoder_depth: int = 3, + atom_encoder_heads: int = 4, + token_transformer_depth: int = 24, + token_transformer_heads: int = 8, + atom_decoder_depth: int = 3, + atom_decoder_heads: int = 4, + atom_feature_dim: int = 128, + conditioning_transition_layers: int = 2, + use_no_atom_char: bool = False, + use_atom_backbone_feat: bool = False, + use_residue_feats_atoms: bool = False, + ) -> None: + super().__init__() + self.pairwise_conditioner = PairwiseConditioning( + token_z=token_z, + dim_token_rel_pos_feats=token_z, + num_transitions=conditioning_transition_layers, + ) + self.atom_encoder = AtomEncoder( + atom_s=atom_s, + atom_z=atom_z, + token_s=token_s, + token_z=token_z, + atoms_per_window_queries=atoms_per_window_queries, + atoms_per_window_keys=atoms_per_window_keys, + atom_feature_dim=atom_feature_dim, + structure_prediction=True, + use_no_atom_char=use_no_atom_char, + use_atom_backbone_feat=use_atom_backbone_feat, + use_residue_feats_atoms=use_residue_feats_atoms, + ) + self.atom_enc_proj_z = _bias_projections( + atom_encoder_depth, + atom_z, + atom_encoder_heads, + ) + self.atom_dec_proj_z = _bias_projections( + atom_decoder_depth, + atom_z, + atom_decoder_heads, + ) + self.token_trans_proj_z = _bias_projections( + token_transformer_depth, + token_z, + token_transformer_heads, + ) + + def forward( + self, + s_trunk: torch.Tensor, + z_trunk: torch.Tensor, + relative_position_encoding: torch.Tensor, + feats: dict[str, torch.Tensor], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """Return conditioned atom tensors and concatenated layer biases. + + ``S`` has shape ``(b, n, d_s)`` and each ``Z`` tensor has shape + ``(b, n, n, d_z)``. Biases are concatenated in transformer-block + order so downstream code can select them by layer. + """ + + # b is batch size, t token count, a atom count, and k the atom-window count. + z_conditioned = self.pairwise_conditioner( + z_trunk, + relative_position_encoding, + ) # (b, t, t, d_z) + q, c, p, to_keys = self.atom_encoder( + feats=feats, + s_trunk=s_trunk, + z=z_conditioned, + ) # q/c: (b, a, d_a); p: (b, k, w, h_k, d_p); to_keys: callable + atom_encoder_bias = _concatenate_biases( + self.atom_enc_proj_z, p + ) # (b, k, w, h_k, depth_enc * heads_enc) + atom_decoder_bias = _concatenate_biases( + self.atom_dec_proj_z, p + ) # (b, k, w, h_k, depth_dec * heads_dec) + token_transformer_bias = _concatenate_biases( + self.token_trans_proj_z, + z_conditioned, + ) # (b, t, t, depth_token * heads_token) + return ( + q, + c, + to_keys, + atom_encoder_bias, + atom_decoder_bias, + token_transformer_bias, + ) # tensor shapes are traced above; to_keys is the atom-key gatherer diff --git a/fastplms/boltz/vb_modules_diffusionv2.py b/src/fastplms/models/boltz/vb_modules_diffusionv2.py similarity index 87% rename from fastplms/boltz/vb_modules_diffusionv2.py rename to src/fastplms/models/boltz/vb_modules_diffusionv2.py index b902476..dbb2762 100644 --- a/fastplms/boltz/vb_modules_diffusionv2.py +++ b/src/fastplms/models/boltz/vb_modules_diffusionv2.py @@ -1,4 +1,5 @@ -# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang +# Based on https://github.com/lucidrains/alphafold3-pytorch. +# MIT License, Copyright (c) 2024 Phil Wang. from __future__ import annotations @@ -6,7 +7,7 @@ import numpy as np import torch -import torch.nn.functional as F # noqa: N812 +import torch.nn.functional as F from einops import rearrange from torch import nn from torch.nn import Module @@ -97,9 +98,7 @@ def __init__( # post_layer_norm=transformer_post_ln, ) - self.a_norm = nn.LayerNorm( - 2 * token_s - ) # if not transformer_post_ln else nn.Identity() + self.a_norm = nn.LayerNorm(2 * token_s) # if not transformer_post_ln else nn.Identity() self.atom_attention_decoder = AtomAttentionDecoder( atom_s=atom_s, @@ -123,7 +122,7 @@ def forward( multiplicity=1, ): if self.activation_checkpointing: - s, normed_fourier = torch.utils.checkpoint.checkpoint( + s, _normed_fourier = torch.utils.checkpoint.checkpoint( self.single_conditioner, times, s_trunk.repeat_interleave(multiplicity, 0), @@ -131,7 +130,7 @@ def forward( use_reentrant=False, ) else: - s, normed_fourier = self.single_conditioner( + s, _normed_fourier = self.single_conditioner( times, s_trunk.repeat_interleave(multiplicity, 0), s_inputs.repeat_interleave(multiplicity, 0), @@ -186,13 +185,15 @@ def __init__( sigma_max: float = 160.0, # max noise level sigma_data: float = 16.0, # standard deviation of data distribution rho: float = 7, # controls the sampling schedule - P_mean: float = -1.2, # mean of log-normal distribution from which noise is drawn for training - P_std: float = 1.5, # standard deviation of log-normal distribution from which noise is drawn for training + # Mean of the log-normal training-noise distribution. + P_mean: float = -1.2, + # Standard deviation of the log-normal training-noise distribution. + P_std: float = 1.5, gamma_0: float = 0.8, gamma_min: float = 1.0, noise_scale: float = 1.003, step_scale: float = 1.5, - step_scale_random: list = None, + step_scale_random: list | None = None, coordinate_augmentation: bool = True, coordinate_augmentation_inference=None, compile_score: bool = False, @@ -204,9 +205,7 @@ def __init__( **score_model_args, ) if compile_score: - self.score_model = torch.compile( - self.score_model, dynamic=False, fullgraph=False - ) + self.score_model = torch.compile(self.score_model, dynamic=False, fullgraph=False) # parameters self.sigma_min = sigma_min @@ -269,8 +268,7 @@ def preconditioned_network_forward( ) denoised_coords = ( - self.c_skip(padded_sigma) * noised_atom_coords - + self.c_out(padded_sigma) * r_update + self.c_skip(padded_sigma) * noised_atom_coords + self.c_out(padded_sigma) * r_update ) return denoised_coords @@ -278,14 +276,10 @@ def sample_schedule(self, num_sampling_steps=None): num_sampling_steps = default(num_sampling_steps, self.num_sampling_steps) inv_rho = 1 / self.rho - steps = torch.arange( - num_sampling_steps, device=self.device, dtype=torch.float32 - ) + steps = torch.arange(num_sampling_steps, device=self.device, dtype=torch.float32) sigmas = ( self.sigma_max**inv_rho - + steps - / (num_sampling_steps - 1) - * (self.sigma_min**inv_rho - self.sigma_max**inv_rho) + + steps / (num_sampling_steps - 1) * (self.sigma_min**inv_rho - self.sigma_max**inv_rho) ) ** self.rho sigmas = sigmas * self.sigma_data @@ -315,10 +309,7 @@ def sample( resample_weights = torch.ones(multiplicity, device=self.device).reshape( -1, steering_args["num_particles"] ) - if ( - steering_args["physical_guidance_update"] - or steering_args["contact_guidance_update"] - ): + if steering_args["physical_guidance_update"] or steering_args["contact_guidance_update"]: scaled_guidance_update = torch.zeros( (multiplicity, *atom_mask.shape[1:], 3), dtype=torch.float32, @@ -332,10 +323,10 @@ def sample( shape = (*atom_mask.shape, 3) - # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma + # Pair each schedule value with the next sigma and gamma. sigmas = self.sample_schedule(num_sampling_steps) gammas = torch.where(sigmas > self.gamma_min, self.gamma_0, 0.0) - sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:])) + sigmas_and_gammas = list(zip(sigmas[:-1], sigmas[1:], gammas[1:], strict=True)) if self.training and self.step_scale_random is not None: step_scale = np.random.choice(self.step_scale_random) else: @@ -353,14 +344,11 @@ def sample( multiplicity, device=atom_coords.device, dtype=atom_coords.dtype ) atom_coords = atom_coords - atom_coords.mean(dim=-2, keepdims=True) - atom_coords = ( - torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr - ) + atom_coords = torch.einsum("bmd,bds->bms", atom_coords, random_R) + random_tr if atom_coords_denoised is not None: atom_coords_denoised -= atom_coords_denoised.mean(dim=-2, keepdims=True) atom_coords_denoised = ( - torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R) - + random_tr + torch.einsum("bmd,bds->bms", atom_coords_denoised, random_R) + random_tr ) if ( steering_args["physical_guidance_update"] @@ -381,9 +369,7 @@ def sample( with torch.no_grad(): atom_coords_denoised = torch.zeros_like(atom_coords_noisy) sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device) - sample_ids_chunks = sample_ids.chunk( - multiplicity % max_parallel_samples + 1 - ) + sample_ids_chunks = sample_ids.chunk(multiplicity % max_parallel_samples + 1) for sample_ids_chunk in sample_ids_chunks: atom_coords_denoised_chunk = self.preconditioned_network_forward( @@ -397,10 +383,7 @@ def sample( atom_coords_denoised[sample_ids_chunk] = atom_coords_denoised_chunk if steering_args["fk_steering"] and ( - ( - step_idx % steering_args["fk_resampling_interval"] == 0 - and noise_var > 0 - ) + (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0) or step_idx == num_sampling_steps - 1 ): # Compute energy of x_0 prediction @@ -427,9 +410,9 @@ def sample( steering_args["physical_guidance_update"] or steering_args["contact_guidance_update"] ) and noise_var > 0: - ll_difference = ( - eps**2 - (eps + scaled_guidance_update) ** 2 - ).sum(dim=(-1, -2)) / (2 * noise_var) + ll_difference = (eps**2 - (eps + scaled_guidance_update) ** 2).sum( + dim=(-1, -2) + ) / (2 * noise_var) else: ll_difference = torch.zeros_like(energy) @@ -453,8 +436,7 @@ def sample( parameters = potential.compute_parameters(steering_t) if ( parameters["guidance_weight"] > 0 - and (guidance_step) % parameters["guidance_interval"] - == 0 + and (guidance_step) % parameters["guidance_interval"] == 0 ): energy_gradient += parameters[ "guidance_weight" @@ -466,26 +448,17 @@ def sample( guidance_update -= energy_gradient atom_coords_denoised += guidance_update scaled_guidance_update = ( - guidance_update - * -1 - * self.step_scale - * (sigma_t - t_hat) - / t_hat + guidance_update * -1 * self.step_scale * (sigma_t - t_hat) / t_hat ) if steering_args["fk_steering"] and ( - ( - step_idx % steering_args["fk_resampling_interval"] == 0 - and noise_var > 0 - ) + (step_idx % steering_args["fk_resampling_interval"] == 0 and noise_var > 0) or step_idx == num_sampling_steps - 1 ): resample_indices = ( torch.multinomial( resample_weights, - resample_weights.shape[1] - if step_idx < num_sampling_steps - 1 - else 1, + resample_weights.shape[1] if step_idx < num_sampling_steps - 1 else 1, replacement=True, ) + resample_weights.shape[1] @@ -504,9 +477,7 @@ def sample( steering_args["physical_guidance_update"] or steering_args["contact_guidance_update"] ): - scaled_guidance_update = scaled_guidance_update[ - resample_indices - ] + scaled_guidance_update = scaled_guidance_update[resample_indices] if token_repr is not None: token_repr = token_repr[resample_indices] @@ -536,10 +507,7 @@ def loss_weight(self, sigma): def noise_distribution(self, batch_size): return ( self.sigma_data - * ( - self.P_mean - + self.P_std * torch.randn((batch_size,), device=self.device) - ).exp() + * (self.P_mean + self.P_std * torch.randn((batch_size,), device=self.device)).exp() ) def forward( @@ -554,9 +522,7 @@ def forward( batch_size = feats["coords"].shape[0] // multiplicity if self.synchronize_sigmas: - sigmas = self.noise_distribution(batch_size).repeat_interleave( - multiplicity, 0 - ) + sigmas = self.noise_distribution(batch_size).repeat_interleave(multiplicity, 0) else: sigmas = self.noise_distribution(batch_size * multiplicity) padded_sigmas = rearrange(sigmas, "b -> b 1 1") @@ -611,9 +577,7 @@ def compute_loss( plddt_mask = feats["plddt"] > filter_by_plddt resolved_atom_mask_uni = resolved_atom_mask_uni * plddt_mask.float() - resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave( - multiplicity, 0 - ) + resolved_atom_mask = resolved_atom_mask_uni.repeat_interleave(multiplicity, 0) align_weights = denoised_atom_coords.new_ones(denoised_atom_coords.shape[:2]) atom_type = ( @@ -636,9 +600,7 @@ def compute_loss( + torch.eq(atom_type_mult, const.chain_type_ids["RNA"]).float() ) + ligand_loss_weight - * torch.eq( - atom_type_mult, const.chain_type_ids["NONPOLYMER"] - ).float() + * torch.eq(atom_type_mult, const.chain_type_ids["NONPOLYMER"]).float() ).float() ) @@ -659,12 +621,10 @@ def compute_loss( ) # weighted MSE loss of denoised atom positions - mse_loss = ( - (denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2 - ).sum(dim=-1) - mse_loss = torch.sum( - mse_loss * align_weights * resolved_atom_mask, dim=-1 - ) / (torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5) + mse_loss = ((denoised_atom_coords - atom_coords_aligned_ground_truth) ** 2).sum(dim=-1) + mse_loss = torch.sum(mse_loss * align_weights * resolved_atom_mask, dim=-1) / ( + torch.sum(3 * align_weights * resolved_atom_mask, dim=-1) + 1e-5 + ) # weight by sigma factor loss_weights = self.loss_weight(sigmas) diff --git a/src/fastplms/models/boltz/vb_modules_encodersv2.py b/src/fastplms/models/boltz/vb_modules_encodersv2.py new file mode 100644 index 0000000..f0e9f0c --- /dev/null +++ b/src/fastplms/models/boltz/vb_modules_encodersv2.py @@ -0,0 +1,647 @@ +"""Token, pair, and atom encoders used by the local Boltz2 runtime. + +The classes retain checkpoint-facing submodule names while separating feature +assembly, window indexing, conditioning, and atom-token aggregation into small +mechanism-specific units. No production import depends on the upstream Boltz +package. + +The Fourier and atom-attention mechanisms derive from AlphaFold 3 community +implementations under MIT terms. See ``THIRD_PARTY_NOTICES.md``. +""" + +from __future__ import annotations + +import torch +from collections.abc import Callable +from functools import partial +from math import pi +from torch import nn +from torch.nn.functional import one_hot + +from . import vb_layers_initialize as init +from .vb_layers_transition import Transition +from .vb_modules_transformersv2 import AtomTransformer +from .vb_modules_utils import LinearNoBias + + +def _transition_stack( + count: int, + dim: int, + hidden_dim: int, +) -> nn.ModuleList: + return nn.ModuleList([Transition(dim=dim, hidden=hidden_dim) for _ in range(count)]) + + +class FourierEmbedding(nn.Module): + """Embed diffusion time with fixed random Fourier frequencies.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.proj = nn.Linear(1, dim) + nn.init.normal_(self.proj.weight, mean=0, std=1) + nn.init.normal_(self.proj.bias, mean=0, std=1) + self.proj.requires_grad_(False) + + def forward(self, times: torch.Tensor) -> torch.Tensor: + """Map times with shape ``(b,)`` to an embedding ``H: (b, d)``.""" + + random_phase = self.proj(times.reshape(-1, 1)) # (b, d) + return torch.cos(2 * pi * random_phase) # (b, d) + + +def _pairwise_difference(values: torch.Tensor) -> torch.Tensor: + # values: (b, n) + return values[:, :, None] - values[:, None, :] # (b, n, n) + + +class RelativePositionEncoder(nn.Module): + """Encode residue, token, entity, and symmetry relationships.""" + + def __init__( + self, + token_z: int, + r_max: int = 32, + s_max: int = 2, + fix_sym_check: bool = False, + cyclic_pos_enc: bool = False, + ) -> None: + super().__init__() + self.r_max = r_max + self.s_max = s_max + input_dim = 4 * (r_max + 1) + 2 * (s_max + 1) + 1 + self.linear_layer = LinearNoBias(input_dim, token_z) + self.fix_sym_check = fix_sym_check + self.cyclic_pos_enc = cyclic_pos_enc + + def forward(self, feats: dict[str, torch.Tensor]) -> torch.Tensor: + """Return relative pair features ``Z: (b, n, n, token_z)``.""" + + same_chain = ( + feats["asym_id"][:, :, None] == feats["asym_id"][:, None, :] + ) # (b, n, n) + same_residue = ( + feats["residue_index"][:, :, None] == feats["residue_index"][:, None, :] + ) # (b, n, n) + same_entity = ( + feats["entity_id"][:, :, None] == feats["entity_id"][:, None, :] + ) # (b, n, n) + + residue_offset = _pairwise_difference(feats["residue_index"]) # (b, n, n) + if self.cyclic_pos_enc and torch.any(feats["cyclic_period"] > 0): + period = torch.where( + feats["cyclic_period"] > 0, + feats["cyclic_period"], + torch.zeros_like(feats["cyclic_period"]) + 10000, + ) # (b, n) + residue_offset = ( + residue_offset - period * torch.round(residue_offset / period) + ).long() # (b, n, n) + residue_offset = torch.clip( + residue_offset + self.r_max, + 0, + 2 * self.r_max, + ) # (b, n, n) + residue_offset = torch.where( + same_chain, + residue_offset, + torch.zeros_like(residue_offset) + 2 * self.r_max + 1, + ) # (b, n, n) + residue_features = one_hot( + residue_offset, 2 * self.r_max + 2 + ) # (b, n, n, 2 * r_max + 2) + + token_offset = torch.clip( + _pairwise_difference(feats["token_index"]) + self.r_max, + 0, + 2 * self.r_max, + ) # (b, n, n) + token_offset = torch.where( + same_chain & same_residue, + token_offset, + torch.zeros_like(token_offset) + 2 * self.r_max + 1, + ) # (b, n, n) + token_features = one_hot( + token_offset, 2 * self.r_max + 2 + ) # (b, n, n, 2 * r_max + 2) + + symmetry_offset = torch.clip( + _pairwise_difference(feats["sym_id"]) + self.s_max, + 0, + 2 * self.s_max, + ) # (b, n, n) + invalid_symmetry = ~same_entity if self.fix_sym_check else same_chain # (b, n, n) + symmetry_offset = torch.where( + invalid_symmetry, + torch.zeros_like(symmetry_offset) + 2 * self.s_max + 1, + symmetry_offset, + ) # (b, n, n) + symmetry_features = one_hot( + symmetry_offset, 2 * self.s_max + 2 + ) # (b, n, n, 2 * s_max + 2) + + pair_features = torch.cat( + ( + residue_features.float(), + token_features.float(), + same_entity.unsqueeze(-1).float(), + symmetry_features.float(), + ), + dim=-1, + ) # (b, n, n, d_rel) + return self.linear_layer(pair_features) # (b, n, n, token_z) + + +class SingleConditioning(nn.Module): + """Condition token features on trunk inputs and diffusion time.""" + + def __init__( + self, + sigma_data: float, + token_s: int = 384, + dim_fourier: int = 256, + num_transitions: int = 2, + transition_expansion_factor: int = 2, + eps: float = 1e-20, + disable_times: bool = False, + ) -> None: + super().__init__() + self.eps = eps + self.sigma_data = sigma_data + self.disable_times = disable_times + conditioning_dim = 2 * token_s + self.norm_single = nn.LayerNorm(conditioning_dim) + self.single_embed = nn.Linear(conditioning_dim, conditioning_dim) + if not disable_times: + self.fourier_embed = FourierEmbedding(dim_fourier) + self.norm_fourier = nn.LayerNorm(dim_fourier) + self.fourier_to_single = LinearNoBias(dim_fourier, conditioning_dim) + self.transitions = _transition_stack( + num_transitions, + conditioning_dim, + transition_expansion_factor * conditioning_dim, + ) + + def forward( + self, + times: torch.Tensor, + s_trunk: torch.Tensor, + s_inputs: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return conditioned token features ``S: (b, n, 2 * token_s)``.""" + + # times: (b,); s_trunk/s_inputs: (b, n, token_s). + single = torch.cat((s_trunk, s_inputs), dim=-1) # (b, n, 2 * token_s) + single = self.single_embed(self.norm_single(single)) # (b, n, 2 * token_s) + normalized_fourier = None # (b, d_fourier) or None + if not self.disable_times: + fourier = self.fourier_embed(times) # (b, d_fourier) + normalized_fourier = self.norm_fourier(fourier) # (b, d_fourier) + time_condition = self.fourier_to_single( + normalized_fourier + ) # (b, 2 * token_s) + single = time_condition[:, None, :] + single # (b, n, 2 * token_s) + for transition in self.transitions: + single = transition(single) + single # (b, n, 2 * token_s) + return single, normalized_fourier # (b, n, 2 * token_s), (b, d_fourier) or None + + +class PairwiseConditioning(nn.Module): + """Fuse trunk pair features with relative-position features.""" + + def __init__( + self, + token_z: int, + dim_token_rel_pos_feats: int, + num_transitions: int = 2, + transition_expansion_factor: int = 2, + ) -> None: + super().__init__() + combined_dim = token_z + dim_token_rel_pos_feats + self.dim_pairwise_init_proj = nn.Sequential( + nn.LayerNorm(combined_dim), + LinearNoBias(combined_dim, token_z), + ) + self.transitions = _transition_stack( + num_transitions, + token_z, + transition_expansion_factor * token_z, + ) + + def forward( + self, + z_trunk: torch.Tensor, + token_rel_pos_feats: torch.Tensor, + ) -> torch.Tensor: + """Return conditioned pair features ``Z: (b, n, n, token_z)``.""" + + pair = self.dim_pairwise_init_proj( + torch.cat((z_trunk, token_rel_pos_feats), dim=-1) + ) # (b, n, n, token_z) + for transition in self.transitions: + pair = transition(pair) + pair # (b, n, n, token_z) + return pair # (b, n, n, token_z) + + +def get_indexing_matrix( + k: int, + w: int, + h: int, + device: torch.device, +) -> torch.Tensor: + """Build the atom-window gather matrix used for local attention keys.""" + + for name, value in (("k", k), ("w", w), ("h", h)): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an int, got {type(value).__name__}.") + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}.") + if w % 2 != 0: + raise ValueError(f"w must be even, got {w}.") + half_window = w // 2 + if h % half_window != 0: + raise ValueError(f"h must be divisible by w // 2 ({half_window}), got {h}.") + key_blocks = h // half_window + if key_blocks % 2 != 0: + raise ValueError( + f"h must contain an even number of half-window key blocks; received {key_blocks}." + ) + + positions = torch.arange(2 * k, device=device) # (2 * k,) + relative_blocks = ( + (positions.unsqueeze(0) - positions.unsqueeze(1)) + key_blocks // 2 + ).clamp(min=0, max=key_blocks + 1) # (2 * k, 2 * k) + relative_blocks = relative_blocks.view(k, 2, 2 * k)[:, 0, :] # (k, 2 * k) + selectors = one_hot(relative_blocks, num_classes=key_blocks + 2)[ + ..., 1:-1 + ].transpose(1, 0) # (2 * k, k, key_blocks) + return selectors.reshape( + 2 * k, key_blocks * k + ).float() # (2 * k, key_blocks * k) + + +def single_to_keys( + single: torch.Tensor, + indexing_matrix: torch.Tensor, + w: int, + h: int, +) -> torch.Tensor: + """Gather a sequence tensor into overlapping local key windows.""" + + b, n, d = single.shape + k = n // w + half_windows = single.view(b, 2 * k, w // 2, d) # (b, 2 * k, w / 2, d) + gathered = torch.einsum( + "b j i d, j k -> b k i d", half_windows, indexing_matrix + ) # (b, key_blocks * k, w / 2, d) + return gathered.reshape(b, k, h, d) # (b, k, h, d) + + +class AtomEncoder(nn.Module): + """Encode reference atoms and local atom-pair geometry.""" + + def __init__( + self, + atom_s: int, + atom_z: int, + token_s: int, + token_z: int, + atoms_per_window_queries: int, + atoms_per_window_keys: int, + atom_feature_dim: int, + structure_prediction: bool = True, + use_no_atom_char: bool = False, + use_atom_backbone_feat: bool = False, + use_residue_feats_atoms: bool = False, + ) -> None: + super().__init__() + self.embed_atom_features = nn.Linear(atom_feature_dim, atom_s) + self.embed_atompair_ref_pos = LinearNoBias(3, atom_z) + self.embed_atompair_ref_dist = LinearNoBias(1, atom_z) + self.embed_atompair_mask = LinearNoBias(1, atom_z) + self.atoms_per_window_queries = atoms_per_window_queries + self.atoms_per_window_keys = atoms_per_window_keys + self.use_no_atom_char = use_no_atom_char + self.use_atom_backbone_feat = use_atom_backbone_feat + self.use_residue_feats_atoms = use_residue_feats_atoms + self.structure_prediction = structure_prediction + + if structure_prediction: + self.s_to_c_trans = nn.Sequential( + nn.LayerNorm(token_s), + LinearNoBias(token_s, atom_s), + ) + init.final_init_(self.s_to_c_trans[1].weight) + self.z_to_p_trans = nn.Sequential( + nn.LayerNorm(token_z), + LinearNoBias(token_z, atom_z), + ) + init.final_init_(self.z_to_p_trans[1].weight) + + self.c_to_p_trans_k = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z)) + init.final_init_(self.c_to_p_trans_k[1].weight) + self.c_to_p_trans_q = nn.Sequential(nn.ReLU(), LinearNoBias(atom_s, atom_z)) + init.final_init_(self.c_to_p_trans_q[1].weight) + self.p_mlp = nn.Sequential( + nn.ReLU(), + LinearNoBias(atom_z, atom_z), + nn.ReLU(), + LinearNoBias(atom_z, atom_z), + nn.ReLU(), + LinearNoBias(atom_z, atom_z), + ) + init.final_init_(self.p_mlp[5].weight) + + def _assemble_atom_features( + self, + feats: dict[str, torch.Tensor], + b: int, + n: int, + ) -> torch.Tensor: + # Feature tensors share leading shape (b, n). + feature_parts = [ + feats["ref_pos"], + feats["ref_charge"].unsqueeze(-1), + feats["ref_element"], + ] # each: (b, n, d_feature) + if not self.use_no_atom_char: + feature_parts.append(feats["ref_atom_name_chars"].reshape(b, n, 4 * 64)) + if self.use_atom_backbone_feat: + feature_parts.append(feats["atom_backbone_feat"]) + if self.use_residue_feats_atoms: + residue_features = torch.cat( + ( + feats["res_type"], + feats["modified"].unsqueeze(-1), + one_hot(feats["mol_type"], num_classes=4).float(), + ), + dim=-1, + ) # (b, t, d_residue) + feature_parts.append( + torch.bmm(feats["atom_to_token"].float(), residue_features) + ) # (b, n, d_residue) + return torch.cat(feature_parts, dim=-1) # (b, n, atom_feature_dim) + + def forward( + self, + feats: dict[str, torch.Tensor], + s_trunk: torch.Tensor | None = None, + z: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + Callable[[torch.Tensor], torch.Tensor], + ]: + """Return atom queries, conditioned singles, pairs, and key gatherer.""" + + with torch.autocast("cuda", enabled=False): + b, n, _ = feats["ref_pos"].shape + atom_mask = feats["atom_pad_mask"].bool() # (b, n) + atom_positions = feats["ref_pos"] # (b, n, 3) + atom_space = feats["ref_space_uid"] # (b, n) + atom_features = self._assemble_atom_features( + feats, b, n + ) # (b, n, atom_feature_dim) + conditioned_single = self.embed_atom_features( + atom_features + ) # (b, n, atom_s) + + w = self.atoms_per_window_queries + h = self.atoms_per_window_keys + b, n = conditioned_single.shape[:2] + k = n // w + indexing = get_indexing_matrix( + k, w, h, conditioned_single.device + ) # (2 * k, key_blocks * k) + to_keys = partial(single_to_keys, indexing_matrix=indexing, w=w, h=h) + + position_queries = atom_positions.view(b, k, w, 1, 3) # (b, k, w, 1, 3) + position_keys = to_keys(atom_positions).view( + b, k, 1, h, 3 + ) # (b, k, 1, h, 3) + displacement = position_keys - position_queries # (b, k, w, h, 3) + inverse_squared_distance = 1 / ( + 1 + torch.sum(displacement * displacement, dim=-1, keepdim=True) + ) # (b, k, w, h, 1) + + mask_queries = atom_mask.view(b, k, w, 1) # (b, k, w, 1) + mask_keys = ( + to_keys(atom_mask.unsqueeze(-1).float()).view(b, k, 1, h).bool() + ) # (b, k, 1, h) + space_queries = atom_space.view(b, k, w, 1) # (b, k, w, 1) + space_keys = ( + to_keys(atom_space.unsqueeze(-1).float()).view(b, k, 1, h).long() + ) # (b, k, 1, h) + valid_pair = ( + (mask_queries & mask_keys & (space_queries == space_keys)).float().unsqueeze(-1) + ) # (b, k, w, h, 1) + + pair = self.embed_atompair_ref_pos(displacement) * valid_pair # (b, k, w, h, atom_z) + pair = ( + pair + self.embed_atompair_ref_dist(inverse_squared_distance) * valid_pair + ) # (b, k, w, h, atom_z) + pair = pair + self.embed_atompair_mask(valid_pair) * valid_pair # (b, k, w, h, atom_z) + query = conditioned_single # (b, n, atom_s) + + if self.structure_prediction: + if s_trunk is None or z is None: + raise ValueError("structure prediction requires S and Z trunk tensors") + atom_to_token = feats["atom_to_token"].float() # (b, n, t) + single_update = self.s_to_c_trans(s_trunk.float()) # (b, t, atom_s) + single_update = torch.bmm( + atom_to_token, single_update + ) # (b, n, atom_s) + conditioned_single = conditioned_single + single_update.to( + conditioned_single + ) # (b, n, atom_s) + + token_queries = atom_to_token.view( + b, k, w, atom_to_token.shape[-1] + ) # (b, k, w, t) + token_keys = to_keys(atom_to_token) # (b, k, h, t) + pair_update = self.z_to_p_trans(z.float()) # (b, t, t, atom_z) + pair_update = torch.einsum( + "bijd,bwki,bwlj->bwkld", + pair_update, + token_queries, + token_keys, + ) # (b, k, w, h, atom_z) + pair = pair + pair_update.to(pair) # (b, k, w, h, atom_z) + + pair = pair + self.c_to_p_trans_q( + conditioned_single.view(b, k, w, 1, conditioned_single.shape[-1]) + ) # (b, k, w, h, atom_z) + pair = pair + self.c_to_p_trans_k( + to_keys(conditioned_single).view( + b, + k, + 1, + h, + conditioned_single.shape[-1], + ) + ) # (b, k, w, h, atom_z) + pair = pair + self.p_mlp(pair) # (b, k, w, h, atom_z) + return ( + query, + conditioned_single, + pair, + to_keys, + ) # (b, n, atom_s), (b, n, atom_s), (b, k, w, h, atom_z), callable + + +class AtomAttentionEncoder(nn.Module): + """Run local atom attention and aggregate atoms to tokens.""" + + def __init__( + self, + atom_s: int, + token_s: int, + atoms_per_window_queries: int, + atoms_per_window_keys: int, + atom_encoder_depth: int = 3, + atom_encoder_heads: int = 4, + structure_prediction: bool = True, + activation_checkpointing: bool = False, + transformer_post_layer_norm: bool = False, + ) -> None: + super().__init__() + self.structure_prediction = structure_prediction + if structure_prediction: + self.r_to_q_trans = LinearNoBias(3, atom_s) + init.final_init_(self.r_to_q_trans.weight) + self.atom_encoder = AtomTransformer( + dim=atom_s, + dim_single_cond=atom_s, + attn_window_queries=atoms_per_window_queries, + attn_window_keys=atoms_per_window_keys, + depth=atom_encoder_depth, + heads=atom_encoder_heads, + activation_checkpointing=activation_checkpointing, + post_layer_norm=transformer_post_layer_norm, + ) + output_dim = 2 * token_s if structure_prediction else token_s + self.atom_to_token_trans = nn.Sequential( + LinearNoBias(atom_s, output_dim), + nn.ReLU(), + ) + + def forward( + self, + feats: dict[str, torch.Tensor], + q: torch.Tensor, + c: torch.Tensor, + atom_enc_bias: torch.Tensor, + to_keys: Callable[[torch.Tensor], torch.Tensor], + r: torch.Tensor | None = None, + multiplicity: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Callable]: + b, n, _ = feats["ref_pos"].shape + del b, n + # q/c: (b, a, atom_s); atom_enc_bias uses local-window attention axes. + atom_mask = feats["atom_pad_mask"].bool() # (b, a) + if self.structure_prediction: + if r is None: + raise ValueError("structure prediction requires atom coordinates R") + q = q.repeat_interleave(multiplicity, 0) # (b * m, a, atom_s) + q = q + self.r_to_q_trans(r) # (b * m, a, atom_s) + c = c.repeat_interleave(multiplicity, 0) # (b * m, a, atom_s) + atom_mask = atom_mask.repeat_interleave(multiplicity, 0) # (b * m, a) + q = self.atom_encoder( + q=q, + mask=atom_mask, + c=c, + bias=atom_enc_bias, + multiplicity=multiplicity, + to_keys=to_keys, + ) # (b * m, a, atom_s) + + with torch.autocast("cuda", enabled=False): + atom_update = self.atom_to_token_trans(q).float() # (b * m, a, d_token) + atom_to_token = feats["atom_to_token"].float() # (b, a, t) + atom_to_token = atom_to_token.repeat_interleave( + multiplicity, 0 + ) # (b * m, a, t) + atom_to_token_mean = atom_to_token / ( + atom_to_token.sum(dim=1, keepdim=True) + 1e-6 + ) # (b * m, a, t) + token_update = torch.bmm( + atom_to_token_mean.transpose(1, 2), atom_update + ) # (b * m, t, d_token) + return ( + token_update.to(q), + q, + c, + to_keys, + ) # (b * m, t, d_token), two (b * m, a, atom_s), callable + + +class AtomAttentionDecoder(nn.Module): + """Map token updates back to atoms and predict coordinate displacements.""" + + def __init__( + self, + atom_s: int, + token_s: int, + attn_window_queries: int, + attn_window_keys: int, + atom_decoder_depth: int = 3, + atom_decoder_heads: int = 4, + activation_checkpointing: bool = False, + transformer_post_layer_norm: bool = False, + ) -> None: + super().__init__() + self.a_to_q_trans = LinearNoBias(2 * token_s, atom_s) + init.final_init_(self.a_to_q_trans.weight) + self.atom_decoder = AtomTransformer( + dim=atom_s, + dim_single_cond=atom_s, + attn_window_queries=attn_window_queries, + attn_window_keys=attn_window_keys, + depth=atom_decoder_depth, + heads=atom_decoder_heads, + activation_checkpointing=activation_checkpointing, + post_layer_norm=transformer_post_layer_norm, + ) + if transformer_post_layer_norm: + self.atom_feat_to_atom_pos_update = LinearNoBias(atom_s, 3) + init.final_init_(self.atom_feat_to_atom_pos_update.weight) + else: + self.atom_feat_to_atom_pos_update = nn.Sequential( + nn.LayerNorm(atom_s), + LinearNoBias(atom_s, 3), + ) + init.final_init_(self.atom_feat_to_atom_pos_update[1].weight) + + def forward( + self, + a: torch.Tensor, + q: torch.Tensor, + c: torch.Tensor, + atom_dec_bias: torch.Tensor, + feats: dict[str, torch.Tensor], + to_keys: Callable[[torch.Tensor], torch.Tensor], + multiplicity: int = 1, + ) -> torch.Tensor: + """Return atom-coordinate updates ``R_update: (b, a, 3)``.""" + + with torch.autocast("cuda", enabled=False): + # a: (b * m, t, 2 * token_s); q/c: (b * m, a, atom_s). + atom_to_token = feats["atom_to_token"].float() # (b, a, t) + atom_to_token = atom_to_token.repeat_interleave( + multiplicity, 0 + ) # (b * m, a, t) + token_update = self.a_to_q_trans(a.float()) # (b * m, t, atom_s) + atom_update = torch.bmm(atom_to_token, token_update) # (b * m, a, atom_s) + q = q + atom_update.to(q) # (b * m, a, atom_s) + atom_mask = feats["atom_pad_mask"].repeat_interleave( + multiplicity, 0 + ) # (b * m, a) + q = self.atom_decoder( + q=q, + mask=atom_mask, + c=c, + bias=atom_dec_bias, + multiplicity=multiplicity, + to_keys=to_keys, + ) # (b * m, a, atom_s) + return self.atom_feat_to_atom_pos_update(q) # (b * m, a, 3) diff --git a/src/fastplms/models/boltz/vb_modules_transformersv2.py b/src/fastplms/models/boltz/vb_modules_transformersv2.py new file mode 100644 index 0000000..87b2e24 --- /dev/null +++ b/src/fastplms/models/boltz/vb_modules_transformersv2.py @@ -0,0 +1,276 @@ +"""Conditioned transformer blocks used by Boltz2 diffusion.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, cast +from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint + +from .vb_layers_attentionv2 import AttentionPairBias +from .vb_modules_utils import LinearNoBias, SwiGLU + + +class AdaLN(nn.Module): + """Normalize activations and apply scale and shift from conditioning S.""" + + def __init__(self, dim: int, dim_single_cond: int) -> None: + super().__init__() + self.a_norm = nn.LayerNorm(dim, elementwise_affine=False, bias=False) + self.s_norm = nn.LayerNorm(dim_single_cond, bias=False) + self.s_scale = nn.Linear(dim_single_cond, dim) + self.s_bias = LinearNoBias(dim_single_cond, dim) + + def forward(self, a: Tensor, s: Tensor) -> Tensor: + # a: (..., d); s: (..., d_c). + normalized = self.a_norm(a) # (..., d) + conditioning = self.s_norm(s) # (..., d_c) + # (..., d) + output = self.s_scale(conditioning).sigmoid() * normalized + self.s_bias(conditioning) + return cast(Tensor, output) # (..., d) + + +class ConditionedTransitionBlock(nn.Module): + """Apply a conditioned gated transition without changing tensor shape.""" + + def __init__( + self, + dim_single: int, + dim_single_cond: int, + expansion_factor: float = 2, + ) -> None: + super().__init__() + self.adaln = AdaLN(dim_single, dim_single_cond) + inner_dim = int(dim_single * expansion_factor) + self.swish_gate = nn.Sequential( + LinearNoBias(dim_single, inner_dim * 2), + SwiGLU(), + ) + self.a_to_b = LinearNoBias(dim_single, inner_dim) + self.b_to_a = LinearNoBias(inner_dim, dim_single) + + projection = nn.Linear(dim_single_cond, dim_single) + nn.init.zeros_(projection.weight) # (d, d_c) + nn.init.constant_(projection.bias, -2.0) # (d,) + self.output_projection = nn.Sequential(projection, nn.Sigmoid()) + + def forward(self, a: Tensor, s: Tensor) -> Tensor: + # a: (..., d); s: (..., d_c). + normalized = self.adaln(a, s) # (..., d) + hidden = self.swish_gate(normalized) * self.a_to_b(normalized) # (..., d_inner) + # (..., d) + return cast(Tensor, self.output_projection(s) * self.b_to_a(hidden)) + + +class DiffusionTransformer(nn.Module): + """Run conditioned transformer layers over token or atom states.""" + + def __init__( + self, + depth: int, + heads: int, + dim: int = 384, + dim_single_cond: int | None = None, + pair_bias_attn: bool = True, + activation_checkpointing: bool = False, + post_layer_norm: bool = False, + ) -> None: + super().__init__() + conditioning_dim = dim if dim_single_cond is None else dim_single_cond + self.activation_checkpointing = activation_checkpointing + self.pair_bias_attn = pair_bias_attn + self.layers = nn.ModuleList( + [ + DiffusionTransformerLayer( + heads, + dim, + conditioning_dim, + post_layer_norm, + ) + for _ in range(depth) + ] + ) + + def _split_pair_bias(self, bias: Tensor | None) -> Tensor | None: + # bias: (b_z, l_q, l_k, n_layer * h) or None. + if not self.pair_bias_attn: + return None + if bias is None: + raise ValueError("pair bias is required when pair_bias_attn=True") + batch_size, query_length, key_length, width = bias.shape # b_z, l_q, l_k, n_layer * h + depth = len(self.layers) # n_layer + if depth == 0 or width % depth: + raise ValueError("pair-bias width must be divisible by transformer depth") + # (b_z, l_q, l_k, n_layer, h) + return bias.view(batch_size, query_length, key_length, depth, width // depth) + + def forward( + self, + a: Tensor, + s: Tensor, + bias: Tensor | None = None, + mask: Tensor | None = None, + to_keys: Callable[[Tensor], Tensor] | None = None, + multiplicity: int = 1, + ) -> Tensor: + """Transform A and preserve shape ``(b * m, l, d)``.""" + + # a: (b * m, l_q, d); s: (b * m, l_q, d_c). + # bias: (b, l_q, l_k, n_layer * h) or None. + # mask: (b * m, l_q) before an optional to_keys mapping; to_keys preserves d. + layer_biases = self._split_pair_bias(bias) # (b, l_q, l_k, n_layer, h) or None + output = a # (b * m, l_q, d) + for index, layer in enumerate(self.layers): + # (b, l_q, l_k, h) or None + layer_bias = None if layer_biases is None else layer_biases[..., index, :] + if self.activation_checkpointing: + output = checkpoint( # (b * m, l_q, d) + layer, + output, + s, + layer_bias, + mask, + to_keys, + multiplicity, + use_reentrant=False, + ) + else: + output = layer( # (b * m, l_q, d) + output, + s, + layer_bias, + mask, + to_keys, + multiplicity, + ) + return output # (b * m, l_q, d) + + +class DiffusionTransformerLayer(nn.Module): + """One adaptive-normalization attention and transition block.""" + + def __init__( + self, + heads: int, + dim: int = 384, + dim_single_cond: int | None = None, + post_layer_norm: bool = False, + ) -> None: + super().__init__() + conditioning_dim = dim if dim_single_cond is None else dim_single_cond + self.adaln = AdaLN(dim, conditioning_dim) + self.pair_bias_attn = AttentionPairBias( + c_s=dim, + num_heads=heads, + compute_pair_bias=False, + ) + self.output_projection_linear = nn.Linear(conditioning_dim, dim) + nn.init.zeros_(self.output_projection_linear.weight) # (d, d_c) + nn.init.constant_(self.output_projection_linear.bias, -2.0) # (d,) + self.output_projection = nn.Sequential( + self.output_projection_linear, + nn.Sigmoid(), + ) + self.transition = ConditionedTransitionBlock(dim, conditioning_dim) + self.post_lnorm = nn.LayerNorm(dim) if post_layer_norm else nn.Identity() + + def forward( + self, + a: Tensor, + s: Tensor, + bias: Tensor | None = None, + mask: Tensor | None = None, + to_keys: Callable[[Tensor], Tensor] | None = None, + multiplicity: int = 1, + ) -> Tensor: + """Update activation tensor A with conditioning S and pair bias.""" + + # a: (b * m, l_q, d); s: (b * m, l_q, d_c). + # bias: (b, l_q, l_k, h) or None. + # mask: (b * m, l_q) or None; to_keys maps l_q to l_k. + if bias is None or mask is None: + raise ValueError("diffusion attention requires pair bias and a key mask") + normalized = self.adaln(a, s) # (b * m, l_q, d) + key_states = normalized # (b * m, l_q, d) + key_mask = mask # (b * m, l_q) + if to_keys is not None: + key_states = to_keys(normalized) # (b * m, l_k, d) + key_mask = to_keys(mask.unsqueeze(-1)).squeeze(-1) # (b * m, l_k) + + attended = self.pair_bias_attn( # (b * m, l_q, d) + s=normalized, + z=bias, + mask=key_mask, + multiplicity=multiplicity, + k_in=key_states, + ) + output = a + self.output_projection(s) * attended # (b * m, l_q, d) + output = output + self.transition(output, s) # (b * m, l_q, d) + return cast(Tensor, self.post_lnorm(output)) # (b * m, l_q, d) + + +class AtomTransformer(nn.Module): + """Apply the diffusion transformer over fixed atom-query windows.""" + + def __init__( + self, + attn_window_queries: int, + attn_window_keys: int, + **diffusion_transformer_kwargs: Any, + ) -> None: + super().__init__() + self.attn_window_queries = attn_window_queries + self.attn_window_keys = attn_window_keys + self.diffusion_transformer = DiffusionTransformer( + **diffusion_transformer_kwargs, + ) + + def forward( + self, + q: Tensor, + c: Tensor, + bias: Tensor, + to_keys: Callable[[Tensor], Tensor], + mask: Tensor, + multiplicity: int = 1, + ) -> Tensor: + """Transform atom tensor Q with shape ``(b, n_atoms, d)``.""" + + query_window = self.attn_window_queries # w_q + key_window = self.attn_window_keys # w_k + batch_size, atom_count, width = q.shape # b * m, n_atom, d + window_count = atom_count // query_window # k + # q, c: (b * m, n_atom, d); bias: (b, k, w_q, w_k, n_layer * h). + # mask: (b * m, n_atom); k = n_atom // w_q. + # (b * m * k, w_q, d) + query_states = q.view(batch_size * window_count, query_window, -1) + # (b * m * k, w_q, d_c) + conditioning = c.view(batch_size * window_count, query_window, -1) + # (b * m * k, w_q) + query_mask = mask.view(batch_size * window_count, query_window) + # (b * m, k, w_q, w_k, n_layer * h) + pair_bias = bias.repeat_interleave(multiplicity, dim=0) + pair_bias = pair_bias.view( # (b * m * k, w_q, w_k, n_layer * h) + pair_bias.shape[0] * window_count, + query_window, + key_window, + -1, + ) + + def windowed_keys(states: Tensor) -> Tensor: + # states: (b * m * k, w_q, d_x). + # (b * m, n_atom, d_x) + merged = states.view(batch_size, window_count * query_window, -1) + # (b * m * k, w_k, d_x) + return to_keys(merged).view(batch_size * window_count, key_window, -1) + + output = self.diffusion_transformer( # (b * m * k, w_q, d) + a=query_states, + s=conditioning, + bias=pair_bias, + mask=query_mask.float(), + multiplicity=1, + to_keys=windowed_keys, + ) + # (b * m, n_atom, d) + return cast(Tensor, output.view(batch_size, window_count * query_window, width)) diff --git a/src/fastplms/models/boltz/vb_modules_trunkv2.py b/src/fastplms/models/boltz/vb_modules_trunkv2.py new file mode 100644 index 0000000..21096cf --- /dev/null +++ b/src/fastplms/models/boltz/vb_modules_trunkv2.py @@ -0,0 +1,623 @@ +"""Input, template, MSA, and prediction-head modules for the Boltz2 trunk.""" + +from __future__ import annotations + +import torch +from typing import Any, cast +from torch import Tensor, nn +from torch.nn import functional as F +from torch.utils.checkpoint import checkpoint + +from . import vb_const as const +from .vb_layers_dropout import get_dropout_mask +from .vb_layers_outer_product_mean import OuterProductMean +from .vb_layers_pair_averaging import PairWeightedAveraging +from .vb_layers_pairformer import PairformerNoSeqLayer, PairformerNoSeqModule +from .vb_layers_transition import Transition +from .vb_modules_encodersv2 import ( + AtomAttentionEncoder, + AtomEncoder, + FourierEmbedding, +) + + +class ContactConditioning(nn.Module): + """Encode selected, unselected, and distance-threshold contact constraints.""" + + def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float) -> None: + super().__init__() + self.fourier_embedding = FourierEmbedding(token_z) + input_width = token_z + len(const.contact_conditioning_info) - 1 + self.encoder = nn.Linear(input_width, token_z) + self.encoding_unspecified = nn.Parameter(torch.zeros(token_z)) # (d_z,) + self.encoding_unselected = nn.Parameter(torch.zeros(token_z)) # (d_z,) + self.cutoff_min = cutoff_min + self.cutoff_max = cutoff_max + + def forward(self, feats: dict[str, Tensor]) -> Tensor: + """Return contact tensor C with shape ``(b, l, l, d_z)``.""" + + if const.contact_conditioning_info["UNSPECIFIED"] != 0: + raise ValueError("UNSPECIFIED contact conditioning must use channel zero") + if const.contact_conditioning_info["UNSELECTED"] != 1: + raise ValueError("UNSELECTED contact conditioning must use channel one") + + # c_contact is the number of contact-conditioning categories. + categories = feats["contact_conditioning"] # (b, l, l, c_contact) + threshold = feats["contact_threshold"] # (b, l, l) + # (b, l, l) + normalized = (threshold - self.cutoff_min) / (self.cutoff_max - self.cutoff_min) + # (b, l, l, d_z) + fourier = self.fourier_embedding(normalized.flatten()).reshape((*normalized.shape, -1)) + selected_features = torch.cat( # (b, l, l, c_contact - 1 + d_z) + [categories[..., 2:], normalized.unsqueeze(-1), fourier], + dim=-1, + ) + selected = self.encoder(selected_features) # (b, l, l, d_z) + special = categories[..., :2] # (b, l, l, 2) + return cast( + Tensor, + selected * (1 - special.sum(dim=-1, keepdim=True)) + + self.encoding_unspecified * special[..., 0:1] + + self.encoding_unselected * special[..., 1:2], + ) # (b, l, l, d_z) + + +class InputEmbedder(nn.Module): + """Combine atom, residue, profile, and optional experimental features.""" + + def __init__( + self, + atom_s: int, + atom_z: int, + token_s: int, + token_z: int, + atoms_per_window_queries: int, + atoms_per_window_keys: int, + atom_feature_dim: int, + atom_encoder_depth: int, + atom_encoder_heads: int, + activation_checkpointing: bool = False, + add_method_conditioning: bool = False, + add_modified_flag: bool = False, + add_cyclic_flag: bool = False, + add_mol_type_feat: bool = False, + use_no_atom_char: bool = False, + use_atom_backbone_feat: bool = False, + use_residue_feats_atoms: bool = False, + ) -> None: + super().__init__() + self.token_s = token_s + self.add_method_conditioning = add_method_conditioning + self.add_modified_flag = add_modified_flag + self.add_cyclic_flag = add_cyclic_flag + self.add_mol_type_feat = add_mol_type_feat + self.atom_encoder = AtomEncoder( + atom_s=atom_s, + atom_z=atom_z, + token_s=token_s, + token_z=token_z, + atoms_per_window_queries=atoms_per_window_queries, + atoms_per_window_keys=atoms_per_window_keys, + atom_feature_dim=atom_feature_dim, + structure_prediction=False, + use_no_atom_char=use_no_atom_char, + use_atom_backbone_feat=use_atom_backbone_feat, + use_residue_feats_atoms=use_residue_feats_atoms, + ) + self.atom_enc_proj_z = nn.Sequential( + nn.LayerNorm(atom_z), + nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False), + ) + self.atom_attention_encoder = AtomAttentionEncoder( + atom_s=atom_s, + token_s=token_s, + atoms_per_window_queries=atoms_per_window_queries, + atoms_per_window_keys=atoms_per_window_keys, + atom_encoder_depth=atom_encoder_depth, + atom_encoder_heads=atom_encoder_heads, + structure_prediction=False, + activation_checkpointing=activation_checkpointing, + ) + self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False) + self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False) + + if add_method_conditioning: + self.method_conditioning_init = nn.Embedding(const.num_method_types, token_s) + self.method_conditioning_init.weight.data.fill_(0) # (n_method, d_s) + if add_modified_flag: + self.modified_conditioning_init = nn.Embedding(2, token_s) + self.modified_conditioning_init.weight.data.fill_(0) # (2, d_s) + if add_cyclic_flag: + self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False) + self.cyclic_conditioning_init.weight.data.fill_(0) # (d_s, 1) + if add_mol_type_feat: + self.mol_type_conditioning_init = nn.Embedding( + len(const.chain_type_ids), + token_s, + ) + self.mol_type_conditioning_init.weight.data.fill_(0) # (n_mol_type, d_s) + + def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor: + """Return embedded sequence tensor S with shape ``(b, l, d_s)``.""" + + # n_atom is the padded atom count; k is the number of atom windows. + residue_type = feats["res_type"].float() # (b, l, n_token_type) + suffix = "_affinity" if affinity else "" + profile = feats[f"profile{suffix}"] # (b, l, n_token_type) + deletion_mean = feats[f"deletion_mean{suffix}"].unsqueeze(-1) # (b, l, 1) + # (b, n_atom, d_a), (b, n_atom, d_a), (b, k, w_q, w_k, d_az), callable + atom_queries, atom_conditioning, atom_pairs, to_keys = self.atom_encoder(feats) + atom_bias = self.atom_enc_proj_z(atom_pairs) # (b, k, w_q, w_k, n_layer * h) + atom_output, _, _, _ = self.atom_attention_encoder( + feats=feats, + q=atom_queries, + c=atom_conditioning, + atom_enc_bias=atom_bias, + to_keys=to_keys, + ) # (b, l, d_s), (b, n_atom, d_a), (b, n_atom, d_a), callable + output = ( # (b, l, d_s) + atom_output + + self.res_type_encoding(residue_type) + + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1)) + ) + if self.add_method_conditioning: + # method_feature: (b, l); output: (b, l, d_s). + output = output + self.method_conditioning_init(feats["method_feature"]) + if self.add_modified_flag: + # modified: (b, l); output: (b, l, d_s). + output = output + self.modified_conditioning_init(feats["modified"]) + if self.add_cyclic_flag: + cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1) # (b, l, 1) + output = output + self.cyclic_conditioning_init(cyclic) # (b, l, d_s) + if self.add_mol_type_feat: + # mol_type: (b, l); output: (b, l, d_s). + output = output + self.mol_type_conditioning_init(feats["mol_type"]) + return cast(Tensor, output) # (b, l, d_s) + + +class _TemplateBase(nn.Module): + def __init__( + self, + token_z: int, + template_dim: int, + template_blocks: int, + dropout: float, + pairwise_head_width: int, + pairwise_num_heads: int, + post_layer_norm: bool, + activation_checkpointing: bool, + min_dist: float, + max_dist: float, + num_bins: int, + ) -> None: + super().__init__() + self.min_dist = min_dist + self.max_dist = max_dist + self.num_bins = num_bins + self.relu = nn.ReLU() + self.z_norm = nn.LayerNorm(token_z) + self.v_norm = nn.LayerNorm(template_dim) + self.z_proj = nn.Linear(token_z, template_dim, bias=False) + feature_width = const.num_tokens * 2 + num_bins + 5 + self.a_proj = nn.Linear(feature_width, template_dim, bias=False) + self.u_proj = nn.Linear(template_dim, token_z, bias=False) + self.pairformer = PairformerNoSeqModule( + template_dim, + num_blocks=template_blocks, + dropout=dropout, + pairwise_head_width=pairwise_head_width, + pairwise_num_heads=pairwise_num_heads, + post_layer_norm=post_layer_norm, + activation_checkpointing=activation_checkpointing, + ) + + def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor: + raise NotImplementedError + + def _template_features( + self, + feats: dict[str, Tensor], + template_pair_mask: Tensor, + ) -> Tensor: + # residue_type: (b, t, l, n_token_type); template_pair_mask: (b, t, l, l). + residue_type = feats["template_restype"] # (b, t, l, n_token_type) + cb_mask = feats["template_mask_cb"] # (b, t, l) + frame_mask = feats["template_mask_frame"] # (b, t, l) + # (b, t, l, l, 1) + cb_pair_mask = (cb_mask[..., :, None] * cb_mask[..., None, :]).unsqueeze(-1) + # (b, t, l, l, 1) + frame_pair_mask = (frame_mask[..., :, None] * frame_mask[..., None, :]).unsqueeze(-1) + with torch.autocast(device_type="cuda", enabled=False): + # template_cb: (b, t, l, 3). + # (b, t, l, l) + cb_distances = torch.cdist(feats["template_cb"], feats["template_cb"]) + boundaries = torch.linspace( # (n_bin - 1,) + self.min_dist, + self.max_dist, + self.num_bins - 1, + ).to(cb_distances.device) + bins = (cb_distances[..., None] > boundaries).sum(dim=-1).long() # (b, t, l, l) + distogram = F.one_hot(bins, num_classes=self.num_bins) # (b, t, l, l, n_bin) + + # (b, t, 1, l, 3, 3) + rotations = feats["template_frame_rot"].unsqueeze(2).transpose(-1, -2) + # (b, t, 1, l, 3, 1) + translations = feats["template_frame_t"].unsqueeze(2).unsqueeze(-1) + # (b, t, l, 1, 3, 1) + ca_coordinates = feats["template_ca"].unsqueeze(3).unsqueeze(-1) + # (b, t, l, l, 3, 1) + vectors = torch.matmul(rotations, ca_coordinates - translations) + norms = torch.norm(vectors, dim=-1, keepdim=True) # (b, t, l, l, 3, 1) + unit_vectors = torch.where( # (b, t, l, l, 3) + norms > 0, + vectors / norms, + torch.zeros_like(vectors), + ).squeeze(-1) + pair_features = torch.cat( # (b, t, l, l, n_bin + 5) + [distogram, cb_pair_mask, unit_vectors, frame_pair_mask], + dim=-1, + ) + # (b, t, l, l, n_bin + 5) + pair_features = pair_features * template_pair_mask.unsqueeze(-1) + residue_i = residue_type[:, :, :, None].expand( # (b, t, l, l, n_token_type) + -1, + -1, + -1, + residue_type.size(2), + -1, + ) + residue_j = residue_type[:, :, None, :].expand( # (b, t, l, l, n_token_type) + -1, + -1, + residue_type.size(2), + -1, + -1, + ) + return cast( + Tensor, + self.a_proj(torch.cat([pair_features, residue_i, residue_j], dim=-1)), + ) # (b, t, l, l, d_t) + + def forward( + self, + z: Tensor, + feats: dict[str, Tensor], + pair_mask: Tensor, + use_kernels: bool = False, + ) -> Tensor: + """Aggregate template pair tensor V into trunk update U.""" + + # z: (b, l, l, d_z); pair_mask: (b, l, l). + residue_type = feats["template_restype"] # (b, t, l, n_token_type) + batch_size, template_count = residue_type.shape[:2] + template_present = feats["template_mask"].any(dim=2).float() # (b, t) + present_count = template_present.sum(dim=1).clamp(min=1) # (b,) + features = self._template_features( # (b, t, l, l, d_t) + feats, + self._template_pair_mask(feats, template_count), + ) + expanded_mask = pair_mask[:, None].expand(-1, template_count, -1, -1) # (b, t, l, l) + expanded_mask = expanded_mask.reshape( # (b * t, l, l) + batch_size * template_count, + *expanded_mask.shape[2:], + ) + template_states = self.z_proj(self.z_norm(z[:, None])) + features # (b, t, l, l, d_t) + template_states = template_states.view( # (b * t, l, l, d_t) + batch_size * template_count, + *template_states.shape[2:], + ) + template_states = template_states + self.pairformer( # (b * t, l, l, d_t) + template_states, + expanded_mask, + use_kernels=use_kernels, + ) + template_states = self.v_norm(template_states).view( # (b, t, l, l, d_t) + batch_size, + template_count, + *template_states.shape[1:], + ) + weights = template_present[:, :, None, None, None] # (b, t, 1, 1, 1) + aggregate = (template_states * weights).sum(dim=1) # (b, l, l, d_t) + # (b, l, l, d_t) + aggregate = aggregate / present_count[:, None, None, None].to(template_states) + return cast(Tensor, self.u_proj(self.relu(aggregate))) # (b, l, l, d_z) + + +class TemplateModule(_TemplateBase): + """Aggregate templates while restricting features to the same chain.""" + + def __init__( + self, + token_z: int, + template_dim: int, + template_blocks: int, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + activation_checkpointing: bool = False, + min_dist: float = 3.25, + max_dist: float = 50.75, + num_bins: int = 38, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__( + token_z, + template_dim, + template_blocks, + dropout, + pairwise_head_width, + pairwise_num_heads, + post_layer_norm, + activation_checkpointing, + min_dist, + max_dist, + num_bins, + ) + + def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor: + asym_id = feats["asym_id"] # (b, l) + same_chain = (asym_id[:, :, None] == asym_id[:, None, :]).float() # (b, l, l) + return same_chain[:, None].expand(-1, count, -1, -1) # (b, t, l, l) + + +class TemplateV2Module(_TemplateBase): + """Aggregate templates under per-template visibility groups.""" + + def __init__( + self, + token_z: int, + template_dim: int, + template_blocks: int, + dropout: float = 0.25, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + post_layer_norm: bool = False, + activation_checkpointing: bool = False, + min_dist: float = 3.25, + max_dist: float = 50.75, + num_bins: int = 38, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__( + token_z, + template_dim, + template_blocks, + dropout, + pairwise_head_width, + pairwise_num_heads, + post_layer_norm, + activation_checkpointing, + min_dist, + max_dist, + num_bins, + ) + + def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor: + del count + visibility = feats["visibility_ids"] # (b, t, l) + return (visibility[..., :, None] == visibility[..., None, :]).float() # (b, t, l, l) + + +class MSAModule(nn.Module): + """Embed and update an MSA before returning its accumulated pair update.""" + + def __init__( + self, + msa_s: int, + token_z: int, + token_s: int, + msa_blocks: int, + msa_dropout: float, + z_dropout: float, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + activation_checkpointing: bool = False, + use_paired_feature: bool = True, + subsample_msa: bool = False, + num_subsampled_msa: int = 1024, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__() + self.msa_blocks = msa_blocks + self.msa_dropout = msa_dropout + self.z_dropout = z_dropout + self.use_paired_feature = use_paired_feature + self.activation_checkpointing = activation_checkpointing + self.subsample_msa = subsample_msa + self.num_subsampled_msa = num_subsampled_msa + self.s_proj = nn.Linear(token_s, msa_s, bias=False) + input_width = const.num_tokens + 2 + int(use_paired_feature) + self.msa_proj = nn.Linear(input_width, msa_s, bias=False) + self.layers = nn.ModuleList( + [ + MSALayer( + msa_s, + token_z, + msa_dropout, + z_dropout, + pairwise_head_width, + pairwise_num_heads, + ) + for _ in range(msa_blocks) + ] + ) + + @staticmethod + def _chunk_configuration( + pair_states: Tensor, + training: bool, + ) -> tuple[bool, int | None, int | None, int | None, int | None]: + # pair_states: (b, l, l, d_z). + if training: + return False, None, None, None, None + if pair_states.shape[1] > const.chunk_size_threshold: + return True, 64, 32, 4, 128 + return False, None, None, None, 512 + + def forward( + self, + z: Tensor, + emb: Tensor, + feats: dict[str, Tensor], + use_kernels: bool = False, + ) -> Tensor: + """Return updated pair tensor Z after every MSA block.""" + + chunking = self._chunk_configuration(z, self.training) + # z: (b, l, l, d_z); emb: (b, l, d_s). + # s is MSA depth; n_token_type is the residue vocabulary size. + msa = feats["msa"] # (b, s, l) or (b, s, l, n_token_type) + if msa.dtype in (torch.long, torch.int32, torch.int64): + msa = F.one_hot(msa, num_classes=const.num_tokens).float() # (b, s, l, n_token_type) + msa_mask = feats["msa_mask"] # (b, s, l) + components = [ # (b, s, l, n_token_type), then two (b, s, l, 1) tensors + msa, + feats["has_deletion"].unsqueeze(-1), # (b, s, l, 1) + feats["deletion_value"].unsqueeze(-1), # (b, s, l, 1) + ] + if self.use_paired_feature: + components.append(feats["msa_paired"].unsqueeze(-1)) # (b, s, l, 1) + msa_input = torch.cat(components, dim=-1) # (b, s, l, n_token_type + 2 or 3) + if self.subsample_msa: + indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa] # (s_sub,) + msa_input = msa_input[:, indices] # (b, s_sub, l, n_token_type + 2 or 3) + msa_mask = msa_mask[:, indices] # (b, s_sub, l) + + msa_states = self.msa_proj(msa_input) + self.s_proj(emb).unsqueeze(1) # (b, s, l, d_m) + token_mask = feats["token_pad_mask"].float() # (b, l) + pair_mask = token_mask[:, :, None] * token_mask[:, None, :] # (b, l, l) + pair_states = z # (b, l, l, d_z) + for layer in self.layers: + # Tensor arguments: pair_states (b, l, l, d_z), msa_states (b, s, l, d_m), + # pair_mask (b, l, l), msa_mask (b, s, l). + arguments = ( + pair_states, + msa_states, + pair_mask, + msa_mask, + *chunking, + use_kernels, + ) + if self.activation_checkpointing and self.training: + pair_states, msa_states = checkpoint( # (b, l, l, d_z), (b, s, l, d_m) + layer, + *arguments, + ) + else: + pair_states, msa_states = layer(*arguments) # (b, l, l, d_z), (b, s, l, d_m) + return pair_states # (b, l, l, d_z) + + +class MSALayer(nn.Module): + """Exchange information between MSA tensor M and pair tensor Z.""" + + def __init__( + self, + msa_s: int, + token_z: int, + msa_dropout: float, + z_dropout: float, + pairwise_head_width: int = 32, + pairwise_num_heads: int = 4, + ) -> None: + super().__init__() + self.msa_dropout = msa_dropout + self.msa_transition = Transition(msa_s, msa_s * 4) + self.pair_weighted_averaging = PairWeightedAveraging( + c_m=msa_s, + c_z=token_z, + c_h=32, + num_heads=8, + ) + self.pairformer_layer = PairformerNoSeqLayer( + token_z=token_z, + dropout=z_dropout, + pairwise_head_width=pairwise_head_width, + pairwise_num_heads=pairwise_num_heads, + ) + self.outer_product_mean = OuterProductMean(msa_s, 32, token_z) + + def forward( + self, + z: Tensor, + m: Tensor, + token_mask: Tensor, + msa_mask: Tensor, + chunk_heads_pwa: bool = False, + chunk_size_transition_z: int | None = None, + chunk_size_transition_msa: int | None = None, + chunk_size_outer_product: int | None = None, + chunk_size_tri_attn: int | None = None, + use_kernels: bool = False, + ) -> tuple[Tensor, Tensor]: + """Return updated Z and M tensors.""" + + del chunk_size_transition_z + # z: (b, l, l, d_z); m: (b, s, l, d_m). + # token_mask: (b, l, l); msa_mask: (b, s, l). + dropout = get_dropout_mask(self.msa_dropout, m, self.training) # (b, s, 1, 1) + msa_states = m + dropout * self.pair_weighted_averaging( # (b, s, l, d_m) + m, + z, + token_mask, + chunk_heads_pwa, + ) + msa_states = msa_states + self.msa_transition( # (b, s, l, d_m) + msa_states, + chunk_size_transition_msa, + ) + pair_states = z + self.outer_product_mean( # (b, l, l, d_z) + msa_states, + msa_mask, + chunk_size_outer_product, + ) + pair_states = self.pairformer_layer( # (b, l, l, d_z) + pair_states, + token_mask, + chunk_size_tri_attn, + use_kernels=use_kernels, + ) + return pair_states, msa_states # (b, l, l, d_z), (b, s, l, d_m) + + +class BFactorModule(nn.Module): + """Predict a per-token B-factor histogram.""" + + def __init__(self, token_s: int, num_bins: int) -> None: + super().__init__() + self.bfactor = nn.Linear(token_s, num_bins) + self.num_bins = num_bins + + def forward(self, s: Tensor) -> Tensor: + # s: (..., d_s). + return cast(Tensor, self.bfactor(s)) # (..., n_bin) + + +class DistogramModule(nn.Module): + """Predict symmetric residue-pair distance histograms.""" + + def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None: + super().__init__() + self.distogram = nn.Linear(token_z, num_distograms * num_bins) + self.num_distograms = num_distograms + self.num_bins = num_bins + + def forward(self, z: Tensor) -> Tensor: + # z: (b, l, l, d_z). + symmetric = z + z.transpose(1, 2) # (b, l, l, d_z) + logits = self.distogram(symmetric) # (b, l, l, n_distogram * n_bin) + return cast( + Tensor, + logits.reshape( + symmetric.shape[0], + symmetric.shape[1], + symmetric.shape[2], + self.num_distograms, + self.num_bins, + ), + ) # (b, l, l, n_distogram, n_bin) diff --git a/src/fastplms/models/boltz/vb_modules_utils.py b/src/fastplms/models/boltz/vb_modules_utils.py new file mode 100644 index 0000000..ab04415 --- /dev/null +++ b/src/fastplms/models/boltz/vb_modules_utils.py @@ -0,0 +1,282 @@ +"""Small geometry, augmentation, activation, and EMA utilities for Boltz2. + +The quaternion behavior follows PyTorch3D's public rotation convention. License +and attribution records are retained in ``THIRD_PARTY_NOTICES.md``. +""" + +from __future__ import annotations + +import torch +from collections.abc import Iterable, Mapping, Sequence +from functools import partial +from typing import Any +from torch import Tensor, nn +from torch.nn import functional as F + + +LinearNoBias = partial(nn.Linear, bias=False) + + +def exists(value: object) -> bool: + """Return whether a value is not ``None``.""" + + return value is not None + + +def default(value: Any, fallback: Any) -> Any: + """Return ``value`` unless it is ``None``.""" + + return fallback if value is None else value + + +def log(tensor: Tensor, eps: float = 1e-20) -> Tensor: + """Compute a finite logarithm by applying a scalar lower bound.""" + + # tensor: (...). + return torch.log(tensor.clamp(min=eps)) # (...) + + +class SwiGLU(nn.Module): + """Split X in half and apply a SiLU-gated linear unit.""" + + def forward(self, x: Tensor) -> Tensor: + # x: (..., 2 * d). + values, gates = x.chunk(2, dim=-1) # each: (..., d) + return F.silu(gates) * values # (..., d) + + +def _masked_center(coordinates: Tensor, mask: Tensor) -> Tensor: + # coordinates: (b, n_atom, 3); mask: (b, n_atom). + weights = mask[:, :, None] # (b, n_atom, 1) + # (b, 1, 3) + return (coordinates * weights).sum(dim=1, keepdim=True) / weights.sum( + dim=1, + keepdim=True, + ) + + +def center(atom_coords: Tensor, atom_mask: Tensor) -> Tensor: + """Center coordinate tensor X with shape ``(b, n_atoms, 3)``.""" + + # atom_coords: (b, n_atom, 3); atom_mask: (b, n_atom). + return atom_coords - _masked_center(atom_coords, atom_mask) # (b, n_atom, 3) + + +def _copysign(magnitudes: Tensor, signs: Tensor) -> Tensor: + """Apply the elementwise sign of S to magnitude tensor M.""" + + # magnitudes, signs: broadcast-compatible (...). + return torch.where((magnitudes < 0) != (signs < 0), -magnitudes, magnitudes) # (...) + + +def quaternion_to_matrix(quaternions: Tensor) -> Tensor: + """Convert real-first quaternion tensor Q from ``(..., 4)`` to ``(..., 3, 3)``.""" + + # quaternions: (..., 4). + real, i_axis, j_axis, k_axis = torch.unbind(quaternions, dim=-1) # each: (...) + scale = 2.0 / (quaternions * quaternions).sum(dim=-1) # (...) + entries = ( # nine tensors, each (...) + 1 - scale * (j_axis * j_axis + k_axis * k_axis), + scale * (i_axis * j_axis - k_axis * real), + scale * (i_axis * k_axis + j_axis * real), + scale * (i_axis * j_axis + k_axis * real), + 1 - scale * (i_axis * i_axis + k_axis * k_axis), + scale * (j_axis * k_axis - i_axis * real), + scale * (i_axis * k_axis - j_axis * real), + scale * (j_axis * k_axis + i_axis * real), + 1 - scale * (i_axis * i_axis + j_axis * j_axis), + ) + # (..., 3, 3) + return torch.stack(entries, dim=-1).reshape((*quaternions.shape[:-1], 3, 3)) + + +def random_quaternions( + n: int, + dtype: torch.dtype | None = None, + device: torch.device | str | None = None, +) -> Tensor: + """Draw ``n`` uniformly distributed unit quaternions with nonnegative real part.""" + + resolved_device = torch.device(device) if isinstance(device, str) else device + samples = torch.randn((n, 4), dtype=dtype, device=resolved_device) # (n, 4) + squared_norm = (samples * samples).sum(dim=1) # (n,) + signed_norm = _copysign(torch.sqrt(squared_norm), samples[:, 0]) # (n,) + return samples / signed_norm[:, None] # (n, 4) + + +def random_rotations( + n: int, + dtype: torch.dtype | None = None, + device: torch.device | str | None = None, +) -> Tensor: + """Draw rotation tensor R with shape ``(n, 3, 3)``.""" + + # (n, 3, 3) + return quaternion_to_matrix(random_quaternions(n, dtype=dtype, device=device)) + + +def compute_random_augmentation( + multiplicity: int, + s_trans: float = 1.0, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> tuple[Tensor, Tensor]: + """Draw independent rotations R and translations T for a replicated batch.""" + + rotations = random_rotations(multiplicity, dtype=dtype, device=device) # (m, 3, 3) + translations = ( # (m, 1, 3) + torch.randn( + (multiplicity, 1, 3), + dtype=dtype, + device=device, + ) + * s_trans + ) + return rotations, translations # (m, 3, 3), (m, 1, 3) + + +def randomly_rotate( + coords: Tensor, + return_second_coords: bool = False, + second_coords: Tensor | None = None, +) -> Tensor | tuple[Tensor, Tensor | None]: + """Rotate X and optionally Y using the same sampled rotation tensor R.""" + + # coords: (b, n, 3); second_coords: (b, n_second, 3) or None. + rotations = random_rotations(len(coords), coords.dtype, coords.device) # (b, 3, 3) + rotated = torch.einsum("bmd,bds->bms", coords, rotations) # (b, n, 3) + if not return_second_coords: + return rotated # (b, n, 3) + rotated_second = ( # (b, n_second, 3) or None + None if second_coords is None else torch.einsum("bmd,bds->bms", second_coords, rotations) + ) + return rotated, rotated_second # (b, n, 3), (b, n_second, 3) or None + + +def center_random_augmentation( + atom_coords: Tensor, + atom_mask: Tensor, + s_trans: float = 1.0, + augmentation: bool = True, + centering: bool = True, + return_second_coords: bool = False, + second_coords: Tensor | None = None, +) -> Tensor | tuple[Tensor, Tensor | None]: + """Center and rigidly augment coordinate tensors X and optional Y.""" + + # atom_coords: (b, n_atom, 3); atom_mask: (b, n_atom). + # second_coords: (b, n_second, 3) or None. + primary = atom_coords # (b, n_atom, 3) + secondary = second_coords # (b, n_second, 3) or None + if centering: + centroid = _masked_center(primary, atom_mask) # (b, 1, 3) + primary = primary - centroid # (b, n_atom, 3) + if secondary is not None: + secondary = secondary - centroid # (b, n_second, 3) + + if augmentation: + primary, secondary = randomly_rotate( # (b, n_atom, 3), (b, n_second, 3) or None + primary, + return_second_coords=True, + second_coords=secondary, + ) + translation = torch.randn_like(primary[:, 0:1, :]) * s_trans # (b, 1, 3) + primary = primary + translation # (b, n_atom, 3) + if secondary is not None: + secondary = secondary + translation # (b, n_second, 3) + + # (b, n_atom, 3), optionally paired with (b, n_second, 3) or None. + return (primary, secondary) if return_second_coords else primary + + +class ExponentialMovingAverage: + """Maintain detached exponential moving averages of trainable parameters.""" + + def __init__( + self, + parameters: Iterable[nn.Parameter], + decay: float, + use_num_updates: bool = True, + ) -> None: + # parameters: one independently shaped tensor (...) per trainable parameter. + if not 0.0 <= decay <= 1.0: + raise ValueError("decay must lie in [0, 1]") + self.decay = decay + self.num_updates: int | None = 0 if use_num_updates else None + self.shadow_params = [ # each: same shape as one trainable parameter, (...) + parameter.clone().detach() for parameter in parameters if parameter.requires_grad + ] + self.collected_params: list[Tensor] = [] + + @staticmethod + def _trainable(parameters: Iterable[nn.Parameter]) -> list[nn.Parameter]: + # Each item is one trainable parameter with an independently varying shape (...). + return [parameter for parameter in parameters if parameter.requires_grad] + + def update(self, parameters: Iterable[nn.Parameter]) -> None: + # parameters: one independently shaped tensor (...) per trainable parameter. + trainable = self._trainable(parameters) # each: one trainable parameter, (...) + if len(trainable) != len(self.shadow_params): + raise ValueError("EMA parameter count changed") + decay = self.decay + if self.num_updates is not None: + self.num_updates += 1 + decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates)) + update_weight = 1.0 - decay + with torch.no_grad(): + for shadow, parameter in zip(self.shadow_params, trainable, strict=True): + shadow.sub_(update_weight * (shadow - parameter)) # (...) unchanged in place + + def compatible(self, parameters: Sequence[Tensor]) -> bool: + """Return whether parameter count and tensor shapes match the EMA state.""" + + # Each shadow and corresponding parameter has an independently varying shape (...). + return len(self.shadow_params) == len(parameters) and all( + shadow.shape == parameter.shape + for shadow, parameter in zip(self.shadow_params, parameters, strict=True) + ) + + def copy_to(self, parameters: Iterable[nn.Parameter]) -> None: + # parameters: one independently shaped tensor (...) per trainable parameter. + trainable = self._trainable(parameters) # each: one trainable parameter, (...) + if len(trainable) != len(self.shadow_params): + raise ValueError("EMA parameter count changed") + for shadow, parameter in zip(self.shadow_params, trainable, strict=True): + parameter.data.copy_(shadow.data) # (...) unchanged in place + + def store(self, parameters: Iterable[nn.Parameter]) -> None: + # parameters: one independently shaped tensor (...) per parameter. + # Each clone has the same shape (...) as its corresponding parameter. + self.collected_params = [parameter.clone() for parameter in parameters] + + def restore(self, parameters: Iterable[nn.Parameter]) -> None: + # parameters: one independently shaped tensor (...) per parameter. + current = list(parameters) # each: one parameter, (...) + if len(current) != len(self.collected_params): + raise ValueError("stored EMA parameter count changed") + for collected, parameter in zip(self.collected_params, current, strict=True): + parameter.data.copy_(collected.data) # (...) unchanged in place + + def state_dict(self) -> dict[str, Any]: + # shadow_params contains independently shaped tensors (...). + return { + "decay": self.decay, + "num_updates": self.num_updates, + "shadow_params": self.shadow_params, + } + + def load_state_dict( + self, + state_dict: Mapping[str, Any], + device: torch.device | str, + ) -> None: + # state_dict["shadow_params"] contains independently shaped tensors (...). + self.decay = float(state_dict["decay"]) + updates = state_dict["num_updates"] + self.num_updates = None if updates is None else int(updates) + # Each tensor retains its shape (...). + self.shadow_params = [tensor.to(device) for tensor in state_dict["shadow_params"]] + + def to(self, device: torch.device | str) -> None: + # Each shadow parameter retains its independently varying shape (...). + self.shadow_params = [tensor.to(device) for tensor in self.shadow_params] diff --git a/src/fastplms/models/boltz/vb_potentials_potentials.py b/src/fastplms/models/boltz/vb_potentials_potentials.py new file mode 100644 index 0000000..c51a6d2 --- /dev/null +++ b/src/fastplms/models/boltz/vb_potentials_potentials.py @@ -0,0 +1,895 @@ +"""Differentiable structure-steering potentials for Boltz2 diffusion. + +Each potential follows the same mechanism: derive constraint indices and +bounds from features, evaluate a geometric variable, map that variable to an +energy, and optionally scatter its analytic derivative back to atom +coordinates. This module is maintained independently from the pinned parity +oracle while preserving the public class names used by converted checkpoints. +""" + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + +from . import vb_const as const +from .vb_loss_diffusionv2 import weighted_rigid_align +from .vb_potentials_schedules import ( + ExponentialInterpolation, + ParameterSchedule, + PiecewiseStepFunction, +) + + +Parameter = ParameterSchedule | float | int | bool +ParameterMap = dict[str, Parameter] + + +@dataclass(frozen=True, slots=True) +class _PreparedPotential: + coordinates: torch.Tensor # (..., a, 3) + index: torch.Tensor # (q, n) + function_args: tuple[Any, ...] + com_index: torch.Tensor | None # (a_input,) + atom_pad_mask: torch.Tensor | None # (a_input,) + ref_coords: torch.Tensor | None # (n_ref, n, 3) + ref_mask: torch.Tensor | None # (n_ref, n) + ref_token_index: torch.Tensor | None # (a_input,) + negation_mask: torch.Tensor | None # (n,) + union_index: torch.Tensor | None # (n,) + + +def _reduce_atoms_to_centers( + coordinates: torch.Tensor, + com_index: torch.Tensor, + atom_pad_mask: torch.Tensor, +) -> torch.Tensor: + # coordinates: (..., a, 3); com_index/atom_pad_mask: (a,). + unpadded_index = com_index[atom_pad_mask] # (a_valid,) + unpadded_coordinates = coordinates[..., atom_pad_mask, :] # (..., a_valid, 3) + center_shape = (*unpadded_coordinates.shape[:-2], unpadded_index.max() + 1, 3) + return torch.zeros(center_shape, device=coordinates.device).scatter_reduce( + -2, + unpadded_index.unsqueeze(-1).expand_as(unpadded_coordinates), + unpadded_coordinates, + "mean", + ) # (..., n_center, 3) + + +def _union_weights( + energy: torch.Tensor, + union_index: torch.Tensor, + union_lambda: float, +) -> torch.Tensor: + # energy: (..., n); union_index: (n,). + unnormalized = torch.exp(-union_lambda * energy) # (..., n) + partition = torch.zeros( + (*energy.shape[:-1], union_index.max() + 1), + device=union_index.device, + ).scatter_reduce( + -1, + union_index.expand_as(unnormalized), + unnormalized, + "sum", + ) # (..., n_union) + weights = unnormalized / partition[..., union_index] # (..., n) + weights[partition[..., union_index] == 0] = 0 # (..., n) + return weights # (..., n) + + +def _scatter_constraint_gradients( + coefficients: torch.Tensor, + variable_gradient: torch.Tensor, + index: torch.Tensor, + coordinates: torch.Tensor, +) -> torch.Tensor: + # coefficients: (..., n); variable_gradient: (..., q, n, 3). + # index: (q, n); coordinates: (b, a, 3). + product = coefficients.tile(variable_gradient.shape[-3]).unsqueeze( + -1 + ) * variable_gradient.flatten(start_dim=-3, end_dim=-2) # (..., q * n, 3) + if product.dim() > 3: + product = product.sum(dim=list(range(1, product.dim() - 2))) # (b, q * n, 3) + scatter_index = ( + index.flatten(start_dim=0, end_dim=1).unsqueeze(-1).expand((*coordinates.shape[:-2], -1, 3)) + ) # (b, q * n, 3) + return torch.zeros_like(coordinates).scatter_reduce( + -2, + scatter_index, + product, + "sum", + ) # (b, a, 3) + + +class Potential(ABC): + """Base contract for energy and analytic-gradient steering potentials.""" + + def __init__(self, parameters: ParameterMap | None = None) -> None: + self.parameters = parameters + + def _prepare( + self, + coordinates: torch.Tensor, + feats: dict[str, torch.Tensor], + parameters: dict[str, Any], + computed_args: tuple[Any, ...] | None = None, + ) -> _PreparedPotential: + # coordinates: (..., a_input, 3). + if computed_args is None: + computed_args = self.compute_args(feats, parameters) + index, args, com_args, ref_args, operator_args = computed_args # index: (q, n) + com_index = atom_pad_mask = None + if com_args is not None: + com_index, atom_pad_mask = com_args # each (a_input,) + coordinates = _reduce_atoms_to_centers( + coordinates, + com_index, + atom_pad_mask, + ) # (..., n_center, 3) + + ref_coords = ref_mask = ref_token_index = None + if ref_args is not None: + # ref_coords: (n_ref, n, 3); ref_mask: (n_ref, n). + # ref_atom_index: (n,); ref_token_index: (a_input,). + ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args + coordinates = coordinates[..., ref_atom_index, :] # (..., n, 3) + + negation_mask = union_index = None + if operator_args is not None: + negation_mask, union_index = operator_args # each (n,) + return _PreparedPotential( + coordinates=coordinates, + index=index, + function_args=args, + com_index=com_index, + atom_pad_mask=atom_pad_mask, + ref_coords=ref_coords, + ref_mask=ref_mask, + ref_token_index=ref_token_index, + negation_mask=negation_mask, + union_index=union_index, + ) + + def compute( + self, + coords: torch.Tensor, + feats: dict[str, torch.Tensor], + parameters: dict[str, Any], + ) -> torch.Tensor: + """Evaluate one energy per coordinate sample.""" + + # coords: (..., a, 3). + computed_args = self.compute_args(feats, parameters) + if computed_args[0].shape[1] == 0: + return torch.zeros(coords.shape[:-2], device=coords.device) # (...) + prepared = self._prepare(coords, feats, parameters, computed_args) + value = self.compute_variable( + prepared.coordinates, + prepared.index, + ref_coords=prepared.ref_coords, + ref_mask=prepared.ref_mask, + compute_gradient=False, + ) # (..., n) + energy = self.compute_function( + value, + *prepared.function_args, + negation_mask=prepared.negation_mask, + compute_derivative=False, + ) # (..., n) + if prepared.union_index is not None: + weights = _union_weights( + energy, + prepared.union_index, + parameters["union_lambda"], + ) # (..., n) + return (energy * weights).sum(dim=-1) # (...) + return energy.sum(dim=tuple(range(1, energy.dim()))) # (b,) + + def compute_gradient( + self, + coords: torch.Tensor, + feats: dict[str, torch.Tensor], + parameters: dict[str, Any], + ) -> torch.Tensor: + """Return the analytic coordinate gradient of the potential energy.""" + + # coords: (b, a, 3). + computed_args = self.compute_args(feats, parameters) + if computed_args[0].shape[1] == 0: + return torch.zeros_like(coords) # (b, a, 3) + prepared = self._prepare(coords, feats, parameters, computed_args) + value, variable_gradient = self.compute_variable( + prepared.coordinates, + prepared.index, + ref_coords=prepared.ref_coords, + ref_mask=prepared.ref_mask, + compute_gradient=True, + ) # value: (..., n); variable_gradient: (..., q, n, 3) + energy, energy_derivative = self.compute_function( + value, + *prepared.function_args, + negation_mask=prepared.negation_mask, + compute_derivative=True, + ) # each (..., n) + if prepared.union_index is not None: + weights = _union_weights( + energy, + prepared.union_index, + parameters["union_lambda"], + ) # (..., n) + union_energy = torch.zeros( + (*energy.shape[:-1], prepared.union_index.max() + 1), + device=prepared.union_index.device, + ).scatter_reduce( + -1, + prepared.union_index.expand_as(energy), + energy * weights, + "sum", + ) # (..., n_union) + coefficients = ( + energy_derivative + * weights + * ( + 1 + + parameters["union_lambda"] + * (energy - union_energy[..., prepared.union_index]) + ) + ) # (..., n) + else: + coefficients = energy_derivative # (..., n) + + atom_gradient = _scatter_constraint_gradients( + coefficients, + variable_gradient, + prepared.index, + prepared.coordinates, + ) # (b, a_prepared, 3) + if prepared.com_index is not None: + atom_gradient = atom_gradient[..., prepared.com_index, :] # (b, a, 3) + elif prepared.ref_token_index is not None: + atom_gradient = atom_gradient[..., prepared.ref_token_index, :] # (b, a, 3) + return atom_gradient # (b, a, 3) + + def compute_parameters(self, t: float) -> dict[str, Any] | None: + """Resolve scheduled parameters at diffusion time ``t``.""" + + if self.parameters is None: + return None + return { + name: parameter.compute(t) if isinstance(parameter, ParameterSchedule) else parameter + for name, parameter in self.parameters.items() + } + + @abstractmethod + def compute_function( + self, + value: torch.Tensor, + *args: Any, + negation_mask: torch.Tensor | None = None, + compute_derivative: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # value: (..., n); negation_mask: (n,). + raise NotImplementedError + + @abstractmethod + def compute_variable( + self, + coords: torch.Tensor, + index: torch.Tensor, + ref_coords: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + compute_gradient: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # coords: (..., a, 3); index: (q, n). + raise NotImplementedError + + @abstractmethod + def compute_args( + self, + feats: dict[str, torch.Tensor], + parameters: dict[str, Any], + ) -> tuple[Any, ...]: + raise NotImplementedError + + def get_reference_coords( + self, + feats: dict[str, torch.Tensor], + parameters: dict[str, Any], + ) -> tuple[None, None]: + del feats, parameters + return None, None + + +class FlatBottomPotential(Potential): + """Linear penalty outside an allowed lower-to-upper interval.""" + + def compute_function( + self, + value: torch.Tensor, + k: torch.Tensor, + lower_bounds: torch.Tensor | None, + upper_bounds: torch.Tensor | None, + negation_mask: torch.Tensor | None = None, + compute_derivative: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # value: (..., n); k/lower_bounds/upper_bounds: (n,) or (..., n). + # negation_mask: (n,) or a shape broadcastable to (..., n). + lower = ( + (torch.full_like(value, float("-inf")) if lower_bounds is None else lower_bounds) + .expand_as(value) + .clone() + ) # (..., n) + upper = ( + (torch.full_like(value, float("inf")) if upper_bounds is None else upper_bounds) + .expand_as(value) + .clone() + ) # (..., n) + + if negation_mask is not None: + if not torch.is_tensor(negation_mask): + raise TypeError( + f"negation_mask must be a boolean tensor, got {type(negation_mask).__name__}." + ) + if negation_mask.dtype != torch.bool: + raise TypeError( + f"negation_mask must be a boolean tensor, got {negation_mask.dtype}." + ) + if negation_mask.device != value.device: + raise ValueError( + "negation_mask must be on the same device as value; " + f"got {negation_mask.device} and {value.device}." + ) + try: + expanded_negation_mask = negation_mask.expand_as(value) # (..., n) + except RuntimeError as error: + raise ValueError( + "negation_mask must be broadcastable to value shape " + f"{tuple(value.shape)}, got {tuple(negation_mask.shape)}." + ) from error + unbounded_below = torch.isneginf(lower) # (..., n) + unbounded_above = torch.isposinf(upper) # (..., n) + # (..., n) + valid_negation = unbounded_below | unbounded_above | expanded_negation_mask + if not bool(torch.all(valid_negation).item()): + raise ValueError( + "negation_mask may be false only where at least one bound is infinite." + ) + select_upper = ~unbounded_above & ~expanded_negation_mask # (..., n) + lower[select_upper] = upper[select_upper] # (..., n) + upper[select_upper] = float("inf") # (..., n) + select_lower = ~unbounded_below & ~expanded_negation_mask # (..., n) + upper[select_lower] = lower[select_lower] # (..., n) + lower[select_lower] = float("-inf") # (..., n) + + below = value < lower # (..., n) + above = value > upper # (..., n) + energy = torch.zeros_like(value) # (..., n) + energy[below] = (k * (lower - value))[below] # (..., n) + energy[above] = (k * (value - upper))[above] # (..., n) + if not compute_derivative: + return energy # (..., n) + derivative = torch.zeros_like(value) # (..., n) + derivative[below] = -k.expand_as(below)[below] # (..., n) + derivative[above] = k.expand_as(above)[above] # (..., n) + return energy, derivative # each (..., n) + + +class ReferencePotential(Potential): + """Measure atom displacement after weighted rigid alignment.""" + + def compute_variable( + self, + coords: torch.Tensor, + index: torch.Tensor, + ref_coords: torch.Tensor, + ref_mask: torch.Tensor, + compute_gradient: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # coords: (b, a, 3); index: (1, n). + # ref_coords: (n_ref, n, 3); ref_mask: (n_ref, n). + aligned_reference = weighted_rigid_align( + ref_coords.float(), + coords[:, index].float(), + ref_mask, + ref_mask, + ) # (b, n_ref, n, 3) + displacement = coords[:, index] - aligned_reference # (b, n_ref, n, 3) + distance = torch.linalg.norm(displacement, dim=-1) # (b, n_ref, n) + if not compute_gradient: + return distance # (b, n_ref, n) + unit_displacement = displacement / distance.unsqueeze(-1) # (b, n_ref, n, 3) + # (b, 1, n_ref, n, 3) + gradient = (unit_displacement * ref_mask.unsqueeze(-1)).unsqueeze(1) + return distance, gradient # (b, n_ref, n); (b, 1, n_ref, n, 3) + + +class DistancePotential(Potential): + """Measure Euclidean distances for indexed atom pairs.""" + + def compute_variable( + self, + coords: torch.Tensor, + index: torch.Tensor, + ref_coords: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + compute_gradient: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # coords: (..., a, 3); index: (2, n). + del ref_coords, ref_mask + displacement = coords.index_select(-2, index[0]) - coords.index_select( + -2, + index[1], + ) # (..., n, 3) + distance = torch.linalg.norm(displacement, dim=-1) # (..., n) + if not compute_gradient: + return distance # (..., n) + unit_displacement = displacement / distance.unsqueeze(-1) # (..., n, 3) + # Returns (..., n); (..., 2, n, 3). + return distance, torch.stack((unit_displacement, -unit_displacement), dim=1) + + +class DihedralPotential(Potential): + """Measure signed torsion angles for indexed atom quartets.""" + + def compute_variable( + self, + coords: torch.Tensor, + index: torch.Tensor, + ref_coords: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + compute_gradient: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # coords: (..., a, 3); index: (4, n). + del ref_coords, ref_mask + # Each displacement has shape (..., n, 3). + r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1]) + r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1]) + r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3]) + n_ijk = torch.cross(r_ij, r_kj, dim=-1) # (..., n, 3) + n_jkl = torch.cross(r_kj, r_kl, dim=-1) # (..., n, 3) + r_kj_norm = torch.linalg.norm(r_kj, dim=-1) # (..., n) + n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1) # (..., n) + n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1) # (..., n) + + orientation = torch.sign( + r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1) + ).squeeze(-1, -2) # (..., n) + cosine = (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2) / ( + n_ijk_norm * n_jkl_norm + ) # (..., n) + # (..., n) + angle = orientation * torch.arccos(torch.clamp(cosine, -1 + 1e-8, 1 - 1e-8)) + if not compute_gradient: + return angle # (..., n) + + projection_i = ( + (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) + ).unsqueeze(-1) # (..., n, 1) + projection_l = ( + (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) + ).unsqueeze(-1) # (..., n, 1) + grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1) # (..., n, 3) + grad_l = -n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1) # (..., n, 3) + grad_j = (projection_i - 1) * grad_i - projection_l * grad_l # (..., n, 3) + grad_k = (projection_l - 1) * grad_l - projection_i * grad_i # (..., n, 3) + # Returns (..., n); (..., 4, n, 3). + return angle, torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1) + + +class AbsDihedralPotential(DihedralPotential): + """Measure the unsigned magnitude of indexed torsion angles.""" + + def compute_variable( + self, + coords: torch.Tensor, + index: torch.Tensor, + ref_coords: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + compute_gradient: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # coords: (..., a, 3); index: (4, n). + del ref_coords, ref_mask + if not compute_gradient: + return torch.abs(super().compute_variable(coords, index)) # (..., n) + angle, gradient = super().compute_variable( + coords, + index, + compute_gradient=True, + ) # (..., n); (..., 4, n, 3) + gradient[(angle < 0)[..., None, :, None].expand_as(gradient)] *= -1 # same + return torch.abs(angle), gradient # (..., n); (..., 4, n, 3) + + +def _element_radii(feats: dict[str, torch.Tensor]) -> torch.Tensor: + # feats["ref_element"]: (1, a, n_element). + element_radii = torch.zeros( + const.num_elements, + dtype=torch.float32, + device=feats["ref_element"].device, + ) # (n_element,) + element_radii[1:119] = torch.tensor( + const.vdw_radii, + dtype=torch.float32, + device=element_radii.device, + ) # (n_element,) + # (a,) + return (feats["ref_element"].float() @ element_radii.unsqueeze(-1)).squeeze(-1)[0] + + +def _atom_chain_ids(feats: dict[str, torch.Tensor]) -> torch.Tensor: + # atom_to_token: (1, a, l); asym_id: (1, l). + return ( + torch.bmm( + feats["atom_to_token"].float(), + feats["asym_id"].unsqueeze(-1).float(), + ) + .squeeze(-1) + .long() + )[0] # (a,) + + +class PoseBustersPotential(FlatBottomPotential, DistancePotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + # rdkit_bounds_index: (1, 2, n); bound and mask features: (1, n). + pair_index = feats["rdkit_bounds_index"][0] # (2, n) + lower = feats["rdkit_lower_bounds"][0].clone() # (n,) + upper = feats["rdkit_upper_bounds"][0].clone() # (n,) + bond = feats["rdkit_bounds_bond_mask"][0] # (n,) + angle = feats["rdkit_bounds_angle_mask"][0] # (n,) + lower[bond * ~angle] *= 1.0 - parameters["bond_buffer"] # (n,) + upper[bond * ~angle] *= 1.0 + parameters["bond_buffer"] # (n,) + lower[~bond * angle] *= 1.0 - parameters["angle_buffer"] # (n,) + upper[~bond * angle] *= 1.0 + parameters["angle_buffer"] # (n,) + shared_buffer = min(parameters["bond_buffer"], parameters["angle_buffer"]) + lower[bond * angle] *= 1.0 - shared_buffer # (n,) + upper[bond * angle] *= 1.0 + shared_buffer # (n,) + lower[~bond * ~angle] *= 1.0 - parameters["clash_buffer"] # (n,) + upper[~bond * ~angle] = float("inf") # (n,) + + atom_radii = _element_radii(feats) # (a,) + bond_cutoff = 0.35 + atom_radii[pair_index].mean(dim=0) # (n,) + lower[~bond] = torch.max(lower[~bond], bond_cutoff[~bond]) # (n,) + upper[bond] = torch.min(upper[bond], bond_cutoff[bond]) # (n,) + # Returns index (2, n) and three function arguments of shape (n,). + return pair_index, (torch.ones_like(lower), lower, upper), None, None, None + + +class ConnectionsPotential(FlatBottomPotential, DistancePotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + # connected_atom_index: (1, 2, n). + pair_index = feats["connected_atom_index"][0] # (2, n) + upper = torch.full( + (pair_index.shape[1],), + parameters["buffer"], + device=pair_index.device, + ) # (n,) + # Returns index (2, n), k (n,), and upper bound (n,). + return pair_index, (torch.ones_like(upper), None, upper), None, None, None + + +class VDWOverlapPotential(FlatBottomPotential, DistancePotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + atom_chain_id = _atom_chain_ids(feats) # (a,) + atom_pad_mask = feats["atom_pad_mask"][0].bool() # (a,) + chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask]) # (n_chain,) + nonion_atom = (chain_sizes > 1)[atom_chain_id] # (a,) + atom_radii = _element_radii(feats) # (a,) + pair_index = torch.triu_indices( + atom_chain_id.shape[0], + atom_chain_id.shape[0], + 1, + device=atom_chain_id.device, + ) # (2, n_pair) + pair_pad_mask = atom_pad_mask[pair_index].all(dim=0) # (n_pair,) + # (n_pair,) + pair_ion_mask = nonion_atom[pair_index[0]] * nonion_atom[pair_index[1]] + + num_chains = atom_chain_id.max() + 1 # () + connected = feats["connected_chain_index"][0] # (2, n_connection) + connected_matrix = torch.eye( + num_chains, + device=atom_chain_id.device, + dtype=torch.bool, + ) # (n_chain, n_chain) + connected_matrix[connected[0], connected[1]] = True # (n_chain, n_chain) + connected_matrix[connected[1], connected[0]] = True # (n_chain, n_chain) + connected_pair = connected_matrix[ + atom_chain_id[pair_index[0]], + atom_chain_id[pair_index[1]], + ] # (n_pair,) + # (2, n) + pair_index = pair_index[:, pair_pad_mask * pair_ion_mask * ~connected_pair] + # (n,) + lower = atom_radii[pair_index].sum(dim=0) * (1.0 - parameters["buffer"]) + # Returns index (2, n), k (n,), and lower bound (n,). + return pair_index, (torch.ones_like(lower), lower, None), None, None, None + + +class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + atom_chain_id = _atom_chain_ids(feats) # (a,) + atom_pad_mask = feats["atom_pad_mask"][0].bool() # (a,) + nonion_chain = torch.bincount(atom_chain_id[atom_pad_mask]) > 1 # (n_chain,) + pair_index = feats["symmetric_chain_index"][0] # (2, n_candidate) + pair_index = pair_index[ + :, + nonion_chain[pair_index[0]] * nonion_chain[pair_index[1]], + ] # (2, n) + lower = torch.full( + (pair_index.shape[1],), + parameters["buffer"], + dtype=torch.float32, + device=pair_index.device, + ) # (n,) + # Returns center index (2, n), k/lower (n,), and atom-to-center maps (a,). + return ( + pair_index, + (torch.ones_like(lower), lower, None), + (atom_chain_id, atom_pad_mask), + None, + None, + ) + + +def _oriented_bounds( + orientations: torch.Tensor, + positive_lower: float, + negative_upper: float, +) -> tuple[torch.Tensor, torch.Tensor]: + # orientations: (n,). + lower = torch.zeros(orientations.shape, device=orientations.device) # (n,) + upper = torch.zeros(orientations.shape, device=orientations.device) # (n,) + lower[orientations] = positive_lower # (n,) + upper[orientations] = float("inf") # (n,) + lower[~orientations] = float("-inf") # (n,) + upper[~orientations] = negative_upper # (n,) + return lower, upper # each (n,) + + +class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + # stereo_bond_index: (1, 4, n); orientations: (1, n). + index = feats["stereo_bond_index"][0] # (4, n) + orientation = feats["stereo_bond_orientations"][0].bool() # (n,) + lower, upper = _oriented_bounds( + orientation, + torch.pi - parameters["buffer"], + parameters["buffer"], + ) # each (n,) + # Returns index (4, n) and three function arguments of shape (n,). + return index, (torch.ones_like(lower), lower, upper), None, None, None + + +class ChiralAtomPotential(FlatBottomPotential, DihedralPotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + # chiral_atom_index: (1, 4, n); orientations: (1, n). + index = feats["chiral_atom_index"][0] # (4, n) + orientation = feats["chiral_atom_orientations"][0].bool() # (n,) + lower, upper = _oriented_bounds( + orientation, + parameters["buffer"], + -parameters["buffer"], + ) # each (n,) + # Returns index (4, n) and three function arguments of shape (n,). + return index, (torch.ones_like(lower), lower, upper), None, None, None + + +class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + # The feature stores two atom-index rows for each six-entry bond pattern. + bond_index = feats["planar_bond_index"][0].T # (2, n_bond_entry) + improper_pattern = torch.tensor( + [[1, 2, 3, 0], [4, 5, 0, 3]], + device=bond_index.device, + ).T # (4, 2) + # (4, n_improper) + improper_index = bond_index[:, improper_pattern].swapaxes(0, 1).flatten(start_dim=1) + upper = torch.full( + (improper_index.shape[1],), + parameters["buffer"], + device=improper_index.device, + ) # (n_improper,) + # Returns index (4, n_improper), k (n_improper,), and upper (n_improper,). + return ( + improper_index, + (torch.ones_like(upper), None, upper), + None, + None, + None, + ) + + +class TemplateReferencePotential(FlatBottomPotential, ReferencePotential): + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + del parameters + if "template_mask_cb" not in feats or "template_force" not in feats: + # Empty sentinel index: (1, 0). + return torch.empty((1, 0)), None, None, None, None + # template_mask_cb/template_force: (b, n_template, n)/(b, n_template). + template_mask = feats["template_mask_cb"][feats["template_force"]] # (n_ref, n) + if template_mask.shape[0] == 0: + # Empty sentinel index: (1, 0). + return torch.empty((1, 0)), None, None, None, None + + ref_coords = feats["template_cb"][feats["template_force"]].clone() # (n_ref, n, 3) + ref_mask = feats["template_mask_cb"][feats["template_force"]].clone() # (n_ref, n) + atom_indices = torch.arange( + feats["atom_pad_mask"].shape[1], + device=feats["atom_pad_mask"].device, + dtype=torch.float32, + )[None, :, None] # (1, a, 1) + ref_atom_index = ( + torch.bmm( + feats["token_to_rep_atom"].float(), + atom_indices, + ) + .squeeze(-1) + .long()[0] # (n,) + ) + ref_token_index = ( + torch.bmm( + feats["atom_to_token"].float(), + feats["token_index"].unsqueeze(-1).float(), + ) + .squeeze(-1) + .long()[0] # (a,) + ) + + index = torch.arange( + template_mask.shape[-1], + dtype=torch.long, + device=template_mask.device, + )[None] # (1, n) + upper = torch.full( + template_mask.shape, + float("inf"), + device=index.device, + dtype=torch.float32, + ) # (n_ref, n) + reference_indices = torch.argwhere(template_mask).T # (2, n_active) + upper[reference_indices.unbind()] = feats["template_force_threshold"][ + feats["template_force"] + ][reference_indices[0]] # (n_ref, n) + # Returns index (1, n), function args (n_ref, n), and reference mappings. + return ( + index, + (torch.ones_like(upper), None, upper), + None, + (ref_coords, ref_mask, ref_atom_index, ref_token_index), + None, + ) + + +class ContactPotentital(FlatBottomPotential, DistancePotential): + """Contact-union potential retaining the historical checkpoint name.""" + + def compute_args(self, feats: dict[str, torch.Tensor], parameters: dict[str, Any]): + del parameters + # pair index: (1, 2, n); threshold/operator features: (1, n). + index = feats["contact_pair_index"][0] # (2, n) + upper = feats["contact_thresholds"][0].clone() # (n,) + # Returns index (2, n), k/upper (n,), and two operator tensors (n,). + return ( + index, + (torch.ones_like(upper), None, upper), + None, + None, + ( + feats["contact_negation_mask"][0], + feats["contact_union_index"][0], + ), + ) + + +ContactPotential = ContactPotentital + + +def get_potentials( + steering_args: dict[str, bool], + boltz2: bool = False, +) -> list[Potential]: + """Build the ordered potential stack for requested steering modes.""" + + use_fk = steering_args["fk_steering"] + use_physical = steering_args["physical_guidance_update"] + use_contacts = steering_args.get("contact_guidance_update", False) + potentials: list[Potential] = [] + if use_fk or use_physical: + potentials.extend( + ( + SymmetricChainCOMPotential( + { + "guidance_interval": 4, + "guidance_weight": 0.5 if use_physical else 0.0, + "resampling_weight": 0.5, + "buffer": ExponentialInterpolation(1.0, 5.0, -2.0), + } + ), + VDWOverlapPotential( + { + "guidance_interval": 5, + "guidance_weight": PiecewiseStepFunction( + [0.4], + [0.125, 0.0], + ) + if use_physical + else 0.0, + "resampling_weight": PiecewiseStepFunction( + [0.6], + [0.01, 0.0], + ), + "buffer": 0.225, + } + ), + ConnectionsPotential( + { + "guidance_interval": 1, + "guidance_weight": 0.15 if use_physical else 0.0, + "resampling_weight": 1.0, + "buffer": 2.0, + } + ), + PoseBustersPotential( + { + "guidance_interval": 1, + "guidance_weight": 0.01 if use_physical else 0.0, + "resampling_weight": 0.1, + "bond_buffer": 0.125, + "angle_buffer": 0.125, + "clash_buffer": 0.10, + } + ), + ChiralAtomPotential( + { + "guidance_interval": 1, + "guidance_weight": 0.1 if use_physical else 0.0, + "resampling_weight": 1.0, + "buffer": 0.52360, + } + ), + StereoBondPotential( + { + "guidance_interval": 1, + "guidance_weight": 0.05 if use_physical else 0.0, + "resampling_weight": 1.0, + "buffer": 0.52360, + } + ), + PlanarBondPotential( + { + "guidance_interval": 1, + "guidance_weight": 0.05 if use_physical else 0.0, + "resampling_weight": 1.0, + "buffer": 0.26180, + } + ), + ) + ) + if boltz2 and (use_fk or use_contacts): + potentials.extend( + ( + ContactPotentital( + { + "guidance_interval": 4, + "guidance_weight": PiecewiseStepFunction( + [0.25, 0.75], + [0.0, 0.5, 1.0], + ) + if use_contacts + else 0.0, + "resampling_weight": 1.0, + "union_lambda": ExponentialInterpolation(8.0, 0.0, -2.0), + } + ), + TemplateReferencePotential( + { + "guidance_interval": 2, + "guidance_weight": 0.1 if use_contacts else 0.0, + "resampling_weight": 1.0, + } + ), + ) + ) + return potentials diff --git a/src/fastplms/models/boltz/vb_potentials_schedules.py b/src/fastplms/models/boltz/vb_potentials_schedules.py new file mode 100644 index 0000000..80a6658 --- /dev/null +++ b/src/fastplms/models/boltz/vb_potentials_schedules.py @@ -0,0 +1,67 @@ +"""Scalar schedules for structure-steering potentials.""" + +from __future__ import annotations + +import math +from abc import ABC, abstractmethod +from collections.abc import Sequence + + +class ParameterSchedule(ABC): + """Map normalized diffusion time ``t`` to a potential parameter.""" + + @abstractmethod + def compute(self, t: float) -> float: + """Evaluate the schedule at ``t``.""" + + +class ExponentialInterpolation(ParameterSchedule): + """Interpolate from ``start`` to ``end`` with exponential curvature.""" + + def __init__(self, start: float, end: float, alpha: float) -> None: + self.start = start + self.end = end + self.alpha = alpha + + def compute(self, t: float) -> float: + span = self.end - self.start + if self.alpha == 0: + return self.start + span * t + numerator = math.exp(self.alpha * t) - 1 + denominator = math.exp(self.alpha) - 1 + return self.start + span * numerator / denominator + + +class PiecewiseStepFunction(ParameterSchedule): + """Select values separated by strict upper thresholds. + + A time exactly equal to a threshold remains in the lower interval. This + boundary convention is part of the steering-input contract. + """ + + def __init__( + self, + thresholds: Sequence[float], + values: Sequence[float], + ) -> None: + self.thresholds = tuple(thresholds) + self.values = tuple(values) + if not self.thresholds: + raise ValueError("PiecewiseStepFunction requires at least one threshold.") + if len(self.values) != len(self.thresholds) + 1: + raise ValueError( + "PiecewiseStepFunction requires exactly one more value than threshold; " + f"received {len(self.values)} values and {len(self.thresholds)} thresholds." + ) + if any( + current >= following + for current, following in zip(self.thresholds, self.thresholds[1:], strict=False) + ): + raise ValueError("PiecewiseStepFunction thresholds must be strictly increasing.") + + def compute(self, t: float) -> float: + interval = next( + (index for index, threshold in enumerate(self.thresholds) if t <= threshold), + len(self.thresholds), + ) + return self.values[interval] diff --git a/src/fastplms/models/boltz/vb_tri_attn_attention.py b/src/fastplms/models/boltz/vb_tri_attn_attention.py new file mode 100644 index 0000000..5270a5b --- /dev/null +++ b/src/fastplms/models/boltz/vb_tri_attn_attention.py @@ -0,0 +1,150 @@ +"""Triangular row and column attention for Boltz2 pair states.""" + +from __future__ import annotations + +import torch +from typing import cast +from torch import Tensor, nn + +from .vb_tri_attn_primitives import Attention, LayerNorm, Linear +from .vb_tri_attn_utils import chunk_layer, permute_final_dims + + +class TriangleAttention(nn.Module): + """Apply pair-biased attention along one axis of pair tensor X.""" + + def __init__( + self, + c_in: int, + c_hidden: int, + no_heads: int, + starting: bool = True, + inf: float = 1e9, + ) -> None: + super().__init__() + self.c_in = c_in + self.c_hidden = c_hidden + self.no_heads = no_heads + self.starting = starting + self.inf = inf + self.layer_norm = LayerNorm(c_in) + self.linear = Linear(c_in, no_heads, bias=False, init="normal") + self.mha = Attention(c_in, c_in, c_in, c_hidden, no_heads) + + @torch.jit.ignore # type: ignore[untyped-decorator] + def _chunk( + self, + x: Tensor, + tri_bias: Tensor, + mask_bias: Tensor, + mask: Tensor, + chunk_size: int, + use_kernels: bool = False, + ) -> Tensor: + """Evaluate attention in slices of the batch-like leading axes.""" + + # x: (..., l, l, d); biases/mask retain broadcastable triangular axes. + def attention_call(**inputs: Tensor) -> Tensor: + return cast(Tensor, self.mha(**inputs, use_kernels=use_kernels)) + + inputs = { + "q_x": x, + "kv_x": x, + "tri_bias": tri_bias, + "mask_bias": mask_bias, + "mask": mask, + } # tensor values retain the input shapes above + return cast( + Tensor, + chunk_layer( + attention_call, + inputs, + chunk_size=chunk_size, + no_batch_dims=len(x.shape[:-2]), + _out=None, + ), + ) + + def _run_attention( + self, + states: Tensor, + triangle_bias: Tensor, + mask_bias: Tensor, + expanded_mask: Tensor, + chunk_size: int | None, + use_kernels: bool, + ) -> Tensor: + # states: (..., l, l, d); returned tensor has the same shape. + if chunk_size is not None and not use_kernels: + return cast( + Tensor, + self._chunk( + states, + triangle_bias, + mask_bias, + expanded_mask, + chunk_size, + use_kernels=False, + ), + ) + return cast( + Tensor, + self.mha( + states, + states, + triangle_bias, + mask_bias, + expanded_mask, + use_kernels=use_kernels, + ), + ) + + def forward( + self, + x: Tensor, + mask: Tensor | None = None, + chunk_size: int | None = None, + use_kernels: bool = False, + ) -> Tensor: + """Transform X with shape ``(..., l, l, d)`` under mask M.""" + + pair_mask = x.new_ones(x.shape[:-1]) if mask is None else mask # (..., l, l) + pair_states = x # (..., l, l, d) + if not self.starting: + pair_states = pair_states.transpose(-2, -3) # (..., l, l, d) + pair_mask = pair_mask.transpose(-1, -2) # (..., l, l) + + normalized = self.layer_norm(pair_states) # (..., l, l, d) + expanded_mask = pair_mask[..., :, None, None, :] # (..., l, 1, 1, l) + mask_bias = self.inf * (expanded_mask - 1) # (..., l, 1, 1, l) + triangle_bias = permute_final_dims( + self.linear(normalized), (2, 0, 1) + ) # (..., h, l, l) + triangle_bias = triangle_bias.unsqueeze(-4) # (..., 1, h, l, l) + output = self._run_attention( + normalized, + triangle_bias, + mask_bias, + expanded_mask, + chunk_size, + use_kernels, + ) # (..., l, l, d) + return ( + output if self.starting else output.transpose(-2, -3) + ) # (..., l, l, d) + + +TriangleAttentionStartingNode = TriangleAttention + + +class TriangleAttentionEndingNode(TriangleAttention): + """Apply triangular attention around each pair's ending node.""" + + def __init__( + self, + c_in: int, + c_hidden: int, + no_heads: int, + inf: float = 1e9, + ) -> None: + super().__init__(c_in, c_hidden, no_heads, starting=False, inf=inf) diff --git a/src/fastplms/models/boltz/vb_tri_attn_primitives.py b/src/fastplms/models/boltz/vb_tri_attn_primitives.py new file mode 100644 index 0000000..38b67a4 --- /dev/null +++ b/src/fastplms/models/boltz/vb_tri_attn_primitives.py @@ -0,0 +1,265 @@ +"""Attention primitives for the Boltz pair stack. + +The module retains the converted checkpoint names while providing a compact +implementation around PyTorch operations. Optional triangle kernels are +loaded only at the call site, so importing FastPLMs does not initialize CUDA. + +Portions of the numerical contract follow OpenFold and DeepMind AlphaFold +under Apache-2.0. See ``THIRD_PARTY_NOTICES.md`` for provenance. +""" + +from __future__ import annotations + +import importlib +import math +import torch +from collections.abc import Callable, Sequence +from importlib.util import find_spec +from torch import nn + +from . import vb_layers_initialize as initialize +from .vb_tri_attn_utils import flatten_final_dims, permute_final_dims + + +def _initialize_weight( + weight: torch.Tensor, + bias: torch.Tensor | None, + method: str, +) -> None: + # weight: (d_out, d_in); bias: (d_out,) or None; both retain shape in place. + initializers = { + "default": initialize.lecun_normal_init_, + "relu": initialize.he_normal_init_, + "glorot": initialize.glorot_uniform_init_, + "normal": initialize.normal_init_, + "final": initialize.final_init_, + "gating": initialize.gating_init_, + } + try: + initializers[method](weight) + except KeyError as error: + raise ValueError(f"unknown linear initialization {method!r}") from error + if method == "gating" and bias is not None: + bias.fill_(1.0) + + +class Linear(nn.Linear): + """Linear projection with the initializers used by converted checkpoints.""" + + def __init__( + self, + in_dim: int, + out_dim: int, + bias: bool = True, + init: str = "default", + init_fn: Callable[[torch.Tensor, torch.Tensor | None], None] | None = None, + precision: torch.dtype | None = None, + ) -> None: + super().__init__(in_dim, out_dim, bias=bias) + with torch.no_grad(): + if self.bias is not None: + self.bias.zero_() + if init_fn is None: + _initialize_weight(self.weight, self.bias, init) + else: + init_fn(self.weight, self.bias) + self.precision = precision + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Project ``input`` while preserving explicit precision policies.""" + + # input: (..., d_in); every branch returns (..., d_out). + input_dtype = input.dtype + if self.precision is not None: + with torch.autocast("cuda", enabled=False): + bias = ( + None if self.bias is None else self.bias.to(self.precision) + ) # (d_out,) or None + projected = nn.functional.linear( + input.to(self.precision), + self.weight.to(self.precision), + bias, + ) # (..., d_out) + return projected.to(input_dtype) # (..., d_out) + if input_dtype is torch.bfloat16: + with torch.autocast("cuda", enabled=False): + bias = ( + None if self.bias is None else self.bias.to(input_dtype) + ) # (d_out,) or None + return nn.functional.linear( + input, self.weight.to(input_dtype), bias + ) # (..., d_out) + return nn.functional.linear(input, self.weight, self.bias) # (..., d_out) + + +class LayerNorm(nn.Module): + """Layer normalization that avoids an autocast upcast for BF16 inputs.""" + + def __init__(self, c_in: int, eps: float = 1e-5) -> None: + super().__init__() + self.c_in = (c_in,) + self.eps = eps + self.weight = nn.Parameter(torch.ones(c_in)) # (c_in,) + self.bias = nn.Parameter(torch.zeros(c_in)) # (c_in,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (..., c_in); normalization preserves shape. + if x.dtype is torch.bfloat16: + with torch.autocast("cuda", enabled=False): + return nn.functional.layer_norm( + x, + self.c_in, + self.weight.to(x.dtype), + self.bias.to(x.dtype), + self.eps, + ) + return nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps) + + +@torch.jit.ignore +def softmax_no_cast(tensor: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Apply softmax without promoting BF16 inputs under CUDA autocast.""" + + if tensor.dtype is torch.bfloat16: + with torch.autocast("cuda", enabled=False): + return nn.functional.softmax(tensor, dim=dim) + return nn.functional.softmax(tensor, dim=dim) + + +def _attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + biases: Sequence[torch.Tensor], +) -> torch.Tensor: + """Compute biased scaled attention after query scaling.""" + + # query/key/value: (..., h, n_q_or_k, d_h). + scores = torch.matmul( + query, permute_final_dims(key, (1, 0)) + ) # (..., h, n_q, n_k) + for bias in biases: + scores += bias # (..., h, n_q, n_k) + probabilities = softmax_no_cast(scores, dim=-1) # (..., h, n_q, n_k) + return torch.matmul(probabilities, value) # (..., h, n_q, d_h) + + +@torch.compiler.disable +def kernel_triangular_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + tri_bias: torch.Tensor, + mask: torch.Tensor, + scale: float, +) -> torch.Tensor: + """Call the optional cuEquivariance triangle-attention primitive.""" + + # q/k/v: (..., h, n, d_h); returned tensor uses the kernel's matching layout. + if ( + find_spec("cuequivariance_torch") is None + or find_spec("cuequivariance_ops_torch") is None + ): + raise RuntimeError( + "Boltz2 use_kernels=True requires cuequivariance_torch and the CUDA 13 " + "cuequivariance_ops_torch runtime from the 'structure,cueq' extras." + ) + cueq = importlib.import_module("cuequivariance_torch") + return cueq.triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale) + + +class Attention(nn.Module): + """Multi-head pair attention with optional query-dependent output gates.""" + + def __init__( + self, + c_q: int, + c_k: int, + c_v: int, + c_hidden: int, + no_heads: int, + gating: bool = True, + ) -> None: + super().__init__() + self.c_q = c_q + self.c_k = c_k + self.c_v = c_v + self.c_hidden = c_hidden + self.no_heads = no_heads + self.gating = gating + + projected_dim = c_hidden * no_heads + self.linear_q = Linear(c_q, projected_dim, bias=False, init="glorot") + self.linear_k = Linear(c_k, projected_dim, bias=False, init="glorot") + self.linear_v = Linear(c_v, projected_dim, bias=False, init="glorot") + self.linear_o = Linear(projected_dim, c_q, bias=False, init="final") + self.linear_g = Linear(c_q, projected_dim, bias=False, init="gating") if gating else None + self.sigmoid = nn.Sigmoid() + + def _prep_qkv( + self, + q_x: torch.Tensor, + kv_x: torch.Tensor, + apply_scale: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project Q, K, and V to ``(..., h, n, d)`` tensors.""" + + def split_heads(X: torch.Tensor) -> torch.Tensor: + # X: (..., n, h * d_h) + return X.view(*X.shape[:-1], self.no_heads, -1).transpose( + -2, -3 + ) # (..., h, n, d_h) + + # q_x: (..., n_q, c_q); kv_x: (..., n_k, c_k). + q = split_heads(self.linear_q(q_x)) # (..., h, n_q, d_h) + k = split_heads(self.linear_k(kv_x)) # (..., h, n_k, d_h) + v = split_heads(self.linear_v(kv_x)) # (..., h, n_k, d_h) + if apply_scale: + q /= math.sqrt(self.c_hidden) # (..., h, n_q, d_h) + return q, k, v # (..., h, n_q, d_h), two (..., h, n_k, d_h) + + def _wrap_up(self, output: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor: + """Apply the optional gate, merge heads, and project the update.""" + + # output: (..., n_q, h, d_h); q_x: (..., n_q, c_q). + if self.linear_g is not None: + gate = self.sigmoid(self.linear_g(q_x)) # (..., n_q, h * d_h) + gate = gate.view( + *gate.shape[:-1], self.no_heads, -1 + ) # (..., n_q, h, d_h) + output = output * gate # (..., n_q, h, d_h) + return self.linear_o( + flatten_final_dims(output, 2) + ) # (..., n_q, c_q) + + def forward( + self, + q_x: torch.Tensor, + kv_x: torch.Tensor, + tri_bias: torch.Tensor, + mask_bias: torch.Tensor, + mask: torch.Tensor, + use_kernels: bool = False, + ) -> torch.Tensor: + """Return a gated pair update with shape ``(..., n, c_q)``.""" + + q, k, v = self._prep_qkv( + q_x, kv_x, apply_scale=not use_kernels + ) # q: (..., h, n_q, d_h); k/v: (..., h, n_k, d_h) + if use_kernels: + output = kernel_triangular_attn( + q, + k, + v, + tri_bias=tri_bias, + mask=mask.bool(), + scale=1.0 / math.sqrt(self.c_hidden), + ) # (..., h, n_q, d_h) + else: + output = _attention( + q, k, v, (mask_bias, tri_bias) + ) # (..., h, n_q, d_h) + return self._wrap_up( + output.transpose(-2, -3), q_x + ) # (..., n_q, c_q) + # tensor: (...); softmax preserves shape. diff --git a/src/fastplms/models/boltz/vb_tri_attn_utils.py b/src/fastplms/models/boltz/vb_tri_attn_utils.py new file mode 100644 index 0000000..9fc9d95 --- /dev/null +++ b/src/fastplms/models/boltz/vb_tri_attn_utils.py @@ -0,0 +1,340 @@ +"""Tensor-tree and memory-bounded execution helpers for pair attention. + +These utilities implement the small subset of OpenFold-style chunking needed +by the local Boltz pair stack. The implementation is independent, but keeps +the public function names expected by converted checkpoints. +""" + +from __future__ import annotations + +import torch +from collections.abc import Callable, Sequence +from functools import partial +from math import prod +from typing import Any + + +def add(left: torch.Tensor, right: torch.Tensor, inplace: bool) -> torch.Tensor: + """Add ``right`` to ``left``, optionally reusing ``left`` storage.""" + + # left/right: broadcast-compatible shapes; result has their broadcast shape. + if inplace: + left += right # same shape as left + return left # same shape as left + return left + right # broadcast shape + + +def permute_final_dims(tensor: torch.Tensor, inds: Sequence[int]) -> torch.Tensor: + """Permute only the final ``len(inds)`` axes of ``tensor``.""" + + final_count = len(inds) + leading = list(range(tensor.ndim - final_count)) + final = [tensor.ndim - final_count + index for index in inds] + return tensor.permute(leading + final) # (..., d_0, ..., d_n) + + +def is_fp16_enabled() -> bool: + """Return whether CUDA autocast currently targets ``float16``.""" + + return torch.is_autocast_enabled() and torch.get_autocast_dtype("cuda") == torch.float16 + + +def dict_map( + fn: Callable[[Any], Any], + dic: dict[Any, Any], + leaf_type: type | tuple[type, ...], +) -> dict[Any, Any]: + """Apply ``fn`` to leaves in a nested dictionary tree.""" + + return {key: tree_map(fn, value, leaf_type) for key, value in dic.items()} + + +def tree_map( + fn: Callable[[Any], Any], + tree: Any, + leaf_type: type | tuple[type, ...], +) -> Any: + """Map a function over dict, list, and tuple containers.""" + + if isinstance(tree, leaf_type): + return fn(tree) + if isinstance(tree, dict): + return dict_map(fn, tree, leaf_type) + if isinstance(tree, list): + return [tree_map(fn, item, leaf_type) for item in tree] + if isinstance(tree, tuple): + return tuple(tree_map(fn, item, leaf_type) for item in tree) + raise ValueError(f"tree type {type(tree)!r} is not supported") + + +tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor) + + +def flatten_final_dims(tensor: torch.Tensor, no_dims: int) -> torch.Tensor: + """Collapse the final ``no_dims`` axes into one axis.""" + + # tensor: (..., d_0, ..., d_n) + return tensor.reshape(*tensor.shape[:-no_dims], -1) # (..., prod(final dims)) + + +def _fetch_dims(tree: Any) -> list[torch.Size]: + """Collect tensor shapes from a supported tree.""" + + if isinstance(tree, torch.Tensor): + return [tree.shape] + if isinstance(tree, dict): + children = tree.values() + elif isinstance(tree, (list, tuple)): + children = tree + else: + raise ValueError(f"tree type {type(tree)!r} is not supported") + + shapes: list[torch.Size] = [] + for child in children: + shapes.extend(_fetch_dims(child)) + return shapes + + +@torch.jit.ignore +def _flat_idx_to_idx(flat_idx: int, dims: tuple[int, ...]) -> tuple[int, ...]: + """Convert a row-major flat index into an index tuple.""" + + coordinates = [0] * len(dims) + remainder = flat_idx + for axis in range(len(dims) - 1, -1, -1): + remainder, coordinates[axis] = divmod(remainder, dims[axis]) + return tuple(coordinates) + + +def _ravel_index(index: Sequence[int], dims: Sequence[int]) -> int: + """Convert a row-major index tuple to a flat index.""" + + flat = 0 + for coordinate, size in zip(index, dims, strict=True): + flat = flat * size + coordinate + return flat + + +def _cover_flat_interval( + start: int, + stop: int, + dims: tuple[int, ...], +) -> list[tuple[slice, ...]]: + """Cover ``[start, stop)`` with ordered contiguous tensor slices.""" + + if not dims: + return [tuple()] + if len(dims) == 1: + return [(slice(start, stop),)] + + child_size = prod(dims[1:]) + first_child = start // child_size + last_child = (stop - 1) // child_size + if first_child == last_child: + tail = _cover_flat_interval( + start % child_size, + (stop - 1) % child_size + 1, + dims[1:], + ) + prefix = slice(first_child, first_child + 1) + return [(prefix, *item) for item in tail] + + slices: list[tuple[slice, ...]] = [] + start_offset = start % child_size + first_full_child = first_child + if start_offset: + prefix = slice(first_child, first_child + 1) + slices.extend( + (prefix, *item) for item in _cover_flat_interval(start_offset, child_size, dims[1:]) + ) + first_full_child += 1 + + stop_offset = stop % child_size + full_stop = last_child if stop_offset else last_child + 1 + if first_full_child < full_stop: + slices.append((slice(first_full_child, full_stop),)) + + if stop_offset: + prefix = slice(last_child, last_child + 1) + slices.extend((prefix, *item) for item in _cover_flat_interval(0, stop_offset, dims[1:])) + return slices + + +@torch.jit.ignore +def _get_minimal_slice_set( + start: Sequence[int], + end: Sequence[int], + dims: Sequence[int], + start_edges: Sequence[bool] | None = None, + end_edges: Sequence[bool] | None = None, +) -> list[tuple[slice, ...]]: + """Return ordered slices covering the inclusive row-major interval. + + ``start_edges`` and ``end_edges`` remain accepted for compatibility. The + interval decomposition derives the same information directly. + """ + + del start_edges, end_edges + shape = tuple(dims) + if not shape: + return [tuple()] + flat_start = _ravel_index(start, shape) + flat_stop = _ravel_index(end, shape) + 1 + return _cover_flat_interval(flat_start, flat_stop, shape) + + +@torch.jit.ignore +def _chunk_slice( + tensor: torch.Tensor, + flat_start: int, + flat_end: int, + no_batch_dims: int, +) -> torch.Tensor: + """Slice a flattened batch interval without flattening the full tensor.""" + + batch_shape = tuple(tensor.shape[:no_batch_dims]) + start = _flat_idx_to_idx(flat_start, batch_shape) + end = _flat_idx_to_idx(flat_end - 1, batch_shape) + pieces = [ + tensor[item] for item in _get_minimal_slice_set(start, end, batch_shape) + ] # each: (*covered_batch_shape, *feature_shape) + feature_shape = tuple(tensor.shape[no_batch_dims:]) + return torch.cat( + [piece.reshape(-1, *feature_shape) for piece in pieces] + ) # (flat_end - flat_start, *feature_shape) + + +def _prepare_input( + tensor: torch.Tensor, + batch_shape: tuple[int, ...], + no_batch_dims: int, + *, + low_mem: bool, +) -> torch.Tensor: + # tensor: (*input_batch_shape, *feature_shape) + feature_shape = tuple(tensor.shape[no_batch_dims:]) + if low_mem: + return tensor.expand(batch_shape + feature_shape) # (*batch_shape, *feature_shape) + if any(size != 1 for size in tensor.shape[:no_batch_dims]): + tensor = tensor.expand( + batch_shape + feature_shape + ) # (*batch_shape, *feature_shape) + return tensor.reshape(-1, *feature_shape) # (flat_batch, *feature_shape) + + +def _write_chunk( + destination: Any, + source: Any, + start: int, + stop: int, + *, + add_into_out: bool, +) -> None: + if isinstance(source, dict): + for key, value in source.items(): + _write_chunk( + destination[key], + value, + start, + stop, + add_into_out=add_into_out, + ) + return + if isinstance(source, (tuple, list)): + for destination_item, source_item in zip(destination, source, strict=True): + _write_chunk( + destination_item, + source_item, + start, + stop, + add_into_out=add_into_out, + ) + return + if not isinstance(source, torch.Tensor): + raise ValueError(f"output type {type(source)!r} is not supported") + if add_into_out: + destination[start:stop] += source # (stop - start, *feature_shape) + else: + destination[start:stop] = source # (stop - start, *feature_shape) + + +def chunk_layer( + layer: Callable[..., Any], + inputs: dict[str, Any], + chunk_size: int, + no_batch_dims: int, + low_mem: bool = False, + _out: Any = None, + _add_into_out: bool = False, +) -> Any: + """Apply ``layer`` to flattened batch chunks and reassemble its output.""" + + if not inputs: + raise ValueError("at least one input is required") + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + + leading_shapes = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)] + batch_shape = tuple(max(sizes) for sizes in zip(*leading_shapes, strict=True)) + prepare = partial( + _prepare_input, + batch_shape=batch_shape, + no_batch_dims=no_batch_dims, + low_mem=low_mem, + ) + prepared_inputs = tensor_tree_map( + prepare, inputs + ) # each tensor: (flat_batch, *feature_shape), or broadcast batch in low-memory mode + + output = None + if _out is not None: + output = tensor_tree_map( + lambda tensor: tensor.reshape(-1, *tensor.shape[no_batch_dims:]), + _out, + ) # each tensor: (flat_batch, *feature_shape) + + flat_batch_size = prod(batch_shape) + for start in range(0, flat_batch_size, chunk_size): + stop = min(flat_batch_size, start + chunk_size) + if low_mem: + select = partial( + _chunk_slice, + flat_start=start, + flat_end=stop, + no_batch_dims=len(batch_shape), + ) + else: + + def select( + tensor: torch.Tensor, + start: int = start, + stop: int = stop, + ) -> torch.Tensor: + return tensor if tensor.shape[0] == 1 else tensor[start:stop] + + input_chunk = tensor_tree_map( + select, prepared_inputs + ) # each tensor: (chunk, *feature_shape) + output_chunk = layer( + **input_chunk + ) # each output tensor: (chunk, *output_feature_shape) + + if output is None: + output = tensor_tree_map( + lambda tensor: tensor.new_zeros((flat_batch_size, *tensor.shape[1:])), + output_chunk, + ) # each tensor: (flat_batch, *output_feature_shape) + _write_chunk( + output, + output_chunk, + start, + stop, + add_into_out=_add_into_out, + ) + + return tensor_tree_map( + lambda tensor: tensor.reshape(batch_shape + tuple(tensor.shape[1:])), + output, + ) # each tensor: (*batch_shape, *output_feature_shape) + # tensor: (..., d_0, ..., d_n); only the named final dimensions are reordered. + # tensor: (*batch_shape, *feature_shape) diff --git a/fastplms/__init__.py b/src/fastplms/models/dplm/__init__.py similarity index 100% rename from fastplms/__init__.py rename to src/fastplms/models/dplm/__init__.py diff --git a/src/fastplms/models/dplm/modeling_dplm.py b/src/fastplms/models/dplm/modeling_dplm.py new file mode 100644 index 0000000..f13886c --- /dev/null +++ b/src/fastplms/models/dplm/modeling_dplm.py @@ -0,0 +1,1326 @@ +"""FastPLMs-compatible DPLM implementation.""" + +# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import contextlib +import torch +import torch.nn as nn +from dataclasses import dataclass +from typing import ClassVar +from einops import rearrange +from torch.nn import functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel +from transformers import EsmTokenizer +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.models.esm.configuration_esm import EsmConfig +from transformers.models.esm.modeling_esm import ( + EsmAttention, + EsmClassificationHead, + EsmContactPredictionHead, + EsmEmbeddings, + EsmEncoder, + EsmLayer, + EsmLMHead, + EsmPooler, + EsmPreTrainedModel, + EsmSelfAttention, +) + +from fastplms.models._diffusion_generation import generate_dplm +from fastplms.models._esm_rotary import RotaryEmbedding + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + flex_attention, + get_attention_mask, + kernels_flash_attention_func, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin, select_hidden_state_embeddings + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "_get_flex_attention_fn", + "flex_attention", + "get_attention_mask", + "kernels_flash_attention_func", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +@dataclass +class DPLMMaskedLMOutput(MaskedLMOutput): + """Masked-LM output with DPLM extensions after the HF fields.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + last_hidden_state: torch.Tensor | None = None + + +@dataclass +class DPLMSequenceClassifierOutput(SequenceClassifierOutput): + """Sequence-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class DPLMTokenClassifierOutput(TokenClassifierOutput): + """Token-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class DPLMEncoderOutput(ModelOutput): + last_hidden_state: torch.Tensor | None = None + pooler_output: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +def _reject_unsupported_dplm_arguments(**arguments: object) -> None: + unsupported = [ + name + for name, value in arguments.items() + if value is not None and not (name == "use_cache" and value is False) + ] + if unsupported: + names = ", ".join(sorted(unsupported)) + raise ValueError( + "DPLM is an encoder-only diffusion model and does not support " + f"decoder, cross-attention, or KV-cache arguments: {names}." + ) + + +class DPLMConfig(EsmConfig): + model_type = "dplm" + + def __init__( + self, + attn_backend: str | None = None, + add_pooling_layer: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.attn_backend = attn_backend + self.add_pooling_layer = add_pooling_layer + self.tie_word_embeddings = False + + +_TOKENIZER_LOAD_CONTEXT_KEYS = ( + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", + "trust_remote_code", +) + + +class DPLMPreTrainedModel(FastPLMsAttentionMixin, EsmPreTrainedModel): + config_class = DPLMConfig + # All advertised wrappers install the encoder at ``self.esm``. Keep the + # Hugging Face base-model and checkpoint-prefix contract aligned with that + # actual module path. + base_model_prefix = "esm" + supports_gradient_checkpointing = True + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn = True + _supports_flash_attn_2 = False + _supports_flash_attn_3 = True + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_3", + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs} + if "token" not in load_context and "use_auth_token" in kwargs: + load_context["token"] = kwargs["use_auth_token"] + load_context["source"] = pretrained_model_name_or_path + + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + model = loaded[0] if isinstance(loaded, tuple) else loaded + model.__dict__["_fastplms_tokenizer_load_context"] = load_context + model.__dict__["_fastplms_tokenizer"] = None + return loaded + + @property + def tokenizer(self): + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + load_context = dict(self.__dict__.get("_fastplms_tokenizer_load_context") or {}) + source = load_context.pop("source", None) + if source is None: + source = str(getattr(self.config, "_name_or_path", "")).strip() + if not source: + raise RuntimeError( + "DPLM tokenizer loading requires a model loaded with from_pretrained " + "so checkpoint provenance is available." + ) + tokenizer_kwargs = { + key: value + for key, value in load_context.items() + if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None + } + resolved_revision = getattr(self.config, "_commit_hash", None) + if resolved_revision: + tokenizer_kwargs["revision"] = resolved_revision + tokenizer = EsmTokenizer.from_pretrained(source, **tokenizer_kwargs) + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"DPLM does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.config.attn_backend = backend + resolved = resolve_attention_backend(backend) + for module in self.modules(): + if isinstance(module, ModifiedEsmEncoder): + module.attention_backend = resolved + elif isinstance(module, ModifiedEsmSelfAttention): + module.attn_backend = resolved + + +class ModifiedEsmSelfAttention(EsmSelfAttention): + def __init__(self, config, position_embedding_type=None) -> None: + super().__init__(config, position_embedding_type) + self.config = config + self.scale = self.attention_head_size**-0.5 + self.dropout_prob = float(config.attention_probs_dropout_prob) + self.attn_backend = resolve_attention_backend(config.attn_backend) + if self.position_embedding_type == "rotary": + self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + # x: (b, l, d) + new_x_shape = (*x.size()[:-1], self.num_attention_heads, self.attention_head_size) + x = x.view(new_x_shape) # (b, l, h, d_h) + return x.permute(0, 2, 1, 3) # (b, h, l, d_h) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: object | None = None, + head_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_value: tuple[tuple[torch.FloatTensor]] | None = None, + output_attentions: bool | None = False, + output_s_max: bool | None = False, + past_key_values: tuple[tuple[torch.FloatTensor]] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l_q, d); encoder_hidden_states: (b, l_kv, d) or None + if past_key_values is not None: + past_key_value = past_key_values + + mixed_query_layer = self.query(hidden_states) # (b, l_q, d) + is_cross_attention = encoder_hidden_states is not None + + if is_cross_attention and past_key_value is not None: + key_layer = past_key_value[0] # (b, h, l_kv, d_h) + value_layer = past_key_value[1] # (b, h, l_kv, d_h) + cross_attn_mask = encoder_attention_mask + elif is_cross_attention: + key_layer = self.transpose_for_scores( # (b, h, l_kv, d_h) + self.key(encoder_hidden_states) + ) + value_layer = self.transpose_for_scores( # (b, h, l_kv, d_h) + self.value(encoder_hidden_states) + ) + cross_attn_mask = encoder_attention_mask + elif past_key_value is not None: + key_layer = self.transpose_for_scores(self.key(hidden_states)) # (b, h, l_q, d_h) + value_layer = self.transpose_for_scores(self.value(hidden_states)) # (b, h, l_q, d_h) + key_layer = torch.cat([past_key_value[0], key_layer], dim=2) # (b, h, l_kv, d_h) + value_layer = torch.cat([past_key_value[1], value_layer], dim=2) # (b, h, l_kv, d_h) + cross_attn_mask = None + else: + key_layer = self.transpose_for_scores(self.key(hidden_states)) # (b, h, l_q, d_h) + value_layer = self.transpose_for_scores(self.value(hidden_states)) # (b, h, l_q, d_h) + cross_attn_mask = None + + query_layer = self.transpose_for_scores(mixed_query_layer) * self.scale # (b, h, l_q, d_h) + + if self.position_embedding_type == "rotary": + query_layer, key_layer = self.rotary_embeddings( # Q: (b,h,l_q,d_h), K: (b,h,l_kv,d_h) + query_layer, + key_layer, + ) + + if self.position_embedding_type in ["relative_key", "relative_key_query"]: + raise NotImplementedError + + query_layer = query_layer.contiguous() # (b, h, l_q, d_h) + key_layer = key_layer.contiguous() # (b, h, l_kv, d_h) + value_layer = value_layer.contiguous() # (b, h, l_kv, d_h) + + if is_cross_attention: + if self.attn_backend not in { + AttentionBackend.EAGER, + AttentionBackend.SDPA, + }: + raise RuntimeError( + f"DPLM cross-attention does not implement {self.attn_backend.value!r}. " + "Use eager or SDPA for decoder cross-attention." + ) + if output_attentions: + attn_output, attn_weights, s_max = self._manual_attn( + query_layer, + key_layer, + value_layer, + cross_attn_mask, + output_s_max, + ) + elif self.attn_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_layer, + key_layer, + value_layer, + cross_attn_mask, + output_s_max, + ) + attn_weights = None + elif self.attn_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_layer, + key_layer, + value_layer, + cross_attn_mask, + ) + s_max = self._compute_s_max(query_layer, key_layer) if output_s_max else None + else: + attn_output, attn_weights, s_max = self._attn( + query_layer, + key_layer, + value_layer, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + if head_mask is not None and torch.is_tensor(head_mask): + batch_size, seq_len, _ = attn_output.shape + attn_output = attn_output.view( # (b, l_q, h, d_h) + batch_size, seq_len, self.num_attention_heads, self.attention_head_size + ) + attn_output = attn_output.permute(0, 2, 1, 3) * head_mask # (b, h, l_q, d_h) + attn_output = rearrange(attn_output, "b h s d -> b s (h d)") # (b, l_q, d) + + return attn_output, attn_weights, s_max # (b, l_q, d), optional weights, heads + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + if output_attentions: + return self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + + if ( + self.training + and self.dropout_prob > 0 + and (self.attn_backend.is_flash or self.attn_backend == AttentionBackend.FLEX_ATTENTION) + ): + raise RuntimeError( + f"DPLM {self.attn_backend.value} attention is inference-only when attention " + "dropout is nonzero. Use eager or SDPA for this training configuration." + ) + + if self.attn_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + return attn_output, None, s_max + if self.attn_backend.is_flash: + attn_output, attn_weights = self._kernels_flash_attn( + query_heads, key_heads, value_heads, attention_mask_2d + ) + elif self.attn_backend == AttentionBackend.FLEX: + attn_output, attn_weights = self._flex_attn( + query_heads, + key_heads, + value_heads, + flex_block_mask, + attention_mask_2d, + ) + elif self.attn_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_heads, key_heads, value_heads, attention_mask_4d + ) + else: + raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") + + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, query_heads: torch.Tensor, key_heads: torch.Tensor + ) -> list[torch.Tensor]: + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l_q) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l_kv) + s_max_bound = ( # (h,) + q_norm.max(dim=-1).values * k_norm.max(dim=-1).values + ).max(dim=0).values + return [s_max_bound[h] for h in range(self.num_attention_heads)] + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + attn_weights = torch.matmul( # (b, h, l_q, l_kv) + query_heads, + key_heads.transpose(-1, -2), + ) + if attention_mask_4d is not None: + if attention_mask_4d.dtype == torch.bool: + attn_weights = attn_weights.masked_fill( # (b, h, l_q, l_kv) + attention_mask_4d.logical_not(), + float("-inf"), + ) + else: + attn_weights = attn_weights + attention_mask_4d.to( # (b, h, l_q, l_kv) + device=attn_weights.device, + dtype=attn_weights.dtype, + ) + attn_weights = F.softmax(attn_weights, dim=-1) # (b, h, l_q, l_kv) + if self.dropout_prob > 0 and self.training: + attn_weights = F.dropout( # (b, h, l_q, l_kv) + attn_weights, + p=self.dropout_prob, + training=True, + ) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l_q, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l_q, d) + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + def _kernels_flash_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + query_tokens = query_heads.transpose(1, 2).contiguous() # (b, l_q, h, d_h) + key_tokens = key_heads.transpose(1, 2).contiguous() # (b, l_kv, h, d_h) + value_tokens = value_heads.transpose(1, 2).contiguous() # (b, l_kv, h, d_h) + # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward(). + # Pass softmax_scale=1.0 to prevent double-scaling by the kernel. + attn_output = kernels_flash_attention_func( # (b, l_q, h, d_h) + query_states=query_tokens, + key_states=key_tokens, + value_states=value_tokens, + attention_mask_2d=attention_mask_2d, + causal=False, + softmax_scale=1.0, + implementation=self.attn_backend.value, + ) + return rearrange(attn_output, "b s h d -> b s (h d)"), None # (b, l_q, d), None + + def _flex_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + flex_block_mask: BlockMask | None = None, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if flex_attention is None: + raise RuntimeError("Flex attention is not available in this environment.") + fn = _get_flex_attention_fn( + device=query_heads.device, + dtype=query_heads.dtype, + shape=tuple(query_heads.shape), + mask_semantics="padding", + ) + context_heads = fn( # (b, h, l_q, d_h) + query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0 + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None # (b, l_q, d), None + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + # The pinned official DPLM path uses Torch's efficient SDPA kernel for + # its non-null padding mask. Torch 2.13 otherwise selects cuDNN on H100, + # changing every downstream hidden state. Requiring the same public + # SDPA kernel makes the official FP32-storage/BF16-autocast path exact. + kernel_context = ( + sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION) + if query_heads.is_cuda + else contextlib.nullcontext() + ) + with kernel_context: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + dropout_p=self.dropout_prob if self.training else 0.0, + scale=1.0, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +class ModifiedEsmAttention(EsmAttention): + def __init__(self, config) -> None: + # Reuse Transformers' maintained ESM container layout, replacing only + # the self-attention engine that FastPLMs extends. This preserves the + # checkpoint schema without duplicating an upstream DPLM constructor. + EsmAttention.__init__(self, config) + self.self = ModifiedEsmSelfAttention(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: object | None = None, + head_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_value: tuple[tuple[torch.FloatTensor]] | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + hidden_states_ln = self.LayerNorm(hidden_states) # (b, l, d) + attn_output, attn_weights, s_max = self.self( # (b, l, d), optional weights, heads + hidden_states_ln, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + attention_output = self.output(attn_output, hidden_states) # (b, l, d) + return attention_output, attn_weights, s_max + + +class ModifiedEsmLayer(EsmLayer): + def __init__(self, config) -> None: + # Transformers owns the feed-forward, normalization, and decoder + # plumbing. Only attention dispatch differs for DPLM. + EsmLayer.__init__(self, config) + self.attention = ModifiedEsmAttention(config) + if self.add_cross_attention: + self.crossattention = ModifiedEsmAttention(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: object | None = None, + head_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_value: tuple[tuple[torch.FloatTensor]] | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + attention_output, attn_weights, s_max = self.attention( # (b, l, d), weights, heads + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + head_mask=head_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + past_key_value=past_key_value[:2] if past_key_value is not None else None, + ) + + if self.is_decoder and encoder_hidden_states is not None: + if self.add_cross_attention is False: + raise AttributeError( + f"If `encoder_hidden_states` are passed, {self} has to be " + "instantiated with cross-attention " + "layers by setting `config.add_cross_attention=True`" + ) + cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None + cross_attention_output, _, _ = self.crossattention( # (b, l, d), weights, heads + attention_output, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + output_attentions=output_attentions, + output_s_max=False, + ) + attention_output = cross_attention_output # (b, l, d) + + layer_output = self.feed_forward_chunk(attention_output) # (b, l, d) + return layer_output, attn_weights, s_max + + +class ModifiedEsmEncoder(EsmEncoder): + def __init__(self, config): + # Start from the public Transformers encoder contract, then substitute + # backend-aware layers while retaining every canonical state key. + EsmEncoder.__init__(self, config) + self.attention_backend = resolve_attention_backend(config.attn_backend) + self.layer = nn.ModuleList( + ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers) + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + head_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: list[tuple[tuple[torch.FloatTensor]]] | None = None, + use_cache: bool | None = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + output_s_max: bool = False, + ) -> DPLMEncoderOutput: + first_parameter = next(self.parameters(), None) + if ( + not self.training + and first_parameter is not None + and first_parameter.dtype == torch.bfloat16 + ): + raise RuntimeError( + "DPLM BF16 inference requires FP32-resident parameters under " + "CUDA BF16 autocast; static BF16 parameters do not meet the " + "declared parity contract." + ) + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + full_s_max = () if output_s_max else None + + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( + effective_backend=effective_backend, + batch_size=hidden_states.shape[0], + seq_len=hidden_states.shape[1], + device=hidden_states.device, + attention_mask=attention_mask, + dtype=hidden_states.dtype, + mask_semantics="padding", + ) + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + layer_head_mask = head_mask[i] if head_mask is not None else None + past_key_value = past_key_values[i] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + attention_mask_2d, + attention_mask_4d, + flex_block_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + output_s_max, + ) + else: + hidden_states, attn_weights, s_max = layer_module( + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + head_mask=layer_head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + if all_self_attentions is not None: + all_self_attentions = (*all_self_attentions, attn_weights) + if full_s_max is not None: + full_s_max = (*full_s_max, s_max) + + if self.emb_layer_norm_after: + hidden_states = self.emb_layer_norm_after(hidden_states) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + return DPLMEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + s_max=full_s_max, + ) + + +class FAST_DPLM_ENCODER(DPLMPreTrainedModel, EmbeddingMixin): + """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder, + contact_head) so that the weight keys are prefixed with 'esm.' in the outer DPLMModel, + matching pretrained DPLM checkpoints.""" + + def __init__(self, config, **kwargs): + DPLMPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.embeddings = EsmEmbeddings(config) + self.encoder = ModifiedEsmEncoder(config) + self.contact_head = EsmContactPredictionHead( + in_features=config.num_hidden_layers * config.num_attention_heads, + bias=True, + ) + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + attns = self(input_ids, attention_mask=attention_mask, output_attentions=True).attentions + attns = torch.stack(attns, dim=1) + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) + return self.contact_head(input_ids, attns) + + def _convert_head_mask_to_5d( + self, head_mask: torch.Tensor, num_hidden_layers: int + ) -> torch.Tensor: + if head_mask.dim() == 1: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == 2: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) + if head_mask.dim() != 5: + raise ValueError(f"head_mask.dim != 5, got {head_mask.dim()}") + head_mask = head_mask.to(dtype=self.dtype) + return head_mask + + def get_head_mask( + self, + head_mask: torch.Tensor | None, + num_hidden_layers: int, + is_attention_chunked: bool = False, + ) -> torch.Tensor | list[None]: + if head_mask is None: + return [None] * num_hidden_layers + head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) + if is_attention_chunked: + head_mask = head_mask.unsqueeze(-1) + return head_mask + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + head_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor] | DPLMEncoderOutput: + if self.config.is_decoder or self.config.add_cross_attention: + raise ValueError( + "DPLM is encoder-only; is_decoder and add_cross_attention must be false." + ) + _reject_unsupported_dplm_arguments( + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is not None: + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + + expected_attention_mask_shape = (batch_size, seq_length) + if attention_mask is None: + attention_mask_2d = torch.ones((batch_size, seq_length), device=device).bool() + elif attention_mask.dim() == 4: + raise ValueError( + "DPLM accepts a two-dimensional padding mask. Passing a four-dimensional " + "custom attention mask is unsupported because it cannot be applied to both " + "the embedding and optimized-attention paths without changing semantics." + ) + elif ( + attention_mask.dim() != 2 + or tuple(attention_mask.shape) != expected_attention_mask_shape + ): + raise ValueError( + f"attention_mask must have shape {expected_attention_mask_shape}; " + f"received {tuple(attention_mask.shape)}." + ) + else: + attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool) + + encoder_extended_attention_mask = encoder_attention_mask + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask_2d, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask_2d, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + ) + sequence_output = encoder_outputs.last_hidden_state + + if return_dict is False: + return (sequence_output, *encoder_outputs[1:]) + + result = DPLMEncoderOutput( + last_hidden_state=sequence_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + s_max=encoder_outputs.s_max, + ) + return result + + +class DPLMModel(DPLMPreTrainedModel, EmbeddingMixin): + config_class = DPLMConfig + + def __init__(self, config, add_pooling_layer: bool | None = None) -> None: + DPLMPreTrainedModel.__init__(self, config) + self.config = config + self.esm = FAST_DPLM_ENCODER(config) + if add_pooling_layer is None: + add_pooling_layer = config.add_pooling_layer + config.add_pooling_layer = bool(add_pooling_layer) + self.pooler = EsmPooler(config) if add_pooling_layer else None + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.embeddings.word_embeddings = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + head_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor] | DPLMEncoderOutput: + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs[0] # (b, l, d) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None # (b, d) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + result = DPLMEncoderOutput( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class DPLMForMaskedLM(FastPLMTestTimeTrainingMixin, DPLMPreTrainedModel, EmbeddingMixin): + config_class = DPLMConfig + + def __init__(self, config, dropout: float | None = None) -> None: + if dropout is not None: + config.hidden_dropout_prob = dropout + DPLMPreTrainedModel.__init__(self, config) + self.esm = FAST_DPLM_ENCODER(config) + self.lm_head = EsmLMHead(config) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + self.contact_head = None + self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"}) + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.lm_head.decoder + + def set_output_embeddings(self, new_embeddings): + old_bias = self.lm_head.bias + new_vocab_size = int(new_embeddings.out_features) + if old_bias.shape[0] != new_vocab_size: + resized_bias = old_bias.new_zeros(new_vocab_size) + copy_length = min(old_bias.shape[0], new_vocab_size) + with torch.no_grad(): + resized_bias[:copy_length].copy_(old_bias[:copy_length]) + self.lm_head.bias = nn.Parameter(resized_bias) + # EsmLMHead.forward adds this standalone bias after the decoder. HF's + # generic LM-head resizer may create a biased Linear, which would apply + # the bias twice and introduce an undeclared shared tensor on save. + new_embeddings.bias = None + self.lm_head.decoder = new_embeddings + + def generate( + self, + input_tokens: torch.Tensor, + tokenizer: object | None = None, + max_iter: int | None = None, + temperature: float | None = None, + partial_masks: torch.Tensor | None = None, + sampling_strategy: str = "gumbel_argmax", + disable_resample: bool = False, + resample_ratio: float = 0.25, + show_progress: bool = False, + **kwargs, + ) -> torch.Tensor: + """Generate protein tokens with the official DPLM diffusion schedule. + + ``input_tokens`` is X with shape (b, l). Positions marked ``True`` in + ``partial_masks`` remain fixed. ``max_iter=None`` uses the official + 500-step schedule; shorter schedules are useful for rapid exploration. + """ + + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected DPLM generation arguments: {names}") + return generate_dplm( + self, + input_tokens, + tokenizer=tokenizer, + max_iter=max_iter, + temperature=temperature, + partial_masks=partial_masks, + sampling_strategy=sampling_strategy, + disable_resample=disable_resample, + resample_ratio=resample_ratio, + show_progress=show_progress, + ) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.esm] + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + decoder_input_ids: torch.Tensor | None = None, + decoder_attention_mask: torch.Tensor | None = None, + decoder_inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor] | DPLMMaskedLMOutput: + _reject_unsupported_dplm_arguments( + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + decoder_inputs_embeds=decoder_inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + if attention_mask is None and input_ids is not None: + attention_mask = input_ids.ne(self.config.pad_token_id) + + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + logits = self.lm_head(sequence_output) # (b, l, c) + + loss = None + if labels is not None: + labels = labels.to(logits.device) # (b, l) + loss = self.loss_fct( # () + logits.view(-1, self.config.vocab_size), + labels.view(-1), + ) + + result = DPLMMaskedLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + last_hidden_state=sequence_output, + ) + return result if return_dict else result.to_tuple() + + +class DPLMForSequenceClassification(DPLMPreTrainedModel, EmbeddingMixin): + config_class = DPLMConfig + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def __init__(self, config) -> None: + DPLMPreTrainedModel.__init__(self, config) + self.num_labels = config.num_labels + self.esm = FAST_DPLM_ENCODER(config) + self.classifier = EsmClassificationHead(config) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.post_init() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, ...] | DPLMSequenceClassifierOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + logits = self.classifier(sequence_output) # (b, c) + + loss = None + if labels is not None: + labels = labels.to(logits.device) # (b,) or (b, c) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.squeeze(), labels.squeeze()) + else: + loss = self.mse(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + result = DPLMSequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class DPLMForTokenClassification(DPLMPreTrainedModel, EmbeddingMixin): + config_class = DPLMConfig + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def __init__(self, config) -> None: + DPLMPreTrainedModel.__init__(self, config) + self.num_labels = config.num_labels + self.esm = FAST_DPLM_ENCODER(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, ...] | DPLMTokenClassifierOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = self.dropout(outputs.last_hidden_state) # (b, l, d) + logits = self.classifier(sequence_output) # (b, l, c) + + loss = None + if labels is not None: + labels = labels.to(logits.device) # (b, l) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) # () + + result = DPLMTokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() diff --git a/src/fastplms/models/dplm2/__init__.py b/src/fastplms/models/dplm2/__init__.py new file mode 100644 index 0000000..3578194 --- /dev/null +++ b/src/fastplms/models/dplm2/__init__.py @@ -0,0 +1,6 @@ +"""DPLM2 model and tokenizer implementation.""" + +from .tokenization_dplm2 import DPLM2Tokenizer + + +__all__ = ["DPLM2Tokenizer"] diff --git a/src/fastplms/models/dplm2/modeling_dplm2.py b/src/fastplms/models/dplm2/modeling_dplm2.py new file mode 100644 index 0000000..3895e76 --- /dev/null +++ b/src/fastplms/models/dplm2/modeling_dplm2.py @@ -0,0 +1,1468 @@ +""" +FastPLMs-compatible DPLM2 implementation. +""" + +from __future__ import annotations + +import contextlib +import warnings +import torch +import torch.nn as nn +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, ClassVar +from einops import rearrange +from torch.nn import functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel +from transformers import AutoTokenizer +from transformers.modeling_outputs import ( + BaseModelOutputWithPoolingAndCrossAttentions, + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.models.esm.configuration_esm import EsmConfig +from transformers.models.esm.modeling_esm import ( + EsmAttention, + EsmClassificationHead, + EsmContactPredictionHead, + EsmEmbeddings, + EsmEncoder, + EsmIntermediate, + EsmLayer, + EsmLMHead, + EsmOutput, + EsmPooler, + EsmPreTrainedModel, + EsmSelfAttention, + EsmSelfOutput, +) + +from fastplms.models._diffusion_generation import generate_dplm2 +from fastplms.models._esm_rotary import RotaryEmbedding, apply_rotary_pos_emb +from fastplms.models.dplm2.tokenization_dplm2 import DPLM2Tokenizer + + +try: + from fastplms.attention import ( + AttentionBackend, + FastPLMsAttentionMixin, + get_attention_mask, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin, select_hidden_state_embeddings + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "get_attention_mask", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +def _infer_modality_type(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + input_mask = attention_mask.bool() + modality_type = ((input_ids < 33) & input_mask).int() + modality_type[~input_mask] = 2 + return modality_type + + +def _normalize_dplm2_input_ids(input_ids: torch.Tensor, vocab_size: int) -> torch.Tensor: + if input_ids.numel() == 0: + return input_ids + + normalized_input_ids = input_ids.clone() + generic_to_aa_special_ids = { + vocab_size: 2, + vocab_size + 1: 3, + vocab_size + 2: 0, + vocab_size + 3: 32, + } + for generic_id, aa_id in generic_to_aa_special_ids.items(): + normalized_input_ids[input_ids == generic_id] = aa_id + + valid_token_mask = normalized_input_ids.ge(0) + if valid_token_mask.any(): + max_token_id = int(normalized_input_ids[valid_token_mask].max().item()) + if max_token_id >= vocab_size: + raise ValueError( + f"Found token id {max_token_id} outside the DPLM2 embedding table " + f"(vocab_size={vocab_size}). Tokenizer special tokens must be normalized " + "before embedding." + ) + return normalized_input_ids + + +def _validate_dplm2_model_inputs( + *, + input_ids: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + attention_mask: torch.Tensor | None, + type_ids: torch.Tensor | None, + hidden_size: int, +) -> tuple[int, int]: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds.") + + if input_ids is not None: + if input_ids.ndim != 2: + raise ValueError( + f"input_ids must have shape (batch, seq_len), got {tuple(input_ids.shape)}." + ) + batch_size, seq_len = input_ids.shape + else: + if inputs_embeds is None: # Defensive guard for static narrowing. + raise RuntimeError("inputs_embeds validation reached an invalid state.") + if inputs_embeds.ndim != 3: + raise ValueError( + "inputs_embeds must have shape (batch, seq_len, hidden_size), " + f"got {tuple(inputs_embeds.shape)}." + ) + if inputs_embeds.shape[-1] != hidden_size: + raise ValueError( + f"inputs_embeds hidden size must be {hidden_size}, got {inputs_embeds.shape[-1]}." + ) + batch_size, seq_len = inputs_embeds.shape[:2] + + expected_shape = (batch_size, seq_len) + for name, value in (("attention_mask", attention_mask), ("type_ids", type_ids)): + if value is not None and tuple(value.shape) != expected_shape: + raise ValueError(f"{name} must have shape {expected_shape}, got {tuple(value.shape)}.") + return expected_shape + + +def _has_packed_multimodal_layout( + type_ids: torch.Tensor | None, + aa_type: int, + struct_type: int, + pad_type: int, +) -> bool: + if type_ids is None: + return False + if type_ids.ndim != 2: + raise ValueError( + f"Expected type_ids to have shape (batch, seq_len), got {tuple(type_ids.shape)}" + ) + seq_len = type_ids.shape[-1] + if seq_len % 2 != 0: + return False + + half_len = seq_len // 2 + first_half = type_ids[:, :half_len] + second_half = type_ids[:, half_len:] + + first_is_aa = ((first_half == aa_type) | (first_half == pad_type)).all(dim=-1) + first_is_struct = ((first_half == struct_type) | (first_half == pad_type)).all(dim=-1) + second_is_aa = ((second_half == aa_type) | (second_half == pad_type)).all(dim=-1) + second_is_struct = ((second_half == struct_type) | (second_half == pad_type)).all(dim=-1) + first_count = first_half.ne(pad_type).sum(dim=-1) + second_count = second_half.ne(pad_type).sum(dim=-1) + modalities_are_separate = (first_is_aa & second_is_struct) | (first_is_struct & second_is_aa) + packed_rows = modalities_are_separate & first_count.gt(0) & first_count.eq(second_count) + return bool(packed_rows.all()) + + +@dataclass +class DPLM2MaskedLMOutput(MaskedLMOutput): + """Masked-LM output with DPLM2 extensions after the HF fields.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + last_hidden_state: torch.Tensor | None = None + + +@dataclass +class DPLM2ModelOutput(BaseModelOutputWithPoolingAndCrossAttentions): + """Base-model output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class DPLM2SequenceClassifierOutput(SequenceClassifierOutput): + """Sequence-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class DPLM2TokenClassifierOutput(TokenClassifierOutput): + """Token-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class DPLM2EncoderOutput(ModelOutput): + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +class DPLM2Config(EsmConfig): + model_type = "dplm2" + + def __init__( + self, + attn_backend: str | None = "sdpa", + add_pooling_layer: bool = False, + aa_type: int = 1, + struct_type: int = 0, + pad_type: int = 2, + **kwargs, + ): + if kwargs.get("is_decoder", False) or kwargs.get("add_cross_attention", False): + raise ValueError( + "DPLM2 is encoder-only; is_decoder and add_cross_attention must be false." + ) + + # Published DPLM2 checkpoint configs inherited ``use_cache=true`` from + # EsmConfig even though the FastPLMs encoder has never implemented a KV + # cache. Keep those legacy artifacts loadable, but make the effective + # and newly serialized contract explicit and fail closed. + if kwargs.get("use_cache") is True: + warnings.warn( + "Legacy DPLM2 config requested use_cache=True, but DPLM2 is encoder-only " + "and does not implement KV caching; normalizing use_cache to False.", + UserWarning, + stacklevel=2, + ) + kwargs["is_decoder"] = False + kwargs["add_cross_attention"] = False + kwargs["use_cache"] = False + super().__init__(**kwargs) + # DPLM2's published implementation and manifest expose SDPA only. An + # older checkpoint may omit this FastPLMs field (or serialize it as + # null), so normalize that legacy representation to the same explicit + # backend before Transformers chooses its own generic eager default. + self.attn_backend = "sdpa" if attn_backend is None else attn_backend + self.add_pooling_layer = add_pooling_layer + self.aa_type = aa_type + self.struct_type = struct_type + self.pad_type = pad_type + self.tie_word_embeddings = False + + +_TOKENIZER_LOAD_CONTEXT_KEYS = ( + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", + "trust_remote_code", +) + + +class DPLM2PreTrainedModel(FastPLMsAttentionMixin, EsmPreTrainedModel): + config_class = DPLM2Config + # All advertised wrappers install the encoder at ``self.esm``. Transformers + # uses this name both for ``base_model`` and checkpoint prefix reconciliation. + base_model_prefix = "esm" + supports_gradient_checkpointing = True + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flex_attn = False + _supports_flash_attn = False + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ("sdpa",) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs} + if "token" not in load_context and "use_auth_token" in kwargs: + load_context["token"] = kwargs["use_auth_token"] + load_context["source"] = pretrained_model_name_or_path + + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + model = loaded[0] if isinstance(loaded, tuple) else loaded + model.__dict__["_fastplms_tokenizer_load_context"] = load_context + model.__dict__["_fastplms_tokenizer"] = None + return loaded + + @property + def tokenizer(self): + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + load_context = dict(self.__dict__.get("_fastplms_tokenizer_load_context") or {}) + source = load_context.pop("source", None) + if source is None: + source = str(getattr(self.config, "_name_or_path", "")).strip() + if not source: + raise RuntimeError( + "DPLM2 tokenizer loading requires a model loaded with from_pretrained " + "so checkpoint provenance is available." + ) + tokenizer_kwargs = { + key: value + for key, value in load_context.items() + if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None + } + resolved_revision = getattr(self.config, "_commit_hash", None) + if resolved_revision: + tokenizer_kwargs["revision"] = resolved_revision + tokenizer = DPLM2Tokenizer.from_pretrained(source, **tokenizer_kwargs) + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + def _tokenize_sequence_batch( + self, + sequences: Sequence[str], + *, + tokenizer: Any | None = None, + **kwargs: Any, + ) -> Any: + """Tokenize raw amino-acid sequences with official DPLM2 boundaries.""" + + resolved = tokenizer if tokenizer is not None else self.tokenizer + sequence_list = [sequences] if isinstance(sequences, str) else sequences + formatted = [ + f"{resolved.aa_cls_token}{sequence}{resolved.aa_eos_token}" + for sequence in sequence_list + ] + return resolved(formatted, add_special_tokens=False, **kwargs) + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"DPLM2 does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.config.attn_backend = backend + resolved = resolve_attention_backend(backend) + for module in self.modules(): + if isinstance(module, ModifiedEsmEncoder): + module.attention_backend = resolved + elif isinstance(module, ModifiedEsmSelfAttention): + module.attn_backend = resolved + + +class ModifiedRotaryEmbedding(RotaryEmbedding): + def __init__(self, dim: int, aa_type: int, struct_type: int, pad_type: int) -> None: + super().__init__(dim) + self.aa_type = aa_type + self.struct_type = struct_type + self.pad_type = pad_type + + def _has_multimodal_tokens(self, type_ids: torch.Tensor | None) -> bool: + # The split rotary path only works when the sequence tensor is already packed + # as two equal-length, modality-specific halves. Either track may come first. + # Plain protein batches can still contain high-ID special tokens, so mere + # modality presence is not enough. + return _has_packed_multimodal_layout( + type_ids=type_ids, + aa_type=self.aa_type, + struct_type=self.struct_type, + pad_type=self.pad_type, + ) + + def align_frequency_buffer( + self, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Match the official model-wide ``to(device, dtype)`` conversion. + + Transformers' meta-device loader converts parameters to the requested + dtype but can leave this persistent rotary buffer in FP32. The pinned + official implementation moves the complete module, including + ``inv_freq``. Aligning the buffer before building rotary factors keeps + Q, K, and V in one dtype for every attention backend. + """ + + if self.inv_freq.device == device and self.inv_freq.dtype == dtype: + return + self.inv_freq = self.inv_freq.to(device=device, dtype=dtype) + self._seq_len_cached = None + self._cos_cached = None + self._sin_cached = None + + def _update_cos_sin_tables( + self, + x: torch.Tensor, + type_ids: torch.Tensor | None, + seq_dimension: int = 2, + ) -> tuple[torch.Tensor, torch.Tensor]: + # x: (b, h, l, d) + seq_len = x.shape[seq_dimension] + if self._has_multimodal_tokens(type_ids): + seq_len = seq_len // 2 + + cache_is_stale = ( + self._cos_cached is None + or self._sin_cached is None + or seq_len != self._seq_len_cached + or self._cos_cached.device != x.device + or self._cos_cached.dtype != self.inv_freq.dtype + ) + if cache_is_stale: + self._seq_len_cached = seq_len + t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq) # (l,) + freqs = torch.outer(t, self.inv_freq) # (l, d / 2) + # Match the official DPLM2 operation order: rotary factors inherit + # the frequency-buffer dtype. This keeps them in FP32 under BF16 + # autocast, while a model explicitly converted to BF16 still builds + # BF16 factors and remains usable without autocast. + emb = torch.cat((freqs, freqs), dim=-1).to(device=x.device) # (l, d) + self._cos_cached = emb.cos()[None, None, :, :] # (1, 1, l, d) + self._sin_cached = emb.sin()[None, None, :, :] # (1, 1, l, d) + + return self._cos_cached, self._sin_cached + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + type_ids: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, h, l, d) + self._cos_cached, self._sin_cached = self._update_cos_sin_tables( + k, + type_ids=type_ids, + seq_dimension=-2, + ) + + if self._has_multimodal_tokens(type_ids): + q_1, q_2 = q.chunk(2, dim=-2) # each (b, h, l / 2, d) + k_1, k_2 = k.chunk(2, dim=-2) # each (b, h, l / 2, d) + q_1 = apply_rotary_pos_emb(q_1, self._cos_cached, self._sin_cached) + q_2 = apply_rotary_pos_emb(q_2, self._cos_cached, self._sin_cached) + k_1 = apply_rotary_pos_emb(k_1, self._cos_cached, self._sin_cached) + k_2 = apply_rotary_pos_emb(k_2, self._cos_cached, self._sin_cached) + return torch.cat((q_1, q_2), dim=-2), torch.cat((k_1, k_2), dim=-2) + + return ( + apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), + apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), + ) + + +class ModifiedEsmSelfAttention(EsmSelfAttention): + def __init__(self, config, position_embedding_type=None) -> None: + super().__init__(config, position_embedding_type) + self.config = config + self.scale = self.attention_head_size**-0.5 + self.dropout_prob = config.attention_probs_dropout_prob + self.attn_backend = resolve_attention_backend(config.attn_backend) + self.rotary_embeddings = ModifiedRotaryEmbedding( + dim=self.attention_head_size, + aa_type=config.aa_type, + struct_type=config.struct_type, + pad_type=config.pad_type, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + type_ids: torch.Tensor | None = None, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d) + batch_size, seq_length = hidden_states.shape[:-1] + hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) + query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + + query_heads = query_heads * self.scale + + if self.position_embedding_type == "rotary": + self.rotary_embeddings.align_frequency_buffer( + device=query_heads.device, + dtype=self.query.weight.dtype, + ) + query_heads, key_heads = self.rotary_embeddings(query_heads, key_heads, type_ids) + + attn_output, attn_weights, s_max = self._attn( + query_heads, + key_heads, + value_heads, + attention_mask_4d=attention_mask_4d, + output_attentions=output_attentions, + output_s_max=output_s_max, + effective_backend=effective_backend, + ) + return attn_output, attn_weights, s_max + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + if effective_backend is None: + effective_backend = resolve_attention_backend_for_call( + self.attn_backend, + output_attentions=output_attentions, + ) + + if effective_backend == AttentionBackend.EAGER: + attn_output, attn_weights, s_max = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + return attn_output, attn_weights if output_attentions else None, s_max + + if output_attentions: + raise AssertionError( + "DPLM2 output_attentions=True must resolve to eager attention for this call." + ) + + if effective_backend != AttentionBackend.SDPA: + raise AssertionError(f"Unsupported resolved backend: {effective_backend}") + attn_output, attn_weights = self._sdpa_attn( + query_heads, + key_heads, + value_heads, + attention_mask_4d, + ) + + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, query_heads: torch.Tensor, key_heads: torch.Tensor + ) -> list[torch.Tensor]: + # query_heads, key_heads: (b, h, l, d_h) + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) + s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values + return [s_max_bound[h] for h in range(self.num_attention_heads)] + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = torch.matmul( + query_heads, key_heads.transpose(-1, -2) + ) # (b, h, l, l) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) + attn_weights = F.softmax(attn_weights, dim=-1) + if self.dropout_prob > 0 and self.training: + attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l, d) + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + # Pinned DPLM2 uses PyTorch's efficient SDPA kernel for its non-null + # padding mask. Newer PyTorch releases otherwise select cuDNN on H100, + # which exceeds the fixed deep-BF16 parity target. This is still the + # public SDPA operation and raises if its required CUDA kernel is absent. + kernel_context = ( + sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION) + if query_heads.is_cuda + else contextlib.nullcontext() + ) + with kernel_context: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + dropout_p=self.dropout_prob if self.training else 0.0, + scale=1.0, + ) # (b, h, l, d_h) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +class ModifiedEsmAttention(EsmAttention): + def __init__(self, config) -> None: + nn.Module.__init__(self) + self.self = ModifiedEsmSelfAttention(config) + self.output = EsmSelfOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + type_ids: torch.Tensor | None = None, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d) + hidden_states_ln = self.LayerNorm(hidden_states) # (b, l, d) + attn_output, attn_weights, s_max = self.self( + hidden_states_ln, + attention_mask_4d=attention_mask_4d, + output_attentions=output_attentions, + output_s_max=output_s_max, + type_ids=type_ids, + effective_backend=effective_backend, + ) + attention_output = self.output(attn_output, hidden_states) + return attention_output, attn_weights, s_max + + +class ModifiedEsmLayer(EsmLayer): + def __init__(self, config) -> None: + nn.Module.__init__(self) + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = ModifiedEsmAttention(config) + self.intermediate = EsmIntermediate(config) + self.output = EsmOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + type_ids: torch.Tensor | None = None, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d) + attention_output, attn_weights, s_max = self.attention( + hidden_states, + attention_mask_4d=attention_mask_4d, + output_attentions=output_attentions, + output_s_max=output_s_max, + type_ids=type_ids, + effective_backend=effective_backend, + ) + layer_output = self.feed_forward_chunk(attention_output) + return layer_output, attn_weights, s_max + + +class ModifiedEsmEncoder(EsmEncoder): + def __init__(self, config) -> None: + nn.Module.__init__(self) + self.config = config + self.attention_backend = resolve_attention_backend(config.attn_backend) + self.layer = nn.ModuleList( + [ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool = False, + output_attentions: bool = False, + output_s_max: bool = False, + type_ids: torch.Tensor | None = None, + ) -> DPLM2EncoderOutput: + # hidden_states: (b, l, d); attention_mask, type_ids: (b, l) + first_parameter = next(self.parameters(), None) + if ( + not self.training + and first_parameter is not None + and first_parameter.dtype == torch.bfloat16 + ): + raise RuntimeError( + "DPLM2 BF16 inference requires FP32-resident parameters under " + "CUDA BF16 autocast; static BF16 parameters do not meet the " + "declared parity contract." + ) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + full_s_max = () if output_s_max else None + + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + _, attention_mask_4d, _ = get_attention_mask( + effective_backend=effective_backend, + batch_size=hidden_states.shape[0], + seq_len=hidden_states.shape[1], + device=hidden_states.device, + attention_mask=attention_mask, + dtype=hidden_states.dtype, + mask_semantics="padding", + ) # attention_mask_4d: (b, 1, 1, l) or (b, 1, l, l) + + for layer_module in self.layer: + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + if self.gradient_checkpointing and self.training: + hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + attention_mask_4d, + output_attentions, + output_s_max, + type_ids, + effective_backend, + ) + else: + hidden_states, attn_weights, s_max = layer_module( + hidden_states, + attention_mask_4d=attention_mask_4d, + output_attentions=output_attentions, + output_s_max=output_s_max, + type_ids=type_ids, + effective_backend=effective_backend, + ) + + if all_attentions is not None: + all_attentions = (*all_attentions, attn_weights) + if full_s_max is not None: + full_s_max = (*full_s_max, s_max) + + if self.emb_layer_norm_after: + hidden_states = self.emb_layer_norm_after(hidden_states) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + return DPLM2EncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + s_max=full_s_max, + ) + + +class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin): + """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder) + so that the weight keys are prefixed with 'esm.' in the outer DPLM2Model, + matching pretrained DPLM2 checkpoints.""" + + def __init__(self, config, **kwargs) -> None: + DPLM2PreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.embeddings = EsmEmbeddings(config) + self.encoder = ModifiedEsmEncoder(config) + self.contact_head = EsmContactPredictionHead( + in_features=config.num_hidden_layers * config.num_attention_heads, + bias=True, + ) + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def predict_contacts( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Predict residue contacts with the checkpoint's tied contact head.""" + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + type_ids = self._get_modality_type(input_ids, attention_mask) + attentions = self( + input_ids=input_ids, + attention_mask=attention_mask, + type_ids=type_ids, + output_attentions=True, + ).attentions + if attentions is None: + raise RuntimeError("DPLM2 did not return attention maps for contact prediction.") + # A is the layer/head attention tensor; M marks valid tokens. + attention_tensor = torch.stack(attentions, dim=1) + residue_mask = attention_mask.to(dtype=attention_tensor.dtype) + attention_tensor = ( + attention_tensor + * residue_mask[:, None, None, :, None] + * residue_mask[:, None, None, None, :] + ) + return self.contact_head(input_ids, attention_tensor) + + def _get_modality_type( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return _infer_modality_type(input_ids, attention_mask) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + type_ids = _infer_modality_type(input_ids, attention_mask) + token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + encoder_outputs = self.encoder( + token_embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + type_ids=type_ids, + ) + return select_hidden_state_embeddings( + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + type_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...] | DPLM2EncoderOutput: + _validate_dplm2_model_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + type_ids=type_ids, + hidden_size=self.config.hidden_size, + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + token_embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + token_embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + type_ids=type_ids, + ) + + result = DPLM2EncoderOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + s_max=encoder_outputs.s_max, + ) + if not return_dict: + return result.to_tuple() + return result + + +class DPLM2Model(DPLM2PreTrainedModel, EmbeddingMixin): + config_class = DPLM2Config + + def __init__(self, config, add_pooling_layer: bool | None = None): + DPLM2PreTrainedModel.__init__(self, config) + self.config = config + self.esm = FAST_DPLM2_ENCODER(config) + if add_pooling_layer is None: + add_pooling_layer = config.add_pooling_layer + config.add_pooling_layer = bool(add_pooling_layer) + self.pooler = EsmPooler(config) if add_pooling_layer else None + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.embeddings.word_embeddings = value + + def predict_contacts( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + type_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...] | DPLM2ModelOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm" + _validate_dplm2_model_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + type_ids=type_ids, + hidden_size=self.config.hidden_size, + ) + if inputs_embeds is not None and type_ids is None and not direct_dplm_esm: + raise ValueError( + "type_ids is required for multimodal DPLM2 calls that use inputs_embeds." + ) + if input_ids is not None: + normalized_input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + if attention_mask is None: + attention_mask = normalized_input_ids.ne(self.config.pad_token_id) + if type_ids is None and not direct_dplm_esm: + type_ids = _infer_modality_type(normalized_input_ids, attention_mask) + input_ids = normalized_input_ids + + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + type_ids=type_ids, + ) + sequence_output = outputs.last_hidden_state + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + result = DPLM2ModelOutput( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class DPLM2ForMaskedLM(FastPLMTestTimeTrainingMixin, DPLM2PreTrainedModel, EmbeddingMixin): + config_class = DPLM2Config + + def __init__( + self, + config, + dropout: float | None = None, + vocab_size: int | None = None, + ): + if dropout is not None: + config.hidden_dropout_prob = dropout + config.tie_word_embeddings = False + if vocab_size is not None: + config.vocab_size = vocab_size + DPLM2PreTrainedModel.__init__(self, config) + self.esm = FAST_DPLM2_ENCODER(config) + self.lm_head = EsmLMHead(config) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + self.pad_id = config.pad_token_id + self.contact_head = None + self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"}) + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.lm_head.decoder + + def set_output_embeddings(self, new_embeddings): + old_bias = self.lm_head.bias + new_vocab_size = int(new_embeddings.out_features) + if old_bias.shape[0] != new_vocab_size: + resized_bias = old_bias.new_zeros(new_vocab_size) + copy_length = min(old_bias.shape[0], new_vocab_size) + with torch.no_grad(): + resized_bias[:copy_length].copy_(old_bias[:copy_length]) + self.lm_head.bias = nn.Parameter(resized_bias) + # EsmLMHead.forward adds this standalone bias after the decoder. HF's + # generic LM-head resizer may create a biased Linear, which would apply + # the bias twice and introduce an undeclared shared tensor on save. + new_embeddings.bias = None + self.lm_head.decoder = new_embeddings + + def generate( + self, + input_tokens: torch.Tensor, + max_iter: int | None = None, + temperature: float = 1.0, + partial_masks: torch.Tensor | None = None, + unmasking_strategy: str = "stochastic1.0", + sampling_strategy: str = "annealing@2.0:0.1", + show_progress: bool = False, + **kwargs, + ) -> dict[str, torch.Tensor]: + """Generate packed sequence and structure tokens with DPLM2 diffusion. + + ``input_tokens`` is X with shape (b, l). Positions marked ``True`` in + ``partial_masks`` remain fixed. The returned mapping contains + ``output_tokens``, matching the official DPLM2 public API. + """ + + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected DPLM2 generation arguments: {names}") + return generate_dplm2( + self, + input_tokens, + max_iter=max_iter, + temperature=temperature, + partial_masks=partial_masks, + unmasking_strategy=unmasking_strategy, + sampling_strategy=sampling_strategy, + show_progress=show_progress, + ) + + def predict_contacts( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the official ESM contact head output from the encoder.""" + return self.esm.predict_contacts(input_ids, attention_mask) + + def _get_modality_type( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + return _infer_modality_type(input_ids, attention_mask) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.pad_id) + type_ids = self._get_modality_type(input_ids, attention_mask) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + type_ids=type_ids, + output_attentions=False, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + return select_hidden_state_embeddings( + outputs.last_hidden_state, + outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.esm] + + def _ttt_tokenize( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + if input_ids is not None: + return input_ids + if seq is None: + raise ValueError("Pass either seq or input_ids for TTT.") + sequences = [seq] if isinstance(seq, str) else seq + tokenized = self._tokenize_sequence_batch( + sequences, + return_tensors="pt", + padding=True, + ) + return tokenized["input_ids"] + + def _ttt_mask_token(self) -> int: + return int(self.tokenizer._token_to_id[self.tokenizer.aa_mask_token]) + + def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + tokenizer = self.tokenizer + special_ids = set(tokenizer.all_special_ids) + struct_boundary = int(tokenizer._token_to_id[tokenizer.struct_cls_token]) + residue_ids = [] + for residue in "ACDEFGHIKLMNPQRSTVWY": + token_id = tokenizer._token_to_id.get(residue) + if ( + isinstance(token_id, int) + and 0 <= token_id < struct_boundary + and token_id not in special_ids + and token_id not in residue_ids + ): + residue_ids.append(token_id) + if not residue_ids: + raise RuntimeError("DPLM2 TTT amino-acid replacement set is empty.") + if len(residue_ids) != 20: + raise RuntimeError( + "DPLM2 TTT requires all 20 canonical amino-acid replacement tokens; " + f"resolved {len(residue_ids)}." + ) + return torch.tensor(residue_ids, device=input_ids.device, dtype=input_ids.dtype) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + type_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor] | DPLM2MaskedLMOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm" + _validate_dplm2_model_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + type_ids=type_ids, + hidden_size=self.config.hidden_size, + ) + + if attention_mask is None: + if input_ids is None: + raise ValueError( + "attention_mask is required when DPLM2 is called with inputs_embeds." + ) + attention_mask = input_ids.ne(self.pad_id) + + encoder_input_ids = input_ids + if input_ids is not None: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + encoder_input_ids = input_ids + if type_ids is None and not direct_dplm_esm: + if input_ids is None: + raise ValueError( + "type_ids is required for multimodal DPLM2 calls that use inputs_embeds." + ) + type_ids = self._get_modality_type(input_ids, attention_mask) + + if input_ids is not None and inputs_embeds is None and not direct_dplm_esm: + # The official multimodal wrapper applies the embedding block + # once before entering EsmForDPLM2. The inner ESM model then + # applies it a second time using these intermediate embeddings. + inputs_embeds = self.esm.embeddings( + input_ids=input_ids, + attention_mask=attention_mask, + ) + encoder_input_ids = None + + outputs = self.esm( + input_ids=encoder_input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + type_ids=type_ids, + ) + + sequence_output = outputs.last_hidden_state + logits = self.lm_head(sequence_output) + loss = None + if labels is not None: + labels = _normalize_dplm2_input_ids(labels, self.config.vocab_size) + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) + + result = DPLM2MaskedLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + last_hidden_state=sequence_output, + ) + return result if return_dict else result.to_tuple() + + +class DPLM2ForSequenceClassification(DPLM2PreTrainedModel, EmbeddingMixin): + config_class = DPLM2Config + + def __init__(self, config): + DPLM2PreTrainedModel.__init__(self, config) + self.num_labels = config.num_labels + self.esm = FAST_DPLM2_ENCODER(config) + self.classifier = EsmClassificationHead(config) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + type_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, ...] | DPLM2SequenceClassifierOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm" + _validate_dplm2_model_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + type_ids=type_ids, + hidden_size=self.config.hidden_size, + ) + if inputs_embeds is not None and type_ids is None and not direct_dplm_esm: + raise ValueError( + "type_ids is required for multimodal DPLM2 calls that use inputs_embeds." + ) + if input_ids is not None: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + if type_ids is None and input_ids is not None and not direct_dplm_esm: + type_ids = _infer_modality_type(input_ids, attention_mask) + + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + type_ids=type_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.squeeze(), labels.squeeze()) + else: + loss = self.mse(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + result = DPLM2SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class DPLM2ForTokenClassification(DPLM2PreTrainedModel, EmbeddingMixin): + config_class = DPLM2Config + + def __init__(self, config): + DPLM2PreTrainedModel.__init__(self, config) + self.num_labels = config.num_labels + self.esm = FAST_DPLM2_ENCODER(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.esm.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm.set_input_embeddings(value) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + type_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, ...] | DPLM2TokenClassifierOutput: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm" + _validate_dplm2_model_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + type_ids=type_ids, + hidden_size=self.config.hidden_size, + ) + if inputs_embeds is not None and type_ids is None and not direct_dplm_esm: + raise ValueError( + "type_ids is required for multimodal DPLM2 calls that use inputs_embeds." + ) + if input_ids is not None: + input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size) + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + if type_ids is None and input_ids is not None and not direct_dplm_esm: + type_ids = _infer_modality_type(input_ids, attention_mask) + + outputs = self.esm( + input_ids=input_ids, + attention_mask=attention_mask, + type_ids=type_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = self.dropout(outputs.last_hidden_state) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + result = DPLM2TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +# Importing the DPLM2 model implementation makes its paired tokenizer visible +# to AutoTokenizer. This is registration only; it performs no I/O or downloads. +try: + AutoTokenizer.register( + DPLM2Config, + tokenizer_class=DPLM2Tokenizer, + exist_ok=True, + ) +except TypeError: + # Transformers 4.x used this name; 5.x prefers tokenizer_class. + AutoTokenizer.register( + DPLM2Config, + slow_tokenizer_class=DPLM2Tokenizer, + exist_ok=True, + ) diff --git a/src/fastplms/models/dplm2/tokenization_dplm2.py b/src/fastplms/models/dplm2/tokenization_dplm2.py new file mode 100644 index 0000000..1194072 --- /dev/null +++ b/src/fastplms/models/dplm2/tokenization_dplm2.py @@ -0,0 +1,92 @@ +"""Independent DPLM2 amino-acid and structure tokenizer. + +DPLM2 stores two token tracks in one vocabulary. Amino-acid tokens use their +own boundary, unknown, and mask tokens, while structure codes use a separate +set. The generic ``cls_token`` and ``eos_token`` aliases are deliberately not +defined because a caller must choose the modality-specific boundaries. +""" + +from __future__ import annotations + +from typing import ClassVar +from transformers import AddedToken, EsmTokenizer, PreTrainedTokenizer + + +class DPLM2Tokenizer(EsmTokenizer): + """Tokenize the official DPLM2 amino-acid and structure vocabulary. + + Input text is split against the complete pinned vocabulary. Amino-acid + sequences may be passed as contiguous characters and structure sequences + as whitespace-separated four-digit codes. Callers constructing a model + input add ``aa_*`` or ``struct_*`` boundaries explicitly and then use + ``add_special_tokens=False``. The output token IDs preserve the official + multimodal vocabulary exactly. + """ + + SPECIAL_TOKENS_ATTRIBUTES: ClassVar[list[str]] = [ + "aa_cls_token", + "aa_eos_token", + "aa_unk_token", + "aa_mask_token", + "struct_cls_token", + "struct_eos_token", + "struct_unk_token", + "struct_mask_token", + "pad_token", + ] + # The official tokenizer exposes no generic sequence-boundary aliases. + # Keeping these attributes explicitly set to None preserves that public + # contract on Transformers v5, whose custom special-token lookup is strict. + bos_token: ClassVar[None] = None + cls_token: ClassVar[None] = None + eos_token: ClassVar[None] = None + mask_token: ClassVar[None] = None + sep_token: ClassVar[None] = None + unk_token: ClassVar[None] = None + bos_token_id: ClassVar[None] = None + cls_token_id: ClassVar[None] = None + eos_token_id: ClassVar[None] = None + mask_token_id: ClassVar[None] = None + sep_token_id: ClassVar[None] = None + unk_token_id: ClassVar[None] = None + + def __init__( + self, + vocab_file: str, + aa_cls_token: str | AddedToken = "", + aa_eos_token: str | AddedToken = "", + aa_unk_token: str | AddedToken = "", + aa_mask_token: str | AddedToken = "", + struct_cls_token: str | AddedToken = "", + struct_eos_token: str | AddedToken = "", + struct_unk_token: str | AddedToken = "", + struct_mask_token: str | AddedToken = "", + pad_token: str | AddedToken = "", + **kwargs: object, + ) -> None: + with open(vocab_file, encoding="utf-8") as handle: + self.all_tokens = [line.strip() for line in handle.read().splitlines()] + self._id_to_token = dict(enumerate(self.all_tokens)) + self._token_to_id = {token: token_id for token_id, token in self._id_to_token.items()} + + # EsmTokenizer would install generic ESM boundary aliases. DPLM2 has + # modality-specific boundaries instead, so initialize the common + # tokenizer base with only the nine official special-token fields. + PreTrainedTokenizer.__init__( + self, + aa_cls_token=aa_cls_token, + aa_eos_token=aa_eos_token, + aa_unk_token=aa_unk_token, + aa_mask_token=aa_mask_token, + struct_cls_token=struct_cls_token, + struct_eos_token=struct_eos_token, + struct_unk_token=struct_unk_token, + struct_mask_token=struct_mask_token, + pad_token=pad_token, + **kwargs, + ) + self.unique_no_split_tokens = self.all_tokens + self._update_trie(self.unique_no_split_tokens) + + +__all__ = ["DPLM2Tokenizer"] diff --git a/fastplms/ankh/__init__.py b/src/fastplms/models/e1/__init__.py similarity index 100% rename from fastplms/ankh/__init__.py rename to src/fastplms/models/e1/__init__.py diff --git a/src/fastplms/models/e1/attention.py b/src/fastplms/models/e1/attention.py new file mode 100644 index 0000000..c664380 --- /dev/null +++ b/src/fastplms/models/e1/attention.py @@ -0,0 +1,485 @@ +"""E1 attention mask, unpadding, FlexAttention, and kernel adapters.""" + +from __future__ import annotations + +import os +import torch +from collections.abc import Callable +from torch.nn.attention.flex_attention import _create_sparse_block_from_block_mask + +from fastplms.attention import ( + BlockMask, + _ensure_flash_kernels_loaded, + _get_flex_attention_fn, + _kernels_flash_forward, + _kernels_flash_varlen_forward, + create_block_mask, + flex_attention, + index_first_axis, + pad_input, +) + + +@torch.compiler.disable +def create_block_causal_mask_optimized(sequence_ids: torch.Tensor) -> BlockMask: + # sequence_ids: (b, l) + if create_block_mask is None: + raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.") + # Assumes sequence_ids is sorted in increasing order for each batch item, except for + # the -1 values, which are used to indicate the padding tokens. + def document_mask(b, h, q_idx, kv_idx): # type: ignore[no-untyped-def] + return ( + (sequence_ids[b, q_idx] >= sequence_ids[b, kv_idx]) + & (sequence_ids[b, q_idx] != -1) + & (sequence_ids[b, kv_idx] != -1) + ) + + batch_size, seqlen = sequence_ids.shape + return create_block_mask( + document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device + ) + + +@torch.compiler.disable +def create_within_seq_block_mask(sequence_ids: torch.Tensor) -> BlockMask: + # sequence_ids: (b, l) + if create_block_mask is None: + raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.") + def document_mask(b, h, q_idx, kv_idx): # type: ignore[no-untyped-def] + return ( + (sequence_ids[b, q_idx] == sequence_ids[b, kv_idx]) + & (sequence_ids[b, q_idx] != -1) + & (sequence_ids[b, kv_idx] != -1) + ) + + batch_size, seqlen = sequence_ids.shape + return create_block_mask( + document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device + ) + + +def build_within_seq_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor: + # sequence_ids: (b, l) + not_pad = sequence_ids != -1 # (b, l) + same_seq = sequence_ids.unsqueeze(-1) == sequence_ids.unsqueeze(-2) # (b, l, l) + valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2) # (b, l, l) + return (same_seq & valid).unsqueeze(1) # (b, 1, l, l) + + +def build_block_causal_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor: + # sequence_ids: (b, l) + not_pad = sequence_ids != -1 # (b, l) + causal = sequence_ids.unsqueeze(-1) >= sequence_ids.unsqueeze(-2) # (b, l, l) + valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2) # (b, l, l) + return (causal & valid).unsqueeze(1) # (b, 1, l, l) + + +def flex_attention_func( + query_states: torch.Tensor, # Q has shape (b, l, h, d). + key_states: torch.Tensor, # K has shape (b, l, h_kv, d). + value_states: torch.Tensor, # V has shape (b, l, h_kv, d). + score_mod: Callable | None = None, + block_mask: BlockMask | None = None, + sequence_lengths: tuple[int, ...] | None = None, + mask_semantics: str = "within_sequence", +) -> torch.Tensor: + if flex_attention is None: + raise RuntimeError("Flex Attention is not available in this environment.") + if score_mod is not None: + raise NotImplementedError("E1 Flex Attention does not support score_mod.") + query_states = query_states.transpose(1, 2).contiguous() # (b, h, l, d) + key_states = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d) + value_states = value_states.transpose(1, 2).contiguous() # (b, h_kv, l, d) + + fn = _get_flex_attention_fn( + device=query_states.device, + dtype=query_states.dtype, + shape=tuple(query_states.shape), + sequence_lengths=sequence_lengths, + mask_semantics=mask_semantics, + ) + if fn is None: + raise RuntimeError("Flex Attention is not available in this environment.") + outputs = fn( # (b, h, l, d) + query_states, + key_states, + value_states, + block_mask=block_mask, + score_mod=score_mod, + enable_gqa=query_states.shape[1] != key_states.shape[1], # if nkv != nh + ) + + outputs = outputs.transpose(1, 2) # (b, l, h, d) + return outputs # (b, l, h, d) + + +def kernels_flash_attention_func( + query_states: torch.Tensor, # (b, l_q, h, d) + key_states: torch.Tensor, # (b, l_kv, h_kv, d) + value_states: torch.Tensor, # (b, l_kv, h_kv, d) + q_sequence_ids: torch.Tensor, + k_sequence_ids: torch.Tensor, + causal: bool = False, + implementation: str = "flash_attention_3", +) -> torch.Tensor: # (b, l_q, h, d) + # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv) + _ensure_flash_kernels_loaded(implementation) + + if not causal: + batch_size, q_len = query_states.shape[0], query_states.shape[1] + ( + query_states, + key_states, + value_states, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) = _unpad_input( # Q: (t_q, h, d); K/V: (t_kv, h_kv, d) + query_states, + key_states, + value_states, + q_sequence_ids, + k_sequence_ids, + ) + + attn_output_unpad = _kernels_flash_varlen_forward( # (t_q, h, d) + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_in_batch_q=max_seqlen_in_batch_q, + max_seqlen_in_batch_k=max_seqlen_in_batch_k, + causal=False, + implementation=implementation, + ) + attn_output = pad_input( # (b, l_q, h, d) + attn_output_unpad, + indices_q, + batch_size, + q_len, + ) + + else: + attn_output = _kernels_flash_forward( # (b, l_q, h, d) + query_states, key_states, value_states, causal=True, implementation=implementation + ) + + return attn_output # (b, l_q, h, d) + + +def block_min_max_seq_ids( + sequence_lengths: torch.Tensor, + block_size: int = 128, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map each physical attention block to its first and last sequence.""" + + # sequence_lengths: (n,) + total_tokens = sequence_lengths.sum() # () + block_count = int( + torch.div( + total_tokens + block_size - 1, + block_size, + rounding_mode="floor", + ).item() + ) + padded_tokens = block_count * block_size - total_tokens # () + lengths_with_tail = torch.cat( # (n + 1,) + (sequence_lengths, padded_tokens.to(sequence_lengths).reshape(1)), + ) + sequence_ends = lengths_with_tail.to(torch.long).cumsum(dim=0) # (n + 1,) + block_starts = torch.arange( # (n_blocks,) + start=0, + end=block_count * block_size, + step=block_size, + dtype=torch.long, + device=sequence_lengths.device, + ) + block_last_tokens = block_starts + block_size - 1 # (n_blocks,) + first_sequence = torch.searchsorted(sequence_ends, block_starts, right=True) # (n_blocks,) + last_sequence = torch.searchsorted(sequence_ends, block_last_tokens, right=True) # (n_blocks,) + return first_sequence, last_sequence # (n_blocks,), (n_blocks,) + + +def get_overlapping_blocks( + q_lengths: torch.Tensor, + k_lengths: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Classify query/key block pairs as full, partial, or disjoint.""" + + # q_lengths: (n_q,); k_lengths: (n_k,) + q_first, q_last = block_min_max_seq_ids(q_lengths) # (q_blocks,), (q_blocks,) + k_first, k_last = block_min_max_seq_ids(k_lengths) # (k_blocks,), (k_blocks,) + intersection_start = torch.maximum( # (q_blocks, k_blocks) + q_first[:, None], + k_first[None, :], + ) + intersection_end = torch.minimum( # (q_blocks, k_blocks) + q_last[:, None], + k_last[None, :], + ) + intersects = intersection_start <= intersection_end # (q_blocks, k_blocks) + both_blocks_are_single_sequence = ( # (q_blocks, k_blocks) + (q_first == q_last)[:, None] & (k_first == k_last)[None, :] + ) + full_blocks = intersects & both_blocks_are_single_sequence # (q_blocks, k_blocks) + return full_blocks, intersects & ~both_blocks_are_single_sequence # both (q_blocks, k_blocks) + + +def _document_ids(sequence_lengths: torch.Tensor) -> torch.Tensor: + # sequence_lengths: (n,) + sequence_numbers = torch.arange( # (n,) + sequence_lengths.numel(), + device=sequence_lengths.device, + dtype=torch.long, + ) + return sequence_numbers.repeat_interleave(sequence_lengths.to(torch.long)) # (t,) + + +@torch.compiler.disable +def direct_block_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask: + """Build a packed-sequence mask from preclassified sparse blocks.""" + + full, partial = get_overlapping_blocks(q_lengths, k_lengths) # both (q_blocks, k_blocks) + q_document = _document_ids(q_lengths) # (t_q,) + k_document = _document_ids(k_lengths) # (t_k,) + + def same_document( + _batch: torch.Tensor, + _head: torch.Tensor, + q_index: torch.Tensor, + k_index: torch.Tensor, + ) -> torch.Tensor: + return q_document[q_index].eq(k_document[k_index]) + + return _create_sparse_block_from_block_mask( + (partial[None, None], full[None, None]), + same_document, + seq_lengths=(q_document.numel(), k_document.numel()), + Q_BLOCK_SIZE=128, + KV_BLOCK_SIZE=128, + ) + + +@torch.compiler.disable +def doc_id_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask: + if create_block_mask is None: + raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.") + q_document = _document_ids(q_lengths) # (t_q,) + k_document = _document_ids(k_lengths) # (t_k,) + + def same_document( + _batch: torch.Tensor, + _head: torch.Tensor, + q_index: torch.Tensor, + k_index: torch.Tensor, + ) -> torch.Tensor: + return q_document[q_index].eq(k_document[k_index]) + + return create_block_mask( + same_document, + 1, + 1, + q_document.numel(), + k_document.numel(), + BLOCK_SIZE=128, + device=q_lengths.device, + ) + + +def varlen_flex_attention_func( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + q_sequence_ids: torch.Tensor, + k_sequence_ids: torch.Tensor, +) -> torch.Tensor: + # query_states: (b, l_q, h, d); key_states, value_states: (b, l_kv, h_kv, d) + # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv) + if flex_attention is None: + raise RuntimeError("Flex Attention is not available in this environment.") + batch_size, q_len = query_states.shape[0], query_states.shape[1] + ( + query_states, + key_states, + value_states, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (_max_seqlen_in_batch_q, _max_seqlen_in_batch_k), + ) = _unpad_input( # Q: (t_q, h, d); K/V: (t_kv, h_kv, d) + query_states, + key_states, + value_states, + q_sequence_ids, + k_sequence_ids, + ) + + query_states = query_states.unsqueeze(0).transpose(1, 2).contiguous() # (1, h, t_q, d) + key_states = key_states.unsqueeze(0).transpose(1, 2).contiguous() # (1, h_kv, t_kv, d) + value_states = value_states.unsqueeze(0).transpose(1, 2).contiguous() # (1, h_kv, t_kv, d) + + seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] # (n,) + seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1] # (n,) + block_mask = block_mask_creator(seqlens_q, seqlens_k) + + packed_lengths = ( + *(int(length) for length in seqlens_q.tolist()), + -1, + *(int(length) for length in seqlens_k.tolist()), + ) + fn = _get_flex_attention_fn( + device=query_states.device, + dtype=query_states.dtype, + shape=tuple(query_states.shape) + tuple(key_states.shape), + sequence_lengths=packed_lengths, + mask_semantics="packed_document_equality", + ) + if fn is None: + raise RuntimeError("Flex Attention is not available in this environment.") + attn_output_unpad = fn( # (1, h, t_q, d) + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=query_states.shape[1] != key_states.shape[1], + ) + + attn_output = pad_input( # (b, l_q, h, d) + attn_output_unpad.transpose(1, 2).squeeze(0), indices_q, batch_size, q_len + ) + + return attn_output # (b, l_q, h, d) + + +def _get_unpad_data(sequence_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """Return packed indices and run lengths for the non-padding sequence IDs.""" + + # sequence_ids: (b, l) + flat_ids = sequence_ids.reshape(-1) # (b * l,) + non_pad_indices = torch.where(flat_ids.ne(-1))[0] # (t,) + if non_pad_indices.numel() == 0: + raise ValueError("Packed attention requires at least one non-padding token.") + + valid_ids = flat_ids.index_select(0, non_pad_indices) # (t,) + row_ids = torch.div( # (t,) + non_pad_indices, + sequence_ids.shape[1], + rounding_mode="floor", + ) + run_starts = torch.ones_like(valid_ids, dtype=torch.bool) # (t,) + run_starts[1:] = (valid_ids[1:] != valid_ids[:-1]) | (row_ids[1:] != row_ids[:-1]) + start_indices = torch.where(run_starts)[0] # (n,) + end_indices = torch.cat( # (n,) + (start_indices[1:], start_indices.new_tensor([valid_ids.numel()])) + ) + sequence_lengths = end_indices - start_indices # (n,) + cumulative_lengths = torch.cat( # (n + 1,) + ( + torch.zeros(1, dtype=torch.int32, device=sequence_ids.device), + sequence_lengths.cumsum(dim=0, dtype=torch.int32), + ), + ) + return non_pad_indices, cumulative_lengths, int( # (t,), (n + 1,), scalar + sequence_lengths.max().item() + ) + + +def _unpad_input( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + q_sequence_ids: torch.Tensor, + k_sequence_ids: torch.Tensor, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, torch.Tensor], + tuple[int, int], +]: + # query_layer: (b, l_q, h, d); key_layer, value_layer: (b, l_kv, h_kv, d) + # q_sequence_ids: (b, l_q); k_sequence_ids: (b, l_kv) + for name, layer in ( + ("query_layer", query_layer), + ("key_layer", key_layer), + ("value_layer", value_layer), + ): + if layer.ndim != 4: + raise ValueError( + f"{name} must have shape (batch, sequence, heads, head_dim); " + f"got {tuple(layer.shape)}." + ) + if value_layer.shape != key_layer.shape: + raise ValueError( + "key_layer and value_layer must have identical shapes; " + f"got {tuple(key_layer.shape)} and {tuple(value_layer.shape)}." + ) + if query_layer.shape[0] != key_layer.shape[0]: + raise ValueError( + "Query and KV batch sizes must match; " + f"got {query_layer.shape[0]} and {key_layer.shape[0]}." + ) + if query_layer.shape[-1] != key_layer.shape[-1]: + raise ValueError( + "Query and KV head dimensions must match; " + f"got {query_layer.shape[-1]} and {key_layer.shape[-1]}." + ) + batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape + query_length, num_q_heads = query_layer.shape[1], query_layer.shape[2] + if query_layer.shape[:2] != q_sequence_ids.shape: + raise ValueError( + "Shape mismatch between query layer and query sequence ids: " + f"{query_layer.shape[:2]} != {q_sequence_ids.shape}" + ) + if key_layer.shape[:2] != k_sequence_ids.shape: + raise ValueError( + "Shape mismatch between key layer and key sequence ids: " + f"{key_layer.shape[:2]} != {k_sequence_ids.shape}" + ) + if query_length > kv_seq_len: + raise ValueError( + "Query length must be less than or equal to KV sequence length: " + f"{query_length} > {kv_seq_len}" + ) + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data( # (t_kv,), (n + 1,), scalar + k_sequence_ids + ) + + key_layer = index_first_axis( # (t_kv, h_kv, d) + key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + value_layer = index_first_axis( # (t_kv, h_kv, d) + value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + + if torch.equal(q_sequence_ids, k_sequence_ids): + indices_q = indices_k + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + else: + # (t_q,), (n + 1,), scalar + indices_q, cu_seqlens_q, max_seqlen_in_batch_q = _get_unpad_data(q_sequence_ids) + + query_layer = index_first_axis( # (t_q, h, d) + query_layer.reshape(batch_size * query_length, num_q_heads, head_dim), indices_q + ) + + if cu_seqlens_q.shape != cu_seqlens_k.shape: + raise ValueError( + "Query and KV must have the same number of sequences: " + f"{cu_seqlens_q.shape} != {cu_seqlens_k.shape}" + ) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +block_mask_creator = direct_block_mask if os.getenv("FAST_BLOCK_MASK", "1") == "1" else doc_id_mask diff --git a/src/fastplms/models/e1/cache.py b/src/fastplms/models/e1/cache.py new file mode 100644 index 0000000..044b79a --- /dev/null +++ b/src/fastplms/models/e1/cache.py @@ -0,0 +1,240 @@ +"""Key-value cache implementations used by E1 inference.""" + +from __future__ import annotations + +import torch +from typing import Any +from transformers.modeling_outputs import ModelOutput +from transformers.utils import logging + + +def _get_logger(): + """Resolve the Transformers logger only when a cache path emits a message.""" + + return logging.get_logger(__name__) + + +class DynamicCache: + """A cache that grows K and V along their sequence dimension. + + Each cached tensor has shape (b, l, h, d). + + Args: + key_cache (`list[torch.Tensor]`): The list of key states. + value_cache (`list[torch.Tensor]`): The list of value states. + """ + + def __init__(self) -> None: + self.key_cache: list[torch.Tensor] = [] + self.value_cache: list[torch.Tensor] = [] + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Update the key and value caches in-place, and return the necessary keys and value states. + + Args: + key_states (`torch.Tensor`): K to cache with shape (b, l, h, d). + value_states (`torch.Tensor`): V to cache with shape (b, l, h, d). + layer_idx (`int`): The index of the layer to update. + + Returns: + tuple[`torch.Tensor`, `torch.Tensor`]: Cached K and V, each with shape + (b, l, h, d). + """ + # key_states, value_states: (b, l_new, h, d) + if len(self.key_cache) <= layer_idx: + # Empty tensors preserve skipped layer indices until those layers receive state. + for _ in range(len(self.key_cache), layer_idx): + self.key_cache.append(torch.tensor([])) + self.value_cache.append(torch.tensor([])) + self.key_cache.append(key_states) + self.value_cache.append(value_states) + elif ( + not self.key_cache[ + layer_idx + ].numel() # prefers not t.numel() to len(t) == 0 to export the model + ): # fills previously skipped layers; checking for tensor causes errors + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat( # (b, l_cached + l_new, h, d) + [self.key_cache[layer_idx], key_states], + dim=1, + ) + self.value_cache[layer_idx] = torch.cat( # (b, l_cached + l_new, h, d) + [self.value_cache[layer_idx], value_states], dim=1 + ) + + return ( # (b, l_total, h, d), (b, l_total, h, d) + self.key_cache[layer_idx], + self.value_cache[layer_idx], + ) + + def get_seq_length(self, layer_idx: int = 0) -> int: + """Return the cached sequence length for one layer.""" + is_empty_layer = ( + len(self.key_cache) == 0 # no cache in any layer + or len(self.key_cache) + <= layer_idx # skipped `layer_idx` and hasn't run a layer with cache after it + or not self.key_cache[layer_idx].numel() # the layer has no cache + ) + layer_seq_length = self.key_cache[layer_idx].shape[1] if not is_empty_layer else 0 + return layer_seq_length + + def crop(self, max_length: int) -> None: + """Crop every cached K and V tensor to ``max_length`` tokens.""" + if max_length <= 0: + raise ValueError("max_length must be positive") + + if self.get_seq_length() <= max_length: + return + + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx].numel(): + self.key_cache[layer_idx] = self.key_cache[layer_idx][:, :max_length, ...] + self.value_cache[layer_idx] = self.value_cache[layer_idx][:, :max_length, ...] + + def batch_repeat_interleave(self, repeats: int) -> None: + """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx].numel(): + # (b * repeats, l, h, d) + self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave( + repeats, dim=0 + ) + self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave( + repeats, dim=0 + ) # (b * repeats, l, h, d) + + def batch_select_indices(self, indices: torch.Tensor | list[int]) -> None: + """Keep selected rows of the cache batch dimension.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx].numel(): + self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...] # (n, l, h, d) + self.value_cache[layer_idx] = self.value_cache[layer_idx][ + indices, ... + ] # (n, l, h, d) + + +class KVCache: + def __init__(self, cache_size: int = 4) -> None: + self.cache_size = cache_size + self.tensor_input_field_names = [ + "input_ids", + "within_seq_position_ids", + "global_position_ids", + "sequence_ids", + "labels", + ] + # Upstream E1 called the encoder output ``embeddings``. FastPLMs uses + # the standard Transformers ``last_hidden_state`` name, while keeping + # the aliases here makes the cache safe for either output contract. + self.tensor_output_field_names = [ + "logits", + "last_hidden_state", + "embeddings", + "token_embeddings", + ] + self.cache_dict: dict[str, DynamicCache] = {} + self.cache_queue: list[str] = [] + + def reset(self) -> None: + for k in list(self.cache_dict.keys()): + del self.cache_dict[k] + del self.cache_dict + self.cache_dict = {} + self.cache_queue = [] + + torch.cuda.empty_cache() + + def before_forward(self, batch: dict[str, torch.Tensor]) -> None: + contexts: list[str] | None = batch.get("context") + if contexts is None or "context_len" not in batch: + _get_logger().warning_once( + "KVCache requires both `context` and `context_len`; cache setup was skipped." + ) + return + + context_lens: list[int] = list(set(batch["context_len"])) + contexts: list[str] = list(set(contexts)) # type: ignore[no-redef] + if len(contexts) != 1 or len(context_lens) != 1: + _get_logger().warning( + "SingleContextKVCache requires a single context and context length. " + "Multiple contexts or context lengths found in a single batch. Skipping." + ) + return + + batch_size = batch["input_ids"].shape[0] # b + + unique_context = contexts[0] + unique_context_len = context_lens[0] + batch["use_cache"] = True + + if unique_context not in self.cache_dict: + return + + self.cache_dict[unique_context].batch_repeat_interleave(batch_size) + past_key_values = self.cache_dict[unique_context] + batch["past_key_values"] = past_key_values + + # A cached prefix leaves only query-suffix tokens for the model call. + for field_name in self.tensor_input_field_names: + if batch.get(field_name) is not None: + batch[field_name] = batch[field_name][:, unique_context_len:] # (b, l_suffix, ...) + + def after_forward(self, batch: dict[str, Any], outputs: ModelOutput) -> None: + contexts = batch.get("context") + context_lens = batch.get("context_len", []) + if ( + contexts is None + or len(set(contexts)) != 1 + or len(set(context_lens)) != 1 + or context_lens[0] == 0 + ): + return + + if not batch.get("use_cache", False): + raise ValueError("E1 retrieval cache updates require use_cache=True.") + unique_context = contexts[0] + unique_context_len = context_lens[0] + + past_key_values = getattr(outputs, "past_key_values", None) + if not isinstance(past_key_values, DynamicCache): + _get_logger().warning_once( + "KVCache is incompatible with models that don't return a DynamicCache. Skipping." + ) + return + + if "past_key_values" not in batch: + if len(self.cache_queue) == self.cache_size: + last_context = self.cache_queue.pop(0) + if last_context not in self.cache_queue: + del self.cache_dict[last_context] + torch.cuda.empty_cache() + + self.cache_dict[unique_context] = past_key_values + self.cache_queue.append(unique_context) + + # The first uncached call returns the full sequence; expose its query suffix. + for field_name in self.tensor_input_field_names: + if field_name in batch and batch[field_name] is not None: + batch[field_name] = batch[field_name][ + :, unique_context_len: + ] # (b, l_suffix, ...) + + for field_name in self.tensor_output_field_names: + if field_name in outputs and outputs[field_name] is not None: + outputs[field_name] = outputs[field_name][ + :, unique_context_len: + ] # (b, l_suffix, ...) + if "hidden_states" in outputs and outputs["hidden_states"] is not None: + hidden_states = outputs["hidden_states"] + sliced_hidden_states = tuple( # each: (b, l_suffix, d) + hidden_state[:, unique_context_len:] for hidden_state in hidden_states + ) + outputs["hidden_states"] = sliced_hidden_states + + self.cache_dict[unique_context].crop(unique_context_len) + self.cache_dict[unique_context].batch_select_indices([0]) diff --git a/src/fastplms/models/e1/modeling_e1.py b/src/fastplms/models/e1/modeling_e1.py new file mode 100644 index 0000000..8b372b9 --- /dev/null +++ b/src/fastplms/models/e1/modeling_e1.py @@ -0,0 +1,2328 @@ +from __future__ import annotations + +import hashlib +import os +import sys +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import defaultdict +from contextvars import ContextVar +from dataclasses import dataclass +from enum import Enum +from typing import Any, ClassVar, TypedDict +from tqdm.auto import tqdm +from transformers import PretrainedConfig, PreTrainedModel +from transformers.activations import ACT2FN +from transformers.modeling_outputs import ModelOutput +from transformers.utils import logging + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import ( + EmbeddingBatch, + EmbeddingMixin, + EmbeddingResult, + Pooler, + embed_dataset, + select_hidden_state_embeddings, + ) + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingBatch", + "EmbeddingMixin", + "EmbeddingResult", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "Pooler", + "embed_dataset", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + +from .attention import ( # noqa: F401 + _document_ids, + _get_unpad_data, + _unpad_input, + block_mask_creator, + block_min_max_seq_ids, + build_block_causal_mask_4d, + build_within_seq_mask_4d, + create_block_causal_mask_optimized, + create_within_seq_block_mask, + direct_block_mask, + doc_id_mask, + flex_attention_func, + get_overlapping_blocks, + kernels_flash_attention_func, + varlen_flex_attention_func, +) +from .cache import DynamicCache, KVCache # noqa: F401 +from .preparation import ( # noqa: F401 + BOS_TOKEN_ID, + E1_TOKENIZER_REPO_ID, + E1_VOCAB_SIZE, + EOS_TOKEN_ID, + PAD_TOKEN_ID, + DataPrepConfig, + E1BatchPreparer, + _load_tokenizer_file, + get_context, + get_tokenizer, +) +from .retrieval import ( # noqa: F401 + COLABFOLD_HOST, + DEFAULT_EMBED_MAX_TOKENS, + DEFAULT_EMBED_SIMILARITY, + DEFAULT_MAX_CONTEXT_TOKENS, + DEFAULT_SIMILARITY_THRESHOLDS, + DOCKER_IMAGE, + E1_MSA_SAMPLING_SOURCE_REVISION, + LOWERCASE_CHARS, + ColabFoldSearcher, + ContextCache, + ContextSpecification, + E1Prediction, + HomologueSearcher, + IdSequence, + IndexedSequence, + _ColabFoldResponse, + _E1ContextPredictor, + _forward_for_embedding, + _make_homologue_searcher, + _pool_hidden_states, + _safe_extract_tar, + _sequence_output_dir, + _strip_a3m_insertions, + build_context_specifications, + compute_ppll, + convert_to_tensor, + get_context_id, + get_msa_for_sequence, + get_num_neighbors, + get_query_from_a3m, + get_similarity_to_query, + load_msa_dir, + load_msa_from_hf, + parse_msa, + read_fasta_sequences, + sample_context, + sample_contexts_for_msa, + sample_multiple_contexts, + write_fasta_sequences, +) + + +def _get_logger(): + """Resolve the Transformers logger only when a runtime path emits a message.""" + + return logging.get_logger(__name__) + + +_TOKENIZER_LOAD_CONTEXT: ContextVar[dict[str, Any] | None] = ContextVar( + "fastplms_e1_tokenizer_load_context", + default=None, +) + + +class E1Config(PretrainedConfig): + model_type = "E1" + keys_to_ignore_at_inference: ClassVar[list[str]] = ["past_key_values"] + + def __init__( # type: ignore + self, + # Model architecture/initialization + vocab_size=None, + hidden_size=4096, + intermediate_size=16384, + gated_mlp=False, + num_hidden_layers=40, + num_attention_heads=32, + num_key_value_heads=8, + hidden_act="silu", + rms_norm_eps=1e-5, + initializer_range=0.02, + dtype="bfloat16", + gradient_checkpointing=False, + no_ffn_gradient_checkpointing=False, + use_cache=False, + # Tokenization + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + tie_word_embeddings=False, + # Attention implementation & rotary positional embeddings + global_attention_every_n_layers=0, + max_num_sequences=512, + max_num_positions_within_seq=8192, + max_num_positions_global=1024 * 128, + rope_theta_within_seq=10000.0, + rope_theta_global=100000.0, + clip_qkv=None, + attn_backend=None, + **kwargs, + ) -> None: + super().__init__( + pad_token_id=PAD_TOKEN_ID, + bos_token_id=BOS_TOKEN_ID, + eos_token_id=EOS_TOKEN_ID, + tie_word_embeddings=tie_word_embeddings, + dtype=dtype, + **kwargs, + ) + + self.hidden_size = hidden_size + if intermediate_size is None: + intermediate_size = 3 * hidden_size if gated_mlp else 4 * hidden_size + self.intermediate_size = intermediate_size + self.gated_mlp = gated_mlp + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.max_num_positions_within_seq = max_num_positions_within_seq + self.max_num_positions_global = max_num_positions_global + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.rope_theta_within_seq = rope_theta_within_seq + self.rope_theta_global = rope_theta_global + self.max_num_sequences = max_num_sequences + if clip_qkv is not None and clip_qkv <= 0: + raise ValueError(f"clip_qkv must be positive when provided, got {clip_qkv}.") + self.clip_qkv = clip_qkv + self.global_attention_every_n_layers = global_attention_every_n_layers + + self.vocab_size = E1_VOCAB_SIZE + self.gradient_checkpointing = gradient_checkpointing + self.no_ffn_gradient_checkpointing = no_ffn_gradient_checkpointing + if not isinstance(use_cache, bool): + raise TypeError("use_cache must be a boolean.") + self.use_cache = use_cache + self.attn_backend = attn_backend + + if vocab_size is not None: + if vocab_size < self.vocab_size: + _get_logger().warning( + f"Using vocab_size {vocab_size} smaller than {self.vocab_size} " + "from the tokenizer contract." + ) + self.vocab_size = vocab_size + elif vocab_size > self.vocab_size: + _get_logger().warning( + f"Using vocab_size {vocab_size} instead of smaller {self.vocab_size} " + "from E1 tokenizer contract." + ) + self.vocab_size = vocab_size + if pad_token_id is not None and pad_token_id != self.pad_token_id: + _get_logger().warning( + f"Ignoring pad_token_id. Using {self.pad_token_id} from E1 tokenizer contract" + ) + if bos_token_id is not None and bos_token_id != self.bos_token_id: + _get_logger().warning( + f"Ignoring bos_token_id. Using {self.bos_token_id} from E1 tokenizer contract" + ) + if eos_token_id is not None and eos_token_id != self.eos_token_id: + _get_logger().warning( + f"Ignoring eos_token_id. Using {self.eos_token_id} from E1 tokenizer contract" + ) + + +class AttentionLayerType(Enum): + WITHIN_SEQ = "within_seq" + GLOBAL = "global" + + +class AttentionArgs(TypedDict, total=False): + within_seq_block_mask: BlockMask | None + block_causal_block_mask: BlockMask | None + within_seq_mask_4d: torch.Tensor | None + block_causal_mask_4d: torch.Tensor | None + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). + + The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, + num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class RotaryPositionalEmbedding(nn.Module): + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: int = 10000, + device: torch.device | None = None, + ) -> None: + super().__init__() + + self.dim = dim + self.base = base + self.max_position_embeddings = max_position_embeddings + # Transformers may instantiate modules on the meta device while loading + # a checkpoint. Precomputed non-persistent buffers would then be + # materialized without values. Empty buffers make initialization lazy + # and deterministic on the first real-device forward. + empty = torch.empty(0, dtype=torch.float32, device=device) + self.register_buffer("inv_freq", empty, persistent=False) + self.register_buffer("cos_cached", empty.clone(), persistent=False) + self.register_buffer("sin_cached", empty.clone(), persistent=False) + self.max_seq_len_cached = 0 + + @staticmethod + def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _set_sin_cos_cache(self, seq_len: int, device: torch.device) -> None: + # Compute angles in FP32, matching the official cache constructed before + # the model is converted to its inference dtype. + self.max_seq_len_cached = seq_len + inv_freq = self.base ** -( + torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim + ) # (d / 2,) + self.inv_freq = inv_freq + t = torch.arange(seq_len, device=device, dtype=torch.float32) # (l,) + angles = torch.outer(t, inv_freq) # (l, d / 2) + angles = torch.cat((angles, angles), dim=1) # (l, d) + self.cos_cached = angles.cos() + self.sin_cached = angles.sin() + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: torch.LongTensor, + seq_len: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, l, h, d) + device, dtype = q.device, q.dtype + seq_len = position_ids.max().item() + 1 if seq_len is None else seq_len + + if seq_len > self.max_seq_len_cached: + self._set_sin_cos_cache(seq_len=seq_len, device=device) + + # Selecting by position inserts a head axis for broadcasting. + idxs = position_ids.to(device) + cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] # (b, l, 1, d) + sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] # (b, l, 1, d) + + # Apply the real and imaginary parts of the rotary transform to Q and K. + # Both halves reuse C and S, so rotate_half supplies the cross terms. + q_embed = (q * cos) + (self.rotate_half(q) * sin) + k_embed = (k * cos) + (self.rotate_half(k) * sin) + return q_embed, k_embed + + +class Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper.""" + + def __init__(self, config: E1Config, layer_idx: int) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_kv_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_kv_heads + self.max_num_seqs = config.max_num_sequences + self.clip_qkv = config.clip_qkv + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + if self.config.global_attention_every_n_layers > 0: + self.layer_type = ( + AttentionLayerType.GLOBAL + if (self.layer_idx + 1) % self.config.global_attention_every_n_layers == 0 + else AttentionLayerType.WITHIN_SEQ + ) + else: + self.layer_type = AttentionLayerType.WITHIN_SEQ + + self.rope_theta = ( + config.rope_theta_within_seq + if self.layer_type == AttentionLayerType.WITHIN_SEQ + else config.rope_theta_global + ) + self.max_position_embeddings = ( + config.max_num_positions_within_seq + if self.layer_type == AttentionLayerType.WITHIN_SEQ + else config.max_num_positions_global + ) + + self.rotary_emb = RotaryPositionalEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + self.attn_backend = resolve_attention_backend(config.attn_backend) + + def prepare_qkv( + self, + hidden_states: torch.Tensor, + position_ids: torch.LongTensor, + past_key_value: DynamicCache | None = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # hidden_states: (b, l, d); position_ids: (b, l) + bsz, q_len, _ = hidden_states.size() + query_states: torch.Tensor = self.q_proj(hidden_states) + key_states: torch.Tensor = self.k_proj(hidden_states) + val_states: torch.Tensor = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ) # (b, l, h, d_h) + key_states = key_states.view( + bsz, q_len, self.num_kv_heads, self.head_dim + ) # (b, l, h_kv, d_h) + val_states = val_states.view( + bsz, q_len, self.num_kv_heads, self.head_dim + ) # (b, l, h_kv, d_h) + + if self.clip_qkv is not None: + query_states = query_states.clamp(-self.clip_qkv, self.clip_qkv) + key_states = key_states.clamp(-self.clip_qkv, self.clip_qkv) + val_states = val_states.clamp(-self.clip_qkv, self.clip_qkv) + + query_states, key_states = self.rotary_emb(query_states, key_states, position_ids) + + if use_cache and past_key_value is not None: + key_states, val_states = past_key_value.update(key_states, val_states, self.layer_idx) + + input_dtype = query_states.dtype + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_dtype("cuda") + else: + target_dtype = self.q_proj.weight.dtype + if input_dtype != target_dtype: + _get_logger().warning_once( + f"The input hidden states seems to be silently casted in {input_dtype}. " + f"This might be because you have upcasted embedding or layer norm layers " + f"in {input_dtype}. We will cast back the input in {target_dtype}." + ) + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + val_states = val_states.to(target_dtype) + + return query_states, key_states, val_states + + def forward( + self, + hidden_states: torch.Tensor, + within_seq_position_ids: torch.LongTensor, + global_position_ids: torch.LongTensor, + sequence_ids: torch.LongTensor, + attention_args: AttentionArgs | None = None, + past_key_value: DynamicCache | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + use_cache: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d); position and sequence IDs: (b, l) + is_cache_prefilled = ( + use_cache + and past_key_value is not None + and past_key_value.get_seq_length(self.layer_idx) > 0 + ) + + query_states, key_states, val_states = self.prepare_qkv( + hidden_states=hidden_states, + position_ids=within_seq_position_ids + if self.layer_type == AttentionLayerType.WITHIN_SEQ + else global_position_ids, + past_key_value=past_key_value, + use_cache=use_cache, + ) + + attn_output, attn_weights, s_max = self._attn( + query_states=query_states, + key_states=key_states, + val_states=val_states, + sequence_ids=sequence_ids, + attention_args=attention_args, + output_attentions=output_attentions, + output_s_max=output_s_max, + is_cache_prefilled=is_cache_prefilled, + effective_backend=effective_backend, + ) + + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights, past_key_value, s_max + + def _attn( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + val_states: torch.Tensor, + sequence_ids: torch.Tensor, + attention_args: AttentionArgs | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + is_cache_prefilled: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # A filled cache changes the implementation shape, not the layer's + # biological attention contract. Global layers must retain the cached + # context, while within-sequence layers consume only the newly appended + # sequence. This matches the pinned E1 inference implementation. + effective_layer_type = self.layer_type + + if effective_backend is None: + effective_backend = resolve_attention_backend_for_call( + self.attn_backend, + output_attentions=output_attentions, + ) + if output_attentions: + return self._manual_attn( + query_states, + key_states, + val_states, + sequence_ids=sequence_ids, + attention_args=attention_args, + effective_layer_type=effective_layer_type, + output_s_max=output_s_max, + is_cache_prefilled=is_cache_prefilled, + ) + + if effective_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_states, + key_states, + val_states, + sequence_ids=sequence_ids, + attention_args=attention_args, + effective_layer_type=effective_layer_type, + output_s_max=output_s_max, + is_cache_prefilled=is_cache_prefilled, + ) + return attn_output, None, s_max + if effective_backend.is_flash: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + attn_output, attn_weights = self._kernels_flash_attn( + query_states, + key_states, + val_states, + sequence_ids=sequence_ids, + is_cache_prefilled=is_cache_prefilled, + ) + else: + raise ValueError( + "E1 global attention does not support a kernels Flash backend; " + "use eager, sdpa, or flex_attention." + ) + elif effective_backend == AttentionBackend.FLEX: + attn_output, attn_weights = self._flex_attn( + query_states, + key_states, + val_states, + sequence_ids=sequence_ids, + attention_args=attention_args, + effective_layer_type=effective_layer_type, + is_cache_prefilled=is_cache_prefilled, + ) + elif effective_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_states, + key_states, + val_states, + sequence_ids=sequence_ids, + attention_args=attention_args, + effective_layer_type=effective_layer_type, + is_cache_prefilled=is_cache_prefilled, + ) + else: + raise AssertionError(f"Unsupported resolved backend: {effective_backend}") + + s_max_key_states = key_states + if ( + is_cache_prefilled + and effective_layer_type == AttentionLayerType.WITHIN_SEQ + and query_states.shape[1] < key_states.shape[1] + ): + s_max_key_states = key_states[:, -query_states.shape[1] :] + s_max = self._compute_s_max(query_states, s_max_key_states) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, + query_states: torch.Tensor, # Q has shape (b, l, h, d). + key_states: torch.Tensor, # K has shape (b, l, h_kv, d). + ) -> list[torch.Tensor]: + query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) + key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) + key_heads = repeat_kv(key_heads, self.num_key_value_groups) + scale = 1.0 / (self.head_dim**0.5) + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) + s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max( + dim=0 + ).values * scale + return [s_max_bound[h] for h in range(self.num_heads)] + + def _kernels_flash_attn( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + val_states: torch.Tensor, + sequence_ids: torch.Tensor, + is_cache_prefilled: bool = False, + ) -> tuple[torch.Tensor, None]: + bsz, q_len = query_states.shape[0], query_states.shape[1] + _, kv_len = key_states.shape[0], key_states.shape[1] + + if self.layer_type == AttentionLayerType.GLOBAL: + q_sequence_ids = sequence_ids + if q_len < kv_len: + first_token_id = sequence_ids[:, 0].unsqueeze(1) + k_sequence_ids = torch.cat( + [first_token_id.expand(bsz, kv_len - q_len), sequence_ids], dim=-1 + ) + else: + k_sequence_ids = sequence_ids + else: + if q_len < kv_len: + key_states = key_states[:, -q_len:] + val_states = val_states[:, -q_len:] + q_sequence_ids = k_sequence_ids = sequence_ids + + attn_output = kernels_flash_attention_func( + query_states, + key_states, + val_states, + q_sequence_ids=q_sequence_ids, + k_sequence_ids=k_sequence_ids, + causal=False, + implementation=self.attn_backend.value, + ) + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + return attn_output, None + + def _flex_attn( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + val_states: torch.Tensor, + sequence_ids: torch.Tensor, + attention_args: AttentionArgs | None = None, + effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, + is_cache_prefilled: bool = False, + ) -> tuple[torch.Tensor, None]: + bsz, q_len = query_states.shape[0], query_states.shape[1] + kv_len = key_states.shape[1] + if is_cache_prefilled and q_len < kv_len: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + key_states = key_states[:, -q_len:] + val_states = val_states[:, -q_len:] + block_mask = create_within_seq_block_mask(sequence_ids) + outputs = flex_attention_func( + query_states, + key_states, + val_states, + block_mask=block_mask, + mask_semantics=effective_layer_type.value, + ) + else: + q_sequence_ids, k_sequence_ids = self._cached_global_sequence_ids( + sequence_ids, + kv_len, + ) + outputs = varlen_flex_attention_func( + query_states, + key_states, + val_states, + q_sequence_ids=q_sequence_ids, + k_sequence_ids=k_sequence_ids, + ) + outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous() + return outputs, None + + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + block_mask = ( + attention_args["within_seq_block_mask"] if attention_args is not None else None + ) + else: + block_mask = ( + attention_args["block_causal_block_mask"] if attention_args is not None else None + ) + outputs = flex_attention_func( + query_states, + key_states, + val_states, + block_mask=block_mask, + mask_semantics=effective_layer_type.value, + ) + outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous() + return outputs, None + + @staticmethod + def _cached_global_sequence_ids( + query_sequence_ids: torch.Tensor, + kv_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Assign cached context to the incoming query sequence. + + E1 retrieval cache hits contain one incoming sequence. The pinned + implementation relabels the cached prefix with that sequence ID so its + valid query tokens attend the complete cached context, while padding is + excluded by the equality mask or packed Flex path. + """ + + q_len = query_sequence_ids.shape[1] + cached_len = kv_len - q_len + if cached_len < 0: + raise ValueError(f"E1 cached KV length {kv_len} is shorter than query length {q_len}.") + first_sequence_id = query_sequence_ids[:, :1] + if bool(first_sequence_id.eq(-1).any()): + raise ValueError("E1 cached queries must start with a non-padding sequence token.") + cached_sequence_ids = first_sequence_id.expand(-1, cached_len) + key_sequence_ids = torch.cat((cached_sequence_ids, query_sequence_ids), dim=-1) + return query_sequence_ids, key_sequence_ids + + def _cached_attention_mask_4d( + self, + sequence_ids: torch.Tensor, + kv_len: int, + effective_layer_type: AttentionLayerType, + ) -> torch.Tensor: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + return build_within_seq_mask_4d(sequence_ids) + query_sequence_ids, key_sequence_ids = self._cached_global_sequence_ids( + sequence_ids, + kv_len, + ) + query_valid = query_sequence_ids.ne(-1) + key_valid = key_sequence_ids.ne(-1) + same_sequence = query_sequence_ids.unsqueeze(-1).eq(key_sequence_ids.unsqueeze(-2)) + return (same_sequence & query_valid.unsqueeze(-1) & key_valid.unsqueeze(-2)).unsqueeze(1) + + def _sdpa_attn( + self, + query_states: torch.Tensor, # Q has shape (b, l, h, d). + key_states: torch.Tensor, # K has shape (b, l, h_kv, d). + val_states: torch.Tensor, # V has shape (b, l, h_kv, d). + sequence_ids: torch.Tensor, + attention_args: AttentionArgs | None = None, + effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, + is_cache_prefilled: bool = False, + ) -> tuple[torch.Tensor, None]: + bsz, q_len = query_states.shape[:2] + kv_len = key_states.shape[1] + + if is_cache_prefilled and q_len < kv_len: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + key_states = key_states[:, -q_len:] + val_states = val_states[:, -q_len:] + attention_mask_4d = self._cached_attention_mask_4d( + sequence_ids, + kv_len, + effective_layer_type, + ) + elif attention_args is not None: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + attention_mask_4d = attention_args["within_seq_mask_4d"] + else: + attention_mask_4d = attention_args["block_causal_mask_4d"] + else: + attention_mask_4d = None + + query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) + key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) + value_heads = val_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) + key_heads = repeat_kv(key_heads, self.num_key_value_groups) + value_heads = repeat_kv(value_heads, self.num_key_value_groups) + context_heads = F.scaled_dot_product_attention( + query_heads, key_heads, value_heads, attn_mask=attention_mask_4d + ) # (b, h, l, d_h) + attn_output = ( + context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() + ) + return attn_output, None + + def _manual_attn( + self, + query_states: torch.Tensor, # Q has shape (b, l, h, d). + key_states: torch.Tensor, # K has shape (b, l, h_kv, d). + val_states: torch.Tensor, # V has shape (b, l, h_kv, d). + sequence_ids: torch.Tensor, + attention_args: AttentionArgs | None = None, + effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, + output_s_max: bool = False, + is_cache_prefilled: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + bsz, q_len = query_states.shape[:2] + kv_len = key_states.shape[1] + + if is_cache_prefilled and q_len < kv_len: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + key_states = key_states[:, -q_len:] + val_states = val_states[:, -q_len:] + attention_mask_4d = self._cached_attention_mask_4d( + sequence_ids, + kv_len, + effective_layer_type, + ) + elif attention_args is not None: + if effective_layer_type == AttentionLayerType.WITHIN_SEQ: + attention_mask_4d = attention_args["within_seq_mask_4d"] + else: + attention_mask_4d = attention_args["block_causal_mask_4d"] + else: + attention_mask_4d = None + + query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) + key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) + value_heads = val_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) + key_heads = repeat_kv(key_heads, self.num_key_value_groups) + value_heads = repeat_kv(value_heads, self.num_key_value_groups) + scale = 1.0 / (self.head_dim**0.5) + attn_weights = ( + torch.matmul(query_heads, key_heads.transpose(-2, -1)) * scale + ) # (b, h, l, l) + if attention_mask_4d is not None: + attention_mask_4d = attention_mask_4d.to(dtype=torch.bool) + attn_weights = attn_weights.masked_fill( + attention_mask_4d.logical_not(), + torch.finfo(attn_weights.dtype).min, + ) + attn_weights = F.softmax(attn_weights, dim=-1) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), 0.0) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = ( + context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() + ) + s_max = self._compute_s_max(query_states, key_states) if output_s_max else None + return attn_output, attn_weights, s_max + + +class MLP(nn.Module): + def __init__(self, config: E1Config) -> None: + super().__init__() + self.ffn_dim = config.intermediate_size + self.hidden_dim = config.hidden_size + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.w2(self.act_fn(self.w1(hidden_states))) + + +class GLUMLP(nn.Module): + def __init__(self, config: E1Config) -> None: + super().__init__() + self.ffn_dim = config.intermediate_size + self.hidden_dim = config.hidden_size + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states) + hidden_states = self.w2(hidden_states) + return hidden_states + + +class FFN(nn.Module): + def __init__(self, config: E1Config) -> None: + super().__init__() + mlp_cls = GLUMLP if config.gated_mlp else MLP + self.mlp = mlp_cls(config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.mlp(hidden_states) + + +@dataclass +class E1ModelOutputWithPast(ModelOutput): + """E1 encoder outputs. + + ``last_hidden_state`` is H with shape (b, l, d). Optional hidden states use + the same shape per layer, while attention tensors have shape (b, h, l, l). + ``past_key_values`` stores the reusable K and V tensors for cached decoding. + """ + + last_hidden_state: torch.FloatTensor | None = None + past_key_values: DynamicCache | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class E1MaskedLMOutputWithPast(ModelOutput): + """Masked-LM output with the standard HF fields first, then E1 diagnostics.""" + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + mlm_loss: torch.FloatTensor | None = None + last_hidden_state: torch.FloatTensor | None = None + past_key_values: DynamicCache | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class E1ClassificationOutputWithPast(ModelOutput): + """Sequence-classifier output matching HF ``SequenceClassifierOutputWithPast``.""" + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: DynamicCache | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + last_hidden_state: torch.FloatTensor | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class E1TokenClassificationOutputWithPast(ModelOutput): + """Token-classifier output with the standard HF fields before E1 extensions.""" + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + last_hidden_state: torch.FloatTensor | None = None + past_key_values: DynamicCache | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + return torch.nn.functional.rms_norm( + hidden_states, (self.hidden_size,), self.weight, self.variance_epsilon + ).to(input_dtype) + + +class NormAttentionNorm(nn.Module): + def __init__(self, config: E1Config, layer_idx: int) -> None: + super().__init__() + self.self_attn = Attention(config, layer_idx) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + within_seq_position_ids: torch.LongTensor, + global_position_ids: torch.LongTensor, + sequence_ids: torch.LongTensor, + attention_args: AttentionArgs | None = None, + past_key_value: DynamicCache | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + use_cache: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + DynamicCache | None, + list[torch.Tensor] | None, + ]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, self_attn_weights, present_key_value, s_max = self.self_attn( + hidden_states=hidden_states, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + attention_args=attention_args, + past_key_value=past_key_value, + output_attentions=output_attentions, + output_s_max=output_s_max, + use_cache=use_cache, + effective_backend=effective_backend, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + return hidden_states, residual, self_attn_weights, present_key_value, s_max + + +class DecoderLayer(nn.Module): + def __init__(self, config: E1Config, layer_idx: int) -> None: + super().__init__() + self.initializer_range = config.initializer_range + self.hidden_size = config.hidden_size + self.norm_attn_norm = NormAttentionNorm(config, layer_idx) + self.ffn = FFN(config) + + def forward( + self, + hidden_states: torch.Tensor, + within_seq_position_ids: torch.LongTensor, + global_position_ids: torch.LongTensor, + sequence_ids: torch.LongTensor, + attention_args: AttentionArgs | None = None, + past_key_value: DynamicCache | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + use_cache: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]: + hidden_states, residual, self_attn_weights, present_key_value, s_max = self.norm_attn_norm( + hidden_states=hidden_states, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + attention_args=attention_args, + past_key_value=past_key_value, + output_attentions=output_attentions, + output_s_max=output_s_max, + use_cache=use_cache, + effective_backend=effective_backend, + ) + + # Fully Connected + hidden_states = self.ffn(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states, self_attn_weights, present_key_value, s_max + + +class E1PreTrainedModel(FastPLMsAttentionMixin, PreTrainedModel): + config_class = E1Config + embedding_unsupported_pooling = ("cls", "parti") + config: E1Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["DecoderLayer"] + _transformer_layer_cls: ClassVar[list[type[nn.Module]]] = [DecoderLayer] + _skip_keys_device_placement = "past_key_values" + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ("sdpa", "flex_attention") + _is_internal_encoder = False + + def __init__(self, config: E1Config, *args: Any, **kwargs: Any) -> None: + super().__init__(config, *args, **kwargs) + # The E1 agreement requires this exact attribution when an E1 model is + # launched. Internal encoder construction is excluded so each public + # model launch displays the attribution exactly once. + if not self._is_internal_encoder: + print("Profluent-E1", file=sys.stderr, flush=True) + + @classmethod + def from_pretrained( # type: ignore[override] + cls, + pretrained_model_name_or_path: str | os.PathLike, + *model_args: Any, + **kwargs: Any, + ) -> E1PreTrainedModel: + tokenizer_token = None + if "token" in kwargs: + tokenizer_token = kwargs["token"] + elif "use_auth_token" in kwargs: + tokenizer_token = kwargs["use_auth_token"] + load_context_token = _TOKENIZER_LOAD_CONTEXT.set( + { + "tokenizer_source": pretrained_model_name_or_path, + "local_files_only": bool(kwargs.get("local_files_only", False)), + "cache_dir": kwargs.get("cache_dir"), + "revision": kwargs.get("revision"), + "token": tokenizer_token, + } + ) + try: + return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + finally: + _TOKENIZER_LOAD_CONTEXT.reset(load_context_token) + + @staticmethod + def _tokenizer_kwargs_from_config(config: E1Config) -> dict[str, Any]: + load_context = _TOKENIZER_LOAD_CONTEXT.get() + resolved_revision = getattr(config, "_commit_hash", None) + if not isinstance(resolved_revision, str) or not resolved_revision.strip(): + resolved_revision = None + if load_context is not None: + tokenizer_kwargs = dict(load_context) + if resolved_revision is not None: + tokenizer_kwargs["revision"] = resolved_revision + return tokenizer_kwargs + + tokenizer_source = None + if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0: + tokenizer_source = config._name_or_path + return { + "tokenizer_source": tokenizer_source, + "local_files_only": False, + "cache_dir": None, + "revision": resolved_revision, + "token": None, + } + + @property + def prep_tokens(self) -> E1BatchPreparer: + """Create E1's raw-sequence preparer only when a sequence API uses it.""" + + preparer = self.__dict__.get("_fastplms_prep_tokens") + if preparer is not None: + return preparer + encoder = self._modules.get("model") + if encoder is not None and encoder is not self: + return encoder.prep_tokens + tokenizer_kwargs = self.__dict__.get("_fastplms_tokenizer_kwargs") + if tokenizer_kwargs is None: + raise RuntimeError("E1 tokenizer settings were not initialized.") + preparer = E1BatchPreparer( + data_prep_config=DataPrepConfig( + max_num_sequences=self.config.max_num_sequences, + max_num_positions_within_seq=self.config.max_num_positions_within_seq, + ), + **tokenizer_kwargs, + ) + self.__dict__["_fastplms_prep_tokens"] = preparer + return preparer + + @prep_tokens.setter + def prep_tokens(self, value: E1BatchPreparer | None) -> None: + self.__dict__["_fastplms_prep_tokens"] = value + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, RMSNorm): + nn.init.ones_(module.weight) + return + if not isinstance(module, (nn.Linear, nn.Embedding)): + return + + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if isinstance(module, nn.Linear) and module.bias is not None: + nn.init.zeros_(module.bias) + if isinstance(module, nn.Embedding) and module.padding_idx is not None: + with torch.no_grad(): + module.weight[module.padding_idx].zero_() + + def _backward_compatibility_gradient_checkpointing(self) -> None: + if self.supports_gradient_checkpointing and getattr( + self.config, "gradient_checkpointing", False + ): + self.gradient_checkpointing_enable(dict(use_reentrant=False)) + + def post_init(self) -> None: + super().post_init() + + @property + def _device(self) -> torch.device: + return next(self.parameters()).device + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"E1 does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.config.attn_backend = backend + resolved = resolve_attention_backend(backend) + for module in self.modules(): + if isinstance(module, FAST_E1_ENCODER): + module._attn_backend = resolved + elif isinstance(module, Attention): + module.attn_backend = resolved + + +class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin): + config: E1Config + config_class = E1Config + _is_internal_encoder = True + + def __init__(self, config: E1Config, **kwargs) -> None: + E1PreTrainedModel.__init__(self, config, **kwargs) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.embed_seq_id = nn.Embedding(config.max_num_sequences, config.hidden_size) + self.layers = nn.ModuleList( + [DecoderLayer(config, i) for i in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.gradient_checkpointing = config.gradient_checkpointing + self.__dict__["_fastplms_tokenizer_kwargs"] = ( + E1PreTrainedModel._tokenizer_kwargs_from_config(config) + ) + self.__dict__["_fastplms_prep_tokens"] = None + self._attn_backend = resolve_attention_backend(config.attn_backend) + self.post_init() + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.embed_tokens = value + + def _embed( + self, + sequences: list[str], + return_attention_mask: bool = False, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **kwargs, + ) -> torch.Tensor: + batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device) + # The native preparer also returns training labels plus retrieval + # descriptors. The encoder accepts only its aligned model inputs. + encoder_batch: dict[str, torch.Tensor] = {} + for name in ( + "input_ids", + "within_seq_position_ids", + "global_position_ids", + "sequence_ids", + ): + value = batch[name] + if not isinstance(value, torch.Tensor): + raise TypeError(f"Prepared E1 field {name!r} must be a tensor.") + encoder_batch[name] = value + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.forward( + **encoder_batch, + output_hidden_states=output_hidden_states, + output_attentions=False, + return_dict=True, + ) + embeddings = select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + if return_attention_mask: + attention_mask = (encoder_batch["sequence_ids"] != -1).long() + return embeddings, attention_mask + else: + return embeddings + + def _prepare_hidden_states( + self, + input_ids: torch.LongTensor | None, + inputs_embeds: torch.FloatTensor | None, + within_seq_position_ids: torch.LongTensor | None, + global_position_ids: torch.LongTensor | None, + sequence_ids: torch.LongTensor | None, + ) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]: + if (input_ids is None) == (inputs_embeds is None): + message = ( + "Must specify either input_ids or inputs_embeds" + if input_ids is None + else "Cannot specify both input_ids and inputs_embeds" + ) + raise ValueError(message) + + source = input_ids if input_ids is not None else inputs_embeds + if source is None: + raise RuntimeError("E1 input validation did not resolve an input tensor.") + expected_rank = 2 if input_ids is not None else 3 + if source.ndim != expected_rank: + source_name = "input_ids" if input_ids is not None else "inputs_embeds" + raise ValueError( + f"{source_name} must have rank {expected_rank}; got shape {tuple(source.shape)}." + ) + batch_size, sequence_length = source.shape[:2] + if sequence_length == 0: + raise ValueError("E1 inputs must contain at least one token.") + if inputs_embeds is not None and inputs_embeds.shape[-1] != self.config.hidden_size: + raise ValueError( + "inputs_embeds hidden dimension must match config.hidden_size; " + f"got {inputs_embeds.shape[-1]} and {self.config.hidden_size}." + ) + if inputs_embeds is not None: + default_positions = torch.arange(sequence_length, device=source.device).expand( + batch_size, + -1, + ) + if within_seq_position_ids is None: + within_seq_position_ids = default_positions + if global_position_ids is None: + global_position_ids = default_positions + if sequence_ids is None: + sequence_ids = torch.zeros_like(default_positions) + + if within_seq_position_ids is None or global_position_ids is None or sequence_ids is None: + raise ValueError("Position and sequence IDs are required when input_ids are provided.") + expected_shape = (batch_size, sequence_length) + aligned_inputs = { + "within_seq_position_ids": within_seq_position_ids, + "global_position_ids": global_position_ids, + "sequence_ids": sequence_ids, + } + for name, value in aligned_inputs.items(): + if tuple(value.shape) != expected_shape: + raise ValueError( + f"{name} must have shape {expected_shape}; got {tuple(value.shape)}." + ) + within_positions = within_seq_position_ids.long() + global_positions = global_position_ids.long() + sequence_numbers = sequence_ids.long() + lowest_position, highest_position = torch.aminmax(within_positions) + if ( + lowest_position.item() < -1 + or highest_position.item() >= self.config.max_num_positions_within_seq + ): + raise ValueError( + "Position ids must be in the range " + f"[-1, {self.config.max_num_positions_within_seq}); got max " + f"{highest_position.item()} and min {lowest_position.item()}" + ) + lowest_global, highest_global = torch.aminmax(global_positions) + if ( + lowest_global.item() < -1 + or highest_global.item() >= self.config.max_num_positions_global + ): + raise ValueError( + "Global position ids must be in the range " + f"[-1, {self.config.max_num_positions_global}); got max " + f"{highest_global.item()} and min {lowest_global.item()}" + ) + lowest_sequence, highest_sequence = torch.aminmax(sequence_numbers) + if lowest_sequence.item() < -1 or highest_sequence.item() >= self.config.max_num_sequences: + raise ValueError( + "Sequence ids must be in the range " + f"[-1, {self.config.max_num_sequences}); got max " + f"{highest_sequence.item()} and min {lowest_sequence.item()}" + ) + + if inputs_embeds is None: + if input_ids is None: + raise RuntimeError("E1 input validation lost the token ID tensor.") + token_embeddings = self.embed_tokens(input_ids) + inputs_embeds = token_embeddings + self.embed_seq_id(sequence_numbers.clamp_min(0)) + layer_dtype = self.layers[0].norm_attn_norm.self_attn.q_proj.weight.dtype + target_dtype = ( + torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else layer_dtype + ) + return ( + inputs_embeds.to(target_dtype), + within_positions, + global_positions, + sequence_numbers, + ) + + def _resolve_forward_cache( + self, + past_key_values: DynamicCache | None, + use_cache: bool, + ) -> tuple[DynamicCache | None, bool]: + checkpointing = self.gradient_checkpointing and self.training and torch.is_grad_enabled() + if checkpointing and use_cache: + _get_logger().warning_once( + "`use_cache=True` is incompatible with gradient checkpointing; " + "setting `use_cache=False`." + ) + use_cache = False + if not use_cache: + return None, False + return past_key_values if past_key_values is not None else DynamicCache(), True + + def _build_forward_attention_args( + self, + sequence_ids: torch.LongTensor, + past_key_values: DynamicCache | None, + effective_backend: AttentionBackend, + ) -> AttentionArgs | None: + if past_key_values is not None and past_key_values.get_seq_length() != 0: + return None + + use_flex = effective_backend == AttentionBackend.FLEX + use_dense_mask = effective_backend in { + AttentionBackend.EAGER, + AttentionBackend.SDPA, + } + return AttentionArgs( + block_causal_block_mask=( + create_block_causal_mask_optimized(sequence_ids) + if use_flex and self.config.global_attention_every_n_layers > 0 + else None + ), + within_seq_block_mask=( + create_within_seq_block_mask(sequence_ids) if use_flex else None + ), + within_seq_mask_4d=(build_within_seq_mask_4d(sequence_ids) if use_dense_mask else None), + block_causal_mask_4d=( + build_block_causal_mask_4d(sequence_ids) if use_dense_mask else None + ), + ) + + def _run_decoder_layers( + self, + hidden_states: torch.Tensor, + within_seq_position_ids: torch.LongTensor, + global_position_ids: torch.LongTensor, + sequence_ids: torch.LongTensor, + attention_args: AttentionArgs | None, + past_key_values: DynamicCache | None, + use_cache: bool, + output_attentions: bool, + output_hidden_states: bool, + output_s_max: bool, + effective_backend: AttentionBackend, + ) -> E1ModelOutputWithPast: + hidden_history: list[torch.Tensor] | None = [] if output_hidden_states else None + attention_history: list[torch.Tensor] | None = [] if output_attentions else None + s_max_history: list[list[torch.Tensor]] | None = [] if output_s_max else None + next_cache: DynamicCache | None = None + + for layer in self.layers: + if hidden_history is not None: + hidden_history.append(hidden_states) + if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): + layer_output = self._gradient_checkpointing_func( + layer.__call__, + hidden_states, + within_seq_position_ids, + global_position_ids, + sequence_ids, + attention_args, + past_key_values, + output_attentions, + output_s_max, + use_cache, + effective_backend, + ) + else: + layer_output = layer( + hidden_states, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + attention_args=attention_args, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_s_max=output_s_max, + use_cache=use_cache, + effective_backend=effective_backend, + ) + hidden_states, attention, layer_cache, s_max = layer_output + if use_cache: + past_key_values = layer_cache + next_cache = layer_cache + if attention_history is not None: + if attention is None: + raise RuntimeError( + "An E1 layer did not return attention tensors when requested." + ) + attention_history.append(attention) + if s_max_history is not None: + if s_max is None: + raise RuntimeError( + "An E1 layer did not return s_max diagnostics when requested." + ) + s_max_history.append(s_max) + + hidden_states = self.norm(hidden_states) + if hidden_history is not None: + hidden_history.append(hidden_states) + return E1ModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=tuple(hidden_history) if hidden_history is not None else None, + attentions=tuple(attention_history) if attention_history is not None else None, + s_max=tuple(s_max_history) if s_max_history is not None else None, + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + within_seq_position_ids: torch.LongTensor | None = None, + global_position_ids: torch.LongTensor | None = None, + sequence_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: DynamicCache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool = False, + return_dict: bool | None = None, + ) -> E1ModelOutputWithPast | tuple[Any, ...]: + """Transform token or soft embeddings H with shape (b, l, d).""" + + use_cache = ( + use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + hidden_states, within_positions, global_positions, sequence_numbers = ( + self._prepare_hidden_states( + input_ids, + inputs_embeds, + within_seq_position_ids, + global_position_ids, + sequence_ids, + ) + ) + cache, use_cache = self._resolve_forward_cache(past_key_values, use_cache) + effective_backend = resolve_attention_backend_for_call( + self._attn_backend, + output_attentions=bool(output_attentions), + ) + attention_args = self._build_forward_attention_args( + sequence_numbers, + cache, + effective_backend, + ) + result = self._run_decoder_layers( + hidden_states, + within_positions, + global_positions, + sequence_numbers, + attention_args, + cache, + use_cache, + output_attentions, + output_hidden_states, + output_s_max, + effective_backend, + ) + if not return_dict: + return result.to_tuple() + return result + + +class E1Model(E1PreTrainedModel, EmbeddingMixin): + config: E1Config + config_class = E1Config + + def __init__(self, config: E1Config, **kwargs) -> None: + E1PreTrainedModel.__init__(self, config, **kwargs) + self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) + self.post_init() + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.model.set_input_embeddings(value) + + def _embed( + self, sequences: list[str], return_attention_mask: bool = False, **kwargs + ) -> torch.Tensor: + return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + within_seq_position_ids: torch.LongTensor | None = None, + global_position_ids: torch.LongTensor | None = None, + sequence_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: DynamicCache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool = False, + return_dict: bool | None = None, + ) -> E1ModelOutputWithPast | tuple[Any, ...]: + return self.model( + input_ids=input_ids, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=return_dict, + ) + + +class E1ForMaskedLM(FastPLMTestTimeTrainingMixin, E1PreTrainedModel, EmbeddingMixin): + config: E1Config + config_class = E1Config + + def __init__(self, config: E1Config, **kwargs) -> None: + E1PreTrainedModel.__init__(self, config, **kwargs) + self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) + self.vocab_size = config.vocab_size + self.mlm_head = torch.nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size, bias=True), + nn.GELU(), + nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps), + nn.Linear(config.hidden_size, config.vocab_size, bias=True), + ) + self.gradient_checkpointing = config.gradient_checkpointing + self.post_init() + self.init_ttt({"lora_target_replace_module": "Attention"}) + + @property + def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: + return self.model.device_mesh + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.model.set_input_embeddings(value) + + def _embed( + self, sequences: list[str], return_attention_mask: bool = False, **kwargs + ) -> torch.Tensor: + return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) + + def get_output_embeddings(self) -> nn.Linear: + return self.mlm_head[-1] + + def set_output_embeddings(self, value: nn.Linear) -> None: + self.mlm_head[-1] = value + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.model] + + def _ttt_tokenize( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs, + ) -> dict[str, torch.Tensor]: + if input_ids is not None: + return { + "input_ids": input_ids, + "within_seq_position_ids": kwargs["within_seq_position_ids"], + "global_position_ids": kwargs["global_position_ids"], + "sequence_ids": kwargs["sequence_ids"], + } + if seq is None: + raise ValueError("Pass either seq or E1 token tensors for TTT.") + sequences = [seq] if isinstance(seq, str) else seq + batch = self.prep_tokens.get_batch_kwargs(sequences, device=torch.device("cpu")) + return { + "input_ids": batch["input_ids"], + "within_seq_position_ids": batch["within_seq_position_ids"], + "global_position_ids": batch["global_position_ids"], + "sequence_ids": batch["sequence_ids"], + } + + def _ttt_mask_token(self) -> int: + return int(self.prep_tokens.mask_token_id) + + def _ttt_padding_token(self) -> int: + return int(self.prep_tokens.pad_token_id) + + def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + amino_acids = "ACDEFGHIKLMNPQRSTVWY" + ids = [self.prep_tokens.vocab[aa] for aa in amino_acids] + return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) + + def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: + return ~self.prep_tokens.get_boundary_token_mask(input_ids) + + def _ttt_predict_logits( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + **kwargs, + ) -> torch.Tensor: + del kwargs + if not isinstance(batch, dict): + raise TypeError("E1 TTT expects a tensor dictionary.") + output = self( + input_ids=batch["input_ids"], + within_seq_position_ids=batch["within_seq_position_ids"], + global_position_ids=batch["global_position_ids"], + sequence_ids=batch["sequence_ids"], + return_dict=True, + ) + return output.logits + + def search_homologues( + self, + sequence: str, + output_dir: str, + provider: str = "colabfold", + target_db: str | None = None, + seq_id: str | None = None, + **kwargs, + ) -> str: + searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) + return searcher.search(sequence=sequence, output_dir=output_dir, seq_id=seq_id) + + def batch_search_homologues( + self, + sequences: list[str], + output_dir: str, + provider: str = "colabfold", + target_db: str | None = None, + seq_ids: list[str] | None = None, + continue_on_error: bool = True, + **kwargs, + ) -> dict[str, str]: + searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) + return searcher.batch_search( + sequences=sequences, + output_dir=output_dir, + seq_ids=seq_ids, + continue_on_error=continue_on_error, + ) + + def sample_msa_contexts( + self, + a3m_path: str, + seed: int = 42, + max_context_tokens: list[int] | None = None, + similarity_thresholds: list[float] | None = None, + min_query_similarity: float = 0.3, + context_cache_dir: str | None = None, + ) -> dict[str, str]: + context_specs = build_context_specifications( + max_context_tokens=max_context_tokens, + similarity_thresholds=similarity_thresholds, + min_query_similarity=min_query_similarity, + ) + cache = None + if context_cache_dir is not None: + key = repr((max_context_tokens, similarity_thresholds, min_query_similarity)) + specs_hash = hashlib.md5(key.encode()).hexdigest()[:8] + cache = ContextCache(context_cache_dir, specs_hash, seed) + cached = cache.load(a3m_path) + if cached is not None: + return cached + contexts = sample_contexts_for_msa(a3m_path, context_specs, seed=seed) + if cache is not None: + cache.store(a3m_path, contexts) + return contexts + + @torch.inference_mode() + def score_ppll( + self, + sequences: list[str], + a3m_path: str, + ensemble: bool = True, + seed: int = 42, + max_context_tokens: list[int] | None = None, + similarity_thresholds: list[float] | None = None, + min_query_similarity: float = 0.3, + max_batch_tokens: int = 131072, + cache_size: int = 1, + context_cache_dir: str | None = None, + progress: bool = True, + ) -> list[float] | list[list[float]]: + """Score sequences with FastPLMs PPLL reduction over sampled E1 MSA contexts. + + This intentionally differs from Profluent's official E1Scorer, which scores + mutants against a parent sequence with wildtype or masked marginal log-prob + deltas. Here each sequence is scored by mean correct-token probability and + optionally averaged across sampled contexts. + """ + contexts = self.sample_msa_contexts( + a3m_path=a3m_path, + seed=seed, + max_context_tokens=max_context_tokens, + similarity_thresholds=similarity_thresholds, + min_query_similarity=min_query_similarity, + context_cache_dir=context_cache_dir, + ) + if not contexts: + raise ValueError("At least one sampled MSA context is required for PPLL scoring.") + + predictor = _E1ContextPredictor( + model=self, + data_prep_config=DataPrepConfig(remove_X_tokens=True), + max_batch_tokens=max_batch_tokens, + fields_to_save=["logits"], + save_masked_positions_only=False, + keep_predictions_in_gpu=False, + use_cache=True, + cache_size=cache_size, + progress=progress, + ) + vocab = predictor.batch_preparer.vocab + seq_token_ids = [ + torch.tensor([vocab[aa] for aa in seq if aa != "X"], device=self.device) + for seq in sequences + ] + context_ids = list(contexts.keys()) + all_scores = torch.zeros(len(sequences), len(context_ids), device=self.device) + + iterator = tqdm(context_ids, desc="Scoring with contexts", disable=not progress) + for ctx_idx, ctx_id in enumerate(iterator): + predictions = list( + predictor.predict( + sequences=sequences, + sequence_ids=list(range(len(sequences))), + context_seqs={ctx_id: contexts[ctx_id]}, + ) + ) + for prediction in predictions: + seq_idx = prediction["id"] + if not isinstance(seq_idx, int): + raise TypeError("Expected integer sequence ids for score aggregation.") + all_scores[seq_idx, ctx_idx] = compute_ppll( + prediction["logits"], seq_token_ids[seq_idx] + ) + if predictor.kv_cache is not None: + predictor.kv_cache.reset() + + if ensemble: + return all_scores.mean(dim=1).tolist() + return all_scores.tolist() + + @torch.inference_mode() + def embed_with_msa( + self, + sequences: list[str], + a3m_path: str | None = None, + context: str | None = None, + pooling_types: list[str] | None = None, + pooling: str = "mean", + matrix_embed: bool = False, + seed: int = 42, + max_batch_tokens: int = 131072, + embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, + embed_similarity: float = DEFAULT_EMBED_SIMILARITY, + min_query_similarity: float = 0.3, + progress: bool = True, + ) -> torch.Tensor | list[torch.Tensor]: + if a3m_path is not None and context is None: + spec = ContextSpecification( + max_num_samples=511, + max_token_length=embed_max_tokens, + max_query_similarity=embed_similarity, + min_query_similarity=min_query_similarity, + ) + contexts, _ = sample_multiple_contexts( + msa_path=a3m_path, + context_specifications=[spec], + seed=seed, + ) + context = contexts[0] if contexts else None + + hidden_list = _forward_for_embedding( + model=self, + sequences=sequences, + context=context, + max_batch_tokens=max_batch_tokens, + progress=progress, + ) + if matrix_embed: + return hidden_list + if pooling_types is not None: + return _pool_hidden_states(hidden_list, pooling_types, self.device) + if pooling not in ("mean", "cls"): + raise ValueError("pooling must be 'mean' or 'cls' when pooling_types is not provided") + embeddings = [ + hidden.mean(dim=0) if pooling == "mean" else hidden[0] for hidden in hidden_list + ] + return torch.stack(embeddings) + + @torch.inference_mode() + def embed_dataset_with_msa( + self, + sequences: list[str], + msa_lookup: dict[str, str] | None = None, + msa_dir: str | None = None, + msa_hf_path: str | None = None, + batch_size: int = 2, + max_len: int = 2048, + pooling_types: list[str] | None = None, + pooling: str = "mean", + matrix_embed: bool = False, + embed_dtype: torch.dtype = torch.bfloat16, + embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, + embed_similarity: float = DEFAULT_EMBED_SIMILARITY, + min_query_similarity: float = 0.3, + seed: int = 42, + progress: bool = True, + max_batch_tokens: int = 131072, + batch_window_size: int | None = None, + max_tokens_per_batch: int | None = None, + output: str | os.PathLike[str] | None = None, + format: str = "safetensors", + resume: bool = True, + shard_size: int = 2 * 1024**3, + model_state_fingerprint: str | None = None, + ) -> EmbeddingResult: + """Embed an ordered sequence dataset with optional sampled MSA context. + + Unlike the legacy dictionary return, the result preserves duplicate + sequences and input order. ``output`` uses the same transactional, + resumable SQLite or safetensors persistence as :meth:`embed_dataset`. + ``max_len`` counts biological residues. + """ + + if not sequences: + raise ValueError("sequences must contain at least one protein sequence.") + if any(not isinstance(sequence, str) or not sequence for sequence in sequences): + raise ValueError("sequences must contain non-empty strings.") + if max_len <= 0: + raise ValueError("max_len must be positive.") + if msa_lookup is None: + if msa_dir is not None: + msa_lookup = load_msa_dir(msa_dir) + elif msa_hf_path is not None: + msa_lookup = load_msa_from_hf(msa_hf_path) + else: + msa_lookup = {} + + truncated_sequences = [sequence[:max_len] for sequence in sequences] + unique_seqs = sorted(set(truncated_sequences), key=lambda value: (-len(value), value)) + context_map: dict[str, str | None] = {} + spec = ContextSpecification( + max_num_samples=511, + max_token_length=embed_max_tokens, + max_query_similarity=embed_similarity, + min_query_similarity=min_query_similarity, + ) + for seq in unique_seqs: + a3m_path = get_msa_for_sequence(seq, msa_lookup) + if a3m_path is None: + context_map[seq] = None + continue + contexts, _ = sample_multiple_contexts( + msa_path=a3m_path, + context_specifications=[spec], + seed=seed, + ) + context_map[seq] = contexts[0] if contexts else None + + context_digest = hashlib.sha256() + for sequence in unique_seqs: + for value in (sequence, context_map[sequence] or ""): + encoded = value.encode("utf-8") + context_digest.update(len(encoded).to_bytes(8, "big")) + context_digest.update(encoded) + + def embed_msa_batch(batch_sequences: list[str]) -> EmbeddingBatch: + grouped_positions: dict[str | None, list[int]] = defaultdict(list) + for position, sequence in enumerate(batch_sequences): + grouped_positions[context_map[sequence]].append(position) + hidden_by_position: list[torch.Tensor | None] = [None] * len(batch_sequences) + for context, positions in grouped_positions.items(): + context_sequences = [batch_sequences[position] for position in positions] + hidden_states = _forward_for_embedding( + model=self, + sequences=context_sequences, + context=context, + max_batch_tokens=max_batch_tokens, + progress=progress, + ) + for position, hidden in zip(positions, hidden_states, strict=True): + hidden_by_position[position] = hidden + resolved = [hidden for hidden in hidden_by_position if hidden is not None] + if len(resolved) != len(batch_sequences): + raise RuntimeError("E1 MSA embedding did not return every requested sequence.") + max_residues = max(hidden.shape[0] for hidden in resolved) + hidden_size = resolved[0].shape[-1] + X = resolved[0].new_zeros((len(resolved), max_residues, hidden_size)) + residue_mask = torch.zeros( + (len(resolved), max_residues), + dtype=torch.bool, + device=X.device, + ) + for position, hidden in enumerate(resolved): + residue_count = hidden.shape[0] + X[position, :residue_count] = hidden + residue_mask[position, :residue_count] = True + return EmbeddingBatch(X=X, residue_mask=residue_mask) + + resolved_pooling: str | list[str] | None = ( + None if matrix_embed else pooling_types if pooling_types is not None else pooling + ) + adapter_identity = { + "kind": "e1-msa-v1", + "sampling_source_revision": E1_MSA_SAMPLING_SOURCE_REVISION, + "context_sha256": context_digest.hexdigest(), + "context_count": sum(context is not None for context in context_map.values()), + "seed": seed, + "embed_max_tokens": embed_max_tokens, + "embed_similarity": embed_similarity, + "min_query_similarity": min_query_similarity, + "max_batch_tokens": max_batch_tokens, + } + return embed_dataset( + self, + [(str(position), sequence) for position, sequence in enumerate(sequences)], + batch_size=batch_size, + pooling=resolved_pooling, + full_embeddings=matrix_embed, + output=output, + format=format, + resume=resume, + max_length=max_len, + truncate=True, + dtype=embed_dtype, + shard_size=shard_size, + model_state_fingerprint=model_state_fingerprint, + batch_window_size=batch_window_size, + max_tokens_per_batch=max_tokens_per_batch, + _embedding_batch_fn=embed_msa_batch, + _embedding_batch_identity=adapter_identity, + _allowed_unsupported_pooling=("cls",), + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + within_seq_position_ids: torch.LongTensor | None = None, + global_position_ids: torch.LongTensor | None = None, + sequence_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + past_key_values: DynamicCache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool = False, + return_dict: bool | None = None, + ) -> E1MaskedLMOutputWithPast | tuple[Any, ...]: + """Return hidden states and masked-token logits for E1 inputs. + + Token, position, sequence, and label tensors have shape (b, l). + Callers may instead provide precomputed H with shape (b, l, d). + """ + use_cache = ( + use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs: E1ModelOutputWithPast = self.model( + input_ids=input_ids, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + + last_hidden_state = outputs.last_hidden_state + loss = None + + mlm_logits = self.mlm_head(last_hidden_state).float() + mlm_loss = None + if labels is not None: + mlm_logits_flat = mlm_logits.contiguous().view(-1, self.config.vocab_size) + mlm_labels_flat = labels.to(mlm_logits_flat.device).contiguous().view(-1) + mlm_loss = F.cross_entropy( + mlm_logits_flat, + mlm_labels_flat, + ignore_index=-100, + reduction="none", + ) + mask = mlm_labels_flat.ne(-100) & mlm_labels_flat.ne(self.model.padding_idx) + n_mlm = mask.sum().clamp_min(1) + mlm_loss = (mlm_loss * mask.to(mlm_loss)).sum() / n_mlm + loss = 0.0 + loss += mlm_loss + + result = E1MaskedLMOutputWithPast( + loss=loss, + logits=mlm_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + mlm_loss=mlm_loss, + last_hidden_state=last_hidden_state, + past_key_values=outputs.past_key_values, + s_max=outputs.s_max, + ) + if not return_dict: + return result.to_tuple() + return result + + +class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin): + config: E1Config + config_class = E1Config + + def __init__(self, config: E1Config, **kwargs) -> None: + pooling_types = kwargs.pop("pooling_types", None) + if pooling_types is None: + pooling_types = ["mean", "var"] + elif not isinstance(pooling_types, list): + raise TypeError("pooling_types must be a non-empty list of pooling names") + elif not pooling_types or not all(isinstance(name, str) for name in pooling_types): + raise ValueError("pooling_types must be a non-empty list of pooling names") + + E1PreTrainedModel.__init__(self, config, **kwargs) + self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) + self.vocab_size = config.vocab_size + self.num_labels = config.num_labels + self.pooler = Pooler(pooling_types) + self.classifier = nn.Sequential( + nn.Linear(config.hidden_size * len(pooling_types), config.hidden_size * 4), + nn.GELU(), + nn.LayerNorm(config.hidden_size * 4), + nn.Linear(config.hidden_size * 4, config.num_labels), + ) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.gradient_checkpointing = config.gradient_checkpointing + self.post_init() + + @property + def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: + return self.model.device_mesh + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.model.set_input_embeddings(value) + + def _embed( + self, sequences: list[str], return_attention_mask: bool = False, **kwargs + ) -> torch.Tensor: + return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + within_seq_position_ids: torch.LongTensor | None = None, + global_position_ids: torch.LongTensor | None = None, + sequence_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + past_key_values: DynamicCache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool = False, + return_dict: bool | None = None, + ) -> E1ClassificationOutputWithPast | tuple[Any, ...]: + use_cache = ( + use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs: E1ModelOutputWithPast = self.model( + input_ids=input_ids, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + + attention_mask = ( + (sequence_ids != -1).long() + if sequence_ids is not None + else torch.ones( + outputs.last_hidden_state.shape[:2], + device=outputs.last_hidden_state.device, + dtype=torch.long, + ) + ) + x = outputs.last_hidden_state + features = self.pooler(x, attention_mask) + logits = self.classifier(features) + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.flatten(), labels.flatten()) + else: + loss = self.mse(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + result = E1ClassificationOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + last_hidden_state=x, + s_max=outputs.s_max, + ) + if not return_dict: + return result.to_tuple() + return result + + +class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin): + config: E1Config + config_class = E1Config + + def __init__(self, config: E1Config, **kwargs) -> None: + E1PreTrainedModel.__init__(self, config, **kwargs) + self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) + self.vocab_size = config.vocab_size + self.num_labels = config.num_labels + self.classifier = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size * 4), + nn.GELU(), + nn.LayerNorm(config.hidden_size * 4), + nn.Linear(config.hidden_size * 4, config.num_labels), + ) + self.loss_fct = nn.CrossEntropyLoss() + self.gradient_checkpointing = config.gradient_checkpointing + self.post_init() + + @property + def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: + return self.model.device_mesh + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.model.set_input_embeddings(value) + + def _embed( + self, sequences: list[str], return_attention_mask: bool = False, **kwargs + ) -> torch.Tensor: + return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + within_seq_position_ids: torch.LongTensor | None = None, + global_position_ids: torch.LongTensor | None = None, + sequence_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + past_key_values: DynamicCache | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool = False, + return_dict: bool | None = None, + ) -> E1TokenClassificationOutputWithPast | tuple[Any, ...]: + use_cache = ( + use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) + ) + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs: E1ModelOutputWithPast = self.model( + input_ids=input_ids, + within_seq_position_ids=within_seq_position_ids, + global_position_ids=global_position_ids, + sequence_ids=sequence_ids, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + + x = outputs.last_hidden_state + logits = self.classifier(x) + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + result = E1TokenClassificationOutputWithPast( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + last_hidden_state=x, + past_key_values=outputs.past_key_values, + s_max=outputs.s_max, + ) + if not return_dict: + return result.to_tuple() + return result diff --git a/src/fastplms/models/e1/preparation.py b/src/fastplms/models/e1/preparation.py new file mode 100644 index 0000000..d9aff31 --- /dev/null +++ b/src/fastplms/models/e1/preparation.py @@ -0,0 +1,275 @@ +"""Tokenizer loading and raw-sequence batch preparation for E1.""" + +from __future__ import annotations + +import itertools +import os +import torch +from dataclasses import dataclass +from tokenizers import Tokenizer +from torch.nn.utils.rnn import pad_sequence + + +PAD_TOKEN_ID = 0 +BOS_TOKEN_ID = 1 +EOS_TOKEN_ID = 2 +E1_VOCAB_SIZE = 34 +E1_TOKENIZER_REPO_ID = "Synthyra/Profluent-E1-150M" + + +def _load_tokenizer_file(fname: str) -> Tokenizer: + tokenizer: Tokenizer = Tokenizer.from_file(fname) + padding = tokenizer.padding + actual_pad_id = None if padding is None else padding.get("pad_id") + if actual_pad_id != PAD_TOKEN_ID: + raise ValueError( + f"Padding token id must be {PAD_TOKEN_ID}, but got {actual_pad_id}" + ) + return tokenizer + + +def get_tokenizer( + pretrained_model_name_or_path: str | os.PathLike | None = None, + *, + local_files_only: bool = False, + cache_dir: str | os.PathLike | None = None, + revision: str | None = None, + token: str | bool | None = None, +) -> Tokenizer: + source_path = None + checked_local_source = False + if pretrained_model_name_or_path is not None: + source_path = os.fspath(pretrained_model_name_or_path) + if os.path.isdir(source_path): + checked_local_source = True + fname = os.path.join(source_path, "tokenizer.json") + if os.path.isfile(fname): + return _load_tokenizer_file(fname) + + fname = os.path.join(os.path.dirname(__file__), "tokenizer.json") + if os.path.isfile(fname): + return _load_tokenizer_file(fname) + + if local_files_only and checked_local_source: + raise FileNotFoundError( + f"E1 tokenizer.json was not found in {source_path} or next to {__file__}." + ) + + from huggingface_hub import hf_hub_download + + repo_id = E1_TOKENIZER_REPO_ID + if source_path is not None and not checked_local_source: + repo_id = source_path + try: + fname = hf_hub_download( + repo_id=repo_id, + filename="tokenizer.json", + cache_dir=os.fspath(cache_dir) if cache_dir is not None else None, + revision=revision, + token=token, + local_files_only=local_files_only, + ) + except Exception as error: + raise FileNotFoundError( + f"E1 tokenizer.json was not found locally and could not be loaded from {repo_id}." + ) from error + return _load_tokenizer_file(fname) + + +@dataclass +class DataPrepConfig: + max_num_sequences: int = 512 + max_num_positions_within_seq: int = 8192 + remove_X_tokens: bool = False + + +def get_context(sequence: str) -> str | None: + if "," in sequence: + return sequence.rsplit(",", 1)[0] + return None + + +class E1BatchPreparer: + def __init__( + self, + data_prep_config: DataPrepConfig | None = None, + tokenizer: Tokenizer | None = None, + tokenizer_source: str | os.PathLike | None = None, + local_files_only: bool = False, + cache_dir: str | os.PathLike | None = None, + revision: str | None = None, + token: str | bool | None = None, + preserve_context_labels: bool = False, + ) -> None: + self.tokenizer = tokenizer or get_tokenizer( + tokenizer_source, + local_files_only=local_files_only, + cache_dir=cache_dir, + revision=revision, + token=token, + ) + self.data_prep_config = data_prep_config or DataPrepConfig() + self.pad_token_id = self.tokenizer.token_to_id("") + self.preserve_context_labels = preserve_context_labels + self.boundary_token_ids = torch.tensor( # (5,) + [self.tokenizer.token_to_id(token) for token in ["", "", "1", "2", ""]] + ).long() + self.mask_token = "?" # nosec + self.mask_token_id = self.tokenizer.token_to_id(self.mask_token) + self.X_token_id = self.tokenizer.token_to_id("X") + self.vocab = self.tokenizer.get_vocab() + + def get_batch_kwargs( # type: ignore[override] + self, + sequences: list[str], + device: torch.device | None = None, + non_blocking: bool = False, + ) -> dict[str, torch.Tensor | list[str] | list[int]]: + device = torch.device("cpu") if device is None else device + sequence_encodings = [self.prepare_multiseq(sequence) for sequence in sequences] + return self.pad_encodings(sequence_encodings, device, non_blocking) + + def pad_encodings( + self, + sequence_encodings: list[dict[str, torch.Tensor]], + device: torch.device | None = None, + non_blocking: bool = False, + ) -> dict[str, torch.Tensor | list[str] | list[int]]: + # Each sequence encoding contains aligned one-dimensional tensors of length l_i. + device = torch.device("cpu") if device is None else device + non_blocking = non_blocking and device.type == "cuda" + padded_encodings = {} + # Sequence and position ID zero is valid, so -1 unambiguously marks padding. + for key, padding_value in { + "input_ids": self.pad_token_id, + "sequence_ids": -1, + "within_seq_position_ids": -1, + "global_position_ids": -1, + "labels": self.pad_token_id, + }.items(): + padded_encodings[key] = pad_sequence( # (b, l_max) + [enc[key] for enc in sequence_encodings], + batch_first=True, + padding_value=padding_value, + ).to(device=device, dtype=torch.long, non_blocking=non_blocking) + + padded_encodings["context"] = [enc["context"] for enc in sequence_encodings] + padded_encodings["context_len"] = [enc["context_len"] for enc in sequence_encodings] + + return padded_encodings + + def prepare_multiseq(self, sequence: str) -> dict[str, torch.Tensor | str | int]: + sequences = sequence.split(",") + if len(sequences) > self.data_prep_config.max_num_sequences: + raise ValueError( + f"Number of sequences {len(sequences)} exceeds max number of sequences " + f"{self.data_prep_config.max_num_sequences} in the provided multi-sequence " + "instance. Please remove some homologous sequences before trying again." + ) + + encodings = tuple(self.prepare_singleseq(item) for item in sequences) + token_counts = torch.tensor( # (n,) + [encoding["input_ids"].numel() for encoding in encodings], + dtype=torch.long, + ) + input_ids = torch.cat(tuple(encoding["input_ids"] for encoding in encodings)) # (t,) + labels = torch.cat(tuple(encoding["labels"] for encoding in encodings)) # (t,) + positions = tuple(encoding["position_ids"] for encoding in encodings) + within_seq_position_ids = torch.cat(positions) # (t,) + + # Offsets preserve gaps left by optional X-token removal. + position_spans = torch.tensor( # (n,) + [int(position_ids[-1].item()) + 1 for position_ids in positions], + dtype=torch.long, + ) + position_offsets = torch.cat( # (n,) + (torch.zeros(1, dtype=torch.long), position_spans[:-1]), + ).cumsum(dim=0) + global_position_ids = torch.cat( # (t,) + tuple( + position_ids + offset + for position_ids, offset in zip(positions, position_offsets, strict=True) + ) + ) + sequence_ids = torch.arange( # (t,) + len(encodings), + dtype=torch.long, + ).repeat_interleave( + token_counts + ) + + context_len = int(token_counts[:-1].sum().item()) + context = self.tokenizer.decode(input_ids[:context_len].tolist(), skip_special_tokens=False) + if not self.preserve_context_labels: + labels[:context_len] = self.pad_token_id + + aligned_tensors = ( + sequence_ids, + within_seq_position_ids, + global_position_ids, + labels, + ) + if any(tensor.shape != input_ids.shape for tensor in aligned_tensors): + raise AssertionError( + "Input ids, sequence ids, within seq position ids, global position ids, " + "and labels must have the same shape" + ) + if input_ids.numel() < context_len: + raise AssertionError( + "Input ids must have at least as many tokens as the context length" + ) + + return { + "input_ids": input_ids, + "sequence_ids": sequence_ids, + "within_seq_position_ids": within_seq_position_ids, + "global_position_ids": global_position_ids, + "labels": labels, + "context": context, + "context_len": context_len, + } + + def prepare_singleseq(self, sequence: str) -> dict[str, torch.Tensor]: + if not self.validate_sequence(sequence): + raise ValueError( + f"Invalid sequence: {sequence}; Input sequence should contain " + "[A-Z] or ? characters only" + ) + + if len(sequence) > self.data_prep_config.max_num_positions_within_seq: + raise ValueError( + f"Sequence length {len(sequence)} exceeds max length " + f"{self.data_prep_config.max_num_positions_within_seq}" + ) + + symbols = itertools.chain(("", "1"), sequence, ("2", "")) + tokens = torch.tensor( # (l + 4,) + [self.vocab[symbol] for symbol in symbols], + dtype=torch.long, + ) + position_ids = torch.arange(tokens.numel(), dtype=torch.long) # (l + 4,) + + if self.data_prep_config.remove_X_tokens: + keep = tokens.ne(self.X_token_id) # (l + 4,) + tokens = tokens[keep] # (t,) + position_ids = position_ids[keep] # (t,) + + return { # Each tensor: (t,) + "input_ids": tokens, + "labels": tokens, + "position_ids": position_ids, + } + + def get_boundary_token_mask(self, tokens: torch.Tensor) -> torch.BoolTensor: + # tokens: (...) + return torch.isin(tokens, self.boundary_token_ids.to(tokens.device)) # (...) + + def get_mask_positions_mask(self, tokens: torch.Tensor) -> torch.BoolTensor: + # tokens: (...) + return tokens == self.mask_token_id # (...) + + def validate_sequence(self, sequence: str) -> bool: + if not isinstance(sequence, str): + raise TypeError("Sequence must be a string.") + sequence = sequence.replace(self.mask_token, "") + return sequence.isalpha() and sequence.isupper() diff --git a/src/fastplms/models/e1/retrieval.py b/src/fastplms/models/e1/retrieval.py new file mode 100644 index 0000000..5423822 --- /dev/null +++ b/src/fastplms/models/e1/retrieval.py @@ -0,0 +1,1804 @@ +"""FASTA, MSA, context sampling, and homologue-search utilities for E1.""" + +from __future__ import annotations + +import hashlib +import itertools +import json +import math +import numbers +import os +import platform +import random +import re +import shutil +import subprocess +import tarfile +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +import numpy as np +import torch +from collections import defaultdict, namedtuple +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from email.utils import parsedate_to_datetime +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING, Any, TypedDict +from tqdm.auto import tqdm +from transformers import PreTrainedModel +from transformers.utils import logging + +from fastplms.embeddings import Pooler +from .cache import KVCache +from .preparation import DataPrepConfig, E1BatchPreparer, get_context + + +if TYPE_CHECKING: + from .modeling_e1 import E1MaskedLMOutputWithPast + + +def _get_logger(): + """Resolve the Transformers logger only when a retrieval path emits a message.""" + + return logging.get_logger(__name__) + + +MMSEQS2_IMAGE_REPOSITORY = "ghcr.io/soedinglab/mmseqs2" +MMSEQS2_VERSION = "18-8cc5c" +MMSEQS2_CPU_MANIFEST_DIGEST = ( + "sha256:41b12b0d5f41432fa1b9976123da6e2e06e7fab49a34964f3b54ec038e5845d9" +) +MMSEQS2_CPU_ARM64_CHILD_DIGEST = ( + "sha256:8bec048845f8f20749c2e2ad067a27d67eef839d2bb068e9d6e957113e9a7fba" +) +DOCKER_IMAGE = ( + f"{MMSEQS2_IMAGE_REPOSITORY}:{MMSEQS2_VERSION}@{MMSEQS2_CPU_MANIFEST_DIGEST}" +) +DEFAULT_MMSEQS2_PHASE_TIMEOUT = 1800.0 +COLABFOLD_HOST = "https://api.colabfold.com" +LOWERCASE_CHARS = b"abcdefghijklmnopqrstuvwxyz" +DEFAULT_MAX_CONTEXT_TOKENS = [6144, 12288, 24576] +DEFAULT_SIMILARITY_THRESHOLDS = [1.0, 0.95, 0.9, 0.7, 0.5] +DEFAULT_EMBED_MAX_TOKENS = 8192 +DEFAULT_EMBED_SIMILARITY = 0.95 +E1_MSA_SAMPLING_SOURCE_REVISION = "bfd2620a602248499f3d2583d85a7ecddf0b6e02" + +IdSequence = namedtuple("IdSequence", ["id", "sequence"]) +IndexedSequence = tuple[int, str] + +_SHA256_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$") +_IMAGE_VERSION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +_SAFE_SEQUENCE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +_SAFE_QUERY_SEQUENCE_RE = re.compile(r"^[A-Za-z*.-]+$") + + +@dataclass(frozen=True, slots=True) +class _PinnedImageReference: + repository: str + version: str + digest: str + + +@dataclass(frozen=True, slots=True) +class _DockerImageIdentity: + reference: str + repository: str + version: str + manifest_digest: str + image_id: str + os: str + architecture: str + + def to_dict(self) -> dict[str, str]: + return { + "reference": self.reference, + "repository": self.repository, + "version": self.version, + "manifest_digest": self.manifest_digest, + "image_id": self.image_id, + "os": self.os, + "architecture": self.architecture, + } + + +def _parse_pinned_image_reference(reference: str) -> _PinnedImageReference: + """Parse ``repository:version@sha256:digest`` and reject mutable images.""" + + if not isinstance(reference, str) or not reference or any(char.isspace() for char in reference): + raise ValueError( + "docker_image must be an immutable repository:version@sha256:digest reference" + ) + try: + name_and_version, digest = reference.rsplit("@", maxsplit=1) + except ValueError as error: + raise ValueError( + "docker_image must include an immutable @sha256 digest; mutable tags are rejected" + ) from error + last_slash = name_and_version.rfind("/") + last_colon = name_and_version.rfind(":") + if last_colon <= last_slash: + raise ValueError("docker_image must include an explicit version tag before its digest") + repository = name_and_version[:last_colon] + version = name_and_version[last_colon + 1 :] + if ( + not repository + or repository.endswith("/") + or "@" in repository + or _IMAGE_VERSION_RE.fullmatch(version) is None + ): + raise ValueError("docker_image contains an invalid repository or version tag") + if _SHA256_DIGEST_RE.fullmatch(digest) is None: + raise ValueError("docker_image must include a lowercase sha256 digest") + return _PinnedImageReference(repository=repository, version=version, digest=digest) + + +def _docker_architecture() -> str: + machine = platform.machine().lower() + aliases = { + "aarch64": "arm64", + "arm64": "arm64", + "amd64": "amd64", + "x86_64": "amd64", + } + try: + return aliases[machine] + except KeyError as error: + raise RuntimeError(f"Unsupported Docker host architecture: {machine!r}") from error + + +def _json_sha256(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _file_sha256(path: str) -> str: + hasher = hashlib.sha256() + with open(path, "rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(block) + return hasher.hexdigest() + + +def _sequence_output_dir(output_dir: str, seq_id: str) -> str: + """Return the per-sequence directory after enforcing path containment.""" + + if not isinstance(seq_id, str) or _SAFE_SEQUENCE_ID_RE.fullmatch(seq_id) is None: + raise ValueError( + "seq_id must use only ASCII letters, digits, dot, underscore, and hyphen" + ) + if ( + seq_id in {".", ".."} + or PurePosixPath(seq_id).name != seq_id + or PureWindowsPath(seq_id).name != seq_id + or PureWindowsPath(seq_id).is_absolute() + ): + raise ValueError(f"seq_id must be a single relative filename component: {seq_id!r}") + + output_root = Path(output_dir).resolve() + sequence_dir = Path(output_dir) / seq_id + resolved_sequence_dir = sequence_dir.resolve() + if resolved_sequence_dir.parent != output_root: + raise ValueError(f"seq_id resolves outside output_dir: {seq_id!r}") + return os.fspath(sequence_dir) + + +@dataclass +class ContextSpecification: + max_num_samples: int = 511 + max_token_length: int = 32768 + max_query_similarity: float = 1.0 + min_query_similarity: float = 0.0 + neighbor_similarity_lower_bound: float = 0.8 + + +class E1Prediction(TypedDict, total=False): + id: str | int + context_id: str | int | None + logits: torch.Tensor + token_embeddings: torch.Tensor + mean_token_embeddings: torch.Tensor + + +def read_fasta_sequences(path: str) -> dict[str, str]: + sequences: dict[str, str] = {} + header: str | None = None + parts: list[str] = [] + with open(path, encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line: + continue + if line.startswith(">"): + if header is not None: + sequences[header] = "".join(parts) + header = line[1:].strip() + parts = [] + else: + if header is None: + raise ValueError(f"FASTA sequence found before header in {path}") + parts.append(line) + if header is not None: + sequences[header] = "".join(parts) + return sequences + + +def write_fasta_sequences(path: str, sequences: dict[str, str]) -> None: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + with open(path, "w", encoding="utf-8") as handle: + for header, sequence in sequences.items(): + handle.write(f">{header}\n{sequence}\n") + + +def parse_msa(path: str) -> list[IdSequence]: + records = read_fasta_sequences(path) + sequences = [] + for record_id, record_seq in records.items(): + sequence = str(record_seq).replace("\x00", "").replace(".", "-") + sequences.append(IdSequence(record_id, sequence)) + if not sequences: + raise ValueError(f"No sequences found in MSA file: {path}") + return sequences + + +def convert_to_tensor( + sequences: list[IdSequence], device: torch.device | None = None +) -> torch.ByteTensor: + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + byte_sequences = [ + sequence.sequence.encode("ascii").translate(None, LOWERCASE_CHARS) for sequence in sequences + ] + lengths = {len(byte_sequence) for byte_sequence in byte_sequences} + if len(lengths) != 1: + raise ValueError( + "MSA rows must have equal aligned lengths after removing insertions: " + f"{sorted(lengths)}" + ) + array = np.vstack( # (n, l) + [np.frombuffer(byte_sequence, dtype=np.uint8) for byte_sequence in byte_sequences] + ) + return torch.from_numpy(array).to(device) # (n, l) + + +def get_num_neighbors(byte_seqs: torch.ByteTensor, sim_threshold: float = 0.8) -> list[int]: + # byte_seqs: (n, l) + gap_token_id = np.frombuffer(b"-", np.uint8)[0].item() + seq_lens = (byte_seqs != gap_token_id).sum(dim=1) # (n,) + num_neighbors: list[int] = [] + for i in range(byte_seqs.shape[0]): + query_non_gaps = byte_seqs[i] != gap_token_id # (l,) + seqs_sim = ( # (n,) + byte_seqs[:, query_non_gaps] == byte_seqs[i, query_non_gaps] + ).sum( + dim=1 + ) / seq_lens + num_neighbors.append(int((seqs_sim >= sim_threshold).sum().item())) + return num_neighbors + + +def get_similarity_to_query(byte_seqs: torch.ByteTensor) -> torch.FloatTensor: + # byte_seqs: (n, l) + return (byte_seqs == byte_seqs[0, :]).sum(dim=1) / byte_seqs.shape[1] # (n,) + + +def sample_context( + msa_path: str, + max_num_samples: int, + max_token_length: int, + max_query_similarity: float = 1.0, + min_query_similarity: float = 0.0, + neighbor_similarity_lower_bound: float = 0.8, + use_full_sequences_in_context: bool = False, + full_sequences_path: str | None = None, + seed: int = 0, + device: torch.device | None = None, + cache_num_neighbors_path: str | None = None, +) -> tuple[str, list[str]]: + msa_sequences = parse_msa(msa_path) + msa_as_byte_tensor = convert_to_tensor(msa_sequences, device) # (n, l) + if cache_num_neighbors_path is not None and os.path.exists(cache_num_neighbors_path): + num_neighbors = np.load(cache_num_neighbors_path) # (n,) + else: + num_neighbors = np.array( # (n,) + get_num_neighbors(msa_as_byte_tensor, neighbor_similarity_lower_bound) + ) + if cache_num_neighbors_path is not None: + np.save(cache_num_neighbors_path, num_neighbors) + + sampling_weights = 1.0 / num_neighbors # (n,) + query_similarity = get_similarity_to_query(msa_as_byte_tensor) # (n,) + filtered_mask = (query_similarity <= max_query_similarity) & ( # (n,) + query_similarity >= min_query_similarity + ) + if int(filtered_mask.sum()) < 1: + raise ValueError( + "No sequences found with similarity to query within range " + f"{min_query_similarity} <= query_similarity <= {max_query_similarity}." + ) + + filtered_weights = np.where( # (n,) + filtered_mask.cpu().numpy(), + sampling_weights, + 0.0, + ) + sampled_indices = np.random.default_rng(seed).choice( # (n_sampled,) + len(filtered_weights), + size=min(max_num_samples, int(filtered_mask.sum())), + p=filtered_weights / filtered_weights.sum(), + replace=False, + shuffle=True, + ) + + if use_full_sequences_in_context: + if full_sequences_path is None: + raise ValueError( + "full_sequences_path is required when use_full_sequences_in_context=True" + ) + full_sequences = parse_msa(full_sequences_path) + if len(full_sequences) != len(msa_sequences): + raise ValueError("Number of full sequences must match number of MSA sequences") + for i, (full_seq, msa_seq) in enumerate(zip(full_sequences, msa_sequences, strict=True)): + if full_seq.id != msa_seq.id: + raise ValueError( + "Full sequences and MSA sequences must be in the same order and have the " + f"same ids. Found differing id for sample {i}: " + f"{full_seq.id} != {msa_seq.id}" + ) + sampled_sequences = [full_sequences[int(i)] for i in sampled_indices] + else: + sampled_sequences = [msa_sequences[int(i)] for i in sampled_indices] + + context_sequences: list[str] = [] + context_ids: list[str] = [] + context_length = 0 + for seq in sampled_sequences: + seq_str = seq.sequence.upper().encode("ascii").translate(None, b"-").decode("ascii") + if context_length + len(seq_str) > max_token_length: + break + context_sequences.append(seq_str) + context_ids.append(seq.id) + context_length += len(seq_str) + return ",".join(context_sequences), context_ids + + +def sample_multiple_contexts( + msa_path: str, + context_specifications: list[ContextSpecification], + use_full_sequences_in_context: bool = False, + full_sequences_path: str | None = None, + seed: int = 0, + device: torch.device | None = None, + cache_num_neighbors_path: str | None = None, +) -> tuple[list[str], list[list[str]]]: + with tempfile.TemporaryDirectory() as temp_dir: + if cache_num_neighbors_path is None: + cache_num_neighbors_path = os.path.join(temp_dir, "num_neighbors.npy") + + contexts: list[str] = [] + context_ids: list[list[str]] = [] + for i, context_specification in enumerate(context_specifications): + context, ids = sample_context( + msa_path=msa_path, + max_num_samples=context_specification.max_num_samples, + max_token_length=context_specification.max_token_length, + max_query_similarity=context_specification.max_query_similarity, + min_query_similarity=context_specification.min_query_similarity, + neighbor_similarity_lower_bound=( + context_specification.neighbor_similarity_lower_bound + ), + use_full_sequences_in_context=use_full_sequences_in_context, + full_sequences_path=full_sequences_path, + seed=seed + i, + device=device, + cache_num_neighbors_path=cache_num_neighbors_path, + ) + contexts.append(context) + context_ids.append(ids) + return contexts, context_ids + + +def get_context_id(max_tokens: int, sim_threshold: float) -> str: + return f"identity_{sim_threshold}_tokens_{max_tokens}" + + +def build_context_specifications( + max_context_tokens: list[int] | None = None, + similarity_thresholds: list[float] | None = None, + min_query_similarity: float = 0.3, +) -> list[tuple[ContextSpecification, str]]: + if max_context_tokens is None: + max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS + if similarity_thresholds is None: + similarity_thresholds = DEFAULT_SIMILARITY_THRESHOLDS + + specs = [] + for max_tokens in max_context_tokens: + for sim_threshold in similarity_thresholds: + spec = ContextSpecification( + max_num_samples=511, + max_token_length=max_tokens, + max_query_similarity=sim_threshold, + min_query_similarity=min_query_similarity, + neighbor_similarity_lower_bound=0.8, + ) + specs.append((spec, get_context_id(max_tokens, sim_threshold))) + return specs + + +def sample_contexts_for_msa( + a3m_path: str, + context_specs: list[tuple[ContextSpecification, str]], + seed: int = 42, +) -> dict[str, str]: + specs_only = [spec for spec, _ in context_specs] + context_ids = [context_id for _, context_id in context_specs] + contexts, _ = sample_multiple_contexts( + msa_path=a3m_path, + context_specifications=specs_only, + seed=seed, + ) + return dict(zip(context_ids, contexts, strict=True)) + + +def _strip_a3m_insertions(sequence: str) -> str: + uppercase_or_gap = [char for char in sequence if char.isupper() or char in "-."] + return "".join(uppercase_or_gap).replace("-", "").replace(".", "") + + +def get_query_from_a3m(path: str) -> str: + header_found = False + seq_parts: list[str] = [] + with open(path, encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line: + continue + if line.startswith(">"): + if header_found: + break + header_found = True + continue + if header_found: + seq_parts.append(line) + if not header_found: + raise ValueError(f"No FASTA header found in A3M file: {path}") + return _strip_a3m_insertions("".join(seq_parts)) + + +def load_msa_dir(msa_dir: str) -> dict[str, str]: + msa_lookup: dict[str, str] = {} + a3m_files = list(Path(msa_dir).rglob("*.a3m")) + if not a3m_files: + raise FileNotFoundError(f"No .a3m files found in {msa_dir}") + for a3m_path in tqdm(a3m_files, desc="Loading MSAs"): + query_seq = get_query_from_a3m(str(a3m_path)) + msa_lookup[query_seq] = str(a3m_path) + _get_logger().info("Loaded %d MSAs from %s", len(msa_lookup), msa_dir) + return msa_lookup + + +def _safe_extract_tar(tar: tarfile.TarFile, output_dir: str) -> None: + output_root = Path(output_dir).resolve() + for member in tar.getmembers(): + if member.issym() or member.islnk(): + raise ValueError(f"Tar links are not allowed: {member.name}") + if member.isdev(): + raise ValueError(f"Tar device entries are not allowed: {member.name}") + target = (output_root / member.name).resolve() + if output_root != target and output_root not in target.parents: + raise ValueError(f"Unsafe tar member path: {member.name}") + tar.extractall(output_root, filter="data") + + +def load_msa_from_hf( + hf_path: str, + cache_dir: str | None = None, + token: str | None = None, +) -> dict[str, str]: + from huggingface_hub import snapshot_download + + if cache_dir is None: + cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "fastplms_msa") + os.makedirs(cache_dir, exist_ok=True) + local_dir = os.path.join(cache_dir, hf_path.replace("/", "_")) + if not os.path.exists(local_dir) or not any(Path(local_dir).rglob("*.a3m")): + local_dir = snapshot_download( + repo_id=hf_path, + repo_type="dataset", + local_dir=local_dir, + token=token, + ) + for tar_path in Path(local_dir).rglob("*.tar.gz"): + with tarfile.open(tar_path) as tar: + _safe_extract_tar(tar, str(tar_path.parent)) + return load_msa_dir(local_dir) + + +def get_msa_for_sequence( + sequence: str, msa_lookup: dict[str, str], min_identity: float = 0.95 +) -> str | None: + if sequence in msa_lookup: + return msa_lookup[sequence] + + best_match_path: str | None = None + best_identity = 0.0 + for query_seq, a3m_path in msa_lookup.items(): + if abs(len(query_seq) - len(sequence)) > 10: + continue + min_len = min(len(query_seq), len(sequence)) + if min_len == 0: + continue + matches = sum(a == b for a, b in zip(query_seq[:min_len], sequence[:min_len], strict=True)) + identity = matches / min_len + if identity > best_identity: + best_identity = identity + best_match_path = a3m_path + + if best_identity >= min_identity: + return best_match_path + return None + + +class ContextCache: + """Content-addressed JSON cache for deterministic E1 MSA contexts.""" + + _SCHEMA_VERSION = 1 + + def __init__( + self, + cache_dir: str, + specs_hash: str, + seed: int, + source_revision: str = E1_MSA_SAMPLING_SOURCE_REVISION, + ) -> None: + self.cache_dir = cache_dir + self.specs_hash = specs_hash + self.seed = seed + self.source_revision = source_revision + os.makedirs(cache_dir, exist_ok=True) + + def _cache_path(self, key: str) -> str: + safe_key = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16] + return os.path.join(self.cache_dir, f"{safe_key}_seed{self.seed}_{self.specs_hash}.json") + + def _input_fingerprint(self, key: str) -> str: + descriptor: dict[str, Any] = { + "key": os.path.abspath(key) if os.path.isfile(key) else key, + "seed": self.seed, + "source_revision": self.source_revision, + "specs_hash": self.specs_hash, + } + if os.path.isfile(key): + hasher = hashlib.sha256() + with open(key, "rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(block) + descriptor["content_sha256"] = hasher.hexdigest() + else: + descriptor["literal_key"] = key + payload = json.dumps(descriptor, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def load(self, key: str) -> dict[str, str] | None: + path = self._cache_path(key) + if not os.path.exists(path): + return None + try: + with open(path, encoding="utf-8") as handle: + payload = json.load(handle) + except (OSError, UnicodeError, json.JSONDecodeError): + return None + if not isinstance(payload, dict): + return None + if payload.get("schema_version") != self._SCHEMA_VERSION: + return None + if payload.get("input_fingerprint") != self._input_fingerprint(key): + return None + if payload.get("source_revision") != self.source_revision: + return None + contexts = payload.get("contexts") + if not isinstance(contexts, dict) or not all( + isinstance(name, str) and isinstance(context, str) for name, context in contexts.items() + ): + return None + return contexts + + def store(self, key: str, contexts: dict[str, str]) -> None: + if not all( + isinstance(name, str) and isinstance(context, str) for name, context in contexts.items() + ): + raise TypeError("contexts must map string identifiers to string contexts") + path = self._cache_path(key) + payload = { + "schema_version": self._SCHEMA_VERSION, + "source_revision": self.source_revision, + "input_fingerprint": self._input_fingerprint(key), + "contexts": contexts, + } + temp_path: str | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=self.cache_dir, + prefix=".context-", + suffix=".tmp", + delete=False, + ) as handle: + temp_path = handle.name + json.dump(payload, handle, sort_keys=True, separators=(",", ":")) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp_path, path) + temp_path = None + finally: + if temp_path is not None: + Path(temp_path).unlink(missing_ok=True) + + +def compute_ppll(logits: torch.Tensor, token_ids: torch.Tensor) -> float: + # logits: (l, c); token_ids: (l,) + if token_ids.numel() == 0: + raise ValueError("Cannot score an empty token sequence") + if token_ids.device != logits.device: + token_ids = token_ids.to(logits.device) # (l,) + if logits.shape[0] != token_ids.shape[0]: + raise ValueError( + f"Logits length {logits.shape[0]} != token_ids length {token_ids.shape[0]}" + ) + probs = logits.softmax(dim=-1) # (l, c) + token_probs = probs.gather(dim=1, index=token_ids.unsqueeze(1)).squeeze(1) # (l,) + return float(token_probs.mean().item()) + + +class _E1ContextPredictor: + def __init__( + self, + model: PreTrainedModel, + data_prep_config: DataPrepConfig | None = None, + max_batch_tokens: int = 65536, + use_cache: bool = True, + cache_size: int = 4, + save_masked_positions_only: bool = False, + fields_to_save: list[str] | None = None, + keep_predictions_in_gpu: bool = False, + progress: bool = True, + ) -> None: + self.model = model + self.max_batch_tokens = max_batch_tokens + self.batch_preparer = E1BatchPreparer(data_prep_config=data_prep_config) + self.model.eval() + self.kv_cache = KVCache(cache_size=cache_size) if use_cache else None + self.fields_to_save = fields_to_save or [ + "logits", + "token_embeddings", + "mean_token_embeddings", + ] + self.save_masked_positions_only = save_masked_positions_only + self.keep_predictions_in_gpu = keep_predictions_in_gpu + self.progress = progress + + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + + def group_by_length( + self, indexed_sequences: list[IndexedSequence] + ) -> list[list[IndexedSequence]]: + batches: list[list[IndexedSequence]] = [[]] + for idx, seq in sorted( + indexed_sequences, key=lambda idx_seq: (len(idx_seq[1]), idx_seq[0]) + ): + if len(batches[-1]) > 0 and len(seq) * (len(batches[-1]) + 1) > self.max_batch_tokens: + batches.append([]) + batches[-1].append((idx, seq)) + return batches + + def group_by_context( + self, indexed_sequences: list[IndexedSequence] + ) -> list[list[IndexedSequence]]: + batches: dict[str | None, list[IndexedSequence]] = defaultdict(list) + for idx, seq in indexed_sequences: + batches[get_context(seq)].append((idx, seq)) + return list(batches.values()) + + def batch_sequences(self, sequences: list[str]) -> list[list[int]]: + indexed_sequences: list[IndexedSequence] = list(enumerate(sequences)) + indexed_batches = self.group_by_context(indexed_sequences) + indexed_batches = list( + itertools.chain.from_iterable( + [self.group_by_length(batch) for batch in indexed_batches] + ) + ) + batches = [[item[0] for item in batch] for batch in indexed_batches] + flattened_indices = list(itertools.chain.from_iterable(batches)) + if sorted(flattened_indices) != list(range(len(sequences))): + raise RuntimeError("Batches must contain all indices with no repetition") + return batches + + @torch.no_grad() + def predict_batch( + self, sequences: list[str], sequence_metadata: list[dict[str, str | int]] + ) -> list[E1Prediction]: + outputs = self.predict_batch_padded(sequences) + outputs["logits"] = outputs["logits"].float() # (b, l, c) + outputs["embeddings"] = outputs["embeddings"].float() # (b, l, d) + + token_mask = ( # (b, l) + outputs["non_boundary_token_mask"] & outputs["last_sequence_mask"] + ) + if self.save_masked_positions_only: + token_mask = token_mask & outputs["mask_positions_mask"] # (b, l) + + predictions: list[E1Prediction] = [] + for i in range(len(sequences)): + pred: E1Prediction = {"id": sequence_metadata[i]["id"]} + if "context_id" in sequence_metadata[i]: + pred["context_id"] = sequence_metadata[i]["context_id"] + if "logits" in self.fields_to_save: + pred["logits"] = outputs["logits"][i, token_mask[i]] # (r_i, c) + if not self.keep_predictions_in_gpu: + pred["logits"] = pred["logits"].to("cpu") # (r_i, c) + if "token_embeddings" in self.fields_to_save: + pred["token_embeddings"] = outputs["embeddings"][i, token_mask[i]] # (r_i, d) + if not self.keep_predictions_in_gpu: + pred["token_embeddings"] = pred["token_embeddings"].to("cpu") # (r_i, d) + if "mean_token_embeddings" in self.fields_to_save: + pred["mean_token_embeddings"] = outputs["embeddings"][ + i, token_mask[i] + ].mean(dim=0) # (d,) + if not self.keep_predictions_in_gpu: + pred["mean_token_embeddings"] = pred["mean_token_embeddings"].to( # (d,) + "cpu" + ) + predictions.append(pred) + return predictions + + @torch.no_grad() + def predict_batch_padded(self, sequences: list[str]) -> dict[str, torch.Tensor]: + device = self.device + autocast_enabled = device.type == "cuda" + with torch.autocast(device.type, torch.bfloat16, enabled=autocast_enabled): + batch = self.batch_preparer.get_batch_kwargs(sequences, device=device) + if self.kv_cache is not None: + self.kv_cache.before_forward(batch) + + past_key_values = batch.get("past_key_values") + use_cache = bool(batch["use_cache"]) if "use_cache" in batch else False + output: E1MaskedLMOutputWithPast = self.model( + input_ids=batch["input_ids"], + within_seq_position_ids=batch["within_seq_position_ids"], + global_position_ids=batch["global_position_ids"], + sequence_ids=batch["sequence_ids"], + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=False, + output_hidden_states=False, + ) + if self.kv_cache is not None: + self.kv_cache.after_forward(batch, output) + + padding_mask = batch["input_ids"] == self.batch_preparer.pad_token_id # (b, l) + last_sequence_mask = ( # (b, l) + batch["sequence_ids"] == batch["sequence_ids"].max(dim=1).values[:, None] + ) + boundary_token_mask = self.batch_preparer.get_boundary_token_mask( # (b, l) + batch["input_ids"] + ) + mask_positions_mask = self.batch_preparer.get_mask_positions_mask( # (b, l) + batch["input_ids"] + ) + return { + "logits": output.logits, + "embeddings": output.last_hidden_state, + "last_sequence_mask": last_sequence_mask, + "non_boundary_token_mask": ~boundary_token_mask, + "mask_positions_mask": mask_positions_mask, + "valid_token_mask": ~padding_mask, + } + + @torch.no_grad() + def predict( + self, + sequences: Sequence[str], + sequence_ids: Sequence[int | str] | None = None, + context_seqs: dict[str, str] | None = None, + ) -> Iterator[E1Prediction]: + if sequence_ids is None: + sequence_ids = list(range(len(sequences))) + if context_seqs: + sequences_with_context = [ + (ctx + "," + seq, {"context_id": ctx_id, "id": sequence_id}) + for ctx_id, ctx in context_seqs.items() + for seq, sequence_id in zip(sequences, sequence_ids, strict=True) + ] + else: + sequences_with_context = [ + (seq, {"id": sequence_id}) + for seq, sequence_id in zip(sequences, sequence_ids, strict=True) + ] + + batched_sequences, sequence_metadata = tuple(zip(*sequences_with_context, strict=True)) + batches = self.batch_sequences(list(batched_sequences)) + iterator = tqdm(batches, desc="Predicting batches", disable=not self.progress) + for indices in iterator: + sequence_batch = [batched_sequences[i] for i in indices] + sequence_batch_metadata = [sequence_metadata[i] for i in indices] + yield from self.predict_batch(sequence_batch, sequence_batch_metadata) + + +def _pool_hidden_states( + hidden_list: list[torch.Tensor], + pooling_types: list[str], + device: torch.device, +) -> torch.Tensor: + # Each hidden_list entry: (l_i, d) + pooler = Pooler(pooling_types) + max_len = max(hidden.shape[0] for hidden in hidden_list) + hidden_dim = hidden_list[0].shape[1] + batch_size = len(hidden_list) + padded = torch.zeros(batch_size, max_len, hidden_dim, device=device) # (b, l_max, d) + attention_mask = torch.zeros(batch_size, max_len, device=device) # (b, l_max) + for i, hidden in enumerate(hidden_list): + seq_len = hidden.shape[0] + padded[i, :seq_len] = hidden + attention_mask[i, :seq_len] = 1.0 + return pooler(padded, attention_mask) # (b, n_poolers * d) + + +def _forward_for_embedding( + model: PreTrainedModel, + sequences: list[str], + context: str | None, + max_batch_tokens: int, + progress: bool, +) -> list[torch.Tensor]: + predictor = _E1ContextPredictor( + model=model, + data_prep_config=DataPrepConfig(remove_X_tokens=True), + max_batch_tokens=max_batch_tokens, + fields_to_save=["token_embeddings"], + keep_predictions_in_gpu=True, + use_cache=False, + cache_size=1, + progress=progress, + ) + context_seqs = {"embed_ctx": context} if context else None + predictions = list( + predictor.predict( + sequences=sequences, + sequence_ids=list(range(len(sequences))), + context_seqs=context_seqs, + ) + ) + predictions.sort(key=lambda prediction: prediction["id"]) + return [prediction["token_embeddings"] for prediction in predictions] + + +class HomologueSearcher: + """Run local MMseqs2 searches through one verified, digest-pinned image. + + The default CPU image is multi-architecture and immutable. Pulling and + container networking are separate explicit opt-ins. GPU execution requires + a caller-supplied digest-pinned GPU image because the official CUDA image is + not portable to every supported host architecture. + """ + + _PROVENANCE_SCHEMA_VERSION = 1 + _PROVENANCE_FILENAME = "search-provenance.json" + + def __init__( + self, + target_db: str, + docker_image: str = DOCKER_IMAGE, + sensitivity: float = 7.5, + max_seqs: int = 1000, + min_seq_id: float = 0.0, + coverage: float = 0.8, + split_memory_limit: str | None = None, + use_gpu: bool = False, + allow_pull: bool = False, + allow_network: bool = False, + phase_timeout: float = DEFAULT_MMSEQS2_PHASE_TIMEOUT, + target_db_identity: str | None = None, + ) -> None: + image_reference = _parse_pinned_image_reference(docker_image) + if not isinstance(target_db, str) or not target_db or "\x00" in target_db: + raise ValueError("target_db must be a non-empty path without null bytes") + numeric_values = { + "sensitivity": sensitivity, + "min_seq_id": min_seq_id, + "coverage": coverage, + } + for name, value in numeric_values.items(): + if ( + isinstance(value, bool) + or not isinstance(value, numbers.Real) + or not math.isfinite(float(value)) + ): + raise ValueError(f"{name} must be a finite real number") + if sensitivity <= 0: + raise ValueError("sensitivity must be positive") + if not 0.0 <= min_seq_id <= 1.0: + raise ValueError("min_seq_id must be in [0, 1]") + if not 0.0 <= coverage <= 1.0: + raise ValueError("coverage must be in [0, 1]") + if isinstance(max_seqs, bool) or not isinstance(max_seqs, int) or max_seqs < 1: + raise ValueError("max_seqs must be an integer >= 1") + if split_memory_limit is not None and ( + not isinstance(split_memory_limit, str) + or not split_memory_limit.strip() + or "\x00" in split_memory_limit + ): + raise ValueError("split_memory_limit must be None or a non-empty string") + if type(use_gpu) is not bool: + raise TypeError("use_gpu must be a boolean") + if type(allow_pull) is not bool: + raise TypeError("allow_pull must be a boolean") + if type(allow_network) is not bool: + raise TypeError("allow_network must be a boolean") + if ( + isinstance(phase_timeout, bool) + or not isinstance(phase_timeout, (int, float)) + or not math.isfinite(float(phase_timeout)) + or phase_timeout <= 0 + ): + raise ValueError("phase_timeout must be a finite positive number") + if target_db_identity is not None and ( + not isinstance(target_db_identity, str) or not target_db_identity.strip() + ): + raise ValueError("target_db_identity must be None or a non-empty string") + if ( + use_gpu + and image_reference.repository == MMSEQS2_IMAGE_REPOSITORY + and image_reference.digest == MMSEQS2_CPU_MANIFEST_DIGEST + ): + raise ValueError( + "The default MMseqs2 image is CPU-only. GPU search requires an explicit " + "digest-pinned image compatible with the host architecture." + ) + self.target_db = target_db + self.docker_image = docker_image + self._image_reference = image_reference + self.sensitivity = float(sensitivity) + self.max_seqs = max_seqs + self.min_seq_id = float(min_seq_id) + self.coverage = float(coverage) + self.split_memory_limit = ( + split_memory_limit.strip() if split_memory_limit is not None else None + ) + self.use_gpu = use_gpu + self.allow_pull = allow_pull + self.allow_network = allow_network + self.phase_timeout = float(phase_timeout) + self.target_db_identity = ( + target_db_identity.strip() if target_db_identity is not None else None + ) + self._verified_image_identity: _DockerImageIdentity | None = None + + @staticmethod + def _seq_hash(sequence: str) -> str: + return hashlib.md5(sequence.encode()).hexdigest()[:12] + + def _run_docker_command( + self, + cmd: list[str], + *, + phase: str = "docker command", + **kwargs, + ) -> subprocess.CompletedProcess: + kwargs["timeout"] = self.phase_timeout + try: + return subprocess.run(cmd, **kwargs) + except subprocess.TimeoutExpired as error: + raise TimeoutError( + f"MMseqs2 phase {phase!r} exceeded {self.phase_timeout:g} seconds" + ) from error + + @staticmethod + def _working_root() -> Path: + return Path.cwd().resolve(strict=True) + + def _resolve_path_under_cwd(self, path: str, *, must_exist: bool = False) -> Path: + root = self._working_root() + candidate = Path(path) + if not candidate.is_absolute(): + candidate = root / candidate + try: + resolved = candidate.resolve(strict=must_exist) + except OSError as error: + raise ValueError(f"Path cannot be resolved safely: {path!r}") from error + if resolved != root and root not in resolved.parents: + raise ValueError( + "Path must resolve under the current working directory for the Docker mount. " + f"cwd={os.fspath(root)!r}, path={os.fspath(resolved)!r}" + ) + return resolved + + def _validate_paths_under_cwd(self, *paths: str) -> None: + for path in paths: + self._resolve_path_under_cwd(path) + + def _path_in_container(self, local_path: str) -> str: + resolved = self._resolve_path_under_cwd(local_path) + relative = resolved.relative_to(self._working_root()) + return relative.as_posix() or "." + + def _docker_base_cmd(self) -> list[str]: + root = self._working_root() + cmd = ["docker", "run", "--rm"] + if not self.allow_network: + cmd.extend(["--network", "none"]) + cmd.extend(["-v", f"{os.fspath(root)}:/app", "-w", "/app"]) + if self.use_gpu: + if not torch.cuda.is_available(): + raise RuntimeError( + "use_gpu=True requires CUDA to be available in the FastPLMs host process" + ) + cmd.extend(["--gpus", "all"]) + cmd.append(self.docker_image) + return cmd + + def _inspect_docker_image(self, *, check: bool) -> _DockerImageIdentity | None: + inspect = self._run_docker_command( + ["docker", "image", "inspect", self.docker_image], + phase="image inspection", + capture_output=True, + text=True, + check=check, + ) + if inspect.returncode != 0: + stderr = inspect.stderr if isinstance(inspect.stderr, str) else "" + if "no such image" in stderr.lower() or "not found" in stderr.lower(): + return None + raise subprocess.CalledProcessError( + inspect.returncode, + inspect.args, + output=inspect.stdout, + stderr=inspect.stderr, + ) + try: + payload = json.loads(inspect.stdout) + if ( + not isinstance(payload, list) + or len(payload) != 1 + or not isinstance(payload[0], dict) + ): + raise ValueError("Docker inspect must return exactly one image object") + image = payload[0] + repo_digests = image.get("RepoDigests") + image_id = image.get("Id") + image_os = image.get("Os") + architecture = image.get("Architecture") + if not isinstance(repo_digests, list) or not all( + isinstance(value, str) for value in repo_digests + ): + raise ValueError("Docker inspect did not return RepoDigests") + expected_repo_digest = ( + f"{self._image_reference.repository}@{self._image_reference.digest}" + ) + if expected_repo_digest not in repo_digests: + raise ValueError( + "Docker image RepoDigests do not contain the requested repository " + "and manifest digest" + ) + if not isinstance(image_id, str) or _SHA256_DIGEST_RE.fullmatch(image_id) is None: + raise ValueError("Docker inspect returned an invalid image ID") + if image_os != "linux": + raise ValueError(f"MMseqs2 image OS must be 'linux', got {image_os!r}") + expected_architecture = _docker_architecture() + if architecture != expected_architecture: + raise ValueError( + "MMseqs2 image architecture does not match the host: " + f"expected {expected_architecture!r}, got {architecture!r}" + ) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error: + raise RuntimeError( + f"Docker image identity verification failed for {self.docker_image!r}" + ) from error + return _DockerImageIdentity( + reference=self.docker_image, + repository=self._image_reference.repository, + version=self._image_reference.version, + manifest_digest=self._image_reference.digest, + image_id=image_id, + os=image_os, + architecture=architecture, + ) + + def _ensure_docker_image(self) -> _DockerImageIdentity: + if self._verified_image_identity is not None: + return self._verified_image_identity + self._run_docker_command( + ["docker", "version"], + phase="Docker availability check", + capture_output=True, + text=True, + check=True, + ) + identity = self._inspect_docker_image(check=False) + if identity is None: + if not self.allow_pull: + raise RuntimeError( + "The pinned MMseqs2 image is not present locally and allow_pull=False. " + "Preload the exact image out of band or opt in with allow_pull=True." + ) + self._run_docker_command( + ["docker", "pull", self.docker_image], + phase="image pull", + check=True, + capture_output=True, + text=True, + ) + identity = self._inspect_docker_image(check=True) + if identity is None: + raise RuntimeError("Docker image inspection succeeded without a verified identity") + self._verified_image_identity = identity + return identity + + def _target_db_descriptor(self) -> dict[str, Any]: + prefix = self._resolve_path_under_cwd(self.target_db) + files: list[dict[str, Any]] = [] + for candidate in sorted(prefix.parent.glob(f"{prefix.name}*")): + resolved = self._resolve_path_under_cwd(os.fspath(candidate), must_exist=True) + if not resolved.is_file(): + continue + stat_result = resolved.stat() + files.append( + { + "path": resolved.relative_to(self._working_root()).as_posix(), + "size": stat_result.st_size, + "mtime_ns": stat_result.st_mtime_ns, + } + ) + if not files: + raise FileNotFoundError( + f"No MMseqs2 database files found for target_db prefix {self.target_db!r}" + ) + derived_identity = _json_sha256(files) + return { + "prefix": prefix.relative_to(self._working_root()).as_posix(), + "identity": self.target_db_identity or derived_identity, + "identity_kind": "explicit" if self.target_db_identity is not None else "file-metadata", + "files": files, + } + + def _request_provenance(self, sequence: str) -> dict[str, Any]: + return { + "provider": "mmseqs2", + "sequence_sha256": hashlib.sha256(sequence.encode("utf-8")).hexdigest(), + "image": { + "reference": self.docker_image, + "repository": self._image_reference.repository, + "version": self._image_reference.version, + "manifest_digest": self._image_reference.digest, + }, + "platform": {"os": "linux", "architecture": _docker_architecture()}, + "target_db": self._target_db_descriptor(), + "parameters": { + "sensitivity": self.sensitivity, + "max_seqs": self.max_seqs, + "min_seq_id": self.min_seq_id, + "coverage": self.coverage, + "split_memory_limit": self.split_memory_limit, + "use_gpu": self.use_gpu, + "allow_network": self.allow_network, + }, + } + + def _load_cached_result( + self, + a3m_output: str, + provenance_path: str, + request_provenance: dict[str, Any], + ) -> bool: + if not Path(a3m_output).is_file() or not Path(provenance_path).is_file(): + return False + try: + with open(provenance_path, encoding="utf-8") as handle: + payload = json.load(handle) + if not isinstance(payload, dict): + return False + if payload.get("schema_version") != self._PROVENANCE_SCHEMA_VERSION: + return False + if payload.get("request") != request_provenance: + return False + request_identity = _json_sha256(request_provenance) + if payload.get("request_identity_sha256") != request_identity: + return False + runtime = payload.get("runtime") + if not isinstance(runtime, dict): + return False + if runtime.get("reference") != self.docker_image: + return False + if runtime.get("repository") != self._image_reference.repository: + return False + if runtime.get("version") != self._image_reference.version: + return False + if runtime.get("manifest_digest") != self._image_reference.digest: + return False + if runtime.get("os") != "linux": + return False + if runtime.get("architecture") != _docker_architecture(): + return False + image_id = runtime.get("image_id") + if not isinstance(image_id, str) or _SHA256_DIGEST_RE.fullmatch(image_id) is None: + return False + cache_identity = _json_sha256( + {"request_identity_sha256": request_identity, "runtime": runtime} + ) + if payload.get("cache_identity_sha256") != cache_identity: + return False + result = payload.get("result") + if not isinstance(result, dict): + return False + if result.get("path") != Path(a3m_output).name: + return False + if result.get("size") != Path(a3m_output).stat().st_size: + return False + return result.get("sha256") == _file_sha256(a3m_output) + except (OSError, UnicodeError, json.JSONDecodeError, TypeError, ValueError): + return False + + def _store_result_provenance( + self, + provenance_path: str, + a3m_output: str, + request_provenance: dict[str, Any], + identity: _DockerImageIdentity, + ) -> None: + request_identity = _json_sha256(request_provenance) + runtime = identity.to_dict() + payload = { + "schema_version": self._PROVENANCE_SCHEMA_VERSION, + "request": request_provenance, + "request_identity_sha256": request_identity, + "runtime": runtime, + "cache_identity_sha256": _json_sha256( + {"request_identity_sha256": request_identity, "runtime": runtime} + ), + "result": { + "path": Path(a3m_output).name, + "size": Path(a3m_output).stat().st_size, + "sha256": _file_sha256(a3m_output), + }, + } + output_dir = os.path.dirname(provenance_path) or "." + temporary_path: str | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=output_dir, + prefix=".mmseqs2-provenance-", + suffix=".tmp", + delete=False, + ) as handle: + temporary_path = handle.name + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary_path, provenance_path) + temporary_path = None + finally: + if temporary_path is not None: + Path(temporary_path).unlink(missing_ok=True) + + def create_db(self, fasta_path: str, db_path: str) -> str: + self._validate_paths_under_cwd(fasta_path, db_path) + os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) + if os.path.exists(f"{db_path}.dbtype"): + return db_path + self._ensure_docker_image() + self._run_docker_command( + [ + *self._docker_base_cmd(), + "createdb", + self._path_in_container(fasta_path), + self._path_in_container(db_path), + ], + phase="createdb", + check=True, + capture_output=True, + text=True, + ) + return db_path + + def create_index(self, db_path: str, tmp_dir: str | None = None) -> None: + if tmp_dir is None: + tmp_dir = os.path.join(os.path.dirname(db_path), "tmp_index") + self._validate_paths_under_cwd(db_path, tmp_dir) + os.makedirs(tmp_dir, exist_ok=True) + self._ensure_docker_image() + self._run_docker_command( + [ + *self._docker_base_cmd(), + "createindex", + self._path_in_container(db_path), + self._path_in_container(tmp_dir), + ], + phase="createindex", + check=True, + capture_output=True, + text=True, + ) + + def search(self, sequence: str, output_dir: str, seq_id: str | None = None) -> str: + if ( + not isinstance(sequence, str) + or _SAFE_QUERY_SEQUENCE_RE.fullmatch(sequence) is None + ): + raise ValueError( + "sequence must be a non-empty unaligned ASCII protein sequence" + ) + if seq_id is None: + seq_id = self._seq_hash(sequence) + seq_output_dir = _sequence_output_dir(output_dir, seq_id) + a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m") + provenance_path = os.path.join(seq_output_dir, self._PROVENANCE_FILENAME) + self._validate_paths_under_cwd(seq_output_dir, self.target_db) + request_provenance = self._request_provenance(sequence) + if self._load_cached_result(a3m_output, provenance_path, request_provenance): + return a3m_output + + identity = self._ensure_docker_image() + os.makedirs(seq_output_dir, exist_ok=True) + Path(a3m_output).unlink(missing_ok=True) + Path(provenance_path).unlink(missing_ok=True) + query_fasta = os.path.join(seq_output_dir, "query.fasta") + write_fasta_sequences(query_fasta, {seq_id: sequence}) + query_db = os.path.join(seq_output_dir, "queryDB") + result_db = os.path.join(seq_output_dir, "resultDB") + tmp_dir = os.path.join(seq_output_dir, "tmp") + os.makedirs(tmp_dir, exist_ok=True) + self._validate_paths_under_cwd( + query_fasta, query_db, self.target_db, seq_output_dir, result_db, tmp_dir + ) + + docker_base = self._docker_base_cmd() + self._run_docker_command( + [ + *docker_base, + "createdb", + self._path_in_container(query_fasta), + self._path_in_container(query_db), + ], + phase="query createdb", + check=True, + capture_output=True, + text=True, + ) + search_cmd = [ + *docker_base, + "search", + self._path_in_container(query_db), + self._path_in_container(self.target_db), + self._path_in_container(result_db), + self._path_in_container(tmp_dir), + "-s", + str(self.sensitivity), + "--max-seqs", + str(self.max_seqs), + "--min-seq-id", + str(self.min_seq_id), + "-c", + str(self.coverage), + ] + if self.split_memory_limit is not None: + search_cmd.extend(["--split-memory-limit", self.split_memory_limit]) + if self.use_gpu and torch.cuda.is_available(): + search_cmd.extend(["--gpu", "1"]) + self._run_docker_command( + search_cmd, + phase="search", + check=True, + capture_output=True, + text=True, + ) + self._run_docker_command( + [ + *docker_base, + "result2msa", + self._path_in_container(query_db), + self._path_in_container(self.target_db), + self._path_in_container(result_db), + self._path_in_container(a3m_output), + "--msa-format-mode", + "6", + ], + phase="result2msa", + check=True, + capture_output=True, + text=True, + ) + if not Path(a3m_output).is_file(): + raise RuntimeError("MMseqs2 result2msa did not create a regular A3M file") + resolved_a3m = self._resolve_path_under_cwd(a3m_output, must_exist=True) + if not resolved_a3m.is_file(): + raise RuntimeError("MMseqs2 result2msa did not create a regular A3M file") + self._store_result_provenance( + provenance_path, + a3m_output, + request_provenance, + identity, + ) + for pattern in ["queryDB*", "resultDB*"]: + for path in Path(seq_output_dir).glob(pattern): + path.unlink(missing_ok=True) + tmp_path = Path(tmp_dir) + if tmp_path.exists(): + shutil.rmtree(tmp_path, ignore_errors=True) + return a3m_output + + def batch_search( + self, + sequences: list[str], + output_dir: str, + seq_ids: list[str] | None = None, + continue_on_error: bool = True, + ) -> dict[str, str]: + if seq_ids is None: + seq_ids = [self._seq_hash(seq) for seq in sequences] + if len(seq_ids) != len(sequences): + raise ValueError("seq_ids must contain exactly one identifier per sequence") + self._validate_paths_under_cwd(output_dir) + os.makedirs(output_dir, exist_ok=True) + results: dict[str, str] = {} + for seq, sid in tqdm( + list(zip(sequences, seq_ids, strict=True)), + desc="Searching homologues", + ): + try: + results[seq] = self.search(seq, output_dir, sid) + except Exception as error: + if not continue_on_error: + raise + _get_logger().warning( + "Homologue search failed and was skipped: " + "provider=mmseqs2 seq_id=%s error_type=%s", + sid, + type(error).__name__, + ) + return results + + +@dataclass(frozen=True) +class _ColabFoldResponse: + """Minimal response surface required by the ColabFold API client.""" + + status_code: int + headers: dict[str, str] + content: bytes + + def json(self) -> dict[str, Any]: + value = json.loads(self.content.decode("utf-8")) + if not isinstance(value, dict): + raise ValueError("ColabFold returned a non-object JSON response") + return value + + +class ColabFoldSearcher: + def __init__( + self, + host_url: str = COLABFOLD_HOST, + user_agent: str = "", + mode: str = "env", + timeout: float = 30.0, + max_retries: int = 10, + base_delay: float = 1.0, + max_delay: float = 60.0, + inter_request_delay: tuple[float, float] = (1.0, 3.0), + max_wait_time: int = 600, + ) -> None: + self.host_url = host_url.rstrip("/") + self.mode = mode + self.timeout = timeout + self.max_retries = max_retries + self.base_delay = base_delay + self.max_delay = max_delay + self.inter_request_delay = inter_request_delay + self.max_wait_time = max_wait_time + self.headers = {"User-Agent": user_agent} if user_agent else {} + + @staticmethod + def _seq_hash(sequence: str) -> str: + return hashlib.md5(sequence.encode()).hexdigest()[:12] + + def _backoff_delay(self, attempt: int) -> float: + delay = min(self.base_delay * (2**attempt), self.max_delay) + return min(delay + random.uniform(0, delay * 0.5), self.max_delay) + + def _retry_after_delay(self, headers: dict[str, str], attempt: int) -> float: + raw_value = next( + (value for name, value in headers.items() if name.lower() == "retry-after"), + None, + ) + if raw_value is None: + return self._backoff_delay(attempt) + try: + delay = float(raw_value) + except (TypeError, ValueError): + try: + retry_at = parsedate_to_datetime(raw_value) + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo=UTC) + delay = (retry_at - datetime.now(UTC)).total_seconds() + except (TypeError, ValueError, OverflowError): + return self._backoff_delay(attempt) + if not math.isfinite(delay): + return self._backoff_delay(attempt) + return min(max(0.0, delay), self.max_delay) + + def _remaining_timeout(self, deadline: float | None, context: str) -> float: + if deadline is None: + return self.timeout + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline") + return min(self.timeout, remaining) + + def _sleep_with_deadline( + self, + delay: float, + deadline: float | None, + context: str, + ) -> None: + delay = max(0.0, delay) + if deadline is None: + time.sleep(delay) + return + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline") + if delay >= remaining: + time.sleep(remaining) + raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline") + time.sleep(delay) + + @staticmethod + def _http_error(response: _ColabFoldResponse, url: str) -> RuntimeError: + return RuntimeError(f"ColabFold request to {url} returned HTTP {response.status_code}") + + def _request_with_retries( + self, + method: str, + url: str, + *, + deadline: float | None = None, + **kwargs: Any, + ) -> _ColabFoldResponse: + payload = kwargs.pop("data", None) + if kwargs: + unexpected = ", ".join(sorted(kwargs)) + raise TypeError(f"Unsupported HTTP request options: {unexpected}") + + encoded_payload = None + headers = dict(self.headers) + if payload is not None: + encoded_payload = urllib.parse.urlencode(payload).encode("utf-8") + headers["Content-Type"] = "application/x-www-form-urlencoded" + + last_error: BaseException | None = None + for attempt in range(self.max_retries): + try: + request = urllib.request.Request( + url, + data=encoded_payload, + headers=headers, + method=method.upper(), + ) + try: + timeout = self._remaining_timeout(deadline, f"Request to {url}") + with urllib.request.urlopen(request, timeout=timeout) as stream: + response = _ColabFoldResponse( + status_code=int(stream.status), + headers={name.lower(): value for name, value in stream.headers.items()}, + content=stream.read(), + ) + except urllib.error.HTTPError as error: + response = _ColabFoldResponse( + status_code=int(error.code), + headers={name.lower(): value for name, value in error.headers.items()}, + content=error.read(), + ) + if response.status_code == 429: + last_error = self._http_error(response, url) + if attempt + 1 >= self.max_retries: + break + self._sleep_with_deadline( + self._retry_after_delay(response.headers, attempt), + deadline, + f"Request to {url}", + ) + continue + if response.status_code >= 500: + last_error = self._http_error(response, url) + if attempt + 1 >= self.max_retries: + break + self._sleep_with_deadline( + self._backoff_delay(attempt), + deadline, + f"Request to {url}", + ) + continue + if not 200 <= response.status_code < 300: + raise self._http_error(response, url) + return response + except (TimeoutError, urllib.error.URLError) as error: + last_error = error + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError( + f"Request to {url} exceeded the {self.max_wait_time}s deadline" + ) from error + if attempt + 1 >= self.max_retries: + break + self._sleep_with_deadline( + self._backoff_delay(attempt), + deadline, + f"Request to {url}", + ) + raise RuntimeError( + f"Request to {url} failed after {self.max_retries} attempts" + ) from last_error + + def _submit( + self, + sequence: str, + mode: str | None = None, + deadline: float | None = None, + ) -> dict[str, Any]: + mode = mode or self.mode + query = f">101\n{sequence}\n" + for attempt in range(self.max_retries): + response = self._request_with_retries( + "POST", + f"{self.host_url}/ticket/msa", + data={"q": query, "mode": mode}, + deadline=deadline, + ) + data = response.json() + status = data.get("status", "UNKNOWN") + if status in ("RATELIMIT", "UNKNOWN"): + if attempt + 1 >= self.max_retries: + break + self._sleep_with_deadline( + self._backoff_delay(attempt), + deadline, + "ColabFold job submission", + ) + continue + return data + raise RuntimeError(f"Failed to submit sequence after {self.max_retries} attempts") + + def _poll(self, ticket_id: str, deadline: float | None = None) -> dict[str, Any]: + if deadline is None: + deadline = time.monotonic() + self.max_wait_time + poll_interval = 1.0 + while True: + response = self._request_with_retries( + "GET", + f"{self.host_url}/ticket/{ticket_id}", + deadline=deadline, + ) + data = response.json() + status = data.get("status", "ERROR") + if status in ("COMPLETE", "ERROR"): + return data + if status not in ("RUNNING", "PENDING", "UNKNOWN"): + return data + wait = min(poll_interval + random.uniform(0, 0.5), 5.0) + self._sleep_with_deadline(wait, deadline, f"Job {ticket_id}") + poll_interval = min(poll_interval + 1.0, 5.0) + + def _download( + self, + ticket_id: str, + output_path: str, + deadline: float | None = None, + ) -> None: + response = self._request_with_retries( + "GET", + f"{self.host_url}/result/download/{ticket_id}", + deadline=deadline, + ) + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + with open(output_path, "wb") as handle: + handle.write(response.content) + + def _extract_a3m(self, tar_path: str, output_dir: str, seq_id: str) -> str: + with tarfile.open(tar_path) as tar: + _safe_extract_tar(tar, output_dir) + + uniref_a3m = os.path.join(output_dir, "uniref.a3m") + env_a3m = os.path.join(output_dir, "bfd.mgnify30.metaeuk30.smag30.a3m") + a3m_files: list[str] = [] + if os.path.exists(uniref_a3m): + a3m_files.append(uniref_a3m) + if "env" in self.mode and os.path.exists(env_a3m): + a3m_files.append(env_a3m) + combined_path = os.path.join(output_dir, f"{seq_id}.a3m") + if len(a3m_files) == 1: + os.replace(a3m_files[0], combined_path) + elif len(a3m_files) > 1: + with open(combined_path, "w", encoding="utf-8") as out_handle: + for a3m_file in a3m_files: + with open(a3m_file, encoding="utf-8") as in_handle: + out_handle.write(in_handle.read()) + else: + raise RuntimeError("No .a3m files found in downloaded archive") + if os.path.exists(tar_path): + os.remove(tar_path) + for a3m_file in a3m_files: + if os.path.exists(a3m_file) and a3m_file != combined_path: + os.remove(a3m_file) + return combined_path + + def search(self, sequence: str, output_dir: str, seq_id: str | None = None) -> str: + if seq_id is None: + seq_id = self._seq_hash(sequence) + seq_output_dir = _sequence_output_dir(output_dir, seq_id) + a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m") + if os.path.exists(a3m_output): + return a3m_output + os.makedirs(seq_output_dir, exist_ok=True) + deadline = time.monotonic() + self.max_wait_time + result = self._submit(sequence, deadline=deadline) + status = result.get("status", "UNKNOWN") + if status == "ERROR": + raise RuntimeError(f"ColabFold API error for {seq_id}") + if status == "MAINTENANCE": + raise RuntimeError("ColabFold API is under maintenance") + ticket_id = result["id"] + result = self._poll(ticket_id, deadline=deadline) + status = result.get("status", "UNKNOWN") + if status != "COMPLETE": + raise RuntimeError(f"Job failed for {seq_id}: {status}") + tar_path = os.path.join(seq_output_dir, f"{seq_id}.tar.gz") + self._download(ticket_id, tar_path, deadline=deadline) + return self._extract_a3m(tar_path, seq_output_dir, seq_id) + + def batch_search( + self, + sequences: list[str], + output_dir: str, + seq_ids: list[str] | None = None, + continue_on_error: bool = True, + ) -> dict[str, str]: + if seq_ids is None: + seq_ids = [self._seq_hash(seq) for seq in sequences] + os.makedirs(output_dir, exist_ok=True) + results: dict[str, str] = {} + pairs = list(zip(sequences, seq_ids, strict=True)) + for i, (seq, sid) in enumerate(tqdm(pairs, desc="ColabFold search")): + try: + results[seq] = self.search(seq, output_dir, sid) + except Exception as error: + if not continue_on_error: + raise + _get_logger().warning( + "Homologue search failed and was skipped: " + "provider=colabfold seq_id=%s error_type=%s", + sid, + type(error).__name__, + ) + if i < len(pairs) - 1: + time.sleep(random.uniform(*self.inter_request_delay)) + return results + + +def _make_homologue_searcher( + provider: str, target_db: str | None, **kwargs +) -> HomologueSearcher | ColabFoldSearcher: + if provider == "mmseqs2": + if target_db is None: + raise ValueError("target_db is required for MMseqs2 homologue search") + return HomologueSearcher(target_db=target_db, **kwargs) + if provider == "colabfold": + return ColabFoldSearcher(**kwargs) + raise ValueError(f"Unknown homologue search provider: {provider}") diff --git a/fastplms/e1/tokenizer.json b/src/fastplms/models/e1/tokenizer.json similarity index 100% rename from fastplms/e1/tokenizer.json rename to src/fastplms/models/e1/tokenizer.json diff --git a/fastplms/dplm/__init__.py b/src/fastplms/models/esm2/__init__.py similarity index 100% rename from fastplms/dplm/__init__.py rename to src/fastplms/models/esm2/__init__.py diff --git a/src/fastplms/models/esm2/modeling_fastesm.py b/src/fastplms/models/esm2/modeling_fastesm.py new file mode 100644 index 0000000..ce4b95f --- /dev/null +++ b/src/fastplms/models/esm2/modeling_fastesm.py @@ -0,0 +1,1126 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +from dataclasses import dataclass +from typing import Any, ClassVar +from einops import rearrange +from torch.nn import functional as F +from transformers import EsmTokenizer, PretrainedConfig, PreTrainedModel +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.models.esm.modeling_esm import ( + EsmClassificationHead, + EsmContactPredictionHead, + EsmEmbeddings, + EsmIntermediate, + EsmLMHead, + EsmOutput, + EsmPooler, + EsmSelfOutput, +) + +from fastplms.models._esm_rotary import RotaryEmbedding + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + flex_attention, + get_attention_mask, + kernels_flash_attention_func, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin, select_hidden_state_embeddings + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "_get_flex_attention_fn", + "flex_attention", + "get_attention_mask", + "kernels_flash_attention_func", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +@dataclass +class FastEsmEncoderOutput(ModelOutput): + last_hidden_state: torch.Tensor | None = None + pooler_output: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class EsmMaskedLMOutput(MaskedLMOutput): + """Masked-LM output with FastPLMs diagnostics after the HF fields.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + last_hidden_state: torch.Tensor | None = None + + +@dataclass +class EsmSequenceClassifierOutput(SequenceClassifierOutput): + """Sequence-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class EsmTokenClassifierOutput(TokenClassifierOutput): + """Token-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +class FastEsmConfig(PretrainedConfig): + model_type = "fast_esm" + + def __init__( + self, + vocab_size: int | None = None, + bos_token_id: int | None = 0, + eos_token_id: int | None = 2, + mask_token_id: int | None = None, + pad_token_id: int | None = None, + hidden_size: int = 768, + num_hidden_layers: int = 12, + num_attention_heads: int = 12, + intermediate_size: int = 3072, + hidden_dropout_prob: float = 0.1, + attention_probs_dropout_prob: float = 0.1, + max_position_embeddings: int = 1026, + initializer_range: float = 0.02, + layer_norm_eps: float = 1e-12, + position_embedding_type: str = "rotary", + emb_layer_norm_before: bool | None = None, + token_dropout: bool = True, + add_pooling_layer: bool = False, + attn_backend: str | None = None, + **kwargs, + ): + bos_token_id = 0 if bos_token_id is None else bos_token_id + eos_token_id = 2 if eos_token_id is None else eos_token_id + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + mask_token_id=mask_token_id, + **kwargs, + ) + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.intermediate_size = intermediate_size + self.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.layer_norm_eps = layer_norm_eps + self.position_embedding_type = position_embedding_type + self.emb_layer_norm_before = emb_layer_norm_before + self.tie_word_embeddings = False + self.token_dropout = token_dropout + self.add_pooling_layer = add_pooling_layer + self.attn_backend = attn_backend + + def to_dict(self) -> dict[str, Any]: + """Serialize the complete configuration to a Python dictionary.""" + return super().to_dict() + + +_TOKENIZER_LOAD_CONTEXT_KEYS = ( + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", + "trust_remote_code", +) + + +class FastEsmTokenizer(EsmTokenizer): + """Retain fair-esm's strict handling of residues outside its alphabet.""" + + def __call__( + self, + text: Any = None, + *args: Any, + truncation: Any = None, + max_length: int | None = None, + **kwargs: Any, + ) -> Any: + if truncation and max_length is not None: + residue_limit = max(1, max_length - 2) + if isinstance(text, str): + text = text[:residue_limit] + elif isinstance(text, (list, tuple)) and all( + isinstance(sequence, str) for sequence in text + ): + text = [sequence[:residue_limit] for sequence in text] + return super().__call__( + text, + *args, + truncation=truncation, + max_length=max_length, + **kwargs, + ) + + def _convert_token_to_id(self, token: str) -> int: + try: + return self._token_to_id[token] + except KeyError: + raise KeyError(token) from None + + +class EsmSelfAttention(nn.Module): + def __init__(self, config, position_embedding_type: str | None = None) -> None: + super().__init__() + if config.hidden_size % config.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of " + f"attention heads ({config.num_attention_heads})" + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + self.scale = self.attention_head_size**-0.5 + + self.dropout_prob = config.attention_probs_dropout_prob + self.config = config + self.attn_backend = resolve_attention_backend(config.attn_backend) + self.position_embedding_type = position_embedding_type or config.position_embedding_type + self.rotary_embeddings = None + if self.position_embedding_type == "rotary": + self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d); masks: (b, l) and (b, 1, 1, l) + batch_size, seq_length = hidden_states.shape[:-1] + hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) + query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + + query_heads = query_heads * self.scale # (b, h, l, d_h) + + if self.position_embedding_type == "rotary": + query_heads, key_heads = self.rotary_embeddings( # both (b, h, l, d_h) + query_heads, + key_heads, + ) + + attn_output, attn_weights, s_max = self._attn( # (b, l, d), (b, h, l, l), heads + query_heads, + key_heads, + value_heads, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + return attn_output, attn_weights, s_max # (b, l, d), optional (b, h, l, l), heads + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + if output_attentions: + return self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + + if ( + self.training + and self.dropout_prob > 0 + and (self.attn_backend.is_flash or self.attn_backend == AttentionBackend.FLEX_ATTENTION) + ): + raise RuntimeError( + f"ESM2 {self.attn_backend.value} attention is inference-only when attention " + "dropout is nonzero. Use eager or SDPA for this training configuration." + ) + + if self.attn_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + return attn_output, None, s_max + if self.attn_backend.is_flash: + attn_output, attn_weights = self._kernels_flash_attn( + query_heads, key_heads, value_heads, attention_mask_2d + ) + elif self.attn_backend == AttentionBackend.FLEX: + attn_output, attn_weights = self._flex_attn( + query_heads, + key_heads, + value_heads, + flex_block_mask, + attention_mask_2d, + ) + elif self.attn_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_heads, key_heads, value_heads, attention_mask_4d + ) + else: + raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") + + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, query_heads: torch.Tensor, key_heads: torch.Tensor + ) -> list[torch.Tensor]: + # query_heads, key_heads: (b, h, l, d_h) + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) + s_max_bound = ( # (h,) + q_norm.max(dim=-1).values * k_norm.max(dim=-1).values + ).max(dim=0).values + return [s_max_bound[h] for h in range(self.num_attention_heads)] # h scalars + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = torch.matmul(query_heads, key_heads.transpose(-1, -2)) # (b, h, l, l) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill( # (b, h, l, l) + attention_mask_4d.logical_not(), + float("-inf"), + ) + attn_weights = F.softmax(attn_weights, dim=-1) # (b, h, l, l) + if self.dropout_prob > 0 and self.training: + attn_weights = F.dropout( # (b, h, l, l) + attn_weights, + p=self.dropout_prob, + training=self.training, + ) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l, d) + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max # (b, l, d), (b, h, l, l), heads + + def _kernels_flash_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + query_tokens = query_heads.transpose(1, 2).contiguous() # (b, l, h, d_h) + key_tokens = key_heads.transpose(1, 2).contiguous() # (b, l, h, d_h) + value_tokens = value_heads.transpose(1, 2).contiguous() # (b, l, h, d_h) + # Q has been pre-scaled by self.scale = 1/sqrt(head_dim) in forward(). + # Pass softmax_scale=1.0 to prevent the kernel from applying its default + # 1/sqrt(head_dim) scale on top (which would yield effective scale + # 1/head_dim and break parity vs sdpa). + attn_output = kernels_flash_attention_func( # (b, l, h, d_h) + query_states=query_tokens, + key_states=key_tokens, + value_states=value_tokens, + attention_mask_2d=attention_mask_2d, + causal=False, + softmax_scale=1.0, + implementation=self.attn_backend.value, + ) + return rearrange(attn_output, "b s h d -> b s (h d)"), None # (b, l, d), None + + def _flex_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + flex_block_mask: BlockMask | None = None, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if flex_attention is None: + raise RuntimeError("Flex attention is not available in this environment.") + fn = _get_flex_attention_fn( + device=query_heads.device, + dtype=query_heads.dtype, + shape=tuple(query_heads.shape), + mask_semantics="padding", + ) + context_heads = fn( # (b, h, l, d_h) + query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0 + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None # (b, l, d), None + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + dropout_p=self.dropout_prob if self.training else 0.0, + scale=1.0, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +class EsmAttention(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.self = EsmSelfAttention(config) + self.output = EsmSelfOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # hidden_states: (b, l, d) + hidden_states_ln = self.LayerNorm(hidden_states) # (b, l, d) + attn_output, attn_weights, s_max = self.self( # (b, l, d), optional (b, h, l, l), heads + hidden_states_ln, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + attention_output = self.output(attn_output, hidden_states) # (b, l, d) + return attention_output, attn_weights, s_max # (b, l, d), optional weights, heads + + +class EsmLayer(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = EsmAttention(config) + self.intermediate = EsmIntermediate(config) + self.output = EsmOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + attention_output, attn_weights, s_max = self.attention( # (b, l, d), weights, heads + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + layer_output = self.feed_forward_chunk(attention_output) # (b, l, d) + return layer_output, attn_weights, s_max # (b, l, d), weights, heads + + def feed_forward_chunk(self, attention_output: torch.Tensor) -> torch.Tensor: + # attention_output: (b, l, d) + attention_output_ln = self.LayerNorm(attention_output) # (b, l, d) + intermediate_output = self.intermediate(attention_output_ln) # (b, l, d_ff) + layer_output = self.output(intermediate_output, attention_output) # (b, l, d) + return layer_output # (b, l, d) + + +class EsmEncoder(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.attention_backend = resolve_attention_backend(config.attn_backend) + self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)]) + self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool = False, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> FastEsmEncoderOutput: + # hidden_states: (b, l, d); attention_mask: (b, l) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + full_s_max = () if output_s_max else None + + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( + effective_backend=effective_backend, + batch_size=hidden_states.shape[0], + seq_len=hidden_states.shape[1], + device=hidden_states.device, + attention_mask=attention_mask, + dtype=hidden_states.dtype, + mask_semantics="padding", + ) + + for layer_module in self.layer: + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + if self.gradient_checkpointing and self.training: + # hidden_states: (b, l, d) + hidden_states, attn_weights, s_max = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + attention_mask_2d, + attention_mask_4d, + flex_block_mask, + output_attentions, + output_s_max, + ) + else: + hidden_states, attn_weights, s_max = layer_module( # (b, l, d), weights, heads + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + if all_attentions is not None: + all_attentions = (*all_attentions, attn_weights) + if full_s_max is not None: + full_s_max = (*full_s_max, s_max) + + if self.emb_layer_norm_after: + hidden_states = self.emb_layer_norm_after(hidden_states) # (b, l, d) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + return FastEsmEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + s_max=full_s_max, + ) + + +class FastEsmPreTrainedModel(FastPLMsAttentionMixin, PreTrainedModel): + """Initialize weights and provide the shared pretrained-model interface.""" + + config_class = FastEsmConfig + # Every advertised task wrapper stores the shared encoder at ``self.esm``. + # Transformers uses this name for ``base_model`` and for loading an + # unprefixed base checkpoint into a prefixed task wrapper. + base_model_prefix = "esm" + supports_gradient_checkpointing = True + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_flash_attn_3 = True + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs} + if "token" not in load_context and "use_auth_token" in kwargs: + load_context["token"] = kwargs["use_auth_token"] + load_context["source"] = pretrained_model_name_or_path + + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + model = loaded[0] if isinstance(loaded, tuple) else loaded + model.__dict__["_fastplms_tokenizer_load_context"] = load_context + model.__dict__["_fastplms_tokenizer"] = None + return loaded + + @property + def tokenizer(self): + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + load_context = dict(self.__dict__.get("_fastplms_tokenizer_load_context") or {}) + source = load_context.pop("source", None) + if source is None: + source = str(getattr(self.config, "_name_or_path", "")).strip() + if not source: + raise RuntimeError( + "ESM2 tokenizer loading requires a model loaded with from_pretrained " + "so checkpoint provenance is available." + ) + tokenizer_kwargs = { + key: value + for key, value in load_context.items() + if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None + } + resolved_revision = getattr(self.config, "_commit_hash", None) + if resolved_revision: + tokenizer_kwargs["revision"] = resolved_revision + tokenizer = FastEsmTokenizer.from_pretrained(source, **tokenizer_kwargs) + if getattr(tokenizer, "bos_token_id", None) is None and hasattr(tokenizer, "cls_token"): + tokenizer.bos_token = tokenizer.cls_token + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + @torch.no_grad() + def _init_weights(self, module: nn.Module) -> None: + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def post_init(self) -> None: + super().post_init() + + def get_output_embeddings(self): + # NOTE: get_output_embeddings() must return None to prevent accidental weight tying. + # See e.g. https://github.com/huggingface/transformers/pull/39339#discussion_r2219126400 + return None + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"{type(self).__name__} does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.config.attn_backend = backend + resolved = resolve_attention_backend(backend) + for module in self.modules(): + if isinstance(module, EsmEncoder): + module.attention_backend = resolved + elif isinstance(module, EsmSelfAttention): + module.attn_backend = resolved + + +class FAST_ESM_ENCODER(FastEsmPreTrainedModel, EmbeddingMixin): + def __init__(self, config, add_pooling_layer: bool | None = True, **kwargs): + FastEsmPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.embeddings = EsmEmbeddings(config) + self.encoder = EsmEncoder(config) + self.contact_head = EsmContactPredictionHead( + in_features=config.num_hidden_layers * config.num_attention_heads, bias=True + ) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + encoder_outputs = self.encoder( + token_embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + attns = self( + input_ids, + attention_mask=attention_mask, + output_attentions=True, + return_dict=True, + ).attentions + attns = torch.stack(attns, dim=1) + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) + return self.contact_head(input_ids, attns) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> FastEsmEncoderOutput | tuple[torch.Tensor, ...]: + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + elif inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + token_embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + token_embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + result = FastEsmEncoderOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + s_max=encoder_outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class FastEsmModel(FastEsmPreTrainedModel, EmbeddingMixin): + def __init__(self, config, add_pooling_layer: bool | None = None, **kwargs) -> None: + FastEsmPreTrainedModel.__init__(self, config, **kwargs) + self.config = config + self.esm = FAST_ESM_ENCODER(config) + if add_pooling_layer is None: + add_pooling_layer = config.add_pooling_layer + config.add_pooling_layer = bool(add_pooling_layer) + self.pooler = EsmPooler(config) if add_pooling_layer else None + self.post_init() + + def get_input_embeddings(self): + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.embeddings.word_embeddings = value + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> FastEsmEncoderOutput | tuple[torch.Tensor, ...]: + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.esm( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None # (b, d) + + result = FastEsmEncoderOutput( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class FastEsmForMaskedLM(FastPLMTestTimeTrainingMixin, FastEsmPreTrainedModel, EmbeddingMixin): + def __init__(self, config, **kwargs) -> None: + FastEsmPreTrainedModel.__init__(self, config, **kwargs) + self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) + self.lm_head = EsmLMHead(config) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + self.init_ttt({"lora_target_replace_module": "EsmAttention"}) + + def get_input_embeddings(self): + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.lm_head.decoder + + def set_output_embeddings(self, new_embeddings): + old_bias = self.lm_head.bias + new_vocab_size = int(new_embeddings.out_features) + if old_bias.shape[0] != new_vocab_size: + resized_bias = old_bias.new_zeros(new_vocab_size) + copy_length = min(old_bias.shape[0], new_vocab_size) + with torch.no_grad(): + resized_bias[:copy_length].copy_(old_bias[:copy_length]) + self.lm_head.bias = nn.Parameter(resized_bias) + # EsmLMHead.forward adds this standalone bias after the decoder. HF's + # generic LM-head resizer may create a biased Linear, which would apply + # the bias twice and introduce an undeclared shared tensor on save. + new_embeddings.bias = None + self.lm_head.decoder = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.esm] + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> EsmMaskedLMOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.esm( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + prediction_scores = self.lm_head(sequence_output) # (b, l, c) + + loss = None + if labels is not None: + labels = labels.to(prediction_scores.device) # (b, l) + loss = self.loss_fct( # () + prediction_scores.view(-1, self.config.vocab_size), labels.view(-1) + ) + + result = EsmMaskedLMOutput( + loss=loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + last_hidden_state=sequence_output, + ) + return result if return_dict else result.to_tuple() + + +class FastEsmForSequenceClassification(FastEsmPreTrainedModel, EmbeddingMixin): + def __init__(self, config, **kwargs) -> None: + FastEsmPreTrainedModel.__init__(self, config, **kwargs) + self.num_labels = config.num_labels + self.config = config + self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) + self.classifier = EsmClassificationHead(config) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.post_init() + + def get_input_embeddings(self): + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.set_input_embeddings(value) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> EsmSequenceClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.esm( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + logits = self.classifier(sequence_output) # (b, c) + + loss = None + if labels is not None: + labels = labels.to(logits.device) # (b,) or (b, c) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.squeeze(), labels.squeeze()) # () + else: + loss = self.mse(logits, labels) # () + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) # () + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) # () + + result = EsmSequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() + + +class FastEsmForTokenClassification(FastEsmPreTrainedModel, EmbeddingMixin): + def __init__(self, config, **kwargs) -> None: + FastEsmPreTrainedModel.__init__(self, config, **kwargs) + self.num_labels = config.num_labels + self.esm = FAST_ESM_ENCODER(config, add_pooling_layer=False) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + self.loss_fct = nn.CrossEntropyLoss() + self.post_init() + + def get_input_embeddings(self): + return self.esm.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.esm.set_input_embeddings(value) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + return self.esm._embed( + input_ids, + attention_mask, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def predict_contacts( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + return self.esm.predict_contacts(input_ids, attention_mask=attention_mask) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> EsmTokenClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + outputs = self.esm( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + ) + sequence_output = outputs.last_hidden_state # (b, l, d) + sequence_output = self.dropout(sequence_output) # (b, l, d) + logits = self.classifier(sequence_output) # (b, l, c) + + loss = None + if labels is not None: + labels = labels.to(logits.device) # (b, l) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) # () + + result = EsmTokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + s_max=outputs.s_max, + ) + return result if return_dict else result.to_tuple() diff --git a/src/fastplms/models/esm3/__init__.py b/src/fastplms/models/esm3/__init__.py new file mode 100644 index 0000000..3d466d0 --- /dev/null +++ b/src/fastplms/models/esm3/__init__.py @@ -0,0 +1,4 @@ +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model + + +__all__ = ["FastESM3Config", "FastESM3Model"] diff --git a/src/fastplms/models/esm3/modeling_esm3.py b/src/fastplms/models/esm3/modeling_esm3.py new file mode 100644 index 0000000..1ca57b1 --- /dev/null +++ b/src/fastplms/models/esm3/modeling_esm3.py @@ -0,0 +1,2526 @@ +"""Hugging Face-compatible ESM3 implementation. + +The production module is self-contained. The pinned Biohub repository is used +only by the reference adapter in the parity suite. +""" + +from __future__ import annotations + +import base64 +import functools +import hashlib +import io +import json +import math +import os +import shutil +import stat +import tempfile +import einops +import torch +import torch.nn as nn +import torch.nn.functional as F +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import ClassVar +from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo +from einops import rearrange +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.processors import TemplateProcessing +from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerFast +from transformers.modeling_outputs import ModelOutput + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + create_block_mask, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "_get_flex_attention_fn", + "create_block_mask", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +_SAVED_RUNTIME_SCHEMA_VERSION = 1 +_SAVED_RUNTIME_FILES = ( + "__init__.py", + "attention/__init__.py", + "attention/_core.py", + "attention/_kernel_lock.py", + "attention/interfaces.py", + "embeddings/__init__.py", + "embeddings/pooling.py", + "embeddings/runner.py", + "embeddings/storage.py", + "embeddings/types.py", + "models/__init__.py", + "models/esm3/__init__.py", + "models/esm3/modeling_esm3.py", + "models/ttt.py", + "models.toml", + "registry.py", + "runtime.py", +) +_MAX_SAVED_RUNTIME_FILE_BYTES = 1024 * 1024 +_MAX_SAVED_RUNTIME_TOTAL_BYTES = 4 * 1024 * 1024 +_MAX_SAVED_RUNTIME_ARCHIVE_BYTES = 2 * 1024 * 1024 + + +@contextmanager +def _temporary_eval(model: nn.Module): + """Temporarily disable training behavior without flattening mixed module states.""" + training_states = tuple((module, module.training) for module in model.modules()) + model.eval() + try: + yield + finally: + for module, training in training_states: + module.training = training + + +def _validate_saved_runtime_relative_path(value: str) -> PurePosixPath: + """Return one canonical, fixed-inventory runtime source path.""" + + relative = PurePosixPath(value) + if ( + not value + or "\\" in value + or relative.is_absolute() + or relative.as_posix() != value + or any(part in {"", ".", ".."} or ":" in part or "\0" in part for part in relative.parts) + ): + raise RuntimeError(f"Saved ESM3 runtime path is unsafe: {value!r}.") + return relative + + +def _read_saved_runtime_file(package_root: Path, relative: PurePosixPath) -> bytes: + """Read one allowlisted regular file without following a symlink.""" + + current = package_root + for index, part in enumerate(relative.parts): + current = current / part + try: + metadata = current.lstat() + except OSError as error: + raise RuntimeError( + f"Saved ESM3 runtime file is missing: {relative.as_posix()!r}." + ) from error + if stat.S_ISLNK(metadata.st_mode): + raise RuntimeError( + f"Saved ESM3 runtime path must not contain a symlink: {relative.as_posix()!r}." + ) + if index < len(relative.parts) - 1: + if not stat.S_ISDIR(metadata.st_mode): + raise RuntimeError( + f"Saved ESM3 runtime parent is not a directory: {relative.as_posix()!r}." + ) + continue + if not stat.S_ISREG(metadata.st_mode): + raise RuntimeError( + f"Saved ESM3 runtime entry is not a regular file: {relative.as_posix()!r}." + ) + if metadata.st_size > _MAX_SAVED_RUNTIME_FILE_BYTES: + raise RuntimeError( + f"Saved ESM3 runtime file exceeds its size limit: {relative.as_posix()!r}." + ) + before = metadata + + try: + with current.open("rb") as handle: + payload = handle.read(_MAX_SAVED_RUNTIME_FILE_BYTES + 1) + after = current.lstat() + except OSError as error: + raise RuntimeError( + f"Unable to read saved ESM3 runtime file: {relative.as_posix()!r}." + ) from error + identity_before = ( + before.st_dev, + before.st_ino, + before.st_size, + before.st_mtime_ns, + before.st_ctime_ns, + ) + identity_after = ( + after.st_dev, + after.st_ino, + after.st_size, + after.st_mtime_ns, + after.st_ctime_ns, + ) + if ( + stat.S_ISLNK(after.st_mode) + or not stat.S_ISREG(after.st_mode) + or identity_before != identity_after + or len(payload) != before.st_size + or len(payload) > _MAX_SAVED_RUNTIME_FILE_BYTES + ): + raise RuntimeError( + f"Saved ESM3 runtime file changed while it was validated: {relative.as_posix()!r}." + ) + return payload + + +def _saved_runtime_files(package_root: Path) -> dict[str, bytes]: + """Read exactly the fixed ESM3 runtime inventory into validated bytes.""" + + try: + root_metadata = package_root.lstat() + except OSError as error: + raise RuntimeError(f"Saved ESM3 runtime package is unavailable: {package_root}.") from error + if stat.S_ISLNK(root_metadata.st_mode) or not stat.S_ISDIR(root_metadata.st_mode): + raise RuntimeError("Saved ESM3 runtime package root must be a non-symlink directory.") + + files: dict[str, bytes] = {} + total_size = 0 + for value in _SAVED_RUNTIME_FILES: + relative = _validate_saved_runtime_relative_path(value) + payload = _read_saved_runtime_file(package_root, relative) + total_size += len(payload) + if total_size > _MAX_SAVED_RUNTIME_TOTAL_BYTES: + raise RuntimeError("Saved ESM3 runtime exceeds its total expanded size limit.") + files[relative.as_posix()] = payload + if len(files) != len(_SAVED_RUNTIME_FILES): + raise RuntimeError("Saved ESM3 runtime allowlist contains duplicate paths.") + return files + + +def _saved_runtime_manifest(files: dict[str, bytes]) -> dict[str, object]: + records = { + relative: { + "sha256": hashlib.sha256(payload).hexdigest(), + "size": len(payload), + } + for relative, payload in sorted(files.items()) + } + return { + "schema_version": _SAVED_RUNTIME_SCHEMA_VERSION, + "files": records, + "total_size": sum(record["size"] for record in records.values()), + } + + +def _saved_runtime_tree_hash(manifest: dict[str, object]) -> str: + files = manifest["files"] + if not isinstance(files, dict): + raise RuntimeError("Saved ESM3 runtime manifest files are invalid.") + digest = hashlib.sha256() + for relative, raw_record in sorted(files.items()): + if not isinstance(relative, str) or not isinstance(raw_record, dict): + raise RuntimeError("Saved ESM3 runtime manifest record is invalid.") + digest.update(relative.encode("utf-8")) + digest.update(b"\0") + digest.update(str(raw_record["size"]).encode("ascii")) + digest.update(b"\0") + digest.update(str(raw_record["sha256"]).encode("ascii")) + digest.update(b"\n") + return digest.hexdigest() + + +def _build_saved_runtime_archive( + package_root: Path, +) -> tuple[bytes, dict[str, object], str]: + """Build a deterministic archive directly from validated runtime bytes.""" + + files = _saved_runtime_files(package_root) + manifest = _saved_runtime_manifest(files) + tree_hash = _saved_runtime_tree_hash(manifest) + + buffer = io.BytesIO() + with ZipFile(buffer, mode="w", compression=ZIP_DEFLATED, compresslevel=9) as archive: + for relative, contents in sorted(files.items()): + archive_path = (PurePosixPath("fastplms") / relative).as_posix() + info = ZipInfo(archive_path, date_time=(1980, 1, 1, 0, 0, 0)) + info.create_system = 3 + info.compress_type = ZIP_DEFLATED + info.external_attr = 0o100644 << 16 + archive.writestr(info, contents, compress_type=ZIP_DEFLATED, compresslevel=9) + payload = buffer.getvalue() + if len(payload) > _MAX_SAVED_RUNTIME_ARCHIVE_BYTES: + raise RuntimeError("Saved ESM3 runtime archive exceeds its compressed size limit.") + return payload, manifest, tree_hash + + +def _render_saved_runtime_bundle( + archive: bytes, + manifest: dict[str, object], + tree_hash: str, +) -> tuple[str, bytes]: + archive_hash = hashlib.sha256(archive).hexdigest() + encoded = base64.b85encode(archive).decode("ascii") + chunks = (encoded[index : index + 100] for index in range(0, len(encoded), 100)) + manifest_source = json.dumps(manifest, indent=2, sort_keys=True, ensure_ascii=True) + lines = [ + '"""Deterministic embedded FastPLMs runtime for one saved ESM3 model."""', + "", + f'RUNTIME_HASH = "{archive_hash}"', + f'RUNTIME_TREE_HASH = "{tree_hash}"', + f"RUNTIME_MANIFEST = {manifest_source}", + "RUNTIME_DATA = (", + *(f" {chunk!r}," for chunk in chunks), + ")", + "", + ] + return archive_hash, "\n".join(lines).encode("utf-8") + + +def _render_saved_runtime_bridge(archive_hash: str, tree_hash: str) -> str: + """Render the fail-closed Transformers bridge for one runtime identity.""" + + lines = [ + '"""Bridge to the bundled FastPLMs ESM3 runtime."""', + "", + "import atexit", + "import base64", + "import hashlib", + "import importlib", + "import importlib.util", + "import stat", + "import sys", + "import tempfile", + "from io import BytesIO", + "from pathlib import Path, PurePosixPath", + "from zipfile import BadZipFile, ZIP_DEFLATED, ZipFile", + "", + "from .fastplms_bundle import (", + " RUNTIME_DATA,", + " RUNTIME_HASH,", + " RUNTIME_MANIFEST,", + " RUNTIME_TREE_HASH,", + ")", + "", + f'if RUNTIME_HASH != "{archive_hash}" or RUNTIME_TREE_HASH != "{tree_hash}":', + ' raise RuntimeError("FastPLMs runtime identity differs from the saved ESM3 bridge.")', + "", + f"_MAX_RUNTIME_FILE_BYTES = {_MAX_SAVED_RUNTIME_FILE_BYTES}", + f"_MAX_RUNTIME_TOTAL_BYTES = {_MAX_SAVED_RUNTIME_TOTAL_BYTES}", + f"_MAX_RUNTIME_ARCHIVE_BYTES = {_MAX_SAVED_RUNTIME_ARCHIVE_BYTES}", + "_MAX_RUNTIME_ENCODED_BYTES = (_MAX_RUNTIME_ARCHIVE_BYTES * 5 + 3) // 4", + "_EXPECTED_RUNTIME_FILES = (", + *(f" {relative!r}," for relative in _SAVED_RUNTIME_FILES), + ")", + "_RUNTIME_TEMPORARIES = []", + "", + "def _runtime_tree_hash(files):", + " digest = hashlib.sha256()", + " for relative, record in sorted(files.items()):", + ' digest.update(relative.encode("utf-8"))', + ' digest.update(b"\\0")', + ' digest.update(str(record["size"]).encode("ascii"))', + ' digest.update(b"\\0")', + ' digest.update(record["sha256"].encode("ascii"))', + ' digest.update(b"\\n")', + " return digest.hexdigest()", + "", + "def _validated_manifest():", + " if not isinstance(RUNTIME_MANIFEST, dict) or set(RUNTIME_MANIFEST) != {", + ' "schema_version",', + ' "files",', + ' "total_size",', + " }:", + ' raise RuntimeError("Embedded FastPLMs runtime manifest is invalid.")', + f' if RUNTIME_MANIFEST["schema_version"] != {_SAVED_RUNTIME_SCHEMA_VERSION}:', + ' raise RuntimeError("Embedded FastPLMs runtime manifest schema is unsupported.")', + ' raw_files = RUNTIME_MANIFEST["files"]', + " if not isinstance(raw_files, dict) or set(raw_files) != set(_EXPECTED_RUNTIME_FILES):", + ' raise RuntimeError("Embedded FastPLMs runtime inventory is invalid.")', + " files = {}", + " total_size = 0", + " for relative in _EXPECTED_RUNTIME_FILES:", + " record = raw_files[relative]", + ' if not isinstance(record, dict) or set(record) != {"sha256", "size"}:', + ' raise RuntimeError("Embedded FastPLMs runtime manifest record is invalid.")', + ' size = record["size"]', + ' file_hash = record["sha256"]', + " if (", + " isinstance(size, bool)", + " or not isinstance(size, int)", + " or size < 0", + " or size > _MAX_RUNTIME_FILE_BYTES", + " or not isinstance(file_hash, str)", + " or len(file_hash) != 64", + ' or any(character not in "0123456789abcdef" for character in file_hash)', + " ):", + ' raise RuntimeError("Embedded FastPLMs runtime manifest record is invalid.")', + ' files[relative] = {"sha256": file_hash, "size": size}', + " total_size += size", + " if total_size > _MAX_RUNTIME_TOTAL_BYTES:", + ' raise RuntimeError("Embedded FastPLMs runtime exceeds its size limit.")', + " if (", + ' isinstance(RUNTIME_MANIFEST["total_size"], bool)', + ' or RUNTIME_MANIFEST["total_size"] != total_size', + " ):", + ' raise RuntimeError("Embedded FastPLMs runtime total size is invalid.")', + " if _runtime_tree_hash(files) != RUNTIME_TREE_HASH:", + ' raise RuntimeError("Embedded FastPLMs runtime tree hash mismatch.")', + " return files", + "", + "_EXPECTED_MANIFEST = _validated_manifest()", + "", + "def _archive_relative_path(member):", + " name = member.filename", + " relative_archive = PurePosixPath(name)", + " parts = relative_archive.parts", + " if (", + ' not name or "\\\\" in name', + " or relative_archive.is_absolute()", + " or relative_archive.as_posix() != name", + " or len(parts) < 2", + ' or parts[0] != "fastplms"', + ' or any(part in {"", ".", ".."} or ":" in part or "\\0" in part for part in parts)', + " ):", + ' raise RuntimeError("Embedded FastPLMs archive has an unsafe path.")', + " relative = PurePosixPath(*parts[1:]).as_posix()", + " if relative not in _EXPECTED_MANIFEST:", + ' raise RuntimeError("Embedded FastPLMs archive inventory is unexpected.")', + " return relative", + "", + "def _validated_archive_files(payload):", + " if len(payload) > _MAX_RUNTIME_ARCHIVE_BYTES:", + ( + ' raise RuntimeError("Embedded FastPLMs archive exceeds its compressed ' + 'size limit.")' + ), + " try:", + " with ZipFile(BytesIO(payload)) as archive:", + " members = archive.infolist()", + " if archive.comment or len(members) != len(_EXPECTED_MANIFEST):", + ' raise RuntimeError("Embedded FastPLMs archive inventory is invalid.")', + " files = {}", + " total_size = 0", + " for member in members:", + " relative = _archive_relative_path(member)", + " if relative in files:", + ' raise RuntimeError("Embedded FastPLMs archive repeats a path.")', + " record = _EXPECTED_MANIFEST[relative]", + " if (", + " member.is_dir()", + " or member.flag_bits & 0x1", + " or member.compress_type != ZIP_DEFLATED", + " or member.create_system != 3", + " or member.external_attr >> 16 != 0o100644", + " or member.date_time != (1980, 1, 1, 0, 0, 0)", + " or member.extra", + " or member.comment", + ' or member.filename != f"fastplms/{relative}"', + ' or member.file_size != record["size"]', + " or member.file_size > _MAX_RUNTIME_FILE_BYTES", + " or member.compress_size > _MAX_RUNTIME_ARCHIVE_BYTES", + " ):", + ( + ' raise RuntimeError("Embedded FastPLMs archive member is not ' + 'canonical.")' + ), + ' with archive.open(member, mode="r") as handle:', + ' contents = handle.read(record["size"] + 1)', + " if (", + ' len(contents) != record["size"]', + ' or hashlib.sha256(contents).hexdigest() != record["sha256"]', + " ):", + ' raise RuntimeError("Embedded FastPLMs archive member hash mismatch.")', + " total_size += len(contents)", + " if total_size > _MAX_RUNTIME_TOTAL_BYTES:", + ( + ' raise RuntimeError("Embedded FastPLMs archive exceeds its size ' + 'limit.")' + ), + " files[relative] = contents", + " except RuntimeError:", + " raise", + " except (BadZipFile, KeyError, OSError, ValueError) as error:", + ' raise RuntimeError("Embedded FastPLMs archive is invalid.") from error', + " if set(files) != set(_EXPECTED_MANIFEST):", + ' raise RuntimeError("Embedded FastPLMs archive inventory is incomplete.")', + " return files", + "", + "def _read_runtime_file(package_root, relative):", + " current = package_root", + " parts = PurePosixPath(relative).parts", + " for index, part in enumerate(parts):", + " current = current / part", + " try:", + " metadata = current.lstat()", + " except OSError as error:", + ' raise RuntimeError(f"Runtime file is missing: {relative!r}.") from error', + " if stat.S_ISLNK(metadata.st_mode):", + ' raise RuntimeError(f"Runtime path contains a symlink: {relative!r}.")', + " if index < len(parts) - 1:", + " if not stat.S_ISDIR(metadata.st_mode):", + ' raise RuntimeError(f"Runtime parent is not a directory: {relative!r}.")', + " continue", + " if not stat.S_ISREG(metadata.st_mode):", + ' raise RuntimeError(f"Runtime entry is not a regular file: {relative!r}.")', + " if metadata.st_size > _MAX_RUNTIME_FILE_BYTES:", + ' raise RuntimeError(f"Runtime file exceeds its size limit: {relative!r}.")', + " before = metadata", + " try:", + ' with current.open("rb") as handle:', + " contents = handle.read(_MAX_RUNTIME_FILE_BYTES + 1)", + " after = current.lstat()", + " except OSError as error:", + ' raise RuntimeError(f"Unable to read runtime file: {relative!r}.") from error', + " before_identity = (", + " before.st_dev,", + " before.st_ino,", + " before.st_size,", + " before.st_mtime_ns,", + " before.st_ctime_ns,", + " )", + " after_identity = (", + " after.st_dev,", + " after.st_ino,", + " after.st_size,", + " after.st_mtime_ns,", + " after.st_ctime_ns,", + " )", + " if (", + " stat.S_ISLNK(after.st_mode)", + " or not stat.S_ISREG(after.st_mode)", + " or before_identity != after_identity", + " or len(contents) != before.st_size", + " or len(contents) > _MAX_RUNTIME_FILE_BYTES", + " ):", + ' raise RuntimeError(f"Runtime file changed while validated: {relative!r}.")', + " return contents", + "", + "def _runtime_file_manifest(package_root):", + " try:", + " root_metadata = package_root.lstat()", + " except OSError as error:", + ' raise RuntimeError("Runtime package root is unavailable.") from error', + " if stat.S_ISLNK(root_metadata.st_mode) or not stat.S_ISDIR(root_metadata.st_mode):", + ' raise RuntimeError("Runtime package root must be a non-symlink directory.")', + " files = {}", + " total_size = 0", + " for relative in _EXPECTED_RUNTIME_FILES:", + " contents = _read_runtime_file(package_root, relative)", + " files[relative] = {", + ' "sha256": hashlib.sha256(contents).hexdigest(),', + ' "size": len(contents),', + " }", + " total_size += len(contents)", + " if total_size > _MAX_RUNTIME_TOTAL_BYTES:", + ' raise RuntimeError("Runtime package exceeds its total size limit.")', + " return files", + "", + "def _cleanup_runtime_temporaries():", + " while _RUNTIME_TEMPORARIES:", + " _RUNTIME_TEMPORARIES.pop().cleanup()", + "", + "atexit.register(_cleanup_runtime_temporaries)", + "", + "def _ensure_runtime():", + " if (", + " not isinstance(RUNTIME_DATA, tuple)", + " or not RUNTIME_DATA", + " or any(not isinstance(chunk, str) for chunk in RUNTIME_DATA)", + " ):", + ' raise RuntimeError("Embedded FastPLMs runtime data is invalid.")', + ' encoded = "".join(RUNTIME_DATA)', + " if len(encoded) > _MAX_RUNTIME_ENCODED_BYTES:", + ' raise RuntimeError("Embedded FastPLMs runtime data exceeds its size limit.")', + " try:", + ' payload = base64.b85decode(encoded.encode("ascii"))', + " except (UnicodeEncodeError, ValueError) as error:", + ' raise RuntimeError("Embedded FastPLMs runtime data is invalid.") from error', + " if hashlib.sha256(payload).hexdigest() != RUNTIME_HASH:", + ' raise RuntimeError("Embedded FastPLMs runtime hash mismatch.")', + " files = _validated_archive_files(payload)", + ' temporary = tempfile.TemporaryDirectory(prefix="fastplms-esm3-runtime-")', + " try:", + " runtime_root = Path(temporary.name).resolve()", + " module_root = Path(__file__).resolve().parent", + " if runtime_root == module_root or module_root in runtime_root.parents:", + ( + ' raise RuntimeError("FastPLMs runtime temporary must be outside the saved ' + 'model.")' + ), + ' package_root = runtime_root / "fastplms"', + " for relative in _EXPECTED_RUNTIME_FILES:", + " target = package_root.joinpath(*PurePosixPath(relative).parts)", + " target.parent.mkdir(parents=True, exist_ok=True)", + ' with target.open("xb") as handle:', + " handle.write(files[relative])", + " actual = _runtime_file_manifest(package_root)", + " if (", + " actual != _EXPECTED_MANIFEST", + " or _runtime_tree_hash(actual) != RUNTIME_TREE_HASH", + " ):", + ' raise RuntimeError("Extracted FastPLMs runtime identity mismatch.")', + " except BaseException:", + " temporary.cleanup()", + " raise", + " return package_root, temporary", + "", + "def _verify_loaded_runtime(package):", + ' package_file = getattr(package, "__file__", None)', + " if not isinstance(package_file, str) or not package_file:", + " raise RuntimeError(", + ' "Loaded FastPLMs version/runtime mismatch: source path is unavailable."', + " )", + " package_root = Path(package_file).absolute().parent", + " try:", + " actual = _runtime_file_manifest(package_root)", + " except RuntimeError as error:", + " raise RuntimeError(", + ' "Loaded FastPLMs version/runtime mismatch: sources cannot be verified."', + " ) from error", + " if actual != _EXPECTED_MANIFEST or _runtime_tree_hash(actual) != RUNTIME_TREE_HASH:", + " mismatch = next(", + " (", + " relative", + " for relative in _EXPECTED_RUNTIME_FILES", + " if actual.get(relative) != _EXPECTED_MANIFEST[relative]", + " ),", + ' "unknown",', + " )", + " raise RuntimeError(", + ' f"Loaded FastPLMs version/runtime mismatch at {mismatch!r}. "', + ' "Install the matching FastPLMs release or use a separate Python process."', + " )", + " package.__fastplms_saved_runtime_tree_hash__ = RUNTIME_TREE_HASH", + " package.__fastplms_saved_runtime_manifest__ = _EXPECTED_MANIFEST", + " return package", + "", + "def _install_runtime():", + ' installed = sys.modules.get("fastplms")', + " if installed is not None:", + " return _verify_loaded_runtime(installed)", + ' stale = sorted(name for name in sys.modules if name.startswith("fastplms."))', + " if stale:", + " raise RuntimeError(", + ' "Loaded FastPLMs version/runtime mismatch: orphaned submodules exist."', + " )", + " package_root, temporary = _ensure_runtime()", + " spec = importlib.util.spec_from_file_location(", + ' "fastplms",', + ' package_root / "__init__.py",', + " submodule_search_locations=[str(package_root)],", + " )", + " if spec is None or spec.loader is None:", + " temporary.cleanup()", + ' raise ImportError("Unable to load the embedded FastPLMs runtime.")', + " package = importlib.util.module_from_spec(spec)", + ' sys.modules["fastplms"] = package', + " previous = sys.dont_write_bytecode", + " sys.dont_write_bytecode = True", + " try:", + " spec.loader.exec_module(package)", + " except BaseException:", + ' sys.modules.pop("fastplms", None)', + " temporary.cleanup()", + " raise", + " finally:", + " sys.dont_write_bytecode = previous", + " _RUNTIME_TEMPORARIES.append(temporary)", + " package.__fastplms_saved_runtime_tree_hash__ = RUNTIME_TREE_HASH", + " package.__fastplms_saved_runtime_manifest__ = _EXPECTED_MANIFEST", + " package.__fastplms_saved_runtime_temporary__ = temporary", + " return package", + "", + "def _import_without_bytecode(module_name):", + " previous = sys.dont_write_bytecode", + " sys.dont_write_bytecode = True", + " try:", + " return importlib.import_module(module_name)", + " finally:", + " sys.dont_write_bytecode = previous", + "", + "_install_runtime()", + '_modeling = _import_without_bytecode("fastplms.models.esm3.modeling_esm3")', + "FastESM3Config = _modeling.FastESM3Config", + "FastESM3Model = _modeling.FastESM3Model", + "", + ] + return "\n".join(lines) + + +def _replace_saved_runtime_file(path: Path, payload: bytes) -> None: + """Atomically replace one generated runtime file without following a symlink.""" + + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + if temporary.exists(): + temporary.unlink() + + +def _remove_old_saved_runtime_path(path: Path) -> None: + try: + metadata = path.lstat() + except FileNotFoundError: + return + if stat.S_ISDIR(metadata.st_mode) and not stat.S_ISLNK(metadata.st_mode): + shutil.rmtree(path) + return + path.unlink() + + +def _clean_old_saved_runtime(save_directory: Path) -> None: + _remove_old_saved_runtime_path(save_directory / "fastplms") + for pattern in ("_fastplms_runtime_*", "._fastplms_runtime_*"): + for candidate in save_directory.glob(pattern): + _remove_old_saved_runtime_path(candidate) + + +def _validate_saved_runtime_destination(save_directory: Path) -> None: + if save_directory.is_symlink(): + raise ValueError("ESM3 save directory must not be a symlink.") + package_source = Path(__file__).resolve().parents[2] + destination = save_directory.resolve(strict=False) + if destination == package_source or package_source in destination.parents: + raise ValueError("ESM3 save directory must be outside the FastPLMs source package.") + for name in ("config.json", "fastplms_bundle.py", "modeling_fastplms.py"): + if (save_directory / name).is_symlink(): + raise ValueError(f"ESM3 generated save path must not be a symlink: {name!r}.") + + +def _write_saved_runtime( + save_directory: Path, + prepared_runtime: tuple[bytes, dict[str, object], str] | None = None, +) -> None: + """Make one ESM3 ``save_pretrained`` directory independently loadable.""" + + _validate_saved_runtime_destination(save_directory) + if prepared_runtime is None: + package_source = Path(__file__).resolve().parents[2] + prepared_runtime = _build_saved_runtime_archive(package_source) + archive, manifest, tree_hash = prepared_runtime + archive_hash, bundle = _render_saved_runtime_bundle(archive, manifest, tree_hash) + bridge = _render_saved_runtime_bridge(archive_hash, tree_hash).encode("utf-8") + + _clean_old_saved_runtime(save_directory) + _replace_saved_runtime_file(save_directory / "fastplms_bundle.py", bundle) + _replace_saved_runtime_file(save_directory / "modeling_fastplms.py", bridge) + + config_path = save_directory / "config.json" + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("Saved ESM3 config.json is missing or invalid.") from error + if not isinstance(config, dict): + raise RuntimeError("Saved ESM3 config.json must contain a JSON object.") + config["auto_map"] = { + "AutoConfig": "modeling_fastplms.FastESM3Config", + "AutoModel": "modeling_fastplms.FastESM3Model", + } + config_payload = (json.dumps(config, indent=2, sort_keys=True) + "\n").encode("utf-8") + _replace_saved_runtime_file(config_path, config_payload) + + +ESM3_OPEN_SMALL = "esm3_sm_open_v1" +ESM3_OPEN_SMALL_ALIASES = { + "ESM3_small", + "esm3_small", + "esm3_sm_open_v1", + "esm3-open-2024-03", + "esm3-sm-open-v1", + "esm3-open", +} + +SEQUENCE_BOS_TOKEN = 0 +SEQUENCE_PAD_TOKEN = 1 +SEQUENCE_EOS_TOKEN = 2 +SEQUENCE_CHAINBREAK_TOKEN = 31 +SEQUENCE_MASK_TOKEN = 32 + +VQVAE_CODEBOOK_SIZE = 4096 +STRUCTURE_MASK_TOKEN = VQVAE_CODEBOOK_SIZE +STRUCTURE_EOS_TOKEN = VQVAE_CODEBOOK_SIZE + 1 +STRUCTURE_BOS_TOKEN = VQVAE_CODEBOOK_SIZE + 2 +STRUCTURE_PAD_TOKEN = VQVAE_CODEBOOK_SIZE + 3 +STRUCTURE_CHAINBREAK_TOKEN = VQVAE_CODEBOOK_SIZE + 4 + +SASA_PAD_TOKEN = 0 +SS8_PAD_TOKEN = 0 +INTERPRO_PAD_TOKEN = 0 +RESIDUE_PAD_TOKEN = 0 +MAX_RESIDUE_ANNOTATIONS = 16 +FUNCTION_TOKENS_DEPTH = 8 + +SEQUENCE_VOCAB = [ + "", + "", + "", + "", + "L", + "A", + "G", + "V", + "S", + "E", + "R", + "T", + "I", + "D", + "P", + "K", + "Q", + "N", + "F", + "Y", + "M", + "H", + "W", + "C", + "X", + "B", + "U", + "Z", + "O", + ".", + "-", + "|", + "", +] + +_SUPPORTED_ATTENTION_BACKENDS = ("eager", "sdpa", "flex_attention") + + +class FastESM3Config(PretrainedConfig): + model_type = "fast_esm3" + + def __init__( + self, + vocab_size: int = 64, + hidden_size: int = 1536, + num_attention_heads: int = 24, + num_vector_heads: int = 256, + num_hidden_layers: int = 48, + initializer_range: float = 0.02, + attn_backend: str | None = None, + model_name: str = ESM3_OPEN_SMALL, + **kwargs, + ): + super().__init__(**kwargs) + if hidden_size <= 0: + raise ValueError(f"hidden_size must be positive, got {hidden_size}.") + if num_attention_heads <= 0: + raise ValueError(f"num_attention_heads must be positive, got {num_attention_heads}.") + if hidden_size % FUNCTION_TOKENS_DEPTH != 0: + raise ValueError( + f"hidden_size must be divisible by {FUNCTION_TOKENS_DEPTH}, got {hidden_size}." + ) + if hidden_size % num_attention_heads != 0: + raise ValueError( + "hidden_size must be divisible by num_attention_heads, " + f"got hidden_size={hidden_size} and num_attention_heads={num_attention_heads}." + ) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_vector_heads = num_vector_heads + self.num_hidden_layers = num_hidden_layers + self.initializer_range = initializer_range + self.attn_backend = attn_backend + self.model_name = _resolve_esm3_checkpoint_key(model_name) + self.tie_word_embeddings = False + + +@dataclass +class FastESM3Output(ModelOutput): + loss: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + logits: torch.Tensor | None = None + sequence_logits: torch.Tensor | None = None + structure_logits: torch.Tensor | None = None + secondary_structure_logits: torch.Tensor | None = None + sasa_logits: torch.Tensor | None = None + function_logits: torch.Tensor | None = None + residue_logits: torch.Tensor | None = None + embeddings: torch.Tensor | None = None + + +@dataclass(frozen=True) +class FastESM3GenerationConfig: + """Sequence-track sampling controls for the local ESM3 generation API.""" + + num_steps: int | None = None + temperature: float = 1.0 + seed: int | None = None + + +class EsmSequenceTokenizer(PreTrainedTokenizerFast): + model_input_names: ClassVar[list[str]] = ["input_ids", "attention_mask"] + + def __init__( + self, + unk_token: str = "", + cls_token: str = "", + pad_token: str = "", + mask_token: str = "", + eos_token: str = "", + chain_break_token: str = "|", + **kwargs, + ): + token_to_id = {token: index for index, token in enumerate(SEQUENCE_VOCAB)} + bpe = BPE(token_to_id, merges=[], unk_token=unk_token) + tokenizer = Tokenizer(bpe) + special_tokens = [ + cls_token, + pad_token, + mask_token, + eos_token, + chain_break_token, + ] + self.cb_token = chain_break_token + tokenizer.add_special_tokens(special_tokens) + tokenizer.post_processor = TemplateProcessing( + single=" $A ", + pair=":0 $A:0 :0 $B:1 :1", + special_tokens=[ + ("", tokenizer.token_to_id("")), + ("", tokenizer.token_to_id("")), + ], + ) + super().__init__( + tokenizer_object=tokenizer, + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + additional_special_tokens=[chain_break_token], + **kwargs, + ) + + @property + def bos_token(self) -> str: + return self.cls_token + + @property + def bos_token_id(self) -> int: + return self.cls_token_id + + @property + def chain_break_token(self) -> str: + return self.cb_token + + @property + def chain_break_token_id(self) -> int: + token_id = self.convert_tokens_to_ids(self.chain_break_token) + if not isinstance(token_id, int): + raise RuntimeError("ESM3 chain-break token did not resolve to one token id.") + return token_id + + @property + def all_token_ids(self) -> list[int]: + return list(range(self.vocab_size)) + + @property + def special_token_ids(self) -> list[int]: + return self.all_special_ids + + +def rbf(values: torch.Tensor, v_min: float, v_max: float, n_bins: int = 16) -> torch.Tensor: + # values: (...) + centers = torch.linspace( + v_min, + v_max, + n_bins, + device=values.device, + dtype=values.dtype, + ) + centers = centers.view([1] * len(values.shape) + [-1]) # (..., n) + std = (v_max - v_min) / n_bins + z = (values.unsqueeze(-1) - centers) / std # (..., n) + return torch.exp(-(z**2)) + + +def RegressionHead( + d_model: int, + output_dim: int, + hidden_dim: int | None = None, +) -> nn.Module: + hidden_dim = hidden_dim if hidden_dim is not None else d_model + return nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, output_dim), + ) + + +def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange( + torch.stack((-x2, x1), dim=-1), + "... d two -> ... (d two)", + two=2, + ) + + +def apply_rotary_emb_torch( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = False, +) -> torch.Tensor: + ro_dim = cos.shape[-1] * 2 + if ro_dim > x.shape[-1]: + raise ValueError( + "Rotary embedding width cannot exceed the input head dimension; " + f"got rotary width {ro_dim} and head dimension {x.shape[-1]}." + ) + seqlen = x.size(1) + cos = cos[:seqlen] + sin = sin[:seqlen] + cos = einops.repeat(cos, "s d -> s 1 (2 d)") + sin = einops.repeat(sin, "s d -> s 1 (2 d)") + return torch.cat( + [ + x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, + x[..., ro_dim:], + ], + dim=-1, + ) + + +class RotaryEmbedding(nn.Module): + def __init__( + self, + dim: int, + base: float = 10000.0, + interleaved: bool = False, + scale_base: float | None = None, + scaling_factor: float = 1.0, + pos_idx_in_fp32: bool = True, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.dim = dim + self.base = float(base) + self.pos_idx_in_fp32 = pos_idx_in_fp32 + self.interleaved = interleaved + self.scale_base = scale_base + self.scaling_factor = scaling_factor + self.device = device + self._seq_len_cached = 0 + self._cos_cached = None + self._sin_cached = None + self.reset_parameters() + + def reset_parameters(self) -> None: + inv_freq = self._compute_inv_freq(self.device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + arange = torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32) + scale = ( + (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None + ) + self.register_buffer("scale", scale) + + def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor: + return 1 / ( + self.base + ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) + ) + + def _update_cos_sin_cache( + self, + seqlen: int, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + if ( + seqlen > self._seq_len_cached + or self._cos_cached is None + or self._cos_cached.device != device + or self._cos_cached.dtype != dtype + or (self.training and self._cos_cached.is_inference()) + ): + self._seq_len_cached = seqlen + # ``inv_freq`` is non-persistent and may have been materialized + # without values after Transformers constructs this module on the + # meta device. Recreate it deterministically on the first forward. + self.inv_freq = self._compute_inv_freq(device) + if self.pos_idx_in_fp32: + t = torch.arange(seqlen, device=device, dtype=torch.float32) # (l,) + t /= self.scaling_factor + inv_freq = self.inv_freq + else: + t = torch.arange( + seqlen, device=device, dtype=self.inv_freq.dtype + ) # (l,) + t /= self.scaling_factor + inv_freq = self.inv_freq + freqs = torch.outer(t, inv_freq) # (l, d / 2) + + if self.scale is None: + self._cos_cached = torch.cos(freqs).to(dtype) + self._sin_cached = torch.sin(freqs).to(dtype) + else: + raise NotImplementedError("Scaled rotary embeddings are not used by ESM3.") + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + seqlen_offset: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, l, h, d) + self._update_cos_sin_cache( + q.shape[1] + seqlen_offset, + device=q.device, + dtype=q.dtype, + ) + if self._cos_cached is None or self._sin_cached is None: + raise RuntimeError( + "ESM3 rotary cache initialization did not produce sine/cosine tables." + ) + return ( + apply_rotary_emb_torch( + q, + self._cos_cached[seqlen_offset:], + self._sin_cached[seqlen_offset:], + self.interleaved, + ), + apply_rotary_emb_torch( + k, + self._cos_cached[seqlen_offset:], + self._sin_cached[seqlen_offset:], + self.interleaved, + ), + ) + + +def fp32_autocast_context(device_type: str): + if device_type == "cuda": + return torch.autocast(device_type="cuda", enabled=False) + return torch.autocast(device_type=device_type, enabled=False) + + +class RotationMatrix: + def __init__(self, rots: torch.Tensor) -> None: + if rots.ndim >= 1 and rots.shape[-1] == 9: + rots = rots.unflatten(-1, (3, 3)) + if rots.ndim < 2 or tuple(rots.shape[-2:]) != (3, 3): + raise ValueError( + "Rotation matrices must have trailing shape (3, 3) or flattened " + f"shape (9,); got {tuple(rots.shape)}." + ) + self._rots = rots.to(torch.float32) + + @classmethod + def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationMatrix: + rots = torch.eye(3, **tensor_kwargs) + rots = rots.view(*[1 for _ in range(len(shape))], 3, 3) + rots = rots.expand(*shape, -1, -1) + return cls(rots) + + def __getitem__(self, idx) -> RotationMatrix: + indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) + return RotationMatrix(self._rots[(*indices, slice(None), slice(None))]) + + @property + def shape(self) -> torch.Size: + return self._rots.shape[:-2] + + @property + def tensor(self) -> torch.Tensor: + return self._rots.flatten(-2) + + @property + def device(self) -> torch.device: + return self._rots.device + + def as_matrix(self) -> RotationMatrix: + return self + + def apply(self, p: torch.Tensor) -> torch.Tensor: + with fp32_autocast_context(self.device.type): + p = p.to(self._rots.dtype) + if self._rots.shape[-3] == 1: + return p @ self._rots.transpose(-1, -2).squeeze(-3) + return torch.einsum("...ij,...j", self._rots, p) + + def invert(self) -> RotationMatrix: + return RotationMatrix(self._rots.transpose(-1, -2)) + + @staticmethod + def from_graham_schmidt( + x_axis: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-12, + ) -> RotationMatrix: + with fp32_autocast_context(x_axis.device.type): + e1 = xy_plane + denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps) + x_axis = x_axis / denom + dot = (x_axis * e1).sum(dim=-1, keepdim=True) + e1 = e1 - x_axis * dot + denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps) + e1 = e1 / denom + e2 = torch.cross(x_axis, e1, dim=-1) + return RotationMatrix(torch.stack([x_axis, e1, e2], dim=-1)) + + +@dataclass(frozen=True) +class Affine3D: + trans: torch.Tensor + rot: RotationMatrix + + def __post_init__(self) -> None: + if self.trans.ndim < 1 or self.trans.shape[-1] != 3: + raise ValueError( + "Affine translations must have trailing dimension 3; " + f"got {tuple(self.trans.shape)}." + ) + if self.trans.shape[:-1] != self.rot.shape: + raise ValueError( + "Affine translation and rotation batch shapes must match; " + f"got {tuple(self.trans.shape[:-1])} and {tuple(self.rot.shape)}." + ) + + def __getitem__(self, idx) -> Affine3D: + indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx) + return Affine3D( + trans=self.trans[(*indices, slice(None))], + rot=self.rot[idx], + ) + + @property + def shape(self) -> torch.Size: + return self.trans.shape[:-1] + + @property + def dtype(self) -> torch.dtype: + return self.trans.dtype + + @property + def device(self) -> torch.device: + return self.trans.device + + @property + def tensor(self) -> torch.Tensor: + return torch.cat([self.rot.tensor, self.trans], dim=-1) + + def as_matrix(self) -> Affine3D: + return Affine3D(trans=self.trans, rot=self.rot.as_matrix()) + + def apply(self, p: torch.Tensor) -> torch.Tensor: + return self.rot.apply(p) + self.trans + + @staticmethod + def from_tensor(t: torch.Tensor) -> Affine3D: + match t.shape[-1]: + case 12: + trans = t[..., -3:] + rot = RotationMatrix(t[..., :-3].unflatten(-1, (3, 3))) + case _: + raise RuntimeError( + f"Cannot detect rotation format from {t.shape[-1] - 3}-d flat vector" + ) + return Affine3D(trans, rot) + + @staticmethod + def from_graham_schmidt( + neg_x_axis: torch.Tensor, + origin: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-10, + ) -> Affine3D: + x_axis = origin - neg_x_axis + xy_plane = xy_plane - origin + return Affine3D( + trans=origin, + rot=RotationMatrix.from_graham_schmidt(x_axis, xy_plane, eps), + ) + + +def build_affine3d_from_coordinates(coords: torch.Tensor) -> tuple[Affine3D, torch.Tensor]: + # coords: (b, l, 3, 3) + max_supported_distance = 1e6 + coord_mask = torch.all( + torch.all(torch.isfinite(coords) & (coords < max_supported_distance), dim=-1), + dim=-1, + ) # (b, l) + + def atom3_to_backbone_affine(bb_positions: torch.Tensor) -> Affine3D: + n_atom, ca_atom, c_atom = bb_positions.unbind(dim=-2) + return Affine3D.from_graham_schmidt(c_atom, ca_atom, n_atom) + + coords = coords.clone().float() + coords[~coord_mask] = 0 + average_per_n_ca_c = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / ( + coord_mask.sum(-1)[..., None, None] + 1e-8 + ) # (b, 3, 3) + affine_from_average = atom3_to_backbone_affine(average_per_n_ca_c.float()).as_matrix() + + batch_size, seq_len, _, _ = coords.shape + affine_rot_mats = affine_from_average.rot.tensor[..., None, :].expand( + batch_size, + seq_len, + 9, + ) + affine_trans = affine_from_average.trans[..., None, :].expand(batch_size, seq_len, 3) + identity_rot = RotationMatrix.identity( + (batch_size, seq_len), + dtype=torch.float32, + device=coords.device, + requires_grad=False, + ) + affine_rot_mats = affine_rot_mats.where( + coord_mask.any(-1)[..., None, None], + identity_rot.tensor, + ) + black_hole_affine = Affine3D(affine_trans, RotationMatrix(affine_rot_mats)) + + affine = atom3_to_backbone_affine(coords.float()) + affine = Affine3D.from_tensor( + affine.tensor.where(coord_mask[..., None], black_hole_affine.tensor) + ) + return affine, coord_mask + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + bias: bool = False, + qk_layernorm: bool = True, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.d_model = d_model + self.n_heads = n_heads + self.d_head = self.d_model // self.n_heads + self.scale = self.d_head**-0.5 + self.attn_backend = resolve_attention_backend(attn_backend) + self.layernorm_qkv = nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear(d_model, d_model * 3, bias=bias), + ) + self.out_proj = nn.Linear(d_model, d_model, bias=bias) + if qk_layernorm: + self.q_ln = nn.LayerNorm(d_model, bias=bias) + self.k_ln = nn.LayerNorm(d_model, bias=bias) + else: + self.q_ln = nn.Identity() + self.k_ln = nn.Identity() + self.rotary = RotaryEmbedding(d_model // n_heads) + + def _apply_rotary( + self, + q: torch.Tensor, + k: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, l, d) + q = q.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + k = k.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + q, k = self.rotary(q, k) + q = q.flatten(-2, -1) + k = k.flatten(-2, -1) + return q, k + + def forward( + self, + x: torch.Tensor, + seq_id: torch.Tensor | None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # x: (b, l, d); seq_id, attention_mask: (b, l) + qkv = self.layernorm_qkv(x) # (b, l, 3 * d) + query, key, value = torch.chunk(qkv, 3, dim=-1) + query = self.q_ln(query).to(query.dtype) + key = self.k_ln(key).to(query.dtype) + query, key = self._apply_rotary(query, key) + + reshaper = functools.partial( + einops.rearrange, + pattern="b s (h d) -> b h s d", + h=self.n_heads, + ) + query, key, value = map(reshaper, (query, key, value)) # each (b, h, l, d_h) + + mask = None + if seq_id is not None: + mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1) + if attention_mask is not None: + key_padding_mask = attention_mask[:, None, None, :] + mask = key_padding_mask if mask is None else mask & key_padding_mask + + if effective_backend is None: + effective_backend = resolve_attention_backend_for_call( + self.attn_backend, + output_attentions=output_attentions, + ) + if output_attentions or effective_backend == AttentionBackend.EAGER: + attn_scores = ( + torch.einsum("bhld,bhsd->bhls", query, key) * self.scale + ) # (b, h, l, l) + if mask is not None: + attn_scores = attn_scores.masked_fill( + ~mask, + torch.finfo(attn_scores.dtype).min, + ) + attn_weights = torch.softmax(attn_scores, dim=-1) + if mask is not None: + attn_weights = attn_weights.masked_fill(~mask, 0.0) + context = torch.einsum( + "bhls,bhsd->bhld", attn_weights, value + ) # (b, h, l, d_h) + if not output_attentions: + attn_weights = None + else: + attn_weights = None + if effective_backend == AttentionBackend.FLEX: + block_mask = self._create_flex_block_mask(seq_id, attention_mask, query) + if seq_id is not None and attention_mask is not None: + mask_semantics = "sequence_id_and_padding" + elif seq_id is not None: + mask_semantics = "sequence_id_equality" + elif attention_mask is not None: + mask_semantics = "padding" + else: + mask_semantics = "dense" + fn = _get_flex_attention_fn( + device=query.device, + dtype=query.dtype, + shape=tuple(query.shape), + sequence_lengths=None, + mask_semantics=mask_semantics, + ) + if fn is None: + raise RuntimeError("Flex Attention is not available in this environment.") + context = fn( + query, + key, + value, + block_mask=block_mask, + scale=self.scale, + ) + elif effective_backend == AttentionBackend.SDPA: + context = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + scale=self.scale, + ) + else: + raise RuntimeError(f"Unsupported resolved ESM3 backend: {effective_backend}") + + if mask is not None: + context = context.masked_fill(~mask.any(dim=-1, keepdim=True), 0.0) + context = einops.rearrange(context, "b h s d -> b s (h d)") # (b, l, d) + return self.out_proj(context), attn_weights + + @staticmethod + def _create_flex_block_mask( + seq_id: torch.Tensor | None, + attention_mask: torch.Tensor | None, + query: torch.Tensor, + ) -> BlockMask | None: + if seq_id is None and attention_mask is None: + return None + if create_block_mask is None: + raise RuntimeError( + "Flex Attention requested but torch.create_block_mask is unavailable." + ) + batch_size, _, seq_len, _ = query.shape + + def mask_mod(batch_idx, _head_idx, q_idx, kv_idx): + if seq_id is None: + return attention_mask[batch_idx, kv_idx] + allowed = seq_id[batch_idx, q_idx] == seq_id[batch_idx, kv_idx] + if attention_mask is not None: + allowed = allowed & attention_mask[batch_idx, kv_idx] + return allowed + + return create_block_mask( + mask_mod, + batch_size, + 1, + seq_len, + seq_len, + device=query.device, + ) + + +class GeometricReasoningOriginalImpl(nn.Module): + def __init__( + self, + c_s: int, + v_heads: int, + num_vector_messages: int = 1, + mask_and_zero_frameless: bool = True, + bias: bool = False, + ): + super().__init__() + self.c_s = c_s + self.v_heads = v_heads + self.num_vector_messages = num_vector_messages + self.mask_and_zero_frameless = mask_and_zero_frameless + + coordinate_width = 3 + vector_channels = coordinate_width * v_heads + projection_width = vector_channels * (4 + num_vector_messages) + output_width = vector_channels * num_vector_messages + self.s_norm = nn.LayerNorm(c_s, bias=bias) + self.proj = nn.Linear(c_s, projection_width, bias=bias) + self.out_proj = nn.Linear(output_width, c_s, bias=bias) + self.distance_scale_per_head = nn.Parameter(torch.zeros(v_heads)) + self.rotation_scale_per_head = nn.Parameter(torch.zeros(v_heads)) + + def forward( + self, + s: torch.Tensor, + affine: Affine3D, + affine_mask: torch.Tensor, + sequence_id: torch.Tensor | None, + chain_id: torch.Tensor, + ) -> torch.Tensor: + if sequence_id is None: + sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64) + attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2) + attn_bias = attn_bias.unsqueeze(1).float() + attn_bias = attn_bias.masked_fill( + ~affine_mask[:, None, None, :], + torch.finfo(attn_bias.dtype).min, + ) + chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2) + attn_bias = attn_bias.masked_fill( + chain_id_mask.unsqueeze(1), + torch.finfo(s.dtype).min, + ) + + ns = self.s_norm(s) + vec_rot, vec_dist = self.proj(ns).split( + [ + self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages, + self.v_heads * 2 * 3, + ], + dim=-1, + ) + + query_rot, key_rot, value = ( + affine.rot[..., None] + .apply(rearrange(vec_rot, "... (h c) -> ... h c", c=3)) + .split( + [self.v_heads, self.v_heads, self.v_heads * self.num_vector_messages], + dim=-2, + ) + ) + query_dist, key_dist = ( + affine[..., None] + .apply(rearrange(vec_dist, "... (h c) -> ... h c", c=3)) + .chunk(2, dim=-2) + ) + + query_dist = rearrange(query_dist, "b s h d -> b h s 1 d") + key_dist = rearrange(key_dist, "b s h d -> b h 1 s d") + query_rot = rearrange(query_rot, "b s h d -> b h s d") + key_rot = rearrange(key_rot, "b s h d -> b h d s") + value = rearrange( + value, + "b s (h m) d -> b h s (m d)", + m=self.num_vector_messages, + ) + + distance_term = (query_dist - key_dist).norm(dim=-1) / math.sqrt(3) + rotation_term = query_rot.matmul(key_rot) / math.sqrt(3) + distance_term_weight = rearrange( + F.softplus(self.distance_scale_per_head), + "h -> h 1 1", + ) + rotation_term_weight = rearrange( + F.softplus(self.rotation_scale_per_head), + "h -> h 1 1", + ) + attn_weight = rotation_term * rotation_term_weight - distance_term * distance_term_weight + + s_q = attn_weight.size(2) + s_k = attn_weight.size(3) + offset_q = max(0, attn_bias.size(2) - s_q) + offset_k = max(0, attn_bias.size(3) - s_k) + attn_bias = attn_bias[:, :, offset_q:, offset_k:] + attn_weight = torch.softmax(attn_weight + attn_bias, dim=-1) + + attn_out = attn_weight.matmul(value) + attn_out = ( + affine.rot[..., None] + .invert() + .apply( + rearrange( + attn_out, + "b h s (m d) -> b s (h m) d", + m=self.num_vector_messages, + ) + ) + ) + attn_out = rearrange( + attn_out, + "b s (h m) d -> b s (h m d)", + m=self.num_vector_messages, + ) + if self.mask_and_zero_frameless: + attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0) + attn_out = attn_out.to(self.out_proj.weight.dtype) + return self.out_proj(attn_out) + + +def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int: + return int(((expansion_ratio * d_model) + 255) // 256 * 256) + + +class SwiGLU(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return F.silu(x1) * x2 + + +def swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module: + return nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear( + d_model, + swiglu_correction_fn(expansion_ratio, d_model) * 2, + bias=bias, + ), + SwiGLU(), + nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=bias), + ) + + +def gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module: + hidden_dim = int(expansion_ratio * d_model) + return nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear(d_model, hidden_dim, bias=bias), + nn.GELU(), + nn.Linear(hidden_dim, d_model, bias=bias), + ) + + +class UnifiedTransformerBlock(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + use_geom_attn: bool = False, + use_plain_attn: bool = True, + v_heads: int | None = None, + bias: bool = False, + expansion_ratio: float = 4.0, + residue_scaling_factor: float = 1.0, + mask_and_zero_frameless: bool = False, + qk_layernorm: bool = True, + ffn_type: str = "swiglu", + attn_backend: str = "sdpa", + ): + super().__init__() + self.use_plain_attn = use_plain_attn + if self.use_plain_attn: + self.attn = MultiHeadAttention( + d_model, + n_heads, + bias, + qk_layernorm=qk_layernorm, + attn_backend=attn_backend, + ) + self.use_geom_attn = use_geom_attn + if self.use_geom_attn: + if v_heads is None: + raise ValueError("v_heads is required when geometric attention is enabled.") + self.geom_attn = GeometricReasoningOriginalImpl( + c_s=d_model, + v_heads=v_heads, + bias=bias, + mask_and_zero_frameless=mask_and_zero_frameless, + ) + if ffn_type == "swiglu": + self.ffn = swiglu_ln_ffn(d_model, expansion_ratio, bias) + elif ffn_type == "gelu": + self.ffn = gelu_ln_ffn(d_model, expansion_ratio, bias) + else: + raise ValueError(f"Unknown ffn_type: {ffn_type}") + self.scaling_factor = residue_scaling_factor + + def _add_scaled_residual( + self, hidden_states: torch.Tensor, residual: torch.Tensor + ) -> torch.Tensor: + return hidden_states + residual / self.scaling_factor + + def forward( + self, + x: torch.Tensor, + sequence_id: torch.Tensor | None, + attention_mask: torch.Tensor | None, + frames: Affine3D, + frames_mask: torch.Tensor, + chain_id: torch.Tensor, + output_attentions: bool = False, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + attn_weights: torch.Tensor | None = None + if self.use_plain_attn: + plain_residual, attn_weights = self.attn( + x, + sequence_id, + attention_mask, + output_attentions=output_attentions, + effective_backend=effective_backend, + ) + x = self._add_scaled_residual(x, plain_residual) + + if self.use_geom_attn: + geometric_residual = self.geom_attn( + x, + frames, + frames_mask, + sequence_id, + chain_id, + ) + x = self._add_scaled_residual(x, geometric_residual) + + return self._add_scaled_residual(x, self.ffn(x)), attn_weights + + +class TransformerStack(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + v_heads: int | None, + n_layers: int, + n_layers_geom: int = 1, + scale_residue: bool = True, + mask_and_zero_frameless: bool = False, + bias: bool = False, + qk_layernorm: bool = True, + ffn_type: str = "swiglu", + expansion_ratio: float = 8 / 3, + attn_backend: str = "sdpa", + ): + super().__init__() + self.blocks = nn.ModuleList( + [ + UnifiedTransformerBlock( + d_model, + n_heads, + v_heads=v_heads, + use_geom_attn=index < n_layers_geom, + residue_scaling_factor=(math.sqrt(n_layers / 36) if scale_residue else 1.0), + expansion_ratio=expansion_ratio, + mask_and_zero_frameless=mask_and_zero_frameless, + bias=bias, + qk_layernorm=qk_layernorm, + ffn_type=ffn_type, + attn_backend=attn_backend, + ) + for index in range(n_layers) + ] + ) + self.attention_backend = resolve_attention_backend(attn_backend) + self.norm = nn.LayerNorm(d_model, bias=False) + + def forward( + self, + x: torch.Tensor, + sequence_id: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + affine: Affine3D | None = None, + affine_mask: torch.Tensor | None = None, + chain_id: torch.Tensor | None = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, ...] | None, + tuple[torch.Tensor, ...] | None, + ]: + *batch_dims, _ = x.shape + if chain_id is None: + chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device) + if affine is None or affine_mask is None: + raise ValueError("affine and affine_mask are required for ESM3 transformer calls.") + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + all_hidden_states = [] if output_hidden_states else None + all_attentions = [] + for block in self.blocks: + x, attn_weights = block( + x, + sequence_id, + attention_mask, + affine, + affine_mask, + chain_id, + output_attentions=output_attentions, + effective_backend=effective_backend, + ) + if all_hidden_states is not None: + all_hidden_states.append(x) + if output_attentions and attn_weights is not None: + all_attentions.append(attn_weights) + hidden_states = tuple(all_hidden_states) if all_hidden_states is not None else None + attentions = tuple(all_attentions) if output_attentions else None + return self.norm(x), x, hidden_states, attentions + + +class EncodeInputs(nn.Module): + def __init__(self, d_model: int, sequence_vocab_size: int = 64) -> None: + super().__init__() + + discrete_tracks = ( + ("sequence_embed", sequence_vocab_size), + ("structure_tokens_embed", 4101), + ("ss8_embed", 11), + ("sasa_embed", 19), + ) + for attribute, vocabulary_size in discrete_tracks: + setattr(self, attribute, nn.Embedding(vocabulary_size, d_model)) + + self.plddt_projection, self.structure_per_res_plddt_projection = ( + nn.Linear(16, d_model), + nn.Linear(16, d_model), + ) + function_width = d_model // 8 + self.function_embed = nn.ModuleList( + nn.Embedding(260, function_width, padding_idx=0) for _ in range(8) + ) + self.residue_embed = nn.EmbeddingBag(1478, d_model, mode="sum", padding_idx=0) + + def forward( + self, + sequence_tokens: torch.Tensor, + structure_tokens: torch.Tensor, + average_plddt: torch.Tensor, + per_res_plddt: torch.Tensor, + ss8_tokens: torch.Tensor, + sasa_tokens: torch.Tensor, + function_tokens: torch.Tensor, + residue_annotation_tokens: torch.Tensor, + ) -> torch.Tensor: + sequence_embed = self.sequence_embed(sequence_tokens) + rbf_16_fn = functools.partial(rbf, v_min=0.0, v_max=1.0, n_bins=16) + plddt_embed = self.plddt_projection( + rbf_16_fn(average_plddt).to(self.plddt_projection.weight.dtype) + ) + structure_per_res_plddt = self.structure_per_res_plddt_projection( + rbf_16_fn(per_res_plddt).to(self.structure_per_res_plddt_projection.weight.dtype) + ) + structure_embed = self.structure_tokens_embed(structure_tokens) + ss8_embed = self.ss8_embed(ss8_tokens) + sasa_embed = self.sasa_embed(sasa_tokens) + function_embed = torch.cat( + [ + embed_fn(funcs) + for embed_fn, funcs in zip( + self.function_embed, + function_tokens.unbind(-1), + strict=True, + ) + ], + -1, + ) + + batch_size, seq_len, num_annotations = residue_annotation_tokens.shape + residue_embed = self.residue_embed( + rearrange( + residue_annotation_tokens, + "b l n -> (b l) n", + b=batch_size, + l=seq_len, + n=num_annotations, + ) + ) + residue_embed = rearrange( + residue_embed, + "(b l) d -> b l d", + b=batch_size, + l=seq_len, + ) + + return ( + sequence_embed + + plddt_embed + + structure_per_res_plddt + + structure_embed + + ss8_embed + + sasa_embed + + function_embed + + residue_embed + ) + + +@dataclass +class ESM3CoreOutput: + sequence_logits: torch.Tensor + structure_logits: torch.Tensor + secondary_structure_logits: torch.Tensor + sasa_logits: torch.Tensor + function_logits: torch.Tensor + residue_logits: torch.Tensor + embeddings: torch.Tensor + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +class OutputHeads(nn.Module): + def __init__(self, d_model: int, sequence_vocab_size: int = 64) -> None: + super().__init__() + self.sequence_head = RegressionHead(d_model, sequence_vocab_size) + self.structure_head = RegressionHead(d_model, 4096) + self.ss8_head = RegressionHead(d_model, 8 + 3) + self.sasa_head = RegressionHead(d_model, 16 + 3) + self.function_head = RegressionHead(d_model, 260 * 8) + self.residue_head = RegressionHead(d_model, 1478) + + def forward( + self, + x: torch.Tensor, + embed: torch.Tensor, + hidden_states: tuple[torch.Tensor, ...] | None = None, + attentions: tuple[torch.Tensor, ...] | None = None, + ) -> ESM3CoreOutput: + function_logits = self.function_head(x) + function_logits = rearrange(function_logits, "... (k v) -> ... k v", k=8) + return ESM3CoreOutput( + sequence_logits=self.sequence_head(x), + structure_logits=self.structure_head(x), + secondary_structure_logits=self.ss8_head(x), + sasa_logits=self.sasa_head(x), + function_logits=function_logits, + residue_logits=self.residue_head(x), + embeddings=embed, + hidden_states=hidden_states, + attentions=attentions, + ) + + +class ESM3Core(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + v_heads: int, + n_layers: int, + attn_backend: str = "sdpa", + sequence_vocab_size: int = 64, + ): + super().__init__() + self.encoder = EncodeInputs(d_model, sequence_vocab_size) + self.transformer = TransformerStack( + d_model, + n_heads, + v_heads, + n_layers, + mask_and_zero_frameless=True, + attn_backend=attn_backend, + ) + self.output_heads = OutputHeads(d_model, sequence_vocab_size) + + def forward( + self, + *, + sequence_tokens: torch.Tensor | None = None, + structure_tokens: torch.Tensor | None = None, + ss8_tokens: torch.Tensor | None = None, + sasa_tokens: torch.Tensor | None = None, + function_tokens: torch.Tensor | None = None, + residue_annotation_tokens: torch.Tensor | None = None, + average_plddt: torch.Tensor | None = None, + per_res_plddt: torch.Tensor | None = None, + structure_coords: torch.Tensor | None = None, + chain_id: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + ) -> ESM3CoreOutput: + output_attentions = bool(output_attentions) + output_hidden_states = bool(output_hidden_states) + present_inputs = [ + sequence_tokens, + structure_tokens, + ss8_tokens, + sasa_tokens, + structure_coords, + function_tokens, + residue_annotation_tokens, + ] + try: + seq_len, device = next((x.shape[1], x.device) for x in present_inputs if x is not None) + except StopIteration: + raise ValueError("At least one of the inputs must be non-None") from None + + def defaults(x: torch.Tensor | None, token: int) -> torch.Tensor: + if x is None: + return torch.full( + (1, seq_len), + token, + dtype=torch.long, + device=device, + ) + return x + + sequence_tokens = defaults(sequence_tokens, SEQUENCE_MASK_TOKEN) + ss8_tokens = defaults(ss8_tokens, SS8_PAD_TOKEN) + sasa_tokens = defaults(sasa_tokens, SASA_PAD_TOKEN) + average_plddt = defaults(average_plddt, 1).float() + per_res_plddt = defaults(per_res_plddt, 0).float() + chain_id = defaults(chain_id, 0) + + if residue_annotation_tokens is None: + residue_annotation_tokens = torch.full( + (1, seq_len, MAX_RESIDUE_ANNOTATIONS), + RESIDUE_PAD_TOKEN, + dtype=torch.long, + device=device, + ) + if function_tokens is None: + function_tokens = torch.full( + (1, seq_len, FUNCTION_TOKENS_DEPTH), + INTERPRO_PAD_TOKEN, + dtype=torch.long, + device=device, + ) + if structure_coords is None: + structure_coords = torch.full( + (1, seq_len, 3, 3), + float("nan"), + dtype=torch.float, + device=device, + ) + + structure_coords = structure_coords[..., :3, :] + affine, affine_mask = build_affine3d_from_coordinates(structure_coords) + + structure_tokens = defaults(structure_tokens, STRUCTURE_MASK_TOKEN) + structure_tokens = ( + structure_tokens.masked_fill(structure_tokens == -1, STRUCTURE_MASK_TOKEN) + .masked_fill(sequence_tokens == SEQUENCE_BOS_TOKEN, STRUCTURE_BOS_TOKEN) + .masked_fill(sequence_tokens == SEQUENCE_PAD_TOKEN, STRUCTURE_PAD_TOKEN) + .masked_fill(sequence_tokens == SEQUENCE_EOS_TOKEN, STRUCTURE_EOS_TOKEN) + .masked_fill( + sequence_tokens == SEQUENCE_CHAINBREAK_TOKEN, + STRUCTURE_CHAINBREAK_TOKEN, + ) + ) + + x = self.encoder( + sequence_tokens, + structure_tokens, + average_plddt, + per_res_plddt, + ss8_tokens, + sasa_tokens, + function_tokens, + residue_annotation_tokens, + ) + expected_mask_shape = tuple(x.shape[:2]) + if sequence_id is not None and tuple(sequence_id.shape) != expected_mask_shape: + raise ValueError( + "sequence_id must have shape (batch, sequence); " + f"expected {expected_mask_shape}, received {tuple(sequence_id.shape)}." + ) + if attention_mask is not None: + if tuple(attention_mask.shape) != expected_mask_shape: + raise ValueError( + "attention_mask must have shape (batch, sequence); " + f"expected {expected_mask_shape}, received {tuple(attention_mask.shape)}." + ) + if attention_mask.dtype != torch.bool and not bool( + torch.logical_or(attention_mask == 0, attention_mask == 1).all() + ): + raise ValueError("attention_mask must contain only boolean or 0/1 values.") + attention_mask = attention_mask.to(device=x.device, dtype=torch.bool) + if not bool(attention_mask.any(dim=-1).all()): + raise ValueError("attention_mask must keep at least one valid key per batch row.") + affine_mask = affine_mask & attention_mask + x, embedding, hidden_states, attentions = self.transformer( + x, + sequence_id, + attention_mask, + affine, + affine_mask, + chain_id, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + return self.output_heads( + x, + embedding, + hidden_states=hidden_states, + attentions=attentions, + ) + + +def _resolve_esm3_checkpoint_key(model_name: str) -> str: + if model_name in ESM3_OPEN_SMALL_ALIASES: + return ESM3_OPEN_SMALL + raise ValueError( + f"Unsupported ESM3 checkpoint {model_name}. " + f"Supported names: {sorted(ESM3_OPEN_SMALL_ALIASES)}" + ) + + +def _build_esm3_core(config: FastESM3Config) -> nn.Module: + return ESM3Core( + d_model=config.hidden_size, + n_heads=config.num_attention_heads, + v_heads=config.num_vector_heads, + n_layers=config.num_hidden_layers, + attn_backend=config.attn_backend, + sequence_vocab_size=config.vocab_size, + ) + + +class FastESM3PreTrainedModel(FastPLMsAttentionMixin, PreTrainedModel): + config_class = FastESM3Config + base_model_prefix = "esm3" + main_input_name = "input_ids" + supports_gradient_checkpointing = False + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = _SUPPORTED_ATTENTION_BACKENDS + + @property + def tokenizer(self) -> EsmSequenceTokenizer: + """Construct the sequence tokenizer only when a raw-sequence API needs it.""" + + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + tokenizer = EsmSequenceTokenizer() + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value: EsmSequenceTokenizer | None) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + def _init_weights(self, module: nn.Module) -> None: + for parameter in module.parameters(recurse=False): + if parameter.__dict__.get("_is_hf_initialized"): + return + + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + with torch.no_grad(): + module.weight[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + if module.bias is not None: + nn.init.zeros_(module.bias) + nn.init.ones_(module.weight) + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in _SUPPORTED_ATTENTION_BACKENDS: + raise ValueError( + f"ESM3 currently supports only {_SUPPORTED_ATTENTION_BACKENDS}; got {backend}." + ) + self.set_attn_implementation(backend) + + +class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, EmbeddingMixin): + config_class = FastESM3Config + # Direct ESM3 saves intentionally package an independently loadable remote + # runtime. Register the concrete advertised class explicitly so + # Transformers writes a real AutoModel key instead of a null auto_map key. + _auto_class = "AutoModel" + + def __init__(self, config: FastESM3Config, **kwargs) -> None: + super().__init__(config, **kwargs) + self.esm3 = _build_esm3_core(config) + self.post_init() + self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + @property + def raw_model(self) -> nn.Module: + return self.esm3 + + def get_input_embeddings(self) -> nn.Module: + return self.esm3.encoder.sequence_embed + + def set_input_embeddings(self, value: nn.Module) -> None: + self.esm3.encoder.sequence_embed = value + + def get_output_embeddings(self) -> nn.Module: + return self.esm3.output_heads.sequence_head[-1] + + def set_output_embeddings(self, value: nn.Module) -> None: + self.esm3.output_heads.sequence_head[-1] = value + + def save_pretrained(self, save_directory, *args, **kwargs) -> None: + """Save weights plus the unchanged sources needed for an isolated reload.""" + + save_path = Path(save_directory) + _validate_saved_runtime_destination(save_path) + prepared_runtime = _build_saved_runtime_archive(Path(__file__).resolve().parents[2]) + super().save_pretrained(save_directory, *args, **kwargs) + _write_saved_runtime(save_path, prepared_runtime) + + def tokenize_sequences( + self, + sequences: str | list[str], + padding: bool = True, + return_tensors: str = "pt", + device: torch.device | str | None = None, + add_special_tokens: bool = True, + ) -> dict[str, torch.Tensor]: + tokenized = self.tokenizer( + sequences, + padding=padding, + return_tensors=return_tensors, + add_special_tokens=add_special_tokens, + ) + if device is None: + return tokenized + return {name: tensor.to(device) for name, tensor in tokenized.items()} + + def forward_sequence( + self, + sequences: str | list[str], + device: torch.device | str | None = None, + **kwargs, + ) -> FastESM3Output: + if device is None: + device = self.device + tokenized = self.tokenize_sequences(sequences, device=device) + return self(**tokenized, **kwargs) + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + **kwargs, + ) -> torch.Tensor: + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + return_dict=True, + **kwargs, + ) + if store_all_hidden_states: + if output.hidden_states is None: + raise RuntimeError("store_all_hidden_states requires hidden states.") + return torch.stack(tuple(output.hidden_states), dim=1) + if hidden_state_index == -1: + return output.last_hidden_state + if output.hidden_states is None: + raise RuntimeError("hidden_state_index selection requires hidden states.") + return output.hidden_states[hidden_state_index] + + def encode( + self, + inputs: str | list[str], + *, + device: torch.device | str | None = None, + ) -> dict[str, torch.Tensor]: + """Tokenize raw sequences without importing the Biohub SDK.""" + if isinstance(inputs, str): + inputs = inputs.replace("_", self.tokenizer.mask_token) + else: + inputs = [sequence.replace("_", self.tokenizer.mask_token) for sequence in inputs] + return self.tokenize_sequences(inputs, device=device or self.device) + + def decode(self, inputs: torch.Tensor | dict[str, torch.Tensor]) -> str | list[str]: + """Decode sequence tokens while removing model special tokens.""" + token_ids = inputs["input_ids"] if isinstance(inputs, dict) else inputs + single = token_ids.ndim == 1 + if single: + token_ids = token_ids.unsqueeze(0) + sequences = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True) + sequences = [sequence.replace(" ", "") for sequence in sequences] + return sequences[0] if single else sequences + + @torch.inference_mode() + def generate( + self, + inputs: str | list[str] | torch.Tensor | dict[str, torch.Tensor], + config: FastESM3GenerationConfig | None = None, + ) -> str | list[str] | torch.Tensor: + """Fill sequence-track mask tokens with iterative categorical sampling. + + Raw strings use ``_`` for masked residues. Tensor inputs use token ID 32. + The method samples only amino-acid token IDs and preserves every + unmasked input token. + """ + config = config or FastESM3GenerationConfig() + if config.temperature <= 0: + raise ValueError("temperature must be greater than zero") + if config.num_steps is not None: + if isinstance(config.num_steps, bool) or not isinstance(config.num_steps, int): + raise TypeError("num_steps must be an integer or None") + if config.num_steps <= 0: + raise ValueError("num_steps must be positive") + + return_strings = isinstance(inputs, (str, list)) + single_string = isinstance(inputs, str) + if return_strings: + encoded = self.encode(inputs) + token_ids = encoded["input_ids"] + conditioning = {"attention_mask": encoded["attention_mask"]} + elif isinstance(inputs, dict): + supported_inputs = { + "input_ids", + "attention_mask", + "sequence_tokens", + "structure_tokens", + "ss8_tokens", + "sasa_tokens", + "function_tokens", + "residue_annotation_tokens", + "average_plddt", + "per_res_plddt", + "structure_coords", + "chain_id", + "sequence_id", + } + unsupported = sorted(set(inputs) - supported_inputs) + if unsupported: + names = ", ".join(unsupported) + raise TypeError(f"Unsupported ESM3 generation inputs: {names}") + if "input_ids" in inputs and "sequence_tokens" in inputs: + raise ValueError("Pass only one of input_ids or sequence_tokens to generate().") + sequence_key = "input_ids" if "input_ids" in inputs else "sequence_tokens" + if sequence_key not in inputs: + raise ValueError("ESM3 generation requires input_ids or sequence_tokens.") + token_ids = inputs[sequence_key].to(self.device) + conditioning = { + name: value.to(self.device) + for name, value in inputs.items() + if name != sequence_key + } + else: + token_ids = inputs.to(self.device) + conditioning = {} + + single_tensor = token_ids.ndim == 1 + if single_tensor: + sequence_length = token_ids.shape[0] + token_ids = token_ids.unsqueeze(0) + conditioning = { + name: ( + value.unsqueeze(0) + if value.ndim > 0 and value.shape[0] == sequence_length + else value + ) + for name, value in conditioning.items() + } + sampled_ids = token_ids.clone() + initial_mask = sampled_ids.eq(SEQUENCE_MASK_TOKEN) + n_masked = int(initial_mask.sum().item()) + if n_masked == 0: + result = sampled_ids.squeeze(0) if single_tensor else sampled_ids + if return_strings: + decoded = self.decode(result) + return decoded[0] if single_string and isinstance(decoded, list) else decoded + return result + + n_steps = n_masked if config.num_steps is None else config.num_steps + generator = None + if config.seed is not None: + generator = torch.Generator(device=sampled_ids.device) + generator.manual_seed(config.seed) + + for step in range(n_steps): + remaining = sampled_ids.eq(SEQUENCE_MASK_TOKEN) + if not bool(remaining.any()): + break + with _temporary_eval(self): + output = self( + sequence_tokens=sampled_ids, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + **conditioning, + ) + amino_acid_logits = output.sequence_logits[..., 4:29] / config.temperature + probabilities = amino_acid_logits.softmax(dim=-1) + sampled = ( + torch.multinomial( + probabilities.reshape(-1, probabilities.shape[-1]), + num_samples=1, + generator=generator, + ).reshape_as(sampled_ids) + + 4 + ) + + remaining_count = int(remaining.sum().item()) + steps_left = n_steps - step + fill_count = max(1, (remaining_count + steps_left - 1) // steps_left) + confidence = probabilities.max(dim=-1).values.masked_fill(~remaining, -1.0) + selected = torch.zeros_like(remaining) + flat_selected = selected.reshape(-1) + chosen = confidence.reshape(-1).topk(min(fill_count, remaining_count)).indices + flat_selected[chosen] = True + sampled_ids[selected] = sampled[selected] + + if bool(sampled_ids.eq(SEQUENCE_MASK_TOKEN).any()): + raise RuntimeError("generation ended before all sequence masks were filled") + result = sampled_ids.squeeze(0) if single_tensor else sampled_ids + if return_strings: + decoded = self.decode(result) + return decoded[0] if single_string and isinstance(decoded, list) else decoded + return result + + def batch_generate( + self, + inputs: list[str | torch.Tensor], + configs: list[FastESM3GenerationConfig], + ) -> list[str | torch.Tensor]: + if len(inputs) != len(configs): + raise ValueError("inputs and configs must have equal lengths") + return [self.generate(value, config) for value, config in zip(inputs, configs, strict=True)] + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.esm3] + + def forward_and_sample( + self, + inputs: str | list[str] | torch.Tensor | dict[str, torch.Tensor], + sampling_configuration: FastESM3GenerationConfig | None = None, + ) -> str | list[str] | torch.Tensor: + return self.generate(inputs, sampling_configuration) + + def logits(self, inputs=None, **kwargs) -> FastESM3Output: + if inputs is None: + return self.forward(**kwargs) + if isinstance(inputs, (str, list)): + return self.forward(**self.encode(inputs), **kwargs) + if isinstance(inputs, dict): + return self.forward(**inputs, **kwargs) + if isinstance(inputs, torch.Tensor): + return self.forward(sequence_tokens=inputs, **kwargs) + raise TypeError("inputs must be raw sequences, sequence tokens, or a token mapping") + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_tokens: torch.Tensor | None = None, + structure_tokens: torch.Tensor | None = None, + ss8_tokens: torch.Tensor | None = None, + sasa_tokens: torch.Tensor | None = None, + function_tokens: torch.Tensor | None = None, + residue_annotation_tokens: torch.Tensor | None = None, + average_plddt: torch.Tensor | None = None, + per_res_plddt: torch.Tensor | None = None, + structure_coords: torch.Tensor | None = None, + chain_id: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> FastESM3Output | tuple[torch.Tensor, ...]: + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected ESM3 forward arguments: {names}") + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + if input_ids is not None and sequence_tokens is not None: + raise ValueError("Pass only one of input_ids or sequence_tokens.") + if sequence_tokens is None: + sequence_tokens = input_ids + output = self.esm3( + sequence_tokens=sequence_tokens, + structure_tokens=structure_tokens, + ss8_tokens=ss8_tokens, + sasa_tokens=sasa_tokens, + function_tokens=function_tokens, + residue_annotation_tokens=residue_annotation_tokens, + average_plddt=average_plddt, + per_res_plddt=per_res_plddt, + structure_coords=structure_coords, + chain_id=chain_id, + sequence_id=sequence_id, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + loss = None + if labels is not None: + labels = labels.to(output.sequence_logits.device) + loss = F.cross_entropy( + output.sequence_logits.view(-1, output.sequence_logits.shape[-1]), + labels.view(-1), + ignore_index=-100, + ) + + result = FastESM3Output( + last_hidden_state=output.embeddings, + hidden_states=output.hidden_states, + attentions=output.attentions, + logits=output.sequence_logits, + sequence_logits=output.sequence_logits, + structure_logits=output.structure_logits, + secondary_structure_logits=output.secondary_structure_logits, + sasa_logits=output.sasa_logits, + function_logits=output.function_logits, + residue_logits=output.residue_logits, + embeddings=output.embeddings, + loss=loss, + ) + if not return_dict: + return result.to_tuple() + return result diff --git a/fastplms/dplm2/__init__.py b/src/fastplms/models/esm_plusplus/__init__.py similarity index 100% rename from fastplms/dplm2/__init__.py rename to src/fastplms/models/esm_plusplus/__init__.py diff --git a/src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py b/src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py new file mode 100644 index 0000000..4fd0fbd --- /dev/null +++ b/src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py @@ -0,0 +1,1552 @@ +"""Hugging Face-compatible ESMC models implemented by FastPLMs.""" + +from __future__ import annotations + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass +from functools import partial +from typing import ClassVar +from einops import rearrange +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.processors import TemplateProcessing +from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerFast +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + _get_flex_block_mask, + flex_attention, + get_attention_mask, + kernels_flash_attention_func, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin, Pooler, select_hidden_state_embeddings + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "Pooler", + "_get_flex_attention_fn", + "_get_flex_block_mask", + "flex_attention", + "get_attention_mask", + "kernels_flash_attention_func", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +class ESMplusplusConfig(PretrainedConfig): + """Configuration class for ESM++ model. + + Args: + vocab_size: Size of the vocabulary + hidden_size: Dimension of hidden layers + num_attention_heads: Number of attention heads + num_hidden_layers: Number of transformer layers + num_labels: Number of output labels for classification + problem_type: Type of problem - regression, single/multi label classification + """ + + model_type = "ESMplusplus" + + def __init__( + self, + vocab_size: int = 64, + hidden_size: int = 960, + num_attention_heads: int = 15, + num_hidden_layers: int = 30, + num_labels: int | None = None, + problem_type: str | None = None, + dropout: float = 0.0, + initializer_range: float = 0.02, + classifier_dropout: float = 0.1, + classifier_pooling_types: list[str] | None = None, + attn_backend: str | None = None, + pad_token_id: int = 1, + mask_token_id: int = 32, + **kwargs, + ): + if num_labels is None: + configured_labels = kwargs.get("id2label") + num_labels = len(configured_labels) if configured_labels else 2 + super().__init__( + pad_token_id=pad_token_id, + mask_token_id=mask_token_id, + num_labels=num_labels, + **kwargs, + ) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.problem_type = problem_type + self.dropout = dropout + self.initializer_range = initializer_range + self.classifier_dropout = classifier_dropout + self.classifier_pooling_types = ( + list(classifier_pooling_types) if classifier_pooling_types is not None else None + ) + self.tie_word_embeddings = False + self.attn_backend = attn_backend + + +### Rotary Embeddings +def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: + """Rotate the final axis of X by 90 degrees in each two-dimensional plane.""" + if interleaved: + paired = x.unflatten(-1, (-1, 2)) + return torch.stack((-paired[..., 1], paired[..., 0]), dim=-1).flatten(-2) + + # torch.chunk assigns an odd remainder to the first half. Express the same + # public behavior explicitly while keeping the ESMC path branch-free. + midpoint = (x.shape[-1] + 1) // 2 + return torch.cat((-x[..., midpoint:], x[..., :midpoint]), dim=-1) + + +def apply_rotary_emb_torch( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = False, + _inplace: bool = False, +) -> torch.Tensor: + """Apply cached rotary angles to X while preserving any unrotated features.""" + del _inplace # Kept in the signature for checkpoint remote-code compatibility. + rotary_width = 2 * cos.shape[-1] + if rotary_width > x.shape[-1]: + raise AssertionError("rotary width exceeds the attention head dimension") + + token_count = x.shape[1] + cos_full = torch.cat((cos[:token_count], cos[:token_count]), dim=-1).unsqueeze(1) + sin_full = torch.cat((sin[:token_count], sin[:token_count]), dim=-1).unsqueeze(1) + x_rotary = x[..., :rotary_width] + y_rotary = x_rotary * cos_full + rotate_half(x_rotary, interleaved) * sin_full + if rotary_width == x.shape[-1]: + return y_rotary + return torch.cat((y_rotary, x[..., rotary_width:]), dim=-1) + + +class RotaryEmbedding(torch.nn.Module): + """Rotary position embeddings. + + Based on the paper "RoFormer: Enhanced Transformer with Rotary Position Embedding" + + Args: + dim: Dimension of the embedding + base: Base for computing angular frequencies + interleaved: Whether to use interleaved rotations + scale_base: Base for scaling + scaling_factor: Factor for scaling positions + pos_idx_in_fp32: Whether to compute position indices in fp32 + device: Computation device + """ + + def __init__( + self, + dim: int, + base: float = 10000.0, + interleaved: bool = False, + scale_base: float | None = None, + scaling_factor: float = 1.0, + pos_idx_in_fp32: bool = True, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.dim, self.base = dim, float(base) + self.interleaved, self.scale_base = interleaved, scale_base + self.scaling_factor, self.pos_idx_in_fp32 = scaling_factor, pos_idx_in_fp32 + self.device = device + self._clear_cache() + self.reset_parameters() + + def _clear_cache(self) -> None: + self._seq_len_cached = 0 + self._cos_cached: torch.Tensor | None = None + self._sin_cached: torch.Tensor | None = None + self._cos_k_cached: torch.Tensor | None = None + self._sin_k_cached: torch.Tensor | None = None + + def reset_parameters(self, device: torch.device | str | None = None) -> None: + """Rebuild the non-persistent frequency buffers on ``device``.""" + if device is not None: + buffer_device = torch.device(device) + elif "inv_freq" in self._buffers and isinstance(self._buffers["inv_freq"], torch.Tensor): + buffer_device = self._buffers["inv_freq"].device + else: + buffer_device = self.device + inv_freq = self._compute_inv_freq(buffer_device) + self._clear_cache() + self.register_buffer("inv_freq", inv_freq, persistent=False) + arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32) + scale = ( + (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None + ) + self.register_buffer("scale", scale) + + def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor: + """Compute inverse frequency bands on their execution device.""" + return 1 / ( + self.base + ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) + ) + + def _apply(self, fn, recurse: bool = True): + """Move the module, then regenerate device-specific RoPE frequencies.""" + if self.inv_freq.is_meta: + self.reset_parameters(device="cpu") + result = super()._apply(fn, recurse=recurse) + self.register_buffer( + "inv_freq", + self._compute_inv_freq(self.inv_freq.device), + persistent=False, + ) + self._clear_cache() + return result + + def _cache_is_current( + self, + token_count: int, + device: torch.device | None, + dtype: torch.dtype | None, + ) -> bool: + cached = self._cos_cached + return ( + cached is not None + and self._seq_len_cached >= token_count + and cached.device == device + and cached.dtype == dtype + and not (self.training and cached.is_inference()) + ) + + def _rotary_angles( + self, + token_count: int, + device: torch.device | None, + ) -> torch.Tensor: + position_dtype = torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype + positions = torch.arange(token_count, device=device, dtype=position_dtype) # (l,) + positions.div_(self.scaling_factor) + frequencies = ( + self.inv_freq.to(torch.float32) + if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32 + else self.inv_freq + ) + return torch.outer(positions, frequencies) # (l, d / 2) + + def _update_cos_sin_cache( + self, seqlen: int, device: torch.device | None = None, dtype: torch.dtype | None = None + ) -> None: + """Build angle tables when the requested cache identity has changed.""" + if self._cache_is_current(seqlen, device, dtype): + return + + self._seq_len_cached = seqlen + angles = self._rotary_angles(seqlen, device) # (l, d / 2) + cos_angles = torch.cos(angles) # (l, d / 2) + sin_angles = torch.sin(angles) # (l, d / 2) + if self.scale is None: + self._cos_cached = cos_angles.to(dtype) + self._sin_cached = sin_angles.to(dtype) + return + + centered_positions = ( + torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2 + ) / self.scale_base + scale = self.scale ** centered_positions.unsqueeze(-1) + self._cos_cached = (cos_angles * scale).to(dtype) + self._sin_cached = (sin_angles * scale).to(dtype) + self._cos_k_cached = (cos_angles / scale).to(dtype) + self._sin_k_cached = (sin_angles / scale).to(dtype) + + def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary embeddings to queries and keys. + + Args: + q: Query tensor Q with shape (b, l, h, d). + k: Key tensor K with shape (b, l, h, d). + + Returns: + Tuple of rotated query and key tensors + """ + # The pinned Biohub Transformers oracle recomputes inverse frequencies + # on the execution device. CPU and CUDA differ by about one FP32 ULP in + # some bands, which is immaterial in BF16 but accumulates measurably in + # deep FP32 execution. + self._update_cos_sin_cache(q.shape[1], device=q.device, dtype=q.dtype) + if self._cos_cached is None or self._sin_cached is None: + raise RuntimeError( + "Rotary cache initialization did not produce cosine and sine values." + ) + if self.scale is not None: + raise AssertionError("Scaled rotary embeddings are unsupported for ESMC.") + + cos_angles = self._cos_cached + sin_angles = self._sin_cached + return ( + apply_rotary_emb_torch(q, cos_angles, sin_angles, self.interleaved, True), + apply_rotary_emb_torch(k, cos_angles, sin_angles, self.interleaved, True), + ) + + +def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int: + """Compute corrected dimension for SwiGLU.""" + return int(((expansion_ratio * d_model) + 255) // 256 * 256) + + +class SwiGLU(nn.Module): + """SwiGLU activation function.""" + + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return F.silu(x1) * x2 + + +def swiglu_ln_ffn(d_model: int, expansion_ratio: float) -> nn.Sequential: + """Create SwiGLU feedforward network with layer normalization.""" + return nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear(d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=False), + SwiGLU(), + nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=False), + ) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention with rotary embeddings and configurable backend. + + Args: + d_model: Model dimension + n_heads: Number of attention heads + attn_backend: One of "eager", "sdpa", or "flex_attention". + """ + + def __init__( + self, + d_model: int, + n_heads: int, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.d_model = d_model + self.n_heads = n_heads + self.d_head = self.d_model // self.n_heads + self.scale = 1.0 / math.sqrt(self.d_head) + self.attn_backend = resolve_attention_backend(attn_backend) + self.layernorm_qkv = nn.Sequential( + nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=False) + ) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.q_ln = nn.LayerNorm(d_model, bias=False) + self.k_ln = nn.LayerNorm(d_model, bias=False) + self.reshaper = partial(rearrange, pattern="b s (h d) -> b h s d", h=n_heads) + self.rotary = RotaryEmbedding(d_model // n_heads) + + def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, l, d) + q = q.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + k = k.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + q, k = self.rotary(q, k) + q = q.flatten(-2, -1) + k = k.flatten(-2, -1) + return q, k + + def forward( + self, + x: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # x: (b, l, d) + qkv = self.layernorm_qkv(x) # (b, l, 3 * d) + query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1) + query_sequence, key_sequence = ( + self.q_ln(query_sequence).to(query_sequence.dtype), + self.k_ln(key_sequence).to(query_sequence.dtype), + ) + query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence) + query_heads, key_heads, value_heads = map( + self.reshaper, (query_sequence, key_sequence, value_sequence) + ) # each (b, h, l, d_h) + + attn_output, attn_weights, s_max = self._attn( + query_heads, + key_heads, + value_heads, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + output = self.out_proj(attn_output) + return output, attn_weights, s_max + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + if output_attentions: + return self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + + if self.attn_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + return attn_output, None, s_max + if self.attn_backend.is_flash: + attn_output, attn_weights = self._kernels_flash_attn( + query_heads, key_heads, value_heads, attention_mask_2d + ) + elif self.attn_backend == AttentionBackend.FLEX: + attn_output, attn_weights = self._flex_attn( + query_heads, + key_heads, + value_heads, + flex_block_mask, + attention_mask_2d, + ) + elif self.attn_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_heads, key_heads, value_heads, attention_mask_4d + ) + else: + raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") + + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, query_heads: torch.Tensor, key_heads: torch.Tensor + ) -> list[torch.Tensor]: + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) + s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max( + dim=0 + ).values * self.scale + return [s_max_bound[h] for h in range(self.n_heads)] + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = ( + torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self.scale + ) # (b, h, l, l) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) + attn_weights = F.softmax(attn_weights, dim=-1) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l, d) + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + def _kernels_flash_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + query_tokens = query_heads.transpose(1, 2).contiguous() + key_tokens = key_heads.transpose(1, 2).contiguous() + value_tokens = value_heads.transpose(1, 2).contiguous() + attn_output = kernels_flash_attention_func( + query_states=query_tokens, + key_states=key_tokens, + value_states=value_tokens, + attention_mask_2d=attention_mask_2d, + causal=False, + implementation=self.attn_backend.value, + ) + return rearrange(attn_output, "b s h d -> b s (h d)"), None + + def _flex_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + flex_block_mask: BlockMask | None = None, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if flex_attention is None: + raise RuntimeError("Flex attention is not available in this environment.") + fn = _get_flex_attention_fn( + device=query_heads.device, + dtype=query_heads.dtype, + shape=tuple(query_heads.shape), + mask_semantics="padding", + ) + context_heads = fn( + query_heads, + key_heads, + value_heads, + block_mask=flex_block_mask, + scale=self.scale, + kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32}, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + scale=self.scale, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +def RegressionHead(d_model: int, output_dim: int, hidden_dim: int | None = None) -> nn.Module: + """Create a regression head with optional hidden dimension. + + Args: + d_model: Input dimension + output_dim: Output dimension + hidden_dim: Optional hidden dimension (defaults to d_model) + """ + hidden_dim = hidden_dim if hidden_dim is not None else d_model + return nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, output_dim), + ) + + +class UnifiedTransformerBlock(nn.Module): + """Transformer block with attention and feedforward layers.""" + + def __init__( + self, + d_model: int, + n_heads: int, + residue_scaling_factor: float = 1, + expansion_ratio: float = 8 / 3, + dropout: float = 0.0, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, attn_backend=attn_backend) + self.ffn = swiglu_ln_ffn(d_model, expansion_ratio) + self.scaling_factor = residue_scaling_factor + self.dropout = nn.Dropout(dropout) + + def forward( + self, + x: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + attn_output, attn_weights, s_max = self.attn( + x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + x = x + self.dropout(attn_output) / self.scaling_factor + x = x + self.dropout(self.ffn(x)) / self.scaling_factor + return x, attn_weights, s_max + + +@dataclass +class TransformerOutput(ModelOutput): + """Output type for transformer encoder.""" + + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor] | None = None + attentions: tuple[torch.Tensor] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class ESMplusplusOutput(MaskedLMOutput): + """Masked-LM output with FastPLMs fields after the HF contract.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + last_hidden_state: torch.Tensor | None = None + + +@dataclass +class ESMplusplusSequenceClassifierOutput(SequenceClassifierOutput): + """Sequence-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +@dataclass +class ESMplusplusTokenClassifierOutput(TokenClassifierOutput): + """Token-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + + +class TransformerStack(nn.Module): + """Stack of transformer blocks.""" + + def __init__( + self, + d_model: int, + n_heads: int, + n_layers: int, + dropout: float = 0.0, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.attention_backend = resolve_attention_backend(attn_backend) + self.blocks = nn.ModuleList( + [ + UnifiedTransformerBlock( + d_model, + n_heads, + residue_scaling_factor=math.sqrt(n_layers / 36), + dropout=dropout, + attn_backend=attn_backend, + ) + for i in range(n_layers) + ] + ) + self.norm = nn.LayerNorm(d_model, bias=False) + self.gradient_checkpointing = False + + @property + def attn_backend(self) -> AttentionBackend: + return self.attention_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + resolved = resolve_attention_backend(backend) + self.attention_backend = resolved + for block in self.blocks: + block.attn.attn_backend = resolved + + def forward( + self, + x: torch.Tensor, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_hidden_states: bool | None = False, + output_attentions: bool | None = False, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + ) -> TransformerOutput: + # x: (b, l, d); attention_mask, sequence_id: (b, l) + hidden_states = () if output_hidden_states else None + attentions = () if output_attentions else None + full_s_max = () if output_s_max else None + # Match the pinned Biohub Transformers contract: a supplied sequence_id + # is authoritative and must encode padding as -1. attention_mask is + # ignored in that mode rather than intersected with the chain mask. + if sequence_id is None and attention_mask is not None: + expected_shape = (x.shape[0], x.shape[1]) + if attention_mask.ndim != 2 or tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"attention_mask must have shape {expected_shape}; " + f"received {tuple(attention_mask.shape)}." + ) + attention_mask = attention_mask.to(device=x.device, dtype=torch.bool) + if not bool(attention_mask.any(dim=1).all()): + raise ValueError("attention_mask must keep at least one valid key per batch row.") + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=bool(output_attentions), + ) + + if sequence_id is None and attention_mask is not None: + attention_mask_2d, attention_mask_4d, flex_block_mask = ( + self._sequence_id_attention_masks( + sequence_id=attention_mask.to(device=x.device, dtype=torch.bool), + batch_size=x.shape[0], + seq_len=x.shape[1], + device=x.device, + dtype=x.dtype, + effective_backend=effective_backend, + ) + ) + elif sequence_id is None: + attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( + effective_backend=effective_backend, + batch_size=x.shape[0], + seq_len=x.shape[1], + device=x.device, + attention_mask=attention_mask, + dtype=x.dtype, + mask_semantics="padding", + ) + else: + attention_mask_2d, attention_mask_4d, flex_block_mask = ( + self._sequence_id_attention_masks( + sequence_id=sequence_id, + batch_size=x.shape[0], + seq_len=x.shape[1], + device=x.device, + dtype=x.dtype, + effective_backend=effective_backend, + ) + ) + + for block in self.blocks: + if output_hidden_states: + if hidden_states is None: + raise RuntimeError( + "Hidden-state collection was not initialized for an enabled request." + ) + # Biohub Transformers records the input to each block followed + # by the final normalized state. This gives n_layers + 1 states + # and, for ESMC-6B, the 81-state order consumed by ESMFold2. + hidden_states += (x,) + if self.gradient_checkpointing and self.training: + x, attn_weights, s_max = self._gradient_checkpointing_func( + block.__call__, + x=x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + else: + x, attn_weights, s_max = block( + x=x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + if attentions is not None: + attentions += (attn_weights,) + if full_s_max is not None: + full_s_max += (s_max,) + + last_hidden_state = self.norm(x) + if output_hidden_states: + hidden_states += (last_hidden_state,) + + return TransformerOutput( + last_hidden_state=last_hidden_state, + hidden_states=hidden_states, + attentions=attentions, + s_max=full_s_max, + ) + + def _sequence_id_attention_masks( + self, + sequence_id: torch.Tensor, + batch_size: int, + seq_len: int, + device: torch.device, + dtype: torch.dtype | None = None, + effective_backend: AttentionBackend | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]: + expected_shape = (batch_size, seq_len) + if sequence_id.ndim != 2 or tuple(sequence_id.shape) != expected_shape: + raise ValueError( + f"sequence_id must have shape {expected_shape}; " + f"received {tuple(sequence_id.shape)}." + ) + if sequence_id.device != device: + sequence_id = sequence_id.to(device=device) + backend = ( + self.attention_backend + if effective_backend is None + else resolve_attention_backend(effective_backend) + ) + if sequence_id.dtype == torch.bool: + attention_mask_2d = sequence_id + # Biohub's boolean single-chain form groups biological positions + # together and padding positions together. Padding queries remain + # finite without allowing their states to enter residue attention. + attention_mask_4d = sequence_id[:, None, :, None] == sequence_id[:, None, None, :] + else: + attention_mask_2d = sequence_id != -1 + attention_mask_4d = (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze( + 1 + ) + if not bool(attention_mask_2d.any(dim=1).all()): + raise ValueError("attention_mask must keep at least one valid key per batch row.") + + if backend.is_flash: + if sequence_id.dtype != torch.bool: + raise ValueError( + "ESM++ FlashAttention only supports boolean sequence_id padding masks. " + "Use eager, sdpa, or flex_attention for chain-aware integer sequence_id " + "masks." + ) + return attention_mask_2d, attention_mask_4d, None + + if backend == AttentionBackend.FLEX: + if sequence_id.dtype == torch.bool: + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + del head_idx + return sequence_id[batch_idx, q_idx] == sequence_id[batch_idx, kv_idx] + + else: + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + del head_idx + q_id = sequence_id[batch_idx, q_idx] + kv_id = sequence_id[batch_idx, kv_idx] + return q_id == kv_id + + flex_block_mask = _get_flex_block_mask( + mask_pattern=sequence_id, + batch_size=batch_size, + query_length=seq_len, + key_value_length=seq_len, + device=device, + dtype=dtype, + mask_semantics=( + "boolean_sequence_id" + if sequence_id.dtype == torch.bool + else "integer_sequence_id" + ), + mask_mod=mask_mod, + ) + return attention_mask_2d, attention_mask_4d, flex_block_mask + + return attention_mask_2d, attention_mask_4d, None + + +class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel): + """ + init weights for ESM++ models + """ + + config_class = ESMplusplusConfig + base_model_prefix = "esm++" + supports_gradient_checkpointing = True + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_flash_attn_3 = True + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + + @property + def tokenizer(self) -> EsmSequenceTokenizer: + """Construct the sequence tokenizer only when a raw-sequence API needs it.""" + + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + tokenizer = EsmSequenceTokenizer() + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value: EsmSequenceTokenizer | None) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + def _init_weights(self, module): + """Initialize the weights""" + # HF from_pretrained marks loaded parameters with `_is_hf_initialized`. + # Skip this module if any local parameter is already marked as loaded. + for parameter in module.parameters(recurse=False): + if parameter.__dict__.get("_is_hf_initialized"): + return + + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + with torch.no_grad(): + module.weight[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + if module.bias is not None: + nn.init.zeros_(module.bias) + nn.init.ones_(module.weight) + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"{type(self).__name__} does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.set_attn_implementation(backend) + + def _reset_rotary_embeddings(self): + """Refresh non-persistent rotary buffers after checkpoint loading.""" + for module in self.modules(): + if isinstance(module, RotaryEmbedding): + module.reset_parameters() + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + output_loading_info = ( + bool(kwargs["output_loading_info"]) if "output_loading_info" in kwargs else False + ) + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + if output_loading_info: + model, loading_info = loaded + model._reset_rotary_embeddings() + return model, loading_info + loaded._reset_rotary_embeddings() + return loaded + + +### ESM++ Models +class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin): + """ + ESM++ transformer backbone. + + Official ESM++ checkpoints contain the sequence head even when loaded through + ``AutoModel``. Keep that module in the base class so the checkpoint has one + exact state-dict contract across ``AutoModel`` and ``AutoModelForMaskedLM``; + the base forward path intentionally does not compute or return logits. + """ + + config_class = ESMplusplusConfig + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + PreTrainedESMplusplusModel.__init__(self, config, **kwargs) + self.config = config + self.vocab_size = config.vocab_size + self.embed = nn.Embedding(self.vocab_size, config.hidden_size) + self.transformer = TransformerStack( + d_model=config.hidden_size, + n_heads=config.num_attention_heads, + n_layers=config.num_hidden_layers, + dropout=config.dropout, + attn_backend=config.attn_backend, + ) + self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size) + self.init_weights() + + def get_input_embeddings(self): + return self.embed + + def set_input_embeddings(self, value): + self.embed = value + + def get_output_embeddings(self): + return self.sequence_head[-1] + + def set_output_embeddings(self, new_embeddings): + self.sequence_head[-1] = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + return_dict: bool | None = None, + ) -> TransformerOutput | tuple[torch.Tensor, ...]: + """Run ESMC inference with the pinned Biohub mask precedence. + + ``sequence_id`` is authoritative when supplied: non-negative integers + identify chains and ``-1`` identifies padding. In that mode + ``attention_mask`` is ignored, matching the official implementation. + Without ``sequence_id``, ``attention_mask`` is the ordinary padding + mask and defaults to ``input_ids != pad_token_id``. + """ + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if attention_mask is None and sequence_id is None and input_ids is not None: + attention_mask = input_ids.ne(self.config.pad_token_id) + + x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds + + transformer_output = self.transformer( + x=x, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + esmfold2_hidden_states=esmfold2_hidden_states, + ) + result = TransformerOutput( + last_hidden_state=transformer_output.last_hidden_state, + hidden_states=transformer_output.hidden_states, + attentions=transformer_output.attentions, + s_max=transformer_output.s_max, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForMaskedLM( + FastPLMTestTimeTrainingMixin, PreTrainedESMplusplusModel, EmbeddingMixin +): + """ + ESM++ model for masked language modeling. + Implements the base ESM++ architecture with a masked language modeling head. + """ + + config_class = ESMplusplusConfig + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + PreTrainedESMplusplusModel.__init__(self, config, **kwargs) + self.config = config + self.vocab_size = config.vocab_size + self.embed = nn.Embedding(self.vocab_size, config.hidden_size) + self.transformer = TransformerStack( + d_model=config.hidden_size, + n_heads=config.num_attention_heads, + n_layers=config.num_hidden_layers, + dropout=config.dropout, + attn_backend=config.attn_backend, + ) + self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size) + self.ce_loss = nn.CrossEntropyLoss() + self.init_weights() + self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) + + def get_input_embeddings(self): + return self.embed + + def set_input_embeddings(self, value): + self.embed = value + + def get_output_embeddings(self): + return self.sequence_head[-1] + + def set_output_embeddings(self, new_embeddings): + self.sequence_head[-1] = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.transformer] + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + return_dict: bool | None = None, + compute_logits: bool = True, + ) -> ESMplusplusOutput | tuple[torch.Tensor, ...]: + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if labels is not None and not compute_logits: + raise ValueError("labels require compute_logits=True.") + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + if attention_mask is None and sequence_id is None and input_ids is not None: + attention_mask = input_ids.ne(self.config.pad_token_id) + + x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds + + output = self.transformer( + x=x, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + esmfold2_hidden_states=esmfold2_hidden_states, + ) + + last_hidden_state = output.last_hidden_state + logits = self.sequence_head(last_hidden_state) if compute_logits else None + loss = None + if labels is not None: + if logits is None: + raise ValueError("labels require compute_logits=True.") + labels = labels.to(logits.device) + loss = self.ce_loss(logits.view(-1, self.vocab_size), labels.view(-1)) + + result = ESMplusplusOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + last_hidden_state=last_hidden_state, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixin): + """ + ESM++ model for sequence classification. + Extends the base ESM++ model with a classification head. + """ + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + pooling_types = kwargs.pop("pooling_types", None) + if pooling_types is None: + pooling_types = config.classifier_pooling_types or ["mean", "var"] + elif not isinstance(pooling_types, list): + raise TypeError("pooling_types must be a non-empty list of strings.") + elif not pooling_types: + raise ValueError("pooling_types must contain at least one pooling operation.") + elif not all(isinstance(pooling_type, str) for pooling_type in pooling_types): + raise TypeError("pooling_types must be a non-empty list of strings.") + if "parti" in pooling_types: + raise ValueError( + "pooling_types cannot contain 'parti' for sequence classification " + "because the classifier does not expose layer attentions to its pooler." + ) + config.classifier_pooling_types = list(pooling_types) + + ESMplusplusForMaskedLM.__init__(self, config, **kwargs) + self.config = config + self.num_labels = config.num_labels + self.classifier = RegressionHead( + config.hidden_size * len(pooling_types), + config.num_labels, + config.hidden_size * 4, + ) + # Large intermediate projections help with sequence classification tasks (*4) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.pooler = Pooler(pooling_types) + self.init_weights() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> ESMplusplusSequenceClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + pooling_mask = attention_mask + if pooling_mask is None: + if sequence_id is not None: + pooling_mask = ( + sequence_id if sequence_id.dtype == torch.bool else sequence_id.ne(-1) + ) + elif input_ids is not None: + pooling_mask = input_ids.ne(self.config.pad_token_id) + else: + if inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + pooling_mask = torch.ones( + inputs_embeds.shape[:2], + dtype=torch.bool, + device=inputs_embeds.device, + ) + + output = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + inputs_embeds=inputs_embeds, + labels=None, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + compute_logits=False, + ) + + last_hidden_state = output.last_hidden_state + features = self.pooler(last_hidden_state, pooling_mask) + logits = self.classifier(features) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.flatten(), labels.flatten()) + else: + loss = self.mse(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + result = ESMplusplusSequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin): + """ + ESM++ model for token classification. + Extends the base ESM++ model with a token classification head. + """ + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + ESMplusplusForMaskedLM.__init__(self, config, **kwargs) + self.config = config + self.num_labels = config.num_labels + self.classifier = RegressionHead( + config.hidden_size, config.num_labels, config.hidden_size * 4 + ) + # Large intermediate projections help with sequence classification tasks (*4) + self.loss_fct = nn.CrossEntropyLoss() + self.init_weights() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x, + attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + ) -> ESMplusplusTokenClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + output = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + inputs_embeds=inputs_embeds, + labels=None, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + compute_logits=False, + ) + + last_hidden_state = output.last_hidden_state + logits = self.classifier(last_hidden_state) + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + result = ESMplusplusTokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + ) + return result if return_dict else result.to_tuple() + + +### Tokenization +SEQUENCE_VOCAB = [ + "", + "", + "", + "", + "L", + "A", + "G", + "V", + "S", + "E", + "R", + "T", + "I", + "D", + "P", + "K", + "Q", + "N", + "F", + "Y", + "M", + "H", + "W", + "C", + "X", + "B", + "U", + "Z", + "O", + ".", + "-", + "|", + "", +] + + +def _build_sequence_tokenizer_backend( + *, + unk_token: str, + cls_token: str, + pad_token: str, + mask_token: str, + eos_token: str, + chain_break_token: str, +) -> Tokenizer: + """Build the fixed ESMC character vocabulary and boundary-token policy.""" + vocabulary = dict(zip(SEQUENCE_VOCAB, range(len(SEQUENCE_VOCAB)), strict=True)) + backend = Tokenizer(BPE(vocabulary, merges=[], unk_token=unk_token)) + backend.add_special_tokens([cls_token, pad_token, mask_token, eos_token, chain_break_token]) + backend.post_processor = TemplateProcessing( + single=" $A ", + pair=":0 $A:0 :0 $B:1 :1", + special_tokens=[ + ("", backend.token_to_id("")), + ("", backend.token_to_id("")), + ], + ) + return backend + + +class EsmSequenceTokenizer(PreTrainedTokenizerFast): + model_input_names: ClassVar[list[str]] = ["input_ids", "attention_mask"] + + def __init__( + self, + unk_token="", + cls_token="", + pad_token="", + mask_token="", + eos_token="", + chain_break_token="|", + **kwargs, + ): + backend = _build_sequence_tokenizer_backend( + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + chain_break_token=chain_break_token, + ) + self.cb_token = chain_break_token + super().__init__( + tokenizer_object=backend, + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + additional_special_tokens=[chain_break_token], + **kwargs, + ) + + # ESMC does not use BOS, so expose the sequence-start token through the HF BOS fields. + @property + def bos_token(self): + return self.cls_token + + @property + def bos_token_id(self): + return self.cls_token_id + + @property + def chain_break_token(self): + return self.cb_token + + @property + def chain_break_token_id(self): + return self.convert_tokens_to_ids(self.chain_break_token) + + @property + def all_token_ids(self): + return list(range(self.vocab_size)) + + @property + def special_token_ids(self): + return self.all_special_ids diff --git a/fastplms/e1/__init__.py b/src/fastplms/models/esmfold/__init__.py similarity index 100% rename from fastplms/e1/__init__.py rename to src/fastplms/models/esmfold/__init__.py diff --git a/src/fastplms/models/esmfold/modeling_fast_esmfold.py b/src/fastplms/models/esmfold/modeling_fast_esmfold.py new file mode 100644 index 0000000..88b246b --- /dev/null +++ b/src/fastplms/models/esmfold/modeling_fast_esmfold.py @@ -0,0 +1,875 @@ +"""FastESMFold with FastESM2 attention. + +Usage: + from transformers import AutoModel + model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True).cuda() + + # Basic folding, no TTT + result = model.fold_protein("MKTLLILAVVA...") + print(result["plddt"], result["pdb_string"][:100]) + +The runtime uses public Transformers folding components. It does not import or +depend on the pinned fair-esm or OpenFold parity repositories. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any +from einops import rearrange +from torch.nn import functional as F +from transformers.modeling_outputs import ModelOutput +from transformers.models.esm.configuration_esm import EsmConfig +from transformers.models.esm.modeling_esm import ( + EsmEmbeddings, + EsmIntermediate, + EsmOutput, + EsmSelfOutput, +) +from transformers.models.esm.modeling_esmfold import ( + EsmForProteinFolding, + collate_dense_tensors, +) +from transformers.models.esm.openfold_utils import residue_constants + +from fastplms.models._esm_rotary import RotaryEmbedding + + +# Hub composite artifacts define these shared names earlier in the assembled file. +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + flex_attention, + get_attention_mask, + kernels_flash_attention_func, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "FastPLMsAttentionMixin", + "_get_flex_attention_fn", + "flex_attention", + "get_attention_mask", + "kernels_flash_attention_func", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +@dataclass +class FastEsmEncoderOutput(ModelOutput): + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +@dataclass +class FastEsmForProteinFoldingOutput(ModelOutput): + """Folding output with a standard Transformers AutoModel prefix.""" + + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + frames: torch.Tensor | None = None + sidechain_frames: torch.Tensor | None = None + unnormalized_angles: torch.Tensor | None = None + angles: torch.Tensor | None = None + positions: torch.Tensor | None = None + states: torch.Tensor | None = None + s_s: torch.Tensor | None = None + s_z: torch.Tensor | None = None + distogram_logits: torch.Tensor | None = None + lm_logits: torch.Tensor | None = None + aatype: torch.Tensor | None = None + atom14_atom_exists: torch.Tensor | None = None + residx_atom14_to_atom37: torch.Tensor | None = None + residx_atom37_to_atom14: torch.Tensor | None = None + atom37_atom_exists: torch.Tensor | None = None + residue_index: torch.Tensor | None = None + lddt_head: torch.Tensor | None = None + plddt: torch.Tensor | None = None + ptm_logits: torch.Tensor | None = None + ptm: torch.Tensor | None = None + aligned_confidence_probs: torch.Tensor | None = None + predicted_aligned_error: torch.Tensor | None = None + max_predicted_aligned_error: torch.Tensor | None = None + mlm_targets: torch.Tensor | None = None + + +# ``EsmForProteinFolding.forward`` calls +# ``compute_language_model_representations`` without forwarding output controls. +# Context variables bridge that private call boundary without mutating the model +# instance, so concurrent calls can independently request attention tensors. +_ESMFOLD_OUTPUT_ATTENTIONS: ContextVar[bool] = ContextVar( + "fastplms_esmfold_output_attentions", + default=False, +) +_ESMFOLD_CAPTURED_ATTENTIONS: ContextVar[ + tuple[torch.Tensor, ...] | None +] = ContextVar( + "fastplms_esmfold_captured_attentions", + default=None, +) + + +def _align_internal_esm_attentions( + attentions: tuple[torch.Tensor, ...], + residue_mask: torch.Tensor, +) -> tuple[torch.Tensor, ...]: + """Remove internal BOS/EOS positions while preserving public padding slots.""" + + # attention: (b, h, l + 2, l + 2); residue_mask: (b, l) + batch_size, sequence_length = residue_mask.shape + public_positions = torch.arange(sequence_length, device=residue_mask.device) # (l,) + valid_lengths = residue_mask.to(dtype=torch.int64).sum(dim=-1, keepdim=True) # (b, 1) + # Internal layout is BOS, compact biological residues, EOS, then padding. + # Public padding positions therefore advance by two rather than one. + internal_positions = public_positions.unsqueeze(0) + 1 # (1, l) + internal_positions = internal_positions + ( + public_positions.unsqueeze(0) >= valid_lengths + ).to(dtype=torch.int64) + + aligned: list[torch.Tensor] = [] + for attention in attentions: + if attention.shape[0] != batch_size or attention.shape[-2:] != ( + sequence_length + 2, + sequence_length + 2, + ): + raise RuntimeError( + "FastESM attention shape does not match the folding input: " + f"got {tuple(attention.shape)} for residue mask " + f"{tuple(residue_mask.shape)}." + ) + query_index = internal_positions[:, None, :, None].expand( + batch_size, + attention.shape[1], + sequence_length, + attention.shape[-1], + ) + query_aligned = torch.gather(attention, dim=2, index=query_index) + key_index = internal_positions[:, None, None, :].expand( + batch_size, + attention.shape[1], + sequence_length, + sequence_length, + ) + aligned.append(torch.gather(query_aligned, dim=3, index=key_index)) + return tuple(aligned) + + +class EsmSelfAttention(nn.Module): + def __init__(self, config, position_embedding_type: str | None = None) -> None: + super().__init__() + if config.hidden_size % config.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number " + f"of attention heads ({config.num_attention_heads})." + ) + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + self.scale = self.attention_head_size**-0.5 + + self.dropout_prob = config.attention_probs_dropout_prob + self.config = config + self.attn_backend = resolve_attention_backend(config.attn_backend) + self.position_embedding_type = position_embedding_type or config.position_embedding_type + self.rotary_embeddings = None + if self.position_embedding_type == "rotary": + self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # hidden_states: (b, l, d) + batch_size, seq_length = hidden_states.shape[:-1] + hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) + query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2) # (b, h, l, d_h) + + query_heads = query_heads * self.scale + + if self.position_embedding_type == "rotary": + query_heads, key_heads = self.rotary_embeddings(query_heads, key_heads) + + attn_output, attn_weights = self._attn( + query_heads, + key_heads, + value_heads, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + ) + return attn_output, attn_weights + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_attentions: + return self._manual_attn(query_heads, key_heads, value_heads, attention_mask_4d) + + if ( + self.training + and self.dropout_prob > 0 + and self.attn_backend == AttentionBackend.FLEX_ATTENTION + ): + raise RuntimeError( + "ESMFold flex_attention is inference-only when attention dropout is " + "nonzero. Use eager or SDPA for this training configuration." + ) + + if self.attn_backend == AttentionBackend.EAGER: + attn_output, _ = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d + ) + return attn_output, None + if self.attn_backend.is_flash: + return self._kernels_flash_attn(query_heads, key_heads, value_heads, attention_mask_2d) + elif self.attn_backend == AttentionBackend.FLEX: + return self._flex_attn( + query_heads, + key_heads, + value_heads, + flex_block_mask, + attention_mask_2d, + ) + elif self.attn_backend == AttentionBackend.SDPA: + return self._sdpa_attn(query_heads, key_heads, value_heads, attention_mask_4d) + else: + raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = torch.matmul( + query_heads, key_heads.transpose(-1, -2) + ) # (b, h, l, l) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) + attn_weights = F.softmax(attn_weights, dim=-1) + if self.dropout_prob > 0 and self.training: + attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l, d) + return attn_output, attn_weights + + def _kernels_flash_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + query_tokens = query_heads.transpose(1, 2).contiguous() + key_tokens = key_heads.transpose(1, 2).contiguous() + value_tokens = value_heads.transpose(1, 2).contiguous() + # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0 + # to prevent the kernel from applying its default 1/sqrt(head_dim). + attn_output = kernels_flash_attention_func( + query_states=query_tokens, + key_states=key_tokens, + value_states=value_tokens, + attention_mask_2d=attention_mask_2d, + causal=False, + softmax_scale=1.0, + implementation=self.attn_backend.value, + ) + return rearrange(attn_output, "b s h d -> b s (h d)"), None + + def _flex_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + flex_block_mask: BlockMask | None = None, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if flex_attention is None: + raise RuntimeError("Flex attention is not available in this environment.") + fn = _get_flex_attention_fn( + device=query_heads.device, + dtype=query_heads.dtype, + shape=tuple(query_heads.shape), + mask_semantics="padding", + ) + context_heads = fn( + query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0 + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + dropout_p=self.dropout_prob if self.training else 0.0, + scale=1.0, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +class EsmAttention(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.self = EsmSelfAttention(config) + self.output = EsmSelfOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + hidden_states_ln = self.LayerNorm(hidden_states) + attn_output, attn_weights = self.self( + hidden_states_ln, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + ) + attention_output = self.output(attn_output, hidden_states) + return attention_output, attn_weights + + +class EsmLayer(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.attention = EsmAttention(config) + self.intermediate = EsmIntermediate(config) + self.output = EsmOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + attention_output, attn_weights = self.attention( + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + ) + layer_output = self._feed_forward(attention_output) + return layer_output, attn_weights + + def _feed_forward(self, attention_output: torch.Tensor) -> torch.Tensor: + attention_output_ln = self.LayerNorm(attention_output) + intermediate_output = self.intermediate(attention_output_ln) + return self.output(intermediate_output, attention_output) + + +class FastEsmEncoder(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.attention_backend = resolve_attention_backend(config.attn_backend) + self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)]) + self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool = False, + output_attentions: bool = False, + ) -> FastEsmEncoderOutput: + # hidden_states: (b, l, d); attention_mask: (b, l) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + effective_backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( + effective_backend=effective_backend, + batch_size=hidden_states.shape[0], + seq_len=hidden_states.shape[1], + device=hidden_states.device, + attention_mask=attention_mask, + dtype=hidden_states.dtype, + mask_semantics="padding", + ) + + for layer_module in self.layer: + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + hidden_states, attn_weights = layer_module( + hidden_states, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + ) + + if all_attentions is not None: + all_attentions = (*all_attentions, attn_weights) + + if self.emb_layer_norm_after: + hidden_states = self.emb_layer_norm_after(hidden_states) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + return FastEsmEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + ) + + +class FastEsmBackbone(nn.Module): + """FastESM2 backbone with multi-backend attention. Drop-in replacement for + transformers.EsmModel inside EsmForProteinFolding. + + Folding uses hidden states only. The standalone ESM2 contact regressor and + masked-LM head are therefore omitted from this structure-only backbone. + """ + + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.embeddings = EsmEmbeddings(config) + self.encoder = FastEsmEncoder(config) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> FastEsmEncoderOutput | tuple[Any, ...]: + output_attentions = ( + self.config.output_attentions + if output_attentions is None + else output_attentions + ) + output_hidden_states = ( + self.config.output_hidden_states + if output_hidden_states is None + else output_hidden_states + ) + return_dict = self.config.use_return_dict if return_dict is None else return_dict + + token_embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + token_embedding_output, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + ) + output = FastEsmEncoderOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + return output if return_dict else output.to_tuple() + + +class FastEsmFoldConfig(EsmConfig): + model_type = "fast_esmfold" + + def __init__(self, attn_backend: str | None = None, **kwargs: Any) -> None: + # Earlier mirrors serialized an untrained ESMFold-specific TTT policy. + # It is intentionally ignored because the official checkpoint has no + # trained masked-language-model head. + kwargs.pop("ttt_config", None) + super().__init__(**kwargs) + self.attn_backend = attn_backend + + +class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding): + """ESMFold with FastESM2 attention backends. + + Inherits all folding logic (trunk, structure module, output_to_pdb, infer) + from transformers.EsmForProteinFolding. Replaces the ESM2 backbone with + FastESM2 for selectable attention implementations. + + Key API: + result = model.fold_protein("MKTL...") + # result = {"plddt": float, "ptm": float, "pdb_string": str} + """ + + config_class = FastEsmFoldConfig + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ("eager", "sdpa", "flex_attention") + + def __init__(self, config: FastEsmFoldConfig) -> None: + super().__init__(config) + + # Replace the standard ESM2 backbone with the multi-backend FastESM2 + # implementation while retaining the canonical checkpoint key schema. + self.esm = FastEsmBackbone(config) + self.esm.requires_grad_(False) + if config.esmfold_config.fp16_esm: + self.esm.half() + + def compute_language_model_representations( + self, + esmaa: torch.Tensor, + ) -> torch.Tensor: + """Run the internal ESM stem with a structured output unconditionally. + + The outer folding model still honors ``config.return_dict``. This + internal call must remain structured because the folding stem selects + hidden states by name before constructing the public output. + """ + + device = next(self.parameters()).device + batch_size, sequence_length = esmaa.shape + output_attentions = _ESMFOLD_OUTPUT_ATTENTIONS.get() + if self.config.esmfold_config.bypass_lm: + if output_attentions: + _ESMFOLD_CAPTURED_ATTENTIONS.set(()) + return torch.zeros( + batch_size, + sequence_length, + self.esm_s_combine.size(0), + self.esm_feats, + device=device, + ) + + bos = esmaa.new_full((batch_size, 1), self.esm_dict_cls_idx) + eos = esmaa.new_full((batch_size, 1), self.esm_dict_padding_idx) + residue_mask = esmaa != self.esm_dict_padding_idx + with_special_tokens = torch.cat([bos, esmaa, eos], dim=1) + with_special_tokens[ + range(batch_size), + (with_special_tokens != self.esm_dict_padding_idx).sum(1), + ] = self.esm_dict_eos_idx + esm_output = self.esm( + with_special_tokens, + attention_mask=with_special_tokens != self.esm_dict_padding_idx, + output_attentions=output_attentions, + output_hidden_states=True, + return_dict=True, + ) + if not isinstance(esm_output, FastEsmEncoderOutput): + raise TypeError("FastESM internal backbone did not return FastEsmEncoderOutput.") + if esm_output.hidden_states is None: + raise RuntimeError("FastESM internal backbone omitted requested hidden states.") + if output_attentions: + if esm_output.attentions is None: + raise RuntimeError("FastESM internal backbone omitted requested attentions.") + _ESMFOLD_CAPTURED_ATTENTIONS.set( + _align_internal_esm_attentions(esm_output.attentions, residue_mask) + ) + return torch.stack(esm_output.hidden_states, dim=2)[:, 1:-1] + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + masking_pattern: torch.Tensor | None = None, + num_recycles: int | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> FastEsmForProteinFoldingOutput | tuple[Any, ...]: + """Run folding with Meta ESMFold's 0-to-100 pLDDT convention.""" + + config = getattr(self, "config", None) + resolved_attentions = ( + bool(getattr(config, "output_attentions", False)) + if output_attentions is None + else output_attentions + ) + resolved_hidden_states = ( + bool(getattr(config, "output_hidden_states", False)) + if output_hidden_states is None + else output_hidden_states + ) + resolved_return_dict = ( + bool(getattr(config, "use_return_dict", True)) + if return_dict is None + else return_dict + ) + + request_token = _ESMFOLD_OUTPUT_ATTENTIONS.set(bool(resolved_attentions)) + capture_token = _ESMFOLD_CAPTURED_ATTENTIONS.set(None) + captured_attentions: tuple[torch.Tensor, ...] | None = None + try: + output = super().forward( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + masking_pattern=masking_pattern, + num_recycles=num_recycles, + output_hidden_states=resolved_hidden_states, + ) + captured_attentions = _ESMFOLD_CAPTURED_ATTENTIONS.get() + finally: + _ESMFOLD_CAPTURED_ATTENTIONS.reset(capture_token) + _ESMFOLD_OUTPUT_ATTENTIONS.reset(request_token) + # Transformers 5.13 returns categorical lDDT probabilities on [0, 1], + # while Meta ESMFold's public forward output reports pLDDT on [0, 100]. + output["plddt"] = output["plddt"] * 100 + payload = dict(output) + payload.pop("last_hidden_state", None) + payload.pop("hidden_states", None) + parent_attentions = payload.pop("attentions", None) + if captured_attentions is None: + captured_attentions = parent_attentions + if resolved_attentions and captured_attentions is None: + # A bypassed or injected folding stem has no attention layers, but + # still honors the output contract without fabricating tensors. + captured_attentions = () + sequence_state = payload.get("s_s") + structured = FastEsmForProteinFoldingOutput( + last_hidden_state=sequence_state, + hidden_states=(sequence_state,) + if resolved_hidden_states and sequence_state is not None + else None, + attentions=captured_attentions if resolved_attentions else None, + **payload, + ) + return structured if resolved_return_dict else structured.to_tuple() + + @torch.no_grad() + def infer( + self, + sequences: str | list[str], + residx: torch.Tensor | list[torch.Tensor] | None = None, + masking_pattern: torch.Tensor | None = None, + num_recycles: int | None = None, + residue_index_offset: int | None = 512, + chain_linker: str | None = "G" * 25, + ): + """Fold raw sequences through Meta ESMFold's public input contract. + + Transformers v5 narrows ``infer`` even though ``forward`` retains the + required controls. This adapter restores recycle selection, explicit + residue indices, masking, and colon-delimited multimer preparation. + """ + + sequence_batch = [sequences] if isinstance(sequences, str) else sequences + linker = "" if chain_linker is None else chain_linker + index_offset = 0 if residue_index_offset is None else residue_index_offset + unknown_index = residue_constants.restype_order_with_x["X"] + aatype_batch: list[torch.Tensor] = [] + residx_batch: list[torch.Tensor] = [] + linker_mask_batch: list[torch.Tensor] = [] + chain_index_batch: list[torch.Tensor] = [] + + for sequence in sequence_batch: + chains = sequence.split(":") + joined_sequence = linker.join(chains) + encoded = torch.tensor( + [ + residue_constants.restype_order_with_x.get(residue, unknown_index) + for residue in joined_sequence + ], + dtype=torch.int64, + ) + sequence_residx = torch.arange(len(encoded), dtype=torch.int64) + cursor = 0 + for chain_number, chain in enumerate(chains): + segment_length = len(chain) + len(linker) + sequence_residx[cursor : cursor + segment_length] += chain_number * index_offset + cursor += segment_length + + linker_mask = torch.ones_like(encoded, dtype=torch.float32) + chain_indices: list[int] = [] + cursor = 0 + for chain_number, chain in enumerate(chains): + if chain_number > 0: + chain_indices.extend([chain_number - 1] * len(linker)) + chain_indices.extend([chain_number] * len(chain)) + cursor += len(chain) + linker_mask[cursor : cursor + len(linker)] = 0 + cursor += len(linker) + + aatype_batch.append(encoded) + residx_batch.append(sequence_residx) + linker_mask_batch.append(linker_mask) + chain_index_batch.append(torch.tensor(chain_indices, dtype=torch.int64)) + + aatype = collate_dense_tensors(aatype_batch) + attention_mask = collate_dense_tensors( + [aatype.new_ones(len(encoded)) for encoded in aatype_batch] + ) + prepared_residx = collate_dense_tensors(residx_batch) + linker_mask = collate_dense_tensors(linker_mask_batch) + chain_index = collate_dense_tensors(chain_index_batch, pad_v=-1) + if residx is None: + residx = prepared_residx + elif not isinstance(residx, torch.Tensor): + residx = collate_dense_tensors(residx) + + device = next(self.parameters()).device + aatype = aatype.to(device) + attention_mask = attention_mask.to(device) + residx = residx.to(device) + linker_mask = linker_mask.to(device) + output = self.forward( + aatype, + attention_mask, + position_ids=residx, + masking_pattern=masking_pattern, + num_recycles=num_recycles, + ) + output["atom37_atom_exists"] = output["atom37_atom_exists"] * linker_mask.unsqueeze(2) + output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum( + dim=(1, 2) + ) / output["atom37_atom_exists"].sum(dim=(1, 2)) + output["chain_index"] = chain_index + return output + + @staticmethod + def _ttt_unavailable() -> None: + raise RuntimeError( + "ESMFold TTT is unavailable: the pinned Meta ESMFold checkpoint does " + "not contain a trained masked-language-model head. FastPLMs does not " + "construct or serialize a random replacement head." + ) + + def ttt(self, seq: str, **kwargs: Any) -> None: + """Reject ESMFold-specific TTT because no faithful MLM objective exists.""" + + del seq, kwargs + self._ttt_unavailable() + + def ttt_reset(self) -> None: + """Reject reset because ESMFold does not expose a faithful TTT path.""" + + self._ttt_unavailable() + + def _fold_single(self, sequence: str, return_pdb_string: bool = True) -> dict[str, Any]: + """Fold a sequence once and return pLDDT, ptm, and optionally PDB string.""" + with torch.no_grad(): + output = self.infer(sequence) + if "mean_plddt" in output: + # ``infer`` masks multimer linker atoms before computing this mean. + # Reusing it prevents synthetic linker confidence from affecting the + # public fold_protein summary. + mean_plddt = float(output["mean_plddt"].reshape(-1)[0].item()) + else: + plddt = output["plddt"] + # P has shape (b, l, 37), with confidence for each atom37 position. + # Use CA atom (index 1) only, matching PDB B-factor output. + if plddt.dim() == 3: + mean_plddt = float(plddt[:, :, 1].mean().item()) + elif plddt.dim() == 2: + mean_plddt = float(plddt[:, 1].mean().item()) + else: + mean_plddt = float(plddt.mean().item()) + result = { + "plddt": mean_plddt, + "ptm": float(output["ptm"].item()) if "ptm" in output else None, + } + if return_pdb_string: + pdb_strings = self.output_to_pdb(output) + result["pdb_string"] = pdb_strings[0] if isinstance(pdb_strings, list) else pdb_strings + return result + + def fold_protein( + self, + sequence: str, + return_pdb_string: bool = True, + ttt: bool = False, + ) -> dict[str, Any]: + """Fold a protein sequence. + + Passing ``ttt=True`` fails explicitly because the official Meta + checkpoint contains no trained masked-language-model head. + + Args: + sequence: Protein sequence (single-letter amino acid codes) + return_pdb_string: If True, include PDB string in output + ttt: Reserved rejection flag for unsupported ESMFold TTT + + Returns: + Dict with keys: + - plddt: float, mean pLDDT + - ptm: float, predicted TM-score + - pdb_string: str (if return_pdb_string=True), PDB from best step + - step_plddts: list[float], baseline pLDDT when TTT is disabled + - best_step: int, 0 when TTT is disabled + """ + if ttt: + return self.fold_protein_ttt( + sequence=sequence, + return_pdb_string=return_pdb_string, + ) + result = self._fold_single(sequence, return_pdb_string=return_pdb_string) + return { + "plddt": result["plddt"], + "ptm": result["ptm"], + "pdb_string": result.get("pdb_string"), + "step_plddts": [result["plddt"]], + "best_step": 0, + } + + def fold_protein_ttt( + self, + sequence: str, + return_pdb_string: bool = True, + ) -> None: + """Reject ESMFold TTT because the official checkpoint has no MLM head.""" + + del sequence, return_pdb_string + self._ttt_unavailable() diff --git a/src/fastplms/models/esmfold2/__init__.py b/src/fastplms/models/esmfold2/__init__.py new file mode 100644 index 0000000..5aefbd1 --- /dev/null +++ b/src/fastplms/models/esmfold2/__init__.py @@ -0,0 +1,39 @@ +"""ESMFold2 public classes, imported lazily to keep optional extras isolated.""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .configuration_esmfold2 import ESMFold2Config as ESMFold2Config + from .modeling_esmfold2 import ESMFold2Model as ESMFold2Model + from .modeling_esmfold2 import ESMFold2Output as ESMFold2Output + from .modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel as ESMFold2ExperimentalModel, + ) + from .reproducibility import seed_context as seed_context + +_EXPORT_MODULES = { + "ESMFold2Config": ".configuration_esmfold2", + "ESMFold2ExperimentalModel": ".modeling_esmfold2_experimental", + "ESMFold2Model": ".modeling_esmfold2", + "ESMFold2Output": ".modeling_esmfold2", + "seed_context": ".reproducibility", +} + + +def __getattr__(name: str) -> Any: + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(module_name, __name__), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_EXPORT_MODULES)) + + +__all__ = list(_EXPORT_MODULES) diff --git a/src/fastplms/models/esmfold2/attention.py b/src/fastplms/models/esmfold2/attention.py new file mode 100644 index 0000000..5823fc6 --- /dev/null +++ b/src/fastplms/models/esmfold2/attention.py @@ -0,0 +1,48 @@ +"""Transformers-compatible attention selection for ESMFold2's ESMC backbone.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from ...attention import FastPLMsAttentionMixin, get_attn_implementation + + +class ESMFold2AttentionMixin(FastPLMsAttentionMixin): + """Route the outer Transformers attention API into the loaded ESMC model.""" + + _supports_attention_backend = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + ) + + def __init__(self, config, *args, **kwargs) -> None: + super().__init__(config, *args, **kwargs) + config.esmc_attn_backend = get_attn_implementation(config) + + def set_attn_implementation( + self, + attn_implementation: str | Mapping[str, str], + allow_all_kernels: bool = False, + ) -> None: + """Set one canonical backend on ESMFold2 and its loaded ESMC model.""" + + if allow_all_kernels: + raise ValueError( + "ESMFold2 accepts only its declared built-in attention backends; " + "external attention kernels are not supported." + ) + super().set_attn_implementation(attn_implementation) + resolved = get_attn_implementation(self.config) + self.config.esmc_attn_backend = resolved + esmc = getattr(self, "_esmc", None) + if esmc is not None: + esmc.set_attn_implementation(resolved) + + +__all__ = ["ESMFold2AttentionMixin"] diff --git a/src/fastplms/models/esmfold2/configuration_esmfold2.py b/src/fastplms/models/esmfold2/configuration_esmfold2.py new file mode 100644 index 0000000..b2bbd62 --- /dev/null +++ b/src/fastplms/models/esmfold2/configuration_esmfold2.py @@ -0,0 +1,306 @@ +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration schema for release and experimental ESMFold2 checkpoints.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any, TypeVar, cast + +from transformers.configuration_utils import PretrainedConfig + +_ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"}) +_ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"}) + + +def _esmc_backbone_checkpoint_ids() -> tuple[str, str]: + """Return the manifest-pinned official and FastPLMs ESMC repositories.""" + + from fastplms.registry import RegistryError, get_model_registry + + registry = get_model_registry() + family = registry.families["esmfold2"] + if family.backbone_model is None: + raise RegistryError("families.esmfold2 must declare backbone_model.") + backbone = registry[family.backbone_model] + return backbone.official.repo_id, backbone.fast.repo_id + + +def normalize_esmc_id(esmc_id: str) -> str: + """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror.""" + + official_repo, fast_repo = _esmc_backbone_checkpoint_ids() + return fast_repo if esmc_id == official_repo else esmc_id + + +def normalize_esmc_attention_implementation( + implementation: str | dict[str, str] | None, +) -> str | None: + """Validate the ESMC backend and translate the historical ``flex`` name.""" + + if isinstance(implementation, dict): + if tuple(implementation) != ("",): + raise ValueError( + "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}." + ) + implementation = implementation[""] + canonical = "flex_attention" if implementation == "flex" else implementation + if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS: + expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS) + raise ValueError( + f"Unsupported ESMFold2 attention implementation {canonical!r}; " + f"expected one of {expected}." + ) + return canonical + + +NestedConfig = TypeVar("NestedConfig") + + +def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig: + if isinstance(value, config_type): + return value + return config_type(**value) if isinstance(value, dict) else config_type() + + +def _coerce_nested_field( + value: NestedConfig | dict[str, Any], config_type: type[NestedConfig] +) -> NestedConfig: + """Convert serialized nested dictionaries while retaining supplied objects.""" + + return config_type(**value) if isinstance(value, dict) else value + + +@dataclass +class AtomAttentionConfig: + """Sliding-window atom attention and three-dimensional RoPE settings.""" + + d_atom: int = field(default=128) + d_token: int = field(default=768) + n_blocks: int = field(default=3) + n_heads: int = field(default=4) + swa_window_size: int = field(default=128) + expansion_ratio: int = field(default=2) + spatial_rope_base_frequency: float = field(default=20.0) + n_spatial_rope_pairs_per_axis: int = field(default=2) + n_uid_rope_pairs: int = field(default=10) + uid_rope_base_frequency: float = field(default=10000.0) + + +@dataclass +class DiffusionModuleConfig: + """Dimensions and depth of the coordinate diffusion network.""" + + sigma_data: float = field(default=16.0) + c_atom: int = field(default=128) + c_token: int = field(default=768) + c_z: int = field(default=256) + c_s_inputs: int = field(default=451) + fourier_dim: int = field(default=256) + relpos_r_max: int = field(default=32) + relpos_s_max: int = field(default=2) + atom_num_blocks: int = field(default=3) + atom_num_heads: int = field(default=4) + token_num_blocks: int = field(default=12) + token_num_heads: int = field(default=16) + transition_multiplier: int = field(default=2) + + +@dataclass +class FoldingTrunkConfig: + """Iterative pair/single trunk dimensions.""" + + n_layers: int = field(default=24) + n_heads: int = field(default=8) + dropout: float = field(default=0.0) + + +@dataclass +class InputsEmbedderConfig: + """Input feature width and atom encoder settings.""" + + d_inputs: int = field(default=451) + atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig) + + def __post_init__(self) -> None: + self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig) + + +@dataclass +class DiffusionStructureHeadConfig: + """Training and inference schedules for coordinate denoising.""" + + diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig) + distogram_bins: int = field(default=128) + train_noise_log_mean: float = field(default=-1.2) + train_noise_log_std: float = field(default=1.5) + gamma_0: float = field(default=0.605) + gamma_min: float = field(default=1.107) + noise_scale: float = field(default=0.0) + step_scale: float = field(default=1.0) + inference_s_max: float = field(default=160.0) + inference_s_min: float = field(default=4e-4) + inference_p: float = field(default=8.0) + inference_num_steps: int = field(default=68) + + def __post_init__(self) -> None: + self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig) + + +@dataclass +class ConfidenceHeadConfig: + """Confidence-bin definitions and the compact confidence trunk.""" + + enabled: bool = field(default=True) + num_plddt_bins: int = field(default=50) + num_pde_bins: int = field(default=64) + num_pae_bins: int = field(default=64) + min_dist: float = field(default=2.0) + max_dist: float = field(default=52.0) + distogram_bins: int = field(default=128) + folding_trunk: FoldingTrunkConfig = field( + default_factory=lambda: FoldingTrunkConfig(n_layers=4) + ) + + def __post_init__(self) -> None: + self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig) + + +@dataclass +class MSAEncoderConfig: + """Optional multiple-sequence-alignment encoder settings.""" + + enabled: bool = field(default=False) + d_msa: int = field(default=128) + d_hidden: int = field(default=32) + n_layers: int = field(default=4) + n_heads_msa: int = field(default=8) + msa_head_width: int = field(default=32) + + +@dataclass +class LMEncoderConfig: + """Release-model pair encoder derived from language-model states.""" + + enabled: bool = field(default=True) + n_layers: int = field(default=4) + lm_dropout: float = field(default=0.25) + per_loop_lm_dropout: bool = field(default=True) + + +@dataclass +class ParcaeConfig: + """Release-model diffusion-loop scheduler settings.""" + + enabled: bool = field(default=True) + poisson_mean: float = field(default=3.0) + min_steps: int = field(default=1) + max_steps: int | None = field(default=6) + coda_n_layers: int = field(default=2) + + +_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = ( + ("d_single", 384), + ("d_pair", 256), + ("n_relative_residx_bins", 32), + ("n_relative_chain_bins", 2), + ("num_loops", 10), + ("num_diffusion_samples", 8), + ("disable_msa_features", False), + ("lm_dropout", 0.0), + ("force_lm_dropout_during_inference", False), + ("lm_mask_pct", 0.0), + ("lm_d_model", 2560), + ("lm_num_layers", 80), +) +_NESTED_CONFIGS = ( + ("inputs", InputsEmbedderConfig), + ("folding_trunk", FoldingTrunkConfig), + ("structure_head", DiffusionStructureHeadConfig), + ("confidence_head", ConfidenceHeadConfig), + ("msa_encoder", MSAEncoderConfig), + ("parcae", ParcaeConfig), + ("lm_encoder", LMEncoderConfig), +) + + +class ESMFold2Config(PretrainedConfig): + """Serializable ESMFold2 architecture, runtime, and precision settings.""" + + model_type = "esmfold2" + has_no_defaults_at_init = True + + def __init__(self, **kwargs: Any) -> None: + legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend")) + requested_backend = normalize_esmc_attention_implementation( + kwargs.get("attn_implementation") + ) + resolved_backend = requested_backend or legacy_backend + kwargs["attn_implementation"] = resolved_backend + super().__init__(**kwargs) + + self.type = kwargs.get("type", "release") + if self.type not in {"experimental", "release"}: + raise ValueError( + f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}" + ) + + for name, default in _SCALAR_DEFAULTS: + setattr(self, name, kwargs.get(name, default)) + + _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids() + self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo)) + self.esmc_attn_backend = resolved_backend + self.esmc_precision = str(kwargs.get("esmc_precision", "auto")) + if self.esmc_precision not in _ESMC_PRECISIONS: + raise ValueError( + "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', " + f"got {self.esmc_precision!r}." + ) + + for name, config_type in _NESTED_CONFIGS: + setattr(self, name, _nested_config(kwargs.get(name), config_type)) + if not isinstance(self.msa_encoder.enabled, bool): + raise TypeError("msa_encoder.enabled must be a boolean.") + declared_msa_conditioning = kwargs.get("msa_conditioning") + if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool): + raise TypeError("msa_conditioning must be a boolean when provided.") + self.msa_conditioning = ( + self.msa_encoder.enabled + if "msa_conditioning" not in kwargs + else declared_msa_conditioning + ) + if self.msa_conditioning != self.msa_encoder.enabled: + raise ValueError( + "msa_conditioning must match msa_encoder.enabled; received " + f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}." + ) + self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True)) + + def to_dict(self) -> dict[str, Any]: + output = cast(dict[str, Any], super().to_dict()) + for name, _config_type in _NESTED_CONFIGS: + output[name] = asdict(getattr(self, name)) + return output + + +__all__ = [ + "ESMFold2Config", + "LMEncoderConfig", + "MSAEncoderConfig", + "ParcaeConfig", + "normalize_esmc_attention_implementation", + "normalize_esmc_id", +] diff --git a/src/fastplms/models/esmfold2/embedding.py b/src/fastplms/models/esmfold2/embedding.py new file mode 100644 index 0000000..e22160f --- /dev/null +++ b/src/fastplms/models/esmfold2/embedding.py @@ -0,0 +1,115 @@ +"""ESMFold2 integration for the shared FastPLMs embedding API.""" + +from __future__ import annotations + +import torch +from typing import Any, ClassVar +from torch import Tensor + +from ...embeddings import EmbeddingBatch, EmbeddingResult, embed_dataset +from .esmfold2_constants_esm3 import SEQUENCE_PAD_TOKEN, SEQUENCE_VOCAB + + +_TOKEN_TO_ID = {token: index for index, token in enumerate(SEQUENCE_VOCAB)} +_VALID_RESIDUES = frozenset(SEQUENCE_VOCAB[4:31]) - {".", "-", "|"} + + +def _encode_single_chain(sequence: str) -> list[int]: + normalized = sequence.upper() + if not normalized: + raise ValueError("ESMFold2 dataset embedding requires at least one protein residue.") + invalid = sorted(set(normalized) - _VALID_RESIDUES) + if invalid: + raise ValueError( + "ESMFold2 dataset embedding accepts one ungapped protein chain; " + f"invalid symbols: {invalid}." + ) + return [_TOKEN_TO_ID[residue] for residue in normalized] + + +class ESMFold2EmbeddingMixin: + """Learned ESMC sequence summaries for ESMFold2 models.""" + + embedding_unsupported_pooling = frozenset({"cls", "parti"}) + embedding_layer = "all_81_esmc_states" + embedding_projection = "esmfold2_learned_sequence_summary" + embedding_token_policy: ClassVar[dict[str, object]] = { + "unit": "residue", + "normalization": "uppercase", + "include": ["single-chain protein residues"], + "exclude": [ + "BOS", + "EOS", + "padding", + "chain delimiters", + "non-protein tokens", + ], + } + + def project_esmc_hidden_states( + self, + hidden_states: Tensor, + residue_mask: Tensor | None = None, + ) -> Tensor: + """Project H from ``(b, l, 81, 2560)`` to Z with shape ``(b, l, 256)``.""" + + # hidden_states: (b, l, 81, d_model); residue_mask: (b, l) or None. + # d_z is the learned projection width: 256 for released ESMFold2 checkpoints. + if hidden_states.ndim != 4 or hidden_states.shape[-2] != 81: + raise ValueError( + "ESMFold2 projection requires the official ordered 81-state " + "ESMC tensor H with shape (b, l, 81, d_model)." + ) + return self.language_model.project_sequence( + hidden_states, residue_mask + ) # (b, l, d_z) + + def _embedding_batch(self, sequences: list[str], **kwargs: Any) -> EmbeddingBatch: + if kwargs: + raise TypeError(f"Unexpected ESMFold2 embedding options: {', '.join(sorted(kwargs))}.") + if self._esmc is None: + raise RuntimeError("ESMFold2 embeddings require load_esmc=True.") + encoded = [_encode_single_chain(sequence) for sequence in sequences] + sequence_length = max(map(len, encoded)) # l + b = len(encoded) # b + device = self.device + input_ids = torch.full( + (b, sequence_length), + SEQUENCE_PAD_TOKEN, + dtype=torch.long, + device=device, + ) # (b, l) + residue_mask = torch.zeros( + (b, sequence_length), dtype=torch.bool, device=device + ) # (b, l) + for batch_index, token_ids in enumerate(encoded): + length = len(token_ids) # l_i + input_ids[batch_index, :length] = torch.tensor( + token_ids, dtype=torch.long, device=device + ) # (l_i,) -> input_ids[batch_index, :l_i]: (l_i,) + residue_mask[batch_index, :length] = True # (l_i,) + + residue_index = torch.arange(sequence_length, device=device).expand(b, -1) # (b, l) + asym_id = torch.zeros_like(input_ids) # (b, l) + mol_type = torch.zeros_like(input_ids) # (b, l) + hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + residue_mask, + ) # (b, l, 81, d_model) + projected = self.project_esmc_hidden_states( + hidden_states, residue_mask + ) # (b, l, d_z) + return EmbeddingBatch( + X=projected, residue_mask=residue_mask + ) # X: (b, l, d_z); residue_mask: (b, l) + + def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult: + """Embed single-chain proteins using the learned 256-wide ESMFold2 summary.""" + + return embed_dataset(self, inputs, **kwargs) + + +__all__ = ["ESMFold2EmbeddingMixin"] diff --git a/src/fastplms/models/esmfold2/esmfold2_affine3d.py b/src/fastplms/models/esmfold2/esmfold2_affine3d.py new file mode 100644 index 0000000..17c47f2 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_affine3d.py @@ -0,0 +1,605 @@ +"""Differentiable rigid rotations and affine transforms for ESMFold2.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Self + +import torch +from torch.nn import functional as F + +from .esmfold2_misc import fp32_autocast_context + + +def _index_tuple(index: Any) -> tuple[Any, ...]: + if isinstance(index, int) or index is None: + return (index,) + return tuple(index) + + +def _sqrt_subgradient(values: torch.Tensor) -> torch.Tensor: + """Square root with a zero subgradient for non-positive inputs.""" + + result = torch.zeros_like(values) + positive = values > 0 + result[positive] = torch.sqrt(values[positive]) + return result + + +def _quat_invert(quaternion: torch.Tensor) -> torch.Tensor: + conjugate_sign = torch.tensor([1, -1, -1, -1], device=quaternion.device) + return quaternion * conjugate_sign + + +def _quat_mult(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + """Hamilton product for real-first quaternion tensors.""" + + aw, ax, ay, az = torch.unbind(left, -1) + bw, bx, by, bz = torch.unbind(right, -1) + return torch.stack( + ( + aw * bw - ax * bx - ay * by - az * bz, + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + ), + -1, + ) + + +def _quat_rotation( + quaternion: torch.Tensor, + points: torch.Tensor, +) -> torch.Tensor: + """Rotate points using normalized real-first quaternions.""" + + aw, ax, ay, az = torch.unbind(quaternion, -1) + bx, by, bz = torch.unbind(points, -1) + product = torch.stack( + ( + -ax * bx - ay * by - az * bz, + aw * bx + ay * bz - az * by, + aw * by - ax * bz + az * bx, + aw * bz + ax * by - ay * bx, + ), + -1, + ) + return _quat_mult(product, _quat_invert(quaternion))[..., 1:] + + +def _graham_schmidt( + x_axis: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-12, +) -> torch.Tensor: + """Construct a right-handed orthonormal frame from two directions.""" + + with fp32_autocast_context(x_axis.device.type): + e1 = xy_plane + denominator = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps) + x_axis = x_axis / denominator + projection = (x_axis * e1).sum(dim=-1, keepdim=True) + e1 = e1 - x_axis * projection + denominator = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps) + e1 = e1 / denominator + e2 = torch.cross(x_axis, e1, dim=-1) + return torch.stack([x_axis, e1, e2], dim=-1) + + +class Rotation: + """Common interface for matrix-backed and quaternion-backed rotations.""" + + @classmethod + def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... + + @classmethod + def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... + + def __getitem__(self, idx: Any) -> Self: ... + + @property + def tensor(self) -> torch.Tensor: ... + + @property + def shape(self) -> torch.Size: ... + + def as_matrix(self) -> RotationMatrix: ... + + def as_quat(self, normalize: bool = False) -> RotationQuat: ... + + def compose(self, other: Self) -> Self: ... + + def convert_compose(self, other: Self) -> Self: ... + + def apply(self, points: torch.Tensor) -> torch.Tensor: ... + + def invert(self) -> Self: ... + + @property + def dtype(self) -> torch.dtype: + return self.tensor.dtype + + @property + def device(self) -> torch.device: + return self.tensor.device + + @property + def requires_grad(self) -> bool: + return self.tensor.requires_grad + + @classmethod + def _from_tensor(cls, tensor: torch.Tensor) -> Self: + return cls(tensor) # type: ignore[call-arg] + + def to(self, **kwargs) -> Self: + return self._from_tensor(self.tensor.to(**kwargs)) + + def detach(self, *args, **kwargs) -> Self: + return self._from_tensor(self.tensor.detach(**kwargs)) + + def tensor_apply(self, func) -> Self: + transformed = [func(component) for component in self.tensor.unbind(dim=-1)] + return self._from_tensor(torch.stack(transformed, dim=-1)) + + +class RotationQuat(Rotation): + """A rotation represented by a real-first quaternion.""" + + def __init__(self, quats: torch.Tensor, normalized: bool = False): + if not isinstance(quats, torch.Tensor): + raise TypeError("quats must be a Torch tensor.") + if quats.ndim == 0 or quats.shape[-1] != 4: + raise ValueError( + f"quats must have trailing dimension 4, got shape {tuple(quats.shape)}." + ) + if not isinstance(normalized, bool): + raise TypeError("normalized must be a boolean.") + self._normalized = normalized + if normalized: + quats = F.normalize(quats.to(torch.float32), dim=-1) + self._quats = quats.where(quats[..., :1] >= 0, -quats) + else: + self._quats = quats.to(torch.float32) + + @property + def tensor(self) -> torch.Tensor: + return self._quats + + @property + def shape(self) -> torch.Size: + return self._quats.shape[:-1] + + @classmethod + def identity(cls, shape, **tensor_kwargs) -> RotationQuat: + quaternions = torch.ones((*shape, 4), **tensor_kwargs) + selector = torch.tensor([1, 0, 0, 0], device=quaternions.device) + return cls(quaternions * selector) + + @classmethod + def random(cls, shape, **tensor_kwargs) -> RotationQuat: + return cls(torch.randn((*shape, 4), **tensor_kwargs), normalized=True) + + def __getitem__(self, idx: Any) -> RotationQuat: + indices = _index_tuple(idx) + return RotationQuat(self._quats[(*indices, slice(None))]) + + def normalized(self) -> RotationQuat: + if self._normalized: + return self + return RotationQuat(self._quats, normalized=True) + + def as_quat(self, normalize: bool = False) -> RotationQuat: + return self + + def as_matrix(self) -> RotationMatrix: + quaternion = self.normalized().tensor + r, i, j, k = torch.unbind(quaternion, -1) + scale = 2.0 / torch.linalg.norm(quaternion, dim=-1) + elements = torch.stack( + ( + 1 - scale * (j * j + k * k), + scale * (i * j - k * r), + scale * (i * k + j * r), + scale * (i * j + k * r), + 1 - scale * (i * i + k * k), + scale * (j * k - i * r), + scale * (i * k - j * r), + scale * (j * k + i * r), + 1 - scale * (i * i + j * j), + ), + -1, + ) + return RotationMatrix(elements.reshape((*quaternion.shape[:-1], 3, 3))) + + def compose(self, other: RotationQuat) -> RotationQuat: + with fp32_autocast_context(self.device.type): + return RotationQuat(_quat_mult(self._quats, other._quats)) + + def convert_compose(self, other: Rotation) -> RotationQuat: + return self.compose(other.as_quat()) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + return _quat_rotation(self.normalized()._quats, points) + + def invert(self) -> RotationQuat: + return RotationQuat(_quat_invert(self._quats)) + + +class RotationMatrix(Rotation): + """A rotation represented by a dense FP32 matrix.""" + + def __init__(self, rots: torch.Tensor): + if not isinstance(rots, torch.Tensor): + raise TypeError("rots must be a Torch tensor.") + if rots.ndim > 0 and rots.shape[-1] == 9: + rots = rots.unflatten(-1, (3, 3)) + if rots.ndim < 2 or rots.shape[-2:] != (3, 3): + raise ValueError( + "rots must have trailing shape (3, 3) or flattened width 9, got " + f"shape {tuple(rots.shape)}." + ) + self._rots = rots.to(torch.float32) + + @property + def tensor(self) -> torch.Tensor: + return self._rots.flatten(-2) + + @property + def shape(self) -> torch.Size: + return self._rots.shape[:-2] + + @classmethod + def identity(cls, shape, **tensor_kwargs) -> RotationMatrix: + matrix = torch.eye(3, **tensor_kwargs) + matrix = matrix.view(*(1 for _ in shape), 3, 3) + return cls(matrix.expand(*shape, -1, -1)) + + @classmethod + def random(cls, shape, **tensor_kwargs) -> RotationMatrix: + return RotationQuat.random(shape, **tensor_kwargs).as_matrix() + + @staticmethod + def from_graham_schmidt( + x_axis: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-12, + ) -> RotationMatrix: + return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps)) + + def __getitem__(self, idx: Any) -> RotationMatrix: + indices = _index_tuple(idx) + return RotationMatrix(self._rots[(*indices, slice(None), slice(None))]) + + def as_matrix(self) -> RotationMatrix: + return self + + def to_3x3(self) -> torch.Tensor: + return self._rots + + def as_quat(self, normalize: bool = False) -> RotationQuat: + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + self._rots.flatten(-2), + dim=-1, + ) + q_abs = _sqrt_subgradient( + torch.stack( + ( + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ), + dim=-1, + ) + ) + products = torch.stack( + ( + q_abs[..., 0] ** 2, + m21 - m12, + m02 - m20, + m10 - m01, + m21 - m12, + q_abs[..., 1] ** 2, + m10 + m01, + m02 + m20, + m02 - m20, + m10 + m01, + q_abs[..., 2] ** 2, + m12 + m21, + m10 - m01, + m20 + m02, + m21 + m12, + q_abs[..., 3] ** 2, + ), + dim=-1, + ).unflatten(-1, (4, 4)) + floor = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + candidates = products / (2.0 * q_abs[..., None].max(floor)) + best = torch.zeros_like(q_abs, dtype=torch.bool) + best.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True) + quaternion = candidates[best, :].reshape(q_abs.shape) + return RotationQuat(quaternion) + + def compose(self, other: RotationMatrix) -> RotationMatrix: + with fp32_autocast_context(self.device.type): + return RotationMatrix(self._rots @ other._rots) + + def convert_compose(self, other: Rotation) -> RotationMatrix: + return self.compose(other.as_matrix()) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + with fp32_autocast_context(self.device.type): + if self._rots.shape[-3] == 1: + return points @ self._rots.transpose(-1, -2).squeeze(-3) + return torch.einsum("...ij,...j", self._rots, points) + + def invert(self) -> RotationMatrix: + return RotationMatrix(self._rots.transpose(-1, -2)) + + +@dataclass(frozen=True) +class Affine3D: + """A rigid transform with translation and rotation components.""" + + trans: torch.Tensor + rot: Rotation + + def __post_init__(self) -> None: + if not isinstance(self.trans, torch.Tensor): + raise TypeError("trans must be a Torch tensor.") + if not isinstance(self.rot, Rotation): + raise TypeError("rot must implement the ESMFold2 Rotation interface.") + if self.trans.ndim == 0 or self.trans.shape[-1] != 3: + raise ValueError( + "trans must have trailing dimension 3, got " + f"shape {tuple(self.trans.shape)}." + ) + if self.trans.shape[:-1] != self.rot.shape: + raise ValueError( + "translation and rotation batch shapes must match, got " + f"{tuple(self.trans.shape[:-1])} and {tuple(self.rot.shape)}." + ) + + @property + def shape(self) -> torch.Size: + return self.trans.shape[:-1] + + @property + def dtype(self) -> torch.dtype: + return self.trans.dtype + + @property + def device(self) -> torch.device: + return self.trans.device + + @property + def requires_grad(self) -> bool: + return self.trans.requires_grad + + @property + def tensor(self) -> torch.Tensor: + return torch.cat((self.rot.tensor, self.trans), dim=-1) + + @staticmethod + def identity( + shape_or_affine: tuple[int, ...] | Affine3D, + rotation_type: type[Rotation] = RotationMatrix, + **tensor_kwargs, + ) -> Affine3D: + if isinstance(shape_or_affine, Affine3D): + kwargs = { + "dtype": shape_or_affine.dtype, + "device": shape_or_affine.device, + } + kwargs.update(tensor_kwargs) + shape = shape_or_affine.shape + rotation_type = type(shape_or_affine.rot) + else: + kwargs = tensor_kwargs + shape = shape_or_affine + return Affine3D( + torch.zeros((*shape, 3), **kwargs), + rotation_type.identity(shape, **kwargs), + ) + + @staticmethod + def random( + shape: tuple[int, ...], + std: float = 1, + rotation_type: type[Rotation] = RotationMatrix, + **tensor_kwargs, + ) -> Affine3D: + translation = torch.randn((*shape, 3), **tensor_kwargs).mul(std) + rotation = rotation_type.random(shape, **tensor_kwargs) + return Affine3D(trans=translation, rot=rotation) + + @staticmethod + def from_tensor(tensor: torch.Tensor) -> Affine3D: + if not isinstance(tensor, torch.Tensor): + raise TypeError("tensor must be a Torch tensor.") + if tensor.ndim == 0: + raise ValueError("tensor must have at least one dimension.") + width = tensor.shape[-1] + if width == 4: + if tensor.ndim < 2 or tensor.shape[-2] not in (3, 4): + raise ValueError( + "matrix-form affine tensors must have trailing shape (3, 4) or " + f"(4, 4), got {tuple(tensor.shape)}." + ) + translation = tensor[..., :3, 3] + rotation: Rotation = RotationMatrix(tensor[..., :3, :3]) + elif width == 6: + translation = tensor[..., -3:] + rotation = RotationQuat(F.pad(tensor[..., :3], (1, 0), value=1)) + elif width == 7: + translation = tensor[..., -3:] + rotation = RotationQuat(tensor[..., :4]) + elif width == 12: + translation = tensor[..., -3:] + rotation = RotationMatrix(tensor[..., :-3].unflatten(-1, (3, 3))) + else: + raise RuntimeError( + f"Cannot detect rotation format from {tensor.shape[-1] - 3}-d flat vector" + ) + return Affine3D(translation, rotation) + + @staticmethod + def from_tensor_pair( + translation: torch.Tensor, + rotation: torch.Tensor, + ) -> Affine3D: + return Affine3D(translation, RotationMatrix(rotation)) + + @staticmethod + def from_graham_schmidt( + neg_x_axis: torch.Tensor, + origin: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-10, + ) -> Affine3D: + x_axis = origin - neg_x_axis + plane_direction = xy_plane - origin + rotation = RotationMatrix.from_graham_schmidt( + x_axis, + plane_direction, + eps, + ) + return Affine3D(trans=origin, rot=rotation) + + @staticmethod + def cat(affines: list[Affine3D], dim: int = 0) -> Affine3D: + if not affines: + raise ValueError("affines must contain at least one transform.") + if any(not isinstance(affine, Affine3D) for affine in affines): + raise TypeError("affines must contain only Affine3D instances.") + if dim < 0: + dim = len(affines[0].shape) + dim + return Affine3D.from_tensor(torch.cat([affine.tensor for affine in affines], dim=dim)) + + def __getitem__(self, idx: Any) -> Affine3D: + indices = _index_tuple(idx) + translation = self.trans[(*indices, slice(None))] + return Affine3D(trans=translation, rot=self.rot[idx]) + + def to(self, **kwargs) -> Affine3D: + return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs)) + + def detach(self, *args, **kwargs) -> Affine3D: + return Affine3D( + self.trans.detach(**kwargs), + self.rot.detach(**kwargs), + ) + + def tensor_apply(self, func) -> Affine3D: + components = [func(value) for value in self.tensor.unbind(dim=-1)] + return Affine3D.from_tensor(torch.stack(components, dim=-1)) + + def as_matrix(self) -> Affine3D: + return Affine3D(trans=self.trans, rot=self.rot.as_matrix()) + + def as_quat(self, normalize: bool = False) -> Affine3D: + return Affine3D( + trans=self.trans, + rot=self.rot.as_quat(normalize), + ) + + def compose( + self, + other: Affine3D, + autoconvert: bool = False, + ) -> Affine3D: + compose_rotation = self.rot.convert_compose if autoconvert else self.rot.compose + rotation = compose_rotation(other.rot) + translation = self.rot.apply(other.trans) + self.trans + return Affine3D(trans=translation, rot=rotation) + + def compose_rotation( + self, + other: Rotation, + autoconvert: bool = False, + ) -> Affine3D: + compose = self.rot.convert_compose if autoconvert else self.rot.compose + return Affine3D(trans=self.trans, rot=compose(other)) + + def scale(self, value: torch.Tensor | float) -> Affine3D: + return Affine3D(self.trans * value, self.rot) + + def mask(self, mask: torch.Tensor, with_zero: bool = False) -> Affine3D: + if with_zero: + masked = torch.zeros_like(self.tensor).where( + mask[..., None], + self.tensor, + ) + return Affine3D.from_tensor(masked) + identity = self.identity( + self.shape, + rotation_type=type(self.rot), + device=self.device, + dtype=self.dtype, + ).tensor + return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor)) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + return self.rot.apply(points) + self.trans + + def invert(self) -> Affine3D: + rotation = self.rot.invert() + return Affine3D(trans=-rotation.apply(self.trans), rot=rotation) + + +def build_affine3d_from_coordinates( + coords: torch.Tensor, +) -> tuple[Affine3D, torch.Tensor]: + """Build residue frames from X with shape (b, l, 3, 3).""" + + if not isinstance(coords, torch.Tensor): + raise TypeError("coords must be a Torch tensor.") + if coords.ndim != 4 or coords.shape[-2:] != (3, 3): + raise ValueError( + "coords must have shape (batch, length, 3, 3), got " + f"{tuple(coords.shape)}." + ) + + maximum_distance = 1e6 + coord_mask = torch.all( + torch.all( + torch.isfinite(coords) & (coords < maximum_distance), + dim=-1, + ), + dim=-1, + ) + + def backbone_affine(positions: torch.Tensor) -> Affine3D: + n, ca, c = positions.unbind(dim=-2) + return Affine3D.from_graham_schmidt(c, ca, n) + + coords = coords.clone().float() + coords[~coord_mask] = 0 + average = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / ( + coord_mask.sum(-1)[..., None, None] + 1e-8 + ) + average_affine = backbone_affine(average.float()).as_matrix() + + b, length, _, _ = coords.shape + rotation = average_affine.rot.tensor[..., None, :].expand(b, length, 9) + translation = average_affine.trans[..., None, :].expand(b, length, 3) + identity = RotationMatrix.identity( + (b, length), + dtype=torch.float32, + device=coords.device, + requires_grad=False, + ) + rotation = rotation.where( + coord_mask.any(-1)[..., None, None], + identity.tensor, + ) + missing_frame = Affine3D(translation, RotationMatrix(rotation)) + + residue_frame = backbone_affine(coords.float()) + residue_frame = Affine3D.from_tensor( + residue_frame.tensor.where( + coord_mask[..., None], + missing_frame.tensor, + ) + ) + return residue_frame, coord_mask diff --git a/src/fastplms/models/esmfold2/esmfold2_aligner.py b/src/fastplms/models/esmfold2/esmfold2_aligner.py new file mode 100644 index 0000000..a23b9d6 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_aligner.py @@ -0,0 +1,87 @@ +"""Rigid alignment for structure dataclasses.""" + +from __future__ import annotations + +from dataclasses import Field, replace +from typing import Any, ClassVar, Protocol, TypeVar + +import numpy as np +import torch +from torch import Tensor + +from .esmfold2_protein_structure import compute_affine_and_rmsd + + +class Alignable(Protocol): + """Minimum structure interface accepted by :class:`Aligner`.""" + + __dataclass_fields__: ClassVar[dict[str, Field[Any]]] + + @property + def atom37_positions(self) -> np.ndarray: ... + + @property + def atom37_mask(self) -> np.ndarray: ... + + def __len__(self) -> int: ... + + +AlignableT = TypeVar("AlignableT", bound=Alignable) + + +def _coordinate_batch(structure: Alignable) -> Tensor: + return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0) + + +def _shared_atom_mask(mobile: Alignable, target: Alignable, backbone_only: bool) -> Tensor: + shared = np.asarray(mobile.atom37_mask, dtype=bool) & np.asarray( + target.atom37_mask, + dtype=bool, + ) + if backbone_only: + shared = shared.copy() + shared[:, 3:] = False + return torch.from_numpy(shared).unsqueeze(0) + + +class Aligner: + """Fit a mobile structure onto a target with masked Kabsch alignment.""" + + def __init__( + self, + mobile: Alignable, + target: Alignable, + only_use_backbone: bool = False, + use_reflection: bool = False, + ) -> None: + if len(mobile) != len(target): + raise AssertionError("mobile and target must contain the same residue count") + + mobile_coordinates = _coordinate_batch(mobile) + target_coordinates = _coordinate_batch(target) + if use_reflection: + target_coordinates = -target_coordinates + atom_mask = _shared_atom_mask(mobile, target, only_use_backbone) + self._affine3D, rmsd = compute_affine_and_rmsd( + mobile_coordinates, + target_coordinates, + atom_exists_mask=atom_mask, + ) + self._rmsd = rmsd.item() + + @property + def rmsd(self) -> float: + return self._rmsd + + def apply(self, mobile: AlignableT) -> AlignableT: + """Return a dataclass copy with all present atom coordinates aligned.""" + + present = np.asarray(mobile.atom37_mask, dtype=bool) + packed = torch.as_tensor( + mobile.atom37_positions[present], + dtype=torch.float32, + ).unsqueeze(0) + aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy() + atom37_positions = np.full_like(mobile.atom37_positions, np.nan) + atom37_positions[present] = aligned + return replace(mobile, atom37_positions=atom37_positions) diff --git a/src/fastplms/models/esmfold2/esmfold2_atom_indexer.py b/src/fastplms/models/esmfold2/esmfold2_atom_indexer.py new file mode 100644 index 0000000..676cbde --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_atom_indexer.py @@ -0,0 +1,30 @@ +"""Name-based views into an atom-axis property.""" + +from __future__ import annotations + +from operator import attrgetter +from typing import Any + +import numpy as np + +from .esmfold2_protein_structure import index_by_atom_name + + +class AtomIndexer: + """Select named atoms from one property of a structure-like object. + + The wrapper intentionally remains small because ``ProteinChain.atom37`` and + related public properties expose it directly. + """ + + __slots__ = ("_get_property", "dim", "property", "structure") + + def __init__(self, structure: Any, property: str, dim: int): + self.structure = structure + self.property = property + self.dim = dim + self._get_property = attrgetter(property) + + def __getitem__(self, atom_names: str | list[str]) -> np.ndarray: + values = self._get_property(self.structure) + return index_by_atom_name(values, atom_names, dim=self.dim) diff --git a/src/fastplms/models/esmfold2/esmfold2_conformers.py b/src/fastplms/models/esmfold2/esmfold2_conformers.py new file mode 100644 index 0000000..ddd0e9b --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_conformers.py @@ -0,0 +1,402 @@ +"""Lazy access to Chemical Component Dictionary conformers. + +The feature pipeline depends on atom names, formal charges, bonds, leaving-atom +flags, and one preferred reference conformer. Asset resolution is explicit at +``load_ccd`` time; importing this module performs no download or file access. +""" + +from __future__ import annotations + +import os +import pickle +import stat +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import file_digest +from pathlib import Path +from typing import Any, BinaryIO + +import numpy as np +from huggingface_hub import hf_hub_download +from huggingface_hub.constants import HF_HUB_CACHE + +from fastplms.registry import RuntimeAsset, get_model_registry + +from .esmfold2_constants import RES_TYPE_TO_CCD + +_CCD_ENVIRONMENT_VARIABLE = "ESMCFOLD_CCD_PATH" +_CCD_ASSET_ID = "esmfold2_ccd" + + +@dataclass(frozen=True) +class _ResolvedAsset: + path: Path + trusted_hub_cache_root: Path | None = None + + +def _asset_contract() -> RuntimeAsset: + """Return the manifest-owned identity of the trusted CCD pickle.""" + + try: + asset = get_model_registry().runtime_assets[_CCD_ASSET_ID] + except KeyError as error: + raise RuntimeError( + f"The package manifest does not declare runtime asset {_CCD_ASSET_ID!r}." + ) from error + if asset.trust_kind != "hash_pinned_pickle": + raise RuntimeError( + f"Runtime asset {_CCD_ASSET_ID!r} must use the hash_pinned_pickle trust policy." + ) + return asset + + +@contextmanager +def _open_verified_asset( + asset_path: Path, + contract: RuntimeAsset, + *, + trusted_hub_cache_root: Path | None = None, +) -> Iterator[BinaryIO]: + """Yield a private snapshot containing exactly the verified pickle bytes.""" + + try: + path_state = asset_path.lstat() + except FileNotFoundError as error: + raise FileNotFoundError(f"CCD asset does not exist: {asset_path}") from error + opened_path = asset_path + if stat.S_ISLNK(path_state.st_mode): + if trusted_hub_cache_root is None: + raise ValueError(f"CCD asset must not be a symlink: {asset_path}") + opened_path = _resolve_trusted_hub_snapshot_link( + asset_path, + contract, + trusted_hub_cache_root, + ) + path_state = opened_path.lstat() + if not stat.S_ISREG(path_state.st_mode): + raise ValueError(f"CCD asset must be a regular file: {asset_path}") + + flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + descriptor: int | None = None + try: + descriptor = os.open(opened_path, flags) + opened_state = os.fstat(descriptor) + if not stat.S_ISREG(opened_state.st_mode): + raise ValueError(f"CCD asset must be a regular file: {asset_path}") + if (path_state.st_dev, path_state.st_ino) != ( + opened_state.st_dev, + opened_state.st_ino, + ): + raise ValueError(f"CCD asset changed while it was being opened: {asset_path}") + + source = os.fdopen(descriptor, "rb") + descriptor = None + with source, tempfile.TemporaryFile(mode="w+b") as snapshot: + actual_size = opened_state.st_size + if actual_size != contract.size: + raise ValueError( + "CCD asset size mismatch: " + f"expected {contract.size} bytes, received {actual_size}." + ) + # Copy into a loader-owned OS temporary file. Hashing and + # deserialization then consume the same immutable snapshot, so a + # path replacement or in-place source write cannot substitute + # unverified pickle bytes after validation. + remaining = contract.size + while remaining: + chunk = source.read(min(1024 * 1024, remaining)) + if not chunk: + break + snapshot.write(chunk) + remaining -= len(chunk) + copied_size = snapshot.tell() + extra_byte = source.read(1) + if remaining or extra_byte: + observed_size = copied_size if remaining else copied_size + len(extra_byte) + raise ValueError( + "CCD asset size changed while it was being copied: " + f"expected {contract.size} bytes, received at least {observed_size}." + ) + snapshot.flush() + snapshot.seek(0) + actual_hash = file_digest(snapshot, "sha256").hexdigest() + if actual_hash != contract.sha256: + raise ValueError( + "CCD asset SHA256 mismatch; refusing to cross the " + "trusted-pickle boundary." + ) + snapshot.seek(0) + yield snapshot + finally: + if descriptor is not None: + os.close(descriptor) + + +def _resolve_trusted_hub_snapshot_link( + asset_path: Path, + contract: RuntimeAsset, + cache_root: Path, +) -> Path: + """Resolve only the immutable Hub snapshot link declared by the manifest.""" + + root = cache_root.expanduser().resolve(strict=True) + if len(contract.revision) != 40 or any( + character not in "0123456789abcdef" for character in contract.revision.lower() + ): + raise ValueError("CCD Hub asset revision must be an immutable 40-character commit.") + relative_asset = Path(contract.path) + if relative_asset.is_absolute() or ".." in relative_asset.parts: + raise ValueError(f"CCD Hub asset path is unsafe: {contract.path!r}") + repository_cache = root / f"models--{contract.repository.replace('/', '--')}" + try: + repository_cache.resolve(strict=True).relative_to(root) + except (FileNotFoundError, ValueError) as error: + raise ValueError( + f"CCD Hub repository cache escapes the effective Hub cache root: {repository_cache}" + ) from error + snapshot_root = repository_cache / "snapshots" / contract.revision + expected_path = snapshot_root / relative_asset + lexical_path = Path(os.path.abspath(asset_path)) + if lexical_path != Path(os.path.abspath(expected_path)): + raise ValueError( + "CCD Hub symlink is not the manifest-owned immutable snapshot path: " + f"{asset_path}" + ) + + try: + asset_path.parent.resolve(strict=True).relative_to(root) + except (FileNotFoundError, ValueError) as error: + raise ValueError( + f"CCD Hub snapshot path escapes the effective Hub cache root: {asset_path}" + ) from error + + resolved = asset_path.resolve(strict=True) + blob_root = (repository_cache / "blobs").resolve(strict=True) + try: + blob_root.relative_to(root) + resolved.relative_to(blob_root) + except ValueError as error: + raise ValueError( + f"CCD Hub snapshot link escapes its repository blob cache: {asset_path}" + ) from error + if not resolved.is_file() or resolved.is_symlink(): + raise ValueError(f"CCD Hub snapshot target must be a regular file: {resolved}") + return resolved + + +class _ChemicalComponentStore: + def __init__(self) -> None: + self.molecules: dict[str, Any] | None = None + self.conformers: dict[str, dict[str, np.ndarray]] = {} + self.atoms: dict[str, list[tuple[str, str, int]]] = {} + self.bonds: dict[str, list[tuple[str, str]]] = {} + self.leaving_atoms: dict[str, set[str]] = {} + self.standard_positions: dict[tuple[int, str], np.ndarray | None] = {} + self.ligand_positions: dict[tuple[str, str], np.ndarray | None] = {} + + def load(self, cache_dir: Path | str | None = None) -> dict[str, Any]: + if self.molecules is not None: + return self.molecules + contract = _asset_contract() + resolved = self._resolve_asset_location(cache_dir, contract) + asset = resolved.path + try: + # SECURITY: the private snapshot is both hash-validated and + # deserialized, closing path-replacement and in-place-write races. + with _open_verified_asset( + asset, + contract, + trusted_hub_cache_root=resolved.trusted_hub_cache_root, + ) as handle: + loaded = pickle.load(handle) + except FileNotFoundError: + raise + except Exception as error: + raise ValueError(f"Could not read the CCD asset at {asset}: {error}") from error + if loaded is not None and not isinstance(loaded, dict): + raise TypeError("The CCD asset must contain a component dictionary.") + self.molecules = loaded or {} + return self.molecules + + @staticmethod + def _resolve_asset(cache_dir: Path | str | None) -> Path: + contract = _asset_contract() + return _ChemicalComponentStore._resolve_asset_location(cache_dir, contract).path + + @staticmethod + def _resolve_asset_location( + cache_dir: Path | str | None, + contract: RuntimeAsset, + ) -> _ResolvedAsset: + configured = os.environ.get(_CCD_ENVIRONMENT_VARIABLE) + if configured: + asset = Path(configured).expanduser() + elif cache_dir is not None: + asset = Path(cache_dir).expanduser() / contract.path + else: + try: + asset = Path( + hf_hub_download( + repo_id=contract.repository, + filename=contract.path, + revision=contract.revision, + ) + ) + except Exception as error: + raise FileNotFoundError( + "Could not resolve the ESMFold2 CCD asset. Set " + f"{_CCD_ENVIRONMENT_VARIABLE} or populate the Hugging Face cache." + ) from error + return _ResolvedAsset( + path=asset, + trusted_hub_cache_root=Path(HF_HUB_CACHE), + ) + return _ResolvedAsset(path=asset) + + def _component_with_conformer(self, component_id: str): + molecule = self.load().get(component_id) + if molecule is None or molecule.GetNumConformers() == 0: + return None, None + + conformers = list(molecule.GetConformers()) + priority = {"Computed": 0, "Ideal": 1} + selected_index = min( + range(len(conformers)), + key=lambda index: priority.get(conformers[index].GetPropsAsDict().get("name"), 2), + ) + + from rdkit import Chem + + heavy_molecule = Chem.RemoveHs(molecule, sanitize=False) + if heavy_molecule.GetNumConformers() == 0: + return None, None + conformer_index = min(selected_index, heavy_molecule.GetNumConformers() - 1) + return heavy_molecule, heavy_molecule.GetConformer(conformer_index) + + def conformer(self, component_id: str) -> dict[str, np.ndarray] | None: + if component_id not in self.conformers: + molecule, conformer = self._component_with_conformer(component_id) + positions: dict[str, np.ndarray] = {} + if molecule is not None and conformer is not None: + for atom in molecule.GetAtoms(): + atom_name = atom.GetPropsAsDict().get("name") + if not isinstance(atom_name, str) or not atom_name: + continue + point = conformer.GetAtomPosition(atom.GetIdx()) + positions[atom_name] = np.asarray((point.x, point.y, point.z), dtype=np.float32) + self.conformers[component_id] = positions + result = self.conformers[component_id] + return result or None + + def atom_records(self, component_id: str) -> list[tuple[str, str, int]] | None: + if component_id not in self.atoms: + molecule, _conformer = self._component_with_conformer(component_id) + records: list[tuple[str, str, int]] = [] + if molecule is not None: + for atom in molecule.GetAtoms(): + atom_name = atom.GetPropsAsDict().get("name") + if isinstance(atom_name, str) and atom_name: + records.append((atom_name, atom.GetSymbol(), atom.GetFormalCharge())) + self.atoms[component_id] = records + result = self.atoms[component_id] + return result or None + + def bond_records(self, component_id: str) -> list[tuple[str, str]] | None: + if component_id not in self.bonds: + molecule, _conformer = self._component_with_conformer(component_id) + records: list[tuple[str, str]] = [] + if molecule is not None: + names = { + atom.GetIdx(): atom.GetPropsAsDict().get("name") for atom in molecule.GetAtoms() + } + for bond in molecule.GetBonds(): + first = names.get(bond.GetBeginAtomIdx()) + second = names.get(bond.GetEndAtomIdx()) + if isinstance(first, str) and first and isinstance(second, str) and second: + records.append((first, second)) + self.bonds[component_id] = records + result = self.bonds[component_id] + return result or None + + def component_leaving_atoms(self, component_id: str) -> set[str]: + if component_id not in self.leaving_atoms: + molecule = self.load().get(component_id) + names: set[str] = set() + if molecule is not None: + for atom in molecule.GetAtoms(): + if atom.HasProp("leaving_atom") and atom.GetProp("leaving_atom") == "1": + name = atom.GetProp("name") if atom.HasProp("name") else "" + if name: + names.add(name) + self.leaving_atoms[component_id] = names + return self.leaving_atoms[component_id] + + +_STORE = _ChemicalComponentStore() + + +def load_ccd(cache_dir: Path | str | None = None) -> dict[str, Any]: + """Load and cache the CCD asset, resolving it only when called.""" + + return _STORE.load(cache_dir) + + +def get_ccd_conformer(component_id: str) -> dict[str, np.ndarray] | None: + """Return the preferred heavy-atom conformer by atom name.""" + + return _STORE.conformer(component_id) + + +def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None: + """Return one standard-residue atom position from the preferred conformer.""" + + key = (res_type, atom_name) + if key not in _STORE.standard_positions: + component_id = RES_TYPE_TO_CCD.get(res_type) + conformer = _STORE.conformer(component_id) if component_id is not None else None + _STORE.standard_positions[key] = None if conformer is None else conformer.get(atom_name) + return _STORE.standard_positions[key] + + +def get_ligand_idealized_atom_pos(residue_name: str, atom_name: str) -> np.ndarray | None: + """Return one ligand atom position from the preferred conformer.""" + + key = (residue_name, atom_name) + if key not in _STORE.ligand_positions: + conformer = _STORE.conformer(residue_name) + _STORE.ligand_positions[key] = None if conformer is None else conformer.get(atom_name) + return _STORE.ligand_positions[key] + + +def get_ligand_ccd_atoms_with_charges( + component_id: str, +) -> list[tuple[str, str, int]] | None: + """Return heavy-atom name, element, and formal-charge records.""" + + return _STORE.atom_records(component_id) + + +def get_ligand_ccd_bonds(component_id: str) -> list[tuple[str, str]] | None: + """Return bonds as component atom-name pairs.""" + + return _STORE.bond_records(component_id) + + +def get_ccd_leaving_atoms(component_id: str) -> set[str]: + """Return atoms removed when a CCD component is polymerized.""" + + return _STORE.component_leaving_atoms(component_id) + + +__all__ = [ + "get_ccd_conformer", + "get_ccd_leaving_atoms", + "get_idealized_atom_pos", + "get_ligand_ccd_atoms_with_charges", + "get_ligand_ccd_bonds", + "get_ligand_idealized_atom_pos", + "load_ccd", +] diff --git a/src/fastplms/models/esmfold2/esmfold2_constants.py b/src/fastplms/models/esmfold2/esmfold2_constants.py new file mode 100644 index 0000000..df07096 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_constants.py @@ -0,0 +1,156 @@ +"""Declarative molecular schema for ESMFold2 feature preparation. + +The package manifest owns the upstream revision and license provenance. This +module expresses the corresponding checkpoint-facing integer schema as compact +ordered records, then derives lookup tables from those records. The generated +tables are validated at import without reading files, downloading assets, or +mutating process state. +""" + +from __future__ import annotations + +SCHEMA_PROVENANCE = { + "manifest_family": "esmfold2", + "contract": "biohub_esmfold2_input_v1", +} + + +def _words(value: str) -> list[str]: + return value.split() + + +MOL_TYPE_PROTEIN = 0 +MOL_TYPE_DNA = 1 +MOL_TYPE_RNA = 2 +MOL_TYPE_NONPOLYMER = 3 + +# The record order is part of the checkpoint input contract. Residue indices +# start at two because zero and one are reserved by the model feature schema. +_PROTEIN_SCHEMA = tuple( + tuple(record.split(":")) + for record in ( + "ALA:A:N CA C O CB", + "ARG:R:N CA C O CB CG CD NE CZ NH1 NH2", + "ASN:N:N CA C O CB CG OD1 ND2", + "ASP:D:N CA C O CB CG OD1 OD2", + "CYS:C:N CA C O CB SG", + "GLN:Q:N CA C O CB CG CD OE1 NE2", + "GLU:E:N CA C O CB CG CD OE1 OE2", + "GLY:G:N CA C O", + "HIS:H:N CA C O CB CG ND1 CD2 CE1 NE2", + "ILE:I:N CA C O CB CG1 CG2 CD1", + "LEU:L:N CA C O CB CG CD1 CD2", + "LYS:K:N CA C O CB CG CD CE NZ", + "MET:M:N CA C O CB CG SD CE", + "PHE:F:N CA C O CB CG CD1 CD2 CE1 CE2 CZ", + "PRO:P:N CA C O CB CG CD", + "SER:S:N CA C O CB OG", + "THR:T:N CA C O CB OG1 CG2", + "TRP:W:N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2", + "TYR:Y:N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH", + "VAL:V:N CA C O CB CG1 CG2", + ) +) + +PROTEIN_RESIDUE_TO_RES_TYPE = { + residue: index for index, (residue, _letter, _atoms) in enumerate(_PROTEIN_SCHEMA, 2) +} +PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] = PROTEIN_RESIDUE_TO_RES_TYPE["MET"] +PROTEIN_UNK_RES_TYPE = 22 + +RNA_RESIDUE_TO_RES_TYPE = dict(zip("AGCU", range(23, 27), strict=True)) +RNA_UNK_RES_TYPE = 27 +DNA_RESIDUE_TO_RES_TYPE = dict(zip(("DA", "DG", "DC", "DT"), range(28, 32), strict=True)) +DNA_UNK_RES_TYPE = 32 +GAP_RES_TYPE = DNA_UNK_RES_TYPE + +PROTEIN_3TO1 = {residue: letter for residue, letter, _atoms in _PROTEIN_SCHEMA} +PROTEIN_3TO1["MSE"] = "M" +PROTEIN_1TO3 = {letter: residue for residue, letter, _atoms in _PROTEIN_SCHEMA} +PROTEIN_1TO3["X"] = "UNK" +DNA_1TO3 = dict(zip("ATCG", ("DA", "DT", "DC", "DG"), strict=True)) +RNA_1TO3 = {letter: letter for letter in "AUCG"} + +_ESM_RESIDUE_ORDER = "LAGVSERTIDPKQNFYM HWC".replace(" ", "") +ESM_PROTEIN_VOCAB = {residue: token_id for token_id, residue in enumerate(_ESM_RESIDUE_ORDER, 4)} +ESM_PROTEIN_VOCAB["X"] = 3 +DNA_RNA_LIGAND_INPUT_ID = 24 +MSA_PAD_TOKEN_ID = 0 +MSA_GAP_TOKEN_ID = 1 + +RES_TYPE_TO_CCD = { + **{ + index: residue for residue, index in PROTEIN_RESIDUE_TO_RES_TYPE.items() if residue != "MSE" + }, + 22: "UNK", + **dict(zip(range(23, 28), ("A", "G", "C", "U", "N"), strict=True)), + **dict(zip(range(28, 33), ("DA", "DG", "DC", "DT", "DN"), strict=True)), +} + +_CHARGE_SCHEMA = _words( + "LYS:NZ:1 ARG:NH2:1 HIS:ND1:1 PO4:O2:-1 PO4:O3:-1 PO4:O4:-1 " + "SO4:O3:-1 SO4:O4:-1 MG:MG:2 ZN:ZN:2 CA:CA:2 FE2:FE:2 MN:MN:2 " + "CO:CO:2 NCO:CO:3 CU:CU:2 NI:NI:2 K:K:1 NA:NA:1 CD:CD:2 CL:CL:-1 " + "ACT:OXT:-1 NAD:O2N:-1 NAD:N1N:1 NAP:O2N:-1 NAP:N1N:1 IMD:N3:1 " + "SAM:SD:1 FE:FE:3 A1BH3:N3:1" +) +CHARGED_ATOMS = { + (component, atom): int(charge) + for component, atom, charge in (record.split(":") for record in _CHARGE_SCHEMA) +} + +_PERIODIC_SYMBOLS = _words( + "H HE LI BE B C N O F NE NA MG AL SI P S CL AR K CA SC TI V CR MN FE CO NI CU ZN " + "GA GE AS SE BR KR RB SR Y ZR NB MO TC RU RH PD AG CD IN SN SB TE I XE CS BA LA CE " + "PR ND PM SM EU GD TB DY HO ER TM YB LU HF TA W RE OS IR PT AU HG TL PB BI PO AT RN " + "FR RA AC TH PA U" +) +ELEMENT_TO_ATOMIC_NUM = { + symbol: atomic_number + for atomic_number, symbol in enumerate(_PERIODIC_SYMBOLS, 1) + if symbol != "HE" +} +ELEMENT_NUMBER_TO_SYMBOL = { + atomic_number: symbol for symbol, atomic_number in ELEMENT_TO_ATOMIC_NUM.items() +} + +PROTEIN_HEAVY_ATOMS = { + residue: atom_string.split() for residue, _letter, atom_string in _PROTEIN_SCHEMA +} +PROTEIN_HEAVY_ATOMS["MSE"] = PROTEIN_HEAVY_ATOMS["MET"].copy() +PROTEIN_HEAVY_ATOMS["UNK"] = _words("N CA C O") + +DNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' C1'") +RNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' O2' C1'") +_NUCLEOBASE_ATOMS = { + "A": _words("N9 C8 N7 C5 C6 N6 N1 C2 N3 C4"), + "G": _words("N9 C8 N7 C5 C6 O6 N1 C2 N2 N3 C4"), + "C": _words("N1 C2 O2 N3 C4 N4 C5 C6"), + "U": _words("N1 C2 O2 N3 C4 O4 C5 C6"), + "T": _words("N1 C2 O2 N3 C4 O4 C5 C7 C6"), +} +DNA_HEAVY_ATOMS = { + "DA": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["A"], + "DG": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["G"], + "DC": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["C"], + "DT": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["T"], +} +RNA_HEAVY_ATOMS = {residue: RNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS[residue] for residue in "AGCU"} + + +def _validate_schema() -> None: + if sorted(set(PROTEIN_RESIDUE_TO_RES_TYPE.values())) != list(range(2, 22)): + raise RuntimeError("Protein residue indices must cover the checkpoint interval 2..21.") + if len(_ESM_RESIDUE_ORDER) != 20 or len(set(_ESM_RESIDUE_ORDER)) != 20: + raise RuntimeError("The ESM residue vocabulary must contain 20 canonical residues.") + if RES_TYPE_TO_CCD[14] != "MET" or PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] != 14: + raise RuntimeError("Selenomethionine must share the methionine residue index.") + if ELEMENT_TO_ATOMIC_NUM.get("U") != 92 or 2 in ELEMENT_NUMBER_TO_SYMBOL: + raise RuntimeError("The element schema must preserve the training-time atomic-number map.") + if set(DNA_HEAVY_ATOMS) != {"DA", "DG", "DC", "DT"}: + raise RuntimeError("The DNA atom schema is incomplete.") + + +_validate_schema() + +__all__ = [name for name in globals() if name.isupper()] diff --git a/src/fastplms/models/esmfold2/esmfold2_constants_esm3.py b/src/fastplms/models/esmfold2/esmfold2_constants_esm3.py new file mode 100644 index 0000000..4c04411 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_constants_esm3.py @@ -0,0 +1,98 @@ +"""Token schemas needed by the ESMC encoder inside ESMFold2. + +The values implement the published Biohub ESM sequence-token contract pinned by +``models.toml``. They are generated from ordered schemas so token position and +special-token relationships are explicit and independently testable. This +module performs no downloads and resolves no model assets at import time. +""" + +from __future__ import annotations + +from types import MappingProxyType + + +def _words(value: str) -> list[str]: + return value.split() + + +SEQUENCE_VOCAB = _words( + " L A G V S E R T I D P K Q N F Y M H W C X B U Z O . - | " +) + +_sequence_token_ids = MappingProxyType({token: index for index, token in enumerate(SEQUENCE_VOCAB)}) +SEQUENCE_BOS_TOKEN = _sequence_token_ids[""] +SEQUENCE_PAD_TOKEN = _sequence_token_ids[""] +SEQUENCE_EOS_TOKEN = _sequence_token_ids[""] +SEQUENCE_CHAINBREAK_TOKEN = _sequence_token_ids["|"] +SEQUENCE_MASK_TOKEN = _sequence_token_ids[""] +SEQUENCE_STANDARD_AA_MIN_TOKEN = _sequence_token_ids["L"] +SEQUENCE_STANDARD_AA_MAX_TOKEN = _sequence_token_ids["X"] + +VQVAE_CODEBOOK_SIZE = 4096 +VQVAE_SPECIAL_TOKENS = { + name: VQVAE_CODEBOOK_SIZE + offset + for offset, name in enumerate(("MASK", "EOS", "BOS", "PAD", "CHAINBREAK")) +} +VQVAE_DIRECTION_LOSS_BINS = 16 +VQVAE_PAE_BINS = 64 +VQVAE_MAX_PAE_BIN = 31.0 +VQVAE_PLDDT_BINS = 50 + +STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"] +STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"] +STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"] +STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"] +STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"] +STRUCTURE_UNDEFINED_TOKEN = 955 + +SASA_PAD_TOKEN = 0 +SS8_PAD_TOKEN = 0 +INTERPRO_PAD_TOKEN = 0 +RESIDUE_PAD_TOKEN = 0 + +CHAIN_BREAK_STR = "|" +SEQUENCE_BOS_STR = "" +SEQUENCE_EOS_STR = "" +MASK_STR_SHORT = "_" +SEQUENCE_MASK_STR = "" +SASA_MASK_STR = "" +SS8_MASK_STR = "" + +SSE_8CLASS_VOCAB = "GHITEBSC" +SSE_3CLASS_VOCAB = "HEC" +SSE_8CLASS_TO_3CLASS_MAP = dict(zip(SSE_8CLASS_VOCAB, "HHHCEECC", strict=True)) + +SASA_DISCRETIZATION_BOUNDARIES = [ + 0.8, + 4.0, + 9.6, + 16.4, + 24.5, + 32.9, + 42.0, + 51.5, + 61.2, + 70.9, + 81.6, + 93.3, + 107.2, + 125.4, + 151.4, +] +MAX_RESIDUE_ANNOTATIONS = 16 +TFIDF_VECTOR_SIZE = 58_641 +FUNCTION_TOKENS_DEPTH = 8 + + +def _validate_schema() -> None: + if len(SEQUENCE_VOCAB) != len(set(SEQUENCE_VOCAB)): + raise RuntimeError("The ESM sequence vocabulary contains duplicate tokens.") + if SEQUENCE_STANDARD_AA_MAX_TOKEN - SEQUENCE_STANDARD_AA_MIN_TOKEN != 20: + raise RuntimeError("The canonical residue interval must contain 20 tokens.") + if tuple(VQVAE_SPECIAL_TOKENS.values()) != tuple(range(4096, 4101)): + raise RuntimeError("The structure special-token interval is not contiguous.") + + +_validate_schema() + +__all__ = [name for name in globals() if name.isupper()] diff --git a/src/fastplms/models/esmfold2/esmfold2_input_builder.py b/src/fastplms/models/esmfold2/esmfold2_input_builder.py new file mode 100644 index 0000000..ca51bc2 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_input_builder.py @@ -0,0 +1,244 @@ +"""Typed, JSON-safe inputs for ESMFold2 feature preparation.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, TypeAlias + +import numpy as np + +from .esmfold2_msa import MSA + +MSAInput: TypeAlias = MSA | None + + +@dataclass +class Modification: + """A zero-indexed residue substitution using a CCD component.""" + + position: int + ccd: str + smiles: str | None = None + + +@dataclass +class ProteinInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + msa: MSAInput = None + + +@dataclass +class RNAInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + + +@dataclass +class DNAInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + + +@dataclass +class LigandInput: + id: str | list[str] + smiles: str | None = None + ccd: list[str] | None = None + + +@dataclass +class DistogramConditioning: + chain_id: str + distogram: np.ndarray + + +@dataclass +class PocketConditioning: + binder_chain_id: str + contacts: list[tuple[str, int]] + + +@dataclass +class CovalentBond: + chain_id1: str + res_idx1: int + atom_idx1: int + chain_id2: str + res_idx2: int + atom_idx2: int + + +SequenceInput: TypeAlias = ProteinInput | RNAInput | DNAInput | LigandInput + + +@dataclass +class StructurePredictionInput: + sequences: Sequence[SequenceInput] + pocket: PocketConditioning | None = None + distogram_conditioning: list[DistogramConditioning] | None = None + covalent_bonds: list[CovalentBond] | None = None + + +_CHAIN_TYPE = { + ProteinInput: "protein", + RNAInput: "rna", + DNAInput: "dna", +} + + +def _serialize_modifications( + modifications: list[Modification] | None, +) -> list[dict[str, Any]] | None: + if not modifications: + return None + return [{"position": item.position, "ccd": item.ccd} for item in modifications] + + +def _serialize_chain(chain: SequenceInput) -> dict[str, Any]: + if isinstance(chain, LigandInput): + return { + "smiles": chain.smiles, + "id": chain.id, + "ccd": chain.ccd, + "type": "ligand", + } + + chain_type = _CHAIN_TYPE.get(type(chain)) + if chain_type is None: + raise ValueError(f"Unsupported sequence input type: {type(chain)}") + serialized: dict[str, Any] = { + "sequence": chain.sequence, + "id": chain.id, + "type": chain_type, + } + if modifications := _serialize_modifications(chain.modifications): + serialized["modifications"] = modifications + if isinstance(chain, ProteinInput): + if chain.msa is not None and not isinstance(chain.msa, MSA): + raise AttributeError(f"MSA must be None or MSA. Got {chain.msa} instead.") + serialized["msa"] = None if chain.msa is None else {"sequences": chain.msa.sequences} + return serialized + + +def serialize_structure_prediction_input( + structure_input: StructurePredictionInput, +) -> dict[str, Any]: + """Convert an input object to a JSON-safe mapping.""" + + serialized: dict[str, Any] = { + "sequences": [_serialize_chain(chain) for chain in structure_input.sequences] + } + if structure_input.covalent_bonds is not None: + serialized["covalent_bonds"] = [ + vars(bond).copy() for bond in structure_input.covalent_bonds + ] + if structure_input.pocket is not None: + serialized["pocket"] = { + "binder_chain_id": structure_input.pocket.binder_chain_id, + "contacts": structure_input.pocket.contacts, + } + if structure_input.distogram_conditioning is not None: + serialized["distogram_conditioning"] = [ + {"chain_id": item.chain_id, "distogram": item.distogram.tolist()} + for item in structure_input.distogram_conditioning + ] + return serialized + + +def _deserialize_modifications(chain: dict[str, Any]) -> list[Modification] | None: + raw = chain.get("modifications") + if not raw: + return None + return [Modification(position=item["position"], ccd=item["ccd"]) for item in raw] + + +def _deserialize_msa(chain: dict[str, Any]) -> MSAInput: + raw = chain.get("msa") + if raw is None: + return None + if not isinstance(raw, dict) or not isinstance(raw.get("sequences"), list): + raise ValueError(f"Unexpected MSA value: {raw!r}") + return MSA.from_sequences(raw["sequences"]) + + +def _deserialize_chain(chain: dict[str, Any]) -> SequenceInput: + chain_type = chain.get("type") + common = {"id": chain["id"]} + if chain_type == "protein": + return ProteinInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + msa=_deserialize_msa(chain), + ) + if chain_type == "rna": + return RNAInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + ) + if chain_type == "dna": + return DNAInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + ) + if chain_type == "ligand": + return LigandInput(**common, smiles=chain.get("smiles"), ccd=chain.get("ccd")) + raise ValueError(f"Unsupported sequence type: {chain_type!r}") + + +def deserialize_structure_prediction_input(data: dict[str, Any]) -> StructurePredictionInput: + """Reconstruct the typed input represented by a serialized mapping.""" + + pocket_data = data.get("pocket") + pocket = None + if pocket_data is not None: + pocket = PocketConditioning( + binder_chain_id=pocket_data["binder_chain_id"], + contacts=[tuple(contact) for contact in pocket_data["contacts"]], + ) + + distogram_data = data.get("distogram_conditioning") + distograms = None + if distogram_data is not None: + distograms = [ + DistogramConditioning( + chain_id=item["chain_id"], distogram=np.asarray(item["distogram"]) + ) + for item in distogram_data + ] + + bond_data = data.get("covalent_bonds") + bonds = None + if bond_data is not None: + bonds = [CovalentBond(**item) for item in bond_data] + + return StructurePredictionInput( + sequences=[_deserialize_chain(chain) for chain in data["sequences"]], + pocket=pocket, + distogram_conditioning=distograms, + covalent_bonds=bonds, + ) + + +__all__ = [ + "CovalentBond", + "DNAInput", + "DistogramConditioning", + "LigandInput", + "MSAInput", + "Modification", + "PocketConditioning", + "ProteinInput", + "RNAInput", + "SequenceInput", + "StructurePredictionInput", + "deserialize_structure_prediction_input", + "serialize_structure_prediction_input", +] diff --git a/src/fastplms/models/esmfold2/esmfold2_metrics.py b/src/fastplms/models/esmfold2/esmfold2_metrics.py new file mode 100644 index 0000000..fcd6651 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_metrics.py @@ -0,0 +1,235 @@ +"""Contact, lDDT, RMSD, and GDT-TS metrics for structure validation.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.amp import autocast # type: ignore + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_misc import binpack, unbinpack +from .esmfold2_protein_structure import ( + compute_alignment_tensors, + compute_gdt_ts_no_alignment, + compute_rmsd_no_alignment, +) + + +def _distance_matrix(positions: Tensor, eps: float) -> Tensor: + displacement = positions[..., None, :] - positions[..., None, :, :] + return torch.sqrt(eps + torch.sum(displacement**2, dim=-1)) + + +def compute_lddt_from_dmat( + dmat_pred: Tensor, + dmat_true: Tensor, + pairwise_mask: Tensor, + cutoff: float | Tensor = 15.0, + eps: float = 1e-10, + per_residue: bool = True, +) -> Tensor: + """Score distance matrices ``D_pred`` and ``D_true`` with shape (..., l, l).""" + + sequence_length = dmat_true.size(-1) + identity = torch.eye(sequence_length, device=dmat_true.device) + scored_pairs = (dmat_true < cutoff) * pairwise_mask * (1.0 - identity) + absolute_error = torch.abs(dmat_true - dmat_pred) + score = ( + (absolute_error < 0.5).type(absolute_error.dtype) + + (absolute_error < 1.0).type(absolute_error.dtype) + + (absolute_error < 2.0).type(absolute_error.dtype) + + (absolute_error < 4.0).type(absolute_error.dtype) + ) * 0.25 + dimensions = (-1,) if per_residue else (-2, -1) + normalization = 1.0 / (eps + scored_pairs.sum(dim=dimensions)) + return normalization * (eps + (scored_pairs * score).sum(dim=dimensions)) + + +def compute_lddt( + all_atom_pred_pos: Tensor, + all_atom_positions: Tensor, + all_atom_mask: Tensor, + pairwise_all_atom_mask: Tensor | None = None, + cutoff: float | Tensor = 15.0, + eps: float = 1e-10, + per_residue: bool = True, + sequence_id: Tensor | None = None, +) -> Tensor: + """Compute lDDT from coordinate tensors and atom masks.""" + + expanded_mask = all_atom_mask[..., None] + true_distances = _distance_matrix(all_atom_positions, eps) + predicted_distances = _distance_matrix(all_atom_pred_pos, eps) + pair_mask = expanded_mask * expanded_mask.transpose(-2, -1) + if pairwise_all_atom_mask is not None: + pair_mask = pair_mask * pairwise_all_atom_mask + if sequence_id is not None: + same_sequence = sequence_id[..., None] == sequence_id[..., None, :] + pair_mask = pair_mask * same_sequence.type_as(pair_mask) + return compute_lddt_from_dmat( + predicted_distances, + true_distances, + pair_mask, + cutoff=cutoff, + eps=eps, + per_residue=per_residue, + ) + + +def compute_lddt_ca( + all_atom_pred_pos: Tensor, + all_atom_positions: Tensor, + all_atom_mask: Tensor, + cutoff: float = 15.0, + eps: float = 1e-10, + per_residue: bool = True, + sequence_id: Tensor | None = None, +) -> Tensor: + """Compute lDDT using only C-alpha coordinates.""" + + ca_index = residue_constants.atom_order["CA"] + predicted_ca = ( + all_atom_pred_pos if all_atom_pred_pos.dim() == 3 else all_atom_pred_pos[..., ca_index, :] + ) + return compute_lddt( + predicted_ca, + all_atom_positions[..., ca_index, :], + all_atom_mask[..., ca_index], + cutoff=cutoff, + eps=eps, + per_residue=per_residue, + sequence_id=sequence_id, + ) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_rmsd( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + reduction: str = "batch", +) -> Tensor: + """Align ``X`` to ``Y`` and compute RMSD.""" + + centered_mobile, _, centered_target, _, rotation, counts = compute_alignment_tensors( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + rmsd = compute_rmsd_no_alignment( + torch.matmul(centered_mobile, rotation), + centered_target, + counts, + reduction=reduction, + ) + if reduction == "per_residue" and sequence_id is not None: + return binpack(rmsd, sequence_id, pad_value=0) + return rmsd + + +def compute_gdt_ts( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + reduction: str = "per_sample", +) -> Tensor: + """Align ``X`` to ``Y`` and compute GDT-TS.""" + + if atom_exists_mask is None: + atom_exists_mask = torch.isfinite(target).all(dim=-1) + centered_mobile, _, centered_target, _, rotation, _ = compute_alignment_tensors( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + if sequence_id is not None: + atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False) + return compute_gdt_ts_no_alignment( + torch.matmul(centered_mobile, rotation), + centered_target, + atom_exists_mask, + reduction, + ) + + +def _batched_contacts(predictions: Tensor, targets: Tensor) -> tuple[Tensor, Tensor]: + if predictions.dim() == 2: + predictions = predictions.unsqueeze(0) + if targets.dim() == 2: + targets = targets.unsqueeze(0) + if predictions.size() != targets.size(): + raise ValueError( + f"Size mismatch. Received predictions of size {predictions.size()}, " + f"targets of size {targets.size()}" + ) + return predictions, targets + + +def _valid_contact_mask( + targets: Tensor, + src_lengths: Tensor, + minsep: int, + maxsep: int | None, +) -> Tensor: + sequence_length = targets.shape[-1] + positions = torch.arange(sequence_length, device=targets.device) + separation = (positions.unsqueeze(0) - positions.unsqueeze(1)).unsqueeze(0) + valid = (separation >= minsep) & (targets >= 0) + if maxsep is not None: + valid &= separation < maxsep + within_length = positions.unsqueeze(0) < src_lengths.unsqueeze(1) + return valid & within_length.unsqueeze(1) & within_length.unsqueeze(2) + + +def contact_precision( + predictions: Tensor, + targets: Tensor, + src_lengths: Tensor | None = None, + minsep: int = 6, + maxsep: int | None = None, + override_length: int | None = None, +) -> dict[str, Tensor]: + """Compute P@L, P@L/5, and binned area for contact probabilities.""" + + predictions, targets = _batched_contacts(predictions, targets) + batch_size, sequence_length, _ = predictions.shape + if src_lengths is None: + src_lengths = torch.full( + (batch_size,), + sequence_length, + dtype=torch.long, + device=predictions.device, + ) + valid = _valid_contact_mask(targets, src_lengths, minsep, maxsep) + masked_predictions = predictions.masked_fill(~valid, float("-inf")) + row_index, column_index = np.triu_indices(sequence_length, minsep) + upper_predictions = masked_predictions[:, row_index, column_index] + upper_targets = targets[:, row_index, column_index] + + topk = sequence_length if override_length is None else max(sequence_length, override_length) + ranked_indices = upper_predictions.argsort(dim=-1, descending=True)[:, :topk] + batch_indices = torch.arange(batch_size, device=ranked_indices.device).unsqueeze(1) + ranked_targets = upper_targets[batch_indices, ranked_indices] + if ranked_targets.size(1) < topk: + ranked_targets = F.pad(ranked_targets, [0, topk - ranked_targets.size(1)]) + cumulative_contacts = ranked_targets.type_as(predictions).cumsum(dim=-1) + + gather_lengths = src_lengths.unsqueeze(1) + if override_length is not None: + gather_lengths = override_length * torch.ones_like(gather_lengths) + fractions = torch.arange(0.1, 1.1, 0.1, device=predictions.device).unsqueeze(0) + gather_indices = (fractions * gather_lengths).type(torch.long).sub(1).clamp_min(0) + cumulative_bins = cumulative_contacts.gather(1, gather_indices) + precisions = cumulative_bins / (gather_indices + 1).type_as(cumulative_bins) + return { + "AUC": precisions.mean(dim=-1), + "P@L": precisions[:, 9], + "P@L5": precisions[:, 1], + } diff --git a/src/fastplms/models/esmfold2/esmfold2_misc.py b/src/fastplms/models/esmfold2/esmfold2_misc.py new file mode 100644 index 0000000..022edaf --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_misc.py @@ -0,0 +1,400 @@ +"""Small tensor, sequence, and annotation utilities used by ESMFold2. + +The helpers in this module are deliberately free of model state. Importing the +module therefore performs no device selection, compilation, or remote access. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Generator, Iterable, Sequence +from contextlib import AbstractContextManager, nullcontext +from dataclasses import is_dataclass +from io import BytesIO +from typing import Any, Protocol, TypeVar, runtime_checkable +from warnings import warn + +import numpy as np +import torch +import zstandard + +from .esmfold2_constants_esm3 import CHAIN_BREAK_STR +from .esmfold2_utils_types import FunctionAnnotation + +MAX_SUPPORTED_DISTANCE = 1e6 + +TSequence = TypeVar("TSequence", bound=Sequence) + + +@runtime_checkable +class Concatable(Protocol): + """Protocol for sequence-like records with a class-level concatenator.""" + + @classmethod + def concat(cls, objs: list[Concatable]) -> Concatable: ... + + +def fp32_autocast_context( + device_type: str, +) -> AbstractContextManager[Any]: # type: ignore + """Return a context that keeps numerically sensitive work in FP32.""" + + if device_type == "mps": + return nullcontext() + if device_type == "cpu": + return torch.amp.autocast(device_type, enabled=False) # type: ignore + if device_type == "cuda": + return torch.amp.autocast(device_type, dtype=torch.float32) # type: ignore + raise ValueError(f"Unsupported device type: {device_type}") + + +def maybe_tensor(value, convert_none_to_nan: bool = False) -> torch.Tensor | None: + """Convert an optional array-like value to a tensor.""" + + if value is None: + return None + if isinstance(value, torch.Tensor): + return value + if isinstance(value, list) and all(isinstance(element, torch.Tensor) for element in value): + return torch.stack(value) + if convert_none_to_nan: + value = np.asarray(value, dtype=np.float32) + value = np.where(value is None, np.nan, value) + return torch.tensor(value) + + +def maybe_list(value, convert_nan_to_none: bool = False) -> list | None: + """Convert an optional tensor or NumPy array to nested Python lists.""" + + if value is None: + return None + if not convert_nan_to_none: + return value.tolist() + if isinstance(value, torch.Tensor): + nan_mask = torch.isnan(value).cpu().numpy() + array = value.cpu().numpy().astype(object) + elif isinstance(value, np.ndarray): + nan_mask = np.isnan(value) + array = value.astype(object) + else: + raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.") + array[nan_mask] = None + return array.tolist() + + +def replace_inf(data): + """Replace infinite array values by the ESM API sentinel value.""" + + if data is None: + return None + array = np.asarray(data, dtype=np.float32) + return np.where(np.isinf(array), 1000, array).tolist() + + +def slice_python_object_as_numpy( + obj: TSequence, + idx: int | list[int] | slice | np.ndarray, +) -> TSequence: + """Apply NumPy-style scalar, mask, or index-array slicing to Python data.""" + + normalized_idx: list[int] | slice | np.ndarray = ( + [int(idx)] if np.isscalar(idx) else idx # type: ignore[arg-type] + ) + + if isinstance(normalized_idx, np.ndarray) and normalized_idx.dtype == bool: + selected = [obj[position] for position in np.flatnonzero(normalized_idx)] + elif isinstance(normalized_idx, slice): + selected = obj[normalized_idx] + else: + selected = [obj[position] for position in normalized_idx] + + if isinstance(obj, str) and isinstance(selected, list): + return "".join(selected) # type: ignore[return-value] + return obj.__class__(selected) # type: ignore[call-arg,return-value] + + +def slice_any_object( + obj: TSequence, + idx: int | list[int] | slice | np.ndarray, +) -> TSequence: + """Slice tensors, arrays, dataclasses, and ordinary Python sequences.""" + + if isinstance(obj, (np.ndarray, torch.Tensor)) or is_dataclass(obj): + return obj[idx] # type: ignore[index,return-value] + return slice_python_object_as_numpy(obj, idx) + + +def join_lists( + lists: Sequence[Sequence[Any]], + separator: Sequence[Any] | None = None, +) -> list[Any]: + """Join lists, inserting all elements of ``separator`` between inputs.""" + + if len(lists) == 0: + return [] + joined = list(lists[0]) + for values in lists[1:]: + if separator: + joined.extend(separator) + joined.extend(values) + return joined + + +def iterate_with_intermediate( + lists: Iterable, + intermediate, +) -> Generator[Any, None, None]: + """Yield an intermediate value between consecutive input values.""" + + iterator = iter(lists) + yield next(iterator) + for value in iterator: + yield intermediate + yield value + + +def concat_objects(objs: Sequence[Any], separator: Any | None = None): + """Concatenate one supported homogeneous collection.""" + + if not objs: + raise ValueError("objs must contain at least one value.") + first = objs[0] + if isinstance(first, Concatable): + return first.__class__.concat(objs) + if isinstance(first, str): + if not isinstance(separator, str): + raise TypeError("separator must be a string when joining strings.") + return separator.join(objs) + if isinstance(first, list): + return join_lists(objs, None if separator is None else [separator]) + if isinstance(first, np.ndarray): + pieces = ( + objs + if separator is None + else list(iterate_with_intermediate(objs, np.array([separator]))) + ) + return np.concatenate(pieces) + if isinstance(first, torch.Tensor): + pieces = ( + objs + if separator is None + else list(iterate_with_intermediate(objs, torch.tensor([separator]))) + ) + return torch.cat(pieces) # type: ignore[arg-type] + raise TypeError(type(first)) + + +def rbf(values, v_min, v_max, n_bins=16): + """Encode values against evenly spaced radial basis centers.""" + + centers = torch.linspace( + v_min, + v_max, + n_bins, + dtype=values.dtype, + device=values.device, + ) + centers = centers.reshape((1,) * values.ndim + (-1,)) + standardized = (values.unsqueeze(-1) - centers) / ((v_max - v_min) / n_bins) + return torch.exp(-(standardized**2)) + + +def batched_gather(data, inds, dim=0, no_batch_dims=0): + """Gather along one data dimension while retaining leading batch axes.""" + + batch_indices = [] + index_rank = len(inds.shape) + for axis, size in enumerate(data.shape[:no_batch_dims]): + shape = (1,) * axis + (-1,) + (1,) * (index_rank - axis - 1) + batch_indices.append(torch.arange(size).view(*shape)) + tail = [slice(None)] * (len(data.shape) - no_batch_dims) + tail[dim - no_batch_dims if dim >= 0 else dim] = inds + return data[tuple(batch_indices + tail)] + + +def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor: + """Gather node features for each row of an edge-index tensor.""" + + return batched_gather( + s.unsqueeze(-3), + edges, + -2, + no_batch_dims=len(s.shape) - 1, + ) + + +def knn_graph( + coords: torch.Tensor, + coord_mask: torch.Tensor, + padding_mask: torch.Tensor, + sequence_id: torch.Tensor, + *, + no_knn: int, +): + """Build nearest-neighbor edges, using sequence distance for missing geometry.""" + + length = coords.shape[-2] + coords = coords.nan_to_num() + missing_pair = ~(coord_mask[..., None, :] & coord_mask[..., :, None]) + excluded_pair = padding_mask[..., None, :] | padding_mask[..., :, None] + if sequence_id is not None: + excluded_pair |= sequence_id.unsqueeze(1) != sequence_id.unsqueeze(2) + + distances = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1) + residue_index = torch.arange(length, device=coords.device) + sequence_distance = (residue_index.unsqueeze(-1) - residue_index.unsqueeze(-2)).abs() + if not (distances[~missing_pair] < MAX_SUPPORTED_DISTANCE).all(): + raise ValueError( + "Coordinate pairwise distances exceed max supported distance " + f"({MAX_SUPPORTED_DISTANCE}). " + ) + + rank_distance = sequence_distance.to(distances.dtype).mul(1e2).add(MAX_SUPPORTED_DISTANCE) + rank_distance = rank_distance.where(missing_pair, distances) + rank_distance = rank_distance.masked_fill(excluded_pair, torch.inf) + sorted_distance, sorted_edge = rank_distance.sort(dim=-1, descending=False) + width = min(no_knn, length) + return sorted_edge[..., :width], sorted_distance[..., :width].isfinite() + + +def stack_variable_length_tensors( + sequences: Sequence[torch.Tensor], + constant_value: int | float = 0, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Pad arbitrary tensor dimensions to their maxima, then stack.""" + + output_shape = [ + len(sequences), + *np.max([sequence.shape for sequence in sequences], axis=0).tolist(), + ] + output = torch.full( + output_shape, + constant_value, + dtype=sequences[0].dtype if dtype is None else dtype, + device=sequences[0].device, + ) + for destination, source in zip(output, sequences, strict=True): + destination[tuple(slice(size) for size in source.shape)] = source + return output + + +def binpack( + tensor: torch.Tensor, + sequence_id: torch.Tensor | None, + pad_value: int | float, +): + """Scatter a sequence-major tensor into the packed layout described by IDs.""" + + if sequence_id is None: + return tensor + sequence_counts = sequence_id.max(dim=-1).values + 1 + output = torch.full( + sequence_id.shape + tensor.shape[2:], + fill_value=pad_value, + dtype=tensor.dtype, + device=tensor.device, + ) + source_index = 0 + for batch_index, (batch_ids, count) in enumerate( + zip(sequence_id, sequence_counts, strict=True) + ): + for seqid in range(count): + selection = batch_ids == seqid + output[batch_index, selection] = tensor[source_index, : selection.sum()] + source_index += 1 + return output + + +def unbinpack( + tensor: torch.Tensor, + sequence_id: torch.Tensor | None, + pad_value: int | float, +): + """Restore sequence-major rows from a packed tensor and its sequence IDs.""" + + if sequence_id is None: + return tensor + rows = [] + sequence_counts = sequence_id.max(dim=-1).values + 1 + for batch_index, (batch_ids, count) in enumerate( + zip(sequence_id, sequence_counts, strict=True) + ): + for seqid in range(count): + rows.append(tensor[batch_index, batch_ids == seqid]) + return stack_variable_length_tensors(rows, pad_value) + + +def merge_ranges( + ranges: list[range], + merge_gap_max: int | None = None, +) -> list[range]: + """Merge overlapping or sufficiently close ranges in positional order.""" + + maximum_gap = 0 if merge_gap_max is None else merge_gap_max + if not isinstance(maximum_gap, int) or isinstance(maximum_gap, bool): + raise TypeError("merge_gap_max must be an integer or None.") + if maximum_gap < 0: + raise ValueError(f"merge_gap_max must be non-negative, got {maximum_gap}.") + merged: list[range] = [] + for current in sorted(ranges, key=lambda item: item.start): + if not merged or merged[-1].stop + maximum_gap < current.start: + merged.append(current) + continue + previous = merged[-1] + merged[-1] = range(previous.start, max(previous.stop, current.stop)) + return merged + + +def merge_annotations( + annotations: list[FunctionAnnotation], + merge_gap_max: int | None = None, +) -> list[FunctionAnnotation]: + """Merge overlapping annotations independently for each label.""" + + grouped: dict[str, list[range]] = defaultdict(list) + for annotation in annotations: + grouped[annotation.label].append(range(annotation.start, annotation.end + 1)) + result = [] + for label, spans in grouped.items(): + result.extend( + FunctionAnnotation(label=label, start=span.start, end=span.stop - 1) + for span in merge_ranges(spans, merge_gap_max=merge_gap_max) + ) + return result + + +def get_chainbreak_boundaries_from_sequence( + sequence: Sequence[str], +) -> np.ndarray: + """Return half-open chain intervals split by chain-break tokens.""" + + boundaries = [0] + final_index = len(sequence) - 1 + for index, residue in enumerate(sequence): + if residue != CHAIN_BREAK_STR: + continue + if index == final_index: + raise ValueError( + "Encountered chain break token at end of sequence, this is unexpected." + ) + if index == final_index - 1: + warn( + "Encountered chain break token at penultimate position, this is unexpected.", + stacklevel=2, + ) + boundaries.extend((index, index + 1)) + boundaries.append(len(sequence)) + assert len(boundaries) % 2 == 0 + return np.asarray(boundaries).reshape(-1, 2) + + +def deserialize_tensors(data: bytes) -> Any: + """Decompress a tensor-only Torch payload onto CPU.""" + + decompressed = zstandard.ZstdDecompressor().decompress(data) + return torch.load( + BytesIO(decompressed), + map_location="cpu", + weights_only=True, + ) diff --git a/src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py b/src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py new file mode 100644 index 0000000..caa4a04 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py @@ -0,0 +1,469 @@ +"""Biotite-backed mmCIF parsing used by ESMFold2 structure records.""" + +from __future__ import annotations + +import functools +import io +import os +from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime + +import biotite.structure as bs +import biotite.structure.io.pdbx as pdbx +import numpy as np +from biotite.structure.io.pdbx import CIFColumn, CIFData, CIFFile + +from . import esmfold2_residue_constants as residue_constants + +PathOrBuffer = str | os.PathLike | io.StringIO + +PLDDT_B_FACTOR_SCALE = 100.0 +_MMCIF_COLUMN_DECIMALS = { + "Cartn_x": 3, + "Cartn_y": 3, + "Cartn_z": 3, + "B_iso_or_equiv": 2, +} +_NONPOLYMER_ENTITY_TYPES = frozenset({"NON-POLYMER", "WATER", "BRANCHED"}) + + +class NoProteinError(Exception): + """Raised internally when an mmCIF block contains no model-one atoms.""" + + +@dataclass +class Residue: + residue_number: int | None = None + insertion_code: str = "" + hetflag: bool = False + + +@dataclass +class MmcifHeader: + release_date: datetime | None = None + resolution: float | None = None + structure_method: str = "UNKNOWN" + + +def round_mmcif_columns(cif_file: CIFFile) -> None: + """Round coordinate and confidence columns in place for stable exports.""" + + if "atom_site" not in cif_file.block: + return + atom_site = cif_file.block["atom_site"] + for name, decimals in _MMCIF_COLUMN_DECIMALS.items(): + if name not in atom_site: + continue + original = atom_site[name] + values = original.as_array(np.float64) + strings = np.asarray( + [f"{value:.{decimals}f}" for value in values], + dtype=np.str_, + ) + atom_site[name] = CIFColumn( + data=CIFData(array=strings, dtype=np.str_), + mask=original.mask, + ) + + +def _clean_chain_list(value: str) -> list[str]: + return [chain.strip() for chain in value.split(",") if chain.strip()] + + +def _empty_residue() -> Residue: + return Residue(residue_number=None, insertion_code="", hetflag=False) + + +def _header_from_block( + block, + header: MmcifHeader | None = None, +) -> MmcifHeader: + header = MmcifHeader() if header is None else header + try: + if "pdbx_database_status" in block: + category = block["pdbx_database_status"] + if "recvd_initial_deposition_date" in category: + value = category["recvd_initial_deposition_date"].as_item() + if value and value != "?": + with suppress(ValueError): + header.release_date = datetime.strptime(value, "%Y-%m-%d") + if "refine" in block: + category = block["refine"] + if "ls_d_res_high" in category: + value = category["ls_d_res_high"].as_item() + if value and value != "?": + with suppress(ValueError): + header.resolution = float(value) + if "exptl" in block: + category = block["exptl"] + if "method" in category: + value = category["method"].as_item() + if value and value != "?": + header.structure_method = value.upper() + except Exception: + pass + return header + + +def _entities_from_block( + block, + entities: dict[int, list[str]] | None = None, +) -> dict[int, list[str]]: + entities = {} if entities is None else entities + if "entity" in block: + category = block["entity"] + ids = category["id"].as_array(str) + types = category["type"].as_array(str) + for entity_id, _ in zip(ids, types, strict=False): + entities[int(entity_id)] = [] + if "entity_poly" in block: + category = block["entity_poly"] + ids = category["entity_id"].as_array(str) + chain_lists = category["pdbx_strand_id"].as_array(str) + for raw_id, raw_chains in zip(ids, chain_lists, strict=False): + entity_id = int(raw_id) + if entity_id in entities: + entities[entity_id] = _clean_chain_list(raw_chains) + if "struct_asym" in block: + category = block["struct_asym"] + asym_ids = category["id"].as_array(str) + entity_ids = category["entity_id"].as_array(str) + for asym_id, raw_id in zip(asym_ids, entity_ids, strict=False): + entity_id = int(raw_id) + if entity_id in entities and not entities[entity_id]: + entities[entity_id].append(asym_id) + return entities + + +def _polymer_sequences(block) -> dict[str, str]: + sequences: dict[str, str] = {} + if "entity_poly" not in block: + return sequences + category = block["entity_poly"] + entity_ids = category["entity_id"].as_array(str) + raw_sequences = category["pdbx_seq_one_letter_code_can"].as_array(str) + chain_lists = category["pdbx_strand_id"].as_array(str) + for _, raw_sequence, raw_chains in zip( + entity_ids, + raw_sequences, + chain_lists, + strict=False, + ): + sequence = "".join(raw_sequence.split()) + for chain_id in _clean_chain_list(raw_chains): + sequences[chain_id] = sequence + return sequences + + +def _scheme_columns(category): + asym_ids = category["asym_id"].as_array(str) + insertion_codes = ( + category["pdb_ins_code"].as_array(str) + if "pdb_ins_code" in category + else [""] * len(asym_ids) + ) + hetflags = category["hetflag"].as_array(str) if "hetflag" in category else ["N"] * len(asym_ids) + author_chains = ( + category["pdb_strand_id"].as_array(str) if "pdb_strand_id" in category else asym_ids + ) + return ( + asym_ids, + category["seq_id"].as_array(str), + category["auth_seq_num"].as_array(str), + insertion_codes, + hetflags, + author_chains, + ) + + +def _scheme_residue_map(category): + ( + asym_ids, + sequence_positions, + author_numbers, + insertion_codes, + hetflags, + author_chains, + ) = _scheme_columns(category) + asym_to_author = { + asym_id: author_id for asym_id, author_id in zip(asym_ids, author_chains, strict=False) + } + per_chain: dict[str, dict[int, Residue]] = {} + for asym_id, raw_position, raw_number, raw_code, raw_hetflag in zip( + asym_ids, + sequence_positions, + author_numbers, + insertion_codes, + hetflags, + strict=False, + ): + residues = per_chain.setdefault(asym_id, {}) + try: + position = int(raw_position) - 1 + residue_number = int(raw_number) if raw_number != "?" else None + except ValueError: + continue + if residue_number is None: + insertion_code = "" + else: + insertion_code = "" if raw_code in (".", "?") else raw_code + residues[position] = Residue( + residue_number=residue_number, + insertion_code=insertion_code, + hetflag=raw_hetflag.upper() == "Y", + ) + return per_chain, asym_to_author + + +def _renumber_duplicate_residues( + per_chain: dict[str, dict[int, Residue]], +) -> None: + for residues in per_chain.values(): + positions_by_number: dict[int, list[int]] = {} + for position, residue in residues.items(): + if residue.residue_number is not None: + positions_by_number.setdefault(residue.residue_number, []).append(position) + for number, positions in positions_by_number.items(): + if len(positions) <= 1: + continue + positions.sort() + for offset, position in enumerate(positions): + previous = residues[position] + residues[position] = Residue( + residue_number=number + offset, + insertion_code=previous.insertion_code, + hetflag=previous.hetflag, + ) + + +def _ordered_scheme_mapping( + per_chain: dict[str, dict[int, Residue]], + asym_to_author: dict[str, str], + chain_sequences: dict[str, str], +) -> dict[str, dict[int, Residue]]: + result: dict[str, dict[int, Residue]] = {} + for asym_id, residues in per_chain.items(): + author_chain = asym_to_author.get(asym_id, asym_id) + if author_chain in chain_sequences: + result[author_chain] = { + position: residues.get(position, _empty_residue()) + for position in range(len(chain_sequences[author_chain])) + } + elif residues: + result[author_chain] = { + index: residues[position] for index, position in enumerate(sorted(residues)) + } + return result + + +def _complete_polymer_mappings( + mappings: dict[str, dict[int, Residue]], + chain_sequences: dict[str, str], +) -> None: + for chain_id, sequence in chain_sequences.items(): + mapping = mappings.setdefault(chain_id, {}) + for position in range(len(sequence)): + if position not in mapping: + mapping[position] = _empty_residue() + + +def _add_structure_fallbacks( + mappings: dict[str, dict[int, Residue]], + structure: bs.AtomArray, +) -> None: + if not ( + structure + and hasattr(structure, "chain_id") + and structure.chain_id is not None + and hasattr(structure.chain_id, "__iter__") + ): + return + for chain_id in set(structure.chain_id): + if chain_id in mappings: + continue + chain = structure[structure.chain_id == chain_id] + if not ( + hasattr(chain, "res_id") + and chain.res_id is not None + and hasattr(chain.res_id, "__iter__") + ): + continue + residue_ids = sorted(set(chain.res_id)) + mappings[chain_id] = { + index: Residue( + residue_number=residue_id, + insertion_code="", + hetflag=False, + ) + for index, residue_id in enumerate(residue_ids) + } + + +def _nonpolymer_entity_ids(block) -> set[str]: + result = set() + if "entity" not in block: + return result + category = block["entity"] + ids = category["id"].as_array(str) + types = category["type"].as_array(str) + for entity_id, entity_type in zip(ids, types, strict=False): + if entity_type.upper() in _NONPOLYMER_ENTITY_TYPES: + result.add(entity_id) + return result + + +def _nonpolymer_component_map(block, entity_ids: set[str]) -> dict[str, str]: + result = {} + if "pdbx_entity_nonpoly" not in block: + return result + category = block["pdbx_entity_nonpoly"] + ids = category["entity_id"].as_array(str) + components = category["comp_id"].as_array(str) + for entity_id, component in zip(ids, components, strict=False): + if entity_id in entity_ids: + result[entity_id] = component + return result + + +class MmcifWrapper: + """Parsed model-one structure, metadata, sequences, and residue mappings.""" + + def __init__(self, id: str | None = None): + self.id = id or "" + self.raw: pdbx.CIFFile | None = None + self.structure: bs.AtomArray + self.header = MmcifHeader() + self.entities: dict[int, list[str]] = {} + self.chain_to_seqres: dict[str, str] = {} + self.seqres_to_structure: dict[str, dict[int, Residue]] = {} + + @classmethod + def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper: + wrapper = cls(id=id) + wrapper._load(path) + return wrapper + + def _load(self, path: PathOrBuffer, fileid: str | None = None) -> None: + self.raw = pdbx.CIFFile.read(path) + self._parse_structure() + self._parse_header() + self._parse_entities() + self._parse_sequences() + + def _parse_structure(self) -> None: + try: + structure = pdbx.get_structure(self.raw, model=1) + if structure is None or not isinstance(structure, bs.AtomArray): + raise NoProteinError("No structure found in mmCIF file") + if len(structure) == 0: + raise NoProteinError("Empty structure in mmCIF file") + self.structure = structure + except Exception as error: + raise ValueError(f"Failed to parse structure: {error}") from error + + def _parse_header(self) -> None: + if self.raw: + self.header = _header_from_block(self.raw.block, self.header) + + def _parse_entities(self) -> None: + if not self.raw: + return + try: + self.entities = _entities_from_block(self.raw.block, self.entities) + except Exception: + if ( + self.structure + and hasattr(self.structure, "chain_id") + and self.structure.chain_id is not None + and hasattr(self.structure.chain_id, "__iter__") + ): + self.entities = {1: list(set(self.structure.chain_id))} + + def _parse_sequences(self) -> None: + if not self.raw: + return + block = self.raw.block + self.chain_to_seqres.update(_polymer_sequences(block)) + if "pdbx_poly_seq_scheme" in block: + per_chain, asym_to_author = _scheme_residue_map(block["pdbx_poly_seq_scheme"]) + _renumber_duplicate_residues(per_chain) + self.seqres_to_structure.update( + _ordered_scheme_mapping( + per_chain, + asym_to_author, + self.chain_to_seqres, + ) + ) + _complete_polymer_mappings( + self.seqres_to_structure, + self.chain_to_seqres, + ) + _add_structure_fallbacks(self.seqres_to_structure, self.structure) + + def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]: + assert self.raw is not None + block = self.raw.block + entity_ids = _nonpolymer_entity_ids(block) + _nonpolymer_component_map(block, entity_ids) + groups: dict[tuple[str, str], list[int]] = {} + if "atom_site" in block: + category = block["atom_site"] + chain_ids = category["label_asym_id"].as_array(str) + atom_entity_ids = category["label_entity_id"].as_array(str) + component_ids = category["label_comp_id"].as_array(str) + for index, (chain_id, entity_id, component_id) in enumerate( + zip(chain_ids, atom_entity_ids, component_ids, strict=False) + ): + if entity_id in entity_ids: + groups.setdefault((component_id, chain_id), []).append(index) + + coordinates = {} + for component_id, chain_id in groups: + selection = (self.structure.chain_id == chain_id) & ( + self.structure.res_name == component_id + ) + if not selection.any(): + continue + atoms = self.structure[selection] + if isinstance(atoms, (bs.AtomArray, bs.AtomArrayStack)) and len(atoms) > 0: + coordinates[(component_id, chain_id)] = atoms + return coordinates + + def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]: + result = {} + if not (self.structure and hasattr(self.structure, "chain_id")): + return result + standard_residues = set(residue_constants.resnames[:-1]) + standard_residues.update({"A", "C", "G", "T", "U"}) + if self.structure.chain_id is None: + return result + for chain_id in set(self.structure.chain_id): + chain = self.structure[self.structure.chain_id == chain_id] + if not ( + hasattr(chain, "res_name") + and chain.res_name is not None + and hasattr(chain.res_name, "__iter__") + ): + continue + for residue_name in set(chain.res_name): + if residue_name in standard_residues: + continue + selection = (chain.chain_id == chain_id) & (chain.res_name == residue_name) + if selection.any() and isinstance( + chain, + (bs.AtomArray, bs.AtomArrayStack), + ): + result[(residue_name, chain_id)] = chain[selection] + return result + + @functools.cached_property + def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]: + """Map each non-polymer component and chain to its atoms.""" + + if not self.structure or not self.raw: + return {} + try: + return self._parse_nonpoly_from_mmcif() + except Exception: + return self._parse_nonpoly_fallback() diff --git a/src/fastplms/models/esmfold2/esmfold2_molecular_complex.py b/src/fastplms/models/esmfold2/esmfold2_molecular_complex.py new file mode 100644 index 0000000..3c89215 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_molecular_complex.py @@ -0,0 +1,1016 @@ +"""Flat molecular-complex records used by the ESMFold2 public API. + +The folding model operates on tokens and a single atom table. This module owns +that representation, its protein-only bridge, mmCIF I/O, structure metrics, and +the compact wire format. It deliberately has no dependency on the upstream +Biohub package; the pinned submodule is used only by differential tests. +""" + +from __future__ import annotations + +import io +import os +import re +from dataclasses import asdict, dataclass +from pathlib import Path +from subprocess import check_output +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any + +import biotite.structure as bs +import biotite.structure.io.pdbx as pdbx +import brotli +import msgpack +import numpy as np +import torch +from biotite.structure.io.pdbx import ( + CIFCategory, + CIFColumn, + CIFData, + CIFFile, + set_structure, +) + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_metrics import compute_lddt, compute_rmsd +from .esmfold2_mmcif_parsing import PLDDT_B_FACTOR_SCALE, round_mmcif_columns +from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata + + +@dataclass +class MolecularComplexResult: + """One folded complex and the optional model outputs associated with it.""" + + complex: MolecularComplex + plddt: torch.Tensor | None = None + ptm: float | None = None + iptm: float | None = None + pae: torch.Tensor | None = None + distogram: torch.Tensor | None = None + pair_chains_iptm: torch.Tensor | None = None + output_embedding_sequence: torch.Tensor | None = None + output_embedding_pair_pooled: torch.Tensor | None = None + residue_index: torch.Tensor | None = None + entity_id: torch.Tensor | None = None + sae_features: np.ndarray | None = None # X has shape (l, n_features). + ttt_metrics: dict[str, Any] | None = None + + +@dataclass +class MolecularComplexMetadata: + """Entity and chain labels carried with a molecular complex.""" + + entity_lookup: dict[int, str] + chain_lookup: dict[int, str] + assembly_composition: dict[str, list[str]] | None = None + + +@dataclass +class Molecule: + """The atom slice represented by one model token.""" + + token: str + token_idx: int + atom_positions: np.ndarray # P has shape (n_atoms, 3). + atom_elements: np.ndarray # E has shape (n_atoms,). + atom_names: np.ndarray | None = None # N has shape (n_atoms,) when present. + atom_hetero: np.ndarray | None = None # M has shape (n_atoms,) when present. + residue_type: int = 0 + molecule_type: int = 0 + confidence: float = 0.0 + + +_NUCLEOTIDE_NAMES = frozenset({"A", "T", "G", "C", "U", "DA", "DT", "DG", "DC"}) +_SERIALIZED_ARRAYS = frozenset( + { + "atom_positions", + "atom_elements", + "atom_names", + "atom_hetero", + "token_to_atoms", + "chain_id", + "entity_id", + "sym_id", + "plddt", + } +) + + +def _assert_table_lengths(complex_value: MolecularComplex) -> None: + """Check that token and atom annotations align with their tables.""" + if not isinstance(complex_value.sequence, list) or any( + not isinstance(token, str) for token in complex_value.sequence + ): + raise TypeError("sequence must be a list of token strings.") + n_tokens = len(complex_value.sequence) + if not isinstance(complex_value.atom_positions, np.ndarray): + raise TypeError("atom_positions must be a NumPy array.") + if complex_value.atom_positions.ndim != 2 or complex_value.atom_positions.shape[1:] != ( + 3, + ): + raise ValueError( + "atom_positions must have shape (n_atoms, 3), got " + f"{complex_value.atom_positions.shape}." + ) + if not np.issubdtype(complex_value.atom_positions.dtype, np.number): + raise TypeError("atom_positions must use a numeric dtype.") + n_atoms = len(complex_value.atom_positions) + if not isinstance(complex_value.atom_elements, np.ndarray): + raise TypeError("atom_elements must be a NumPy array.") + if complex_value.atom_elements.shape != (n_atoms,): + raise ValueError( + f"atom_elements shape {complex_value.atom_elements.shape} != {n_atoms} atoms" + ) + token_tables = { + "token_to_atoms": complex_value.token_to_atoms, + "chain_id": complex_value.chain_id, + "plddt": complex_value.plddt, + } + if complex_value.entity_id is not None: + token_tables["entity_id"] = complex_value.entity_id + if complex_value.sym_id is not None: + token_tables["sym_id"] = complex_value.sym_id + for label, values in token_tables.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != n_tokens: + raise ValueError(f"{label} shape {values.shape} != {n_tokens} tokens") + if complex_value.token_to_atoms.shape != (n_tokens, 2): + raise ValueError( + "token_to_atoms must have shape " + f"({n_tokens}, 2), got {complex_value.token_to_atoms.shape}." + ) + if not np.issubdtype(complex_value.token_to_atoms.dtype, np.integer): + raise TypeError("token_to_atoms must use an integer dtype.") + if complex_value.chain_id.shape != (n_tokens,): + raise ValueError(f"chain_id must have shape ({n_tokens},).") + for label, values in ( + ("chain_id", complex_value.chain_id), + ("entity_id", complex_value.entity_id), + ("sym_id", complex_value.sym_id), + ): + if values is not None and values.shape != (n_tokens,): + raise ValueError(f"{label} must have shape ({n_tokens},).") + if values is not None and not np.issubdtype(values.dtype, np.integer): + raise TypeError(f"{label} must use an integer dtype.") + if complex_value.plddt.shape != (n_tokens,): + raise ValueError(f"plddt must have shape ({n_tokens},).") + if not np.issubdtype(complex_value.plddt.dtype, np.number): + raise TypeError("plddt must use a numeric dtype.") + if n_tokens: + starts = complex_value.token_to_atoms[:, 0] + stops = complex_value.token_to_atoms[:, 1] + if np.any(starts < 0) or np.any(stops < starts) or np.any(stops > n_atoms): + raise ValueError("token_to_atoms contains an invalid or out-of-bounds atom span.") + for label, values in ( + ("atom_names", complex_value.atom_names), + ("atom_hetero", complex_value.atom_hetero), + ): + if values is not None and not isinstance(values, np.ndarray): + raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.") + if isinstance(values, np.ndarray) and values.shape != (n_atoms,): + raise ValueError(f"{label} shape {values.shape} != {n_atoms} atoms") + + +def _flat_protein_atoms( + protein: ProteinComplex, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Flatten the populated atom37 entries of a protein complex.""" + positions: list[np.ndarray] = [] + elements: list[str] = [] + names: list[str] = [] + hetero: list[bool] = [] + spans: list[tuple[int, int]] = [] + + for sequence_index, residue in enumerate(protein.sequence): + if residue == "|": + continue + start = len(positions) + mask = protein.atom37_mask[sequence_index] + residue_positions = protein.atom37_positions[sequence_index] + for atom_index in np.flatnonzero(mask): + atom_name = residue_constants.atom_types[int(atom_index)] + positions.append(residue_positions[atom_index]) + elements.append(atom_name[0] if atom_name else "C") + names.append(atom_name) + hetero.append(False) + spans.append((start, len(positions))) + + return ( + np.asarray(positions, dtype=np.float32), + np.asarray(elements, dtype=object), + np.asarray(names, dtype=object), + np.asarray(hetero, dtype=bool), + np.asarray(spans, dtype=np.int32), + ) + + +def _protein_sequence_and_indices( + complex_value: MolecularComplex, +) -> tuple[list[int], str, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + protein_indices = [ + index + for index, token in enumerate(complex_value.sequence) + if token in residue_constants.restype_3to1 + ] + if not protein_indices: + raise ValueError("No protein tokens found in MolecularComplex") + + chain_ids = complex_value.chain_id[protein_indices] + entity_ids = ( + chain_ids + if complex_value.entity_id is None + else complex_value.entity_id[protein_indices] + ) + sym_ids = ( + np.zeros_like(chain_ids) + if complex_value.sym_id is None + else complex_value.sym_id[protein_indices] + ) + confidences = complex_value.plddt[protein_indices] + sequence: list[str] = [] + previous_instance: Any = None + preserve_instances = complex_value.sym_id is not None + for index, chain_id, sym_id in zip( + protein_indices, chain_ids, sym_ids, strict=True + ): + instance = (int(chain_id), int(sym_id)) if preserve_instances else int(chain_id) + if previous_instance is not None and instance != previous_instance: + sequence.append("|") + sequence.append(residue_constants.restype_3to1[complex_value.sequence[index]]) + previous_instance = instance + return protein_indices, "".join(sequence), chain_ids, entity_ids, sym_ids, confidences + + +def _protein_entity_metadata_value(value: int | str) -> int | str: + """Restore the numeric entity labels used by ProteinComplex metadata.""" + if isinstance(value, str): + try: + return int(value) + except ValueError: + pass + return value + + +def _atom37_from_flat( + complex_value: MolecularComplex, protein_indices: list[int] +) -> tuple[np.ndarray, np.ndarray]: + n_residues = len(protein_indices) + positions = np.full((n_residues, 37, 3), np.nan, dtype=np.float32) + mask = np.zeros((n_residues, 37), dtype=bool) + if complex_value.atom_names is None: + return positions, mask + + for residue_index, token_index in enumerate(protein_indices): + start, stop = complex_value.token_to_atoms[token_index] + seen: set[str] = set() + for atom_name, atom_position in zip( + complex_value.atom_names[start:stop], + complex_value.atom_positions[start:stop], + strict=True, + ): + normalized = str(atom_name).upper().strip() + if normalized in seen: + continue + seen.add(normalized) + atom37_index = residue_constants.atom_order.get(normalized) + if atom37_index is not None: + positions[residue_index, atom37_index] = atom_position + mask[residue_index, atom37_index] = True + return positions, mask + + +def _expand_protein_rows( + sequence: str, + protein_chain_ids: np.ndarray, + protein_entity_ids: np.ndarray, + protein_sym_ids: np.ndarray, + confidences: np.ndarray, + compact_positions: np.ndarray, + compact_mask: np.ndarray, +) -> dict[str, np.ndarray]: + """Insert empty rows at chain separators in a protein representation.""" + n_positions = len(sequence) + expanded = { + "chain_id": np.full(n_positions, -1, dtype=np.int64), + "entity_id": np.full(n_positions, -1, dtype=np.int64), + "sym_id": np.zeros(n_positions, dtype=np.int64), + "residue_index": np.zeros(n_positions, dtype=np.int64), + "insertion_code": np.asarray([""] * n_positions, dtype=object), + "confidence": np.zeros(n_positions, dtype=np.float32), + "atom37_positions": np.full((n_positions, 37, 3), np.nan, dtype=np.float32), + "atom37_mask": np.zeros((n_positions, 37), dtype=bool), + } + residue_number = 0 + compact_index = 0 + for sequence_index, residue in enumerate(sequence): + if residue == "|": + residue_number = 0 + continue + chain_id = protein_chain_ids[compact_index] + residue_number += 1 + expanded["chain_id"][sequence_index] = chain_id + expanded["entity_id"][sequence_index] = protein_entity_ids[compact_index] + expanded["sym_id"][sequence_index] = protein_sym_ids[compact_index] + expanded["residue_index"][sequence_index] = residue_number + expanded["confidence"][sequence_index] = confidences[compact_index] + expanded["atom37_positions"][sequence_index] = compact_positions[compact_index] + expanded["atom37_mask"][sequence_index] = compact_mask[compact_index] + compact_index += 1 + return expanded + + +def _read_cif(source: str) -> CIFFile: + if os.path.exists(source): + return pdbx.CIFFile.read(source) + return pdbx.CIFFile.read(io.StringIO(source)) + + +def _read_structure(cif_file: CIFFile) -> Any: + try: + return pdbx.get_structure(cif_file, model=1, extra_fields=["b_factor"]) + except (KeyError, ValueError): + try: + return pdbx.get_structure(cif_file) + except Exception: + return pdbx.get_structure(cif_file, model=None) + + +def _column_array(category: Any, name: str) -> np.ndarray: + column = category[name] + if hasattr(column, "as_array"): + return column.as_array(str) + return np.asarray(list(column), dtype=str) + + +def _label_asym_ids(cif_file: CIFFile, n_structure_atoms: int) -> list[str] | None: + """Return label-asym identifiers after applying Biohub's atom filters.""" + block = cif_file.block + if "atom_site" not in block or "label_asym_id" not in block["atom_site"]: + return None + atom_site = block["atom_site"] + labels = _column_array(atom_site, "label_asym_id") + keep = np.ones(len(labels), dtype=bool) + if "pdbx_PDB_model_num" in atom_site: + keep &= _column_array(atom_site, "pdbx_PDB_model_num") == "1" + if "label_alt_id" in atom_site: + keep &= np.isin(_column_array(atom_site, "label_alt_id"), [".", "?", "", "A"]) + filtered = labels[keep] + return filtered.tolist() if len(filtered) == n_structure_atoms else None + + +def _entity_metadata(cif_file: CIFFile) -> dict[Any, Any]: + result: dict[Any, Any] = {} + try: + category = cif_file.block["entity"] + if "id" not in category or "type" not in category: + return result + for entity_id, entity_type in zip(category["id"], category["type"], strict=False): + result[entity_id] = entity_type + except Exception: + return {} + return result + + +def _group_structure_atoms( + structure: Any, labels: list[str] | None +) -> dict[str, dict[tuple[int, str], dict[str, Any]]]: + grouped: dict[str, dict[tuple[int, str], dict[str, Any]]] = {} + for atom_index, atom in enumerate(structure): + chain = labels[atom_index] if labels is not None else atom.chain_id + residues = grouped.setdefault(chain, {}) + key = (atom.res_id, atom.res_name) + record = residues.setdefault( + key, + {"atoms": [], "res_name": atom.res_name, "is_hetero": atom.hetero}, + ) + record["atoms"].append(atom) + return grouped + + +def _flatten_structure_groups( + grouped: dict[str, dict[tuple[int, str], dict[str, Any]]], +) -> tuple[ + list[str], + list[np.ndarray], + list[str], + list[str], + list[bool], + list[tuple[int, int]], + list[float], + list[int], + dict[str, int], +]: + tokens: list[str] = [] + positions: list[np.ndarray] = [] + elements: list[str] = [] + names: list[str] = [] + hetero: list[bool] = [] + spans: list[tuple[int, int]] = [] + confidences: list[float] = [] + token_chains: list[int] = [] + chain_numbers = {chain: index for index, chain in enumerate(sorted(grouped))} + + for chain in sorted(grouped): + for residue_key in sorted(grouped[chain]): + record = grouped[chain][residue_key] + if record["res_name"] == "HOH": + continue + atoms = record["atoms"] + tokens.append(record["res_name"]) + token_chains.append(chain_numbers[chain]) + start = len(positions) + positions.extend(atom.coord for atom in atoms) + elements.extend(atom.element for atom in atoms) + names.extend(atom.atom_name for atom in atoms) + hetero.extend(atom.hetero for atom in atoms) + spans.append((start, len(positions))) + b_factor = getattr(atoms[0], "b_factor", 50.0) if atoms else 50.0 + confidences.append(min(b_factor / PLDDT_B_FACTOR_SCALE, 1.0)) + return ( + tokens, + positions, + elements, + names, + hetero, + spans, + confidences, + token_chains, + chain_numbers, + ) + + +def _chain_entity_maps( + complex_value: MolecularComplex, +) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]: + chains: dict[str, list[str]] = {} + for token_index, numeric_chain in enumerate(complex_value.chain_id): + numeric = int(numeric_chain) + label = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric)) + chains.setdefault(label, []).append(complex_value.sequence[token_index]) + + sequence_entities: dict[tuple[str, ...], int] = {} + chain_entities: dict[str, int] = {} + entity_sequences: dict[int, tuple[str, ...]] = {} + for label, sequence in chains.items(): + key = tuple(sequence) + entity_id = sequence_entities.get(key) + if entity_id is None: + entity_id = len(sequence_entities) + 1 + sequence_entities[key] = entity_id + entity_sequences[entity_id] = key + chain_entities[label] = entity_id + return chains, chain_entities, entity_sequences + + +def _cif_column(values: list[str]) -> CIFColumn: + return CIFColumn(data=CIFData(array=np.asarray(values), dtype=np.str_)) + + +def _add_entity_categories( + cif_file: CIFFile, + complex_value: MolecularComplex, + entity_sequences: dict[int, tuple[str, ...]], +) -> None: + ids: list[str] = [] + types: list[str] = [] + descriptions: list[str] = [] + for entity_id in sorted(entity_sequences): + sequence = entity_sequences[entity_id] + protein = any(token in residue_constants.restype_3to1 for token in sequence) + nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence) + ids.append(str(entity_id)) + types.append("polymer" if protein or nucleic else "non-polymer") + if protein: + descriptions.append(f"Polymer entity {entity_id} (protein)") + elif nucleic: + descriptions.append(f"Polymer entity {entity_id} (nucleic acid)") + else: + descriptions.append(f"Non-polymer entity {entity_id}") + + if ids: + cif_file.block["entity"] = CIFCategory( + name="entity", + columns={ + "id": _cif_column(ids), + "type": _cif_column(types), + "pdbx_description": _cif_column(descriptions), + }, + ) + + _, chain_entities, _ = _chain_entity_maps(complex_value) + if chain_entities: + labels = sorted(chain_entities) + cif_file.block["struct_asym"] = CIFCategory( + name="struct_asym", + columns={ + "id": _cif_column(labels), + "entity_id": _cif_column([str(chain_entities[label]) for label in labels]), + }, + ) + + entity_chains: dict[int, list[str]] = {} + for chain, entity_id in chain_entities.items(): + entity_chains.setdefault(entity_id, []).append(chain) + polymer_rows: list[tuple[str, str, str, str]] = [] + residue_rows: list[tuple[str, str, str, str]] = [] + for entity_id in sorted(entity_sequences): + sequence = entity_sequences[entity_id] + protein = any(token in residue_constants.restype_3to1 for token in sequence) + nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence) + if not (protein or nucleic): + continue + if protein: + polymer_type = "polypeptide(L)" + canonical = "".join( + residue_constants.restype_3to1.get(token, "(X)") for token in sequence + ) + else: + polymer_type = ( + "polyribonucleotide" + if "U" in sequence + else ( + "polydeoxyribonucleotide" + if any(token in {"DA", "DT", "DG", "DC"} for token in sequence) + else "polyribonucleotide" + ) + ) + nucleotide_letters = {"DA": "A", "DT": "T", "DG": "G", "DC": "C"} + canonical = "".join(nucleotide_letters.get(token, token) for token in sequence) + strand_ids = ",".join(sorted(entity_chains.get(entity_id, []))) or "?" + polymer_rows.append((str(entity_id), polymer_type, strand_ids, canonical)) + residue_rows.extend( + (str(entity_id), str(number), token, "n") + for number, token in enumerate(sequence, start=1) + ) + + if polymer_rows: + columns = list(zip(*polymer_rows, strict=True)) + cif_file.block["entity_poly"] = CIFCategory( + name="entity_poly", + columns={ + "entity_id": _cif_column(list(columns[0])), + "type": _cif_column(list(columns[1])), + "pdbx_strand_id": _cif_column(list(columns[2])), + "pdbx_seq_one_letter_code_can": _cif_column(list(columns[3])), + }, + ) + if residue_rows: + columns = list(zip(*residue_rows, strict=True)) + cif_file.block["entity_poly_seq"] = CIFCategory( + name="entity_poly_seq", + columns={ + "entity_id": _cif_column(list(columns[0])), + "num": _cif_column(list(columns[1])), + "mon_id": _cif_column(list(columns[2])), + "hetero": _cif_column(list(columns[3])), + }, + ) + + +def _fallback_atom_names(token: str, count: int) -> list[str]: + if token in residue_constants.restype_3to1: + names = list(residue_constants.residue_atoms.get(token, ["N", "CA", "C", "O"]))[:count] + names.extend(f"X{index + 1}" for index in range(len(names), count)) + return names + return [f"C{index + 1}" for index in range(count)] + + +def _as_atom_array(complex_value: MolecularComplex, chain_entities: dict[str, int]) -> bs.AtomArray: + n_atoms = len(complex_value.atom_positions) + atom_array = bs.AtomArray(length=n_atoms) + atom_array.coord = complex_value.atom_positions + residue_ids = np.zeros(n_atoms, dtype=np.int32) + chain_labels = np.empty(n_atoms, dtype=object) + residue_names = np.empty(n_atoms, dtype=object) + hetero = np.zeros(n_atoms, dtype=bool) + b_factors = np.zeros(n_atoms, dtype=np.float32) + atom_names = np.empty(n_atoms, dtype=object) + entity_ids = np.zeros(n_atoms, dtype=np.int32) + next_residue: dict[Any, int] = {} + + for token_index, (start, stop) in enumerate(complex_value.token_to_atoms): + token = complex_value.sequence[token_index] + numeric_chain = complex_value.chain_id[token_index] + numeric = int(numeric_chain) + chain = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric)) + residue_id = next_residue.get(numeric_chain, 0) + 1 + next_residue[numeric_chain] = residue_id + count = int(stop - start) + names = ( + list(complex_value.atom_names[start:stop]) + if complex_value.atom_names is not None + else _fallback_atom_names(token, count) + ) + residue_ids[start:stop] = residue_id + chain_labels[start:stop] = chain + residue_names[start:stop] = token + hetero[start:stop] = ( + complex_value.atom_hetero[start:stop] + if complex_value.atom_hetero is not None + else token not in residue_constants.restype_3to1 + ) + b_factors[start:stop] = complex_value.plddt[token_index] * PLDDT_B_FACTOR_SCALE + atom_names[start:stop] = names + entity_ids[start:stop] = chain_entities.get(chain, 1) + + atom_array.res_id = residue_ids + atom_array.chain_id = np.asarray(chain_labels, dtype="U16") + atom_array.res_name = np.asarray(residue_names, dtype="U8") + atom_array.hetero = hetero + atom_array.atom_name = np.asarray(atom_names, dtype="U4") + atom_array.add_annotation("b_factor", dtype=float) + atom_array.b_factor = b_factors + atom_array.add_annotation("occupancy", dtype=float) + atom_array.occupancy = np.ones(n_atoms, dtype=np.float32) + atom_array.add_annotation("entity_id", dtype=int) + atom_array.entity_id = entity_ids + if complex_value.atom_elements is not None and len(complex_value.atom_elements) == n_atoms: + atom_array.element = np.asarray(complex_value.atom_elements, dtype="U4") + else: + atom_array.element = bs.infer_elements(atom_array) + return atom_array + + +def _repair_label_entity_ids(cif_file: CIFFile, chain_entities: dict[str, int]) -> None: + if "atom_site" not in cif_file.block: + return + atom_site = cif_file.block["atom_site"] + if "label_asym_id" not in atom_site or "label_entity_id" not in atom_site: + return + labels = _column_array(atom_site, "label_asym_id").tolist() + if labels: + atom_site["label_entity_id"] = _cif_column( + [str(chain_entities.get(label, 1)) for label in labels] + ) + + +def _centroid_tensors( + mobile: MolecularComplex, + target: MolecularComplex, + *, + retain_missing: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if len(mobile) != len(target): + raise ValueError( + f"Complexes must have the same number of tokens: {len(mobile)} vs {len(target)}" + ) + mobile_centers: list[np.ndarray] = [] + target_centers: list[np.ndarray] = [] + valid: list[bool] = [] + for token_index in range(len(mobile)): + mobile_start, mobile_stop = mobile.token_to_atoms[token_index] + target_start, target_stop = target.token_to_atoms[token_index] + mobile_atoms = mobile.atom_positions[mobile_start:mobile_stop] + target_atoms = target.atom_positions[target_start:target_stop] + present = len(mobile_atoms) > 0 and len(target_atoms) > 0 + if not present and not retain_missing: + continue + if present: + mobile_centers.append(mobile_atoms.mean(axis=0)) + target_centers.append(target_atoms.mean(axis=0)) + else: + mobile_centers.append(np.full(3, np.nan)) + target_centers.append(np.full(3, np.nan)) + valid.append(present) + if not any(valid): + metric = "LDDT" if retain_missing else "RMSD" + raise ValueError(f"No valid atoms found for {metric} computation") + return ( + torch.from_numpy(np.stack(mobile_centers)).unsqueeze(0), + torch.from_numpy(np.stack(target_centers)).unsqueeze(0), + torch.as_tensor(valid, dtype=torch.bool).unsqueeze(0), + ) + + +@dataclass(frozen=True) +class MolecularComplex: + """A token sequence backed by one contiguous atom table. + + P stores atom coordinates with shape (n_atoms, 3). Token span ``i`` is + ``P[token_to_atoms[i, 0]:token_to_atoms[i, 1]]``. ``chain_id`` identifies + the author chain, while optional ``entity_id`` and ``sym_id`` distinguish + biological entities and repeated chain instances. + """ + + id: str + sequence: list[str] + atom_positions: np.ndarray # P has shape (n_atoms, 3). + atom_elements: np.ndarray # E has shape (n_atoms,). + token_to_atoms: np.ndarray # I has shape (n_tokens, 2). + chain_id: np.ndarray # C has shape (n_tokens,). + plddt: np.ndarray # S has shape (n_tokens,). + metadata: MolecularComplexMetadata + atom_names: np.ndarray | None = None # N has shape (n_atoms,) when present. + atom_hetero: np.ndarray | None = None # M has shape (n_atoms,) when present. + # These token-aligned IDs are optional for compatibility with older blobs. + # ProteinComplex adapters populate them so homomers and repeated author-chain + # labels survive a MolecularComplex round trip. + entity_id: np.ndarray | None = None + sym_id: np.ndarray | None = None + + def __post_init__(self) -> None: + _assert_table_lengths(self) + + def __len__(self) -> int: + return len(self.sequence) + + def __getitem__(self, idx: int) -> Molecule: + if idx < 0 or idx >= len(self): + raise IndexError(f"Token index {idx} out of range for {len(self)} tokens") + start, stop = self.token_to_atoms[idx] + return Molecule( + token=self.sequence[idx], + token_idx=idx, + atom_positions=self.atom_positions[start:stop], + atom_elements=self.atom_elements[start:stop], + atom_names=None if self.atom_names is None else self.atom_names[start:stop], + atom_hetero=(None if self.atom_hetero is None else self.atom_hetero[start:stop]), + residue_type=0, + molecule_type=0, + confidence=self.plddt[idx], + ) + + @property + def atom_coordinates(self) -> np.ndarray: + """Return P, the flat atom-coordinate table with shape (n_atoms, 3).""" + return self.atom_positions + + @classmethod + def from_protein_complex(cls, pc: ProteinComplex) -> MolecularComplex: + positions, elements, names, hetero, spans = _flat_protein_atoms(pc) + residue_positions = [index for index, value in enumerate(pc.sequence) if value != "|"] + metadata = MolecularComplexMetadata( + entity_lookup={key: str(value) for key, value in pc.metadata.entity_lookup.items()}, + chain_lookup=dict(pc.metadata.chain_lookup), + assembly_composition=pc.metadata.assembly_composition, + ) + return cls( + id=pc.id, + sequence=[ + residue_constants.restype_1to3.get(pc.sequence[index], "UNK") + for index in residue_positions + ], + atom_positions=positions, + atom_elements=elements, + token_to_atoms=spans, + chain_id=np.asarray(pc.chain_id[residue_positions], dtype=np.int64), + plddt=np.asarray(pc.confidence[residue_positions], dtype=np.float32), + metadata=metadata, + atom_names=names, + atom_hetero=hetero, + entity_id=np.asarray(pc.entity_id[residue_positions], dtype=np.int64), + sym_id=np.asarray(pc.sym_id[residue_positions], dtype=np.int64), + ) + + def to_protein_complex(self) -> ProteinComplex: + ( + protein_indices, + sequence, + chain_ids, + entity_ids, + sym_ids, + confidences, + ) = _protein_sequence_and_indices(self) + compact_positions, compact_mask = _atom37_from_flat(self, protein_indices) + arrays = _expand_protein_rows( + sequence, + chain_ids, + entity_ids, + sym_ids, + confidences, + compact_positions, + compact_mask, + ) + unique_chains = np.unique(chain_ids) + unique_entities = np.unique(entity_ids) + metadata = ProteinComplexMetadata( + entity_lookup={ + int(entity): _protein_entity_metadata_value( + self.metadata.entity_lookup.get(int(entity), int(entity)) + ) + for entity in unique_entities + }, + chain_lookup={ + int(chain): self.metadata.chain_lookup.get(int(chain), chr(65 + int(chain))) + for chain in unique_chains + }, + assembly_composition=self.metadata.assembly_composition, + ) + return ProteinComplex( + id=self.id, + sequence=sequence, + entity_id=arrays["entity_id"], + chain_id=arrays["chain_id"], + sym_id=arrays["sym_id"], + residue_index=arrays["residue_index"], + insertion_code=arrays["insertion_code"], + atom37_positions=arrays["atom37_positions"], + atom37_mask=arrays["atom37_mask"], + confidence=arrays["confidence"], + metadata=metadata, + ) + + @classmethod + def from_mmcif(cls, inp: str, id: str | None = None) -> MolecularComplex: + cif_file = _read_cif(inp) + structure = _read_structure(cif_file) + if TYPE_CHECKING: + structure: Any = structure + labels = _label_asym_ids(cif_file, len(structure)) + grouped = _group_structure_atoms(structure, labels) + ( + tokens, + positions, + elements, + names, + hetero, + spans, + confidences, + token_chains, + chain_numbers, + ) = _flatten_structure_groups(grouped) + n_tokens = len(tokens) + if positions: + position_array = np.asarray(positions, dtype=np.float32) + element_array = np.asarray(elements, dtype=object) + name_array = np.asarray(names, dtype=object) + hetero_array = np.asarray(hetero, dtype=bool) + span_array = np.asarray(spans, dtype=np.int32) + chain_array = np.asarray(token_chains, dtype=np.int64) + else: + position_array = np.zeros((0, 3), dtype=np.float32) + element_array = np.zeros(0, dtype=object) + name_array = np.zeros(0, dtype=object) + hetero_array = np.zeros(0, dtype=bool) + span_array = np.zeros((n_tokens, 2), dtype=np.int32) + chain_array = ( + np.asarray(token_chains, dtype=np.int64) + if token_chains + else np.zeros(n_tokens, dtype=np.int64) + ) + complex_id = id or (Path(inp).stem if os.path.exists(inp) else "complex_from_string") + return cls( + id=complex_id, + sequence=tokens, + atom_positions=position_array, + atom_elements=element_array, + token_to_atoms=span_array, + chain_id=chain_array, + plddt=np.asarray(confidences, dtype=np.float32), + metadata=MolecularComplexMetadata( + entity_lookup=_entity_metadata(cif_file), + chain_lookup={number: chain for chain, number in chain_numbers.items()}, + assembly_composition=None, + ), + atom_names=name_array, + atom_hetero=hetero_array, + ) + + def _get_entity_mapping( + self, + ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]: + return _chain_entity_maps(self) + + def _add_entity_information( + self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]] + ) -> None: + _add_entity_categories(cif_file, self, entity_sequences) + + def to_mmcif(self) -> str: + _, chain_entities, entity_sequences = _chain_entity_maps(self) + atom_array = _as_atom_array(self, chain_entities) + cif_file = CIFFile() + set_structure(cif_file, atom_array, data_block=self.id) + _repair_label_entity_ids(cif_file, chain_entities) + _add_entity_categories(cif_file, self, entity_sequences) + round_mmcif_columns(cif_file) + output = io.StringIO() + cif_file.write(output) + return output.getvalue() + + def dockq(self, native: MolecularComplex) -> Any: + try: + mobile = self.to_protein_complex().normalize_chain_ids_for_pdb() + target = native.to_protein_complex().normalize_chain_ids_for_pdb() + except ValueError as error: + raise ValueError( + f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}" + ) from None + try: + return mobile.dockq(target) + except Exception: + return self._compute_dockq_manual(native) + + def _compute_dockq_manual(self, native: MolecularComplex) -> Any: + try: + mobile = self.to_protein_complex().normalize_chain_ids_for_pdb() + target = native.to_protein_complex().normalize_chain_ids_for_pdb() + except ValueError as error: + raise ValueError( + f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}" + ) from None + with TemporaryDirectory() as directory: + mobile_path = Path(directory) / "self.pdb" + target_path = Path(directory) / "native.pdb" + mobile.to_pdb(mobile_path) + target.to_pdb(target_path) + try: + raw_output = check_output(["DockQ", str(mobile_path), str(target_path)]) + output = raw_output.decode() + score: float | None = None + for line in output.split("\n"): + if "Total DockQ" in line: + match = re.search(r"Total DockQ.*: ([\d.]+)", line) + if match: + score = float(match.group(1)) + break + if score is None: + for line in output.split("\n"): + if line.startswith("DockQ") and ":" in line: + try: + score = float(line.split(":")[1].strip()) + break + except (ValueError, IndexError): + continue + if score is None: + raise ValueError("Could not parse DockQ score from output") + return {"total_dockq": score, "raw_output": output, "aligned": self} + except FileNotFoundError: + raise RuntimeError( + "DockQ is not installed. Please install DockQ to use this method." + ) from None + except Exception as error: + raise RuntimeError(f"DockQ computation failed: {error}") from error + + def rmsd(self, target: MolecularComplex, **kwargs: Any) -> float: + mobile, reference, mask = _centroid_tensors(self, target, retain_missing=False) + value = compute_rmsd( + mobile=mobile, + target=reference, + atom_exists_mask=mask, + reduction="batch", + **kwargs, + ) + return float(value) + + def lddt_ca(self, target: MolecularComplex, **kwargs: Any) -> float: + mobile, reference, mask = _centroid_tensors(self, target, retain_missing=True) + value = compute_lddt( + all_atom_pred_pos=mobile, + all_atom_positions=reference, + all_atom_mask=mask, + per_residue=False, + **kwargs, + ) + return float(value) + + def state_dict(self) -> dict[str, Any]: + state = dict(vars(self)) + for optional_identity in ("entity_id", "sym_id"): + if state[optional_identity] is None: + state.pop(optional_identity) + for key, value in tuple(state.items()): + if isinstance(value, MolecularComplexMetadata): + state[key] = asdict(value) + elif isinstance(value, np.ndarray): + if value.dtype == np.int64: + value = value.astype(np.int32) + elif value.dtype in (np.dtype(np.float64), np.dtype(np.float32)): + value = value.astype(np.float16) + state[key] = value.tolist() + return state + + def to_blob(self) -> bytes: + return brotli.compress(msgpack.dumps(self.state_dict()), quality=5) + + @classmethod + def from_state_dict(cls, dct: dict[str, Any]) -> MolecularComplex: + dct = dict(dct) + for key, value in tuple(dct.items()): + if isinstance(value, list) and key in _SERIALIZED_ARRAYS: + dct[key] = np.asarray(value) + for key, value in tuple(dct.items()): + if not isinstance(value, np.ndarray): + continue + if key in {"atom_positions", "plddt"}: + dct[key] = value.astype(np.float32) + elif key == "token_to_atoms": + dct[key] = value.astype(np.int32) + elif key in {"chain_id", "entity_id", "sym_id"}: + dct[key] = value.astype(np.int64) + dct["metadata"] = MolecularComplexMetadata(**dct["metadata"]) + if "chain_id" not in dct: + dct["chain_id"] = np.zeros(len(dct["sequence"]), dtype=np.int64) + return cls(**dct) + + @classmethod + def from_blob(cls, input: Path | str | io.BytesIO | bytes) -> MolecularComplex: + if isinstance(input, (Path, str)): + payload = Path(input).read_bytes() + elif isinstance(input, io.BytesIO): + payload = input.getvalue() + else: + payload = input + state = msgpack.loads(brotli.decompress(payload), strict_map_key=False) + return cls.from_state_dict(state) diff --git a/src/fastplms/models/esmfold2/esmfold2_msa.py b/src/fastplms/models/esmfold2/esmfold2_msa.py new file mode 100644 index 0000000..56c6255 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_msa.py @@ -0,0 +1,577 @@ +"""Multiple-sequence-alignment value objects and lossless encodings.""" + +from __future__ import annotations + +import dataclasses +import string +from collections.abc import Sequence +from dataclasses import dataclass +from functools import cached_property +from itertools import islice +from typing import Any + +import numpy as np +from Bio import SeqIO +from scipy.spatial.distance import cdist + +from .esmfold2_misc import slice_any_object +from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter +from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences +from .esmfold2_sequential_dataclass import SequentialDataclass +from .esmfold2_system import PathOrBuffer + +_A3M_INSERTION_DELETE_TABLE = str.maketrans( + dict.fromkeys(string.ascii_lowercase + ".") +) +_SERIALIZATION_VERSION = 1 +_UINT32_BYTES = 4 + + +def is_a3m_insertion(character: str) -> bool: + """Return whether a character is an A3M insertion marker.""" + + return character == "." or character.islower() + + +def remove_insertions_from_sequence(sequence: str) -> str: + """Remove lowercase residues and dot insertion markers from an A3M row.""" + + return sequence.translate(_A3M_INSERTION_DELETE_TABLE) + + +def a3m_deletion_counts(sequence: str) -> np.ndarray: + """Count insertions preceding each A3M match column.""" + + codes = np.frombuffer(sequence.encode("ascii"), dtype=np.uint8) + lowercase = (codes >= ord("a")) & (codes <= ord("z")) + insertion_mask = lowercase | (codes == ord(".")) + prefix_counts = np.concatenate(([0], np.cumsum(insertion_mask))) + match_positions = np.flatnonzero(~insertion_mask) + return np.diff(prefix_counts[match_positions], prepend=0) + + +def _parse_full_payload(data: bytes) -> tuple[np.ndarray, list[str]]: + version = int.from_bytes(data[:1], "little") + if version != _SERIALIZATION_VERSION: + raise ValueError(f"Unsupported version: {version}") + seqlen = int.from_bytes(data[1:5], "little") + depth = int.from_bytes(data[5:9], "little") + body = data[9:] + split = seqlen * depth + array = np.frombuffer(body[:split], dtype="|S1").reshape(depth, seqlen) + headers = [header for header in body[split:].decode().split("\n") if header] + if not headers and depth > 0: + headers = [""] * depth + return array, headers + + +def _parse_sequence_payload(data: bytes) -> np.ndarray: + seqlen = int.from_bytes(data[:_UINT32_BYTES], "little") + return np.frombuffer(data[_UINT32_BYTES:], dtype="|S1").reshape(-1, seqlen) + + +def _full_payload(array: np.ndarray, headers: Sequence[str]) -> bytes: + depth, seqlen = array.shape + prefix = b"".join( + ( + _SERIALIZATION_VERSION.to_bytes(1, "little"), + seqlen.to_bytes(_UINT32_BYTES, "little"), + depth.to_bytes(_UINT32_BYTES, "little"), + ) + ) + return prefix + array.tobytes() + "\n".join(headers).encode() + + +def _sequence_payload(array: np.ndarray) -> bytes: + return array.shape[1].to_bytes(_UINT32_BYTES, "little") + array.tobytes() + + +def _random_row_indices(depth: int, count: int) -> np.ndarray: + sampled = np.random.choice(depth - 1, count - 1, replace=False) + 1 + return np.sort(np.append(0, sampled)) + + +@dataclass(frozen=True) +class FastMSA(SequentialDataclass): + """An MSA stored as a two-dimensional NumPy byte array.""" + + array: np.ndarray + headers: list[str] | None = None + + def __post_init__(self) -> None: + if not isinstance(self.array, np.ndarray): + raise TypeError("FastMSA array must be a NumPy array.") + if self.array.ndim != 2 or self.array.shape[0] == 0 or self.array.shape[1] == 0: + raise ValueError( + f"FastMSA array must have non-empty shape (depth, length), got {self.array.shape}." + ) + if self.headers is not None and len(self.headers) != self.depth: + raise ValueError("Number of headers must match depth.") + + @property + def depth(self) -> int: + return self.array.shape[0] + + @property + def seqlen(self) -> int: + return self.array.shape[1] + + def __len__(self) -> int: + return self.seqlen + + @classmethod + def from_bytes(cls, data: bytes) -> FastMSA: + array, headers = _parse_full_payload(data) + return cls(array, headers) + + @classmethod + def from_sequence_bytes(cls, data: bytes) -> FastMSA: + return cls(_parse_sequence_payload(data)) + + def __getitem__( + self, + indices: int | list[int] | slice | np.ndarray, + ) -> FastMSA: + column_indices = [indices] if isinstance(indices, int) else indices + return dataclasses.replace(self, array=self.array[:, column_indices]) + + def select_sequences( + self, + indices: Sequence[int] | np.ndarray, + ) -> FastMSA: + headers = None + if self.headers is not None: + headers = [self.headers[index] for index in indices] + return dataclasses.replace( + self, + array=self.array[indices], + headers=headers, + ) + + def select_random_sequences(self, num_seqs: int) -> FastMSA: + if num_seqs >= self.depth: + return self + return self.select_sequences(_random_row_indices(self.depth, num_seqs)) + + def pad_to_depth(self, depth: int) -> FastMSA: + if depth < self.depth: + raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") + if depth == self.depth: + return self + row_count = depth - self.depth + pad_value = ord("-") if self.array.dtype == np.uint8 else b"-" + array = np.pad( + self.array, + ((0, row_count), (0, 0)), + constant_values=pad_value, + ) + headers = None if self.headers is None else self.headers + [""] * row_count + return dataclasses.replace(self, array=array, headers=headers) + + @classmethod + def concat( + cls, + msas: Sequence[FastMSA], + join_token: str | None = None, + allow_depth_mismatch: bool = False, + ) -> FastMSA: + if not msas: + raise ValueError("Cannot concatenate an empty list of MSAs") + if join_token not in (None, ""): + raise NotImplementedError("join_token is not supported for FastMSA") + depths = [msa.depth for msa in msas] + if len(set(depths)) != 1: + if not allow_depth_mismatch: + raise ValueError("Depth mismatch in concatenating MSAs") + maximum_depth = max(depths) + msas = [msa.pad_to_depth(maximum_depth) for msa in msas] + header_columns = ( + msa.headers if msa.headers is not None else [""] * msa.depth for msa in msas + ) + headers = [ + "|".join(str(header) for header in row) for row in zip(*header_columns, strict=False) + ] + return cls( + np.concatenate([msa.array for msa in msas], axis=1), + headers, + ) + + @classmethod + def stack( + cls, + msas: Sequence[FastMSA], + remove_query_from_later_msas: bool = True, + ) -> FastMSA: + if not msas: + raise ValueError("Cannot stack an empty list of MSAs") + arrays: list[np.ndarray] = [] + headers: list[str] | None = [] if any(msa.headers is not None for msa in msas) else None + for index, msa in enumerate(msas): + start = 1 if index > 0 and remove_query_from_later_msas else 0 + arrays.append(msa.array[start:]) + if headers is not None: + source_headers = msa.headers or [""] * msa.depth + headers.extend(source_headers[start:]) + return cls(np.concatenate(arrays, axis=0), headers) + + def to_msa(self) -> MSA: + headers = self.headers + if headers is None: + headers = [f"seq{index}" for index in range(self.depth)] + entries = [ + FastaEntry(header, b"".join(row).decode()) + for header, row in zip(headers, self.array, strict=False) + ] + return MSA(entries) + + +@dataclass(frozen=True) +class MSA(SequentialDataclass): + """An ordered set of aligned protein sequences and optional A3M metadata.""" + + entries: list[FastaEntry] + deletions: np.ndarray | None = dataclasses.field(default=None, compare=False) + + def __post_init__(self) -> None: + if not isinstance(self.entries, list): + raise TypeError("MSA entries must be a list of FastaEntry rows.") + if not self.entries: + raise ValueError("MSA requires at least one aligned sequence.") + if any(not isinstance(entry, FastaEntry) for entry in self.entries): + raise TypeError("Every MSA entry must be a FastaEntry.") + expected_length = len(self.entries[0].sequence) + if expected_length == 0: + raise ValueError("MSA sequences must be non-empty.") + for row, entry in enumerate(self.entries[1:], start=1): + if len(entry.sequence) != expected_length: + raise ValueError( + "MSA row length mismatch: " + f"row 0 has {expected_length} columns, row {row} has " + f"{len(entry.sequence)}." + ) + deletions = self.deletions + if deletions is not None and not isinstance(deletions, np.ndarray): + raise TypeError("MSA deletions must be a NumPy array when provided.") + if isinstance(deletions, np.ndarray) and deletions.shape != ( + len(self.entries), + expected_length, + ): + raise ValueError( + "MSA deletion matrix must have shape " + f"({len(self.entries)}, {expected_length}), got {deletions.shape}." + ) + + @cached_property + def sequences(self) -> list[str]: + return [entry.sequence for entry in self.entries] + + @cached_property + def headers(self) -> list[str]: + return [entry.header for entry in self.entries] + + @property + def depth(self) -> int: + return len(self.entries) + + @property + def seqlen(self) -> int: + return len(self.entries[0].sequence) + + @property + def query(self) -> str: + return self.entries[0].sequence + + @cached_property + def array(self) -> np.ndarray: + return np.array([list(sequence) for sequence in self.sequences], dtype="|S1") + + @cached_property + def seqid(self) -> np.ndarray: + byte_array = self.array.view(np.uint8) + return (1 - cdist(byte_array[0][None], byte_array, "hamming"))[0] + + def __len__(self) -> int: + return self.seqlen + + def __repr__(self) -> str: + return f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})" + + @classmethod + def from_a3m( + cls, + path: PathOrBuffer, + remove_insertions: bool = True, + max_sequences: int | None = None, + ) -> MSA: + entries = [] + deletion_rows = [] + for header, raw_sequence in islice(read_sequences(path), max_sequences): + if remove_insertions: + deletion_rows.append(a3m_deletion_counts(raw_sequence)) + sequence = ( + remove_insertions_from_sequence(raw_sequence) if remove_insertions else raw_sequence + ) + if entries: + expected_length = len(entries[0].sequence) + if len(sequence) != expected_length: + raise ValueError( + "Sequence length mismatch. " + f"Expected: {expected_length}, Received: {len(sequence)}" + ) + entries.append(FastaEntry(header, sequence)) + deletions = None + if remove_insertions and deletion_rows: + deletions = np.stack(deletion_rows).astype(np.float32) + return cls(entries, deletions=deletions) + + @classmethod + def from_stockholm( + cls, + path: PathOrBuffer, + remove_insertions: bool = True, + max_sequences: int | None = None, + ) -> MSA: + entries = [] + for record in islice(SeqIO.parse(path, "stockholm"), max_sequences): + sequence = str(record.seq) + if entries: + expected_length = len(entries[0].sequence) + if len(sequence) != expected_length: + raise ValueError( + "Sequence length mismatch. " + f"Expected: {expected_length}, Received: {len(sequence)}" + ) + entries.append(FastaEntry(f"{record.id} {record.description}", sequence)) + msa = cls(entries) + if remove_insertions: + msa = msa.select_positions( + [index for index, residue in enumerate(msa.query) if residue != "-"] + ) + return msa + + @classmethod + def from_sequences( + cls, + sequences: list[str], + remove_insertions: bool = False, + ) -> MSA: + transform = ( + remove_insertions_from_sequence if remove_insertions else lambda sequence: sequence + ) + return cls([FastaEntry("", transform(sequence)) for sequence in sequences]) + + @classmethod + def from_bytes(cls, data: bytes) -> MSA: + array, headers = _parse_full_payload(data) + return cls( + [ + FastaEntry(header, b"".join(row).decode()) + for header, row in zip(headers, array, strict=False) + ] + ) + + @classmethod + def from_sequence_bytes(cls, data: bytes) -> MSA: + array = _parse_sequence_payload(data) + return cls([FastaEntry("", b"".join(row).decode()) for row in array]) + + @classmethod + def from_state_dict(cls, dct: dict[str, Any]) -> MSA: + deletions = dct.get("deletions") + return cls( + [FastaEntry("", sequence) for sequence in dct["sequences"]], + deletions=(None if deletions is None else np.asarray(deletions, dtype=np.float32)), + ) + + def to_a3m(self, path: PathOrBuffer) -> None: + write_sequences(self.entries, path) + + def to_fast_msa(self) -> FastMSA: + return FastMSA(self.array, self.headers) + + def to_bytes(self) -> bytes: + return _full_payload(self.array, self.headers) + + def to_sequence_bytes(self) -> bytes: + """Serialize aligned sequences without their headers.""" + + return _sequence_payload(self.array) + + def state_dict(self, json_serializable: bool = False) -> dict[str, Any]: + result: dict[str, Any] = {"sequences": self.sequences} + if self.deletions is not None: + result["deletions"] = self.deletions.tolist() if json_serializable else self.deletions + return result + + def _aligned_deletions(self) -> np.ndarray | None: + if self.deletions is None: + return None + if self.deletions.shape != (self.depth, self.seqlen): + return None + return self.deletions + + def _select_deletion_columns(self, indices) -> np.ndarray | None: + if self.deletions is None or self.deletions.shape[1] != self.seqlen: + return None + return self.deletions[:, indices] + + def select_sequences( + self, + indices: Sequence[int] | np.ndarray, + ) -> MSA: + deletions = None if self.deletions is None else self.deletions[np.asarray(indices)] + return dataclasses.replace( + self, + entries=[self.entries[index] for index in indices], + deletions=deletions, + ) + + def select_positions( + self, + indices: Sequence[int] | np.ndarray, + ) -> MSA: + entries = [ + FastaEntry( + entry.header, + "".join(entry.sequence[index] for index in indices), + ) + for entry in self.entries + ] + return dataclasses.replace( + self, + entries=entries, + deletions=self._select_deletion_columns(indices), + ) + + def __getitem__( + self, + indices: int | list[int] | slice | np.ndarray, + ) -> MSA: + column_indices = [indices] if isinstance(indices, int) else indices + entries = [ + FastaEntry( + entry.header, + slice_any_object(entry.sequence, column_indices), + ) + for entry in self.entries + ] + return dataclasses.replace( + self, + entries=entries, + deletions=self._select_deletion_columns(column_indices), + ) + + def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA: + if mode not in ("max", "min"): + raise ValueError(f"Unsupported MSA selection mode: {mode!r}.") + if self.depth <= num_seqs: + return self + return self.select_sequences(greedy_select_indices(self.array, num_seqs, mode)) + + def hhfilter( + self, + seqid: int = 90, + diff: int = 0, + cov: int = 0, + qid: int = 0, + qsc: float = -20.0, + binary: str = "hhfilter", + ) -> MSA: + indices = hhfilter( + self.sequences, + seqid=seqid, + diff=diff, + cov=cov, + qid=qid, + qsc=qsc, + binary=binary, + ) + return self.select_sequences(indices) + + def select_random_sequences(self, num_seqs: int) -> MSA: + if num_seqs >= self.depth: + return self + return self.select_sequences(_random_row_indices(self.depth, num_seqs)) + + def select_diverse_sequences(self, num_seqs: int) -> MSA: + if num_seqs >= self.depth: + return self + filtered = self.hhfilter(diff=num_seqs) + if num_seqs < filtered.depth: + filtered = filtered.select_random_sequences(num_seqs) + return filtered + + def pad_to_depth(self, depth: int) -> MSA: + if depth < self.depth: + raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") + if depth == self.depth: + return self + count = depth - self.depth + extra = [FastaEntry("", "-" * self.seqlen) for _ in range(count)] + deletions = self._aligned_deletions() + if deletions is not None: + zero_rows = np.zeros((count, self.seqlen), dtype=deletions.dtype) + deletions = np.concatenate((deletions, zero_rows), axis=0) + return dataclasses.replace( + self, + entries=self.entries + extra, + deletions=deletions, + ) + + @classmethod + def stack( + cls, + msas: Sequence[MSA], + remove_query_from_later_msas: bool = True, + ) -> MSA: + entries = [] + deletion_arrays = [] + for index, msa in enumerate(msas): + start = 1 if index > 0 and remove_query_from_later_msas else 0 + entries.extend(msa.entries[start:]) + aligned = msa._aligned_deletions() + if aligned is not None: + deletion_arrays.append(aligned[start:]) + deletions = None + if ( + len(deletion_arrays) == len(msas) + and len({array.shape[1] for array in deletion_arrays}) == 1 + ): + deletions = np.concatenate(deletion_arrays, axis=0) + return cls(entries=entries, deletions=deletions) + + @classmethod + def concat( + cls, + msas: Sequence[MSA], + join_token: str | None = "|", + allow_depth_mismatch: bool = False, + ) -> MSA: + if not msas: + raise ValueError("Cannot concatenate an empty list of MSAs") + depths = [msa.depth for msa in msas] + if len(set(depths)) != 1: + if not allow_depth_mismatch: + raise ValueError("Depth mismatch in concatenating MSAs") + maximum_depth = max(depths) + msas = [msa.pad_to_depth(maximum_depth) for msa in msas] + headers = [ + "|".join(str(header) for header in row) + for row in zip(*(msa.headers for msa in msas), strict=False) + ] + separator = "" if join_token is None else join_token + sequences = [ + separator.join(row) for row in zip(*(msa.sequences for msa in msas), strict=False) + ] + deletions = None + if separator == "": + arrays = [msa._aligned_deletions() for msa in msas] + if all(array is not None for array in arrays): + deletions = np.concatenate(arrays, axis=1) # type: ignore[arg-type] + return cls( + [ + FastaEntry(header, sequence) + for header, sequence in zip(headers, sequences, strict=False) + ], + deletions=deletions, + ) diff --git a/src/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py b/src/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py new file mode 100644 index 0000000..bc13db7 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py @@ -0,0 +1,111 @@ +"""Sequence selection for multiple-sequence alignments.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +import numpy as np + +from .esmfold2_system import run_subprocess_with_errorcheck + + +def _byte_matrix(array: np.ndarray) -> np.ndarray: + """Return a two-dimensional byte view used for Hamming comparisons.""" + + matrix = np.asarray(array).view(np.uint8) + return matrix.reshape(matrix.shape[0], -1) + + +def _hamming_to_all(query: np.ndarray, sequences: np.ndarray) -> np.ndarray: + return np.not_equal(sequences, query).mean(axis=1, dtype=np.float64) + + +def greedy_select_indices(array: np.ndarray, num_seqs: int, mode: str = "max") -> list[int]: + """Select MSA rows by greedy mean Hamming distance from the query row. + + Row zero is always retained. At each step the selector chooses the remaining + row with greatest distance for ``mode="max"`` or least distance for + ``mode="min"``. Returned indices follow source order. + """ + + if not isinstance(array, np.ndarray): + raise TypeError("array must be a NumPy array") + if array.ndim != 2 or array.shape[0] == 0 or array.shape[1] == 0: + raise ValueError( + f"array must have non-empty shape (depth, length), got {array.shape}" + ) + if isinstance(num_seqs, bool) or not isinstance(num_seqs, int): + raise TypeError("num_seqs must be an integer") + if num_seqs <= 0: + raise ValueError("num_seqs must be greater than zero") + if not isinstance(mode, str): + raise TypeError("mode must be a string") + if mode not in {"max", "min"}: + raise ValueError(f"unsupported selection mode: {mode}") + depth = array.shape[0] + if depth <= num_seqs: + return list(range(depth)) + + sequences = _byte_matrix(array) + selected = [0] + available = np.ones(depth, dtype=bool) + available[0] = False + distance_sum = _hamming_to_all(sequences[0], sequences) + choose = np.argmax if mode == "max" else np.argmin + + while len(selected) < num_seqs: + candidates = np.flatnonzero(available) + candidate_scores = distance_sum[candidates] / len(selected) + next_index = int(candidates[int(choose(candidate_scores))]) + selected.append(next_index) + available[next_index] = False + distance_sum += _hamming_to_all(sequences[next_index], sequences) + return sorted(selected) + + +def _temporary_root() -> str | None: + shared_memory = Path("/dev/shm") + return os.fspath(shared_memory) if shared_memory.is_dir() else None + + +def hhfilter( + sequences: list[str], + seqid: int = 90, + diff: int = 0, + cov: int = 0, + qid: int = 0, + qsc: float = -20.0, + binary: str = "hhfilter", +) -> list[int]: + """Run HH-suite filtering and return source indices from its FASTA headers.""" + + with tempfile.TemporaryDirectory(dir=_temporary_root()) as directory: + work = Path(directory) + source_path = work / "input.fasta" + result_path = work / "output.fasta" + records = (f">{index}\n{sequence}" for index, sequence in enumerate(sequences)) + source_path.write_text("\n".join(records), encoding="utf-8") + command = [ + binary, + "-i", + os.fspath(source_path), + "-M", + "a3m", + "-o", + os.fspath(result_path), + "-id", + str(seqid), + "-diff", + str(diff), + "-cov", + str(cov), + "-qid", + str(qid), + "-qsc", + str(qsc), + ] + run_subprocess_with_errorcheck(command, capture_output=True) + headers = result_path.read_text(encoding="utf-8").splitlines() + return [int(line[1:].strip()) for line in headers if line.startswith(">")] diff --git a/src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py b/src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py new file mode 100644 index 0000000..45e0652 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py @@ -0,0 +1,67 @@ +"""Rigid-frame normalization for atom37 coordinates.""" + +from __future__ import annotations + +from typing import TypeVar + +import numpy as np +import torch +from torch import Tensor + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_affine3d import Affine3D + +ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) + + +def atom3_to_backbone_frames(bb_positions: Tensor) -> Affine3D: + """Construct a frame from N, C-alpha, and C positions in ``X``.""" + + n_position, ca_position, c_position = bb_positions.unbind(dim=-2) + return Affine3D.from_graham_schmidt(c_position, ca_position, n_position) + + +def index_by_atom_name( + atom37: ArrayOrTensor, + atom_names: str | list[str], + dim: int = -2, +) -> ArrayOrTensor: + """Select one or more named atoms along an atom37 axis.""" + + single_atom = isinstance(atom_names, str) + names = [atom_names] if single_atom else atom_names + indices = [residue_constants.atom_order[name] for name in names] + axis = dim % atom37.ndim + if isinstance(atom37, Tensor): + index = torch.tensor(indices, dtype=torch.long, device=atom37.device) + selected = torch.index_select(atom37, axis, index) + else: + selected = np.take(atom37, indices, axis=axis) + return selected.squeeze(axis) if single_atom else selected # type: ignore[return-value] + + +def get_protein_normalization_frame(coords: Tensor) -> Affine3D: + """Build one frame from backbone coordinates ``X`` with shape (l, 37, 3).""" + + backbone = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2) + residue_is_valid = torch.isfinite(backbone).all(dim=-1).all(dim=-1) + weights = residue_is_valid[..., None, None] + coordinate_sum = backbone.masked_fill(~weights, 0).sum(dim=-3) + count = residue_is_valid.sum(dim=-1)[..., None, None] + mean_backbone = coordinate_sum / (count + 1e-8) + return atom3_to_backbone_frames(mean_backbone.float()) + + +def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor: + """Express atom coordinates ``X`` in the inverse of ``frame``.""" + + transformed = frame[..., None, None].invert().apply(coords) + frame_is_valid = frame.trans.norm(dim=-1) > 0 + normalized = torch.where(frame_is_valid[..., None, None, None], transformed, coords) + return normalized.masked_fill(torch.isinf(coords), torch.inf) + + +def normalize_coordinates(coords: Tensor) -> Tensor: + """Normalize ``X`` with shape (..., l, 37, 3) to its backbone frame.""" + + return apply_frame_to_coords(coords, get_protein_normalization_frame(coords)) diff --git a/src/fastplms/models/esmfold2/esmfold2_output.py b/src/fastplms/models/esmfold2/esmfold2_output.py new file mode 100644 index 0000000..97c846a --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_output.py @@ -0,0 +1,201 @@ +"""Convert ESMFold2 coordinate tensors into molecular-complex records.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, field +from itertools import groupby +from typing import Any + +import numpy as np +import torch + +from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER +from .esmfold2_molecular_complex import MolecularComplex, MolecularComplexMetadata + + +def get_element_symbol(atomic_number: int) -> str: + """Map a training-time atomic number to a chemical symbol.""" + + return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_number, "X") + + +def _decode_atom_name(encoded_name: Any) -> str: + values = encoded_name.tolist() if hasattr(encoded_name, "tolist") else encoded_name + return "".join(chr(int(value) + 32) for value in values if int(value)).strip() + + +@dataclass +class _ComplexRecords: + sequence: list[str] = field(default_factory=list) + chain_ids: list[int] = field(default_factory=list) + token_to_atoms: list[list[int]] = field(default_factory=list) + confidence: list[float] = field(default_factory=list) + positions: list[list[float]] = field(default_factory=list) + elements: list[str] = field(default_factory=list) + atom_names: list[str] = field(default_factory=list) + atom_hetero: list[bool] = field(default_factory=list) + chain_lookup: dict[int, str] = field(default_factory=dict) + entity_lookup: dict[int, str] = field(default_factory=dict) + + def add_token( + self, + *, + residue_name: str, + asym_id: int, + plddt: float, + atoms: Iterable[tuple[list[float], str, str]], + hetero: bool, + ) -> None: + atom_start = len(self.positions) + for position, element, atom_name in atoms: + self.positions.append(position) + self.elements.append(element) + self.atom_names.append(atom_name) + self.atom_hetero.append(hetero) + self.sequence.append(residue_name) + self.chain_ids.append(asym_id) + self.confidence.append(plddt) + self.token_to_atoms.append([atom_start, len(self.positions)]) + + def build(self, complex_id: str) -> MolecularComplex: + return MolecularComplex( + id=complex_id, + sequence=self.sequence, + atom_positions=np.asarray(self.positions, dtype=np.float32).reshape(-1, 3), + atom_elements=np.asarray(self.elements, dtype=object), + token_to_atoms=np.asarray(self.token_to_atoms, dtype=np.int32).reshape(-1, 2), + chain_id=np.asarray(self.chain_ids, dtype=np.int64), + plddt=np.asarray(self.confidence, dtype=np.float32), + atom_names=np.asarray(self.atom_names, dtype=object), + atom_hetero=np.asarray(self.atom_hetero, dtype=bool), + metadata=MolecularComplexMetadata( + entity_lookup=self.entity_lookup, + chain_lookup=self.chain_lookup, + assembly_composition=None, + ), + ) + + +def build_molecular_complex_from_features( + coords: torch.Tensor, + plddt: torch.Tensor, + atom_mask: torch.Tensor, + ref_element: torch.Tensor, + ref_atom_name_chars: torch.Tensor, + chain_infos: list[Any], + complex_id: str, +) -> MolecularComplex: + """Decode model features into one complex without intermediate structure files. + + Protein, DNA, and RNA tokens are grouped by residue index. Ligand atom + tokens are collapsed into one non-polymer residue per chain. + """ + + M = atom_mask.bool().cpu().numpy() + X = coords.float().cpu().numpy() + atom_names = ref_atom_name_chars.cpu().numpy() + elements = ref_element.cpu().numpy() + confidence = plddt.float().cpu().numpy() + records = _ComplexRecords() + + def decode_atoms(tokens: Iterable[Any]): + for token in tokens: + for atom_index in range(token.atom_start, token.atom_start + token.atom_count): + if M[atom_index]: + yield ( + X[atom_index].tolist(), + get_element_symbol(int(elements[atom_index])), + _decode_atom_name(atom_names[atom_index]), + ) + + for chain in chain_infos: + is_nonpolymer = chain.mol_type == MOL_TYPE_NONPOLYMER + records.chain_lookup[chain.asym_id] = chain.chain_id + records.entity_lookup[chain.entity_id] = "non-polymer" if is_nonpolymer else "polymer" + + if is_nonpolymer: + mean_confidence = ( + float(np.mean([confidence[token.token_index] for token in chain.tokens])) + if chain.tokens + else 0.0 + ) + records.add_token( + residue_name=chain.tokens[0].residue_name if chain.tokens else "LIG", + asym_id=chain.asym_id, + plddt=mean_confidence, + atoms=decode_atoms(chain.tokens), + hetero=True, + ) + continue + + residue_groups = groupby(chain.tokens, key=lambda token: token.residue_index) + for _residue_index, group in residue_groups: + residue_tokens = list(group) + records.add_token( + residue_name=residue_tokens[0].residue_name, + asym_id=chain.asym_id, + plddt=float(np.mean([confidence[token.token_index] for token in residue_tokens])), + atoms=decode_atoms(residue_tokens), + hetero=False, + ) + + return records.build(complex_id) + + +def build_molecular_complex( + structure: Any, + coords: torch.Tensor, + plddt: torch.Tensor, + complex_id: str, +) -> MolecularComplex: + """Decode coordinates using the atom and residue arrays of a prepared structure.""" + + records = _ComplexRecords() + coordinate_index = 0 + confidence_index = 0 + + for chain in structure.chains: + asym_id = int(chain["asym_id"]) + mol_type = int(chain["mol_type"]) + is_nonpolymer = mol_type == MOL_TYPE_NONPOLYMER + records.chain_lookup[asym_id] = str(chain["name"]) + records.entity_lookup[int(chain["entity_id"])] = ( + "non-polymer" if is_nonpolymer else "polymer" + ) + + residue_start = int(chain["res_idx"]) + residue_stop = residue_start + int(chain["res_num"]) + for residue in structure.residues[residue_start:residue_stop]: + atom_start = int(residue["atom_idx"]) + atom_stop = atom_start + int(residue["atom_num"]) + decoded_atoms: list[tuple[list[float], str, str]] = [] + for atom in structure.atoms[atom_start:atom_stop]: + if not atom["is_present"]: + continue + decoded_atoms.append( + ( + coords[coordinate_index].tolist(), + get_element_symbol(int(atom["element"].item())), + _decode_atom_name(atom["name"]), + ) + ) + coordinate_index += 1 + + records.add_token( + residue_name=str(residue["name"]), + asym_id=asym_id, + plddt=float(plddt[confidence_index].item()), + atoms=decoded_atoms, + hetero=is_nonpolymer, + ) + confidence_index += 1 + + return records.build(complex_id) + + +__all__ = [ + "build_molecular_complex", + "build_molecular_complex_from_features", + "get_element_symbol", +] diff --git a/src/fastplms/models/esmfold2/esmfold2_paired_msa.py b/src/fastplms/models/esmfold2/esmfold2_paired_msa.py new file mode 100644 index 0000000..44f050b --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_paired_msa.py @@ -0,0 +1,282 @@ +"""Construct taxonomy-paired MSA features for multichain folding.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import numpy as np + +from .esmfold2_constants import ( + MSA_GAP_TOKEN_ID, + PROTEIN_3TO1, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, +) +from .esmfold2_msa import MSA + +_TAXONOMY_PATTERN = re.compile(r"key=(-?\d+)") + + +def protein_letter_to_res_type() -> dict[str, int]: + """Return the one-letter residue vocabulary used by the MSA encoder.""" + + vocabulary = { + one_letter: PROTEIN_RESIDUE_TO_RES_TYPE[three_letter] + for three_letter, one_letter in PROTEIN_3TO1.items() + if three_letter in PROTEIN_RESIDUE_TO_RES_TYPE + } + vocabulary.update({"-": MSA_GAP_TOKEN_ID, "X": PROTEIN_UNK_RES_TYPE}) + return vocabulary + + +def _taxonomy_from_header(header: str) -> int: + match = _TAXONOMY_PATTERN.search(header) if header else None + return int(match.group(1)) if match is not None else -1 + + +def _emitted_length(sequence: str) -> int: + return sum(character != "." and not character.islower() for character in sequence) + + +def _decode_a3m_row( + sequence: str, + sequence_length: int, + vocabulary: dict[str, int], +) -> tuple[np.ndarray, np.ndarray]: + residues = np.full(sequence_length, MSA_GAP_TOKEN_ID, dtype=np.int64) + deletions = np.zeros(sequence_length, dtype=np.float32) + column = 0 + insertion_count = 0 + for character in sequence: + if character == "." or character.islower(): + insertion_count += 1 + continue + if column == sequence_length: + break + residues[column] = ( + MSA_GAP_TOKEN_ID + if character == "-" + else vocabulary.get(character.upper(), PROTEIN_UNK_RES_TYPE) + ) + if insertion_count: + deletions[column] = float(insertion_count) + insertion_count = 0 + column += 1 + return residues, deletions + + +def msa_to_res_type_and_deletions( + msa: MSA, + letter_to_res_type: dict[str, int], +) -> tuple[np.ndarray, np.ndarray]: + """Decode an A3M alignment into arrays ``X`` and ``D`` with shape (m, l).""" + + sequence_length = _emitted_length(msa.entries[0].sequence) + residue_rows: list[np.ndarray] = [] + deletion_rows: list[np.ndarray] = [] + for entry in msa.entries: + residues, deletions = _decode_a3m_row( + entry.sequence, + sequence_length, + letter_to_res_type, + ) + residue_rows.append(residues) + deletion_rows.append(deletions) + return np.stack(residue_rows), np.stack(deletion_rows) + + +@dataclass(frozen=True) +class _ChainAlignment: + residues: np.ndarray + deletions: np.ndarray + taxonomies: list[int] + + +def _chain_alignment( + msa: MSA | None, + query_res_types: np.ndarray, + vocabulary: dict[str, int], +) -> _ChainAlignment: + if msa is None or msa.depth == 0: + return _ChainAlignment( + residues=query_res_types[None, :], + deletions=np.zeros((1, query_res_types.shape[0]), dtype=np.float32), + taxonomies=[-1], + ) + residues, deletions = msa_to_res_type_and_deletions(msa, vocabulary) + taxonomies = [_taxonomy_from_header(entry.header) for entry in msa.entries] + return _ChainAlignment(residues, deletions, taxonomies) + + +def _taxonomy_groups( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], +) -> dict[int, list[tuple[int, int]]]: + groups: dict[int, list[tuple[int, int]]] = {} + for chain_id in chain_ids: + for row, taxonomy in enumerate(alignments[chain_id].taxonomies): + if row and taxonomy != -1: + groups.setdefault(taxonomy, []).append((chain_id, row)) + return {taxonomy: rows for taxonomy, rows in groups.items() if len(rows) > 1} + + +def _available_rows( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + groups: dict[int, list[tuple[int, int]]], +) -> dict[int, list[int]]: + used = {row for group in groups.values() for row in group} + return { + chain_id: [ + row + for row in range(1, len(alignments[chain_id].taxonomies)) + if (chain_id, row) not in used + ] + for chain_id in chain_ids + } + + +def _append_taxonomy_rows( + rows: list[dict[int, int]], + paired_flags: list[dict[int, int]], + chain_ids: list[int], + groups: dict[int, list[tuple[int, int]]], + available: dict[int, list[int]], + max_pairs: int, +) -> None: + ordered_groups = sorted( + groups.values(), + key=lambda group: len({chain_id for chain_id, _row in group}), + reverse=True, + ) + for group in ordered_groups: + rows_by_chain: dict[int, list[int]] = {} + for chain_id, row in group: + rows_by_chain.setdefault(chain_id, []).append(row) + for occurrence in range(max(map(len, rows_by_chain.values()))): + selected: dict[int, int] = {} + flags: dict[int, int] = {} + for chain_id, candidates in rows_by_chain.items(): + selected[chain_id] = candidates[occurrence % len(candidates)] + flags[chain_id] = 1 + for chain_id in chain_ids: + if chain_id not in selected: + flags[chain_id] = 0 + selected[chain_id] = available[chain_id].pop(0) if available[chain_id] else -1 + rows.append(selected) + paired_flags.append(flags) + if len(rows) >= max_pairs: + break + if len(rows) >= max_pairs: + break + + +def _append_unpaired_rows( + rows: list[dict[int, int]], + paired_flags: list[dict[int, int]], + chain_ids: list[int], + available: dict[int, list[int]], + max_total: int, +) -> None: + max_remaining = max((len(indices) for indices in available.values()), default=0) + for _ in range(min(max_total - len(rows), max_remaining)): + rows.append( + { + chain_id: available[chain_id].pop(0) if available[chain_id] else -1 + for chain_id in chain_ids + } + ) + paired_flags.append({chain_id: 0 for chain_id in chain_ids}) + if len(rows) >= max_total: + break + + +def _pairing_plan( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + max_pairs: int, + max_total: int, + max_seqs: int, +) -> tuple[list[dict[int, int]], list[dict[int, int]]]: + groups = _taxonomy_groups(chain_ids, alignments) + available = _available_rows(chain_ids, alignments, groups) + rows = [{chain_id: 0 for chain_id in chain_ids}] + flags = [{chain_id: 1 for chain_id in chain_ids}] + _append_taxonomy_rows(rows, flags, chain_ids, groups, available, max_pairs) + _append_unpaired_rows(rows, flags, chain_ids, available, max_total) + return rows[:max_seqs], flags[:max_seqs] + + +def _project_alignment_rows( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + rows: list[dict[int, int]], + flags: list[dict[int, int]], + token_asym_ids: np.ndarray, + token_res_ids: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + m, t = len(rows), len(token_asym_ids) + residues = np.full((m, t), MSA_GAP_TOKEN_ID, dtype=np.int64) + deletions = np.zeros((m, t), dtype=np.float32) + paired_mask = np.zeros((m, t), dtype=np.float32) + for chain_id in chain_ids: + alignment = alignments[chain_id] + selected_rows = np.asarray([row[chain_id] for row in rows], dtype=np.int64) + chain_flags = np.asarray([row[chain_id] for row in flags], dtype=np.float32) + token_mask = token_asym_ids == chain_id + if not token_mask.any(): + continue + columns = np.minimum(token_res_ids[token_mask], alignment.residues.shape[1] - 1) + valid_rows = selected_rows >= 0 + if valid_rows.any(): + output_rows = np.flatnonzero(valid_rows) + output_columns = np.flatnonzero(token_mask) + residues[np.ix_(output_rows, output_columns)] = alignment.residues[ + selected_rows[valid_rows] + ][:, columns] + deletions[np.ix_(output_rows, output_columns)] = alignment.deletions[ + selected_rows[valid_rows] + ][:, columns] + paired_mask[:, token_mask] = chain_flags[:, None] + return residues, deletions, paired_mask + + +def construct_paired_msa( + chain_msas: dict[int, MSA | None], + chain_query_res_types: dict[int, np.ndarray], + token_asym_ids: np.ndarray, + token_res_ids: np.ndarray, + letter_to_res_type: dict[str, int] | None = None, + *, + max_pairs: int = 8192, + max_total: int = 16384, + max_seqs: int = 16384, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return residue, deletion, and pairing arrays with shape (m, t).""" + + vocabulary = protein_letter_to_res_type() if letter_to_res_type is None else letter_to_res_type + chain_ids = sorted(chain_msas) + alignments = { + chain_id: _chain_alignment( + chain_msas[chain_id], + chain_query_res_types[chain_id], + vocabulary, + ) + for chain_id in chain_ids + } + rows, flags = _pairing_plan( + chain_ids, + alignments, + max_pairs, + max_total, + max_seqs, + ) + return _project_alignment_rows( + chain_ids, + alignments, + rows, + flags, + token_asym_ids, + token_res_ids, + ) diff --git a/src/fastplms/models/esmfold2/esmfold2_parsing.py b/src/fastplms/models/esmfold2/esmfold2_parsing.py new file mode 100644 index 0000000..f906d1f --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_parsing.py @@ -0,0 +1,126 @@ +"""FASTA parsing and writing with explicit stream ownership.""" + +from __future__ import annotations + +import gzip +import io +from collections.abc import Generator, Iterable +from contextlib import nullcontext +from pathlib import Path +from typing import NamedTuple, TextIO + +from .esmfold2_utils_types import PathOrBuffer + + +class FastaEntry(NamedTuple): + """One FASTA record in source order.""" + + header: str + sequence: str + + +def parse_fasta(text: str) -> Generator[FastaEntry, None, None]: + """Yield records from FASTA text without normalizing sequence symbols.""" + + header: str | None = None + sequence_lines: list[str] = [] + found_record = False + + for line in text.splitlines(): + if not line or line.startswith("#"): + continue + if line.startswith(">"): + if header is not None: + found_record = True + yield FastaEntry(header, "".join(sequence_lines)) + header = line[1:].strip() + sequence_lines.clear() + elif header is not None: + sequence_lines.append(line) + + if header is not None: + found_record = True + yield FastaEntry(header, "".join(sequence_lines)) + if not found_record: + raise ValueError("Found no sequences in input") + + +def _open_reader(source: PathOrBuffer): + if isinstance(source, io.TextIOBase): + return nullcontext(source) + path = Path(source) + if path.suffix.lower() == ".gz": + return gzip.open(path, mode="rt", encoding="utf-8") + return path.open(mode="r", encoding="utf-8") + + +def read_sequences(source: PathOrBuffer) -> Generator[FastaEntry, None, None]: + """Read FASTA records while leaving caller-owned streams open.""" + + with _open_reader(source) as handle: + yield from parse_fasta(handle.read()) + + +def read_first_sequence(source: PathOrBuffer) -> FastaEntry: + """Return the first FASTA record from a path or text stream.""" + + return next(read_sequences(source)) + + +def count_fasta_sequences(path: str | Path) -> int: + """Count FASTA headers without parsing sequence bodies.""" + + source = Path(path) + if not source.exists(): + return 0 + with source.open(encoding="utf-8") as handle: + return sum(line.startswith(">") for line in handle) + + +def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None: + """Append one record, inserting a separator if the file lacks a final newline.""" + + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + needs_separator = ( + destination.exists() + and destination.stat().st_size > 0 + and destination.read_bytes()[-1:] != b"\n" + ) + with destination.open(mode="a", encoding="utf-8") as handle: + if needs_separator: + handle.write("\n") + handle.write(f">{header}\n{sequence}\n") + + +def _open_writer(destination: PathOrBuffer): + if isinstance(destination, io.TextIOBase): + return nullcontext(destination) + path = Path(destination) + path.parent.mkdir(parents=True, exist_ok=True) + return path.open(mode="w", encoding="utf-8") + + +def write_sequences(sequences: Iterable[tuple[str, str]], destination: PathOrBuffer) -> None: + """Write records with one blank-line-free separator between entries.""" + + with _open_writer(destination) as handle: + _write_records(handle, sequences) + + +def _write_records(handle: TextIO, sequences: Iterable[tuple[str, str]]) -> None: + for index, (header, sequence) in enumerate(sequences): + if index: + handle.write("\n") + handle.write(f">{header}\n{sequence}") + + +__all__ = [ + "FastaEntry", + "append_fasta_sequence", + "count_fasta_sequences", + "parse_fasta", + "read_first_sequence", + "read_sequences", + "write_sequences", +] diff --git a/src/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py b/src/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py new file mode 100644 index 0000000..d6ab3e7 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py @@ -0,0 +1,127 @@ +"""Predicted-aligned-error scores and training loss.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import Tensor + +from .esmfold2_affine3d import Affine3D + +_CPU_DEVICE = torch.device("cpu") + + +def _compute_pae_masks(mask: Tensor) -> Tensor: + residue_mask = mask.bool() + return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2) + + +def _pae_bins( + max_bin: float = 31, + num_bins: int = 64, + device: torch.device = _CPU_DEVICE, +) -> Tensor: + """Return the representative distance for each PAE probability bin.""" + + boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device) + width = max_bin / (num_bins - 2) + centers = boundaries + width / 2 + overflow_center = centers[-1:] + width + return torch.cat((centers, overflow_center)) + + +def _masked_probabilities(logits: Tensor, pair_mask: Tensor) -> Tensor: + masked_logits = logits.masked_fill( + ~pair_mask.unsqueeze(-1), + torch.finfo(logits.dtype).min, + ) + return masked_logits.softmax(dim=-1) + + +def masked_mean( + mask: Tensor, + value: Tensor, + dim: int | tuple[int, ...] | None = None, + eps: float = 1e-10, +) -> Tensor: + """Average values over true entries of a broadcast-compatible mask.""" + + weights = mask.expand_as(value) + weighted_sum = torch.sum(weights * value, dim=dim) + weight_sum = torch.sum(weights, dim=dim) + return weighted_sum / (weight_sum + eps) + + +def compute_predicted_aligned_error( + logits: Tensor, + aa_mask: Tensor, + sequence_id: Tensor | None = None, + max_bin: float = 31, +) -> Tensor: + """Convert PAE logits ``X`` with shape (..., l, l, n) to distances.""" + + del sequence_id + pair_mask = _compute_pae_masks(aa_mask) + probabilities = _masked_probabilities(logits, pair_mask) + centers = _pae_bins(max_bin, logits.shape[-1], logits.device) + return torch.sum(probabilities * centers, dim=-1) + + +@torch.no_grad() +def compute_tm(logits: Tensor, aa_mask: Tensor, max_bin: float = 31.0) -> Tensor: + """Estimate TM score from pairwise PAE logits.""" + + pair_mask = _compute_pae_masks(aa_mask) + sequence_lengths = aa_mask.sum(dim=-1, keepdim=True) + centers = _pae_bins(max_bin, logits.shape[-1], logits.device) + distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8 + tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2) + probabilities = _masked_probabilities(logits, pair_mask) + score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1) + score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1) + return score_per_anchor.max(dim=-1).values + + +def _local_coordinates(frames: Affine3D) -> Tensor: + origins = frames.trans[..., None, :, :] + return frames.invert()[..., None].apply(origins) + + +def tm_loss( + logits: Tensor, + pred_affine: Tensor, + targ_affine: Tensor, + targ_mask: Tensor, + tm_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + max_bin: float = 31, +) -> Tensor: + """Cross-entropy loss for discretized aligned-position errors.""" + + del sequence_id + predicted_frames = Affine3D.from_tensor(pred_affine) + target_frames = Affine3D.from_tensor(targ_affine) + with torch.no_grad(): + squared_error = ( + (_local_coordinates(predicted_frames) - _local_coordinates(target_frames)) + .square() + .sum(dim=-1) + ) + boundaries = torch.linspace( + 0, + max_bin, + logits.shape[-1] - 1, + device=logits.device, + ).square() + target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long() + + cross_entropy = F.cross_entropy( + logits.movedim(3, 1), + target_bins, + reduction="none", + ) + pair_mask = _compute_pae_masks(targ_mask) + loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2)) + if tm_mask is None: + return loss_per_sample.mean() + return masked_mean(tm_mask, loss_per_sample) diff --git a/src/fastplms/models/esmfold2/esmfold2_prepare_input.py b/src/fastplms/models/esmfold2/esmfold2_prepare_input.py new file mode 100644 index 0000000..870c457 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_prepare_input.py @@ -0,0 +1,1130 @@ +"""Translate typed sequence inputs into the tensors consumed by ESMFold2. + +The conversion has four explicit stages: entity and chain assignment, residue +tokenization, structural feature construction, and atom-table padding. Keeping +those stages separate makes the biological indexing rules testable without +loading model weights. +""" + +from __future__ import annotations + +import math +import warnings +from collections import defaultdict +from contextlib import suppress +from dataclasses import dataclass, field +from itertools import combinations +from typing import Any + +import numpy as np +import torch + +from .esmfold2_conformers import ( + get_ccd_leaving_atoms, + get_idealized_atom_pos, + get_ligand_ccd_atoms_with_charges, + get_ligand_ccd_bonds, + get_ligand_idealized_atom_pos, +) +from .esmfold2_constants import ( + CHARGED_ATOMS, + DNA_1TO3, + DNA_BACKBONE_ATOMS, + DNA_HEAVY_ATOMS, + DNA_RESIDUE_TO_RES_TYPE, + DNA_RNA_LIGAND_INPUT_ID, + DNA_UNK_RES_TYPE, + ELEMENT_TO_ATOMIC_NUM, + ESM_PROTEIN_VOCAB, + MOL_TYPE_DNA, + MOL_TYPE_NONPOLYMER, + MOL_TYPE_PROTEIN, + MOL_TYPE_RNA, + MSA_GAP_TOKEN_ID, + PROTEIN_1TO3, + PROTEIN_3TO1, + PROTEIN_HEAVY_ATOMS, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, + RNA_1TO3, + RNA_BACKBONE_ATOMS, + RNA_HEAVY_ATOMS, + RNA_RESIDUE_TO_RES_TYPE, + RNA_UNK_RES_TYPE, +) +from .esmfold2_types import ( + MSA, + DNAInput, + LigandInput, + Modification, + ProteinInput, + RNAInput, + StructurePredictionInput, +) + +_ZERO_POS = np.zeros(3, dtype=np.float32) +_ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {} +_ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {} +_TWO_LETTER_ELEMENTS = frozenset({"FE", "ZN", "MG", "MN", "CO", "NI", "CU", "SE", "BR"}) + + +@dataclass +class AtomInfo: + """One row in the unpadded atom table.""" + + name: str + element: str + charge: int + ref_pos: np.ndarray # R has shape (3,). + pos: np.ndarray # X has shape (3,). + token_index: int = -1 + atom_index: int = -1 + space_uid: int = -1 + is_valid: bool = True + + +@dataclass +class TokenInfo: + """Biological and atom-span annotations for one model token.""" + + token_index: int + residue_index: int + residue_name: str + mol_type: int + res_type: int + input_id: int + asym_id: int + sym_id: int + entity_id: int + atom_start: int + atom_count: int + + +@dataclass +class ChainInfo: + """One input chain after entity and symmetry assignment.""" + + chain_id: str + asym_id: int + entity_id: int + sym_id: int + mol_type: int + tokens: list[TokenInfo] = field(default_factory=list) + ligand_bonds: list[tuple[str, str]] = field(default_factory=list) + + +@dataclass +class _TokenizationState: + """Mutable cursor shared by residue tokenizers.""" + + token_index: int + atom_index: int + space_uid: int + tokens: list[TokenInfo] = field(default_factory=list) + atoms: list[AtomInfo] = field(default_factory=list) + + def _append_atom( + self, + name: str, + element: str, + charge: int, + ref_pos: np.ndarray | None, + ) -> None: + self.atoms.append( + AtomInfo( + name=name, + element=element, + charge=charge, + ref_pos=(ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy()), + pos=_ZERO_POS.copy(), + token_index=self.token_index, + atom_index=self.atom_index, + space_uid=self.space_uid, + ) + ) + self.atom_index += 1 + + def _append_token( + self, + *, + residue_index: int, + residue_name: str, + mol_type: int, + res_type: int, + input_id: int, + asym_id: int, + sym_id: int, + entity_id: int, + atom_start: int, + atom_count: int, + ) -> None: + self.tokens.append( + TokenInfo( + token_index=self.token_index, + residue_index=residue_index, + residue_name=residue_name, + mol_type=mol_type, + res_type=res_type, + input_id=input_id, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + atom_start=atom_start, + atom_count=atom_count, + ) + ) + self.token_index += 1 + + def add_residue_token( + self, + atom_specs: list[tuple[str, str, int, np.ndarray | None]], + **token_fields: Any, + ) -> None: + start = self.atom_index + for atom_spec in atom_specs: + self._append_atom(*atom_spec) + self._append_token( + atom_start=start, + atom_count=len(atom_specs), + **token_fields, + ) + self.space_uid += 1 + + def add_atom_tokens( + self, + atom_specs: list[tuple[str, str, int, np.ndarray | None]], + **token_fields: Any, + ) -> None: + for atom_spec in atom_specs: + start = self.atom_index + self._append_atom(*atom_spec) + self._append_token(atom_start=start, atom_count=1, **token_fields) + self.space_uid += 1 + + +def encode_atom_name(name: str) -> list[int]: + """Encode a four-character atom name with the model's ASCII offset.""" + cached = _ENCODE_ATOM_NAME_CACHE.get(name) + if cached is None: + cached = [0 if char == " " else ord(char) - 32 for char in name.ljust(4)[:4]] + _ENCODE_ATOM_NAME_CACHE[name] = cached + return cached + + +def get_element_atomic_num(element: str) -> int: + """Map an element symbol to the model's atomic-number vocabulary.""" + cached = _ELEMENT_ATOMIC_NUM_CACHE.get(element) + if cached is None: + cached = ELEMENT_TO_ATOMIC_NUM.get(element.upper(), 0) + _ELEMENT_ATOMIC_NUM_CACHE[element] = cached + return cached + + +def _infer_element(atom_name: str) -> str: + normalized = atom_name.strip() + if not normalized: + return "C" + if normalized[0].isdigit(): + return normalized[1] if len(normalized) > 1 else "H" + if len(normalized) == 2 and normalized in _TWO_LETTER_ELEMENTS: + return normalized + return normalized[0] + + +def _compute_res_type(name: str, mol_type: int) -> int: + if mol_type == MOL_TYPE_PROTEIN: + return PROTEIN_RESIDUE_TO_RES_TYPE.get(name, PROTEIN_UNK_RES_TYPE) + if mol_type == MOL_TYPE_DNA: + return DNA_RESIDUE_TO_RES_TYPE.get( + name, RNA_RESIDUE_TO_RES_TYPE.get(name, DNA_UNK_RES_TYPE) + ) + if mol_type == MOL_TYPE_RNA: + return RNA_RESIDUE_TO_RES_TYPE.get( + name, DNA_RESIDUE_TO_RES_TYPE.get(name, RNA_UNK_RES_TYPE) + ) + return PROTEIN_UNK_RES_TYPE + + +def _compute_esm_input_id(name: str, mol_type: int) -> int: + if mol_type != MOL_TYPE_PROTEIN: + return DNA_RNA_LIGAND_INPUT_ID + letter = PROTEIN_3TO1.get(name) + return ( + DNA_RNA_LIGAND_INPUT_ID + if letter is None + else ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"]) + ) + + +def _apply_modifications(residues: list[str], modifications: list[Modification] | None) -> set[int]: + changed: set[int] = set() + for modification in modifications or (): + residues[modification.position] = modification.ccd + changed.add(modification.position) + return changed + + +def _ideal_atom_specs( + residue_name: str, + residue_type: int, + atom_names: list[str], + *, + charges: bool = True, +) -> list[tuple[str, str, int, np.ndarray | None]]: + return [ + ( + atom_name, + _infer_element(atom_name), + CHARGED_ATOMS.get((residue_name, atom_name), 0) if charges else 0, + get_idealized_atom_pos(residue_type, atom_name), + ) + for atom_name in atom_names + ] + + +def _ccd_atom_specs( + residue_name: str, + atoms: list[tuple[str, str, int]], + excluded: set[str], + *, + force_zero: bool = False, +) -> list[tuple[str, str, int, np.ndarray | None]]: + return [ + ( + atom_name, + element, + charge, + None if force_zero else get_ligand_idealized_atom_pos(residue_name, atom_name), + ) + for atom_name, element, charge in atoms + if atom_name not in excluded + ] + + +def tokenize_protein( + sequence: str, + modifications: list[Modification] | None, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize protein residues, atom-tokenizing modified CCD components.""" + residues = [PROTEIN_1TO3.get(letter, "UNK") for letter in sequence] + modified = _apply_modifications(residues, modifications) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + + for residue_index, residue_name in enumerate(residues): + canonical_name = "MET" if residue_name == "MSE" else residue_name + common_fields = { + "residue_index": residue_index, + "mol_type": MOL_TYPE_PROTEIN, + "asym_id": asym_id, + "sym_id": sym_id, + "entity_id": entity_id, + } + if residue_index not in modified and canonical_name in PROTEIN_HEAVY_ATOMS: + residue_type = _compute_res_type(canonical_name, MOL_TYPE_PROTEIN) + state.add_residue_token( + _ideal_atom_specs( + canonical_name, + residue_type, + PROTEIN_HEAVY_ATOMS[canonical_name], + ), + residue_name=canonical_name, + res_type=residue_type, + input_id=_compute_esm_input_id(canonical_name, MOL_TYPE_PROTEIN), + **common_fields, + ) + continue + + ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name) + if ccd_atoms is None: + ccd_atoms = [ + (_infer_element(name), _infer_element(name), 0) for name in ("N", "CA", "C", "O") + ] + excluded = ( + set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name) + ) + retained = [atom for atom in ccd_atoms if atom[0] not in excluded] + state.add_atom_tokens( + _ccd_atom_specs( + residue_name, + retained, + set(), + force_zero=len(retained) == 1, + ), + residue_name=residue_name, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + **common_fields, + ) + return state.tokens, state.atoms + + +def tokenize_nucleotide( + sequence: str, + modifications: list[Modification] | None, + mol_type: int, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize DNA or RNA, retaining backbone atoms for unknown bases.""" + dna = mol_type == MOL_TYPE_DNA + letter_map = DNA_1TO3 if dna else RNA_1TO3 + heavy_atoms = DNA_HEAVY_ATOMS if dna else RNA_HEAVY_ATOMS + backbone_atoms = DNA_BACKBONE_ATOMS if dna else RNA_BACKBONE_ATOMS + unknown_type = DNA_UNK_RES_TYPE if dna else RNA_UNK_RES_TYPE + residues = [letter_map.get(letter, "UNK") for letter in sequence] + modified = _apply_modifications(residues, modifications) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + + for residue_index, residue_name in enumerate(residues): + common_fields = { + "residue_index": residue_index, + "residue_name": residue_name, + "mol_type": mol_type, + "asym_id": asym_id, + "sym_id": sym_id, + "entity_id": entity_id, + "input_id": DNA_RNA_LIGAND_INPUT_ID, + } + if residue_index not in modified and residue_name in heavy_atoms: + residue_type = _compute_res_type(residue_name, mol_type) + state.add_residue_token( + _ideal_atom_specs(residue_name, residue_type, heavy_atoms[residue_name]), + res_type=residue_type, + **common_fields, + ) + continue + if residue_index not in modified and residue_name == "UNK": + state.add_residue_token( + [(atom_name, _infer_element(atom_name), 0, None) for atom_name in backbone_atoms], + res_type=unknown_type, + **common_fields, + ) + continue + + ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name) + if ccd_atoms is None: + ccd_atoms = [(_infer_element(name), _infer_element(name), 0) for name in backbone_atoms] + excluded = ( + set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name) + ) + state.add_atom_tokens( + _ccd_atom_specs(residue_name, ccd_atoms, excluded), + res_type=PROTEIN_UNK_RES_TYPE, + **common_fields, + ) + return state.tokens, state.atoms + + +def tokenize_ligand_ccd( + ccd_codes: list[str], + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + has_covalent_bond: bool, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize CCD ligands with one model token per retained atom.""" + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + for residue_index, code in enumerate(ccd_codes): + ccd_atoms = get_ligand_ccd_atoms_with_charges(code) + if ccd_atoms is None: + raise ValueError(f"CCD component {code} not found") + excluded = get_ccd_leaving_atoms(code) if has_covalent_bond else set() + state.add_atom_tokens( + _ccd_atom_specs(code, ccd_atoms, excluded), + residue_index=residue_index, + residue_name=code, + mol_type=MOL_TYPE_NONPOLYMER, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + ) + return state.tokens, state.atoms + + +def tokenize_ligand_smiles( + smiles: str, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + seed: int | None = None, +) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]: + """Generate a conformer and tokenize each heavy atom of a SMILES ligand.""" + from rdkit import Chem + from rdkit.Chem import AllChem + + molecule = Chem.MolFromSmiles(smiles) + if molecule is None: + raise ValueError(f"Failed to parse SMILES: {smiles}") + molecule = Chem.AddHs(molecule) + canonical_order = AllChem.CanonicalRankAtoms(molecule) # type: ignore[attr-defined] + for atom, canonical_index in zip(molecule.GetAtoms(), canonical_order, strict=True): + name = atom.GetSymbol().upper() + str(canonical_index + 1) + if len(name) > 4: + raise ValueError(f"SMILES {smiles} has atom name longer than 4 chars: {name}") + atom.SetProp("name", name) + + options = AllChem.ETKDGv3() # type: ignore[attr-defined] + options.clearConfs = False + if seed is not None: + options.randomSeed = seed + conformer_id = AllChem.EmbedMolecule(molecule, options) # type: ignore[attr-defined] + if conformer_id == -1: + options.useRandomCoords = True + conformer_id = AllChem.EmbedMolecule(molecule, options) # type: ignore[attr-defined] + if conformer_id != -1: + with suppress(RuntimeError, ValueError): + AllChem.UFFOptimizeMolecule( # type: ignore[attr-defined] + molecule, confId=conformer_id, maxIters=1000 + ) + + molecule = Chem.RemoveHs(molecule) + if molecule.GetNumConformers() == 0: + raise ValueError(f"Failed to generate conformer for SMILES: {smiles}") + conformer = molecule.GetConformer(0) + atom_specs: list[tuple[str, str, int, np.ndarray | None]] = [] + for atom in molecule.GetAtoms(): + position = conformer.GetAtomPosition(atom.GetIdx()) + atom_specs.append( + ( + atom.GetProp("name"), + atom.GetSymbol(), + atom.GetFormalCharge(), + np.asarray([position.x, position.y, position.z], dtype=np.float32), + ) + ) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + state.add_atom_tokens( + atom_specs, + residue_index=0, + residue_name="LIG", + mol_type=MOL_TYPE_NONPOLYMER, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + ) + bonds = [ + ( + bond.GetBeginAtom().GetProp("name"), + bond.GetEndAtom().GetProp("name"), + ) + for bond in molecule.GetBonds() + ] + return state.tokens, state.atoms, bonds + + +def _get_sequence_key(item: Any) -> str: + if isinstance(item, ProteinInput): + return f"PROTEIN:{item.sequence}" + if isinstance(item, DNAInput): + return f"DNA:{item.sequence}" + if isinstance(item, RNAInput): + return f"RNA:{item.sequence}" + if isinstance(item, LigandInput): + return f"LIGAND_CCD:{','.join(item.ccd)}" if item.ccd else f"LIGAND_SMILES:{item.smiles}" + raise ValueError(f"Unknown input type: {type(item)}") + + +def _tokenize_chain( + item: Any, + chain_id: str, + *, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + covalent_chains: set[str], + seed: int | None, +) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]: + common = { + "entity_id": entity_id, + "asym_id": asym_id, + "sym_id": sym_id, + "token_offset": token_offset, + "atom_offset": atom_offset, + "space_uid_offset": space_uid_offset, + } + if isinstance(item, ProteinInput): + if item.msa is None: + warnings.warn( + f"No MSA provided for {item.id}, using single sequence mode", + stacklevel=2, + ) + tokens, atoms = tokenize_protein(item.sequence, item.modifications, **common) + return tokens, atoms, [] + if isinstance(item, (DNAInput, RNAInput)): + mol_type = MOL_TYPE_DNA if isinstance(item, DNAInput) else MOL_TYPE_RNA + tokens, atoms = tokenize_nucleotide( + item.sequence, item.modifications, mol_type=mol_type, **common + ) + return tokens, atoms, [] + if not isinstance(item, LigandInput): + raise ValueError(f"Unknown input type: {type(item)}") + if item.ccd is not None: + if item.smiles is not None: + warnings.warn("Both ccd and smiles provided, using ccd", stacklevel=2) + tokens, atoms = tokenize_ligand_ccd( + item.ccd, + has_covalent_bond=chain_id in covalent_chains, + **common, + ) + return tokens, atoms, [] + if item.smiles is not None: + return tokenize_ligand_smiles(item.smiles, seed=seed, **common) + raise ValueError("LigandInput must have either ccd or smiles") + + +def build_chains_from_input( + input: StructurePredictionInput, seed: int | None = None +) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]: + """Assign entities and symmetry copies, then tokenize every input chain.""" + chains: list[ChainInfo] = [] + tokens: list[TokenInfo] = [] + atoms: list[AtomInfo] = [] + entity_for_sequence: dict[str, int] = {} + next_symmetry: dict[int, int] = {} + covalent_chains = { + chain_id + for bond in input.covalent_bonds or () + for chain_id in (bond.chain_id1, bond.chain_id2) + } + space_uid_offset = 0 + + for item in input.sequences: + key = _get_sequence_key(item) + entity_id = entity_for_sequence.setdefault(key, len(entity_for_sequence)) + chain_ids = [item.id] if isinstance(item.id, str) else item.id + for chain_id in chain_ids: + sym_id = next_symmetry.get(entity_id, 0) + next_symmetry[entity_id] = sym_id + 1 + asym_id = len(chains) + new_tokens, new_atoms, ligand_bonds = _tokenize_chain( + item, + chain_id, + entity_id=entity_id, + asym_id=asym_id, + sym_id=sym_id, + token_offset=len(tokens), + atom_offset=len(atoms), + space_uid_offset=space_uid_offset, + covalent_chains=covalent_chains, + seed=seed, + ) + chains.append( + ChainInfo( + chain_id=chain_id, + asym_id=asym_id, + entity_id=entity_id, + sym_id=sym_id, + mol_type=(new_tokens[0].mol_type if new_tokens else MOL_TYPE_PROTEIN), + tokens=new_tokens, + ligand_bonds=ligand_bonds, + ) + ) + tokens.extend(new_tokens) + atoms.extend(new_atoms) + space_uid_offset += len({atom.space_uid for atom in new_atoms}) + return chains, tokens, atoms + + +def _atom_indices_by_name(atoms: list[AtomInfo]) -> dict[int, dict[str, int]]: + result: dict[int, dict[str, int]] = defaultdict(dict) + for atom in atoms: + if atom.is_valid: + result[atom.token_index][atom.name] = atom.atom_index + return result + + +def _ligand_frames( + tokens: list[TokenInfo], + atoms: list[AtomInfo], + atom_indices: dict[int, dict[str, int]], +) -> dict[int, tuple[int, int, int]]: + atom_for_token: dict[int, int] = {} + tokens_by_residue: dict[tuple[int, int], list[int]] = defaultdict(list) + for token in tokens: + if token.mol_type != MOL_TYPE_NONPOLYMER: + continue + named_atoms = atom_indices.get(token.token_index) + if named_atoms: + atom_for_token[token.token_index] = next(iter(named_atoms.values())) + tokens_by_residue[(token.asym_id, token.residue_index)].append(token.token_index) + + frames: dict[int, tuple[int, int, int]] = {} + for residue_tokens in tokens_by_residue.values(): + residue_atoms = [ + atom_for_token[token] for token in residue_tokens if token in atom_for_token + ] + if len(residue_atoms) < 3: + for token in residue_tokens: + if token in atom_for_token: + atom_index = atom_for_token[token] + frames[token] = (atom_index, atom_index, atom_index) + continue + R = np.asarray([atoms[index].ref_pos for index in residue_atoms]) + distances = np.sqrt(((R[:, None] - R[None]) ** 2).sum(-1)) + nearest = np.argsort(distances, axis=1) + local = np.column_stack((nearest[:, 1], nearest[:, 0], nearest[:, 2])) + local_index = {atom_index: index for index, atom_index in enumerate(residue_atoms)} + for token in residue_tokens: + atom_index = atom_for_token.get(token) + if atom_index is None: + continue + selected = local[local_index[atom_index]] + frames[token] = tuple(residue_atoms[int(index)] for index in selected) + return frames + + +def _frame_for_token( + token: TokenInfo, + named_atoms: dict[str, int], + ligand_frames: dict[int, tuple[int, int, int]], +) -> tuple[int, int, int]: + fallback = next(iter(named_atoms.values()), 0) + if token.mol_type == MOL_TYPE_PROTEIN: + return ( + (fallback, fallback, fallback) + if token.res_type == PROTEIN_UNK_RES_TYPE + else ( + named_atoms.get("N", 0), + named_atoms.get("CA", 0), + named_atoms.get("C", 0), + ) + ) + if token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): + return ( + (fallback, fallback, fallback) + if token.res_type == PROTEIN_UNK_RES_TYPE + else ( + named_atoms.get("C1'", 0), + named_atoms.get("C3'", 0), + named_atoms.get("C4'", 0), + ) + ) + if token.mol_type == MOL_TYPE_NONPOLYMER: + return ligand_frames.get(token.token_index, (fallback, fallback, fallback)) + return fallback, fallback, fallback + + +def _resolved_frames( + frames: np.ndarray, tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> np.ndarray: + if not tokens: + return np.zeros(0, dtype=bool) + X = ( + np.asarray([atom.pos for atom in atoms], dtype=np.float32) + if atoms + else np.zeros((0, 3), dtype=np.float32) + ) + valid_atoms = ( + np.asarray([atom.is_valid for atom in atoms], dtype=bool) + if atoms + else np.zeros(0, dtype=bool) + ) + resolved_atoms = valid_atoms & np.any(X != 0, axis=1) + origin = X[frames[:, 1]] + left = X[frames[:, 0]] - origin + right = X[frames[:, 2]] - origin + left_norm = np.linalg.norm(left, axis=1) + right_norm = np.linalg.norm(right, axis=1) + valid_norms = (left_norm >= 1e-6) & (right_norm >= 1e-6) + cosine = np.zeros(len(tokens), dtype=np.float32) + if np.any(valid_norms): + cosine[valid_norms] = np.sum(left[valid_norms] * right[valid_norms], axis=1) / ( + left_norm[valid_norms] * right_norm[valid_norms] + ) + angle = np.degrees(np.arccos(np.abs(np.clip(cosine, -1, 1)))) + all_resolved = resolved_atoms[frames].all(axis=1) + repeated = (frames[:, 0] == frames[:, 1]) & (frames[:, 1] == frames[:, 2]) + return all_resolved & ~repeated & valid_norms & (angle >= 25) + + +def compute_frame_indices( + tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> tuple[np.ndarray, np.ndarray]: + """Return frame atom indices F with shape (l, 3) and validity M with shape (l,).""" + named_atoms = _atom_indices_by_name(atoms) + ligand_frames = _ligand_frames(tokens, atoms, named_atoms) + frames = np.asarray( + [ + _frame_for_token(token, named_atoms.get(token.token_index, {}), ligand_frames) + for token in tokens + ], + dtype=np.int64, + ) + return frames, _resolved_frames(frames, tokens, atoms) + + +def _atom_tokenized_residues( + tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> dict[tuple[int, int], list[tuple[str, int]]]: + grouped: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list) + for atom in atoms: + if not atom.is_valid or atom.token_index >= len(tokens): + continue + token = tokens[atom.token_index] + if token.mol_type == MOL_TYPE_NONPOLYMER or token.res_type == PROTEIN_UNK_RES_TYPE: + grouped[(token.asym_id, token.residue_index)].append((atom.name, atom.token_index)) + return grouped + + +def _backbone_token( + residue_tokens: list[TokenInfo], atom_name: str, atoms: list[AtomInfo] +) -> int | None: + if len(residue_tokens) == 1 and residue_tokens[0].res_type != PROTEIN_UNK_RES_TYPE: + return residue_tokens[0].token_index + for token in residue_tokens: + for atom_index in range(token.atom_start, token.atom_start + token.atom_count): + if atom_index < len(atoms) and atoms[atom_index].name == atom_name: + return token.token_index + return residue_tokens[0].token_index if residue_tokens else None + + +def compute_token_bonds( + tokens: list[TokenInfo], + atoms: list[AtomInfo], + input: StructurePredictionInput, + chains: list[ChainInfo], +) -> torch.Tensor: + """Build the symmetric token-bond matrix M with shape (l, l, 1).""" + edges: set[tuple[int, int]] = set() + + def connect(left: int | None, right: int | None) -> None: + if left is not None and right is not None and left != right: + edges.add((min(left, right), max(left, right))) + + explicit_bonds = { + (chain.asym_id, 0): chain.ligand_bonds for chain in chains if chain.ligand_bonds + } + for residue_key, atom_list in _atom_tokenized_residues(tokens, atoms).items(): + if not atom_list: + continue + residue_name = tokens[atom_list[0][1]].residue_name + token_for_name = {name: token_index for name, token_index in atom_list} + bonds = explicit_bonds.get(residue_key) + if bonds is None: + bonds = get_ligand_ccd_bonds(residue_name) + if bonds: + for left_name, right_name in bonds: + if left_name in token_for_name and right_name in token_for_name: + connect(token_for_name[left_name], token_for_name[right_name]) + else: + for left, right in combinations([token_index for _, token_index in atom_list], 2): + connect(left, right) + + if input.covalent_bonds: + chain_for_id = {chain.chain_id: chain for chain in chains} + residue_atoms: dict[tuple[int, int], list[AtomInfo]] = defaultdict(list) + for atom in atoms: + if atom.is_valid and atom.token_index < len(tokens): + token = tokens[atom.token_index] + residue_atoms[(token.asym_id, token.residue_index)].append(atom) + for bond in input.covalent_bonds: + left_chain = chain_for_id.get(bond.chain_id1) + right_chain = chain_for_id.get(bond.chain_id2) + if left_chain is None or right_chain is None: + continue + left_atoms = residue_atoms.get((left_chain.asym_id, bond.res_idx1), []) + right_atoms = residue_atoms.get((right_chain.asym_id, bond.res_idx2), []) + if bond.atom_idx1 < len(left_atoms) and bond.atom_idx2 < len(right_atoms): + connect( + left_atoms[bond.atom_idx1].token_index, + right_atoms[bond.atom_idx2].token_index, + ) + + protein_residues: dict[tuple[int, int], list[TokenInfo]] = defaultdict(list) + for token in tokens: + if token.mol_type == MOL_TYPE_PROTEIN: + protein_residues[(token.asym_id, token.residue_index)].append(token) + for (asym_id, residue_index), residue_tokens in protein_residues.items(): + if not any(token.res_type == PROTEIN_UNK_RES_TYPE for token in residue_tokens): + continue + previous = protein_residues.get((asym_id, residue_index - 1)) + following = protein_residues.get((asym_id, residue_index + 1)) + if previous: + connect( + _backbone_token(previous, "C", atoms), + _backbone_token(residue_tokens, "N", atoms), + ) + if following: + connect( + _backbone_token(residue_tokens, "C", atoms), + _backbone_token(following, "N", atoms), + ) + + matrix = torch.zeros(len(tokens), len(tokens), 1, dtype=torch.float32) + for left, right in edges: + matrix[left, right, 0] = 1.0 + matrix[right, left, 0] = 1.0 + return matrix + + +def compute_representative_atoms(tokens: list[TokenInfo], atoms: list[AtomInfo]) -> torch.Tensor: + """Choose one distogram atom per token and return indices I with shape (l,).""" + named_atoms = _atom_indices_by_name(atoms) + representatives = torch.zeros(len(tokens), dtype=torch.int64) + for token in tokens: + names = named_atoms.get(token.token_index, {}) + fallback = next(iter(names.values()), 0) + if token.mol_type == MOL_TYPE_PROTEIN: + representative = names.get("CB", names.get("CA", fallback)) + elif token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): + if token.res_type in (27, 32): + representative = names.get("C1'", fallback) + elif token.res_type in (23, 24, 28, 29): + representative = names.get("C4", names.get("C1'", fallback)) + else: + representative = names.get("C2", names.get("C1'", fallback)) + else: + representative = fallback + representatives[token.token_index] = representative + return representatives + + +def _msa_assignments( + input: StructurePredictionInput, chains: list[ChainInfo] +) -> dict[int, MSA | None]: + chain_msas: dict[int, MSA | None] = {} + chain_index = 0 + for item in input.sequences: + chain_ids = [item.id] if isinstance(item.id, str) else list(item.id) + for _ in chain_ids: + chain = chains[chain_index] + if isinstance(item, ProteinInput): + chain_msas[chain.asym_id] = ( + MSA.from_sequences([item.sequence]) if item.msa is None else item.msa + ) + else: + chain_msas[chain.asym_id] = None + chain_index += 1 + return chain_msas + + +def compute_msa_features( + input: StructurePredictionInput, + chains: list[ChainInfo], + tokens: list[TokenInfo], + max_seqs: int = 16384, +) -> dict[str, torch.Tensor]: + """Pair per-chain MSAs and return row features with shape (m, l).""" + from .esmfold2_paired_msa import ( + construct_paired_msa, + protein_letter_to_res_type, + ) + + chain_msas = _msa_assignments(input, chains) + query_types = { + chain.asym_id: np.asarray( + [token.res_type for token in tokens if token.asym_id == chain.asym_id], + dtype=np.int64, + ) + for chain in chains + } + msa_residues, deletion_counts, _ = construct_paired_msa( + chain_msas, + query_types, + np.asarray([token.asym_id for token in tokens], dtype=np.int64), + np.asarray([token.residue_index for token in tokens], dtype=np.int64), + letter_to_res_type=protein_letter_to_res_type(), + max_seqs=max_seqs, + ) + for token in tokens: + if chain_msas.get(token.asym_id) is None: + msa_residues[:, token.token_index] = MSA_GAP_TOKEN_ID + msa_residues[0, token.token_index] = token.res_type + if msa_residues.shape[0] == 0: + msa_residues = np.full((1, len(tokens)), MSA_GAP_TOKEN_ID, dtype=np.int64) + deletion_counts = np.zeros((1, len(tokens)), dtype=np.float32) + + msa = torch.from_numpy(msa_residues) + deletion_count = torch.from_numpy(deletion_counts) + deletion_value = (np.pi / 2) * torch.arctan(deletion_count / 3) + return { + "msa": msa, + "deletion_value": deletion_value, + "has_deletion": deletion_count > 0, + "deletion_mean": deletion_value.mean(dim=0), + "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool), + } + + +def compute_distogram_conditioning( + input: StructurePredictionInput, + chains: list[ChainInfo], + tokens: list[TokenInfo], + disto_center: torch.Tensor, + min_dist: float = 2.0, + max_dist: float = 22.0, + num_bins: int = 64, +) -> tuple[torch.Tensor, torch.Tensor]: + """Bin user distances into D and return D plus its Boolean mask M.""" + del disto_center + n_tokens = len(tokens) + bins = torch.zeros((n_tokens, n_tokens), dtype=torch.long) + mask = torch.zeros((n_tokens, n_tokens), dtype=torch.bool) + if not input.distogram_conditioning: + return bins, mask + asym_for_chain = {chain.chain_id: chain.asym_id for chain in chains} + tokens_for_asym: dict[int, list[int]] = defaultdict(list) + for token in tokens: + tokens_for_asym[token.asym_id].append(token.token_index) + boundaries = torch.linspace(min_dist, max_dist, num_bins + 1) + + for conditioning in input.distogram_conditioning: + asym_id = asym_for_chain.get(conditioning.chain_id) + if asym_id is None: + continue + indices = tokens_for_asym[asym_id] + distances = torch.as_tensor(conditioning.distogram, dtype=torch.float32) + expected_shape = (len(indices), len(indices)) + if distances.shape != expected_shape: + raise ValueError( + f"Distogram shape {distances.shape} doesn't match chain length {len(indices)}" + ) + selected = torch.bucketize(distances, boundaries[:-1]).sub(1).clamp(0, num_bins - 1) + token_indices_tensor = torch.as_tensor(indices, dtype=torch.long) + bins[token_indices_tensor[:, None], token_indices_tensor[None, :]] = selected + mask[token_indices_tensor[:, None], token_indices_tensor[None, :]] = True + return bins, mask + + +def _padded_atoms(atoms: list[AtomInfo]) -> list[AtomInfo]: + target = math.ceil(len(atoms) / 32) * 32 if atoms else 32 + padding = [ + AtomInfo( + name="", + element="", + charge=0, + ref_pos=_ZERO_POS.copy(), + pos=_ZERO_POS.copy(), + token_index=0, + atom_index=index, + space_uid=0, + is_valid=False, + ) + for index in range(len(atoms), target) + ] + return [*atoms, *padding] + + +def _token_tensors(tokens: list[TokenInfo]) -> dict[str, torch.Tensor]: + fields = { + "token_index": "token_index", + "residue_index": "residue_index", + "asym_id": "asym_id", + "sym_id": "sym_id", + "entity_id": "entity_id", + "mol_type": "mol_type", + "res_type": "res_type", + "input_ids": "input_id", + } + return { + output_name: torch.from_numpy( + np.asarray([getattr(token, attribute) for token in tokens], dtype=np.int64) + ) + for output_name, attribute in fields.items() + } + + +def _atom_tensors(atoms: list[AtomInfo]) -> dict[str, torch.Tensor]: + n_atoms = len(atoms) + ref_pos = np.zeros((n_atoms, 3), dtype=np.float32) + ref_element = np.zeros(n_atoms, dtype=np.int64) + ref_charge = np.zeros(n_atoms, dtype=np.int8) + ref_name = np.zeros((n_atoms, 4), dtype=np.int64) + ref_space = np.zeros(n_atoms, dtype=np.int64) + atom_mask = np.zeros(n_atoms, dtype=np.bool_) + atom_to_token = np.zeros(n_atoms, dtype=np.int64) + positions = np.zeros((n_atoms, 3), dtype=np.float64) + valid = np.zeros(n_atoms, dtype=np.bool_) + for index, atom in enumerate(atoms): + if atom.ref_pos is not None: + ref_pos[index] = atom.ref_pos + ref_charge[index] = atom.charge + ref_space[index] = atom.space_uid if atom.space_uid >= 0 else atom.token_index + atom_mask[index] = atom.is_valid + valid[index] = atom.is_valid + positions[index] = atom.pos + if atom.is_valid: + ref_element[index] = get_element_atomic_num(atom.element) + ref_name[index] = encode_atom_name(atom.name) + atom_to_token[index] = atom.token_index + + resolved = valid & np.any(positions != 0, axis=1) + X = torch.from_numpy(positions) + resolved_mask = torch.from_numpy(resolved) + valid_mask = torch.from_numpy(valid) + if resolved_mask.any(): + X = X - X[resolved_mask].mean(dim=0, keepdim=True) + X[~valid_mask] = 0.0 + return { + "ref_pos": torch.from_numpy(ref_pos), + "ref_element": torch.from_numpy(ref_element), + "ref_charge": torch.from_numpy(ref_charge), + "ref_atom_name_chars": torch.from_numpy(ref_name), + "ref_space_uid": torch.from_numpy(ref_space), + "gt_coords": X.float().unsqueeze(0), + "atom_attention_mask": torch.from_numpy(atom_mask), + "atom_to_token": torch.from_numpy(atom_to_token), + "is_resolved": torch.tensor(resolved, dtype=torch.bool), + } + + +def build_feature_tensors( + chains: list[ChainInfo], + tokens: list[TokenInfo], + atoms: list[AtomInfo], + input: StructurePredictionInput, +) -> dict[str, torch.Tensor]: + """Assemble the complete unbatched ESMFold2 feature dictionary.""" + token_features = _token_tensors(tokens) + atom_features = _atom_tensors(_padded_atoms(atoms)) + frames, _ = compute_frame_indices(tokens, atoms) + msa_features = compute_msa_features(input, chains, tokens) + distogram, distogram_mask = compute_distogram_conditioning( + input, + chains, + tokens, + torch.zeros(len(tokens), 3, dtype=torch.float32), + ) + return { + **token_features, + "token_bonds": compute_token_bonds(tokens, atoms, input, chains), + "token_attention_mask": torch.ones(len(tokens), dtype=torch.bool), + "pocket_feature": torch.zeros(len(tokens), dtype=torch.long), + **atom_features, + "distogram_atom_idx": compute_representative_atoms(tokens, atoms), + "frames_idx": torch.from_numpy(frames).to(torch.int64), + "disto_cond": distogram, + "disto_cond_mask": distogram_mask, + **msa_features, + } + + +def prepare_esmfold2_input( + input: StructurePredictionInput, seed: int | None = None +) -> tuple[dict[str, torch.Tensor], list[ChainInfo]]: + """Convert one typed request to model features and output-chain metadata.""" + chains, tokens, atoms = build_chains_from_input(input, seed) + return build_feature_tensors(chains, tokens, atoms, input), chains diff --git a/src/fastplms/models/esmfold2/esmfold2_processor.py b/src/fastplms/models/esmfold2/esmfold2_processor.py new file mode 100644 index 0000000..d11f398 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_processor.py @@ -0,0 +1,332 @@ +"""Input preparation and output decoding for ESMFold2 inference.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import Tensor + +from .esmfold2_conformers import load_ccd +from .esmfold2_molecular_complex import MolecularComplexResult +from .esmfold2_output import build_molecular_complex_from_features +from .esmfold2_prepare_input import ChainInfo, prepare_esmfold2_input +from .esmfold2_types import MSA, Modification, ProteinInput, StructurePredictionInput +from .modeling_esmfold2_common import MSA_CONDITIONING_INPUT_NAMES +from .reproducibility import seed_context + +# Backward-compatible private alias for the pinned parity helpers. New callers +# should import ``seed_context`` from the public ``fastplms.models.esmfold2`` +# package instead of reaching into implementation modules. +_seed_context = seed_context + + +@dataclass(frozen=True) +class _SplitProteinState: + ids: dict[str, list[str]] + modifications: dict[str, list[Modification]] + msas: dict[str, MSA | None] + + +@dataclass(frozen=True) +class _PendingProtein: + source: ProteinInput + sequence: str + state: _SplitProteinState + + +def _chain_starts(chains: list[str]) -> list[int]: + starts: list[int] = [] + position = 0 + for chain in chains: + starts.append(position) + position += len(chain) + 1 + return starts + + +def _split_modifications( + item: ProteinInput, + chains: list[str], + starts: list[int], +) -> dict[str, list[Modification]]: + grouped: dict[str, list[Modification]] = {} + if item.modifications is None: + return grouped + for chain, start in zip(chains, starts, strict=True): + end = start + len(chain) + adjusted = [ + Modification(position=modification.position - start, ccd=modification.ccd) + for modification in item.modifications + if start <= modification.position < end + ] + grouped.setdefault(chain, []).extend(adjusted) + return grouped + + +def _split_msas( + item: ProteinInput, + chains: list[str], + starts: list[int], +) -> dict[str, MSA | None]: + grouped: dict[str, MSA | None] = {} + if item.msa is None: + return grouped + for chain, start in zip(chains, starts, strict=True): + if chain not in grouped: + grouped[chain] = item.msa.select_positions(np.arange(start, start + len(chain))) + return grouped + + +def _split_protein(item: ProteinInput) -> tuple[list[_PendingProtein], _SplitProteinState]: + chains = ":".join(item.sequence.split("|")).split(":") + starts = _chain_starts(chains) + base_id = item.id[0] if isinstance(item.id, list) else item.id + ids: dict[str, list[str]] = {} + for index, chain in enumerate(chains): + chain_ids = ids.setdefault(chain, []) + chain_ids.append(f"{base_id}_{index}") + state = _SplitProteinState( + ids=ids, + modifications=_split_modifications(item, chains, starts), + msas=_split_msas(item, chains, starts), + ) + pending = [ + _PendingProtein(item, chain, state) + for chain, chain_ids in ids.items() + if chain_ids + ] + return pending, state + + +def _resolve_pending(pending: _PendingProtein) -> ProteinInput: + item = pending.source + sequence = pending.sequence + state = pending.state + return ProteinInput( + id=state.ids[sequence], + sequence=sequence, + msa=state.msas.get(sequence) if item.msa else None, + modifications=(state.modifications.get(sequence) if item.modifications else None), + ) + + +def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePredictionInput: + """Expand chain delimiters and group repeated protein sequences by entity.""" + + if input.pocket is not None: + raise NotImplementedError( + "ESMFold2 pocket conditioning is present in the upstream input schema but " + "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this " + "input instead of silently emitting an all-zero pocket feature." + ) + + cleaned: list[Any] = [] + for item in input.sequences: + if not isinstance(item, ProteinInput): + cleaned.append(item) + continue + sequence = ":".join(item.sequence.split("|")) + if ":" not in sequence: + cleaned.append(item) + continue + if input.covalent_bonds is not None: + raise ValueError( + "Covalent bonds are not supported when using chainbreaks. " + "Chains must be separated into multiple ProteinInput objects." + ) + pending, _state = _split_protein(item) + cleaned.extend(pending) + + resolved = [ + _resolve_pending(item) if isinstance(item, _PendingProtein) else item + for item in cleaned + ] + return StructurePredictionInput( + sequences=resolved, + pocket=input.pocket, + distogram_conditioning=input.distogram_conditioning, + covalent_bonds=input.covalent_bonds, + ) + + +def _batch_features( + features: dict[str, Any], + device: torch.device | str | None, +) -> dict[str, Any]: + return { + name: (value[None].to(device) if device is not None else value[None]) + if isinstance(value, Tensor) + else value + for name, value in features.items() + } + + +def _sampler_overrides( + noise_scale: float | None, + step_scale: float | None, + max_inference_sigma: int | None, +) -> dict[str, Any]: + values = { + "noise_scale": noise_scale, + "step_scale": step_scale, + "max_inference_sigma": max_inference_sigma, + } + return {name: value for name, value in values.items() if value is not None} + + +class ESMFold2InputBuilder: + """Prepare public input objects, run folding, and decode model tensors.""" + + def __init__(self, ccd_cache: Path | None = None) -> None: + load_ccd(ccd_cache) + + def prepare_input( + self, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + cleaned = clean_esmfold2_input(input) + with seed_context(seed): + features, chain_infos = prepare_esmfold2_input(cleaned, seed=seed) + return _batch_features(features, device), chain_infos + + def prepare_model_input( + self, + model: Any, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + """Prepare features while enforcing the checkpoint's MSA contract.""" + + msa_conditioning = getattr(model.config, "msa_conditioning", None) + if not isinstance(msa_conditioning, bool): + raise RuntimeError("The ESMFold2 config has no Boolean msa_conditioning contract.") + if not msa_conditioning: + explicit_msa_ids = [ + item.id + for item in input.sequences + if isinstance(item, ProteinInput) and item.msa is not None + ] + if explicit_msa_ids: + raise ValueError( + "This ESMFold2 checkpoint was trained without MSA conditioning and " + f"rejects explicit MSAs for protein inputs {explicit_msa_ids!r}." + ) + features, chain_infos = self.prepare_input(input, seed=seed, device=device) + if not msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) + return features, chain_infos + + def __call__( + self, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + return self.prepare_input(input, seed=seed, device=device) + + def _decode_sample( + self, + output: Mapping[str, Tensor], + features: dict[str, Tensor], + chain_infos: list[ChainInfo], + sample: int, + complex_id: str, + ) -> MolecularComplexResult: + plddt = output["plddt"][sample] + molecular_complex = build_molecular_complex_from_features( + coords=output["sample_atom_coords"][sample], + plddt=plddt, + atom_mask=features["atom_attention_mask"][0], + ref_element=features["ref_element"][0], + ref_atom_name_chars=features["ref_atom_name_chars"][0], + chain_infos=chain_infos, + complex_id=complex_id, + ) + + def sample_tensor(name: str) -> Tensor | None: + value = output.get(name) + return None if value is None else value[sample].detach().cpu() + + def shared_tensor(name: str) -> Tensor | None: + value = output.get(name) + return None if value is None else value[0].detach().cpu() + + ptm = output.get("ptm") + iptm = output.get("iptm") + return MolecularComplexResult( + complex=molecular_complex, + plddt=plddt.detach().cpu(), + ptm=float(ptm[sample].item()) if ptm is not None else None, + iptm=float(iptm[sample].item()) if iptm is not None else None, + pae=sample_tensor("pae"), + distogram=shared_tensor("distogram_logits"), + pair_chains_iptm=sample_tensor("pair_chains_iptm"), + residue_index=shared_tensor("residue_index"), + entity_id=shared_tensor("entity_id"), + ) + + def decode( + self, + output: Mapping[str, Tensor], + features: dict[str, Tensor], + chain_infos: list[ChainInfo], + *, + num_diffusion_samples: int = 1, + complex_id: str = "pred", + ) -> MolecularComplexResult | list[MolecularComplexResult]: + results = [ + self._decode_sample(output, features, chain_infos, sample, complex_id) + for sample in range(output["sample_atom_coords"].shape[0]) + ] + return results[0] if num_diffusion_samples == 1 and len(results) == 1 else results + + def fold( + self, + model: Any, + input: StructurePredictionInput, + *, + num_loops: int = 3, + num_sampling_steps: int = 200, + num_diffusion_samples: int = 1, + seed: int | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: int | None = None, + early_exit: bool = False, + complex_id: str = "pred", + ) -> MolecularComplexResult | list[MolecularComplexResult]: + features, chain_infos = self.prepare_model_input( + model, + input, + seed=seed, + device=model.device, + ) + overrides = _sampler_overrides(noise_scale, step_scale, max_inference_sigma) + with torch.no_grad(), seed_context(seed): + output = model( + **features, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + early_exit=early_exit, + return_dict=True, + **overrides, + ) + return self.decode( + output, + features, + chain_infos, + num_diffusion_samples=num_diffusion_samples, + complex_id=complex_id, + ) + + +__all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input", "seed_context"] diff --git a/fastplms/esmfold2/esmfold2_protein_chain.py b/src/fastplms/models/esmfold2/esmfold2_protein_chain.py similarity index 71% rename from fastplms/esmfold2/esmfold2_protein_chain.py rename to src/fastplms/models/esmfold2/esmfold2_protein_chain.py index 8ab2613..acaaaf5 100644 --- a/fastplms/esmfold2/esmfold2_protein_chain.py +++ b/src/fastplms/models/esmfold2/esmfold2_protein_chain.py @@ -1,11 +1,14 @@ +"""Protein-chain data, geometry, and serialization for ESMFold2.""" + from __future__ import annotations import io import warnings +from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass, replace from functools import cached_property from pathlib import Path -from typing import Any, Mapping, Sequence +from typing import Any import biotite.structure as bs import brotli @@ -21,12 +24,17 @@ from scipy.spatial.distance import cdist, pdist, squareform from . import esmfold2_residue_constants as residue_constants -from .esmfold2_misc import slice_python_object_as_numpy from .esmfold2_affine3d import Affine3D from .esmfold2_aligner import Aligner from .esmfold2_atom_indexer import AtomIndexer from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca -from .esmfold2_mmcif_parsing import MmcifWrapper, Residue +from .esmfold2_misc import slice_python_object_as_numpy +from .esmfold2_mmcif_parsing import ( + PLDDT_B_FACTOR_SCALE, + MmcifWrapper, + Residue, + round_mmcif_columns, +) from .esmfold2_normalize_coordinates import ( apply_frame_to_coords, get_protein_normalization_frame, @@ -34,53 +42,63 @@ from .esmfold2_protein_structure import index_by_atom_name from .esmfold2_utils_types import PathOrBuffer -msgpack_numpy.patch() CHAIN_ID_CONST = "A" -def _str_key_to_int_key(dct: dict, ignore_keys: list[str] | None = None) -> dict: - new_dict = {} - for k, v in dct.items(): - v_new = v - if k not in ignore_keys and isinstance(v, dict): - v_new = _str_key_to_int_key(v, ignore_keys=ignore_keys) - # Note assembly_composition is *supposed* to have string keys. - if isinstance(k, str) and k.isdigit(): - new_dict[int(k)] = v_new - else: - new_dict[k] = v_new - return new_dict +def _str_key_to_int_key(values: dict, ignore_keys: list[str] | None = None) -> dict: + """Restore integer dictionary keys after JSON-compatible serialization.""" + ignored = frozenset(ignore_keys or ()) + restored = {} + for key, value in values.items(): + if isinstance(value, dict) and key not in ignored: + value = _str_key_to_int_key(value, ignore_keys=ignore_keys) + restored_key = int(key) if isinstance(key, str) and key.isdigit() else key + restored[restored_key] = value + return restored def _num_non_null_residues(seqres_to_structure_chain: Mapping[int, Residue]) -> int: - return sum( - residue.residue_number is not None - for residue in seqres_to_structure_chain.values() - ) + return sum(residue.residue_number is not None for residue in seqres_to_structure_chain.values()) -def infer_CB(C, N, Ca, L: float = 1.522, A: float = 1.927, D: float = -2.143): - """ - Inspired by a util in trDesign: - https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L92 - - input: 3 coords (a,b,c), (L)ength, (A)ngle, and (D)ihedral - output: 4th coord - """ - norm = lambda x: x / np.sqrt(np.square(x).sum(-1, keepdims=True) + 1e-8) - with np.errstate(invalid="ignore"): # inf - inf = nan is ok here - vec_bc = N - Ca - vec_ba = N - C - bc = norm(vec_bc) - n = norm(np.cross(vec_ba, bc)) - m = [bc, np.cross(n, bc), n] - d = [L * np.cos(A), L * np.sin(A) * np.cos(D), -L * np.sin(A) * np.sin(D)] - return Ca + sum([m * d for m, d in zip(m, d)]) +def infer_cb( + C, + N, + Ca, + bond_length: float = 1.522, + bond_angle: float = 1.927, + dihedral: float = -2.143, +): + """Infer C-beta coordinates from C, N, and C-alpha coordinates.""" + + def normalize(X: np.ndarray) -> np.ndarray: + return X / np.sqrt(np.square(X).sum(-1, keepdims=True) + 1e-8) + + with np.errstate(invalid="ignore"): + n_to_ca = N - Ca + n_to_c = N - C + axis = normalize(n_to_ca) + normal = normalize(np.cross(n_to_c, axis)) + basis = (axis, np.cross(normal, axis), normal) + offsets = ( + bond_length * np.cos(bond_angle), + bond_length * np.sin(bond_angle) * np.cos(dihedral), + -bond_length * np.sin(bond_angle) * np.sin(dihedral), + ) + return Ca + sum(vector * offset for vector, offset in zip(basis, offsets, strict=False)) def chain_to_ndarray( atom_array: bs.AtomArray, mmcif: MmcifWrapper, chain_id: str, is_predicted=False ): + if not isinstance(atom_array, bs.AtomArray): + raise TypeError("atom_array must be a biotite AtomArray.") + if not isinstance(mmcif, MmcifWrapper): + raise TypeError("mmcif must be an MmcifWrapper.") + if not isinstance(chain_id, str) or not chain_id: + raise ValueError("chain_id must be a non-empty string.") + if chain_id not in mmcif.chain_to_seqres or chain_id not in mmcif.seqres_to_structure: + raise ValueError(f"mmCIF data does not contain sequence mappings for chain {chain_id!r}.") entity_id = None for entity, chains in mmcif.entities.items(): if chain_id in chains: @@ -95,9 +113,10 @@ def chain_to_ndarray( confidence = np.ones([num_res], dtype=np.float32) + chain = atom_array[atom_array.chain_id == chain_id] + if not isinstance(chain, bs.AtomArray): + raise RuntimeError("Biotite selection did not return an AtomArray.") for res_index in range(num_res): - chain = atom_array[atom_array.chain_id == chain_id] - assert isinstance(chain, bs.AtomArray) res_at_position = mmcif.seqres_to_structure[chain_id][res_index] if res_at_position.residue_number is None: @@ -110,7 +129,8 @@ def chain_to_ndarray( & (chain.ins_code == res_at_position.insertion_code) & (chain.hetero == res_at_position.hetflag) ] - assert isinstance(res, bs.AtomArray) + if not isinstance(res, bs.AtomArray): + raise RuntimeError("Biotite residue selection did not return an AtomArray.") # Atom level features for atom in res: @@ -120,14 +140,13 @@ def chain_to_ndarray( atom_name = "SD" if atom_name in residue_constants.atom_order: - atom_positions[res_index, residue_constants.atom_order[atom_name]] = ( - atom.coord - ) + atom_positions[res_index, residue_constants.atom_order[atom_name]] = atom.coord atom_mask[res_index, residue_constants.atom_order[atom_name]] = True if is_predicted and atom_name == "CA": - confidence[res_index] = atom.b_factor + confidence[res_index] = atom.b_factor / PLDDT_B_FACTOR_SCALE - assert all(sequence), "Some residue name was not specified correctly" + if not sequence or not all(sequence): + raise ValueError("Some residue name was not specified correctly.") return ( sequence, atom_positions, @@ -153,1223 +172,1279 @@ class ProteinChain: atom37_mask: np.ndarray confidence: np.ndarray mmcif: MmcifWrapper | None = None - atom37_confidence: np.ndarray | None = None # [L, 37] per-atom pLDDT + atom37_confidence: np.ndarray | None = None # P has shape (l, 37). - def __post_init__(self): - assert self.atom37_mask.dtype == bool, self.atom37_mask.dtype - assert self.atom37_positions.shape[0] == len(self.sequence), ( - self.atom37_positions.shape, - len(self.sequence), - ) - assert self.atom37_mask.shape[0] == len(self.sequence), ( - self.atom37_mask.shape, - len(self.sequence), - ) - assert self.residue_index.shape[0] == len(self.sequence), ( - self.residue_index.shape, - len(self.sequence), - ) - assert self.insertion_code.shape[0] == len(self.sequence), ( - self.insertion_code.shape, - len(self.sequence), - ) - assert self.confidence.shape[0] == len(self.sequence), ( - self.confidence.shape, - len(self.sequence), - ) - if self.atom37_confidence is not None: - assert self.atom37_confidence.shape == self.atom37_mask.shape, ( - self.atom37_confidence.shape, - self.atom37_mask.shape, + # Construction and parsing + @classmethod + def chain_iterable_from_mmcif( + cls, + path: PathOrBuffer | MmcifWrapper, + id: str | None = None, + is_predicted: bool = False, + keep_source: bool = False, + ): + """Yield every protein chain represented in an mmCIF structure.""" + mmcif = path if isinstance(path, MmcifWrapper) else MmcifWrapper.read(path, id) + for chain in bs.chain_iter(mmcif.structure): + chain = chain[bs.filter_amino_acids(chain) & ~chain.hetero] + if len(chain) == 0: + continue + chain_id = chain.chain_id[0] + entity_id = None + for entity, chains in mmcif.entities.items(): + if chain_id in chains: + entity_id = entity + if entity_id is None: + raise ValueError( + f"Failed to resolve entity identity for mmCIF chain {chain_id!r}." + ) + ( + sequence, + atom_positions, + atom_mask, + residue_index, + insertion_code, + confidence, + _, + ) = chain_to_ndarray(chain, mmcif, chain_id, is_predicted) + if not all(sequence): + raise ValueError("Some residue name was not specified correctly.") + + yield cls( + id=mmcif.id, + sequence=sequence, + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=atom_positions, + atom37_mask=atom_mask, + residue_index=residue_index, + insertion_code=insertion_code, + confidence=confidence, + mmcif=mmcif if keep_source else None, ) - @cached_property - def atoms(self) -> AtomIndexer: - return AtomIndexer(self, property="atom37_positions", dim=-2) + @classmethod + def from_mmcif( + cls, + path: PathOrBuffer | MmcifWrapper, + chain_id: str | None = None, + entity_id: int | None = None, + id: str | None = None, + is_predicted: bool = False, + keep_source: bool = False, + ): + """Return a ProteinChain object from an mmcif file. - @cached_property - def atom_mask(self) -> AtomIndexer: - return AtomIndexer(self, property="atom37_mask", dim=-1) + Args: + path: Uncompressed mmCIF path, buffer, or parsed wrapper. + id: Optional structure identifier when parsing a path or buffer. + is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. + chain_id (str, optional): Select a chain corresponding to (author) chain id. + entity_id (int, optional): Select a chain corresponding to a particular entity. - @cached_property - def atom_array(self) -> bs.AtomArray: - atoms = [] - for res_idx_i, ( - res_name, - res_idx, - ins_code, - positions, - mask, - conf, - ) in enumerate( - zip( - self.sequence, - self.residue_index, - self.insertion_code, - self.atom37_positions, - self.atom37_mask.astype(bool), - self.confidence, - ) - ): - for i, pos in zip(np.where(mask)[0], positions[mask]): - b_factor = ( - self.atom37_confidence[res_idx_i, i] - if self.atom37_confidence is not None - else conf - ) - atom = bs.Atom( - coord=pos, - chain_id="A" if self.chain_id is None else self.chain_id, - res_id=res_idx, - ins_code=ins_code, - res_name=residue_constants.restype_1to3.get(res_name, "UNK"), - hetero=False, - atom_name=residue_constants.atom_types[i], - element=residue_constants.atom_types[i][0], - b_factor=float(b_factor), - ) - atoms.append(atom) - return bs.array(atoms) + If neither `chain_id` nor `entity_id` is specified, defaults to the first entity. + """ + mmcif = path if isinstance(path, MmcifWrapper) else MmcifWrapper.read(path, id) - @cached_property - def residue_index_no_insertions(self) -> np.ndarray: - return self.residue_index + np.cumsum(self.insertion_code != "") + if chain_id is not None and entity_id is not None: + raise ValueError("Pass at most one of chain_id or entity_id.") - @cached_property - def atom_array_no_insertions(self) -> bs.AtomArray: - atoms = [] - for res_idx, (res_name, positions, mask, conf) in enumerate( - zip( - self.sequence, - self.atom37_positions, - self.atom37_mask.astype(bool), - self.confidence, - ) - ): - for i, pos in zip(np.where(mask)[0], positions[mask]): - b_factor = ( - self.atom37_confidence[res_idx, i] - if self.atom37_confidence is not None - else conf - ) - atom = bs.Atom( - coord=pos, - # hard coded to as we currently only support single chain structures - chain_id=CHAIN_ID_CONST, - res_id=res_idx + 1, - res_name=residue_constants.restype_1to3.get(res_name, "UNK"), - hetero=False, - atom_name=residue_constants.atom_types[i], - element=residue_constants.atom_types[i][0], - b_factor=float(b_factor), + # If neither chain_id nor entity_id is specified, default to the first entity + if chain_id is None and entity_id is None: + if not mmcif.entities: + raise ValueError("Structure contains no entities") + entity_id = min(mmcif.entities.keys()) # Pick the first entity by ID + + if entity_id is not None: + if entity_id not in mmcif.entities: + raise ValueError( + f"Structure does not contain entity `{entity_id}`. " + f"Valid entities: {mmcif.entities.keys()}" ) - atoms.append(atom) - return bs.array(atoms) + chains = mmcif.entities[entity_id] - def __getitem__(self, idx: int | list[int] | slice | np.ndarray | torch.Tensor): - if isinstance(idx, int): - idx = [idx] - if isinstance(idx, torch.Tensor): - idx = idx.cpu().numpy() + # Prefer the chain with the most resolved residues; ties preserve source order. + chain_id = max( + chains, + key=lambda chain: _num_non_null_residues(mmcif.seqres_to_structure[chain]), + ) + else: + if chain_id is None: + raise RuntimeError("Failed to resolve an mmCIF chain selection.") + for entity, chains in mmcif.entities.items(): + if chain_id in chains: + entity_id = entity + if entity_id is None: + warnings.warn( + "Failed to detect entity_id from mmcif file, it may be malformed.", + stacklevel=2, + ) - sequence = slice_python_object_as_numpy(self.sequence, idx) - return replace( - self, + atom_array = mmcif.structure + ( + sequence, + atom_positions, + atom_mask, + residue_index, + insertion_code, + confidence, + _, + ) = chain_to_ndarray(atom_array, mmcif, chain_id, is_predicted) + if not all(sequence): + raise ValueError("Some residue name was not specified correctly.") + + return cls( + id=mmcif.id, sequence=sequence, - residue_index=self.residue_index[..., idx], - insertion_code=self.insertion_code[..., idx], - atom37_positions=self.atom37_positions[..., idx, :, :], - atom37_mask=self.atom37_mask[..., idx, :], - confidence=self.confidence[..., idx], - atom37_confidence=self.atom37_confidence[..., idx, :] - if self.atom37_confidence is not None - else None, + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=atom_positions, + atom37_mask=atom_mask.astype(bool), + residue_index=residue_index, + insertion_code=insertion_code, + confidence=confidence, + mmcif=mmcif if keep_source else None, ) - def __len__(self): - return len(self.sequence) + @classmethod + def from_atom37( + cls, + atom37_positions: np.ndarray | torch.Tensor, + *, + id: str | None = None, + sequence: str | None = None, + chain_id: str | None = None, + entity_id: int | None = None, + residue_index: np.ndarray | torch.Tensor | None = None, + insertion_code: np.ndarray | None = None, + confidence: np.ndarray | torch.Tensor | None = None, + ): + if isinstance(atom37_positions, torch.Tensor): + atom37_positions = atom37_positions.cpu().numpy() + if atom37_positions.ndim == 4: + if atom37_positions.shape[0] != 1: + raise ValueError( + "Cannot handle batched inputs, atom37_positions has shape " + f"{atom37_positions.shape}" + ) + atom37_positions = atom37_positions[0] - def cbeta_contacts(self, distance_threshold: float = 8.0) -> np.ndarray: - distance = self.pdist_CB - contacts = (distance < distance_threshold).astype(np.int64) - contacts[np.isnan(distance)] = -1 - np.fill_diagonal(contacts, -1) - return contacts + if not isinstance(atom37_positions, np.ndarray): + raise TypeError("atom37_positions must be a NumPy array or Torch tensor.") + if atom37_positions.ndim != 3 or atom37_positions.shape[1:] != (37, 3): + raise ValueError( + "atom37_positions must have shape (length, 37, 3), got " + f"{atom37_positions.shape}." + ) + seqlen = atom37_positions.shape[0] - def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): - """Dssp works better w/o insertions.""" - f = PDBFile() - if not include_insertions: - f.set_structure(self.atom_array_no_insertions) - else: - f.set_structure(self.atom_array) - f.write(path) + atom_mask = np.isfinite(atom37_positions).all(-1) - def to_pdb_string(self, include_insertions: bool = True) -> str: - buf = io.StringIO() - self.to_pdb(buf, include_insertions=include_insertions) - buf.seek(0) - return buf.read() + if id is None: + id = "" - def to_mmcif(self, path: PathOrBuffer): - f = CIFFile() - set_structure_pdbx(f, self.atom_array, data_block=self.id) + if sequence is None: + sequence = "A" * seqlen - # incantations molstar needs to render pLDDT / confidence onto - # the structure with "alphafold-view" - f.block["ma_qa_metric"] = CIFCategory( - name="ma_qa_metric", - columns={ - "id": CIFColumn(data=CIFData(array=np.array([1, 2]), dtype=np.int64)), - "mode": CIFColumn( - data=CIFData(array=np.array(["global", "local"]), dtype=np.str_) - ), - "name": CIFColumn( - data=CIFData(array=np.array(["pLDDT", "pLDDT"]), dtype=np.str_) - ), - }, - ) - - # table is a duplicate of data already in the atom array, but - # needed by molstar to render pLDDT / confidence - resid_pldd_table = { - # hard coded to as we currently only support single chain structures - "label_asym_id": CIFColumn( - data=CIFData( - array=[CHAIN_ID_CONST] * len(self.residue_index), dtype=np.str_ - ) - ), - "label_comp_id": CIFColumn( - data=CIFData( - array=[ - residue_constants.restype_1to3.get(c, "UNK") - for c in self.sequence - ], - dtype=np.str_, - ) - ), - "label_seq_id": CIFColumn( - data=CIFData(array=self.residue_index, dtype=np.int64) - ), - "ordinal_id": CIFColumn( - data=CIFData(array=self.residue_index, dtype=np.int64) - ), - # hard coded to show these are all local plDDT values - "metric_id": CIFColumn( - data=CIFData(array=["2"] * len(self.residue_index), dtype=np.str_) - ), - "metric_value": CIFColumn( - data=CIFData(array=self.confidence, dtype=np.float32) - ), - # hard coded to show there are the initial version, there are no revisions - "model_id": CIFColumn( - data=CIFData(array=["1"] * len(self.residue_index), dtype=np.str_) - ), - } - f.block["ma_qa_metric_local"] = CIFCategory( - name="ma_qa_metric_local", columns=resid_pldd_table - ) - f.write(path) + if chain_id is None: + chain_id = "A" - def to_mmcif_string(self) -> str: - buf = io.StringIO() - self.to_mmcif(buf) - buf.seek(0) - return buf.read() + if residue_index is None: + residue_index = np.arange(1, seqlen + 1) + elif isinstance(residue_index, torch.Tensor): + residue_index = residue_index.cpu().numpy() + if residue_index.ndim == 2: + if residue_index.shape[0] != 1: + raise ValueError( + "Cannot handle batched inputs, residue_index has shape " + f"{residue_index.shape}" + ) + residue_index = residue_index[0] + if not isinstance(residue_index, np.ndarray): + raise TypeError("residue_index must be a NumPy array or Torch tensor.") - def state_dict(self, backbone_only=False, json_serializable=False): - """This state dict is optimized for storage, so it turns things to fp16 whenever - possible. Note that we also only support int32 residue indices, I'm hoping we don't - need more than 2**32 residues...""" - dct = {k: v for k, v in asdict(self).items() if k not in ["mmcif"]} - if backbone_only: - dct["atom37_mask"][:, 3:] = False - dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] - if dct.get("atom37_confidence") is not None: - dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] - else: - dct.pop("atom37_confidence", None) + if insertion_code is None: + insertion_code = np.array(["" for _ in range(seqlen)]) - for k, v in dct.items(): - if isinstance(v, np.ndarray): - match v.dtype: - case np.int64: - dct[k] = v.astype(np.int32) - case np.float64 | np.float32: - dct[k] = v.astype(np.float16) - case _: - pass - if json_serializable: - dct[k] = v.tolist() - return dct + if confidence is None: + confidence = np.ones(seqlen, dtype=np.float32) + elif isinstance(confidence, torch.Tensor): + confidence = confidence.cpu().numpy() + if confidence.ndim == 2: + if confidence.shape[0] != 1: + raise ValueError( + f"Cannot handle batched inputs, confidence has shape {confidence.shape}" + ) + confidence = confidence[0] + if not isinstance(confidence, np.ndarray): + raise TypeError("confidence must be a NumPy array or Torch tensor.") - def to_blob(self, backbone_only=False) -> bytes: - return brotli.compress(msgpack.dumps(self.state_dict(backbone_only)), quality=5) + return cls( + id=id, + sequence=sequence, # type: ignore + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=atom37_positions, + atom37_mask=atom_mask.astype(bool), + residue_index=residue_index, + insertion_code=insertion_code, + confidence=confidence, + ) @classmethod - def from_open_source(cls, pc: ProteinChain): - return cls(**vars(pc)) + def from_backbone_atom_coordinates( + cls, backbone_atom_coordinates: np.ndarray | torch.Tensor, **kwargs + ): + """Create a ProteinChain from a set of backbone atom coordinates. - @classmethod - def from_state_dict(cls, dct): - # Note: assembly_composition is *supposed* to have string keys. - dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) + This function simply expands the seqlen x 3 x 3 array of backbone atom + coordinates to a seqlen x 37 x 3 array of all atom coordinates, with the padded + positions set to infinity. This allows us to use from_atom37 to create the + appropriate ProteinChain object with the appropriate atom37_mask. - for k, v in dct.items(): - if isinstance(v, list): - dct[k] = np.array(v) + This function passes all kwargs to from_atom37. + """ + if isinstance(backbone_atom_coordinates, torch.Tensor): + backbone_atom_coordinates = backbone_atom_coordinates.cpu().numpy() + if backbone_atom_coordinates.ndim == 4: + if backbone_atom_coordinates.shape[0] != 1: + raise ValueError( + f"Cannot handle batched inputs, backbone_atom_coordinates has " + f"shape {backbone_atom_coordinates.shape}" + ) + backbone_atom_coordinates = backbone_atom_coordinates[0] - atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) - atom37[dct["atom37_mask"]] = dct["atom37_positions"] - dct["atom37_positions"] = atom37 - if "atom37_confidence" in dct: - atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) - atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] - dct["atom37_confidence"] = atom37_conf - dct = { - k: ( - v.astype(np.float32) - if k in ["atom37_positions", "confidence", "atom37_confidence"] - else v + if not isinstance(backbone_atom_coordinates, np.ndarray): + raise TypeError( + "backbone_atom_coordinates must be a NumPy array or Torch tensor." ) - for k, v in dct.items() - if not (k == "atom37_confidence" and v is None) - } - return cls(**dct, mmcif=None) - - @classmethod - def from_blob(cls, input: Path | str | io.BytesIO | bytes): - """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory - of chains from 52G/1M chains to 20G/1M chains, I think this is a good first - shot at compressing and dumping chains to disk. I'm sure there's better ways.""" - match input: - case Path() | str(): - bytes = Path(input).read_bytes() - case io.BytesIO(): - bytes = input.getvalue() - case _: - bytes = input - return cls.from_state_dict(msgpack.loads(brotli.decompress(bytes))) - - def sasa(self, by_residue: bool = True): - arr = self.atom_array_no_insertions - sasa_per_atom = bs.sasa(arr) # type: ignore - if by_residue: - # Sum per-atom SASA into residue "bins", with np.bincount. - assert arr.res_id is not None - # NOTE(rverkuil): arr.res_id is 1-indexed, but np.bincount returns a sum for bin 0, so we strip. - # NOTE(aderry): We compute only for residues with coordinates, return NaN otherwise. - num_trailing_residues = len(self) - arr.res_id.max() - sasa_per_residue = np.concatenate( - [ - np.bincount(arr.res_id, weights=sasa_per_atom)[1:], - np.zeros(num_trailing_residues), - ] + if backbone_atom_coordinates.ndim != 3 or backbone_atom_coordinates.shape[-2:] != ( + 3, + 3, + ): + raise ValueError( + "backbone_atom_coordinates must have shape (length, 3, 3), got " + f"{backbone_atom_coordinates.shape}." ) - sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan - assert len(sasa_per_residue) == len(self) - return sasa_per_residue - return sasa_per_atom - - def sap_score(self, aggregation: str = "atom") -> np.ndarray: - """Computes per-atom SAP score. - Can optionally aggregate by residue (by averaging over atoms. NOTE: this returns values only for residues that have coordinates!) - or full-protein (sum of SAP score for atoms with SAP > 0, as in Lauer et al. 2011).""" - sap_radius = 5.0 - arr = self.atom_array_no_insertions - - # asserts to avoid type errors - assert arr.res_id is not None - assert arr.res_name is not None - assert arr.atom_name is not None - assert arr.coord is not None - - # compute SASA and residue-specific properties - sasa_per_atom = self.sasa(by_residue=False) - resid_to_resname = dict(zip(arr.res_id, arr.res_name)) - max_side_chain_asa = np.full(len(self), np.nan) - res_hydrophobicity = np.full(len(self), np.nan) - resolved_res_mask = self.atom37_mask.any(-1) - num_trailing_residues = len(self) - arr.res_id.max() - - max_side_chain_asa[resolved_res_mask] = np.array( - [ - residue_constants.side_chain_asa[resid_to_resname[i]] - for i in np.unique(arr.res_id) - ] - ) - res_hydrophobicity[resolved_res_mask] = np.array( - [ - residue_constants.hydrophobicity[resid_to_resname[i]] - for i in np.unique(arr.res_id) - ] + atom37_positions = np.full( + (backbone_atom_coordinates.shape[0], 37, 3), + np.inf, + dtype=backbone_atom_coordinates.dtype, ) - assert len(max_side_chain_asa) == len(self) - assert len(res_hydrophobicity) == len(self) - - # compute SAP score - is_side_chain = ~bs.filter_peptide_backbone(arr) - sasa_per_atom[is_side_chain] = 0 - kdtree = KDTree(arr.coord) - neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0) - sap_by_atom = np.zeros_like(sasa_per_atom) - for i, nn_list in enumerate(neighbors): - saa_nn = np.zeros_like(sasa_per_atom) - saa_nn[nn_list] = sasa_per_atom[nn_list] - sasa_within_r = np.concatenate( - [ - np.bincount(arr.res_id, weights=saa_nn)[1:], - np.zeros(num_trailing_residues), - ] - ) - sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity) - sap_by_atom[i] = sap - - match aggregation: - case "atom": - return sap_by_atom - case "residue": - sap_by_residue = np.concatenate( - [ - np.bincount(arr.res_id, weights=sap_by_atom)[1:], - np.zeros(num_trailing_residues), - ] - ) / ( - np.concatenate( - [np.bincount(arr.res_id)[1:], np.zeros(num_trailing_residues)] - ) - + 1e-8 - ) - sap_by_residue[~resolved_res_mask] = np.nan - assert len(sap_by_residue) == len(self) - return sap_by_residue - case "protein": - return sum(sap_by_atom[sap_by_atom > 0]) # pyright: ignore[reportReturnType] - case _: - raise ValueError( - f"Invalid aggregation method: {aggregation}. Must be one of 'atom', 'residue', or 'protein'" - ) - - def globularity(self) -> float: - # Computes globularity using total volumes divided by MVEE. - # We make the simplifying approximation that atoms never overlap. - # The globularity is only computed where structure exists. - # Besides the approximation above, this is inspired by: - - # https://www.mdpi.com/2073-4352/11/12/1539 - # NOTE(@zeming): due to the approximation we make here, that atoms never overlap, you might get >1 globularity - mask = self.atom37_mask.any(-1) - points = self.atom37_positions[self.atom37_mask] - sequence = [aa for aa, m in zip(self.sequence, mask) if m] # type: ignore - A, _ = self._mvee(points, tol=1e-3) - mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A))) - volume = sum(residue_constants.amino_acid_volumes[x] for x in sequence) - ratio = volume / mvee_volume - - # The paper says you must compare the ellipsoidal profile with T, a measurement of - # how elongated the ellipsoid is. We want a single number, so we multiply by 1/2T, so - # that value is normalized between 0-1 - eigenvalues = np.linalg.eigvals(A) - R = 1 / np.sqrt(eigenvalues) - # ellipsoid radii length triangle inequality coefficient - T = max(R[0] / (R[1] + R[2]), R[1] / (R[0] + R[2]), R[2] / (R[0] + R[1])) - elongation_metric = 1 / max(T, 1) - return ratio * elongation_metric - - @staticmethod - def _mvee(P: np.ndarray, tol, max_iter=10000): - # Finds minimum volume enclosing ellipsoid of a set of points. - # Returns A, c where the ellipse is defined as: - # (x-c).T @ A @ (x-c) = 1 - hull = ConvexHull(P) - P = P[hull.vertices] - P = P.T + atom37_positions[:, :3, :] = backbone_atom_coordinates - # Data points - d, N = P.shape - Q = np.zeros((d + 1, N)) - Q[:d, :] = P[:d, :N] - Q[d, :] = np.ones((1, N)) + return cls.from_atom37(atom37_positions=atom37_positions, **kwargs) - # Initializations - count = 1 - err = 1.0 - u = np.full((N, 1), 1 / N) # 1st iteration + @classmethod + def from_pdb( + cls, + path: PathOrBuffer, + chain_id: str = "detect", + id: str | None = None, + is_predicted: bool = False, + ) -> ProteinChain: + """Return a ProteinChain object from an pdb file. NOTE: prefer mmcif for rcsb PDB files. + This function is mostly to interface with old PDB files and predicted structures - + it will not fill out the entity id correctly - # Khachiyan Algorithm - for i in range(max_iter): - X = Q.dot(np.diag(u.squeeze())) @ Q.T - M = np.diag(Q.T @ np.linalg.inv(X) @ Q) - maximum, j = np.max(M), np.argmax(M) - step_size = (maximum - d - 1) / ((d + 1) * (maximum - 1)) - new_u = (1 - step_size) * u - new_u[j] += step_size - count += 1 - err = np.linalg.norm(new_u - u) - u = new_u - if err < tol: - break + Args: + path: PDB path or text buffer. + id: Optional structure identifier. + is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. + chain_id: Author chain identifier. ``"detect"`` selects the first chain. + """ + + if id is not None: + file_id = id else: - raise ValueError("MVEE did not converge") + match path: + case Path() | str(): + file_id = Path(path).with_suffix("").name + case _: + file_id = "null" - d = P.shape[0] # Fixed: use P.shape[0] instead of P.shape - U = np.diag(u.squeeze()) + atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"]) + if len(atom_array) == 0: + raise ValueError("PDB contains no atoms.") + if chain_id == "detect": + chain_id = atom_array.chain_id[0] + atom_array = atom_array[ + bs.filter_amino_acids(atom_array) + & ~atom_array.hetero + & (atom_array.chain_id == chain_id) + ] + if len(atom_array) == 0: + raise ValueError(f"PDB contains no amino-acid atoms for chain {chain_id!r}.") - # The A matrix for the ellipse - A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T) + entity_id = 1 # Not supplied in PDBfiles - # Center of the ellipse - c = P @ u + sequence = "".join( + residue_constants.restype_3to1.get(monomer[0].res_name, "X") + for monomer in bs.residue_iter(atom_array) + ) + num_res = len(sequence) - return A, c + atom_positions = np.full( + [num_res, residue_constants.atom_type_num, 3], np.nan, dtype=np.float32 + ) + atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool) + residue_index = np.full([num_res], -1, dtype=np.int64) + insertion_code = np.full([num_res], "", dtype=" ProteinChain: + return cls( + id=data["id"], + chain_id=data["chain_id"], + entity_id=data["entity_id"], + sequence=data["sequence"], + residue_index=data["residue_index"], + insertion_code=np.asarray(data["insertion_code"]), + atom37_positions=data["atom37_positions"], + atom37_mask=data["atom37_mask"].astype(bool), + confidence=data["confidence"], + mmcif=None, ) - avg_rmsd_neg = aligner.rmsd - return min(avg_rmsd, avg_rmsd_neg) + @classmethod + def from_rcsb( + cls, + pdb_id: str, + chain_id: str | None = None, + entity_id: int | None = None, + keep_source: bool = False, + ) -> ProteinChain: + f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore + return cls.from_mmcif( + f, + id=pdb_id, + chain_id=chain_id, + entity_id=entity_id, + keep_source=keep_source, + is_predicted=False, + ) - def lddt_ca( - self, - native: ProteinChain, - mobile_inds: list[int] | np.ndarray | None = None, - target_inds: list[int] | np.ndarray | None = None, - **kwargs, - ) -> float | np.ndarray: - """Compute the LDDT between this protein chain and another. NOTE: LDDT IS NOT SYMMETRIC. - The call should always be prediction.lddt_ca(native). + @classmethod + def from_atomarray( + cls, atom_array: bs.AtomArray, id: str | None = None, is_predicted: bool = False + ) -> ProteinChain: + """A simple converter from bs.AtomArray -> ProteinChain. + Uses PDB file format as intermediate.""" + atom_array = atom_array.copy() + atom_array.box = None # remove surrounding box, from_pdb won't handle this + pdb_file = PDBFile() # pyright: ignore + pdb_file.set_structure(atom_array) - Arguments: - native (ProteinChain): The ground truth protein chain - mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices - target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices + buf = io.StringIO() + pdb_file.write(buf) + buf.seek(0) + return cls.from_pdb(buf, id=id, is_predicted=is_predicted) - Returns: - float | np.ndarray: The LDDT score between the two protein chains, either - a single float or per-residue LDDT scores if `per_residue` is True. - """ - lddt = compute_lddt_ca( - torch.tensor(self.atom37_positions[mobile_inds]).unsqueeze(0), - torch.tensor(native.atom37_positions[target_inds]).unsqueeze(0), - torch.tensor(native.atom37_mask[mobile_inds]).unsqueeze(0), - **kwargs, - ) - return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten() + # Object invariants and atom views + def __post_init__(self): + if not isinstance(self.id, str): + raise TypeError("id must be a string.") + if not isinstance(self.sequence, str): + raise TypeError("sequence must be a string.") + if not isinstance(self.chain_id, str) or not self.chain_id: + raise ValueError("chain_id must be a non-empty string.") + if self.entity_id is not None and ( + not isinstance(self.entity_id, int) or isinstance(self.entity_id, bool) + ): + raise TypeError("entity_id must be an integer or None.") + sequence_length = len(self.sequence) + aligned = { + "atom37_positions": self.atom37_positions, + "atom37_mask": self.atom37_mask, + "residue_index": self.residue_index, + "insertion_code": self.insertion_code, + "confidence": self.confidence, + } + for name, values in aligned.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != sequence_length: + raise ValueError( + f"{name} shape {values.shape} does not align with " + f"sequence length {sequence_length}." + ) + if self.atom37_positions.shape != (sequence_length, 37, 3): + raise ValueError( + "atom37_positions must have shape " + f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}." + ) + if self.atom37_mask.shape != (sequence_length, 37): + raise ValueError( + "atom37_mask must have shape " + f"({sequence_length}, 37), got {self.atom37_mask.shape}." + ) + if self.atom37_mask.dtype != bool: + raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.") + if not np.issubdtype(self.atom37_positions.dtype, np.number): + raise TypeError("atom37_positions must use a numeric dtype.") + if not np.issubdtype(self.residue_index.dtype, np.integer): + raise TypeError("residue_index must use an integer dtype.") + if self.insertion_code.dtype.kind not in {"U", "S", "O"}: + raise TypeError("insertion_code must use a string-compatible dtype.") + if any(not isinstance(value, str) for value in self.insertion_code.tolist()): + raise TypeError("insertion_code must contain only strings.") + for name, values in ( + ("residue_index", self.residue_index), + ("insertion_code", self.insertion_code), + ("confidence", self.confidence), + ): + if values.shape != (sequence_length,): + raise ValueError( + f"{name} must have shape ({sequence_length},), got {values.shape}." + ) + if not np.issubdtype(self.confidence.dtype, np.number): + raise TypeError("confidence must use a numeric dtype.") + atom37_confidence = self.atom37_confidence + if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray): + raise TypeError("atom37_confidence must be a NumPy array when provided.") + if ( + isinstance(atom37_confidence, np.ndarray) + and atom37_confidence.shape != self.atom37_mask.shape + ): + raise ValueError( + "atom37_confidence shape must match atom37_mask: " + f"{atom37_confidence.shape} != {self.atom37_mask.shape}." + ) + if isinstance(atom37_confidence, np.ndarray) and not np.issubdtype( + atom37_confidence.dtype, np.number + ): + raise TypeError("atom37_confidence must use a numeric dtype.") - def gdt_ts( - self, - target: ProteinChain, - mobile_inds: list[int] | np.ndarray | None = None, - target_inds: list[int] | np.ndarray | None = None, - **kwargs, - ) -> float | np.ndarray: - """Compute the GDT_TS between this protein chain and another. + @cached_property + def atoms(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_positions", dim=-2) - Arguments: - target (ProteinChain): The other protein chain to compare to. - mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices - target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices + @cached_property + def atom_mask(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_mask", dim=-1) - Returns: - float: The GDT_TS score between the two protein chains. - """ - gdt_ts = compute_gdt_ts( - mobile=torch.tensor( - index_by_atom_name(self.atom37_positions[mobile_inds], "CA"), - dtype=torch.float32, - ).unsqueeze(0), - target=torch.tensor( - index_by_atom_name(target.atom37_positions[target_inds], "CA"), - dtype=torch.float32, - ).unsqueeze(0), - atom_exists_mask=torch.tensor( - index_by_atom_name(self.atom37_mask[mobile_inds], "CA", dim=-1) - & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1) - ).unsqueeze(0), - **kwargs, - ) - return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten() + @cached_property + def atom_array(self) -> bs.AtomArray: + atoms = [] + for res_idx_i, ( + res_name, + res_idx, + ins_code, + positions, + mask, + conf, + ) in enumerate( + zip( + self.sequence, + self.residue_index, + self.insertion_code, + self.atom37_positions, + self.atom37_mask.astype(bool), + self.confidence, + strict=False, + ) + ): + for i, pos in zip(np.where(mask)[0], positions[mask], strict=False): + b_factor = ( + self.atom37_confidence[res_idx_i, i] + if self.atom37_confidence is not None + else conf + ) + atom = bs.Atom( + coord=pos, + chain_id="A" if self.chain_id is None else self.chain_id, + res_id=res_idx, + ins_code=ins_code, + res_name=residue_constants.restype_1to3.get(res_name, "UNK"), + hetero=False, + atom_name=residue_constants.atom_types[i], + element=residue_constants.atom_types[i][0], + b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE, + occupancy=1.0, + ) + atoms.append(atom) + return bs.array(atoms) - @classmethod - def chain_iterable_from_mmcif( - cls, - path: PathOrBuffer | MmcifWrapper, - id: str | None = None, - is_predicted: bool = False, - keep_source: bool = False, - ): - """Return a list[ProteinChain] object from an mmcif file, a iterable list of all protein chain - from an mmcif file + # Coordinate transformations and dataset adapters + def get_normalization_frame(self) -> Affine3D: + """Given a set of coordinates, compute a single frame. + The frame is built from the mean N, C-alpha, and C coordinates with + Gram-Schmidt orthogonalization. Its origin is the mean C-alpha position. + + Returns: + Affine3D: [] tensor of Affine3D frame """ - if isinstance(path, MmcifWrapper): - mmcif = path - else: - mmcif = MmcifWrapper.read(path, id) - for chain in bs.chain_iter(mmcif.structure): - chain = chain[bs.filter_amino_acids(chain) & ~chain.hetero] - if len(chain) == 0: - continue - chain_id = chain.chain_id[0] - entity_id = None - for entity, chains in mmcif.entities.items(): - if chain_id in chains: - entity_id = entity - assert entity_id is not None - ( - sequence, - atom_positions, - atom_mask, - residue_index, - insertion_code, - confidence, - _, - ) = chain_to_ndarray(chain, mmcif, chain_id, is_predicted) - assert all(sequence), "Some residue name was not specified correctly" + coords = torch.from_numpy(self.atom37_positions) + frame = get_protein_normalization_frame(coords) - yield cls( - id=mmcif.id, - sequence=sequence, - chain_id=chain_id, - entity_id=entity_id, - atom37_positions=atom_positions, - atom37_mask=atom_mask, - residue_index=residue_index, - insertion_code=insertion_code, - confidence=confidence, - mmcif=mmcif if keep_source else None, - ) + return frame - @classmethod - def from_mmcif( - cls, - path: PathOrBuffer | MmcifWrapper, - chain_id: str | None = None, - entity_id: int | None = None, - id: str | None = None, - is_predicted: bool = False, - keep_source: bool = False, - ): - """Return a ProteinChain object from an mmcif file. + def apply_frame(self, frame: Affine3D) -> ProteinChain: + """Given a frame, apply the frame to the protein's coordinates. Args: - path (str | Path | io.TextIO): Path or buffer to read mmcif file from. Should be uncompressed. - id (str, optional): String identifier to assign to structure. Will attempt to infer otherwise. - is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. - chain_id (str, optional): Select a chain corresponding to (author) chain id. - entity_id (int, optional): Select a chain corresponding to a particular entity. + frame (Affine3D): [] tensor of Affine3D frame - If neither `chain_id` nor `entity_id` is specified, defaults to the first entity. + Returns: + ProteinChain: Transformed protein chain """ - if isinstance(path, MmcifWrapper): - mmcif = path - else: - mmcif = MmcifWrapper.read(path, id) + coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype) + coords = apply_frame_to_coords(coords, frame) + atom37_positions = coords.numpy() + return replace(self, atom37_positions=atom37_positions) - # If neither chain_id nor entity_id is specified, default to the first entity - if chain_id is None and entity_id is None: - if not mmcif.entities: - raise ValueError("Structure contains no entities") - entity_id = min(mmcif.entities.keys()) # Pick the first entity by ID + def normalize_coordinates(self) -> ProteinChain: + """Normalize the coordinates of the protein chain.""" + return self.apply_frame(self.get_normalization_frame()) - if entity_id is not None: - assert chain_id is None - if entity_id not in mmcif.entities: - raise ValueError( - f"Structure does not contain entity `{entity_id}`. Valid entities: {mmcif.entities.keys()}" - ) - chains = mmcif.entities[entity_id] + def infer_oxygen(self) -> ProteinChain: + """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" + O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze() - # Select the chain id corresponding to the longest chain. If all are equal length, selects the first. - chain_id = max( - chains, - key=lambda chain: _num_non_null_residues( - mmcif.seqres_to_structure[chain] - ), - ) - else: - assert chain_id is not None - for entity, chains in mmcif.entities.items(): - if chain_id in chains: - entity_id = entity - if entity_id is None: - warnings.warn( - "Failed to detect entity_id from mmcif file, it may be malformed." - ) + O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) + N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) + N = torch.roll(N, -3) + N[..., -1, :] = torch.nan - atom_array = mmcif.structure - ( - sequence, - atom_positions, - atom_mask, - residue_index, - insertion_code, - confidence, - _, - ) = chain_to_ndarray(atom_array, mmcif, chain_id, is_predicted) - assert all(sequence), "Some residue name was not specified correctly" + # Get the frame defined by the CA-C-N atom + frames = Affine3D.from_graham_schmidt(CA, C, N) + oxygen_coordinates = frames.apply(O_vector) + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() - return cls( - id=mmcif.id, - sequence=sequence, - chain_id=chain_id, - entity_id=entity_id, - atom37_positions=atom_positions, - atom37_mask=atom_mask.astype(bool), - residue_index=residue_index, - insertion_code=insertion_code, - confidence=confidence, - mmcif=mmcif if keep_source else None, - ) + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[ + O_missing_indices + ].numpy() + atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain - @classmethod - def from_atom37( - cls, - atom37_positions: np.ndarray | torch.Tensor, - *, - id: str | None = None, - sequence: str | None = None, - chain_id: str | None = None, - entity_id: int | None = None, - residue_index: np.ndarray | torch.Tensor | None = None, - insertion_code: np.ndarray | None = None, - confidence: np.ndarray | torch.Tensor | None = None, - ): - if isinstance(atom37_positions, torch.Tensor): - atom37_positions = atom37_positions.cpu().numpy() - if atom37_positions.ndim == 4: - if atom37_positions.shape[0] != 1: - raise ValueError( - f"Cannot handle batched inputs, atom37_positions has shape {atom37_positions.shape}" - ) - atom37_positions = atom37_positions[0] + @cached_property + def inferred_cbeta(self) -> np.ndarray: + """Infer cbeta positions based on N, C, CA.""" + N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) + # See usage in trDesign codebase. + # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 + CB = infer_cb(C, N, CA, 1.522, 1.927, -2.143) + return CB - assert isinstance(atom37_positions, np.ndarray) - seqlen = atom37_positions.shape[0] + def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain: + """Return a new chain with inferred CB atoms at all residues except GLY. - atom_mask = np.isfinite(atom37_positions).all(-1) + Args: + infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine + residues, even though that residue doesn't have one. Default off. - if id is None: - id = "" + NOTE(rverkuil): The reason for having this switch in the first place + is that sometimes we want a (inferred) CB coordinate for every residue, + for example for making a pairwise distance matrix, or doing an RMSD + calculation between two designs for a given structural template, w/ + CB atoms. + """ + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() - if sequence is None: - sequence = "A" * seqlen + inferred_cbeta_positions = self.inferred_cbeta + if not infer_cbeta_for_glycine: + inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan - if chain_id is None: - chain_id = "A" + atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions + atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( + atom37_positions[:, residue_constants.atom_order["CB"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain - if residue_index is None: - residue_index = np.arange(1, seqlen + 1) - elif isinstance(residue_index, torch.Tensor): - residue_index = residue_index.cpu().numpy() - assert isinstance(residue_index, np.ndarray) - if residue_index.ndim == 2: - if residue_index.shape[0] != 1: - raise ValueError( - f"Cannot handle batched inputs, residue_index has shape {residue_index.shape}" - ) - residue_index = residue_index[0] - assert isinstance(residue_index, np.ndarray) + @cached_property + def pdist_CA(self) -> np.ndarray: + CA = self.atoms["CA"] + pdist_CA = squareform(pdist(CA)) + return pdist_CA - if insertion_code is None: - insertion_code = np.array(["" for _ in range(seqlen)]) + @cached_property + def pdist_CB(self) -> np.ndarray: + pdist_CB = squareform(pdist(self.inferred_cbeta)) + return pdist_CB - if confidence is None: - confidence = np.ones(seqlen, dtype=np.float32) - elif isinstance(confidence, torch.Tensor): - confidence = confidence.cpu().numpy() - assert isinstance(confidence, np.ndarray) - if confidence.ndim == 2: - if confidence.shape[0] != 1: - raise ValueError( - f"Cannot handle batched inputs, confidence has shape {confidence.shape}" - ) - confidence = confidence[0] - assert isinstance(confidence, np.ndarray) + @classmethod + def as_complex(cls, chains: Sequence[ProteinChain]): + raise RuntimeError( + ".as_complex() has been deprecated in favor of .concat(). " + ".concat() will eventually be deprecated in favor of ProteinComplex..." + ) + + @classmethod + def concat(cls, chains: Sequence[ProteinChain], use_chainbreak: bool = True): + if not chains: + raise ValueError("chains must contain at least one ProteinChain.") + if any(not isinstance(chain, ProteinChain) for chain in chains): + raise TypeError("chains must contain only ProteinChain instances.") + sep_tokens = { + "residue_index": np.array([-1]), + "insertion_code": np.array([""]), + "atom37_positions": np.full([1, 37, 3], np.inf), + "atom37_mask": np.zeros([1, 37], dtype=bool), + "confidence": np.array([0]), + } + + def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): + if use_chainbreak: + full_array = [] + for array in arrays: + full_array.append(array) + full_array.append(sep) + full_array = full_array[:-1] + return np.concatenate(full_array, 0) + else: + return np.concatenate(arrays, 0) + + array_args: dict[str, np.ndarray] = { + name: join_arrays([getattr(chain, name) for chain in chains], sep) + for name, sep in sep_tokens.items() + } + chain_break = residue_constants.CHAIN_BREAK_TOKEN if use_chainbreak else "" return cls( - id=id, - sequence=sequence, # type: ignore - chain_id=chain_id, - entity_id=entity_id, - atom37_positions=atom37_positions, - atom37_mask=atom_mask.astype(bool), - residue_index=residue_index, - insertion_code=insertion_code, - confidence=confidence, + id=chains[0].id, + sequence=chain_break.join(chain.sequence for chain in chains), + chain_id="A", + entity_id=None, + mmcif=None, + **array_args, ) - @classmethod - def from_backbone_atom_coordinates( - cls, backbone_atom_coordinates: np.ndarray | torch.Tensor, **kwargs - ): - """Create a ProteinChain from a set of backbone atom coordinates. + def find_nonpolymer_contacts(self): + if self.mmcif is None: + raise ValueError( + "find_nonpolymer_contacts requires a chain loaded with keep_source=True." + ) + nonpolymer_and_chain_id_to_array = self.mmcif.non_polymer_coords - This function simply expands the seqlen x 3 x 3 array of backbone atom - coordinates to a seqlen x 37 x 3 array of all atom coordinates, with the padded - positions set to infinity. This allows us to use from_atom37 to create the - appropriate ProteinChain object with the appropriate atom37_mask. + results = [] + for ( + nonpolymer, + _, + ), nonpolymer_array in nonpolymer_and_chain_id_to_array.items(): + if nonpolymer_array.coord is None: + raise ValueError( + f"Non-polymer {nonpolymer.comp_id!r} has no coordinate table." + ) + chain_coords = self.atom37_positions[self.atom37_mask] + distance = cdist(nonpolymer_array.coord, chain_coords) - This function passes all kwargs to from_atom37. - """ - if isinstance(backbone_atom_coordinates, torch.Tensor): - backbone_atom_coordinates = backbone_atom_coordinates.cpu().numpy() - if backbone_atom_coordinates.ndim == 4: - if backbone_atom_coordinates.shape[0] != 1: - raise ValueError( - f"Cannot handle batched inputs, backbone_atom_coordinates has " - f"shape {backbone_atom_coordinates.shape}" - ) - backbone_atom_coordinates = backbone_atom_coordinates[0] + is_contact = distance < 5 + if not is_contact.any(): + continue + contacting_atoms = np.where(is_contact.any(0))[0] + chain_index = np.where(self.atom37_mask)[0] + contacting_residues = np.unique(chain_index[contacting_atoms]) - assert isinstance(backbone_atom_coordinates, np.ndarray) - assert backbone_atom_coordinates.ndim == 3 - assert backbone_atom_coordinates.shape[-2] == 3 - assert backbone_atom_coordinates.shape[-1] == 3 + result = { + "ligand": nonpolymer.name, + "ligand_id": nonpolymer.comp_id, + "contacting_residues": contacting_residues.tolist(), + } + results.append(result) + return results - atom37_positions = np.full( - (backbone_atom_coordinates.shape[0], 37, 3), - np.inf, - dtype=backbone_atom_coordinates.dtype, - ) - atom37_positions[:, :3, :] = backbone_atom_coordinates + def select_residue_indices( + self, indices: list[int | str], ignore_x_mismatch: bool = False + ) -> ProteinChain: + numeric_indices = [idx if isinstance(idx, int) else int(idx[1:]) for idx in indices] + mask = np.isin(self.residue_index, numeric_indices) + new = self[mask] + mismatches = [] + for aa, idx in zip(new.sequence, indices, strict=False): + if isinstance(idx, int): + continue + if aa == "X" and ignore_x_mismatch: + continue + if aa != idx[0]: + mismatches.append((aa, idx)) + if mismatches: + mismatch_str = "; ".join( + f"Position {idx[1:]}, Expected: {idx[0]}, Received: {aa}" for aa, idx in mismatches + ) + raise RuntimeError(mismatch_str) - return cls.from_atom37(atom37_positions=atom37_positions, **kwargs) + return new - @classmethod - def from_pdb( - cls, - path: PathOrBuffer, - chain_id: str = "detect", - id: str | None = None, - is_predicted: bool = False, - ) -> "ProteinChain": - """Return a ProteinChain object from an pdb file. NOTE: prefer mmcif for rcsb PDB files. - This function is mostly to interface with old PDB files and predicted structures - - it will not fill out the entity id correctly + def to_structure_encoder_inputs( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert protein chain to structure encoder inputs. - Args: - path (str | Path | io.TextIO): Path or buffer to read mmcif file from. Should be uncompressed. - id (str, optional): String identifier to assign to structure. Will attempt to infer otherwise. - is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. - chain_id (str, optional): Select a chain corresponding to (author) chain id. "detect" uses the - first detected chain + Returns: + tuple: (coordinates, plddt, residue_index) where: + - coordinates: X with shape (1, l, 37, 3), containing atom positions + - plddt: P with shape (1, l), containing confidence scores + - residue_index: R with shape (1, l), containing residue indices """ + # Convert to tensors and add batch dimension + coordinates = ( + torch.from_numpy(self.atom37_positions).float().unsqueeze(0) + ) # X has shape (1, l, 37, 3). + plddt = torch.from_numpy(self.confidence).float().unsqueeze(0) # P: (1, l) + residue_index = ( + torch.from_numpy(self.residue_index).long().unsqueeze(0) + ) # R has shape (1, l). - if id is not None: - file_id = id - else: - match path: - case Path() | str(): - file_id = Path(path).with_suffix("").name - case _: - file_id = "null" + return coordinates, plddt, residue_index + + # Sequence access, interchange, and compact storage + def __getitem__(self, idx: int | list[int] | slice | np.ndarray | torch.Tensor): + if isinstance(idx, int): + idx = [idx] + if isinstance(idx, torch.Tensor): + idx = idx.cpu().numpy() - atom_array = PDBFile.read(path).get_structure( - model=1, extra_fields=["b_factor"] + sequence = slice_python_object_as_numpy(self.sequence, idx) + return replace( + self, + sequence=sequence, + residue_index=self.residue_index[..., idx], + insertion_code=self.insertion_code[..., idx], + atom37_positions=self.atom37_positions[..., idx, :, :], + atom37_mask=self.atom37_mask[..., idx, :], + confidence=self.confidence[..., idx], + atom37_confidence=self.atom37_confidence[..., idx, :] + if self.atom37_confidence is not None + else None, ) - if chain_id == "detect": - chain_id = atom_array.chain_id[0] - atom_array = atom_array[ - bs.filter_amino_acids(atom_array) - & ~atom_array.hetero - & (atom_array.chain_id == chain_id) - ] - entity_id = 1 # Not supplied in PDBfiles + def __len__(self): + return len(self.sequence) - sequence = "".join( - residue_constants.restype_3to1.get(monomer[0].res_name, "X") - for monomer in bs.residue_iter(atom_array) - ) - num_res = len(sequence) + def cbeta_contacts(self, distance_threshold: float = 8.0) -> np.ndarray: + distance = self.pdist_CB + contacts = (distance < distance_threshold).astype(np.int64) + contacts[np.isnan(distance)] = -1 + np.fill_diagonal(contacts, -1) + return contacts - atom_positions = np.full( - [num_res, residue_constants.atom_type_num, 3], np.nan, dtype=np.float32 - ) - atom_mask = np.full( - [num_res, residue_constants.atom_type_num], False, dtype=bool - ) - residue_index = np.full([num_res], -1, dtype=np.int64) - insertion_code = np.full([num_res], "", dtype=" str: + buf = io.StringIO() + self.to_pdb(buf, include_insertions=include_insertions) + buf.seek(0) + return buf.read() - for i, res in enumerate(bs.residue_iter(atom_array)): - chain = atom_array[atom_array.chain_id == chain_id] - assert isinstance(chain, bs.AtomArray) + def to_mmcif(self, path: PathOrBuffer): + f = CIFFile() + set_structure_pdbx(f, self.atom_array, data_block=self.id) - res_index = res[0].res_id - residue_index[i] = res_index - insertion_code[i] = res[0].ins_code + # incantations molstar needs to render pLDDT / confidence onto + # the structure with "alphafold-view" + f.block["ma_qa_metric"] = CIFCategory( + name="ma_qa_metric", + columns={ + "id": CIFColumn(data=CIFData(array=np.array([1, 2]), dtype=np.int64)), + "mode": CIFColumn(data=CIFData(array=np.array(["global", "local"]), dtype=np.str_)), + "name": CIFColumn(data=CIFData(array=np.array(["pLDDT", "pLDDT"]), dtype=np.str_)), + }, + ) - # Atom level features - for atom in res: - atom_name = atom.atom_name - if atom_name == "SE" and atom.res_name == "MSE": - # Put the coords of the selenium atom in the sulphur column - atom_name = "SD" + # table is a duplicate of data already in the atom array, but + # needed by molstar to render pLDDT / confidence + resid_pldd_table = { + # hard coded to as we currently only support single chain structures + "label_asym_id": CIFColumn( + data=CIFData(array=[CHAIN_ID_CONST] * len(self.residue_index), dtype=np.str_) + ), + "label_comp_id": CIFColumn( + data=CIFData( + array=[residue_constants.restype_1to3.get(c, "UNK") for c in self.sequence], + dtype=np.str_, + ) + ), + "label_seq_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)), + "ordinal_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)), + # hard coded to show these are all local plDDT values + "metric_id": CIFColumn( + data=CIFData(array=["2"] * len(self.residue_index), dtype=np.str_) + ), + "metric_value": CIFColumn( + data=CIFData( + array=self.confidence * PLDDT_B_FACTOR_SCALE, + dtype=np.float32, + ) + ), + # hard coded to show there are the initial version, there are no revisions + "model_id": CIFColumn( + data=CIFData(array=["1"] * len(self.residue_index), dtype=np.str_) + ), + } + f.block["ma_qa_metric_local"] = CIFCategory( + name="ma_qa_metric_local", columns=resid_pldd_table + ) + round_mmcif_columns(f) + f.write(path) - if atom_name in residue_constants.atom_order: - atom_positions[i, residue_constants.atom_order[atom_name]] = ( - atom.coord - ) - atom_mask[i, residue_constants.atom_order[atom_name]] = True - if is_predicted and atom_name == "CA": - confidence[i] = atom.b_factor + def to_mmcif_string(self) -> str: + buf = io.StringIO() + self.to_mmcif(buf) + buf.seek(0) + return buf.read() + + def state_dict(self, backbone_only=False, json_serializable=False): + """This state dict is optimized for storage, so it turns things to fp16 whenever + possible. Note that we also only support int32 residue indices, I'm hoping we don't + need more than 2**32 residues...""" + dct = {k: v for k, v in asdict(self).items() if k not in ["mmcif"]} + if backbone_only: + dct["atom37_mask"][:, 3:] = False + dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] + if dct.get("atom37_confidence") is not None: + dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] + else: + dct.pop("atom37_confidence", None) - assert all(sequence), "Some residue name was not specified correctly" + for k, v in dct.items(): + if isinstance(v, np.ndarray): + match v.dtype: + case np.int64: + dct[k] = v.astype(np.int32) + case np.float64 | np.float32: + dct[k] = v.astype(np.float16) + case _: + pass + if json_serializable: + dct[k] = v.tolist() + return dct - return cls( - id=file_id, - sequence=sequence, - chain_id=chain_id, - entity_id=entity_id, - atom37_positions=atom_positions, - atom37_mask=atom_mask.astype(bool), - residue_index=residue_index, - insertion_code=insertion_code, - confidence=confidence, - mmcif=None, - ) + def to_blob(self, backbone_only=False) -> bytes: + payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode) + return brotli.compress(payload, quality=5) @classmethod - def from_mds(cls, data: dict[str, Any]) -> "ProteinChain": - return cls( - id=data["id"], - chain_id=data["chain_id"], - entity_id=data["entity_id"], - sequence=data["sequence"], - residue_index=data["residue_index"], - insertion_code=np.asarray(data["insertion_code"]), - atom37_positions=data["atom37_positions"], - atom37_mask=data["atom37_mask"].astype(bool), - confidence=data["confidence"], - mmcif=None, - ) + def from_open_source(cls, pc: ProteinChain): + return cls(**vars(pc)) @classmethod - def from_rcsb( - cls, - pdb_id: str, - chain_id: str | None = None, - entity_id: int | None = None, - keep_source: bool = False, - ) -> ProteinChain: - f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore - return cls.from_mmcif( - f, - id=pdb_id, - chain_id=chain_id, - entity_id=entity_id, - keep_source=keep_source, - is_predicted=False, - ) + def from_state_dict(cls, dct): + # Note: assembly_composition is *supposed* to have string keys. + dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) + + for k, v in dct.items(): + if isinstance(v, list): + dct[k] = np.array(v) + + atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) + atom37[dct["atom37_mask"]] = dct["atom37_positions"] + dct["atom37_positions"] = atom37 + if "atom37_confidence" in dct: + atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) + atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] + dct["atom37_confidence"] = atom37_conf + dct = { + k: ( + v.astype(np.float32) + if k in ["atom37_positions", "confidence", "atom37_confidence"] + else v + ) + for k, v in dct.items() + if not (k == "atom37_confidence" and v is None) + } + return cls(**dct, mmcif=None) @classmethod - def from_atomarray( - cls, atom_array: bs.AtomArray, id: str | None = None, is_predicted: bool = False - ) -> "ProteinChain": - """A simple converter from bs.AtomArray -> ProteinChain. - Uses PDB file format as intermediate.""" - atom_array = atom_array.copy() - atom_array.box = None # remove surrounding box, from_pdb won't handle this - pdb_file = PDBFile() # pyright: ignore - pdb_file.set_structure(atom_array) + def from_blob(cls, input: Path | str | io.BytesIO | bytes): + """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory + of chains from 52G/1M chains to 20G/1M chains, I think this is a good first + shot at compressing and dumping chains to disk. I'm sure there's better ways.""" + match input: + case Path() | str(): + bytes = Path(input).read_bytes() + case io.BytesIO(): + bytes = input.getvalue() + case _: + bytes = input + state = msgpack.loads(brotli.decompress(bytes), object_hook=msgpack_numpy.decode) + return cls.from_state_dict(state) - buf = io.StringIO() - pdb_file.write(buf) - buf.seek(0) - return cls.from_pdb(buf, id=id, is_predicted=is_predicted) + # Surface and structural comparison metrics + def sasa(self, by_residue: bool = True): + arr = self.atom_array_no_insertions + if len(arr) == 0: + raise ValueError("SASA requires at least one resolved atom.") + sasa_per_atom = bs.sasa(arr) # type: ignore + if by_residue: + # Sum per-atom SASA into residue "bins", with np.bincount. + if arr.res_id is None: + raise RuntimeError("Biotite AtomArray is missing residue identifiers.") + # Residue IDs are one-indexed, so discard the unused zero bin. + # NOTE(aderry): We compute only for residues with coordinates, return NaN otherwise. + num_trailing_residues = len(self) - arr.res_id.max() + sasa_per_residue = np.concatenate( + [ + np.bincount(arr.res_id, weights=sasa_per_atom)[1:], + np.zeros(num_trailing_residues), + ] + ) + sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan + if len(sasa_per_residue) != len(self): + raise RuntimeError("Residue SASA output does not align with the protein chain.") + return sasa_per_residue + return sasa_per_atom - def get_normalization_frame(self) -> Affine3D: - """Given a set of coordinates, compute a single frame. - Specifically, we compute the average position of the N, CA, and C atoms use those 3 points to construct a frame using the Gram-Schmidt algorithm. The average CA position is used as the origin of the frame. + def sap_score(self, aggregation: str = "atom") -> np.ndarray: + """Compute per-atom spatial aggregation propensity (SAP). - Returns: - Affine3D: [] tensor of Affine3D frame + Residue aggregation averages resolved atoms and omits unresolved residues. + Protein aggregation sums positive atom scores, following Lauer et al. 2011. """ - coords = torch.from_numpy(self.atom37_positions) - frame = get_protein_normalization_frame(coords) + sap_radius = 5.0 + arr = self.atom_array_no_insertions + if len(arr) == 0: + raise ValueError("SAP requires at least one resolved atom.") - return frame + for name in ("res_id", "res_name", "atom_name", "coord"): + if getattr(arr, name) is None: + raise RuntimeError(f"Biotite AtomArray is missing required {name!r} data.") - def apply_frame(self, frame: Affine3D) -> ProteinChain: - """Given a frame, apply the frame to the protein's coordinates. + # compute SASA and residue-specific properties + sasa_per_atom = self.sasa(by_residue=False) + resid_to_resname = dict(zip(arr.res_id, arr.res_name, strict=False)) - Args: - frame (Affine3D): [] tensor of Affine3D frame + max_side_chain_asa = np.full(len(self), np.nan) + res_hydrophobicity = np.full(len(self), np.nan) + resolved_res_mask = self.atom37_mask.any(-1) + num_trailing_residues = len(self) - arr.res_id.max() - Returns: - ProteinChain: Transformed protein chain - """ - coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype) - coords = apply_frame_to_coords(coords, frame) - atom37_positions = coords.numpy() - return replace(self, atom37_positions=atom37_positions) + max_side_chain_asa[resolved_res_mask] = np.array( + [residue_constants.side_chain_asa[resid_to_resname[i]] for i in np.unique(arr.res_id)] + ) + res_hydrophobicity[resolved_res_mask] = np.array( + [residue_constants.hydrophobicity[resid_to_resname[i]] for i in np.unique(arr.res_id)] + ) - def normalize_coordinates(self) -> ProteinChain: - """Normalize the coordinates of the protein chain.""" - return self.apply_frame(self.get_normalization_frame()) + # compute SAP score + is_side_chain = ~bs.filter_peptide_backbone(arr) + sasa_per_atom[is_side_chain] = 0 + kdtree = KDTree(arr.coord) + neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0) + sap_by_atom = np.zeros_like(sasa_per_atom) + for i, nn_list in enumerate(neighbors): + saa_nn = np.zeros_like(sasa_per_atom) + saa_nn[nn_list] = sasa_per_atom[nn_list] + sasa_within_r = np.concatenate( + [ + np.bincount(arr.res_id, weights=saa_nn)[1:], + np.zeros(num_trailing_residues), + ] + ) + sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity) + sap_by_atom[i] = sap - def infer_oxygen(self) -> ProteinChain: - """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" - O_missing_indices = np.argwhere( - ~np.isfinite(self.atoms["O"]).all(axis=1) - ).squeeze() + match aggregation: + case "atom": + return sap_by_atom + case "residue": + sap_by_residue = np.concatenate( + [ + np.bincount(arr.res_id, weights=sap_by_atom)[1:], + np.zeros(num_trailing_residues), + ] + ) / ( + np.concatenate([np.bincount(arr.res_id)[1:], np.zeros(num_trailing_residues)]) + + 1e-8 + ) + sap_by_residue[~resolved_res_mask] = np.nan + if len(sap_by_residue) != len(self): + raise RuntimeError("Residue SAP output does not align with the protein chain.") + return sap_by_residue + case "protein": + return sum(sap_by_atom[sap_by_atom > 0]) # pyright: ignore[reportReturnType] + case _: + raise ValueError( + f"Invalid aggregation method: {aggregation}. Must be one of " + "'atom', 'residue', or 'protein'" + ) - O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) - N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) - N = torch.roll(N, -3) - N[..., -1, :] = torch.nan + def globularity(self) -> float: + # Computes globularity using total volumes divided by MVEE. + # We make the simplifying approximation that atoms never overlap. + # The globularity is only computed where structure exists. + # Besides the approximation above, this is inspired by: - # Get the frame defined by the CA-C-N atom - frames = Affine3D.from_graham_schmidt(CA, C, N) - O = frames.apply(O_vector) - atom37_positions = self.atom37_positions.copy() - atom37_mask = self.atom37_mask.copy() + # https://www.mdpi.com/2073-4352/11/12/1539 + # The non-overlapping-atom approximation can produce globularity above one. + mask = self.atom37_mask.any(-1) + points = self.atom37_positions[self.atom37_mask] + sequence = [aa for aa, m in zip(self.sequence, mask, strict=False) if m] # type: ignore + A, _ = self._mvee(points, tol=1e-3) + mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A))) + volume = sum(residue_constants.amino_acid_volumes[x] for x in sequence) + ratio = volume / mvee_volume - atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = O[ - O_missing_indices - ].numpy() - atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( - atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] - ).any(-1) - new_chain = replace( - self, atom37_positions=atom37_positions, atom37_mask=atom37_mask - ) - return new_chain + # The paper compares the ellipsoidal profile with scalar t, a measurement + # of elongation. We want a single number, so we multiply by 1/(2t), so + # that value is normalized between 0-1 + eigenvalues = np.linalg.eigvals(A) + R = 1 / np.sqrt(eigenvalues) + # ellipsoid radii length triangle inequality coefficient + t = max(R[0] / (R[1] + R[2]), R[1] / (R[0] + R[2]), R[2] / (R[0] + R[1])) + elongation_metric = 1 / max(t, 1) + return ratio * elongation_metric - @cached_property - def inferred_cbeta(self) -> np.ndarray: - """Infer cbeta positions based on N, C, CA.""" - N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) - # See usage in trDesign codebase. - # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 - CB = infer_CB(C, N, CA, 1.522, 1.927, -2.143) - return CB + @staticmethod + def _mvee(P: np.ndarray, tol, max_iter=10000): + # Finds minimum volume enclosing ellipsoid of a set of points. + # Returns A, c where the ellipse is defined as: + # (x-c).T @ A @ (x-c) = 1 + hull = ConvexHull(P) + P = P[hull.vertices] + P = P.T - def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain: - """Return a new chain with inferred CB atoms at all residues except GLY. + # Data points + d, n = P.shape + Q = np.zeros((d + 1, n)) + Q[:d, :] = P[:d, :n] + Q[d, :] = np.ones((1, n)) + + # Initializations + count = 1 + err = 1.0 + u = np.full((n, 1), 1 / n) # First iteration. + + # Khachiyan Algorithm + for _ in range(max_iter): + X = Q.dot(np.diag(u.squeeze())) @ Q.T + M = np.diag(Q.T @ np.linalg.inv(X) @ Q) + maximum, j = np.max(M), np.argmax(M) + step_size = (maximum - d - 1) / ((d + 1) * (maximum - 1)) + new_u = (1 - step_size) * u + new_u[j] += step_size + count += 1 + err = np.linalg.norm(new_u - u) + u = new_u + if err < tol: + break + else: + raise ValueError("MVEE did not converge") - Args: - infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine - residues, even though that residue doesn't have one. Default off. + d = P.shape[0] # Fixed: use P.shape[0] instead of P.shape + U = np.diag(u.squeeze()) - NOTE(rverkuil): The reason for having this switch in the first place - is that sometimes we want a (inferred) CB coordinate for every residue, - for example for making a pairwise distance matrix, or doing an RMSD - calculation between two designs for a given structural template, w/ - CB atoms. - """ - atom37_positions = self.atom37_positions.copy() - atom37_mask = self.atom37_mask.copy() + # The A matrix for the ellipse + A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T) - inferred_cbeta_positions = self.inferred_cbeta - if not infer_cbeta_for_glycine: - inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan + # Center of the ellipse + c = P @ u - atom37_positions[:, residue_constants.atom_order["CB"]] = ( - inferred_cbeta_positions - ) - atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( - atom37_positions[:, residue_constants.atom_order["CB"]] - ).any(-1) - new_chain = replace( - self, atom37_positions=atom37_positions, atom37_mask=atom37_mask - ) - return new_chain + return A, c - @cached_property - def pdist_CA(self) -> np.ndarray: - CA = self.atoms["CA"] - pdist_CA = squareform(pdist(CA)) - return pdist_CA + def radius_of_gyration(self): + arr = self.atom_array_no_insertions + return bs.gyration_radius(arr) - @cached_property - def pdist_CB(self) -> np.ndarray: - pdist_CB = squareform(pdist(self.inferred_cbeta)) - return pdist_CB + def align( + self, + target: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + only_use_backbone: bool = False, + ): + """ + Aligns the current protein to the provided target. - @classmethod - def as_complex(cls, chains: Sequence[ProteinChain]): - raise RuntimeError( - ".as_complex() has been deprecated in favor of .concat(). " - ".concat() will eventually be deprecated in favor of ProteinComplex..." + Args: + target (ProteinChain): The target protein to align to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_use_backbone (bool, optional): If True, only align the backbone atoms. + """ + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_use_backbone, ) - @classmethod - def concat(cls, chains: Sequence[ProteinChain], use_chainbreak: bool = True): - sep_tokens = { - "residue_index": np.array([-1]), - "insertion_code": np.array([""]), - "atom37_positions": np.full([1, 37, 3], np.inf), - "atom37_mask": np.zeros([1, 37], dtype=bool), - "confidence": np.array([0]), - } - - def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): - if use_chainbreak: - full_array = [] - for array in arrays: - full_array.append(array) - full_array.append(sep) - full_array = full_array[:-1] - return np.concatenate(full_array, 0) - else: - return np.concatenate(arrays, 0) + return aligner.apply(self) - array_args: dict[str, np.ndarray] = { - name: join_arrays([getattr(chain, name) for chain in chains], sep) - for name, sep in sep_tokens.items() - } + def rmsd( + self, + target: ProteinChain, + also_check_reflection: bool = False, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + only_compute_backbone_rmsd: bool = False, + ): + """ + Compute the RMSD between this protein chain and another. - chain_break = residue_constants.CHAIN_BREAK_TOKEN if use_chainbreak else "" - return cls( - id=chains[0].id, - sequence=chain_break.join(chain.sequence for chain in chains), - chain_id="A", - entity_id=None, - mmcif=None, - **array_args, + Args: + target (ProteinChain): The target (other) protein chain to compare to. + also_check_reflection: Compare the reflected mobile coordinates too. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_compute_backbone_rmsd: Restrict the score to backbone atoms. + """ + if isinstance(target, bs.AtomArray): + raise ValueError( + "Support for bs.AtomArray removed, use ProteinChain.from_atomarry for ProteinChain." + ) + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, ) + avg_rmsd = aligner.rmsd - def find_nonpolymer_contacts(self): - assert self.mmcif is not None - nonpolymer_and_chain_id_to_array = self.mmcif.non_polymer_coords + if not also_check_reflection: + return avg_rmsd - results = [] - for ( - nonpolymer, - _, - ), nonpolymer_array in nonpolymer_and_chain_id_to_array.items(): - assert nonpolymer_array.coord is not None - chain_coords = self.atom37_positions[self.atom37_mask] - distance = cdist(nonpolymer_array.coord, chain_coords) + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, + use_reflection=True, + ) + avg_rmsd_neg = aligner.rmsd - is_contact = distance < 5 - if not is_contact.any(): - continue - contacting_atoms = np.where(is_contact.any(0))[0] - chain_index = np.where(self.atom37_mask)[0] - contacting_residues = np.unique(chain_index[contacting_atoms]) + return min(avg_rmsd, avg_rmsd_neg) - result = { - "ligand": nonpolymer.name, - "ligand_id": nonpolymer.comp_id, - "contacting_residues": contacting_residues.tolist(), - } - results.append(result) - return results + def lddt_ca( + self, + native: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + **kwargs, + ) -> float | np.ndarray: + """Compute the LDDT between this protein chain and another. NOTE: LDDT IS NOT SYMMETRIC. + The call should always be prediction.lddt_ca(native). - def select_residue_indices( - self, indices: list[int | str], ignore_x_mismatch: bool = False - ) -> ProteinChain: - numeric_indices = [ - idx if isinstance(idx, int) else int(idx[1:]) for idx in indices - ] - mask = np.isin(self.residue_index, numeric_indices) - new = self[mask] - mismatches = [] - for aa, idx in zip(new.sequence, indices): - if isinstance(idx, int): - continue - if aa == "X" and ignore_x_mismatch: - continue - if aa != idx[0]: - mismatches.append((aa, idx)) - if mismatches: - mismatch_str = "; ".join( - f"Position {idx[1:]}, Expected: {idx[0]}, Received: {aa}" - for aa, idx in mismatches - ) - raise RuntimeError(mismatch_str) + Arguments: + native (ProteinChain): The ground truth protein chain + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. - return new + Returns: + float | np.ndarray: The LDDT score between the two protein chains, either + a single float or per-residue LDDT scores if `per_residue` is True. + """ + lddt = compute_lddt_ca( + torch.tensor(self.atom37_positions[mobile_inds]).unsqueeze(0), + torch.tensor(native.atom37_positions[target_inds]).unsqueeze(0), + torch.tensor(native.atom37_mask[mobile_inds]).unsqueeze(0), + **kwargs, + ) + return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten() - def to_structure_encoder_inputs( + def gdt_ts( self, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert protein chain to structure encoder inputs. + target: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + **kwargs, + ) -> float | np.ndarray: + """Compute the GDT_TS between this protein chain and another. + + Arguments: + target (ProteinChain): The other protein chain to compare to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. Returns: - tuple: (coordinates, plddt, residue_index) where: - - coordinates: (1, L, 37, 3) tensor of atom positions - - plddt: (1, L) tensor of confidence scores - - residue_index: (1, L) tensor of residue indices + float: The GDT_TS score between the two protein chains. """ - # Convert to tensors and add batch dimension - coordinates = ( - torch.from_numpy(self.atom37_positions).float().unsqueeze(0) - ) # (1, L, 37, 3) - plddt = torch.from_numpy(self.confidence).float().unsqueeze(0) # (1, L) - residue_index = ( - torch.from_numpy(self.residue_index).long().unsqueeze(0) - ) # (1, L) + gdt_ts = compute_gdt_ts( + mobile=torch.tensor( + index_by_atom_name(self.atom37_positions[mobile_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + target=torch.tensor( + index_by_atom_name(target.atom37_positions[target_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + atom_exists_mask=torch.tensor( + index_by_atom_name(self.atom37_mask[mobile_inds], "CA", dim=-1) + & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1) + ).unsqueeze(0), + **kwargs, + ) + return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten() - return coordinates, plddt, residue_index + @cached_property + def residue_index_no_insertions(self) -> np.ndarray: + return self.residue_index + np.cumsum(self.insertion_code != "") + + @cached_property + def atom_array_no_insertions(self) -> bs.AtomArray: + atoms = [] + for res_idx, (res_name, positions, mask, conf) in enumerate( + zip( + self.sequence, + self.atom37_positions, + self.atom37_mask.astype(bool), + self.confidence, + strict=False, + ) + ): + for i, pos in zip(np.where(mask)[0], positions[mask], strict=False): + b_factor = ( + self.atom37_confidence[res_idx, i] + if self.atom37_confidence is not None + else conf + ) + atom = bs.Atom( + coord=pos, + # hard coded to as we currently only support single chain structures + chain_id=CHAIN_ID_CONST, + res_id=res_idx + 1, + res_name=residue_constants.restype_1to3.get(res_name, "UNK"), + hetero=False, + atom_name=residue_constants.atom_types[i], + element=residue_constants.atom_types[i][0], + b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE, + occupancy=1.0, + ) + atoms.append(atom) + return bs.array(atoms) diff --git a/fastplms/esmfold2/esmfold2_protein_complex.py b/src/fastplms/models/esmfold2/esmfold2_protein_complex.py similarity index 79% rename from fastplms/esmfold2/esmfold2_protein_complex.py rename to src/fastplms/models/esmfold2/esmfold2_protein_complex.py index 1e0eb6a..688ea5c 100644 --- a/fastplms/esmfold2/esmfold2_protein_complex.py +++ b/src/fastplms/models/esmfold2/esmfold2_protein_complex.py @@ -1,3 +1,5 @@ +"""Protein-complex data, assembly expansion, and geometry for ESMFold2.""" + from __future__ import annotations import io @@ -5,12 +7,13 @@ import random import re import warnings +from collections.abc import Iterable, Sequence from dataclasses import asdict, dataclass, replace from functools import cached_property from pathlib import Path from subprocess import check_output from tempfile import TemporaryDirectory -from typing import Any, Iterable, Sequence +from typing import Any import biotite.structure as bs import brotli @@ -28,581 +31,117 @@ from scipy.spatial import KDTree from . import esmfold2_residue_constants as residue_constants -from .esmfold2_misc import slice_python_object_as_numpy from .esmfold2_affine3d import Affine3D from .esmfold2_aligner import Aligner from .esmfold2_atom_indexer import AtomIndexer from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca -from .esmfold2_mmcif_parsing import MmcifWrapper, NoProteinError +from .esmfold2_misc import slice_python_object_as_numpy +from .esmfold2_mmcif_parsing import ( + MmcifWrapper, + NoProteinError, + round_mmcif_columns, +) from .esmfold2_protein_chain import ( ProteinChain, _str_key_to_int_key, chain_to_ndarray, index_by_atom_name, - infer_CB, + infer_cb, ) from .esmfold2_utils_types import PathOrBuffer -msgpack_numpy.patch() +SINGLE_LETTER_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" -SINGLE_LETTER_CHAIN_IDS = ( - "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" -) +def _parse_operation_expression(expression: str) -> list[tuple[str, ...]]: + """Expand an mmCIF operation expression in application order.""" -def _parse_operation_expression(expression): - """ - Get successive operation steps (IDs) for the given - ``oper_expression``. - Form the cartesian product, if necessary. - Copied from biotite and fixed a bug - """ - # Split groups by parentheses: - # use the opening parenthesis as delimiter - # and just remove the closing parenthesis - expressions_per_step = expression.replace(")", "").split("(") - expressions_per_step = [e for e in expressions_per_step if len(e) > 0] - # Important: Operations are applied from right to left - expressions_per_step.reverse() - - operations = [] - for expr in expressions_per_step: - cur_expr = expr.split(",") - cur_op = [] - # Deal with e='1-10,20-30,40-50' type expressions - for e in cur_expr: - if "-" in e: - first, last = e.split("-") - cur_op.extend(str(id) for id in range(int(first), int(last) + 1)) - else: - cur_op.append(e) - operations.append(cur_op) + def expand_group(group: str) -> list[str]: + operation_ids: list[str] = [] + for term in group.split(","): + if "-" not in term: + operation_ids.append(term) + continue + first, last = (int(value) for value in term.split("-")) + operation_ids.extend(str(value) for value in range(first, last + 1)) + return operation_ids - # Cartesian product of operations - return list(itertools.product(*operations)) + groups = [group for group in expression.replace(")", "").split("(") if group] + groups.reverse() + return list(itertools.product(*(expand_group(group) for group in groups))) def _apply_transformations_fast(chains, transformation_dict, operations): - """ - Get subassembly by applying the given operations to the input - structure containing affected asym IDs. - """ - # Additional first dimesion for 'structure.repeat()' - results = [] - - # Apply corresponding transformation for each copy in the assembly - for c in chains: + """Return transformed copies of each affected protein chain.""" + transformed_chains = [] + for chain in chains: for operation in operations: - coord = c.atom37_positions.copy() - # Execute for each transformation step - # in the operation expression + coordinates = chain.atom37_positions.copy() for op_step in operation: - T = transformation_dict[op_step] - # Rotate - coord = matrix_rotate(coord, T.rotation) - # Translate - coord += T.target_translation - new_chain = replace(c, atom37_positions=coord) - results.append(new_chain) - - return results - - -@dataclass -class ProteinComplexMetadata: - entity_lookup: dict[int, int] - chain_lookup: dict[int, str] - mmcif: MmcifWrapper | None = None - # This is a dictionary that maps assembly ids to the list of unique chains - # in that assembly. Allows for usage of `switch_assembly`. - assembly_composition: dict[str, list[str]] | None = None - - -@dataclass -class DockQSingleScore: - native_chains: tuple[str, str] - DockQ: float - interface_rms: float - ligand_rms: float - fnat: float - fnonnat: float - clashes: float - F1: float - DockQ_F1: float - - -@dataclass -class DockQResult: - total_dockq: float - native_interfaces: int - chain_mapping: dict[str, str] - interfaces: dict[tuple[str, str], DockQSingleScore] - # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing - # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling - aligned: ProteinComplex - aligned_rmsd: float - - -@dataclass(frozen=True) -class ProteinComplex: - """Dataclass with atom37 representation of an entire protein complex.""" - - id: str - sequence: str - entity_id: np.ndarray # entities map to unique sequences - chain_id: np.ndarray # multiple chains might share an entity id - sym_id: np.ndarray # complexes might be copies of the same chain - residue_index: np.ndarray - insertion_code: np.ndarray - atom37_positions: np.ndarray - atom37_mask: np.ndarray - confidence: np.ndarray - # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort. - metadata: ProteinComplexMetadata - atom37_confidence: np.ndarray | None = None # [L, 37] per-atom pLDDT - - def __post_init__(self): - l = len(self.sequence) - assert self.atom37_positions.shape[0] == l, (self.atom37_positions.shape, l) - assert self.atom37_mask.shape[0] == l, (self.atom37_mask.shape, l) - assert self.residue_index.shape[0] == l, (self.residue_index.shape, l) - assert self.insertion_code.shape[0] == l, (self.insertion_code.shape, l) - assert self.confidence.shape[0] == l, (self.confidence.shape, l) - assert self.entity_id.shape[0] == l, (self.entity_id.shape, l) - assert self.chain_id.shape[0] == l, (self.chain_id.shape, l) - assert self.sym_id.shape[0] == l, (self.sym_id.shape, l) - if self.atom37_confidence is not None: - assert self.atom37_confidence.shape == self.atom37_mask.shape, ( - self.atom37_confidence.shape, - self.atom37_mask.shape, - ) - - def __getitem__(self, idx: int | list[int] | slice | np.ndarray): - """This function slices protein complexes without consideration of chain breaks - NOTE: When slicing with a boolean mask, it's possible that the output array won't - be the expected length. This is because we do our best to preserve chainbreak tokens. - """ - - if isinstance(idx, int): - idx = [idx] - if isinstance(idx, list): - raise ValueError( - "ProteinComplex doesn't supports indexing with lists of indices" - ) - - if isinstance(idx, np.ndarray): - is_chainbreak = np.asarray([s == "|" for s in self.sequence]) - idx = idx.astype(bool) | is_chainbreak - - complex = self._unsafe_slice(idx) - if len(complex) == 0: - return complex - - # detect runs of chainbreaks by searching for instances of '||' in complex.sequence - chainbreak_runs = np.asarray( - [ - complex.sequence[i : i + 2] == "||" - for i in range(len(complex.sequence) - 1) - ] - + [complex.sequence[-1] == "|"] - ) - # We should remove as many chainbreaks as possible from the start of the sequence - for i in range(len(chainbreak_runs)): - if complex.sequence[i] == "|": - chainbreak_runs[i] = True - else: - break - complex = complex._unsafe_slice(~chainbreak_runs) - return complex - - def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray): - sequence = slice_python_object_as_numpy(self.sequence, idx) - return replace( - self, - sequence=sequence, - entity_id=self.entity_id[..., idx], - chain_id=self.chain_id[..., idx], - sym_id=self.sym_id[..., idx], - residue_index=self.residue_index[..., idx], - insertion_code=self.insertion_code[..., idx], - atom37_positions=self.atom37_positions[..., idx, :, :], - atom37_mask=self.atom37_mask[..., idx, :], - confidence=self.confidence[..., idx], - atom37_confidence=self.atom37_confidence[..., idx, :] - if self.atom37_confidence is not None - else None, - ) - - def __len__(self): - return len(self.sequence) - - @property - def num_chains(self): - return len(self.chain_boundaries) - - @cached_property - def atoms(self) -> AtomIndexer: - return AtomIndexer(self, property="atom37_positions", dim=-2) - - @cached_property - def atom_mask(self) -> AtomIndexer: - return AtomIndexer(self, property="atom37_mask", dim=-1) - - @cached_property - def chain_lengths(self) -> np.ndarray: - return np.diff(self.chain_boundaries, axis=1).flatten() - - @cached_property - def chain_boundaries(self) -> list[tuple[int, int]]: - cb = [-1] - for i, s in enumerate(self.sequence): - if s == "|": - cb.append(i) - cb.append(len(self)) - return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)] - - def get_chain_by_index(self, index: int) -> ProteinChain: - try: - start, end = self.chain_boundaries[index] - return self[start:end].as_chain() - except IndexError: - raise IndexError(f"Chain index {index} out of bounds") - - def get_chain_by_id( - self, chain_id: str, sample_chain_if_duplicate: bool = True - ) -> ProteinChain: - valid_indices = [ - index - for index, id_of_index in self.metadata.chain_lookup.items() - if id_of_index == chain_id - ] - if not valid_indices: - raise KeyError(f"Chain ID {chain_id} not found") - if sample_chain_if_duplicate: - index_to_return = random.choice(valid_indices) - return self.get_chain_by_index(index_to_return) - else: - if len(valid_indices) > 1: - raise ValueError(f"Multiple chains with chain ID {chain_id} found") - return self.get_chain_by_index(valid_indices[0]) - - def chain_iter(self) -> Iterable[ProteinChain]: - for start, end in self.chain_boundaries: - c = self[start:end] - yield c.as_chain() - - def as_chain(self, force_conversion: bool = False) -> ProteinChain: - """Convert the ProteinComplex to a ProteinChain. - - Args: - force_conversion (bool): Forces the conversion into a protein chain even if the complex has multiple chains. - The purpose of this is to use ProteinChain specific functions (like cbeta_contacts). - - """ - if not force_conversion: - assert len(np.unique(self.chain_id)) == 1, f"{self.id}" - assert len(np.unique(self.entity_id)) == 1, f"{self.id}" - if self.chain_id[0] not in self.metadata.chain_lookup: - warnings.warn("Chain ID not found in metadata, using 'A' as default") - if self.entity_id[0] not in self.metadata.entity_lookup: - warnings.warn("Entity ID not found in metadata, using None as default") - chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A") - entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None) - else: - chain_id = "A" - entity_id = None - - return ProteinChain( - id=self.id, - sequence=self.sequence, - chain_id=chain_id, - entity_id=entity_id, - atom37_positions=self.atom37_positions, - atom37_mask=self.atom37_mask, - residue_index=self.residue_index, - insertion_code=self.insertion_code, - confidence=self.confidence, - mmcif=self.metadata.mmcif, - atom37_confidence=self.atom37_confidence, - ) - - @classmethod - def from_pdb( - cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False - ) -> "ProteinComplex": - atom_array = PDBFile.read(path).get_structure( - model=1, extra_fields=["b_factor"] - ) - - chains = [] - for chain in bs.chain_iter(atom_array): - chain = chain[~chain.hetero] - if len(chain) == 0: - continue - chains.append(ProteinChain.from_atomarray(chain, id, is_predicted)) - return ProteinComplex.from_chains(chains) - - def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): - atom_array = None - for chain in self.chain_iter(): - carr = ( - chain.atom_array - if include_insertions - else chain.atom_array_no_insertions - ) - atom_array = carr if atom_array is None else atom_array + carr - f = PDBFile() - f.set_structure(atom_array) - f.write(path) - - def to_pdb_string(self, include_insertions: bool = True) -> str: - buf = io.StringIO() - self.to_pdb(buf, include_insertions=include_insertions) - buf.seek(0) - return buf.read() - - def normalize_chain_ids_for_pdb(self): - # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index, - # we can normalize it instead which might be necessary for DockQ and to_pdb. - ids = SINGLE_LETTER_CHAIN_IDS - chains = [] - for i, chain in enumerate(self.chain_iter()): - chain = replace(chain, chain_id=ids[i]) - if i > len(ids): - raise RuntimeError("Too many chains to write to PDB file") - chains.append(chain) - - return ProteinComplex.from_chains(chains) - - def find_assembly_ids_with_chain(self, id: str) -> list[str]: - good_chains = [] - if (comp := self.metadata.assembly_composition) is not None: - for assembly_id, chain_ids in comp.items(): - if id in chain_ids: - good_chains.append(assembly_id) - else: - raise ValueError( - "Cannot switch assemblies on this ProteinComplex, you must create the assembly from mmcif to support this" - ) - return good_chains - - def switch_assembly(self, id: str): - assert self.metadata.mmcif is not None - return get_assembly_fast(self.metadata.mmcif, assembly_id=id) - - def state_dict(self, backbone_only=False, json_serializable=False): - """This state dict is optimized for storage, so it turns things to fp16 whenever - possible. Note that we also only support int32 residue indices, I'm hoping we don't - need more than 2**32 residues...""" - dct = {k: v for k, v in vars(self).items()} - if backbone_only: - dct["atom37_mask"][:, 3:] = False - dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] - if dct.get("atom37_confidence") is not None: - dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] - else: - dct.pop("atom37_confidence", None) - for k, v in dct.items(): - if isinstance(v, np.ndarray): - match v.dtype: - case np.int64: - dct[k] = v.astype(np.int32) - case np.float64 | np.float32: - dct[k] = v.astype(np.float16) - case _: - pass - if json_serializable: - dct[k] = v.tolist() - elif isinstance(v, ProteinComplexMetadata): - dct[k] = asdict(v) - dct["metadata"]["mmcif"] = None - # These can be populated with non-serializable objects and are not needed for reconstruction - dct.pop("atoms", None) - dct.pop("atom_mask", None) - dct.pop("per_chain_kd_trees", None) - return dct - - def to_blob(self, backbone_only=False) -> bytes: - return brotli.compress(msgpack.dumps(self.state_dict(backbone_only)), quality=5) - - @classmethod - def from_state_dict(cls, dct): - # Note: assembly_composition is *supposed* to have string keys. - dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) - - for k, v in dct.items(): - if isinstance(v, list): - dct[k] = np.array(v) - - atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) - atom37[dct["atom37_mask"]] = dct["atom37_positions"] - dct["atom37_positions"] = atom37 - if "atom37_confidence" in dct: - atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) - atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] - dct["atom37_confidence"] = atom37_conf - dct = { - k: ( - v.astype(np.float32) - if k in ["atom37_positions", "confidence", "atom37_confidence"] - else v - ) - for k, v in dct.items() - } - if "chain_boundaries" in dct: - del dct["chain_boundaries"] - if "chain_boundaries" in dct["metadata"]: - del dct["metadata"]["chain_boundaries"] - dct["metadata"] = ProteinComplexMetadata(**dct["metadata"]) - return cls(**dct) - - @classmethod - def from_blob(cls, input: Path | str | io.BytesIO | bytes): - """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory - of chains from 52G/1M chains to 20G/1M chains, I think this is a good first - shot at compressing and dumping chains to disk. I'm sure there's better ways.""" - match input: - case Path() | str(): - bytes = Path(input).read_bytes() - case io.BytesIO(): - bytes = input.getvalue() - case _: - bytes = input - return cls.from_state_dict( - msgpack.loads(brotli.decompress(bytes), strict_map_key=False) - ) - - @classmethod - def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex: - f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore - return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False) - - @classmethod - def from_mmcif( - cls, - path: PathOrBuffer, - id: str | None = None, - assembly_id: str | None = None, - is_predicted: bool = False, - keep_source: bool = False, - ): - """Return a ProteinComplex object from an mmcif file. - TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation, - we only consider the first defined complex! - - Args: - path (str | Path | io.TextIO): Path or buffer to read mmcif file from. Should be uncompressed. - id (str, optional): String identifier to assign to structure. Will attempt to infer otherwise. - is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. - chain_id (str, optional): Select a chain corresponding to (author) chain id. - """ - mmcif = MmcifWrapper.read(path, id) - return get_assembly_fast(mmcif, assembly_id=assembly_id) - - @classmethod - def from_chains( - cls, - chains: Sequence[ProteinChain], - mmcif: MmcifWrapper | None = None, - all_assembly_metadata_dictionary: dict[str, list[str]] | None = None, - ): - if not chains: - raise ValueError( - "Cannot create a ProteinComplex from an empty list of chains" - ) - - # TODO(roshan): Make a proper protein complex class - def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): - full_array = [] - for array in arrays: - full_array.append(array) - full_array.append(sep) - full_array = full_array[:-1] - return np.concatenate(full_array, 0) - - sep_tokens = { - "residue_index": np.array([-1]), - "insertion_code": np.array([""]), - "atom37_positions": np.full([1, 37, 3], np.nan), - "atom37_mask": np.zeros([1, 37], dtype=bool), - "confidence": np.array([0]), - } - - any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains) - if any_has_atom37_conf: - sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32) - - def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray: - val = getattr(chain, name) - if val is None and name == "atom37_confidence": - return np.full([len(chain), 37], np.nan, dtype=np.float32) - return val - - array_args: dict[str, np.ndarray] = { - name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep) - for name, sep in sep_tokens.items() - } + transform = transformation_dict[op_step] + coordinates = matrix_rotate(coordinates, transform.rotation) + coordinates += transform.target_translation + transformed_chains.append(replace(chain, atom37_positions=coordinates)) + return transformed_chains - multimer_arrays = [] - chain2num_max = -1 - chain2num = {} - ent2num_max = -1 - ent2num = {} - total_index = 0 - for i, c in enumerate(chains): - num_res = c.residue_index.shape[0] - if c.chain_id not in chain2num: - chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1) - chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64) - if c.entity_id is None: - entity_num = (ent2num_max := ent2num_max + 1) - else: - if c.entity_id not in ent2num: - ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1) - entity_num = ent2num[c.entity_id] - entity_id_array = np.full([num_res], entity_num, dtype=np.int64) +@dataclass +class ProteinComplexMetadata: + entity_lookup: dict[int, int | str] + chain_lookup: dict[int, str] + mmcif: MmcifWrapper | None = None + # This is a dictionary that maps assembly ids to the list of unique chains + # in that assembly. Allows for usage of `switch_assembly`. + assembly_composition: dict[str, list[str]] | None = None - sym_id_array = np.full([num_res], i, dtype=np.int64) - multimer_arrays.append( - { - "chain_id": chain_id_array, - "entity_id": entity_id_array, - "sym_id": sym_id_array, - } - ) +@dataclass +class DockQSingleScore: + native_chains: tuple[str, str] + DockQ: float + interface_rms: float + ligand_rms: float + fnat: float + fnonnat: float + clashes: float + F1: float + DockQ_F1: float - total_index += num_res + 1 - sep = np.array([-1]) - update = { - name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep) - for name in ["chain_id", "entity_id", "sym_id"] - } - array_args.update(update) +@dataclass +class DockQResult: + total_dockq: float + native_interfaces: int + chain_mapping: dict[str, str] + interfaces: dict[tuple[str, str], DockQSingleScore] + # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing + # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling + aligned: ProteinComplex + aligned_rmsd: float - metadata = ProteinComplexMetadata( - mmcif=mmcif, - chain_lookup={v: k for k, v in chain2num.items()}, - entity_lookup={v: k for k, v in ent2num.items()}, - assembly_composition=all_assembly_metadata_dictionary, - ) - return cls( - id=chains[0].id, - sequence=residue_constants.CHAIN_BREAK_TOKEN.join( - chain.sequence for chain in chains - ), - metadata=metadata, - **array_args, - ) +@dataclass(frozen=True) +class ProteinComplex: + """Dataclass with atom37 representation of an entire protein complex.""" + + id: str + sequence: str + entity_id: np.ndarray # entities map to unique sequences + chain_id: np.ndarray # multiple chains might share an entity id + sym_id: np.ndarray # complexes might be copies of the same chain + residue_index: np.ndarray + insertion_code: np.ndarray + atom37_positions: np.ndarray + atom37_mask: np.ndarray + confidence: np.ndarray + # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort. + metadata: ProteinComplexMetadata + atom37_confidence: np.ndarray | None = None # P has shape (l, 37). + # Coordinate completion, concatenation, and comparison def infer_oxygen(self) -> ProteinComplex: """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" - O_missing_indices = np.argwhere( - ~np.isfinite(self.atoms["O"]).all(axis=1) - ).squeeze() + O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze() O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) @@ -611,19 +150,17 @@ def infer_oxygen(self) -> ProteinComplex: # Get the frame defined by the CA-C-N atom frames = Affine3D.from_graham_schmidt(CA, C, N) - O = frames.apply(O_vector) + oxygen_coordinates = frames.apply(O_vector) atom37_positions = self.atom37_positions.copy() atom37_mask = self.atom37_mask.copy() - atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = O[ + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[ O_missing_indices ].numpy() atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] ).any(-1) - new_chain = replace( - self, atom37_positions=atom37_positions, atom37_mask=atom37_mask - ) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) return new_chain def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex: @@ -645,19 +182,15 @@ def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex: N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) # See usage in trDesign codebase. # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 - inferred_cbeta_positions = infer_CB(C, N, CA, 1.522, 1.927, -2.143) + inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143) if not infer_cbeta_for_glycine: inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan - atom37_positions[:, residue_constants.atom_order["CB"]] = ( - inferred_cbeta_positions - ) + atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( atom37_positions[:, residue_constants.atom_order["CB"]] ).any(-1) - new_chain = replace( - self, atom37_positions=atom37_positions, atom37_mask=atom37_mask - ) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) return new_chain @classmethod @@ -677,10 +210,10 @@ def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex: ) def _sanity_check_complexes_are_comparable(self, other: ProteinComplex): - assert len(self) == len(other), "Protein complexes must have the same length" - assert len(list(self.chain_iter())) == len( - list(other.chain_iter()) - ), "Protein complexes must have the same number of chains" + if len(self) != len(other): + raise ValueError("Protein complexes must have the same length") + if len(list(self.chain_iter())) != len(list(other.chain_iter())): + raise ValueError("Protein complexes must have the same number of chains") def rmsd( self, @@ -696,15 +229,12 @@ def rmsd( Args: target (ProteinComplex): The target (other) protein complex to compare to. - also_check_reflection (bool, optional): If True, also check if the reflection of the mobile atoms has a lower RMSD. - mobile_inds (list[int], optional): The indices of the mobile atoms to align. These are NOT residue indices - target_inds (list[int], optional): The indices of the target atoms to align. These are NOT residue indices - only_compute_backbone_rmsd (bool, optional): If True, only compute the RMSD of the backbone atoms. + also_check_reflection: Compare the reflected mobile coordinates too. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_compute_backbone_rmsd: Restrict the score to backbone atoms. """ - if compute_chain_assignment: - aligned = self.dockq(target).aligned - else: - aligned = self + aligned = self.dockq(target).aligned if compute_chain_assignment else self aligner = Aligner( aligned if mobile_inds is None else aligned[mobile_inds], @@ -738,17 +268,14 @@ def lddt_ca( Arguments: target (ProteinComplex): The other protein complex to compare to. - mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices - target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. Returns: float | np.ndarray: The LDDT score between the two protein chains, either a single float or per-residue LDDT scores if `per_residue` is True. """ - if compute_chain_assignment: - aligned = self.dockq(target).aligned - else: - aligned = self + aligned = self.dockq(target).aligned if compute_chain_assignment else self lddt = compute_lddt_ca( torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0), torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0), @@ -769,16 +296,13 @@ def gdt_ts( Arguments: target (ProteinComplex): The other protein complex to compare to. - mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices - target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. Returns: float: The GDT_TS score between the two protein chains. """ - if compute_chain_assignment: - aligned = self.dockq(target).aligned - else: - aligned = self + aligned = self.dockq(target).aligned if compute_chain_assignment else self gdt_ts = compute_gdt_ts( mobile=torch.tensor( index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"), @@ -803,14 +327,12 @@ def dockq(self, native: ProteinComplex): # # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif, # if your protein has multi-letter or repeated chain IDs, this will fail. Please call - # pc = pc.normalize_chain_ids_for_pdb() before calling this function in that case (limit is 62 chains) + # Normalize chain IDs before DockQ when IDs repeat or use multiple letters. try: pass except BaseException: - raise RuntimeError( - "DockQ is not installed. Please update your environment." - ) + raise RuntimeError("DockQ is not installed. Please update your environment.") from None self._sanity_check_complexes_are_comparable(native) def sanity_check_chain_ids(pc: ProteinComplex): @@ -819,9 +341,7 @@ def sanity_check_chain_ids(pc: ProteinComplex): if i > len(SINGLE_LETTER_CHAIN_IDS): raise ValueError("Too many chains to write to PDB file") if len(chain.chain_id) > 1: - raise ValueError( - "We only supports single letter chain IDs for DockQ" - ) + raise ValueError("We only supports single letter chain IDs for DockQ") ids.append(chain.chain_id) if len(set(ids)) != len(ids): raise ValueError(f"Duplicate chain IDs in protein complex: {ids}") @@ -839,9 +359,7 @@ def sanity_check_chain_ids(pc: ProteinComplex): lines = output.decode().split("\n") # Remove the header comments - start_index = next( - i for i, line in enumerate(lines) if line.startswith("Model") - ) + start_index = next(i for i, line in enumerate(lines) if line.startswith("Model")) lines = lines[start_index:] result = {} @@ -853,20 +371,19 @@ def sanity_check_chain_ids(pc: ProteinComplex): if not line: continue - if line.startswith("Model :"): - pass # Tmp pdb file location, it's useless... - elif line.startswith("Native :"): + if line.startswith(("Model :", "Native :")): pass # Tmp pdb file location, it's useless... elif line.startswith("Total DockQ"): total_dockq_match = re.search( - r"Total DockQ over (\d+) native interfaces: ([\d.]+) with (.*) model:native mapping", + r"Total DockQ over (\d+) native interfaces: ([\d.]+) with " + r"(.*) model:native mapping", line, ) if total_dockq_match: result["value"] = float(total_dockq_match.group(2)) result["native interfaces"] = int(total_dockq_match.group(1)) native_chains, self_chains = total_dockq_match.group(3).split(":") - result["mapping"] = dict(zip(native_chains, self_chains)) + result["mapping"] = dict(zip(native_chains, self_chains, strict=False)) else: raise RuntimeError( "Failed to parse DockQ output, maybe your DockQ version is wrong?" @@ -874,13 +391,9 @@ def sanity_check_chain_ids(pc: ProteinComplex): elif line.startswith("Native chains:"): if current_interface: interfaces.append(current_interface) - current_interface = { - "Native chains": line.split(":")[1].strip().split(", ") - } + current_interface = {"Native chains": line.split(":")[1].strip().split(", ")} elif line.startswith("Model chains:"): - current_interface["Model chains"] = ( - line.split(":")[1].strip().split(", ") - ) + current_interface["Model chains"] = line.split(":")[1].strip().split(", ") elif ":" in line: key, value = line.split(":", 1) current_interface[key.strip()] = float(value.strip()) @@ -917,8 +430,7 @@ def parse_dict(d: dict[str, Any]) -> DockQSingleScore: native_interfaces=result["native interfaces"], chain_mapping=result["mapping"], interfaces={ - (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) - for i in interfaces + (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) for i in interfaces }, aligned=realigned, aligned_rmsd=aligner.rmsd, @@ -926,6 +438,227 @@ def parse_dict(d: dict[str, Any]) -> DockQSingleScore: return result + # Object invariants, slicing, and chain views + def __post_init__(self): + if not isinstance(self.sequence, str): + raise TypeError("sequence must be a string.") + sequence_length = len(self.sequence) + aligned = { + "atom37_positions": self.atom37_positions, + "atom37_mask": self.atom37_mask, + "residue_index": self.residue_index, + "insertion_code": self.insertion_code, + "confidence": self.confidence, + "entity_id": self.entity_id, + "chain_id": self.chain_id, + "sym_id": self.sym_id, + } + for name, values in aligned.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != sequence_length: + raise ValueError( + f"{name} shape {values.shape} does not align with " + f"sequence length {sequence_length}." + ) + if self.atom37_positions.shape != (sequence_length, 37, 3): + raise ValueError( + "atom37_positions must have shape " + f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}." + ) + if self.atom37_mask.shape != (sequence_length, 37): + raise ValueError( + "atom37_mask must have shape " + f"({sequence_length}, 37), got {self.atom37_mask.shape}." + ) + if self.atom37_mask.dtype != bool: + raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.") + if not np.issubdtype(self.atom37_positions.dtype, np.number): + raise TypeError("atom37_positions must use a numeric dtype.") + for name, values in ( + ("residue_index", self.residue_index), + ("insertion_code", self.insertion_code), + ("confidence", self.confidence), + ("entity_id", self.entity_id), + ("chain_id", self.chain_id), + ("sym_id", self.sym_id), + ): + if values.shape != (sequence_length,): + raise ValueError( + f"{name} must have shape ({sequence_length},), got {values.shape}." + ) + if not np.issubdtype(self.confidence.dtype, np.number): + raise TypeError("confidence must use a numeric dtype.") + atom37_confidence = self.atom37_confidence + if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray): + raise TypeError("atom37_confidence must be a NumPy array when provided.") + if ( + isinstance(atom37_confidence, np.ndarray) + and atom37_confidence.shape != self.atom37_mask.shape + ): + raise ValueError( + "atom37_confidence shape must match atom37_mask: " + f"{atom37_confidence.shape} != {self.atom37_mask.shape}." + ) + + def __getitem__(self, idx: int | list[int] | slice | np.ndarray): + """This function slices protein complexes without consideration of chain breaks + NOTE: When slicing with a boolean mask, it's possible that the output array won't + be the expected length. This is because we do our best to preserve chainbreak tokens. + """ + + if isinstance(idx, int): + idx = [idx] + if isinstance(idx, list): + raise ValueError("ProteinComplex doesn't supports indexing with lists of indices") + + if isinstance(idx, np.ndarray): + is_chainbreak = np.asarray([s == "|" for s in self.sequence]) + idx = idx.astype(bool) | is_chainbreak + + complex = self._unsafe_slice(idx) + if len(complex) == 0: + return complex + + # detect runs of chainbreaks by searching for instances of '||' in complex.sequence + chainbreak_runs = np.asarray( + [complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1)] + + [complex.sequence[-1] == "|"] + ) + # We should remove as many chainbreaks as possible from the start of the sequence + for i in range(len(chainbreak_runs)): + if complex.sequence[i] == "|": + chainbreak_runs[i] = True + else: + break + complex = complex._unsafe_slice(~chainbreak_runs) + return complex + + def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray): + sequence = slice_python_object_as_numpy(self.sequence, idx) + return replace( + self, + sequence=sequence, + entity_id=self.entity_id[..., idx], + chain_id=self.chain_id[..., idx], + sym_id=self.sym_id[..., idx], + residue_index=self.residue_index[..., idx], + insertion_code=self.insertion_code[..., idx], + atom37_positions=self.atom37_positions[..., idx, :, :], + atom37_mask=self.atom37_mask[..., idx, :], + confidence=self.confidence[..., idx], + atom37_confidence=self.atom37_confidence[..., idx, :] + if self.atom37_confidence is not None + else None, + ) + + def __len__(self): + return len(self.sequence) + + @property + def num_chains(self): + return len(self.chain_boundaries) + + @cached_property + def atoms(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_positions", dim=-2) + + @cached_property + def atom_mask(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_mask", dim=-1) + + @cached_property + def chain_lengths(self) -> np.ndarray: + return np.diff(self.chain_boundaries, axis=1).flatten() + + @cached_property + def chain_boundaries(self) -> list[tuple[int, int]]: + cb = [-1] + for i, s in enumerate(self.sequence): + if s == "|": + cb.append(i) + cb.append(len(self)) + return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)] + + def get_chain_by_index(self, index: int) -> ProteinChain: + try: + start, end = self.chain_boundaries[index] + return self[start:end].as_chain() + except IndexError: + raise IndexError(f"Chain index {index} out of bounds") from None + + def get_chain_by_id( + self, chain_id: str, sample_chain_if_duplicate: bool = True + ) -> ProteinChain: + valid_indices = [ + index + for index, id_of_index in self.metadata.chain_lookup.items() + if id_of_index == chain_id + ] + if not valid_indices: + raise KeyError(f"Chain ID {chain_id} not found") + if sample_chain_if_duplicate: + index_to_return = random.choice(valid_indices) + return self.get_chain_by_index(index_to_return) + else: + if len(valid_indices) > 1: + raise ValueError(f"Multiple chains with chain ID {chain_id} found") + return self.get_chain_by_index(valid_indices[0]) + + def chain_iter(self) -> Iterable[ProteinChain]: + for start, end in self.chain_boundaries: + c = self[start:end] + yield c.as_chain() + + def as_chain(self, force_conversion: bool = False) -> ProteinChain: + """Convert the ProteinComplex to a ProteinChain. + + Args: + force_conversion: Flatten multiple chains to access chain-only utilities. + + """ + if not force_conversion: + if len(np.unique(self.chain_id)) != 1: + raise ValueError( + f"Protein complex {self.id!r} has multiple chains; " + "pass force_conversion=True to flatten it." + ) + if len(np.unique(self.entity_id)) != 1: + raise ValueError( + f"Protein complex {self.id!r} has multiple entities; " + "pass force_conversion=True to flatten it." + ) + if self.chain_id[0] not in self.metadata.chain_lookup: + warnings.warn( + "Chain ID not found in metadata, using 'A' as default", + stacklevel=2, + ) + if self.entity_id[0] not in self.metadata.entity_lookup: + warnings.warn( + "Entity ID not found in metadata, using None as default", + stacklevel=2, + ) + chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A") + entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None) + else: + chain_id = "A" + entity_id = None + + return ProteinChain( + id=self.id, + sequence=self.sequence, + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=self.atom37_positions, + atom37_mask=self.atom37_mask, + residue_index=self.residue_index, + insertion_code=self.insertion_code, + confidence=self.confidence, + mmcif=self.metadata.mmcif, + atom37_confidence=self.atom37_confidence, + ) + + # Contact topology and mmCIF export @cached_property def per_chain_kd_trees(self): # Iterate over chains, build KDTree for each chain @@ -1010,6 +743,7 @@ def to_mmcif_string(self) -> str: # Add entity information for proper mmCIF structure self._add_entity_information(f) + round_mmcif_columns(f) # Write to string output = io.StringIO() @@ -1045,13 +779,9 @@ def _add_entity_information(self, cif_file: CIFFile) -> None: cif_file.block["entity"] = CIFCategory( name="entity", - columns={ - "id": CIFColumn( - data=CIFData(array=np.array(entity_ids), dtype=np.str_) - ), - "type": CIFColumn( - data=CIFData(array=np.array(entity_types), dtype=np.str_) - ), + columns={ + "id": CIFColumn(data=CIFData(array=np.array(entity_ids), dtype=np.str_)), + "type": CIFColumn(data=CIFData(array=np.array(entity_types), dtype=np.str_)), "pdbx_description": CIFColumn( data=CIFData(array=np.array(entity_descriptions), dtype=np.str_) ), @@ -1076,9 +806,7 @@ def _add_entity_information(self, cif_file: CIFFile) -> None: "entity_id": CIFColumn( data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_) ), - "type": CIFColumn( - data=CIFData(array=np.array(poly_types), dtype=np.str_) - ), + "type": CIFColumn(data=CIFData(array=np.array(poly_types), dtype=np.str_)), "nstd_linkage": CIFColumn( data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_) ), @@ -1105,13 +833,288 @@ def _add_entity_information(self, cif_file: CIFFile) -> None: "entity_id": CIFColumn( data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_) ), - "details": CIFColumn( - data=CIFData(array=np.array(asym_details), dtype=np.str_) - ), + "details": CIFColumn(data=CIFData(array=np.array(asym_details), dtype=np.str_)), }, ) + # Construction, PDB interchange, and compact storage + @classmethod + def from_pdb( + cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False + ) -> ProteinComplex: + atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"]) + + chains = [] + for chain in bs.chain_iter(atom_array): + chain = chain[~chain.hetero] + if len(chain) == 0: + continue + chains.append(ProteinChain.from_atomarray(chain, id, is_predicted)) + return ProteinComplex.from_chains(chains) + + def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): + atom_array = None + for chain in self.chain_iter(): + carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions + atom_array = carr if atom_array is None else atom_array + carr + f = PDBFile() + f.set_structure(atom_array) + f.write(path) + + def to_pdb_string(self, include_insertions: bool = True) -> str: + buf = io.StringIO() + self.to_pdb(buf, include_insertions=include_insertions) + buf.seek(0) + return buf.read() + + def normalize_chain_ids_for_pdb(self): + # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index, + # we can normalize it instead which might be necessary for DockQ and to_pdb. + ids = SINGLE_LETTER_CHAIN_IDS + chains = [] + for i, chain in enumerate(self.chain_iter()): + chain = replace(chain, chain_id=ids[i]) + if i > len(ids): + raise RuntimeError("Too many chains to write to PDB file") + chains.append(chain) + + return ProteinComplex.from_chains(chains) + + def find_assembly_ids_with_chain(self, id: str) -> list[str]: + good_chains = [] + if (comp := self.metadata.assembly_composition) is not None: + for assembly_id, chain_ids in comp.items(): + if id in chain_ids: + good_chains.append(assembly_id) + else: + raise ValueError( + "Cannot switch assemblies on this ProteinComplex; construct it from " + "mmCIF to retain assembly metadata" + ) + return good_chains + + def switch_assembly(self, id: str): + if self.metadata.mmcif is None: + raise ValueError( + "Cannot switch assemblies without retained mmCIF source metadata." + ) + return get_assembly_fast(self.metadata.mmcif, assembly_id=id) + + def state_dict(self, backbone_only=False, json_serializable=False): + """This state dict is optimized for storage, so it turns things to fp16 whenever + possible. Note that we also only support int32 residue indices, I'm hoping we don't + need more than 2**32 residues...""" + dct = {k: v for k, v in vars(self).items()} + if backbone_only: + # Frozen dataclasses do not make their NumPy members immutable. Work on a + # private mask so requesting a compact backbone payload cannot clear the + # caller's side-chain atoms in-place. + atom37_mask = dct["atom37_mask"].copy() + atom37_mask[:, 3:] = False + dct["atom37_mask"] = atom37_mask + dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] + if dct.get("atom37_confidence") is not None: + dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] + else: + dct.pop("atom37_confidence", None) + for k, v in dct.items(): + if isinstance(v, np.ndarray): + match v.dtype: + case np.int64: + dct[k] = v.astype(np.int32) + case np.float64 | np.float32: + dct[k] = v.astype(np.float16) + case _: + pass + if json_serializable: + dct[k] = v.tolist() + elif isinstance(v, ProteinComplexMetadata): + dct[k] = asdict(v) + dct["metadata"]["mmcif"] = None + # These can be populated with non-serializable objects and are not needed for reconstruction + dct.pop("atoms", None) + dct.pop("atom_mask", None) + dct.pop("per_chain_kd_trees", None) + return dct + + def to_blob(self, backbone_only=False) -> bytes: + payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode) + return brotli.compress(payload, quality=5) + + @classmethod + def from_state_dict(cls, dct): + # Note: assembly_composition is *supposed* to have string keys. + dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) + + for k, v in dct.items(): + if isinstance(v, list): + dct[k] = np.array(v) + + atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) + atom37[dct["atom37_mask"]] = dct["atom37_positions"] + dct["atom37_positions"] = atom37 + if "atom37_confidence" in dct: + atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) + atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] + dct["atom37_confidence"] = atom37_conf + dct = { + k: ( + v.astype(np.float32) + if k in ["atom37_positions", "confidence", "atom37_confidence"] + else v + ) + for k, v in dct.items() + } + if "chain_boundaries" in dct: + del dct["chain_boundaries"] + if "chain_boundaries" in dct["metadata"]: + del dct["metadata"]["chain_boundaries"] + dct["metadata"] = ProteinComplexMetadata(**dct["metadata"]) + return cls(**dct) + + @classmethod + def from_blob(cls, input: Path | str | io.BytesIO | bytes): + """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory + of chains from 52G/1M chains to 20G/1M chains, I think this is a good first + shot at compressing and dumping chains to disk. I'm sure there's better ways.""" + match input: + case Path() | str(): + bytes = Path(input).read_bytes() + case io.BytesIO(): + bytes = input.getvalue() + case _: + bytes = input + state = msgpack.loads( + brotli.decompress(bytes), + object_hook=msgpack_numpy.decode, + strict_map_key=False, + ) + return cls.from_state_dict(state) + + @classmethod + def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex: + f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore + return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False) + + @classmethod + def from_mmcif( + cls, + path: PathOrBuffer, + id: str | None = None, + assembly_id: str | None = None, + is_predicted: bool = False, + keep_source: bool = False, + ): + """Return a ProteinComplex object from an mmcif file. + TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation, + we only consider the first defined complex! + + Args: + path: Uncompressed mmCIF path or text buffer. + id: Optional structure identifier. + is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. + chain_id (str, optional): Select a chain corresponding to (author) chain id. + """ + mmcif = MmcifWrapper.read(path, id) + return get_assembly_fast(mmcif, assembly_id=assembly_id) + + @classmethod + def from_chains( + cls, + chains: Sequence[ProteinChain], + mmcif: MmcifWrapper | None = None, + all_assembly_metadata_dictionary: dict[str, list[str]] | None = None, + ): + if not chains: + raise ValueError("Cannot create a ProteinComplex from an empty list of chains") + + # TODO(roshan): Make a proper protein complex class + def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): + full_array = [] + for array in arrays: + full_array.append(array) + full_array.append(sep) + full_array = full_array[:-1] + return np.concatenate(full_array, 0) + + sep_tokens = { + "residue_index": np.array([-1]), + "insertion_code": np.array([""]), + "atom37_positions": np.full([1, 37, 3], np.nan), + "atom37_mask": np.zeros([1, 37], dtype=bool), + "confidence": np.array([0]), + } + + any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains) + if any_has_atom37_conf: + sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32) + + def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray: + val = getattr(chain, name) + if val is None and name == "atom37_confidence": + return np.full([len(chain), 37], np.nan, dtype=np.float32) + return val + + array_args: dict[str, np.ndarray] = { + name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep) + for name, sep in sep_tokens.items() + } + + multimer_arrays = [] + chain2num_max = -1 + chain2num = {} + ent2num_max = -1 + ent2num = {} + total_index = 0 + for i, c in enumerate(chains): + num_res = c.residue_index.shape[0] + if c.chain_id not in chain2num: + chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1) + chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64) + + if c.entity_id is None: + entity_num = (ent2num_max := ent2num_max + 1) + else: + if c.entity_id not in ent2num: + ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1) + entity_num = ent2num[c.entity_id] + entity_id_array = np.full([num_res], entity_num, dtype=np.int64) + + sym_id_array = np.full([num_res], i, dtype=np.int64) + + multimer_arrays.append( + { + "chain_id": chain_id_array, + "entity_id": entity_id_array, + "sym_id": sym_id_array, + } + ) + + total_index += num_res + 1 + + sep = np.array([-1]) + update = { + name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep) + for name in ["chain_id", "entity_id", "sym_id"] + } + array_args.update(update) + + metadata = ProteinComplexMetadata( + mmcif=mmcif, + chain_lookup={v: k for k, v in chain2num.items()}, + entity_lookup={v: k for k, v in ent2num.items()}, + assembly_composition=all_assembly_metadata_dictionary, + ) + + return cls( + id=chains[0].id, + sequence=residue_constants.CHAIN_BREAK_TOKEN.join(chain.sequence for chain in chains), + metadata=metadata, + **array_args, + ) + +# Biological-assembly expansion def get_assembly_fast( mmcif: MmcifWrapper, assembly_id=None, @@ -1189,6 +1192,7 @@ def get_assembly_fast( assembly_gen_category["assembly_id"].data.array, assembly_gen_category["oper_expression"].data.array, assembly_gen_category["asym_id_list"].data.array, + strict=False, ): if aid == assembly_id: # Parse operations and asym IDs for this specific entry @@ -1196,14 +1200,10 @@ def get_assembly_fast( asym_ids = asym_id_expr.split(",") # Filter affected asym IDs to only protein chains, preserving order - sub_structures = [ - asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain - ] + sub_structures = [asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain] # Apply transformations - sub_assembly = _apply_transformations_fast( - sub_structures, transformations, operations - ) + sub_assembly = _apply_transformations_fast(sub_structures, transformations, operations) assembly.extend(sub_assembly) # Build assembly_id_dict for this entry @@ -1222,7 +1222,7 @@ def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex: chain_breaks = np.array(list(chain.sequence)) == "|" chain_break_inds = np.where(chain_breaks)[0] chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]]) - chain_break_inds = np.array(list(zip(chain_break_inds[:-1], chain_break_inds[1:]))) + chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds))) complex_chains = [] for start, end in chain_break_inds: if start != 0: diff --git a/src/fastplms/models/esmfold2/esmfold2_protein_structure.py b/src/fastplms/models/esmfold2/esmfold2_protein_structure.py new file mode 100644 index 0000000..1d41330 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_protein_structure.py @@ -0,0 +1,251 @@ +"""Atom selection, rigid alignment, RMSD, and GDT-TS primitives.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TypeVar + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.amp import autocast # type: ignore + +from .esmfold2_affine3d import Affine3D +from .esmfold2_misc import unbinpack +from .esmfold2_normalize_coordinates import index_by_atom_name + +ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) + + +def _coordinate_operations( + coordinates: ArrayOrTensor, +) -> tuple[Callable[[ArrayOrTensor], ArrayOrTensor], Callable[..., ArrayOrTensor]]: + if isinstance(coordinates, np.ndarray): + + def normalize(X: ArrayOrTensor) -> ArrayOrTensor: + return X / np.linalg.norm(X, axis=-1, keepdims=True) + + return normalize, np.cross + return F.normalize, torch.cross # type: ignore[return-value] + + +def infer_cbeta_from_atom37( + atom37: ArrayOrTensor, + bond_length: float = 1.522, + bond_angle: float = 1.927, + dihedral: float = -2.143, +) -> ArrayOrTensor: + """Infer C-beta coordinates from backbone tensor ``X``. + + The scalar keyword arguments encode the bond length, bond angle, and + dihedral in radians used by the checkpoint's training geometry. + """ + + n_position = index_by_atom_name(atom37, "N", dim=-2) + ca_position = index_by_atom_name(atom37, "CA", dim=-2) + c_position = index_by_atom_name(atom37, "C", dim=-2) + normalize, cross = _coordinate_operations(atom37) + with np.errstate(invalid="ignore"): + n_to_ca = n_position - ca_position + n_to_c = n_position - c_position + unit_n_to_ca = normalize(n_to_ca) + normal = normalize(cross(n_to_c, unit_n_to_ca)) + basis = [unit_n_to_ca, cross(normal, unit_n_to_ca), normal] + coefficients = [ + bond_length * np.cos(bond_angle), + bond_length * np.sin(bond_angle) * np.cos(dihedral), + -bond_length * np.sin(bond_angle) * np.sin(dihedral), + ] + offset = sum( + vector * coefficient for vector, coefficient in zip(basis, coefficients, strict=True) + ) + return ca_position + offset + + +def _unpack_alignment_inputs( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, + sequence_id: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor | None]: + if sequence_id is None: + return mobile, target, atom_mask + unpacked_mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan) + unpacked_target = unbinpack(target, sequence_id, pad_value=torch.nan) + if atom_mask is None: + unpacked_mask = torch.isfinite(unpacked_target).all(dim=-1) + else: + unpacked_mask = unbinpack(atom_mask, sequence_id, pad_value=0) + return unpacked_mobile, unpacked_target, unpacked_mask + + +def _flatten_atom_axes( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor | None]: + b = mobile.shape[0] + flat_mobile = mobile.view(b, -1, 3) if mobile.dim() == 4 else mobile + flat_target = target.view(b, -1, 3) if target.dim() == 4 else target + flat_mask = atom_mask + if flat_mask is not None and flat_mask.dim() == 3: + flat_mask = flat_mask.view(b, -1) + return flat_mobile, flat_target, flat_mask + + +def _masked_coordinates( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor]: + if atom_mask is None: + atom_mask = torch.ones( + mobile.shape[:2], + dtype=torch.bool, + device=mobile.device, + ) + return mobile, target, atom_mask + expanded_mask = atom_mask.unsqueeze(-1) + return ( + mobile.masked_fill(~expanded_mask, 0), + target.masked_fill(~expanded_mask, 0), + atom_mask, + ) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_alignment_tensors( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Center and align coordinate tensors ``X`` and ``Y``. + + Inputs have shape (b, n, 3), or (b, l, n_atoms, 3). The returned rotation + tensor ``R`` has shape (b, 3, 3), and atom counts have shape (b, 1). + """ + + mobile, target, atom_exists_mask = _unpack_alignment_inputs( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + if mobile.shape != target.shape: + raise AssertionError("Batch structure shapes do not match!") + mobile, target, atom_exists_mask = _flatten_atom_axes( + mobile, + target, + atom_exists_mask, + ) + mobile, target, atom_exists_mask = _masked_coordinates( + mobile, + target, + atom_exists_mask, + ) + + num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True) + centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) + centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) + centroid_mobile[num_valid_atoms == 0] = 0 + centroid_target[num_valid_atoms == 0] = 0 + + expanded_mask = atom_exists_mask.unsqueeze(-1) + centered_mobile = (mobile - centroid_mobile).masked_fill(~expanded_mask, 0) + centered_target = (target - centroid_target).masked_fill(~expanded_mask, 0) + covariance = torch.matmul(centered_mobile.transpose(1, 2), centered_target) + left_vectors, _, right_vectors = torch.svd(covariance) + rotation = torch.matmul(left_vectors, right_vectors.transpose(1, 2)) + return ( + centered_mobile, + centroid_mobile, + centered_target, + centroid_target, + rotation, + num_valid_atoms, + ) + + +def _validate_reduction(reduction: str, allowed: tuple[str, ...]) -> None: + if reduction not in allowed: + raise ValueError("Unrecognized reduction: '{reduction}'") + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_rmsd_no_alignment( + aligned: Tensor, + target: Tensor, + num_valid_atoms: Tensor, + reduction: str = "batch", +) -> Tensor: + """Measure RMSD after alignment using a declared reduction.""" + + _validate_reduction(reduction, ("per_residue", "per_sample", "batch")) + difference = aligned - target + if reduction == "per_residue": + mean_squared_error = difference.square().view(difference.size(0), -1, 9).mean(-1) + else: + mean_squared_error = difference.square().sum(dim=(1, 2)) / num_valid_atoms.squeeze(-1) + rmsd = torch.sqrt(mean_squared_error) + if reduction in {"per_residue", "per_sample"}: + return rmsd + valid_samples = num_valid_atoms.squeeze(-1) > 0 + return rmsd.masked_fill(~valid_samples, 0).sum() / (valid_samples.sum() + 1e-8) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_affine_and_rmsd( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, +) -> tuple[Affine3D, Tensor]: + """Fit ``X`` onto ``Y`` and return the rigid transform and batch RMSD.""" + + ( + centered_mobile, + centroid_mobile, + centered_target, + centroid_target, + rotation, + num_valid_atoms, + ) = compute_alignment_tensors(mobile, target, atom_exists_mask, sequence_id) + translation = torch.matmul(-centroid_mobile, rotation) + centroid_target + affine = Affine3D.from_tensor_pair( + translation, + rotation.unsqueeze(dim=-3).transpose(-2, -1), + ) + rotated_mobile = torch.matmul(centered_mobile, rotation) + rmsd = compute_rmsd_no_alignment( + rotated_mobile, + centered_target, + num_valid_atoms, + reduction="batch", + ) + return affine, rmsd + + +def compute_gdt_ts_no_alignment( + aligned: Tensor, + target: Tensor, + atom_exists_mask: Tensor, + reduction: str = "batch", +) -> Tensor: + """Compute GDT-TS for already aligned coordinate tensors.""" + + _validate_reduction(reduction, ("per_sample", "batch")) + if atom_exists_mask is None: + atom_exists_mask = torch.isfinite(target).all(dim=-1) + deviation = torch.linalg.vector_norm(aligned - target, dim=-1) + counts = atom_exists_mask.sum(dim=-1) + score_1 = ((deviation < 1) * atom_exists_mask).sum(dim=-1) / counts + score_2 = ((deviation < 2) * atom_exists_mask).sum(dim=-1) / counts + score_4 = ((deviation < 4) * atom_exists_mask).sum(dim=-1) / counts + score_8 = ((deviation < 8) * atom_exists_mask).sum(dim=-1) / counts + score = (score_1 + score_2 + score_4 + score_8) * 0.25 + return score.mean() if reduction == "batch" else score diff --git a/src/fastplms/models/esmfold2/esmfold2_residue_constants.py b/src/fastplms/models/esmfold2/esmfold2_residue_constants.py new file mode 100644 index 0000000..42655c5 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_residue_constants.py @@ -0,0 +1,1017 @@ +# Copyright 2025 EvolutionaryScale +# Copyright 2021 AlQuraishi Laboratory +# Copyright 2021 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canonical amino-acid geometry tables used by ESMFold2. + +The literal chemistry measurements are kept visible for review. Derived masks, +indices, frames, and atom mappings are built locally below and are checked +exactly against the pinned Biohub implementation. +""" + +from __future__ import annotations + +import functools +from collections import defaultdict, namedtuple +from collections.abc import Mapping +from pathlib import Path +from typing import Any, cast + +import numpy as np + +ca_ca = 3.80209737096 +chi_angles_atoms = { + "ALA": [], + "ARG": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "NE"], + ["CG", "CD", "NE", "CZ"], + ], + "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], + "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], + "CYS": [["N", "CA", "CB", "SG"]], + "GLN": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "OE1"], + ], + "GLU": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "OE1"], + ], + "GLY": [], + "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]], + "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]], + "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "LYS": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "CE"], + ["CG", "CD", "CE", "NZ"], + ], + "MET": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "SD"], + ["CB", "CG", "SD", "CE"], + ], + "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]], + "SER": [["N", "CA", "CB", "OG"]], + "THR": [["N", "CA", "CB", "OG1"]], + "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "VAL": [["N", "CA", "CB", "CG1"]], + "UNK": [], +} +chi_angles_mask = [ + [1.0] * len(groups) + [0.0] * (4 - len(groups)) for groups in chi_angles_atoms.values() +] +_PI_PERIODIC_CHI = {"ASP": {1}, "GLU": {2}, "PHE": {1}, "TYR": {1}} +chi_pi_periodic = [ + [1.0 if chi_index in _PI_PERIODIC_CHI.get(residue_name, ()) else 0.0 for chi_index in range(4)] + for residue_name in chi_angles_atoms +] +rigid_group_atom_positions: dict[str, list[list[Any]]] = { + "ALA": [ + ["N", 0, (-0.525, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.529, -0.774, -1.205)], + ["O", 3, (0.627, 1.062, 0.0)], + ], + "ARG": [ + ["N", 0, (-0.524, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.524, -0.778, -1.209)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG", 4, (0.616, 1.39, -0.0)], + ["CD", 5, (0.564, 1.414, 0.0)], + ["NE", 6, (0.539, 1.357, -0.0)], + ["NH1", 7, (0.206, 2.301, 0.0)], + ["NH2", 7, (2.078, 0.978, -0.0)], + ["CZ", 7, (0.758, 1.093, -0.0)], + ], + "ASN": [ + ["N", 0, (-0.536, 1.357, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.531, -0.787, -1.2)], + ["O", 3, (0.625, 1.062, 0.0)], + ["CG", 4, (0.584, 1.399, 0.0)], + ["ND2", 5, (0.593, -1.188, 0.001)], + ["OD1", 5, (0.633, 1.059, 0.0)], + ], + "ASP": [ + ["N", 0, (-0.525, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, 0.0, -0.0)], + ["CB", 0, (-0.526, -0.778, -1.208)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.593, 1.398, -0.0)], + ["OD1", 5, (0.61, 1.091, 0.0)], + ["OD2", 5, (0.592, -1.101, -0.003)], + ], + "CYS": [ + ["N", 0, (-0.522, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, 0.0, 0.0)], + ["CB", 0, (-0.519, -0.773, -1.212)], + ["O", 3, (0.625, 1.062, -0.0)], + ["SG", 4, (0.728, 1.653, 0.0)], + ], + "GLN": [ + ["N", 0, (-0.526, 1.361, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, 0.0)], + ["CB", 0, (-0.525, -0.779, -1.207)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.615, 1.393, 0.0)], + ["CD", 5, (0.587, 1.399, -0.0)], + ["NE2", 6, (0.593, -1.189, -0.001)], + ["OE1", 6, (0.634, 1.06, 0.0)], + ], + "GLU": [ + ["N", 0, (-0.528, 1.361, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.526, -0.781, -1.207)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG", 4, (0.615, 1.392, 0.0)], + ["CD", 5, (0.6, 1.397, 0.0)], + ["OE1", 6, (0.607, 1.095, -0.0)], + ["OE2", 6, (0.589, -1.104, -0.001)], + ], + "GLY": [ + ["N", 0, (-0.572, 1.337, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.517, -0.0, -0.0)], + ["O", 3, (0.626, 1.062, -0.0)], + ], + "HIS": [ + ["N", 0, (-0.527, 1.36, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, 0.0, 0.0)], + ["CB", 0, (-0.525, -0.778, -1.208)], + ["O", 3, (0.625, 1.063, 0.0)], + ["CG", 4, (0.6, 1.37, -0.0)], + ["CD2", 5, (0.889, -1.021, 0.003)], + ["ND1", 5, (0.744, 1.16, -0.0)], + ["CE1", 5, (2.03, 0.851, 0.002)], + ["NE2", 5, (2.145, -0.466, 0.004)], + ], + "ILE": [ + ["N", 0, (-0.493, 1.373, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, -0.0)], + ["CB", 0, (-0.536, -0.793, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG1", 4, (0.534, 1.437, -0.0)], + ["CG2", 4, (0.54, -0.785, -1.199)], + ["CD1", 5, (0.619, 1.391, 0.0)], + ], + "LEU": [ + ["N", 0, (-0.52, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.522, -0.773, -1.214)], + ["O", 3, (0.625, 1.063, -0.0)], + ["CG", 4, (0.678, 1.371, 0.0)], + ["CD1", 5, (0.53, 1.43, -0.0)], + ["CD2", 5, (0.535, -0.774, 1.2)], + ], + "LYS": [ + ["N", 0, (-0.526, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, 0.0)], + ["CB", 0, (-0.524, -0.778, -1.208)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.619, 1.39, 0.0)], + ["CD", 5, (0.559, 1.417, 0.0)], + ["CE", 6, (0.56, 1.416, 0.0)], + ["NZ", 7, (0.554, 1.387, 0.0)], + ], + "MET": [ + ["N", 0, (-0.521, 1.364, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, 0.0, 0.0)], + ["CB", 0, (-0.523, -0.776, -1.21)], + ["O", 3, (0.625, 1.062, -0.0)], + ["CG", 4, (0.613, 1.391, -0.0)], + ["SD", 5, (0.703, 1.695, 0.0)], + ["CE", 6, (0.32, 1.786, -0.0)], + ], + "PHE": [ + ["N", 0, (-0.518, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, 0.0, -0.0)], + ["CB", 0, (-0.525, -0.776, -1.212)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.607, 1.377, 0.0)], + ["CD1", 5, (0.709, 1.195, -0.0)], + ["CD2", 5, (0.706, -1.196, 0.0)], + ["CE1", 5, (2.102, 1.198, -0.0)], + ["CE2", 5, (2.098, -1.201, -0.0)], + ["CZ", 5, (2.794, -0.003, -0.001)], + ], + "PRO": [ + ["N", 0, (-0.566, 1.351, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, 0.0)], + ["CB", 0, (-0.546, -0.611, -1.293)], + ["O", 3, (0.621, 1.066, 0.0)], + ["CG", 4, (0.382, 1.445, 0.0)], + ["CD", 5, (0.477, 1.424, 0.0)], + ], + "SER": [ + ["N", 0, (-0.529, 1.36, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.518, -0.777, -1.211)], + ["O", 3, (0.626, 1.062, -0.0)], + ["OG", 4, (0.503, 1.325, 0.0)], + ], + "THR": [ + ["N", 0, (-0.517, 1.364, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, -0.0)], + ["CB", 0, (-0.516, -0.793, -1.215)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG2", 4, (0.55, -0.718, -1.228)], + ["OG1", 4, (0.472, 1.353, 0.0)], + ], + "TRP": [ + ["N", 0, (-0.521, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, 0.0)], + ["CB", 0, (-0.523, -0.776, -1.212)], + ["O", 3, (0.627, 1.062, 0.0)], + ["CG", 4, (0.609, 1.37, -0.0)], + ["CD1", 5, (0.824, 1.091, 0.0)], + ["CD2", 5, (0.854, -1.148, -0.005)], + ["CE2", 5, (2.186, -0.678, -0.007)], + ["CE3", 5, (0.622, -2.53, -0.007)], + ["NE1", 5, (2.14, 0.69, -0.004)], + ["CH2", 5, (3.028, -2.89, -0.013)], + ["CZ2", 5, (3.283, -1.543, -0.011)], + ["CZ3", 5, (1.715, -3.389, -0.011)], + ], + "TYR": [ + ["N", 0, (-0.522, 1.362, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, -0.0, -0.0)], + ["CB", 0, (-0.522, -0.776, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG", 4, (0.607, 1.382, -0.0)], + ["CD1", 5, (0.716, 1.195, -0.0)], + ["CD2", 5, (0.713, -1.194, -0.001)], + ["CE1", 5, (2.107, 1.2, -0.002)], + ["CE2", 5, (2.104, -1.201, -0.003)], + ["OH", 5, (4.168, -0.002, -0.005)], + ["CZ", 5, (2.791, -0.001, -0.003)], + ], + "VAL": [ + ["N", 0, (-0.494, 1.373, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, -0.0)], + ["CB", 0, (-0.533, -0.795, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG1", 4, (0.54, 1.429, -0.0)], + ["CG2", 4, (0.533, -0.776, 1.203)], + ], + "UNK": [ + ["N", 0, (-0.525, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ], +} +residue_atoms = { + "ALA": ["C", "CA", "CB", "N", "O"], + "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"], + "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"], + "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"], + "CYS": ["C", "CA", "CB", "N", "O", "SG"], + "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"], + "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"], + "GLY": ["C", "CA", "N", "O"], + "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"], + "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"], + "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"], + "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"], + "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"], + "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"], + "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"], + "SER": ["C", "CA", "CB", "N", "O", "OG"], + "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"], + "TRP": [ + "C", + "CA", + "CB", + "CG", + "CD1", + "CD2", + "CE2", + "CE3", + "CZ2", + "CZ3", + "CH2", + "N", + "NE1", + "O", + ], + "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"], + "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"], + "UNK": ["C", "CA", "N"], +} +residue_atom_renaming_swaps = { + "ASP": {"OD1": "OD2"}, + "GLU": {"OE1": "OE2"}, + "PHE": {"CD1": "CD2", "CE1": "CE2"}, + "TYR": {"CD1": "CD2", "CE1": "CE2"}, +} +van_der_waals_radius = {"C": 1.7, "N": 1.55, "O": 1.52, "S": 1.8} +Bond = namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"]) +BondAngle = namedtuple( + "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"] +) + +_STEREO_CHEMICAL_PROPS_PATH = Path("evolutionaryscale/structure/stereo_chemical_props.txt") + + +def _bond_key(atom_a: str, atom_b: str) -> tuple[str, str]: + """Return an order-independent key for a covalent bond.""" + + return (atom_a, atom_b) if atom_a <= atom_b else (atom_b, atom_a) + + +def _read_stereo_sections(text: str) -> tuple[list[str], list[str]]: + """Split the two tabular sections while discarding their headers.""" + + lines = iter(text.splitlines()) + next(lines) + bond_rows = list(iter(lambda: next(lines).strip(), "-")) + next(lines) + next(lines) + angle_rows = list(iter(lambda: next(lines).strip(), "-")) + return bond_rows, angle_rows + + +@functools.cache +def load_stereo_chemical_props() -> tuple[ + dict[str, list[Any]], dict[str, list[Any]], dict[str, list[Any]] +]: + """Load covalent geometry and derive virtual bonds for bond angles. + + The returned dictionaries are keyed by three-letter residue name. Virtual + bond lengths and uncertainties use the same operation order as the + checkpoint's reference feature pipeline so their floating-point values are + bitwise reproducible. + """ + + bond_rows, angle_rows = _read_stereo_sections(_STEREO_CHEMICAL_PROPS_PATH.read_text()) + residue_bonds: dict[str, list[Any]] = {} + for row in bond_rows: + atom_pair, residue_name, length, stddev = row.split() + atom_a, atom_b = atom_pair.split("-") + residue_bonds.setdefault(residue_name, []).append( + Bond(atom_a, atom_b, float(length), float(stddev)) + ) + residue_bonds["UNK"] = [] + + residue_bond_angles: dict[str, list[Any]] = {} + for row in angle_rows: + atom_triple, residue_name, angle_degrees, stddev_degrees = row.split() + atom_a, atom_b, atom_c = atom_triple.split("-") + residue_bond_angles.setdefault(residue_name, []).append( + BondAngle( + atom_a, + atom_b, + atom_c, + float(angle_degrees) / 180.0 * np.pi, + float(stddev_degrees) / 180.0 * np.pi, + ) + ) + residue_bond_angles["UNK"] = [] + + residue_virtual_bonds: dict[str, list[Any]] = {} + for residue_name, angles in residue_bond_angles.items(): + lookup = { + _bond_key(bond.atom1_name, bond.atom2_name): bond + for bond in residue_bonds[residue_name] + } + derived = residue_virtual_bonds.setdefault(residue_name, []) + for angle in angles: + left = lookup[_bond_key(angle.atom1_name, angle.atom2_name)] + right = lookup[_bond_key(angle.atom2_name, angle.atom3name)] + theta = angle.angle_rad + length = np.sqrt( + left.length**2 + right.length**2 - 2 * left.length * right.length * np.cos(theta) + ) + dl_outer = 0.5 / length + dl_dgamma = 2 * left.length * right.length * np.sin(theta) * dl_outer + dl_db1 = (2 * left.length - 2 * right.length * np.cos(theta)) * dl_outer + dl_db2 = (2 * right.length - 2 * left.length * np.cos(theta)) * dl_outer + stddev = np.sqrt( + (dl_dgamma * angle.stddev) ** 2 + + (dl_db1 * left.stddev) ** 2 + + (dl_db2 * right.stddev) ** 2 + ) + derived.append(Bond(angle.atom1_name, angle.atom3name, length, stddev)) + return residue_bonds, residue_virtual_bonds, residue_bond_angles + + +between_res_bond_length_c_n = [1.329, 1.341] +between_res_bond_length_stddev_c_n = [0.014, 0.016] +between_res_cos_angles_c_n_ca = [-0.5203, 0.0353] +between_res_cos_angles_ca_c_n = [-0.4473, 0.0311] +atom_types = [ + "N", + "CA", + "C", + "CB", + "O", + "CG", + "CG1", + "CG2", + "OG", + "OG1", + "SG", + "CD", + "CD1", + "CD2", + "ND1", + "ND2", + "OD1", + "OD2", + "SD", + "CE", + "CE1", + "CE2", + "CE3", + "NE", + "NE1", + "NE2", + "OE1", + "OE2", + "CH2", + "NH1", + "NH2", + "OH", + "CZ", + "CZ2", + "CZ3", + "NZ", + "OXT", +] +atom_order = {atom_type: i for (i, atom_type) in enumerate(atom_types)} +atom_type_num = len(atom_types) +restype_name_to_atom14_names = { + "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""], + "ARG": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD", + "NE", + "CZ", + "NH1", + "NH2", + "", + "", + "", + ], + "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""], + "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""], + "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""], + "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""], + "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""], + "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""], + "HIS": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "ND1", + "CD2", + "CE1", + "NE2", + "", + "", + "", + "", + ], + "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""], + "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""], + "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""], + "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""], + "PHE": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "CE1", + "CE2", + "CZ", + "", + "", + "", + ], + "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""], + "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""], + "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""], + "TRP": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "NE1", + "CE2", + "CE3", + "CZ2", + "CZ3", + "CH2", + ], + "TYR": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "CE1", + "CE2", + "CZ", + "OH", + "", + "", + ], + "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""], + "UNK": ["N", "CA", "C", "", "", "", "", "", "", "", "", "", "", ""], +} +restypes = [ + "A", + "R", + "N", + "D", + "C", + "Q", + "E", + "G", + "H", + "I", + "L", + "K", + "M", + "F", + "P", + "S", + "T", + "W", + "Y", + "V", +] +restype_order = {restype: i for (i, restype) in enumerate(restypes)} +restype_num = len(restypes) +unk_restype_index = restype_num +restypes_with_x = [*restypes, "X"] +restype_order_with_x = {restype: i for (i, restype) in enumerate(restypes_with_x)} +bb_atoms = ["N", "CA", "C", "O"] +hydrophobicity = { + "ALA": 0.116, + "ARG": -0.5, + "ASN": -0.264, + "ASP": -0.472, + "CYS": 0.18, + "GLN": -0.249, + "GLU": -0.457, + "GLY": 0.001, + "HIS": -0.335, + "ILE": 0.443, + "LEU": 0.443, + "LYS": -0.217, + "MET": 0.238, + "PHE": 0.5, + "PRO": 0.211, + "SER": -0.141, + "THR": -0.05, + "TRP": 0.378, + "TYR": 0.38, + "VAL": 0.325, +} +side_chain_asa = { + "ALA": 64.7809, + "ARG": 210.02, + "ASN": 113.187, + "ASP": 110.209, + "CYS": 95.2439, + "GLN": 147.855, + "GLU": 143.924, + "GLY": 23.1338, + "HIS": 146.449, + "ILE": 151.242, + "LEU": 139.524, + "LYS": 177.366, + "MET": 164.674, + "PHE": 186.7, + "PRO": 111.533, + "SER": 81.2159, + "THR": 111.597, + "TRP": 229.619, + "TYR": 200.306, + "VAL": 124.237, +} +amino_acid_volumes = { + "A": 88.6, + "R": 173.4, + "N": 114.1, + "D": 111.1, + "C": 108.5, + "Q": 143.8, + "E": 138.4, + "G": 60.1, + "H": 153.2, + "I": 166.7, + "L": 166.7, + "K": 168.6, + "M": 162.9, + "F": 189.9, + "P": 112.7, + "S": 89.0, + "T": 116.1, + "W": 227.8, + "Y": 193.6, + "V": 140.0, + "X": 88.6, +} + + +def sequence_to_onehot( + sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False +) -> np.ndarray: + """Encode a sequence as X with shape ``(l, n_alphabet)``. + + Unknown uppercase letters map to ``X`` only when ``map_unknown_to_x`` is + enabled. Non-alphabetic or lowercase input remains invalid in that mode. + """ + + n_alphabet = max(mapping.values()) + 1 + observed_indices = sorted(set(mapping.values())) + if observed_indices != list(range(n_alphabet)): + raise ValueError( + "The mapping must have values from 0 to num_unique_aas-1 without any gaps. " + f"Got: {sorted(mapping.values())}" + ) + + encoded: np.ndarray = np.empty(len(sequence), dtype=np.intp) + for position, symbol in enumerate(sequence): + if map_unknown_to_x: + if not symbol.isalpha() or not symbol.isupper(): + raise ValueError(f"Invalid character in the sequence: {symbol}") + encoded[position] = mapping.get(symbol, mapping["X"]) + else: + encoded[position] = mapping[symbol] + + one_hot: np.ndarray = np.zeros((len(sequence), n_alphabet), dtype=np.int32) + one_hot[np.arange(len(sequence)), encoded] = 1 + return one_hot + + +restype_1to3 = { + "A": "ALA", + "R": "ARG", + "N": "ASN", + "D": "ASP", + "C": "CYS", + "Q": "GLN", + "E": "GLU", + "G": "GLY", + "H": "HIS", + "I": "ILE", + "L": "LEU", + "K": "LYS", + "M": "MET", + "F": "PHE", + "P": "PRO", + "S": "SER", + "T": "THR", + "W": "TRP", + "Y": "TYR", + "V": "VAL", + "X": "UNK", +} +restype_3to1 = {v: k for (k, v) in restype_1to3.items()} +unk_restype = "UNK" +resnames = [restype_1to3[r] for r in restypes] + [unk_restype] +resname_to_idx = {resname: i for (i, resname) in enumerate(resnames)} +hydrophobic_resnames = {"VAL", "ILE", "LEU", "PHE", "MET", "TRP"} +HHBLITS_AA_TO_ID = { + "A": 0, + "B": 2, + "C": 1, + "D": 2, + "E": 3, + "F": 4, + "G": 5, + "H": 6, + "I": 7, + "J": 20, + "K": 8, + "L": 9, + "M": 10, + "N": 11, + "O": 20, + "P": 12, + "Q": 13, + "R": 14, + "S": 15, + "T": 16, + "U": 1, + "V": 17, + "W": 18, + "X": 20, + "Y": 19, + "Z": 3, + "-": 21, +} +ID_TO_HHBLITS_AA = { + 0: "A", + 1: "C", + 2: "D", + 3: "E", + 4: "F", + 5: "G", + 6: "H", + 7: "I", + 8: "K", + 9: "L", + 10: "M", + 11: "N", + 12: "P", + 13: "Q", + 14: "R", + 15: "S", + 16: "T", + 17: "V", + 18: "W", + 19: "Y", + 20: "X", + 21: "-", +} +restypes_with_x_and_gap = [*restypes, "X", "-"] +MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple( + restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) for i in range(len(restypes_with_x_and_gap)) +) + + +def _make_standard_atom_mask() -> np.ndarray: + """Return M with shape ``(n_residue_types, n_atom_types)``.""" + + mask: np.ndarray = np.zeros((restype_num + 1, atom_type_num), dtype=np.int32) + for residue_index, residue_code in enumerate(restypes): + residue_name = restype_1to3[residue_code] + atom_indices = [atom_order[name] for name in residue_atoms[residue_name]] + mask[residue_index, atom_indices] = 1 + return mask + + +STANDARD_ATOM_MASK = _make_standard_atom_mask() + + +def chi_angle_atom(atom_index: int) -> np.ndarray: + """Return chi-group selectors X with shape ``(21, 37, 4)``.""" + + selectors: list[np.ndarray] = [] + identity = np.eye(atom_type_num) + for residue_code in restypes: + groups = chi_angles_atoms[restype_1to3[residue_code]] + indices = [atom_order[group[atom_index]] for group in groups] + indices += [-1] * (4 - len(indices)) + selectors.append(identity[indices]) + selectors.append(np.zeros((4, atom_type_num))) + return cast(np.ndarray, np.stack(selectors).transpose(0, 2, 1)) + + +chi_atom_1_one_hot = chi_angle_atom(1) +chi_atom_2_one_hot = chi_angle_atom(2) +chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes] +chi_angles_atom_indices = np.array( + [chi_atoms + [[0, 0, 0, 0]] * (4 - len(chi_atoms)) for chi_atoms in chi_angles_atom_indices] +) +_chi_groups_for_atom: defaultdict[tuple[str, str], list[tuple[int, int]]] = defaultdict(list) +for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items(): + for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res): + for atom_i, atom in enumerate(chi_group): + _chi_groups_for_atom[res_name, atom].append((chi_group_i, atom_i)) +chi_groups_for_atom = dict(_chi_groups_for_atom) + + +def _make_rigid_transformation_4x4( + ex: np.ndarray, ey: np.ndarray, translation: np.ndarray +) -> np.ndarray: + """Construct homogeneous transform M from two basis vectors and an origin.""" + + unit_x = ex / np.linalg.norm(ex) + orthogonal_y = ey - np.dot(ey, unit_x) * unit_x + unit_y = orthogonal_y / np.linalg.norm(orthogonal_y) + unit_z = np.cross(unit_x, unit_y) + rotation_and_origin = np.stack((unit_x, unit_y, unit_z, translation), axis=0).T + homogeneous_row = np.asarray(((0.0, 0.0, 0.0, 1.0),)) + return cast(np.ndarray, np.concatenate((rotation_and_origin, homogeneous_row), axis=0)) + + +restype_atom37_to_rigid_group: np.ndarray = np.zeros((21, 37), dtype=int) +restype_atom37_mask: np.ndarray = np.zeros((21, 37), dtype=np.float32) +restype_atom37_rigid_group_positions: np.ndarray = np.zeros((21, 37, 3), dtype=np.float32) +restype_atom14_to_rigid_group: np.ndarray = np.zeros((21, 14), dtype=int) +restype_atom14_mask: np.ndarray = np.zeros((21, 14), dtype=np.float32) +restype_atom14_rigid_group_positions: np.ndarray = np.zeros((21, 14, 3), dtype=np.float32) +restype_rigid_group_default_frame: np.ndarray = np.zeros((21, 8, 4, 4), dtype=np.float32) + + +def _make_rigid_group_constants() -> None: + """Populate atom-to-frame assignments and default rigid transforms.""" + + for residue_index, residue_code in enumerate(restypes_with_x): + residue_name = restype_1to3[residue_code] + atom14_names = restype_name_to_atom14_names[residue_name] + for atom_name, group_index, coordinates in rigid_group_atom_positions[residue_name]: + atom37_index = atom_order[atom_name] + atom14_index = atom14_names.index(atom_name) + restype_atom37_to_rigid_group[residue_index, atom37_index] = group_index + restype_atom37_mask[residue_index, atom37_index] = 1 + restype_atom37_rigid_group_positions[residue_index, atom37_index] = coordinates + restype_atom14_to_rigid_group[residue_index, atom14_index] = group_index + restype_atom14_mask[residue_index, atom14_index] = 1 + restype_atom14_rigid_group_positions[residue_index, atom14_index] = coordinates + + positions = { + atom_name: np.asarray(coordinates) + for atom_name, _group_index, coordinates in rigid_group_atom_positions[residue_name] + } + restype_rigid_group_default_frame[residue_index, :2] = np.eye(4) + restype_rigid_group_default_frame[residue_index, 2] = _make_rigid_transformation_4x4( + positions["N"] - positions["CA"], + np.asarray((1.0, 0.0, 0.0)), + positions["N"], + ) + restype_rigid_group_default_frame[residue_index, 3] = _make_rigid_transformation_4x4( + positions["C"] - positions["CA"], + positions["CA"] - positions["N"], + positions["C"], + ) + + groups = chi_angles_atoms[residue_name] + if groups: + first_group = [positions[name] for name in groups[0]] + restype_rigid_group_default_frame[residue_index, 4] = _make_rigid_transformation_4x4( + first_group[2] - first_group[1], + first_group[0] - first_group[1], + first_group[2], + ) + for chi_index, group in enumerate(groups[1:], start=1): + axis_end = positions[group[2]] + restype_rigid_group_default_frame[residue_index, 4 + chi_index] = ( + _make_rigid_transformation_4x4( + axis_end, + np.asarray((-1.0, 0.0, 0.0)), + axis_end, + ) + ) + + +_make_rigid_group_constants() + + +def make_atom14_dists_bounds( + overlap_tolerance: float = 1.5, + bond_length_tolerance_factor: float = 15.0, +) -> dict[str, np.ndarray]: + """Build lower, upper, and uncertainty tensors with shape ``(21, 14, 14)``.""" + + lower_bounds: np.ndarray = np.zeros((21, 14, 14), np.float32) + upper_bounds: np.ndarray = np.zeros((21, 14, 14), np.float32) + stddevs: np.ndarray = np.zeros((21, 14, 14), np.float32) + residue_bonds, residue_virtual_bonds, _angles = load_stereo_chemical_props() + for residue_index, residue_code in enumerate(restypes): + residue_name = restype_1to3[residue_code] + atom_names = restype_name_to_atom14_names[residue_name] + for atom_a_index, atom_a_name in enumerate(atom_names): + if not atom_a_name: + continue + radius_a = van_der_waals_radius[atom_a_name[0]] + for atom_b_index, atom_b_name in enumerate(atom_names): + if not atom_b_name or atom_a_index == atom_b_index: + continue + clash_floor = radius_a + van_der_waals_radius[atom_b_name[0]] - overlap_tolerance + lower_bounds[residue_index, atom_a_index, atom_b_index] = clash_floor + lower_bounds[residue_index, atom_b_index, atom_a_index] = clash_floor + upper_bounds[residue_index, atom_a_index, atom_b_index] = 1e10 + upper_bounds[residue_index, atom_b_index, atom_a_index] = 1e10 + + for bond in residue_bonds[residue_name] + residue_virtual_bonds[residue_name]: + atom_a_index = atom_names.index(bond.atom1_name) + atom_b_index = atom_names.index(bond.atom2_name) + lower = bond.length - bond_length_tolerance_factor * bond.stddev + upper = bond.length + bond_length_tolerance_factor * bond.stddev + for row, column in ( + (atom_a_index, atom_b_index), + (atom_b_index, atom_a_index), + ): + lower_bounds[residue_index, row, column] = lower + upper_bounds[residue_index, row, column] = upper + stddevs[residue_index, row, column] = bond.stddev + return { + "lower_bound": lower_bounds, + "upper_bound": upper_bounds, + "stddev": stddevs, + } + + +restype_atom14_ambiguous_atoms: np.ndarray = np.zeros((21, 14), dtype=np.float32) +restype_atom14_ambiguous_atoms_swap_idx = np.tile(np.arange(14, dtype=int), (21, 1)) + + +def _make_atom14_ambiguity_feats() -> None: + """Mark symmetry-equivalent atom names and their exchange indices.""" + + for residue_name, swaps in residue_atom_renaming_swaps.items(): + residue_index = restype_order[restype_3to1[residue_name]] + atom_names = restype_name_to_atom14_names[residue_name] + for atom_a, atom_b in swaps.items(): + atom_a_index = atom_names.index(atom_a) + atom_b_index = atom_names.index(atom_b) + restype_atom14_ambiguous_atoms[residue_index, (atom_a_index, atom_b_index)] = 1 + restype_atom14_ambiguous_atoms_swap_idx[residue_index, (atom_a_index, atom_b_index)] = ( + atom_b_index, + atom_a_index, + ) + + +_make_atom14_ambiguity_feats() + + +def aatype_to_str_sequence(aatype: np.ndarray) -> str: + """Decode residue-type indices without changing their order.""" + + return "".join(restypes_with_x[index] for index in aatype) + + +CA_TO_N_NORM = 1.4591 +CA_TO_C_NORM = 1.5252 + + +def _make_restype_atom37_to_atom14() -> np.ndarray: + """Return atom37-to-atom14 lookup M with shape ``(21, 37)``.""" + + rows: list[list[int]] = [] + for residue_code in restypes: + names = restype_name_to_atom14_names[restype_1to3[residue_code]] + atom14_index = {name: index for index, name in enumerate(names)} + rows.append([atom14_index.get(name, 0) for name in atom_types]) + rows.append([0] * atom_type_num) + return np.asarray(rows, dtype=np.int32) + + +def _make_restype_atom14_to_atom37() -> np.ndarray: + """Return atom14-to-atom37 lookup M with shape ``(21, 14)``.""" + + rows = [ + [atom_order.get(name, 0) for name in restype_name_to_atom14_names[residue_name]] + for residue_name in resnames[:-1] + ] + rows.append([0] * 14) + return np.asarray(rows, dtype=np.int32) + + +RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37() +RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14() +CHAIN_BREAK_TOKEN = "|" diff --git a/src/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py b/src/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py new file mode 100644 index 0000000..97b1a47 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py @@ -0,0 +1,112 @@ +"""Dataclass support for aligned residue-level fields.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterable +from dataclasses import Field, dataclass, fields, replace +from typing import Any, Self + +import numpy as np + +from .esmfold2_misc import concat_objects, slice_any_object + +Index = int | list[int] | slice | np.ndarray + + +def _is_sequential(field: Field[Any]) -> bool: + return bool(field.metadata.get("sequence", False)) + + +def _sequence_axis(field: Field[Any]) -> int: + axis = int(field.metadata.get("sequence_dim", 0)) + if axis not in (0, 1): + raise NotImplementedError("SequentialDataclass supports sequence_dim values zero and one.") + return axis + + +def _slice_value(value: Any, index: Index, axis: int) -> Any: + if axis == 0: + return slice_any_object(value, index) + sliced = [slice_any_object(track, index) for track in value] + return value.__class__(sliced) + + +def _iter_sequence_lengths(value: Any, axis: int) -> Iterable[int]: + if axis == 0: + yield len(value) + else: + yield from (len(track) for track in value) + + +@dataclass(frozen=True) +class SequentialDataclass(ABC): + """Keep dataclass fields aligned along a shared residue dimension. + + A subclass marks aligned fields with ``metadata={"sequence": True}``. + ``sequence_dim`` may be zero for a direct sequence or one for a collection + of aligned tracks. ``join_token`` is passed to the package concatenation + helper when instances are joined. + """ + + def __post_init__(self) -> None: + expected = len(self) + for field in fields(self): + if not _is_sequential(field) or field.name == "complex": + continue + value = getattr(self, field.name) + if value is None: + continue + for actual in _iter_sequence_lengths(value, _sequence_axis(field)): + if actual != expected: + raise ValueError( + f"Mismatch in sequence length for field: {field.name}. " + f"Expected {expected}, received {actual}" + ) + + @abstractmethod + def __len__(self) -> int: + """Return the shared sequence length.""" + + raise NotImplementedError + + def __getitem__(self, index: Index) -> Self: + """Apply one sequence index to every aligned field.""" + + normalized_index: Index = [index] if isinstance(index, int) else index + updates: dict[str, Any] = {} + for field in fields(self): + if not _is_sequential(field): + continue + value = getattr(self, field.name) + if value is not None: + updates[field.name] = _slice_value(value, normalized_index, _sequence_axis(field)) + return replace(self, **updates) + + @classmethod + def concat(cls, items: list[Self], **overrides: Any) -> Self: + """Join aligned fields and retain non-sequential values from the first item.""" + + if not items: + raise ValueError("SequentialDataclass.concat requires at least one item.") + + updates: dict[str, Any] = {} + for field in fields(cls): + if not _is_sequential(field): + continue + first_value = getattr(items[0], field.name) + if first_value is None: + continue + values = [getattr(item, field.name) for item in items] + join_token = field.metadata.get("join_token") + if _sequence_axis(field) == 0: + updates[field.name] = concat_objects(values, join_token) + else: + tracks = [concat_objects(track, join_token) for track in zip(*values, strict=True)] + updates[field.name] = first_value.__class__(tracks) + + updates.update(overrides) + return replace(items[0], **updates) + + +__all__ = ["SequentialDataclass"] diff --git a/src/fastplms/models/esmfold2/esmfold2_system.py b/src/fastplms/models/esmfold2/esmfold2_system.py new file mode 100644 index 0000000..23b69dc --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_system.py @@ -0,0 +1,56 @@ +"""Filesystem and subprocess types used by optional structure utilities.""" + +from __future__ import annotations + +import io +import subprocess +from pathlib import Path +from typing import Any, TypeAlias + +PathLike: TypeAlias = str | Path +PathOrBuffer: TypeAlias = PathLike | io.StringIO + + +def _stdout_destination(*, capture_output: bool, quiet: bool) -> int | None: + if capture_output: + return subprocess.PIPE + if quiet: + return subprocess.DEVNULL + return None + + +def _stderr_text(error: subprocess.CalledProcessError) -> str: + stderr = error.stderr + if stderr is None: + return "" + if isinstance(stderr, bytes): + return stderr.decode(errors="replace") + return str(stderr) + + +def run_subprocess_with_errorcheck( + *popenargs: Any, + capture_output: bool = False, + quiet: bool = False, + env: dict[str, str] | None = None, + shell: bool = False, + executable: str | None = None, + **kwargs: Any, +) -> subprocess.CompletedProcess[Any]: + """Run a command and include captured standard error in failures.""" + + stdout = _stdout_destination(capture_output=capture_output, quiet=quiet) + try: + return subprocess.run( + *popenargs, + check=True, + env=env, + executable=executable, + shell=shell, + stderr=subprocess.PIPE, + stdout=stdout, + **kwargs, + ) + except subprocess.CalledProcessError as error: + message = f"Command failed with errorcode {error.returncode}.\n\n{_stderr_text(error)}" + raise RuntimeError(message) from error diff --git a/src/fastplms/models/esmfold2/esmfold2_types.py b/src/fastplms/models/esmfold2/esmfold2_types.py new file mode 100644 index 0000000..38d581e --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_types.py @@ -0,0 +1,31 @@ +"""Stable namespace for ESMFold2 input schema types.""" + +from __future__ import annotations + +from . import esmfold2_input_builder as _input_schema +from .esmfold2_msa import MSA +from .esmfold2_parsing import FastaEntry + +Modification = _input_schema.Modification +ProteinInput = _input_schema.ProteinInput +RNAInput = _input_schema.RNAInput +DNAInput = _input_schema.DNAInput +LigandInput = _input_schema.LigandInput +DistogramConditioning = _input_schema.DistogramConditioning +PocketConditioning = _input_schema.PocketConditioning +CovalentBond = _input_schema.CovalentBond +StructurePredictionInput = _input_schema.StructurePredictionInput + +__all__ = [ + "MSA", + "CovalentBond", + "DNAInput", + "DistogramConditioning", + "FastaEntry", + "LigandInput", + "Modification", + "PocketConditioning", + "ProteinInput", + "RNAInput", + "StructurePredictionInput", +] diff --git a/src/fastplms/models/esmfold2/esmfold2_utils_types.py b/src/fastplms/models/esmfold2/esmfold2_utils_types.py new file mode 100644 index 0000000..1359024 --- /dev/null +++ b/src/fastplms/models/esmfold2/esmfold2_utils_types.py @@ -0,0 +1,38 @@ +"""Small public types shared by the ESMFold2 structure utilities. + +These definitions are intentionally independent of cloud-storage packages. Any +path object implementing :class:`os.PathLike` is accepted by the runtime file +helpers, including cloud-path implementations installed by an application. +""" + +from __future__ import annotations + +import io +import os +from dataclasses import dataclass +from typing import TypeAlias + +PathLike: TypeAlias = str | os.PathLike[str] +PathOrBuffer: TypeAlias = PathLike | io.TextIOBase + + +@dataclass(slots=True) +class FunctionAnnotation: + """A residue-range annotation using one-based inclusive coordinates.""" + + label: str + start: int + end: int + + def to_tuple(self) -> tuple[str, int, int]: + """Return the serialization order used by annotation tokenizers.""" + + return (self.label, self.start, self.end) + + def __len__(self) -> int: + """Return the number of annotated residues.""" + + return self.end - self.start + 1 + + +__all__ = ["FunctionAnnotation", "PathLike", "PathOrBuffer"] diff --git a/fastplms/esmfold2/modeling_esmfold2.py b/src/fastplms/models/esmfold2/modeling_esmfold2.py similarity index 59% rename from fastplms/esmfold2/modeling_esmfold2.py rename to src/fastplms/models/esmfold2/modeling_esmfold2.py index daa88ff..2a9b475 100644 --- a/fastplms/esmfold2/modeling_esmfold2.py +++ b/src/fastplms/models/esmfold2/modeling_esmfold2.py @@ -1,4 +1,4 @@ -"""PyTorch ESMFold2 model — the standard released architecture. +"""PyTorch ESMFold2 model: the standard released architecture. Quickstart:: @@ -11,41 +11,49 @@ with ``model.fold(...)`` or ``model.prepare_structure_input(...)``. """ +from __future__ import annotations + +import gc import importlib +import importlib.metadata import math from contextlib import contextmanager +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, cast +from typing import Any, ClassVar, Literal, cast import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor - -try: - te = importlib.import_module("transformer_engine.pytorch") - te_recipe = importlib.import_module("transformer_engine.common.recipe") - DelayedScaling = te_recipe.DelayedScaling - Format = te_recipe.Format - - TE_AVAILABLE = True -except ImportError: - te = None # type: ignore[assignment] - DelayedScaling = None # type: ignore[assignment] - Format = None # type: ignore[assignment] - TE_AVAILABLE = False - +from transformers.modeling_outputs import ModelOutput from transformers.modeling_utils import PreTrainedModel +from ...attention import get_attn_implementation, set_config_attn_implementation + try: - from fastplms.test_time_training import FastPLMTestTimeTrainingMixin, TTTConfig -except ImportError: - from .test_time_training import FastPLMTestTimeTrainingMixin, TTTConfig + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin, TTTConfig +except ModuleNotFoundError as error: + if error.name != "fastplms": + raise + from ..ttt import FastPLMTestTimeTrainingMixin, TTTConfig +from .attention import ESMFold2AttentionMixin from .configuration_esmfold2 import ESMFold2Config, normalize_esmc_id +from .embedding import ESMFold2EmbeddingMixin +from .esmfold2_constants_esm3 import ( + SEQUENCE_BOS_TOKEN, + SEQUENCE_EOS_TOKEN, + SEQUENCE_MASK_TOKEN, + SEQUENCE_PAD_TOKEN, + SEQUENCE_STANDARD_AA_MAX_TOKEN, + SEQUENCE_STANDARD_AA_MIN_TOKEN, + SEQUENCE_VOCAB, +) from .modeling_esmfold2_common import ( CHAR_VOCAB_SIZE, MAX_ATOMIC_NUMBER, + MSA_CONDITIONING_INPUT_NAMES, NUM_RES_TYPES, DiffusionStructureHead, FoldingTrunk, @@ -64,56 +72,100 @@ gather_token_to_atom, maybe_apply_msa_column_masking, maybe_subsample_msa, + validate_kernel_backend, + validate_msa_conditioning_inputs, ) -from .esmfold2_affine3d import Affine3D as _FastPLMSESMFold2Affine3D -from .esmfold2_aligner import Aligner as _FastPLMSESMFold2Aligner -from .esmfold2_atom_indexer import AtomIndexer as _FastPLMSESMFold2AtomIndexer -from .esmfold2_conformers import load_ccd as _fastplms_esmfold2_load_ccd -from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL as _FASTPLMS_ESMFOLD2_ELEMENT_NUMBER_TO_SYMBOL -from .esmfold2_constants_esm3 import ( - CHAIN_BREAK_STR as _FASTPLMS_ESMFOLD2_CHAIN_BREAK_STR, - SEQUENCE_BOS_TOKEN, - SEQUENCE_EOS_TOKEN, - SEQUENCE_MASK_TOKEN, - SEQUENCE_PAD_TOKEN, - SEQUENCE_STANDARD_AA_MAX_TOKEN, - SEQUENCE_STANDARD_AA_MIN_TOKEN, - SEQUENCE_VOCAB, -) -from .esmfold2_input_builder import StructurePredictionInput as _FastPLMSESMFold2StructurePredictionInput -from .esmfold2_metrics import compute_rmsd as _fastplms_esmfold2_compute_rmsd -from .esmfold2_misc import slice_any_object as _fastplms_esmfold2_slice_any_object -from .esmfold2_mmcif_parsing import MmcifWrapper as _FastPLMSESMFold2MmcifWrapper -from .esmfold2_molecular_complex import MolecularComplex as _FastPLMSESMFold2MolecularComplex -from .esmfold2_msa import MSA as _FastPLMSESMFold2MSA -from .esmfold2_msa_filter_sequences import greedy_select_indices as _fastplms_esmfold2_greedy_select_indices -from .esmfold2_normalize_coordinates import normalize_coordinates as _fastplms_esmfold2_normalize_coordinates -from .esmfold2_output import build_molecular_complex_from_features as _fastplms_esmfold2_build_molecular_complex_from_features -from .esmfold2_paired_msa import construct_paired_msa as _fastplms_esmfold2_construct_paired_msa -from .esmfold2_parsing import FastaEntry as _FastPLMSESMFold2FastaEntry -from .esmfold2_predicted_aligned_error import compute_tm as _fastplms_esmfold2_compute_tm -from .esmfold2_prepare_input import prepare_esmfold2_input as _fastplms_esmfold2_prepare_esmfold2_input -from .esmfold2_processor import ESMFold2InputBuilder as _FastPLMSESMFold2InputBuilder -from .esmfold2_protein_chain import ProteinChain as _FastPLMSESMFold2ProteinChain -from .esmfold2_protein_complex import ProteinComplex as _FastPLMSESMFold2ProteinComplex -from .esmfold2_protein_structure import index_by_atom_name as _fastplms_esmfold2_index_by_atom_name -from .esmfold2_residue_constants import restypes as _FASTPLMS_ESMFOLD2_RESTYPES -from .esmfold2_sequential_dataclass import SequentialDataclass as _FastPLMSESMFold2SequentialDataclass -from .esmfold2_system import run_subprocess_with_errorcheck as _fastplms_esmfold2_run_subprocess_with_errorcheck -from .esmfold2_types import ProteinInput as _FastPLMSESMFold2ProteinInput -from .esmfold2_utils_types import PathOrBuffer as _FastPLMSESMFold2PathOrBuffer +_ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj" +_ESMC_FP8_EXPECTED_PROJECTIONS = 80 _EPS = 1e-6 _NONPOLYMER_ID = 4 -# Default for the triangle / OPM / pair-transition L² ops. Caps peak memory -# so L≈2k folds on an 80 GB GPU (~76 GB peak at chunk=128 for L=1438; +# Default for the triangle, OPM, and pair-transition l^2 operations. Caps peak +# memory so l around 2k folds on an 80 GB GPU (about 76 GB at chunk=128 for +# l=1438; # chunk=64 leaves headroom for the largest foldbench targets). Override via # ``model.set_chunk_size(...)``; pass None to disable chunking (faster for -# short L but OOM-prone past ~600). +# short l but OOM-prone past approximately 600). _DEFAULT_CHUNK_SIZE = 64 +@dataclass +class ESMFold2Output(ModelOutput): + """Transformers-compatible output shared by released and experimental folds. + + ``last_hidden_state`` is the final pair representation. When requested, + ``hidden_states`` contains the token-input representation followed by the + final pair representation. The structure trunks do not expose normalized + post-softmax attention tensors, so ``output_attentions=True`` fails + explicitly instead of returning incomplete data. + """ + + last_hidden_state: Tensor | None = None + hidden_states: tuple[Tensor, ...] | None = None + attentions: tuple[Tensor, ...] | None = None + distogram_logits: Tensor | None = None + sample_atom_coords: Tensor | None = None + representative_atom_coords: Tensor | None = None + atom_pad_mask: Tensor | None = None + residue_index: Tensor | None = None + entity_id: Tensor | None = None + plddt_logits: Tensor | None = None + plddt: Tensor | None = None + plddt_per_atom: Tensor | None = None + plddt_ca: Tensor | None = None + complex_plddt: Tensor | None = None + complex_iplddt: Tensor | None = None + pae_logits: Tensor | None = None + pae: Tensor | None = None + pde_logits: Tensor | None = None + pde: Tensor | None = None + resolved_logits: Tensor | None = None + ptm: Tensor | None = None + iptm: Tensor | None = None + pair_chains_iptm: Tensor | None = None + + +def _resolve_structure_output_controls( + config: ESMFold2Config, + *, + output_attentions: bool | None, + output_hidden_states: bool | None, + return_dict: bool | None, +) -> tuple[bool, bool]: + resolved_attentions = ( + config.output_attentions if output_attentions is None else output_attentions + ) + if resolved_attentions: + raise NotImplementedError( + "ESMFold2 does not expose normalized attention tensors from its structure " + "trunk. output_attentions=True is unsupported." + ) + resolved_hidden_states = ( + config.output_hidden_states + if output_hidden_states is None + else output_hidden_states + ) + resolved_return_dict = config.use_return_dict if return_dict is None else return_dict + return bool(resolved_hidden_states), bool(resolved_return_dict) + + +def _finalize_structure_output( + output: dict[str, Tensor], + *, + token_input_state: Tensor, + pair_state: Tensor, + output_hidden_states: bool, + return_dict: bool, +) -> ESMFold2Output | tuple[Any, ...]: + model_output = ESMFold2Output( + last_hidden_state=pair_state, + hidden_states=(token_input_state, pair_state) if output_hidden_states else None, + **output, + ) + return model_output if return_dict else model_output.to_tuple() + + class _ESMFold2ESMplusplusAdapter(nn.Module): def __init__(self, model: nn.Module) -> None: super().__init__() @@ -123,6 +175,11 @@ def __init__(self, model: nn.Module) -> None: def config(self): return self.model.config + def set_attn_implementation(self, attn_implementation: str) -> None: + """Update ESMC through its Transformers-compatible attention API.""" + + self.model.set_attn_implementation(attn_implementation) + def forward( self, input_ids: Tensor, @@ -146,7 +203,8 @@ def forward( ) if output_hidden_states: hidden_states = output.hidden_states - assert hidden_states is not None, "ESM++ did not return hidden states." + if hidden_states is None: + raise RuntimeError("ESM++ did not return requested hidden states.") if isinstance(hidden_states, torch.Tensor): output.hidden_states = hidden_states else: @@ -159,23 +217,298 @@ def _load_fastplms_esmplusplus_for_esmfold2( attn_backend: str, device: torch.device, dtype: torch.dtype, + local_files_only: bool = False, ) -> _ESMFold2ESMplusplusAdapter: + from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, + ) + + normalized_path = normalize_esmc_id(esmc_model_path) + source_revision, _ = _manifest_esmc_checkpoint_contract(normalized_path) + revision_kwargs: dict[str, Any] = { + "local_files_only": local_files_only, + } + if source_revision is not None: + revision_kwargs["revision"] = source_revision + esmc_config = ESMplusplusConfig.from_pretrained(normalized_path, **revision_kwargs) + set_config_attn_implementation(esmc_config, attn_backend) + load_kwargs: dict[str, Any] = { + "config": esmc_config, + "torch_dtype": dtype, + **revision_kwargs, + } + if device.type == "cuda": + # Device mapping constructs parameters on the destination GPU instead + # of materializing the 6B backbone in host memory first. + load_kwargs["device_map"] = {"": str(device)} + esmc = ESMplusplusModel.from_pretrained(normalized_path, **load_kwargs) + if device.type != "cuda": + esmc = esmc.to(device=device, dtype=dtype) + else: + loaded_device = next(esmc.parameters()).device + if loaded_device != device: + raise RuntimeError( + f"ESMC loaded on {loaded_device}, expected direct loading on {device}." + ) + return _ESMFold2ESMplusplusAdapter(esmc).eval() + + +def _manifest_esmc_checkpoint_contract( + esmc_model_path: str, +) -> tuple[str | None, dict[str, str]]: + """Return the immutable manifest identity for a registered ESMC source. + + Local checkpoint directories deliberately return no Hub revision. A known + Hub repository is always loaded at the revision and file identities in + ``models.toml`` instead of following a mutable branch. + """ + + normalized_path = normalize_esmc_id(esmc_model_path) try: - from fastplms.esm_plusplus.modeling_esm_plusplus import ( - ESMplusplusConfig, - ESMplusplusModel, + if Path(normalized_path).exists(): + return None, {} + except OSError: + # A repository ID may be too long or otherwise invalid as a local path. + pass + + from fastplms.registry import get_model_registry + + registry = get_model_registry() + backbone_model = registry.families["esmfold2"].backbone_model + if backbone_model is None: + raise RuntimeError("families.esmfold2 must declare backbone_model.") + spec = registry[backbone_model] + for checkpoint in (spec.fast, spec.official): + if checkpoint.repo_id == normalized_path: + return checkpoint.revision, { + item.path: item.encoded for item in checkpoint.files + } + if "/" in normalized_path: + raise ValueError( + f"Remote ESMC source {normalized_path!r} is not the manifest-declared " + f"ESMFold2 backbone {spec.fast.repo_id!r}." ) - except ImportError: - from .modeling_esm_plusplus import ESMplusplusConfig, ESMplusplusModel + return None, {} - normalized_path = normalize_esmc_id(esmc_model_path) - esmc_config = ESMplusplusConfig.from_pretrained(normalized_path) - esmc_config.attn_backend = attn_backend - esmc = ESMplusplusModel.from_pretrained( - normalized_path, - config=esmc_config, + +ESMCPrecision = Literal["auto", "bf16", "fp32", "fp8"] + + +@dataclass(frozen=True, slots=True) +class ESMCPrecisionStatus: + """Resolved ESMC precision and the evidence used to choose it.""" + + requested: str + resolved: str + reason: str + device: str + transformer_engine_version: str | None + + def as_dict(self) -> dict[str, str | None]: + return asdict(self) + + +def _transformer_engine_version() -> str | None: + try: + return importlib.metadata.version("transformer-engine") + except importlib.metadata.PackageNotFoundError: + return None + + +def _load_transformer_engine() -> tuple[Any, Any]: + """Load Transformer Engine lazily so core imports stay dependency-free.""" + + try: + te = importlib.import_module("transformer_engine.pytorch") + recipe = importlib.import_module("transformer_engine.common.recipe") + except (ImportError, OSError, RuntimeError) as error: + raise RuntimeError( + f"Transformer Engine could not be imported: {type(error).__name__}: {error}" + ) from error + if not hasattr(recipe, "Float8CurrentScaling"): + raise RuntimeError( + "Transformer Engine does not expose Float8CurrentScaling, which is " + "required by the validated ESMC FP8 path." + ) + return te, recipe + + +def _te_fp8_capability(device: torch.device) -> tuple[bool, str]: + """Return whether the validated Transformer Engine FP8 path can run.""" + + if device.type != "cuda": + return False, "FP8 requires direct ESMC loading onto a CUDA device." + if not torch.cuda.is_available(): + return False, "CUDA is unavailable." + try: + major, minor = torch.cuda.get_device_capability(device) + except (AssertionError, RuntimeError, ValueError) as error: + return False, f"CUDA capability query failed: {error}" + if not (major >= 9 or (major == 8 and minor >= 9)): + return False, f"CUDA capability {major}.{minor} does not support FP8." + try: + te, _ = _load_transformer_engine() + except RuntimeError as error: + return False, str(error) + + probe = getattr(te, "is_fp8_available", None) + if probe is None: + try: + probe = importlib.import_module("transformer_engine.pytorch.fp8").is_fp8_available + except (ImportError, AttributeError, OSError, RuntimeError) as error: + return False, f"Transformer Engine has no usable FP8 probe: {error}" + try: + try: + result = probe(return_reason=True) + except TypeError: + result = probe() + except (OSError, RuntimeError) as error: + return False, f"Transformer Engine FP8 probe failed: {error}" + if isinstance(result, tuple): + available = bool(result[0]) + detail = str(result[1]) if len(result) > 1 and result[1] else "" + else: + available = bool(result) + detail = "" + if not available: + return False, detail or "Transformer Engine reports FP8 unavailable." + return True, ( + "Transformer Engine reports FP8 availability; FastPLMs will convert " + "the validated ESMC attention output projections." + ) + + +def _resolve_esmc_precision(requested: str, device: torch.device) -> ESMCPrecisionStatus: + allowed = {"auto", "bf16", "fp32", "fp8"} + if requested not in allowed: + raise ValueError(f"precision must be one of {sorted(allowed)}, got {requested!r}.") + if requested in {"auto", "bf16", "fp32"}: + resolved = "bf16" if requested == "auto" else requested + reason = ( + "Automatic precision defaults to BF16; select esmc_precision='fp8' " + "explicitly to opt in to the validated Transformer Engine path." + if requested == "auto" + else "Precision was selected explicitly." + ) + return ESMCPrecisionStatus( + requested=requested, + resolved=resolved, + reason=reason, + device=str(device), + transformer_engine_version=_transformer_engine_version(), + ) + available, reason = _te_fp8_capability(device) + if not available: + raise RuntimeError(f"esmc_precision='fp8' is unavailable: {reason}") + return ESMCPrecisionStatus( + requested=requested, + resolved="fp8", + reason=reason, + device=str(device), + transformer_engine_version=_transformer_engine_version(), + ) + + +def _install_esmc_backbone( + model: Any, + esmc_model_path: str, + *, + precision: str, + device: str | torch.device | None = None, + local_files_only: bool = False, +) -> None: + target_device = torch.device(device) if device is not None else model.device + if target_device.type == "cuda" and target_device.index is None and torch.cuda.is_available(): + target_device = torch.device("cuda", torch.cuda.current_device()) + model_device = torch.device(model.device) + if model_device.type == "cuda" and model_device.index is None and torch.cuda.is_available(): + model_device = torch.device("cuda", torch.cuda.current_device()) + if target_device != model_device: + raise ValueError( + f"ESMC target device {target_device} must match the ESMFold2 device " + f"{model_device}. Move ESMFold2 before loading or reloading ESMC." + ) + status = _resolve_esmc_precision(precision, target_device) + normalized_source = normalize_esmc_id(esmc_model_path) + source_revision, source_files = _manifest_esmc_checkpoint_contract(normalized_source) + attention_implementation = get_attn_implementation(model.config) + model.config.esmc_attn_backend = attention_implementation + dtype = torch.float32 if status.resolved == "fp32" else torch.bfloat16 + esmc = _load_fastplms_esmplusplus_for_esmfold2( + esmc_model_path=esmc_model_path, + attn_backend=attention_implementation, + device=target_device, + dtype=dtype, + local_files_only=local_files_only, + ) + if esmc.config.hidden_size != model.config.lm_d_model: + raise ValueError( + f"ESMFold2 expected lm_d_model={model.config.lm_d_model}, " + f"but loaded ESMC hidden_size={esmc.config.hidden_size}." + ) + if esmc.config.num_hidden_layers != model.config.lm_num_layers: + raise ValueError( + f"ESMFold2 expected lm_num_layers={model.config.lm_num_layers}, " + f"but loaded ESMC num_hidden_layers={esmc.config.num_hidden_layers}." + ) + esmc.eval().requires_grad_(False) + fp8_module_paths: tuple[str, ...] = () + if status.resolved == "fp8": + fp8_module_paths = _convert_esmc_attention_outputs_to_te(esmc) + status = ESMCPrecisionStatus( + requested=status.requested, + resolved=status.resolved, + reason=( + f"{status.reason} Converted {len(fp8_module_paths)} projections; " + "canonical checkpoint weights remain BF16." + ), + device=status.device, + transformer_engine_version=status.transformer_engine_version, + ) + model._esmc_source = normalized_source + model._esmc_source_revision = source_revision + model._esmc_source_files = source_files + model._esmc_local_files_only = local_files_only + model._esmc_precision_policy = precision + model._esmc_precision_status = status + model._esmc_fp8 = status.resolved == "fp8" + model._esmc_fp8_module_paths = fp8_module_paths + model.config.esmc_precision = precision + model._esmc = esmc + model._ttt_lm_head = None + + +def _drop_transient_esmc_state( + module: nn.Module, + state_dict: dict[str, Tensor], + prefix: str, + local_metadata: dict[str, Any], +) -> None: + """Exclude runtime ESMC/TTT modules from canonical folding checkpoints.""" + + del module, local_metadata + transient_prefixes = (f"{prefix}_esmc.", f"{prefix}_ttt_lm_head.") + for key in tuple(state_dict): + if key.startswith(transient_prefixes): + del state_dict[key] + + +def _reload_esmc_bf16_for_gradients(model: Any, *, reason: str) -> None: + """Use BF16 temporarily without overwriting the persisted serving policy.""" + + policy = model._esmc_precision_policy + model.reload_esmc(precision="bf16", device=model.device) + model._esmc_precision_policy = policy + model.config.esmc_precision = policy + status = model._esmc_precision_status + model._esmc_precision_status = ESMCPrecisionStatus( + requested=policy, + resolved="bf16", + reason=reason, + device=status.device, + transformer_engine_version=status.transformer_engine_version, ) - return _ESMFold2ESMplusplusAdapter(esmc).to(device=device, dtype=dtype).eval() class PairTransition(nn.Module): @@ -206,7 +539,7 @@ class ConfidenceHead(nn.Module): boundaries: Tensor - def __init__(self, config: "ESMFold2Config") -> None: + def __init__(self, config: ESMFold2Config) -> None: super().__init__() ch = config.confidence_head d_single = config.d_single @@ -228,14 +561,10 @@ def __init__(self, config: "ESMFold2Config") -> None: self.s_inputs_norm = nn.LayerNorm(d_inputs) self.z_norm = nn.LayerNorm(d_pair) - self.row_attention_pooling = RowAttentionPooling( - d_pair=d_pair, d_single=d_single - ) + self.row_attention_pooling = RowAttentionPooling(d_pair=d_pair, d_single=d_single) pf = ch.folding_trunk - self.folding_trunk = FoldingTrunk( - n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4 - ) + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) # Heads. self.plddt_ln = nn.LayerNorm(d_single) @@ -252,9 +581,7 @@ def __init__(self, config: "ESMFold2Config") -> None: self.resolved_ln = nn.LayerNorm(d_single) # 2 = resolved logits ([unresolved, resolved]). - self.resolved_weight = nn.Parameter( - torch.zeros(max_atoms_per_token, d_single, 2) - ) + self.resolved_weight = nn.Parameter(torch.zeros(max_atoms_per_token, d_single, 2)) def set_kernel_backend(self, backend: str | None) -> None: self.folding_trunk.set_kernel_backend(backend) @@ -264,11 +591,7 @@ def set_chunk_size(self, chunk_size: int | None) -> None: @staticmethod def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor: - return ( - x - if num_diffusion_samples == 1 - else x.repeat_interleave(num_diffusion_samples, 0) - ) + return x if num_diffusion_samples == 1 else x.repeat_interleave(num_diffusion_samples, 0) @staticmethod def _flatten_sample_axis(x: Tensor) -> Tensor: @@ -312,15 +635,13 @@ def forward( atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) - Bm = pair.shape[0] + expanded_batch_size = pair.shape[0] rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) rep_distances = torch.cdist( rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" ) - distogram_bins = ( - (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() - ) + distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() pair = pair + self.dist_bin_pairwise_embed(distogram_bins) pair_mask = mask[:, :, None].float() * mask[:, None, :].float() @@ -344,10 +665,18 @@ def forward( plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt) plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0) - L = single.shape[1] - plddt_sum = torch.zeros(Bm, L, device=single.device, dtype=plddt_per_atom.dtype) + sequence_length = single.shape[1] + plddt_sum = torch.zeros( + expanded_batch_size, + sequence_length, + device=single.device, + dtype=plddt_per_atom.dtype, + ) atom_count = torch.zeros( - Bm, L, device=single.device, dtype=plddt_per_atom.dtype + expanded_batch_size, + sequence_length, + device=single.device, + dtype=plddt_per_atom.dtype, ) atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype) plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t) @@ -361,13 +690,11 @@ def forward( expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) is_ligand = (expanded_type == _NONPOLYMER_ID).float() - inter_chain = ( - expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) - ).float() + inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float() near_contact = (rep_distances < 8).float() - interface_per_token = ( - near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1) - ).amax(dim=-1) + interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax( + dim=-1 + ) iplddt_weight = torch.where( is_ligand.bool(), torch.full_like(interface_per_token, 2.0), @@ -399,20 +726,16 @@ def forward( # pTM / ipTM from pae_logits. n_bins = pae_logits.shape[-1] bin_width = 32.0 / n_bins - bin_centers = torch.arange( - 0.5 * bin_width, 32.0, bin_width, device=pae_logits.device - ) + bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device) mask_f = mask.float() - N_res = mask_f.sum(dim=-1, keepdim=True) - d0 = 1.24 * (N_res.clamp(min=19) - 15) ** (1 / 3) - 1.8 + n_residues = mask_f.sum(dim=-1, keepdim=True) + d0 = 1.24 * (n_residues.clamp(min=19) - 15) ** (1 / 3) - 1.8 tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2) pae_probs = F.softmax(pae_logits, dim=-1) tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1) pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) - ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / ( - pair_mask_2d.sum(dim=-1) + _EPS - ) + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS) ptm = ptm_per_row.max(dim=-1).values inter_chain_mask = ( @@ -423,10 +746,14 @@ def forward( ) iptm = iptm_per_row.max(dim=-1).values - max_chain_id = int(expanded_asym.max().item()) if Bm > 0 else 0 + max_chain_id = int(expanded_asym.max().item()) if expanded_batch_size > 0 else 0 n_chains = max_chain_id + 1 pair_chains_iptm = torch.zeros( - Bm, n_chains, n_chains, device=tm_expected.device, dtype=tm_expected.dtype + expanded_batch_size, + n_chains, + n_chains, + device=tm_expected.device, + dtype=tm_expected.dtype, ) for c1 in range(n_chains): chain_c1 = (expanded_asym == c1).float() * mask_f @@ -436,9 +763,7 @@ def forward( chain_c2 = (expanded_asym == c2).float() * mask_f pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) denom = pair_m.sum(dim=(-1, -2)) + _EPS - pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum( - dim=(-1, -2) - ) / denom + pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom return { "plddt_logits": plddt_logits, @@ -462,145 +787,76 @@ def _inverse_softplus(value: float) -> float: return value + math.log(-math.expm1(-value)) -def _convert_te_modules_to_fp8_inplace(module: nn.Module) -> None: - """Re-init each TE module via quantized_model_init so weights live as fp8. +def _convert_esmc_attention_outputs_to_te(module: nn.Module) -> tuple[str, ...]: + """Replace the 80 ESMC attention output projections with TE linears. - Must be called inside torch.no_grad(); covers nn.Linear, te.Linear, - te.LayerNormLinear, te.LayerNormMLP — the last two hold 99% of ESMC weight. + Converting every ESMC linear compounds FP8 error across the 80-layer + network. The validated inference path limits FP8 GEMMs to each layer's + attention output projection. Transformer Engine retains canonical BF16 + parameters and creates runtime quantization workspaces during autocast. """ - if not TE_AVAILABLE: - raise RuntimeError("transformer_engine is not available; cannot use fp8.") - quantized_model_init = importlib.import_module( - "transformer_engine.pytorch" - ).quantized_model_init - - def _walk(mod: nn.Module) -> None: - for name, child in list(mod.named_children()): - replaced = False - if isinstance(child, nn.Linear): - in_f, out_f = child.in_features, child.out_features - has_bias = child.bias is not None - device = child.weight.device - dtype = child.weight.dtype - w = child.weight.data - b = child.bias.data if has_bias else None - setattr(mod, name, nn.Identity()) - del child - torch.cuda.empty_cache() - with quantized_model_init(enabled=True): - new_mod = te.Linear( # type: ignore[union-attr] - in_f, out_f, bias=has_bias, params_dtype=dtype - ).to(device) - new_mod.weight.quantize_(w) # type: ignore[attr-defined,operator] - if has_bias: - assert b is not None - new_mod.bias.data.copy_(b) # type: ignore[union-attr] - del w, b - replaced = True - elif isinstance(child, te.Linear): # type: ignore[union-attr] - # te.Linear with bf16 weight → re-init inside quantized_model_init for fp8. - in_f, out_f = child.in_features, child.out_features - has_bias = child.bias is not None - device = child.weight.device - dtype = ( - child.weight.dtype - if not hasattr(child.weight, "_data") - else torch.bfloat16 - ) - state = {k: v.detach().clone() for k, v in child.state_dict().items()} - setattr(mod, name, nn.Identity()) - del child - torch.cuda.empty_cache() - with quantized_model_init(enabled=True): - new_mod = te.Linear( # type: ignore[union-attr] - in_f, - out_f, - bias=has_bias, - params_dtype=dtype, # type: ignore[arg-type] - ).to(device) # type: ignore[arg-type] - new_mod.load_state_dict(state, strict=False) - replaced = True - elif ( - hasattr(te, "LayerNormLinear") and isinstance(child, te.LayerNormLinear) # type: ignore[union-attr] - ): - state = {k: v.detach().clone() for k, v in child.state_dict().items()} - hidden_size = child.in_features - out_features = child.out_features - has_bias = child.use_bias - device = next(child.parameters()).device - setattr(mod, name, nn.Identity()) - del child - torch.cuda.empty_cache() - with quantized_model_init(enabled=True): - new_mod = te.LayerNormLinear( # type: ignore[union-attr] - hidden_size, - out_features, - bias=has_bias, - params_dtype=torch.bfloat16, - ).to(device) - new_mod.load_state_dict(state, strict=False) - replaced = True - elif ( - hasattr(te, "LayerNormMLP") and isinstance(child, te.LayerNormMLP) # type: ignore[union-attr] - ): - state = {k: v.detach().clone() for k, v in child.state_dict().items()} - fc1_weight: Tensor = child.fc1_weight # type: ignore[attr-defined] - hidden_size = int(fc1_weight.shape[1]) - # fc1 packed as (2*ffn_hidden_size, hidden_size) for swiglu. - ffn_hidden_size = int(fc1_weight.shape[0]) // 2 - has_bias = ( - getattr(child, "fc1_bias", None) is not None - and child.fc1_bias is not None # type: ignore[attr-defined] + + te, _ = _load_transformer_engine() + converted: list[str] = [] + + def walk(owner: nn.Module, prefix: str = "") -> None: + for name, child in tuple(owner.named_children()): + path = f"{prefix}.{name}" if prefix else name + if isinstance(child, nn.Linear) and path.endswith(_ESMC_FP8_LINEAR_SUFFIX): + replacement = te.Linear( + child.in_features, + child.out_features, + bias=child.bias is not None, + params_dtype=child.weight.dtype, + device=child.weight.device, ) - device = fc1_weight.device - setattr(mod, name, nn.Identity()) - del child - torch.cuda.empty_cache() - with quantized_model_init(enabled=True): - new_mod = te.LayerNormMLP( # type: ignore[union-attr] - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - bias=has_bias, - activation="swiglu", - params_dtype=torch.bfloat16, - ).to(device) # type: ignore[arg-type] - new_mod.load_state_dict(state, strict=False) - replaced = True - - if replaced: - # Freeze via .eval()+.requires_grad_(False); per-param ops would unwrap Float8Tensor. - new_mod.eval().requires_grad_(False) - setattr(mod, name, new_mod) - torch.cuda.empty_cache() + with torch.no_grad(): + replacement.weight.copy_(child.weight) + if child.bias is not None: + replacement.bias.copy_(child.bias) + replacement.eval().requires_grad_(False) + setattr(owner, name, replacement) + converted.append(path) else: - _walk(child) - - _walk(module) - torch.cuda.empty_cache() + walk(child, path) + + walk(module) + if len(converted) != _ESMC_FP8_EXPECTED_PROJECTIONS: + raise RuntimeError( + "ESMC FP8 conversion expected exactly " + f"{_ESMC_FP8_EXPECTED_PROJECTIONS} attention output projections, " + f"found {len(converted)}." + ) + return tuple(converted) @contextmanager -def _lm_precision_context(fp8: bool): - """bf16 autocast (+ optional TE fp8 autocast) around the LM forward. +def _lm_precision_context(precision: str, device: torch.device): + """Apply the resolved ESMC inference precision.""" - te.autocast keeps te.Linear outputs bf16 instead of the fp32 default - (~425 MB at L=1024 in the hidden-state cache). - """ + if device.type != "cuda" or precision == "fp32": + yield + return with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - if fp8 and TE_AVAILABLE: - fp8_recipe = DelayedScaling( # type: ignore[misc] - fp8_format=Format.HYBRID, # type: ignore[union-attr] - amax_history_len=1, - amax_compute_algo="most_recent", + if precision == "fp8": + te, recipe = _load_transformer_engine() + fp8_recipe = recipe.Float8CurrentScaling( + use_power_2_scales=False, + fp8_format=recipe.Format.HYBRID, ) - with te.autocast(enabled=True, recipe=fp8_recipe): # type: ignore[union-attr] + with te.autocast(enabled=True, recipe=fp8_recipe): yield else: yield -class ESMFold2Model(FastPLMTestTimeTrainingMixin, PreTrainedModel): - """ESMFold2 — all-atom structure prediction with an ESMC PLM backbone. +class ESMFold2Model( + FastPLMTestTimeTrainingMixin, + ESMFold2EmbeddingMixin, + ESMFold2AttentionMixin, + PreTrainedModel, +): + """ESMFold2: all-atom structure prediction with an ESMC PLM backbone. This is the standard released ESMFold2 architecture (uses a linear- recurrent trunk, internally referred to as "parcae"). @@ -617,15 +873,16 @@ class ESMFold2Model(FastPLMTestTimeTrainingMixin, PreTrainedModel): Memory / perf knobs: - * ``model.set_chunk_size(int|None)``: caps L² ops (triangle / OPM / - pair transition) at this token-axis chunk. Default 64 — fits - L≈2k on an 80 GB GPU. Pass ``None`` for faster inference at L<600. + * ``model.set_chunk_size(int|None)``: caps l^2 ops (triangle / OPM / + pair transition) at this token-axis chunk. Default 64: fits + l approximately 2k on an 80 GB GPU. Pass ``None`` for faster inference + when l is below 600. * ``model.set_kernel_backend(None | "fused" | "cuequivariance")``: select kernel backend (None = reference path). """ config_class = ESMFold2Config - _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"\._extra_state$"] def __init__(self, config: ESMFold2Config) -> None: super().__init__(config) @@ -646,13 +903,25 @@ def __init__(self, config: ESMFold2Config) -> None: ) self._esmc: nn.Module | None = None self._esmc_fp8: bool = False + self._esmc_fp8_module_paths: tuple[str, ...] = () + self._esmc_source: str = config.esmc_id + self._esmc_source_revision: str | None = None + self._esmc_source_files: dict[str, str] = {} + self._esmc_local_files_only = False + self._esmc_precision_policy: str = str(getattr(config, "esmc_precision", "auto")) + self._esmc_precision_status = ESMCPrecisionStatus( + requested=self._esmc_precision_policy, + resolved="unloaded", + reason="ESMC has not been loaded.", + device=str(self.device), + transformer_engine_version=_transformer_engine_version(), + ) self._ttt_lm_head: nn.Module | None = None self._esmfold2_input_builder: Any | None = None + self._kernel_backend: str | None = None pf = config.folding_trunk - self.folding_trunk = FoldingTrunk( - n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4 - ) + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) if config.lm_encoder.enabled: self.lm_encoder: FoldingTrunk | None = FoldingTrunk( n_layers=config.lm_encoder.n_layers, d_pair=d_pair, expansion_ratio=4 @@ -665,9 +934,7 @@ def __init__(self, config: ESMFold2Config) -> None: parcae_decay_init = math.sqrt(1.0 / 5.0) parcae_delta_init = -math.log(parcae_decay_init) self.parcae_log_delta = nn.Parameter( - torch.full( - (d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32 - ) + torch.full((d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32) ) self.parcae_b_cont = nn.Parameter(torch.eye(d_pair)) self.parcae_readout = nn.Linear(d_pair, d_pair, bias=False) @@ -678,9 +945,7 @@ def __init__(self, config: ESMFold2Config) -> None: # Heads -------------------------------------------------------------- self.structure_head = DiffusionStructureHead(config) - self.distogram_head = nn.Linear( - d_pair, config.structure_head.distogram_bins, bias=True - ) + self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True) self.confidence_head = ConfidenceHead(config) msa_cfg = config.msa_encoder @@ -697,84 +962,106 @@ def __init__(self, config: ESMFold2Config) -> None: ) self.post_init() + self._register_state_dict_hook(_drop_transient_esmc_state) self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) - def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None: - """Load the FastPLMs ESM++ LM used as the ESMFold2 PLM backbone. + @property + def esmc_precision_status(self) -> ESMCPrecisionStatus: + return self._esmc_precision_status - ``precision``: ``"bf16"`` (default), ``"fp32"``, or opt-in ``"fp8"``. - """ - dtype_map = { - "bf16": torch.bfloat16, - "fp32": torch.float32, - "fp8": torch.bfloat16, - } - if precision not in dtype_map: - raise ValueError(f"precision must be one of {list(dtype_map)}, got {precision!r}") - if precision == "fp8" and not TE_AVAILABLE: - raise RuntimeError( - "esmc_precision='fp8' requires transformer_engine.pytorch." - ) - dtype = dtype_map[precision] + def load_esmc( + self, + esmc_model_path: str, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool = False, + ) -> None: + """Load canonical ESMC weights and resolve the inference precision.""" - esmc = _load_fastplms_esmplusplus_for_esmfold2( - esmc_model_path=esmc_model_path, - attn_backend=self.config.esmc_attn_backend, - device=self.device, - dtype=dtype, - ) - assert esmc.config.hidden_size == self.config.lm_d_model, ( - f"ESMFold2 expected lm_d_model={self.config.lm_d_model}, " - f"but loaded ESM++ hidden_size={esmc.config.hidden_size}." - ) - assert esmc.config.num_hidden_layers == self.config.lm_num_layers, ( - f"ESMFold2 expected lm_num_layers={self.config.lm_num_layers}, " - f"but loaded ESM++ num_hidden_layers={esmc.config.num_hidden_layers}." + _install_esmc_backbone( + self, + esmc_model_path, + precision=precision, + device=device, + local_files_only=local_files_only, ) - for p in esmc.parameters(): - p.requires_grad_(False) - if precision == "fp8": - with torch.no_grad(): - _convert_te_modules_to_fp8_inplace(esmc) - - self._esmc_fp8 = precision == "fp8" - self._esmc = esmc + def reload_esmc( + self, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool | None = None, + ) -> None: + """Reload canonical weights with the requested precision policy.""" + + source = self._esmc_source or self.config.esmc_id + old_esmc = self._esmc + old_head = self._ttt_lm_head + self._esmc = None + self._esmc_fp8 = False + self._esmc_fp8_module_paths = () self._ttt_lm_head = None + del old_esmc, old_head + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.load_esmc( + source, + precision=precision, + device=device, + local_files_only=( + self._esmc_local_files_only + if local_files_only is None + else local_files_only + ), + ) - def _ensure_ttt_lm_head(self) -> None: - assert self._esmc is not None, "ESMFold2 TTT requires load_esmc=True." + def _ensure_ttt_bf16(self) -> None: if self._esmc_fp8: - raise RuntimeError("ESMFold2 TTT is not supported with fp8 ESM++.") + _reload_esmc_bf16_for_gradients( + self, + reason="TTT requires BF16; the persisted serving policy is unchanged.", + ) + + def _ensure_ttt_lm_head(self) -> None: + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") if self._ttt_lm_head is not None: return - try: - from fastplms.esm_plusplus.modeling_esm_plusplus import ( - ESMplusplusConfig, - ESMplusplusForMaskedLM, - ) - except ImportError: - from .modeling_esm_plusplus import ( - ESMplusplusConfig, - ESMplusplusForMaskedLM, - ) + from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ) - esmc_config = ESMplusplusConfig.from_pretrained(self.config.esmc_id) - esmc_config.attn_backend = self.config.esmc_attn_backend + source = self._esmc_source or self.config.esmc_id + source_revision = self._esmc_source_revision + if source_revision is None: + source_revision, _ = _manifest_esmc_checkpoint_contract(source) + revision_kwargs: dict[str, Any] = { + "local_files_only": self._esmc_local_files_only, + } + if source_revision is not None: + revision_kwargs["revision"] = source_revision + esmc_config = ESMplusplusConfig.from_pretrained( + source, + **revision_kwargs, + ) + set_config_attn_implementation(esmc_config, get_attn_implementation(self.config)) mlm, loading_info = ESMplusplusForMaskedLM.from_pretrained( - self.config.esmc_id, + source, config=esmc_config, output_loading_info=True, + **revision_kwargs, ) missing_head_keys = [ - key - for key in loading_info["missing_keys"] - if key.startswith("sequence_head") + key for key in loading_info["missing_keys"] if key.startswith("sequence_head") ] - assert len(missing_head_keys) == 0, ( - f"ESMFold2 TTT could not load a pretrained ESM++ MLM head from " - f"{self.config.esmc_id}: missing {missing_head_keys}" - ) + if missing_head_keys: + raise RuntimeError( + "ESMFold2 TTT could not load a pretrained ESM++ MLM head from " + f"{source}: missing {missing_head_keys}" + ) dtype = next(self._esmc.parameters()).dtype mlm = mlm.to(device=self.device, dtype=dtype).eval() self._ttt_lm_head = mlm.sequence_head @@ -782,9 +1069,9 @@ def _ensure_ttt_lm_head(self) -> None: del mlm def _ttt_get_trainable_modules(self) -> list[nn.Module]: - assert self._esmc is not None, "ESMFold2 TTT requires load_esmc=True." - if self._esmc_fp8: - raise RuntimeError("ESMFold2 TTT is not supported with fp8 ESM++.") + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") return [self._esmc] def _ttt_tokenize( @@ -796,16 +1083,17 @@ def _ttt_tokenize( del kwargs if input_ids is not None: return input_ids - assert seq is not None, "Pass either seq or input_ids for ESMFold2 TTT." + if seq is None: + raise ValueError("Pass either seq or input_ids for ESMFold2 TTT.") sequences = [seq] if isinstance(seq, str) else seq + if not sequences: + raise ValueError("ESMFold2 TTT requires at least one protein sequence.") token_to_id = {token: idx for idx, token in enumerate(SEQUENCE_VOCAB)} encoded = [] for sequence in sequences: token_ids = [SEQUENCE_BOS_TOKEN] for amino_acid in sequence: - token_ids.append( - token_to_id[amino_acid if amino_acid in token_to_id else "X"] - ) + token_ids.append(token_to_id[amino_acid if amino_acid in token_to_id else "X"]) token_ids.append(SEQUENCE_EOS_TOKEN) encoded.append(token_ids) max_len = max(len(token_ids) for token_ids in encoded) @@ -846,14 +1134,14 @@ def _ttt_predict_logits( **kwargs, ) -> torch.Tensor: del kwargs - assert isinstance(batch, torch.Tensor), ( - "ESMFold2 TTT expects input_ids tensors." - ) - assert self._esmc is not None, "ESMFold2 TTT requires load_esmc=True." - if self._esmc_fp8: - raise RuntimeError("ESMFold2 TTT is not supported with fp8 ESM++.") + if not isinstance(batch, torch.Tensor): + raise TypeError("ESMFold2 TTT expects input_ids tensors.") + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") self._ensure_ttt_lm_head() - assert self._ttt_lm_head is not None + if self._ttt_lm_head is None: + raise RuntimeError("ESMFold2 TTT MLM head initialization failed.") attention_mask = batch.ne(SEQUENCE_PAD_TOKEN) output = self._esmc( input_ids=batch, @@ -868,9 +1156,7 @@ def from_pretrained( cls, pretrained_model_name_or_path, *args, load_esmc: bool = True, **kwargs ): if cls is ESMFold2Model and "config" not in kwargs: - config = ESMFold2Config.from_pretrained( - pretrained_model_name_or_path, **kwargs - ) + config = ESMFold2Config.from_pretrained(pretrained_model_name_or_path, **kwargs) if config.type == "experimental": raise ValueError( "FastPLMs ESMFold2 supports the released ESMFold2 and " @@ -879,42 +1165,49 @@ def from_pretrained( ) kwargs["config"] = config # Pop the precision knob before forwarding to the HF loader. - esmc_precision = kwargs.pop("esmc_precision", "bf16") - model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + esmc_precision = kwargs.pop("esmc_precision", None) + local_files_only = bool(kwargs.get("local_files_only", False)) + output_loading_info = bool(kwargs.get("output_loading_info", False)) + loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + if output_loading_info: + model, loading_info = loaded + else: + model = loaded if load_esmc: - model.load_esmc(model.config.esmc_id, precision=esmc_precision) - return model + model.load_esmc( + model.config.esmc_id, + precision=esmc_precision or model.config.esmc_precision, + local_files_only=local_files_only, + ) + return (model, loading_info) if output_loading_info else model def set_kernel_backend(self, backend: str | None) -> None: """Select kernel backend. Args: - backend: ``None`` (reference path), ``"fused"`` (vendored Triton - kernels), or ``"cuequivariance"`` (cuequivariance kernels - where applicable; vanilla python fallback otherwise). + backend: ``None`` (reference path), ``"fused"`` (requires the + unavailable source-built Triton bundle), or + ``"cuequivariance"`` (requires the ``structure,cueq`` extras + on a supported Linux CUDA 13 host). """ + validate_kernel_backend(backend) self.folding_trunk.set_kernel_backend(backend) if self.lm_encoder is not None: self.lm_encoder.set_kernel_backend(backend) self.parcae_coda.set_kernel_backend(backend) self.confidence_head.set_kernel_backend(backend) self.structure_head.set_kernel_backend(backend) + self._kernel_backend = backend - def apply_torch_compile( - self, mode: str = "fixed_seqlen", dynamic: bool | None = None - ) -> None: - """Compile L²-heavy blocks. ``mode='fixed_seqlen'`` recompiles per L; ``'dynamic_seqlen'`` compiles once. + def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None: + """Compile l^2-heavy blocks. + + ``mode='fixed_seqlen'`` recompiles per l; ``'dynamic_seqlen'`` compiles + once. - Does NOT stack with our Triton kernels — call ``set_kernel_backend(None)`` + Does NOT stack with our Triton kernels: call ``set_kernel_backend(None)`` before compiling. """ - import torch._dynamo - - torch._dynamo.config.cache_size_limit = 512 # type: ignore[attr-defined] - torch._dynamo.config.accumulated_cache_size_limit = 512 # type: ignore[attr-defined] - # capture_scalar_outputs avoids graph breaks at .item() in atom-attention path. - torch._dynamo.config.capture_scalar_outputs = True # type: ignore[attr-defined] - if dynamic is None: dynamic = mode == "dynamic_seqlen" kwargs: dict = {"dynamic": dynamic} @@ -956,10 +1249,19 @@ def _compute_lm_hidden_states( tok_mask: Tensor, lm_mask_pct: float = 0.0, ) -> Tensor: - assert self._esmc is not None - # fp8 TE kernels require prod(shape[:-1]) % 8 == 0. - pad_to = 8 if self._esmc_fp8 else None - with _lm_precision_context(self._esmc_fp8): + if self._esmc_fp8 and torch.is_grad_enabled(): + _reload_esmc_bf16_for_gradients( + self, + reason=( + "Gradient-enabled ESMC execution requires BF16; the persisted " + "serving policy is unchanged." + ), + ) + if self._esmc is None: + raise RuntimeError("ESMFold2 requires load_esmc=True for LM feature extraction.") + # Transformer Engine FP8 kernels require l to be a multiple of 16. + pad_to = 16 if self._esmc_fp8 else None + with _lm_precision_context(self._esmc_precision_status.resolved, self.device): return compute_lm_hidden_states( self._esmc, input_ids, @@ -996,8 +1298,8 @@ def _run_one_loop( tok_mask: Tensor, total_steps: int, ) -> Tensor: - # Helper method (not inline) so per-iter locals free on return — - # otherwise leaks ~2 GB L²×c_z into distogram/sample scope. + # Helper method (not inline) so per-iter locals free on return: + # otherwise leaks about 2 GB of l^2 * c_z data into distogram/sample scope. # training=True forces dropout under eval(), matching the per-loop # dropout strategy used at train time. lm_cfg = self.config.lm_encoder @@ -1010,7 +1312,8 @@ def _run_one_loop( for _ in range(total_steps): if _per_loop_lm_dropout: - assert lm_z is not None # narrowed by _per_loop_lm_dropout + if lm_z is None: + raise RuntimeError("Per-loop LM dropout requires LM pair features.") lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True) else: lm_z_i = lm_z @@ -1034,26 +1337,24 @@ def _run_one_loop( max_depth=_msa_inputs["max_depth"], enabled=_msa_inputs["subsample_enabled"], ) - B_msa, M, L_msa = msa_i.shape - msa_oh = F.one_hot( - msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES - ).float() + b_msa, m, l_msa = msa_i.shape + msa_oh = F.one_hot(msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES).float() msa_attn = ( mask_i.permute(0, 2, 1).float() if mask_i is not None - else tok_mask[:, :, None].expand(-1, -1, M).float() + else tok_mask[:, :, None].expand(-1, -1, m).float() ) # Bias-free MSAEncoder.embed requires zeroed padding. msa_oh = msa_oh * msa_attn.unsqueeze(-1) hd = ( hd_i.permute(0, 2, 1).float() if hd_i is not None - else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + else torch.zeros(b_msa, l_msa, m, device=msa_i.device) ) dv = ( dv_i.permute(0, 2, 1).float() if dv_i is not None - else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + else torch.zeros(b_msa, l_msa, m, device=msa_i.device) ) msa_pair = self.msa_encoder( x_pair=z_inject_pair, @@ -1064,9 +1365,7 @@ def _run_one_loop( msa_attention_mask=msa_attn, ).to(z_inject_pair.dtype) z_inject_pair = ( - msa_pair - if self.config.msa_encoder_overwrite - else (z_inject_pair + msa_pair) + msa_pair if self.config.msa_encoder_overwrite else (z_inject_pair + msa_pair) ) if refined_lm_z is not None: @@ -1078,7 +1377,6 @@ def _run_one_loop( return z - @torch.inference_mode() def forward( self, token_index: Tensor, @@ -1112,8 +1410,28 @@ def forward( msa_max_depth: int = 1024, msa_column_mask_rate: float = 0.1, msa_subsample_at_inference: bool = True, - **kwargs, - ) -> dict[str, Tensor]: + early_exit: bool = False, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: float | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> ESMFold2Output | tuple[Any, ...]: + output_hidden_states, return_dict = _resolve_structure_output_controls( + self.config, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + validate_msa_conditioning_inputs( + self.config, + msa=msa, + msa_attention_mask=msa_attention_mask, + has_deletion=has_deletion, + deletion_value=deletion_value, + deletion_mean=deletion_mean, + ) tok_mask = token_attention_mask atm_mask = atom_attention_mask disto_idx = distogram_atom_idx @@ -1149,18 +1467,14 @@ def forward( res_type.shape[0], res_type.shape[1], device=res_type.device ) - ref_element_oh = F.one_hot( - ref_element.long(), num_classes=MAX_ATOMIC_NUMBER - ).float() + ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float() ref_atom_name_chars_oh = F.one_hot( ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE ).float() # Bias-free downstream Linears require zeroed padding. atm_mask_f = atm_mask.float() ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1) - ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze( - -1 - ).unsqueeze(-1) + ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1) atom_to_token = atom_to_token * atm_mask.long() use_amp = ref_pos.device.type == "cuda" @@ -1178,9 +1492,7 @@ def forward( atom_to_token=atom_to_token, ) - z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( - x_inputs - ).unsqueeze(1) + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1) relative_position_encoding = self.rel_pos( residue_index=residue_index, @@ -1192,22 +1504,14 @@ def forward( token_bonds_encoding = self.token_bonds(token_bonds.float()) z_init = z_init + relative_position_encoding + token_bonds_encoding - if ( - lm_hidden_states is None - and input_ids is not None - and self._esmc is not None - ): + if lm_hidden_states is None and input_ids is not None and self._esmc is not None: lm_hidden_states = self._compute_lm_hidden_states( input_ids, asym_id, residue_index, mol_type, tok_mask, - lm_mask_pct=( - self.config.lm_mask_pct - if lm_mask_pct is None - else lm_mask_pct - ), + lm_mask_pct=(self.config.lm_mask_pct if lm_mask_pct is None else lm_mask_pct), ) lm_z: Tensor | None = None if lm_hidden_states is not None: @@ -1238,7 +1542,7 @@ def forward( subsample_enabled=msa_subsample_at_inference, ) - # Method call (not inline loop) frees per-iter L²×c_z locals. + # Method call (not inline loop) frees per-iteration l^2 * c_z locals. z = self._run_one_loop( z=z, z_init=z_init, @@ -1278,12 +1582,16 @@ def forward( token_attention_mask=tok_mask, num_diffusion_samples=n_samples, num_sampling_steps=num_sampling_steps, + max_inference_sigma=max_inference_sigma, + noise_scale=noise_scale, + step_scale=step_scale, return_atom_repr=False, - denoising_early_exit_rmsd=None, + denoising_early_exit_rmsd=(0.10 if early_exit else None), ) sample_coords = structure_output["sample_atom_coords"] - assert sample_coords is not None + if sample_coords is None: + raise RuntimeError("ESMFold2 structure sampling did not return coordinates.") output: dict[str, Tensor] = {"distogram_logits": distogram_logits} output["sample_atom_coords"] = sample_coords @@ -1302,20 +1610,31 @@ def forward( token_bonds_encoding=token_bonds_encoding.detach(), ) output.update(confidence_output) - output["atom_pad_mask"] = ( - atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask - ) + output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask output["residue_index"] = residue_index output["entity_id"] = entity_id - return output + return _finalize_structure_output( + output, + token_input_state=x_inputs, + pair_state=z, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) @torch.no_grad() - def infer_protein(self, seq: str, **forward_kwargs) -> dict: + def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output: from .protein_utils import prepare_protein_features + if forward_kwargs.pop("return_dict", True) is not True: + raise ValueError( + "infer_protein always returns a mapping; return_dict=False is invalid." + ) features = prepare_protein_features(seq) + if not self.config.msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) features = {k: v.to(self.device) for k, v in features.items()} - return self(**features, **forward_kwargs) + return self(**features, **forward_kwargs, return_dict=True) @property def input_builder(self): @@ -1332,7 +1651,12 @@ def input_types(self): return esmfold2_types def prepare_structure_input(self, input, seed: int | None = None): - return self.input_builder.prepare_input(input, seed=seed, device=self.device) + return self.input_builder.prepare_model_input( + self, + input, + seed=seed, + device=self.device, + ) def fold( self, @@ -1378,16 +1702,17 @@ def _fold_protein_no_ttt( ): from .esmfold2_types import MSA, ProteinInput, StructurePredictionInput - assert not ( - msa is not None and msa_path is not None - ), "Pass at most one of msa or msa_path." + if msa is not None and msa_path is not None: + raise ValueError("Pass at most one of msa or msa_path.") if msa_path is not None: msa = MSA.from_a3m(msa_path, max_sequences=msa_max_sequences) if msa is not None: query = str(msa.query).replace("-", "").upper() - assert query == sequence.upper(), ( - f"MSA query does not match sequence: expected {sequence.upper()!r}, got {query!r}" - ) + if query != sequence.upper(): + raise ValueError( + "MSA query does not match sequence: " + f"expected {sequence.upper()!r}, got {query!r}" + ) input = StructurePredictionInput( sequences=[ProteinInput(id=chain_id, sequence=sequence, msa=msa)] @@ -1403,12 +1728,14 @@ def _fold_protein_no_ttt( @staticmethod def _ttt_mean_plddt(result) -> float: - assert result.plddt is not None, "ESMFold2 result has no pLDDT tensor." + if result.plddt is None: + raise RuntimeError("ESMFold2 result has no pLDDT tensor.") return float(result.plddt.float().mean().item()) def _ttt_select_result(self, result): if isinstance(result, list): - assert len(result) > 0, "ESMFold2 fold returned an empty result list." + if not result: + raise RuntimeError("ESMFold2 fold returned an empty result list.") return max(result, key=self._ttt_mean_plddt) return result @@ -1421,9 +1748,8 @@ def _ttt_eval_step( **kwargs, ) -> tuple[dict[str, Any], float | None]: del input_ids - assert isinstance(seq, str), ( - "ESMFold2 fold TTT is protein-only and sequence-string only." - ) + if not isinstance(seq, str): + raise TypeError("ESMFold2 fold TTT is protein-only and sequence-string only.") fold_kwargs = kwargs["fold_kwargs"] was_training = self.training self.eval() @@ -1498,9 +1824,9 @@ def fold_protein_ttt( complex_id: str = "pred", ttt_config: TTTConfig | dict[str, Any] | None = None, ): - assert self._esmc is not None, "ESMFold2 TTT requires load_esmc=True." - if self._esmc_fp8: - raise RuntimeError("ESMFold2 TTT is not supported with fp8 ESM++.") + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") fold_kwargs = { "chain_id": chain_id, "msa": msa, @@ -1512,9 +1838,7 @@ def fold_protein_ttt( "seed": seed, "complex_id": complex_id, } - baseline = self._ttt_select_result( - self._fold_protein_no_ttt(sequence, **fold_kwargs) - ) + baseline = self._ttt_select_result(self._fold_protein_no_ttt(sequence, **fold_kwargs)) baseline_plddt = self._ttt_mean_plddt(baseline) best_result = baseline best_plddt = baseline_plddt @@ -1551,12 +1875,14 @@ def fold_protein_ttt( @staticmethod def result_to_cif(result) -> str: - assert not isinstance(result, list), "Pass one MolecularComplexResult at a time." + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") return result.complex.to_mmcif() @staticmethod def result_to_pdb(result) -> str: - assert not isinstance(result, list), "Pass one MolecularComplexResult at a time." + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") return result.complex.to_protein_complex().to_pdb_string() def save_as_cif(self, result, output_path: str | Path) -> None: @@ -1664,7 +1990,7 @@ def forward( deletion_value: Tensor, msa_attention_mask: Tensor, ) -> Tensor: - # All inputs are pre-transposed to [B, L, M, ...] before calling. + # Every input tensor is pre-transposed to shape (b, l, m, ...) before this call. m_feat = torch.cat( [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 ) diff --git a/fastplms/esmfold2/modeling_esmfold2_common.py b/src/fastplms/models/esmfold2/modeling_esmfold2_common.py similarity index 82% rename from fastplms/esmfold2/modeling_esmfold2_common.py rename to src/fastplms/models/esmfold2/modeling_esmfold2_common.py index 8bb1ef3..c59005b 100644 --- a/fastplms/esmfold2/modeling_esmfold2_common.py +++ b/src/fastplms/models/esmfold2/modeling_esmfold2_common.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2026 Biohub. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -10,65 +9,104 @@ from __future__ import annotations -import random import importlib -from contextlib import contextmanager from functools import partial -from typing import cast +from importlib.util import find_spec +from typing import ClassVar, cast -import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor from torch.utils.checkpoint import checkpoint -try: - flash_attn_module = importlib.import_module("flash_attn") - flash_bert_padding = importlib.import_module("flash_attn.bert_padding") - flash_attn_func = flash_attn_module.flash_attn_func - flash_attn_varlen_func = flash_attn_module.flash_attn_varlen_func - index_first_axis = flash_bert_padding.index_first_axis - pad_input = flash_bert_padding.pad_input - - FLASH_ATTN_AVAILABLE = True -except ImportError: - flash_attn_func = None # type: ignore[assignment] - flash_attn_varlen_func = None # type: ignore[assignment] - index_first_axis = None # type: ignore[assignment] - pad_input = None # type: ignore[assignment] - FLASH_ATTN_AVAILABLE = False +from .configuration_esmfold2 import ESMFold2Config +from .reproducibility import seed_context + +_seed_context = seed_context try: + if find_spec("cuequivariance_ops_torch") is None: + raise ImportError("cuequivariance_ops_torch is unavailable") cue_module = importlib.import_module("cuequivariance_torch") - cue_triangle = importlib.import_module("cuequivariance_torch.primitives.triangle") _cue_attn_pair_bias = cue_module.attention_pair_bias - _cue_tri_mul = cue_triangle.triangle_multiplicative_update + _cue_tri_mul = cue_module.triangle_multiplicative_update CUE_AVAILABLE = True -except ImportError: +except (AttributeError, ImportError): _cue_attn_pair_bias = None # type: ignore[assignment] _cue_tri_mul = None # type: ignore[assignment] CUE_AVAILABLE = False -# The Biohub release includes optional Triton kernels. FastPLMs keeps the -# reference path enabled by default so Hugging Face remote-code loading can stay -# flat and self-contained. +# Biohub ships optional source-built Triton helpers. FastPLMs does not bundle or +# compile them; these placeholders retain checkpoint-compatible control flow +# while the portable PyTorch path remains flat and self-contained. _fused_pair_bias = None _fused_trimul_with_residual = None _FusedLNLinearSwiGLU = None _FusedDropoutResidual = None TRITON_KERNELS_AVAILABLE = False -from .configuration_esmfold2 import ESMFold2Config - BACKEND_FUSED = "fused" BACKEND_CUEQ = "cuequivariance" _VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ) +MSA_CONDITIONING_INPUT_NAMES = ( + "msa", + "msa_attention_mask", + "has_deletion", + "deletion_value", + "deletion_mean", +) + + +def validate_kernel_backend(backend: str | None) -> None: + """Fail before mutating modules when a named kernel cannot execute.""" + + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + if backend == BACKEND_FUSED and not TRITON_KERNELS_AVAILABLE: + raise RuntimeError( + "backend='fused' is unavailable because FastPLMs does not bundle the " + "source-built ESMFold2 Triton kernels." + ) + if backend == BACKEND_CUEQ and not CUE_AVAILABLE: + raise RuntimeError( + "backend='cuequivariance' requires cuequivariance_torch and the CUDA 13 " + "cuequivariance_ops_torch runtime. Install FastPLMs with the " + "'structure,cueq' extras on a supported Linux CUDA 13 host." + ) + + +def validate_msa_conditioning_inputs( + config: ESMFold2Config, + *, + msa: Tensor | None, + msa_attention_mask: Tensor | None, + has_deletion: Tensor | None, + deletion_value: Tensor | None, + deletion_mean: Tensor | None, +) -> None: + """Reject MSA-derived tensors for checkpoints trained without MSA conditioning.""" + + if config.msa_conditioning: + return + values = { + "msa": msa, + "msa_attention_mask": msa_attention_mask, + "has_deletion": has_deletion, + "deletion_value": deletion_value, + "deletion_mean": deletion_mean, + } + provided = sorted(name for name, value in values.items() if value is not None) + if provided: + raise ValueError( + "This ESMFold2 checkpoint was trained without MSA conditioning and rejects " + f"MSA-derived inputs: {', '.join(provided)}." + ) def _fused_active(module: nn.Module, tensor: Tensor) -> bool: - """Common preconditions for the vendored fused Triton inference kernels.""" + """Return whether an optional fused implementation can handle this call.""" return ( TRITON_KERNELS_AVAILABLE and getattr(module, "_kernel_backend", None) == BACKEND_FUSED @@ -89,11 +127,10 @@ class DropoutResidual(nn.Module): in-place residual add). Falls back to unfused otherwise. """ - def __init__( - self, r: float, batch_dim: int, use_fused_kernels: bool = False - ) -> None: + def __init__(self, r: float, batch_dim: int, use_fused_kernels: bool = False) -> None: super().__init__() - assert batch_dim in (1, 2), f"batch_dim must be 1 or 2, got {batch_dim}" + if isinstance(batch_dim, bool) or batch_dim not in (1, 2): + raise ValueError(f"batch_dim must be 1 or 2, got {batch_dim}") self._use_fused_kernels = ( use_fused_kernels and batch_dim == 1 and _FusedDropoutResidual is not None ) @@ -108,7 +145,7 @@ def __init__( def forward(self, residual: Tensor, delta: Tensor) -> Tensor: if self._use_fused_kernels: return self._impl(residual, delta) - # Unfused: row/col-shared dropout via [1, ...] mask broadcast. + # The unfused path broadcasts a row/column-shared mask M with shape (1, ...). if self._r == 0.0 or not self.training: return residual + delta shape = list(delta.shape) @@ -126,20 +163,18 @@ def forward(self, residual: Tensor, delta: Tensor) -> Tensor: MAX_ATOMIC_NUMBER: int = 128 # Input feature dim = 3 + 1 + 1 + 128 + 64*4 = 389 -ATOM_FEATURE_DIM: int = ( - XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS -) +ATOM_FEATURE_DIM: int = XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS NUM_RES_TYPES: int = 33 _EPS = 1e-5 -# Default for the triangle / OPM / pair-transition L² ops. Caps peak memory -# so L≈2k folds on an 80 GB GPU (~76 GB peak at chunk=128 for L=1438; -# chunk=64 leaves headroom for the largest foldbench targets). Override via -# ``model.set_chunk_size(...)``; pass None to disable chunking (faster for -# short L but OOM-prone past ~600). +# Default for the quadratic triangle, OPM, and pair-transition operations. +# It caps peak memory so l around 2,000 fits on an 80 GiB GPU. At l=1,438, +# chunk=128 uses roughly 76 GiB, while chunk=64 leaves headroom for the +# largest foldbench targets. Pass None to disable chunking; this is faster for +# short sequences but prone to out-of-memory errors beyond l around 600. _DEFAULT_CHUNK_SIZE = 64 @@ -201,11 +236,11 @@ def gather_token_to_atom(token_features: Tensor, atom_to_token_idx: Tensor) -> T """Broadcast per-token features to per-atom features using gather. Args: - token_features: [B, L, d] - atom_to_token_idx: [B, A] int64 + token_features: X with shape (b, l, d). + atom_to_token_idx: I with shape (b, a), int64. Returns: - [B, A, d] + X with shape (b, a, d). """ idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1)) return torch.gather(token_features, 1, idx) @@ -220,27 +255,29 @@ def scatter_atom_to_token( """Aggregate per-atom features to per-token features (mean). Args: - atom_features: [B, A, d] - atom_to_token_idx: [B, A] int64 - n_tokens: L - atom_mask: [B, A] bool + atom_features: X with shape (b, a, d). + atom_to_token_idx: I with shape (b, a), int64. + n_tokens: Token count l. + atom_mask: M with shape (b, a), Boolean. Returns: - [B, L, d] + X with shape (b, l, d). """ - B, A, d = atom_features.shape + batch_size, n_atoms, d_model = atom_features.shape n_out = n_tokens idx = atom_to_token_idx if atom_mask is not None: idx = torch.where(atom_mask, atom_to_token_idx, n_tokens) n_out = n_tokens + 1 - idx_expanded = idx.unsqueeze(-1).expand(B, A, d) + idx_expanded = idx.unsqueeze(-1).expand(batch_size, n_atoms, d_model) out = torch.zeros( - B, n_out, d, device=atom_features.device, dtype=atom_features.dtype - ) - out.scatter_reduce_( - 1, idx_expanded, atom_features, reduce="mean", include_self=False + batch_size, + n_out, + d_model, + device=atom_features.device, + dtype=atom_features.dtype, ) + out.scatter_reduce_(1, idx_expanded, atom_features, reduce="mean", include_self=False) return out[:, :n_tokens, :] @@ -248,11 +285,11 @@ def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor: """Gather representative atom coordinates for each token. Args: - coords: [B, A, 3] - rep_atom_idx: [B, L] int64 + coords: X with shape (b, a, 3). + rep_atom_idx: I with shape (b, l), int64. Returns: - [B, L, 3] + X with shape (b, l, 3). """ idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1)) return torch.gather(coords, 1, idx) @@ -265,14 +302,13 @@ def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor: running count that resets at each token boundary. Args: - atom_to_token: [B, A] flat index mapping each atom to its token. + atom_to_token: I with shape (b, a), mapping each atom to its token. Returns: - [B, A] tensor with values in [0, max_atoms_per_token - 1]. + Index tensor I with shape (b, a) and values from zero through + ``max_atoms_per_token - 1``. """ - same_as_prev = F.pad( - atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False - ) + same_as_prev = F.pad(atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False) ones = torch.ones_like(atom_to_token) cumsum = torch.cumsum(ones, dim=-1) group_start = cumsum.masked_fill(same_as_prev, 0) @@ -286,1676 +322,1320 @@ def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor: Equivalent to ``CategoricalMixture(logits, bins=logits.shape[-1], start, end).mean()``. Args: - logits: [..., n_bins] + logits: Logit tensor X with shape (..., n_bins). start: left boundary end: right boundary Returns: - [...] expected value + Expected value tensor Y with shape (...). """ n_bins = logits.shape[-1] - edges = torch.linspace( - start, end, n_bins + 1, device=logits.device, dtype=torch.float32 - ) - v_bins = (edges[:-1] + edges[1:]) / 2 # [n_bins] + edges = torch.linspace(start, end, n_bins + 1, device=logits.device, dtype=torch.float32) + v_bins = (edges[:-1] + edges[1:]) / 2 # V_bin has shape (n_bins,). return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1) # =========================================================================== -# TransitionLayer (used in DiffusionConditioning) +# Feature preparation and language-model projection # =========================================================================== -class TransitionLayer(nn.Module): - """SwiGLU transition: norm -> a_proj, b_proj -> silu(a)*b -> out_proj.""" +class RowAttentionPooling(nn.Module): + """Row-wise attention pooling: attn_proj, out_proj.""" - def __init__(self, d_model: int, n: int, eps: float = 1e-5) -> None: + def __init__(self, d_pair: int, d_single: int) -> None: super().__init__() - hidden = n * d_model - self.norm = nn.LayerNorm(d_model, eps=eps) - self.a_proj = nn.Linear(d_model, hidden, bias=False) - self.b_proj = nn.Linear(d_model, hidden, bias=False) - self.out_proj = nn.Linear(hidden, d_model, bias=False) + self.attn_proj = nn.Linear(d_pair, 1, bias=False) + self.out_proj = nn.Linear(d_pair, d_single, bias=False) - def forward(self, x: Tensor) -> Tensor: - x = self.norm(x) - a = self.a_proj(x) - b = self.b_proj(x) - return self.out_proj(F.silu(a) * b) + def forward(self, z: Tensor, mask: Tensor) -> Tensor: + scores = self.attn_proj(z).squeeze(-1) + mask_bias = torch.where( + mask[:, None, :].bool(), + torch.zeros_like(scores), + torch.full_like(scores, -1e9), + ) + scores = scores + mask_bias + weights = F.softmax(scores, dim=-1) + pooled = torch.einsum("bnm,bnmd->bnd", weights, z) + return self.out_proj(pooled) # =========================================================================== -# AdaptiveLayerNorm (used in DiffusionTransformer) +# InputsEmbedder # =========================================================================== -class AdaptiveLayerNorm(nn.Module): - """Adaptive layer normalization (adaLN-Zero).""" +class InputsEmbedder(nn.Module): + """Embeds input features including atom-level encoding via SWA attention.""" - def __init__(self, d_model: int, d_cond: int, eps: float = 1e-5) -> None: + def __init__(self, config: ESMFold2Config) -> None: super().__init__() - self.d_model = d_model - self.d_cond = d_cond - self.eps = eps - self.s_scale = nn.Parameter(torch.ones(d_cond)) - self.s_gate = nn.Linear(d_cond, d_model, bias=True) - self.s_shift = nn.Linear(d_cond, d_model, bias=False) - - def forward(self, a: Tensor, s: Tensor) -> Tensor: - a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps) - s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps) - return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm) - - -# =========================================================================== -# FourierEmbedding -# =========================================================================== - - -class FourierEmbedding(nn.Module): - """Fourier embedding: cos(2*pi*(t*w + b)).""" + swa_cfg = config.inputs.atom_encoder - w: Tensor - b: Tensor + self.atom_attention_encoder = ESMFold2AtomEncoder( + d_atom=swa_cfg.d_atom, + d_token=swa_cfg.d_token, + n_blocks=swa_cfg.n_blocks, + n_heads=swa_cfg.n_heads, + swa_window_size=swa_cfg.swa_window_size, + expansion_ratio=swa_cfg.expansion_ratio, + structure_prediction=False, # no coords_linear + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + ) - def __init__(self, c: int) -> None: - super().__init__() - self.c = c - self.register_buffer("w", torch.randn(c)) - self.register_buffer("b", torch.randn(c)) + def forward( + self, + aatype: Tensor, + profile: Tensor, + deletion_mean: Tensor, + ref_pos: Tensor, + atom_attention_mask: Tensor, + ref_space_uid: Tensor, + ref_charge: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + atom_to_token: Tensor, + ) -> Tensor: + """Embed inputs into per-token features. - def forward(self, t_hat: Tensor) -> Tensor: - t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1) - return torch.cos( - 2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :]) + Returns: + X with shape (b, l, d_inputs), concatenating atom encoding, + aatype, profile, and deletion mean. + """ + a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder( + ref_pos=ref_pos, + atom_attention_mask=atom_attention_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=atom_to_token, ) + return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1) # =========================================================================== -# SwiGLU / SwiGLUMLP +# ResIdxAsymIdSymIdEntityIdEncoding (trunk relative position) # =========================================================================== -def _compute_swiglu_hidden_size(d_model: int, expansion_ratio: int) -> int: - return expansion_ratio * d_model +class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module): + """Embedding weight W has shape (d_pair, n_features). + Here ``n_features = 2 * (2 * r_bins + 2) + 1 + (2 * c_bins + 2)``. -class SwiGLU(nn.Module): - """SwiGLU with packed w12 and output w3.""" + For default r_bins=32, c_bins=2: 2*66 + 1 + 6 = 139. + """ def __init__( self, - in_features: int, - hidden_features: int, - out_features: int | None = None, - bias: bool = True, + n_relative_residx_bins: int = 32, + n_relative_chain_bins: int = 2, + d_pair: int = 256, ) -> None: super().__init__() - out_features = out_features or in_features - self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) - self.w3 = nn.Linear(hidden_features, out_features, bias=bias) - self.hidden_features = hidden_features + self.n_relative_residx_bins = n_relative_residx_bins + self.n_relative_chain_bins = n_relative_chain_bins + self.d_pair = d_pair - def forward(self, x: Tensor) -> Tensor: - x12 = self.w12(x) - x1, x2 = x12.split(self.hidden_features, dim=-1) - hidden = F.silu(x1) * x2 - return self.w3(hidden) + n_feats_residue = 2 * n_relative_residx_bins + 2 + n_feats_token = 2 * n_relative_residx_bins + 2 + n_feats_chain = 2 * n_relative_chain_bins + 2 + n_feats_same_entity = 1 + total_feats = n_feats_residue + n_feats_token + n_feats_chain + n_feats_same_entity + self.embed = nn.Linear(total_feats, d_pair, bias=False) + def forward( + self, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + token_index: Tensor, + ) -> Tensor: + bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1) + bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1) + bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1) -class SwiGLUMLP(SwiGLU): - """SwiGLU MLP with packed weights, no bias.""" + dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1) + dij_residue = torch.clip( + dij_residue + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, + ) + dij_residue = torch.where(bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1) + aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2) - def __init__( - self, d_model: int, expansion_ratio: int = 4, bias: bool = False - ) -> None: - hidden = _compute_swiglu_hidden_size(d_model, expansion_ratio) - super().__init__( - in_features=d_model, hidden_features=hidden, out_features=d_model, bias=bias + dij_token = torch.clip( + token_index.unsqueeze(2) - token_index.unsqueeze(1) + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, ) + dij_token = torch.where( + bij_same_chain & bij_same_residue, + dij_token, + 2 * self.n_relative_residx_bins + 1, + ) + aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2) + + dij_chain = torch.clip( + sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins, + 0, + 2 * self.n_relative_chain_bins, + ) + dij_chain = torch.where(bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain) + aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2) + + feats = torch.cat( + [ + aij_rel_pos.float(), + aij_rel_token.float(), + bij_same_entity.float().unsqueeze(-1), + aij_rel_chain.float(), + ], + dim=-1, + ) + + return self.embed(feats) # =========================================================================== -# SWA Atom Attention components +# SingleToPair (for LanguageModelShim) # =========================================================================== -def _rotate_half(x: Tensor) -> Tensor: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) +class SingleToPair(nn.Module): + """downproject, output_mlp (Sequential of Linear, GELU, Linear).""" + def __init__(self, input_dim: int, downproject_dim: int, output_dim: int) -> None: + super().__init__() + self.downproject = nn.Linear(input_dim, downproject_dim) + self.output_mlp = nn.Sequential( + nn.Linear(2 * downproject_dim, output_dim), + nn.GELU(), + nn.Linear(output_dim, output_dim), + ) -def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - """Apply RoPE with batch-dependent cos/sin. + def forward(self, x: Tensor) -> Tensor: + x = self.downproject(x) + x = torch.cat( + [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))], + dim=3, + ) + return self.output_mlp(x) - Args: - x: [B, L, H, D] - cos: [B, L, D/2] - sin: [B, L, D/2] - """ - ro_dim = cos.shape[-1] * 2 - cos = cos.unsqueeze(2).repeat(1, 1, 1, 2) - sin = sin.unsqueeze(2).repeat(1, 1, 1, 2) - return torch.cat( - [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]], - dim=-1, - ) +# =========================================================================== +# LanguageModelShim +# =========================================================================== -@torch.compiler.disable -def build_3d_rope( - ref_pos: Tensor, - ref_space_uid: Tensor, - head_dim: int, - n_spatial_per_axis: int = 4, - n_uid_pairs: int = 2, - spatial_base_freq: float = 10000.0, - uid_base_freq: float = 10.0, -) -> tuple[Tensor, Tensor]: - """Build cos/sin for 3D RoPE + UID RoPE.""" - device = ref_pos.device - B, N = ref_pos.shape[:2] - half_dim = head_dim // 2 - n_spatial_total = 3 * n_spatial_per_axis - - spatial_inv_freq = 1.0 / ( - spatial_base_freq - ** ( - torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device) - / n_spatial_per_axis - ) - ) - uid_inv_freq = 1.0 / ( - uid_base_freq - ** ( - torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) - / n_uid_pairs - ) - ) - pos_f32 = ref_pos.float() - spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq) - spatial_freqs = spatial_freqs.reshape(B, N, n_spatial_total) +class LanguageModelShim(nn.Module): + """Shim holding the trainable projection weights for LM integration. - uid_f32 = ref_space_uid.float() - uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq) + Contains: + - base_z_combine: nn.Parameter with shape ``(n_layers + 1,)`` + - base_z_linear: Sequential(LayerNorm(d_model), Linear(d_model, d_z, bias=False)) + - base_z_mlp: Sequential(SingleToPair(d_z, d_z, d_z), LayerNorm(d_z)) + """ - n_active = n_spatial_total + n_uid_pairs - freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1) + def __init__(self, d_z: int = 256, d_model: int = 2560, num_layers: int = 80) -> None: + super().__init__() - if n_active < half_dim: - padding = torch.zeros( - B, N, half_dim - n_active, device=device, dtype=torch.float32 + self.base_z_mlp = nn.Sequential(SingleToPair(d_z, d_z, d_z), nn.LayerNorm(d_z)) + self.base_z_linear = nn.Sequential( + nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False) ) - freqs = torch.cat([freqs, padding], dim=-1) - - cos = freqs.cos().to(torch.bfloat16) - sin = freqs.sin().to(torch.bfloat16) - return cos, sin - + self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1)) -def qk_norm(x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),)).to(x.dtype) + def project_sequence( + self, + hidden_states: Tensor, + residue_mask: Tensor | None = None, + ) -> Tensor: + """Project all ESMC layer states into the learned sequence summary. + Args: + hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``. + residue_mask: Optional M with shape ``(b, l)``. Non-residue rows + are set to zero in the returned tensor. -# =========================================================================== -# SwiGLUFFN (atom transformer blocks) -# =========================================================================== + Returns: + Z with shape ``(b, l, d_z)``. + """ + if hidden_states.ndim != 4: + raise ValueError( + "H must have shape (b, l, n_layers + 1, d_model), " + f"got {tuple(hidden_states.shape)}." + ) + expected_layers = self.base_z_combine.numel() + if hidden_states.shape[-2] != expected_layers: + raise ValueError( + f"H contains {hidden_states.shape[-2]} states; expected " + f"{expected_layers} in the official ESMC ordering." + ) + expected_width = cast(nn.LayerNorm, self.base_z_linear[0]).normalized_shape[0] + if hidden_states.shape[-1] != expected_width: + raise ValueError(f"H has width {hidden_states.shape[-1]}; expected {expected_width}.") + + # H can be FP32 even when the folding checkpoint is loaded in BF16. + # Match the learned projection parameters at this explicit boundary; + # this preserves the official BF16 path and leaves FP32 models exact. + projection_dtype = cast(nn.LayerNorm, self.base_z_linear[0]).weight.dtype + hidden_states = hidden_states.to(dtype=projection_dtype) + projected_states = self.base_z_linear(hidden_states) + layer_weights = self.base_z_combine.softmax(dim=0) + # Preserve Biohub's matmul path exactly so checkpoint inference does + # not change through a different reduction order. + projected = layer_weights @ projected_states + if residue_mask is not None: + if residue_mask.shape != hidden_states.shape[:2]: + raise ValueError( + "M must have shape (b, l), got " + f"{tuple(residue_mask.shape)} for H {tuple(hidden_states.shape)}." + ) + projected = projected * residue_mask.to( + device=projected.device, dtype=projected.dtype + ).unsqueeze(-1) + return projected -class SwiGLUFFN(nn.Module): - """SwiGLU FFN with rounded hidden size for hardware alignment.""" + def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor: + """Project pre-computed ESMC hidden states to pair representation. - def __init__(self, d_model: int, expansion_ratio: int = 2) -> None: - super().__init__() - hidden_size = ((expansion_ratio * (d_model // 3) * 2) + 255) // 256 * 256 - self.w_up = nn.Linear(d_model, 2 * hidden_size, bias=False) - self.w_down = nn.Linear(hidden_size, d_model, bias=False) + Args: + hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``. + lm_dropout: Dropout probability applied to the pair + representation after ``base_z_mlp``. - def forward(self, x: Tensor) -> Tensor: - x = x.to(self.w_up.weight.dtype) - x1, x2 = self.w_up(x).chunk(2, dim=-1) - return self.w_down(F.silu(x1) * x2) + Returns: + Z_pair with shape ``(b, l, l, d_pair)``. + """ + lm_z = self.project_sequence(hidden_states) + lm_z = self.base_z_mlp(lm_z) + if lm_dropout > 0: + lm_z = F.dropout(lm_z, p=lm_dropout, training=True) + return lm_z # =========================================================================== -# SWA3DRoPEAttention +# ESMFold2ExperimentalModel: the top-level PreTrainedModel # =========================================================================== -class SWA3DRoPEAttention(nn.Module): - """Sliding window attention with 3D RoPE. Has Wqkv, gate_proj, out_proj.""" +def compute_lm_hidden_states( + esmc: nn.Module, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + token_mask: Tensor, + pad_to_multiple: int | None = None, + lm_mask_pct: float = 0.0, + mask_token_id: int = 32, +) -> Tensor: + """Run ESMC and return H with shape ``(b, l, n_states, d_model)``. - def __init__(self, d_model: int, n_heads: int, half_window: int = 64) -> None: - super().__init__() - self.n_heads = n_heads - self.head_dim = d_model // n_heads - self.scale = self.head_dim**-0.5 - self.half_window = half_window + Atom-tokenized modified residues (HYP, MSE, ACE, NH2, ...) span multiple + structure tokens but share a single ``(asym_id, residue_index)`` key: + collapse them to one LM token per residue before running the LM (the LM + was trained on per-residue inputs, not per-atom), then scatter the + hidden states back to the per-token layout. + """ + b_size, l_size = input_ids.shape + device = input_ids.device + protein_mask = (mol_type == 0) & token_mask - self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False) - self.out_proj = nn.Linear(d_model, d_model, bias=False) - self.gate_proj = nn.Linear(d_model, d_model, bias=False) + lm_input_list = [] + lm_lengths = [] + # Per-batch maps from (original protein-token index) to (LM input position). + expand_maps: list[Tensor] = [] + for batch_index in range(b_size): + mask_b = protein_mask[batch_index] + ids_b = input_ids[batch_index][mask_b] + asym_b = asym_id[batch_index][mask_b] + res_b = residue_index[batch_index][mask_b] - def forward(self, x: Tensor, attention_params: tuple) -> Tensor: - B, N = x.shape[:2] - cos, sin = attention_params[0], attention_params[1] + # Collapse: keep first token per (asym_id, residue_index) key, in + # input order. ``inverse`` maps each original protein-token to its + # collapsed residue index. + keys = torch.stack((asym_b, res_b), dim=1) + unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True) + n_unique = unique_keys.size(0) + token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long) + first_pos = torch.full((n_unique,), keys.size(0), device=device, dtype=torch.long) + first_pos.scatter_reduce_(0, inverse, token_positions, reduce="amin", include_self=True) + ordered = torch.argsort(first_pos) + first_pos_ordered = first_pos[ordered] + ids_collapsed = ids_b[first_pos_ordered] + asym_collapsed = asym_b[first_pos_ordered] + remap = torch.empty_like(ordered) + remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long) + inverse_ordered = remap[inverse] - x_input = x - qkv = self.Wqkv(x) - qkv = qkv.view(B, N, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4) - q, k, v = qkv.unbind(0) - q, k = qk_norm(q), qk_norm(k) + chain_ids = asym_collapsed.unique(sorted=True) + # [BOS] chain1 [EOS BOS] chain2 ... [EOS] + parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)] + # Per-chain LM positions accumulate; track them for the expand map. + per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long) + cursor = 1 # position 0 is the leading BOS + for i, cid in enumerate(chain_ids): + in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0] + parts.append(ids_collapsed[in_chain]) + per_token_lm_pos[in_chain] = torch.arange( + cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long + ) + cursor += in_chain.shape[0] + if i < len(chain_ids) - 1: + parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype)) + cursor += 2 # EOS + BOS + parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype)) + lm_seq = torch.cat(parts) + lm_input_list.append(lm_seq) + lm_lengths.append(lm_seq.shape[0]) - q = apply_rotary_emb_3d(q, cos, sin) - k = apply_rotary_emb_3d(k, cos, sin) + # Map each original protein-token position to its LM input position. + prot_pos_b = mask_b.nonzero(as_tuple=True)[0] + expand_map = torch.full((l_size,), -1, device=device, dtype=torch.long) + expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered] + expand_maps.append(expand_map) - input_dtype = q.dtype - if q.dtype not in (torch.float16, torch.bfloat16): - q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16() + # Pad the language-model input to its longest sequence. FP8 callers round + # l to a multiple of 16 for Transformer Engine kernels. + max_len = max(lm_lengths) + if pad_to_multiple is not None and pad_to_multiple > 1: + max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple + lm_input_ids = torch.full( + (b_size, max_len), + 1, + device=device, + dtype=input_ids.dtype, # PAD=1 + ) + for batch_index in range(b_size): + lm_input_ids[batch_index, : lm_lengths[batch_index]] = lm_input_list[batch_index] - if len(attention_params) > 2 and FLASH_ATTN_AVAILABLE: - indices, cu_seqlens, max_seqlen = ( - attention_params[2], - attention_params[3], - attention_params[4], - ) - q_unpad = index_first_axis( # type: ignore[misc] - q.reshape(-1, self.n_heads, self.head_dim), indices - ) - k_unpad = index_first_axis( # type: ignore[misc] - k.reshape(-1, self.n_heads, self.head_dim), indices - ) - v_unpad = index_first_axis( # type: ignore[misc] - v.reshape(-1, self.n_heads, self.head_dim), indices - ) - out_unpad = flash_attn_varlen_func( # type: ignore[misc] - q_unpad, - k_unpad, - v_unpad, - cu_seqlens, - cu_seqlens, - max_seqlen, - max_seqlen, - softmax_scale=self.scale, - window_size=(self.half_window, self.half_window), - ) - out = pad_input(out_unpad, indices, B, N) # type: ignore[misc] - elif FLASH_ATTN_AVAILABLE: - out = flash_attn_func( # type: ignore[misc] - q, - k, - v, - softmax_scale=self.scale, - window_size=(self.half_window, self.half_window), - ) - else: - # Fallback: standard attention (no SWA) - q_t = q.transpose(1, 2) - k_t = k.transpose(1, 2) - v_t = v.transpose(1, 2) - attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale - attn = F.softmax(attn, dim=-1) - out = torch.matmul(attn, v_t).transpose(1, 2) - - out = out.to(input_dtype).reshape(B, N, -1) # type: ignore[union-attr] - out = out * torch.sigmoid(self.gate_proj(x_input)) - return self.out_proj(out) + # sequence_id for chain-aware attention; PAD tokens get -1 (no attention). + sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1 # BOS=0 + sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1) # PAD=1 + if lm_mask_pct > 0.0: + special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2) + do_mask = (torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct) & ~special + lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id) -# =========================================================================== -# SWAAtomBlock, SWAAtomTransformer -# =========================================================================== + with torch.inference_mode(): + esmc_out = esmc(input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True) + hidden_stack = esmc_out.hidden_states + n_states, _, _, d_model = hidden_stack.shape + result = torch.zeros(b_size, l_size, n_states, d_model, device=device, dtype=hidden_stack.dtype) + for batch_index in range(b_size): + M_i = protein_mask[batch_index] + positions = expand_maps[batch_index][M_i] + gathered = hidden_stack[:, batch_index, positions, :].permute(1, 0, 2) + result[batch_index, M_i.nonzero(as_tuple=True)[0]] = gathered -def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: - return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift + return result.detach() -def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor: - return x + gate * y +# =========================================================================== +# TriangleMultiplicativeUpdate +# =========================================================================== -class SWAAtomBlock(nn.Module): - """adaLN-Zero + SWA attention + SwiGLU FFN. +class TriangleMultiplicativeBlock(nn.Module): + """Triangle multiplicative update block with gated signal routing.""" - Creates adaln_modulation = Sequential(SiLU(), Linear) -> keys like adaln_modulation.1.weight - """ + _FLOW_TO_EINSUM: ClassVar[dict[str, str]] = { + "outgoing": "bikd,bjkd->bijd", + "incoming": "bkid,bkjd->bijd", + } + _VALID_FLOWS = ("outgoing", "incoming") - def __init__( - self, - d_atom: int, - n_heads: int, - half_window: int = 64, - expansion_ratio: int = 2, - use_compile_fusions: bool = False, - ) -> None: + def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None: super().__init__() - self.attn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) - self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + if flow not in self._FLOW_TO_EINSUM: + raise ValueError(f"Invalid flow={flow!r}. Expected one of {self._VALID_FLOWS}.") - adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False) - nn.init.zeros_(adaln_linear.weight) - self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear) + self.input_channels = input_channels + self.latent_channels = latent_channels + self.flow = flow + self._einsum_equation = self._FLOW_TO_EINSUM[flow] + self.norm_start = nn.LayerNorm(self.input_channels, eps=_EPS) + self.norm_mix = nn.LayerNorm(self.latent_channels, eps=_EPS) + self.proj_bundle = nn.Linear(self.input_channels, 4 * self.latent_channels, bias=False) + self.proj_emit = nn.Linear(self.latent_channels, self.input_channels, bias=False) + self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False) - self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window) - self.ffn = SwiGLUFFN(d_atom, expansion_ratio) + self._use_kernels: bool = False + # Default chunked for memory on long sequences; tests override with + # ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16 + # parity checks. + self._chunk_size: int | None = 64 - self._rms_adaln = ( - torch.compile(_rms_adaln_raw) if use_compile_fusions else _rms_adaln_raw - ) - self._gated_residual = ( - torch.compile(_gated_residual_raw) - if use_compile_fusions - else _gated_residual_raw + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def split_kernel_weights(self) -> tuple[Tensor, Tensor]: + return ( + self.proj_bundle.weight[: 2 * self.latent_channels, :], + self.proj_bundle.weight[2 * self.latent_channels :, :], ) - def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor: - mod = self.adaln_modulation(c_l) - if mod.dim() == 2: - mod = mod.unsqueeze(1) - shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1) + def _kernel_flow_direction(self) -> str: + return self.flow - attn_input = self._rms_adaln(x, scale_a, shift_a) - attn_out = self.attn(attn_input, attention_params) - x = self._gated_residual(x, gate_a, attn_out) + def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor: + return torch.einsum(self._einsum_equation, left_stream, right_stream) - ffn_input = self._rms_adaln(x, scale_f, shift_f) - ffn_out = self.ffn(ffn_input) - x = self._gated_residual(x, gate_f, ffn_out) - return x + def _triangular_contract_chunked( + self, left_stream: Tensor, right_stream: Tensor, chunk_size: int + ) -> Tensor: + """Compute the triangular einsum in chunks along the output i-dimension.""" + length = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2] + chunks = [] + for start in range(0, length, chunk_size): + end = min(start + chunk_size, length) + if self.flow == "outgoing": + chunk = torch.einsum(self._einsum_equation, left_stream[:, start:end], right_stream) + else: + chunk = torch.einsum( + self._einsum_equation, left_stream[:, :, start:end], right_stream + ) + chunks.append(chunk) + return torch.cat(chunks, dim=1) + def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor: + if visibility is None: + visibility = pair_grid.new_ones(pair_grid.shape[:-1]) -class SWAAtomTransformer(nn.Module): - """Stack of SWAAtomBlocks.""" + if self._use_kernels: + p_in_weight, g_in_weight = self.split_kernel_weights() + return _cue_tri_mul( # type: ignore[misc] + pair_grid, + direction=self._kernel_flow_direction(), + mask=visibility, + norm_in_weight=self.norm_start.weight, + norm_in_bias=self.norm_start.bias, + p_in_weight=p_in_weight, + g_in_weight=g_in_weight, + norm_out_weight=self.norm_mix.weight, + norm_out_bias=self.norm_mix.bias, + p_out_weight=self.proj_emit.weight, + g_out_weight=self.proj_gate.weight, + eps=_EPS, + ) - def __init__( - self, - d_atom: int = 128, - n_blocks: int = 3, - n_heads: int = 4, - swa_window_size: int = 128, - expansion_ratio: int = 2, - spatial_rope_base_frequency: float = 20.0, - n_spatial_rope_pairs_per_axis: int = 2, - n_uid_rope_pairs: int = 10, - uid_rope_base_frequency: float = 10000.0, - ) -> None: - super().__init__() - self.swa_window_size = swa_window_size - self.head_dim = d_atom // n_heads - self.spatial_rope_base_frequency = spatial_rope_base_frequency - self.n_spatial_rope_pairs_per_axis = n_spatial_rope_pairs_per_axis - self.n_uid_rope_pairs = n_uid_rope_pairs - self.uid_rope_base_frequency = uid_rope_base_frequency + normalized_grid = self.norm_start(pair_grid) + bundled = self.proj_bundle(normalized_grid) + signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1) + routed = signal * torch.sigmoid(gate_logits) + routed = routed * visibility.unsqueeze(-1) - self.blocks = nn.ModuleList( - [ - SWAAtomBlock( - d_atom=d_atom, - n_heads=n_heads, - half_window=swa_window_size // 2, - expansion_ratio=expansion_ratio, - ) - for _ in range(n_blocks) - ] - ) + left_stream, right_stream = routed.float().chunk(2, dim=-1) + if self._chunk_size is not None: + contracted = self._triangular_contract_chunked( + left_stream, right_stream, self._chunk_size + ) + else: + contracted = self._triangular_contract(left_stream, right_stream) + mixed = self.proj_emit(self.norm_mix(contracted)) + output_gate = torch.sigmoid(self.proj_gate(normalized_grid)) + return mixed * output_gate - def _build_3d_rope( - self, ref_pos: Tensor, ref_space_uid: Tensor - ) -> tuple[Tensor, Tensor]: - return build_3d_rope( - ref_pos=ref_pos, - ref_space_uid=ref_space_uid, - head_dim=self.head_dim, - n_spatial_per_axis=self.n_spatial_rope_pairs_per_axis, - n_uid_pairs=self.n_uid_rope_pairs, - spatial_base_freq=self.spatial_rope_base_frequency, - uid_base_freq=self.uid_rope_base_frequency, + +class TriangleMultiplicativeUpdate(nn.Module): + """Thin wrapper exposing the triangular mixer with explicit orientation (v3).""" + + def __init__(self, dim: int = 128, _outgoing: bool = True) -> None: + super().__init__() + flow = "outgoing" if _outgoing else "incoming" + self._engine = TriangleMultiplicativeBlock( + input_channels=dim, latent_channels=dim, flow=flow ) - def forward( - self, - q_l: Tensor, - c_l: Tensor, - attention_params: tuple, - return_intermediates: bool = False, - ) -> Tensor | tuple[Tensor, list[Tensor]]: - intermediates: list[Tensor] = [] - for block in self.blocks: - q_l = block(q_l, c_l, attention_params) - if return_intermediates: - intermediates.append(q_l) - if return_intermediates: - return q_l, intermediates - return q_l + def set_kernel_backend(self, backend: str | None) -> None: + # Engine uses cueq when backend=="cuequivariance"; the "fused" backend + # routes through the parent PairUpdateBlock's fused path (bypassing this). + validate_kernel_backend(backend) + self._engine._use_kernels = backend == BACKEND_CUEQ + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._engine.set_chunk_size(chunk_size) + + def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor: + return self._engine(z, visibility=mask) # =========================================================================== -# ESMFold2AtomEncoder (for both inputs_embedder and diffusion_module) +# FoldingTrunk: Transition, PairUpdateBlock, FoldingTrunk # =========================================================================== -class ESMFold2AtomEncoder(nn.Module): - """SWA atom encoder with atom_linear, atom_norm, atom_to_token_linear, [coords_linear], atom_transformer. - - Args: - d_atom: atom hidden dim - d_token: token dim for atom_to_token aggregation - n_blocks, n_heads, swa_window_size, expansion_ratio: transformer params - structure_prediction: if True, creates coords_linear and uses full d_token - spatial_rope_base_frequency, n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs, uid_rope_base_frequency: 3D RoPE config - """ +class Transition(nn.Module): + """LN + SwiGLU FFN with addmm-fused residual; optional Triton LN+w12+SwiGLU kernel.""" - def __init__( - self, - d_atom: int = 128, - d_token: int = 768, - n_blocks: int = 3, - n_heads: int = 4, - swa_window_size: int = 128, - expansion_ratio: int = 2, - structure_prediction: bool = True, - spatial_rope_base_frequency: float = 20.0, - n_spatial_rope_pairs_per_axis: int = 2, - n_uid_rope_pairs: int = 10, - uid_rope_base_frequency: float = 10000.0, - ) -> None: + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: super().__init__() - self.d_atom = d_atom - self.d_token = d_token - self.structure_prediction = structure_prediction + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + # Default chunked; set_chunk_size(None) disables for bit-exact parity tests. + self._chunk_size: int | None = 64 + self._fused_swiglu: nn.Module | None = None + self._kernel_backend: str | None = None - self.atom_linear = nn.Linear(ATOM_FEATURE_DIM, d_atom, bias=False) - self.atom_norm = nn.LayerNorm(d_atom) + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size - if structure_prediction: - self.coords_linear = nn.Linear(6, d_atom, bias=False) + def set_kernel_backend(self, backend: str | None) -> None: + """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent).""" + validate_kernel_backend(backend) + self._kernel_backend = backend + if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE: + assert _FusedLNLinearSwiGLU is not None + d_model = self.norm.normalized_shape[0] + d_inner = self.ffn.hidden_features + has_ln_bias = self.norm.bias is not None + device = self.ffn.w12.weight.device + dtype = self.ffn.w12.weight.dtype + fused = _FusedLNLinearSwiGLU( + d_model=d_model, + d_inner=d_inner, + has_ln_bias=has_ln_bias, + device=device, + dtype=dtype, + ) + with torch.no_grad(): + fused.LN_W.copy_(self.norm.weight) + if has_ln_bias: + fused.LN_B.copy_(self.norm.bias) # type: ignore[union-attr] + # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once. + fused.W12.copy_(self.ffn.w12.weight.t().contiguous()) + self._fused_swiglu = fused.eval().requires_grad_(False) + else: + self._fused_swiglu = None - self.atom_transformer = SWAAtomTransformer( - d_atom=d_atom, - n_blocks=n_blocks, - n_heads=n_heads, - swa_window_size=swa_window_size, - expansion_ratio=expansion_ratio, - spatial_rope_base_frequency=spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=n_uid_rope_pairs, - uid_rope_base_frequency=uid_rope_base_frequency, + def _can_use_fused_path(self, x: Tensor) -> bool: + return ( + _fused_active(self, x) and self._fused_swiglu is not None and x.dtype == torch.bfloat16 ) - # Output aggregation: d_token for structure prediction, d_token//2 for inputs - out_dim = d_token if structure_prediction else d_token // 2 - self.atom_to_token_linear = nn.Linear(d_atom, out_dim, bias=False) - - def forward( - self, - ref_pos: Tensor, - atom_attention_mask: Tensor, - ref_space_uid: Tensor, - ref_charge: Tensor, - ref_element: Tensor, - ref_atom_name_chars: Tensor, - atom_to_token: Tensor, - r_l: Tensor | None = None, - pred_r1: Tensor | None = None, - s_i: Tensor | None = None, - z_ij: Tensor | None = None, - num_diffusion_samples: int = 1, - return_intermediates: bool = False, - inference_cache: dict | None = None, - ) -> tuple[Tensor, Tensor, Tensor, tuple, list[Tensor]]: - """Returns (a, q, c, attention_params, intermediates). - - ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE, - attention indices, n_tokens) across diffusion steps. - """ - B, N = ref_pos.shape[:2] + def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor: + """SwiGLU through silu(x1)*x2, before the final w3.""" + ffn = self.ffn + x12 = ffn.w12(x_normed) + x1, x2 = x12.split(ffn.hidden_features, dim=-1) + return F.silu(x1) * x2 - layer_cache = None - if inference_cache is not None: - layer_cache = inference_cache.setdefault("atomencoder", {}) + def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor: + """x + w3(hidden) via single cuBLAS addmm: avoids transition-output allocation.""" + ffn = self.ffn + x_shape = x.shape + out = torch.addmm( + x.contiguous().view(-1, x_shape[-1]), + hidden.view(-1, hidden.shape[-1]), + ffn.w3.weight.t(), + ) + return out.view(x_shape) - if layer_cache is None or len(layer_cache) == 0: - atom_feats = torch.cat( - [ - ref_pos, - ref_charge.unsqueeze(-1), - atom_attention_mask.unsqueeze(-1), - ref_element, - ref_atom_name_chars.reshape(B, N, MAX_CHARS * CHAR_VOCAB_SIZE), - ], - dim=-1, - ) - c_base = self.atom_norm(self.atom_linear(atom_feats)) - cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid) - cos = cos.repeat_interleave(num_diffusion_samples, 0) - sin = sin.repeat_interleave(num_diffusion_samples, 0) - mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0) - seqlens = mask_exp.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten() - max_seqlen = int(seqlens.max().item()) - cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) - attention_params = (cos, sin, indices, cu_seqlens, max_seqlen) - n_tokens = int(atom_to_token.max().item()) + 1 - if layer_cache is not None: - layer_cache["c_base"] = c_base - layer_cache["attention_params"] = attention_params - layer_cache["mask_exp"] = mask_exp - layer_cache["n_tokens"] = n_tokens - layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave( - num_diffusion_samples, 0 - ) - else: - c_base = layer_cache["c_base"] - attention_params = layer_cache["attention_params"] - mask_exp = layer_cache["mask_exp"] - n_tokens = layer_cache["n_tokens"] - - c = c_base + def forward(self, x: Tensor) -> Tensor: + # Inference-only fast path (addmm-fused residual + pre-alloc out) + #: diverges bit-exactly from ``x + ffn(norm(x))`` so we only use + # it when grad is disabled (binder-design / bit-exact tests run + # with grad on and need the reference path). + if not torch.is_grad_enabled() and self._can_use_fused_path(x): + fused = self._fused_swiglu + assert fused is not None + pre_w3 = fused + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + hidden = pre_w3(x) + return self._addmm_residual(x, hidden) + out = torch.empty_like(x) + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + hidden = pre_w3(sl) + out[:, s:e] = self._addmm_residual(sl, hidden) + return out + # Reference path: bit-exact with main: x + ffn(norm(x)). + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + return x + self.ffn(self.norm(x)) + out_list: list[Tensor] = [] + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + out_list.append(sl + self.ffn(self.norm(sl))) + return torch.cat(out_list, dim=1) - q = c - if self.structure_prediction and r_l is not None: - q = q.repeat_interleave(num_diffusion_samples, 0) - if pred_r1 is None: - pred_r1 = torch.zeros_like(r_l) - r_input = torch.cat([r_l, pred_r1], dim=-1) - r_to_q = self.coords_linear(r_input) - q = q + r_to_q +class PairUpdateBlock(nn.Module): + """tri_mul_out, tri_mul_in, pair_transition.""" - c = c.repeat_interleave(num_diffusion_samples, 0) + def __init__(self, d_pair: int = 256, expansion_ratio: int = 4) -> None: + super().__init__() + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = Transition(d_pair, expansion_ratio=expansion_ratio) + self._kernel_backend: str | None = None + # Row-shared dropout-residual; r=0 for inference (HF model is inference-only). + # backend='fused' swaps in the FusedDropoutResidual Triton kernel. + self.row_drop = DropoutResidual(0.0, batch_dim=1, use_fused_kernels=False) - result = self.atom_transformer( - q_l=q, - c_l=c, - attention_params=attention_params, - return_intermediates=return_intermediates, + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + self.tri_mul_out.set_kernel_backend(backend) + self.tri_mul_in.set_kernel_backend(backend) + self.pair_transition.set_kernel_backend(backend) + self._kernel_backend = backend + self.row_drop = DropoutResidual( + 0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED) ) - if return_intermediates: - q, intermediates = result - else: - q = result - intermediates = [] - q_to_a = F.relu(self.atom_to_token_linear(q)) - if layer_cache is not None and "atom_to_token_exp" in layer_cache: - atom_to_token_exp = layer_cache["atom_to_token_exp"] - else: - atom_to_token_exp = atom_to_token.repeat_interleave( - num_diffusion_samples, 0 - ) - a = scatter_atom_to_token( - q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool() - ) + def set_chunk_size(self, chunk_size: int | None) -> None: + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + self.pair_transition.set_chunk_size(chunk_size) - return a, q, c, attention_params, intermediates + def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool: + return _fused_active(self, pair) and pair.dtype == torch.bfloat16 + def _fused_trimul_with_residual( + self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None + ) -> Tensor: + """Fused TriMul+residual call; weights from the corresponding engine.""" + tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in + engine: TriangleMultiplicativeBlock = tri._engine # type: ignore[assignment] + p_in_weight, g_in_weight = engine.split_kernel_weights() -# =========================================================================== -# ESMFold2AtomDecoder -# =========================================================================== + def _bf16(t: Tensor) -> Tensor: + return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16) + return _fused_trimul_with_residual( # type: ignore[misc] + pair, + direction, + residual=pair, + drop_mask=None, # inference: no dropout, matches internal's eval path + norm_in_weight=_bf16(engine.norm_start.weight), + norm_in_bias=_bf16(engine.norm_start.bias), + p_in_weight=_bf16(p_in_weight), + g_in_weight=_bf16(g_in_weight), + norm_out_weight=_bf16(engine.norm_mix.weight), + norm_out_bias=_bf16(engine.norm_mix.bias), + p_out_weight=_bf16(engine.proj_emit.weight), + g_out_weight=_bf16(engine.proj_gate.weight), + mask=pair_attention_mask, + eps=_EPS, + ) -class ESMFold2AtomDecoder(nn.Module): - """SWA atom decoder with token_to_atom_linear, atom_transformer, norm, output_linear.""" + def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor: + if self._can_use_fused_trimul_with_residual(pair): + pair = self._fused_trimul_with_residual(pair, "outgoing", pair_attention_mask) + pair = self._fused_trimul_with_residual(pair, "incoming", pair_attention_mask) + else: + pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask)) + pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask)) + pair = self.pair_transition(pair) + return pair - def __init__( - self, - d_atom: int = 128, - d_token: int = 768, - n_blocks: int = 3, - n_heads: int = 4, - swa_window_size: int = 128, - expansion_ratio: int = 2, - spatial_rope_base_frequency: float = 20.0, - n_spatial_rope_pairs_per_axis: int = 2, - n_uid_rope_pairs: int = 10, - uid_rope_base_frequency: float = 10000.0, - ) -> None: - super().__init__() - self.token_to_atom_linear = nn.Linear(d_token, d_atom, bias=False) - self.atom_transformer = SWAAtomTransformer( - d_atom=d_atom, - n_blocks=n_blocks, - n_heads=n_heads, - swa_window_size=swa_window_size, - expansion_ratio=expansion_ratio, - spatial_rope_base_frequency=spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=n_uid_rope_pairs, - uid_rope_base_frequency=uid_rope_base_frequency, +class FoldingTrunk(nn.Module): + """ModuleList of PairUpdateBlocks.""" + + def __init__(self, n_layers: int = 24, d_pair: int = 256, expansion_ratio: int = 4) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [ + PairUpdateBlock(d_pair=d_pair, expansion_ratio=expansion_ratio) + for _ in range(n_layers) + ] ) - self.norm = nn.LayerNorm(d_atom) - self.output_linear = nn.Linear(d_atom, XYZ_DIMS, bias=False) + def set_kernel_backend(self, backend: str | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_kernel_backend(backend) - def forward( - self, - a_i: Tensor, - q_l: Tensor, - c_l: Tensor, - p_lm: tuple, - atom_to_token: Tensor, - atom_attention_mask: Tensor, - num_diffusion_samples: int = 1, - return_intermediates: bool = False, - ) -> tuple[Tensor, list[Tensor]]: - """Returns (r_update, intermediates).""" - atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) - a_to_q = self.token_to_atom_linear(a_i) - a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp) - q_l = q_l + a_to_q + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_chunk_size(chunk_size) - result = self.atom_transformer( - q_l=q_l, - c_l=c_l, - attention_params=p_lm, - return_intermediates=return_intermediates, + def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor: + # Cast the pair tensor to BF16 when the fused triangle backend is enabled + # (its bwd kernel requires bf16). Other backends keep the input dtype. + orig_dtype = pair.dtype + fused_on = ( + len(self.blocks) > 0 + and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED ) - if return_intermediates: - q_l, intermediates = result - else: - q_l = result - intermediates = [] - - r_l = self.output_linear(self.norm(q_l)) - return r_l, intermediates + if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16: + pair = pair.to(torch.bfloat16) + for block in self.blocks: + fn = partial(block, pair_attention_mask=pair_attention_mask) + if torch.is_grad_enabled(): + pair = checkpoint(fn, pair, use_reentrant=False) # pyright: ignore + else: + pair = fn(pair) + if pair.dtype != orig_dtype: + pair = pair.to(orig_dtype) + return pair # =========================================================================== -# AttentionPairBias (DiffusionTransformer attention block) +# MSA Encoder # =========================================================================== -class AttentionPairBias(nn.Module): - """Gated multi-head attention with pair bias conditioning.""" +class OuterProductMean(nn.Module): + """Outer-product mean: maps an MSA representation into a pair update. + + The order of the ``/ n_valid`` divide vs. the ``Wout`` projection is + selectable via ``divide_outer_before_proj`` because different ESMFold2 + checkpoints were trained with different orderings: + + * ``False`` (default): ``Wout(outer) / n_valid``: the projection bias + is scaled by 1/n_valid alongside the outer product. + * ``True``: ``Wout(outer / n_valid)``: the projection bias is added + unscaled, post-divide. + """ def __init__( self, - d_model: int, + d_msa: int, + d_hidden: int, d_pair: int, - num_heads: int, - d_cond: int | None = None, - use_conditioning: bool = True, + divide_outer_before_proj: bool = False, ) -> None: super().__init__() - self.d_model = d_model - self.num_heads = num_heads - self.head_dim = d_model // num_heads - self.scale = self.head_dim**-0.5 - d_cond = d_cond or d_model + self.d_hidden = d_hidden + self.divide_outer_before_proj = divide_outer_before_proj + self.norm = nn.LayerNorm(d_msa) + self.W = nn.Linear(d_msa, 2 * d_hidden, bias=False) + self.Wout = nn.Linear(d_hidden * d_hidden, d_pair, bias=True) + # Off for bit-exact bf16; ``set_chunk_size(64)`` for long sequences. + self._chunk_size: int | None = None - if use_conditioning: - self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) - self.out_gate = nn.Linear(d_cond, d_model, bias=True) - # adaln init: weight=0, bias=-2 - nn.init.zeros_(self.out_gate.weight) - nn.init.constant_(self.out_gate.bias, -2.0) - else: - self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) - - self.q_proj = nn.Linear(d_model, d_model, bias=True) - self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False) - self.g_proj = nn.Linear(d_model, d_model, bias=False) - self.out_proj = nn.Linear(d_model, d_model, bias=False) - - if d_pair > 0: - self.pair_norm = nn.LayerNorm(d_pair, eps=1e-5) - self.pair_bias_proj = nn.Linear(d_pair, num_heads, bias=False) - - self._kernel_backend: str | None = None + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size - def set_kernel_backend(self, backend: str | None) -> None: - if backend not in _VALID_BACKENDS: - raise ValueError( - f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" - ) - self._kernel_backend = backend + def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor: + m_norm = self.norm(m) + x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype) + a, b = x.chunk(2, dim=-1) + mask_f = msa_attention_mask.to(a.dtype) + n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0) + if self._chunk_size is None: + outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2) + if self.divide_outer_before_proj: + return self.Wout(outer / n_valid) + return self.Wout(outer) / n_valid + # Chunk along the left (i) axis so the peak einsum intermediate is + # X uses shape (b, chunk, l, c, d) instead of (b, l, l, c, d). + length = a.shape[1] + out_chunks: list[Tensor] = [] + for start in range(0, length, self._chunk_size): + end = min(start + self._chunk_size, length) + outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, start:end], b).flatten(-2) + if self.divide_outer_before_proj: + out_chunks.append(self.Wout(outer_chunk / n_valid[:, start:end])) + else: + out_chunks.append(self.Wout(outer_chunk) / n_valid[:, start:end]) + return torch.cat(out_chunks, dim=1) - def _is_zero_beta(self, beta: Tensor | float) -> bool: - if isinstance(beta, (int, float)): - return beta == 0.0 - return bool((beta == 0).all()) - def _can_use_fused_pair_bias( - self, z: Tensor, n_queries: int, beta: Tensor | float - ) -> bool: - return ( - _fused_active(self, z) - and z.dim() == 4 - and self._is_zero_beta(beta) - and hasattr(self, "pair_bias_proj") - and hasattr(self, "pair_norm") - ) +class MSAPairWeightedAveraging(nn.Module): + """Pair-biased MSA row update (AF3 Supplement Algorithm 10).""" - def _can_use_cueq_pair_bias( - self, z: Tensor, n_queries: int, beta: Tensor | float - ) -> bool: - return ( - _cueq_active(self) - and n_queries > 750 - and z.dim() == 4 - and self._is_zero_beta(beta) - and hasattr(self, "pair_bias_proj") + def __init__(self, d_msa: int, d_pair: int, n_heads: int = 8, head_width: int = 32) -> None: + super().__init__() + self.n_heads = n_heads + self.head_width = head_width + self.norm_single = nn.LayerNorm(d_msa) + self.compute_bias = nn.Sequential( + nn.LayerNorm(d_pair), nn.Linear(d_pair, n_heads, bias=False) ) + self.Wv = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wgate = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wout = nn.Linear(n_heads * head_width, d_msa, bias=False) - def forward( - self, - a: Tensor, - s: Tensor | None, - z: Tensor, - beta: Tensor | float = 0.0, - attention_mask: Tensor | None = None, - num_diffusion_samples: int = 1, - ) -> Tensor: - bsz, n_queries, d_model = a.shape - - if s is not None: - x = self.adaln(a, s) - else: - x = self.pre_norm(a) + def forward(self, msa_repr: Tensor, pair_repr: Tensor, pair_attention_mask: Tensor) -> Tensor: + """ + Args: + msa_repr: X with shape (b, l, m, d_msa). + pair_repr: Z with shape (b, l, l, d_pair). + pair_attention_mask: M with shape (b, l, l). + Returns: + X with shape (b, l, m, d_msa). + """ + batch_size, length, depth, _ = msa_repr.shape + n_heads, head_width = self.n_heads, self.head_width - n_keys = x.shape[1] - q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim) - kv = self.kv_proj(x) - k, v = kv.chunk(2, dim=-1) - k = k.view(bsz, n_keys, self.num_heads, self.head_dim) - v = v.view(bsz, n_keys, self.num_heads, self.head_dim) + msa_normed = self.norm_single(msa_repr) + bias = self.compute_bias(pair_repr) # A has shape (b, l, l, n_heads). + bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5) + attn = torch.softmax(bias, dim=-2) # softmax over j - # Expand z for num_diffusion_samples - if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1: - z = z.repeat_interleave(num_diffusion_samples, dim=0) - if ( - attention_mask is not None - and attention_mask.shape[0] != bsz - and num_diffusion_samples > 1 - ): - attention_mask = attention_mask.repeat_interleave( - num_diffusion_samples, dim=0 - ) + v = self.Wv(msa_normed).reshape(batch_size, length, depth, n_heads, head_width) + gate = torch.sigmoid(self.Wgate(msa_normed)).reshape( + batch_size, length, depth, n_heads, head_width + ) - if self._can_use_fused_pair_bias(z, n_queries, beta): - kernel_mask = ( - attention_mask - if attention_mask is not None - else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) - ) - pair_norm_w = self.pair_norm.weight - pair_norm_b = ( - self.pair_norm.bias - if self.pair_norm.bias is not None - else torch.zeros_like(pair_norm_w) - ) - z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16) - bias = _fused_pair_bias( # type: ignore[misc] - z_bf, - kernel_mask, - self.pair_bias_proj.weight, - num_heads=self.num_heads, - pair_norm_w=pair_norm_w, - pair_norm_b=pair_norm_b, - ) # (B, H, Q, K) - q_bhqd = q.transpose(1, 2) - k_bhqd = k.transpose(1, 2) - v_bhqd = v.transpose(1, 2) - attn_out = F.scaled_dot_product_attention( - q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype) - ) - g = torch.sigmoid(self.g_proj(x)).view( - bsz, n_queries, self.num_heads, self.head_dim - ) - ctx = g * attn_out.transpose(1, 2) - out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) - if s is not None: - out = torch.sigmoid(self.out_gate(s)) * out - return out + output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate) + return self.Wout(output.reshape(batch_size, length, depth, n_heads * head_width)) - if self._can_use_cueq_pair_bias(z, n_queries, beta): - kernel_mask = ( - attention_mask - if attention_mask is not None - else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) - ) - out, _ = _cue_attn_pair_bias( # type: ignore[misc] - s=x, - q=q.transpose(1, 2), - k=k.transpose(1, 2), - v=v.transpose(1, 2), - z=z, - mask=kernel_mask, - num_heads=self.num_heads, - w_proj_z=self.pair_bias_proj.weight, - w_proj_g=self.g_proj.weight, - w_proj_o=self.out_proj.weight, - w_ln_z=self.pair_norm.weight, - b_ln_z=self.pair_norm.bias, - return_z_proj=False, - is_cached_z_proj=False, - ) - else: - # Standard attention with pair bias - g = torch.sigmoid(self.g_proj(x)).view( - bsz, n_queries, self.num_heads, self.head_dim - ) - logits = ( - torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale - ) +# =========================================================================== +# Atom and diffusion stack +# =========================================================================== - if z.dim() == 4: - pair_bias = self.pair_bias_proj(self.pair_norm(z)) - else: - pair_bias = z.unsqueeze(-1) - logits = logits + pair_bias.to(dtype=logits.dtype) - if attention_mask is not None: - min_val = torch.finfo(logits.dtype).min - mask_bias = torch.where( - attention_mask.bool()[:, None, :, None], 0.0, min_val - ) - logits = logits + mask_bias.to(dtype=logits.dtype) +class TransitionLayer(nn.Module): + """SwiGLU transition: norm -> a_proj, b_proj -> silu(a)*b -> out_proj.""" - attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype) - ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v) - ctx = g * ctx - out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + def __init__(self, d_model: int, n: int, eps: float = 1e-5) -> None: + super().__init__() + hidden = n * d_model + self.norm = nn.LayerNorm(d_model, eps=eps) + self.a_proj = nn.Linear(d_model, hidden, bias=False) + self.b_proj = nn.Linear(d_model, hidden, bias=False) + self.out_proj = nn.Linear(hidden, d_model, bias=False) - if s is not None: - out = torch.sigmoid(self.out_gate(s)) * out - return out + def forward(self, x: Tensor) -> Tensor: + x = self.norm(x) + a = self.a_proj(x) + b = self.b_proj(x) + return self.out_proj(F.silu(a) * b) # =========================================================================== -# ConditionedTransitionBlock +# AdaptiveLayerNorm (used in DiffusionTransformer) # =========================================================================== -class ConditionedTransitionBlock(nn.Module): - """Conditioned SwiGLU transition with adaptive layer norm.""" +class AdaptiveLayerNorm(nn.Module): + """Adaptive layer normalization (adaLN-Zero).""" - def __init__( - self, - d_model: int, - d_cond: int | None = None, - transition_multiplier: int = 2, - use_conditioning: bool = True, - ) -> None: + def __init__(self, d_model: int, d_cond: int, eps: float = 1e-5) -> None: super().__init__() - d_cond = d_cond or d_model - hidden = transition_multiplier * d_model + self.d_model = d_model + self.d_cond = d_cond + self.eps = eps + self.s_scale = nn.Parameter(torch.ones(d_cond)) + self.s_gate = nn.Linear(d_cond, d_model, bias=True) + self.s_shift = nn.Linear(d_cond, d_model, bias=False) - if use_conditioning: - self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) - self.output_gate = nn.Linear(d_cond, d_model, bias=True) - nn.init.zeros_(self.output_gate.weight) - nn.init.constant_(self.output_gate.bias, -2.0) - else: - self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + def forward(self, a: Tensor, s: Tensor) -> Tensor: + a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps) + s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps) + return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm) - self.lin_swish = nn.Linear(d_model, 2 * hidden, bias=False) - self.lin_out = nn.Linear(hidden, d_model, bias=False) - def forward(self, a: Tensor, s: Tensor | None) -> Tensor: - if s is not None: - x = self.adaln(a, s) - else: - x = self.pre_norm(a) +# =========================================================================== +# FourierEmbedding +# =========================================================================== - swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1) - b = F.silu(swish_a) * swish_b - out = self.lin_out(b) - if s is not None: - out = torch.sigmoid(self.output_gate(s)) * out - return out +class FourierEmbedding(nn.Module): + """Fourier embedding: cos(2*pi*(t*w + b)).""" + + w: Tensor + b: Tensor + + def __init__(self, c: int) -> None: + super().__init__() + self.c = c + self.register_buffer("w", torch.randn(c)) + self.register_buffer("b", torch.randn(c)) + + def forward(self, t_hat: Tensor) -> Tensor: + t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1) + return torch.cos(2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :])) # =========================================================================== -# DiffusionTransformer (token transformer) +# SwiGLU / SwiGLUMLP # =========================================================================== -class DiffusionTransformer(nn.Module): - """Diffusion denoising transformer with attention pair bias.""" +def _compute_swiglu_hidden_size(d_model: int, expansion_ratio: int) -> int: + return expansion_ratio * d_model + + +class SwiGLU(nn.Module): + """SwiGLU with packed w12 and output w3.""" def __init__( self, - d_model: int, - d_pair: int, - num_heads: int, - num_blocks: int, - d_cond: int | None = None, - transition_multiplier: int = 2, - use_conditioning: bool = True, + in_features: int, + hidden_features: int, + out_features: int | None = None, + bias: bool = True, ) -> None: super().__init__() - d_cond = d_cond or d_model + out_features = out_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + self.hidden_features = hidden_features - self.attn_blocks = nn.ModuleList( - [ - AttentionPairBias( - d_model=d_model, - d_pair=d_pair, - num_heads=num_heads, - d_cond=d_cond, - use_conditioning=use_conditioning, - ) - for _ in range(num_blocks) - ] - ) - self.transition_blocks = nn.ModuleList( - [ - ConditionedTransitionBlock( - d_model=d_model, - d_cond=d_cond, - transition_multiplier=transition_multiplier, - use_conditioning=use_conditioning, - ) - for _ in range(num_blocks) - ] - ) + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.split(self.hidden_features, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) - def set_kernel_backend(self, backend: str | None) -> None: - for attn in self.attn_blocks: - cast(AttentionPairBias, attn).set_kernel_backend(backend) - def forward( - self, - a: Tensor, - s: Tensor | None, - z: Tensor, - beta: Tensor | float = 0.0, - attention_mask: Tensor | None = None, - num_diffusion_samples: int = 1, - return_intermediates: bool = False, - ) -> tuple[Tensor, list[Tensor]]: - intermediates: list[Tensor] = [] - x = a - for attn, transition in zip(self.attn_blocks, self.transition_blocks): - x = x + attn( - x, - s, - z, - beta, - attention_mask=attention_mask, - num_diffusion_samples=num_diffusion_samples, - ) - x = x + transition(x, s) - if return_intermediates: - intermediates.append(x) - return x, intermediates +class SwiGLUMLP(SwiGLU): + """SwiGLU MLP with packed weights, no bias.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4, bias: bool = False) -> None: + hidden = _compute_swiglu_hidden_size(d_model, expansion_ratio) + super().__init__( + in_features=d_model, hidden_features=hidden, out_features=d_model, bias=bias + ) # =========================================================================== -# DiffusionConditioning +# SWA Atom Attention components # =========================================================================== -class DiffusionConditioning(nn.Module): - """Conditions pair and single representations on noise timestep.""" +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) - def __init__( - self, - c_z: int = 256, - c_s: int = 768, - c_s_inputs: int = 451, - sigma_data: float = 16.0, - fourier_dim: int = 256, - transition_multiplier: int = 2, - layer_norm_eps: float = 1e-5, - ) -> None: - super().__init__() - self.sigma_data = float(sigma_data) - self.c_z = c_z - self.c_s = c_s - self.c_s_inputs = c_s_inputs - self.z_input_norm = nn.LayerNorm(2 * c_z, eps=layer_norm_eps) - self.z_proj = nn.Linear(2 * c_z, c_z, bias=False) - self.z_transitions = nn.ModuleList( - [ - TransitionLayer(c_z, n=transition_multiplier, eps=layer_norm_eps) - for _ in range(2) - ] - ) +def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: + """Apply RoPE with batch-dependent cos/sin. - self.s_input_norm = nn.LayerNorm(c_s_inputs, eps=layer_norm_eps) - self.s_proj = nn.Linear(c_s_inputs, c_s, bias=False) - self.fourier = FourierEmbedding(fourier_dim) - self.noise_norm = nn.LayerNorm(fourier_dim, eps=layer_norm_eps) - self.noise_proj = nn.Linear(fourier_dim, c_s, bias=False) - self.s_transitions = nn.ModuleList( - [ - TransitionLayer(c_s, n=transition_multiplier, eps=layer_norm_eps) - for _ in range(2) - ] + Args: + x: X with shape (b, l, h, d). + cos: C with shape (b, l, d / 2). + sin: S with shape (b, l, d / 2). + """ + ro_dim = cos.shape[-1] * 2 + cos = cos.unsqueeze(2).repeat(1, 1, 1, 2) + sin = sin.unsqueeze(2).repeat(1, 1, 1, 2) + return torch.cat( + [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]], + dim=-1, + ) + + +@torch.compiler.disable +def build_3d_rope( + ref_pos: Tensor, + ref_space_uid: Tensor, + head_dim: int, + n_spatial_per_axis: int = 4, + n_uid_pairs: int = 2, + spatial_base_freq: float = 10000.0, + uid_base_freq: float = 10.0, +) -> tuple[Tensor, Tensor]: + """Build cos/sin for 3D RoPE + UID RoPE.""" + device = ref_pos.device + batch_size, n_atoms = ref_pos.shape[:2] + half_dim = head_dim // 2 + n_spatial_total = 3 * n_spatial_per_axis + + spatial_inv_freq = 1.0 / ( + spatial_base_freq + ** ( + torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device) + / n_spatial_per_axis ) + ) + uid_inv_freq = 1.0 / ( + uid_base_freq + ** (torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) / n_uid_pairs) + ) - def forward( - self, - t_hat: Tensor, - s_inputs: Tensor, - s_trunk: Tensor | None, - z_trunk: Tensor, - relative_position_encoding: Tensor, - sigma_data: float | None = None, - num_diffusion_samples: int = 1, - inference_cache: dict[str, Tensor] | None = None, - ) -> tuple[Tensor, Tensor]: - sigma = self.sigma_data if sigma_data is None else float(sigma_data) - base_batch = z_trunk.shape[0] - target_batch = base_batch * num_diffusion_samples + pos_f32 = ref_pos.float() + spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq) + spatial_freqs = spatial_freqs.reshape(batch_size, n_atoms, n_spatial_total) - # z conditioning (cached across diffusion steps — independent of t_hat) - if inference_cache is not None and "z" in inference_cache: - z = inference_cache["z"] - else: - z_rel = relative_position_encoding.to(dtype=torch.float32) - z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1) - z = self.z_proj(self.z_input_norm(z)) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - for block in self.z_transitions: - z = z + block(z) - if inference_cache is not None: - inference_cache["z"] = z + uid_f32 = ref_space_uid.float() + uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq) - # s conditioning - s_inputs_eff = s_inputs - if s_inputs_eff.shape[0] != target_batch: - s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0) + n_active = n_spatial_total + n_uid_pairs + freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1) - s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32))) + if n_active < half_dim: + padding = torch.zeros( + batch_size, + n_atoms, + half_dim - n_active, + device=device, + dtype=torch.float32, + ) + freqs = torch.cat([freqs, padding], dim=-1) - # Noise embedding - t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1) - if t.numel() == 1: - t = t.expand(target_batch) - elif t.shape[0] != target_batch: - t = t.repeat_interleave(num_diffusion_samples, 0) - t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20)) - n = self.fourier(t_noise) - n = self.noise_proj(self.noise_norm(n)) - s = s + n.unsqueeze(1) + cos = freqs.cos().to(torch.bfloat16) + sin = freqs.sin().to(torch.bfloat16) + return cos, sin - for block in self.s_transitions: - s = s + block(s) - return s, z +def qk_norm(x: Tensor) -> Tensor: + return F.rms_norm(x, (x.size(-1),)).to(x.dtype) # =========================================================================== -# DiffusionModule +# SwiGLUFFN (atom transformer blocks) # =========================================================================== -class DiffusionModule(nn.Module): - """Diffusion denoising module for structure prediction.""" +class SwiGLUFFN(nn.Module): + """SwiGLU FFN with rounded hidden size for hardware alignment.""" - def __init__( - self, - c_atom: int = 128, - c_token: int = 768, - c_z: int = 256, - c_s_inputs: int = 451, - sigma_data: float = 16.0, - fourier_dim: int = 256, - atom_num_blocks: int = 3, - atom_num_heads: int = 4, - token_num_blocks: int = 12, - token_num_heads: int = 16, - transition_multiplier: int = 2, - swa_window_size: int = 128, - spatial_rope_base_frequency: float = 20.0, - n_spatial_rope_pairs_per_axis: int = 2, - n_uid_rope_pairs: int = 10, - uid_rope_base_frequency: float = 10000.0, - ) -> None: + def __init__(self, d_model: int, expansion_ratio: int = 2) -> None: super().__init__() - self.sigma_data = float(sigma_data) - - self.conditioning = DiffusionConditioning( - c_z=c_z, - c_s=c_token, # conditioning s output is c_token - c_s_inputs=c_s_inputs, - sigma_data=sigma_data, - fourier_dim=fourier_dim, - transition_multiplier=transition_multiplier, - ) + hidden_size = ((expansion_ratio * (d_model // 3) * 2) + 255) // 256 * 256 + self.w_up = nn.Linear(d_model, 2 * hidden_size, bias=False) + self.w_down = nn.Linear(hidden_size, d_model, bias=False) - # Atom encoder (structure_prediction=True, with coords_linear) - self.atom_encoder = ESMFold2AtomEncoder( - d_atom=c_atom, - d_token=c_token, - n_blocks=atom_num_blocks, - n_heads=atom_num_heads, - swa_window_size=swa_window_size, - expansion_ratio=2, - structure_prediction=True, - spatial_rope_base_frequency=spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=n_uid_rope_pairs, - uid_rope_base_frequency=uid_rope_base_frequency, - ) + def forward(self, x: Tensor) -> Tensor: + x = x.to(self.w_up.weight.dtype) + x1, x2 = self.w_up(x).chunk(2, dim=-1) + return self.w_down(F.silu(x1) * x2) - # Atom decoder - self.atom_decoder = ESMFold2AtomDecoder( - d_atom=c_atom, - d_token=c_token, - n_blocks=atom_num_blocks, - n_heads=atom_num_heads, - swa_window_size=swa_window_size, - expansion_ratio=2, - spatial_rope_base_frequency=spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=n_uid_rope_pairs, - uid_rope_base_frequency=uid_rope_base_frequency, - ) - self.s_to_token = nn.Linear(c_token, c_token, bias=False) - nn.init.zeros_(self.s_to_token.weight) +# =========================================================================== +# SWA3DRoPEAttention +# =========================================================================== - # Token transformer (DiffusionTransformer with pair bias) - self.token_transformer = DiffusionTransformer( - d_model=c_token, - d_pair=c_z, - num_heads=token_num_heads, - num_blocks=token_num_blocks, - d_cond=c_token, - transition_multiplier=transition_multiplier, - use_conditioning=True, - ) - self.s_step_norm = nn.LayerNorm(c_token) - self.token_norm = nn.LayerNorm(c_token) +class SWA3DRoPEAttention(nn.Module): + """Sliding window attention with 3D RoPE. Has Wqkv, gate_proj, out_proj.""" - def set_kernel_backend(self, backend: str | None) -> None: - self.token_transformer.set_kernel_backend(backend) + def __init__(self, d_model: int, n_heads: int, half_window: int = 64) -> None: + super().__init__() + self.n_heads = n_heads + self.head_dim = d_model // n_heads + self.scale = self.head_dim**-0.5 + self.half_window = half_window - def forward( - self, - x_noisy: Tensor, - t_hat: Tensor, - ref_pos: Tensor, - ref_charge: Tensor, - ref_mask: Tensor, - ref_element: Tensor, - ref_atom_name_chars: Tensor, - ref_space_uid: Tensor, - tok_idx: Tensor, - s_inputs: Tensor, - s_trunk: Tensor | None, - z_trunk: Tensor, - relative_position_encoding: Tensor, - asym_id: Tensor, - residue_index: Tensor, - entity_id: Tensor, - token_index: Tensor, - sym_id: Tensor, - sigma_data: float | None = None, - token_attention_mask: Tensor | None = None, - num_diffusion_samples: int = 1, - return_token_repr: bool = False, - return_atom_repr: bool = False, - inference_cache: dict[str, Tensor] | None = None, - ) -> dict[str, Tensor | None]: - bsz = x_noisy.shape[0] - sigma = self.sigma_data if sigma_data is None else float(sigma_data) - t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape( - -1 - ) - if t.numel() == 1: - t = t.expand(bsz) + self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.gate_proj = nn.Linear(d_model, d_model, bias=False) - # Step 1: conditioning (pair z is cached across diffusion steps) - s, z = self.conditioning( - t_hat=t, - s_inputs=s_inputs, - s_trunk=s_trunk, - z_trunk=z_trunk, - relative_position_encoding=relative_position_encoding, - sigma_data=sigma, - num_diffusion_samples=num_diffusion_samples, - inference_cache=inference_cache, - ) + def forward(self, x: Tensor, attention_params: tuple) -> Tensor: + batch_size, n_atoms = x.shape[:2] + cos, sin = attention_params[0], attention_params[1] - # Step 2: normalize noisy coords - denom = torch.sqrt(t * t + sigma * sigma) - r_noisy = x_noisy / denom[:, None, None] + x_input = x + qkv = self.Wqkv(x) + qkv = qkv.view(batch_size, n_atoms, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4) + q, k, v = qkv.unbind(0) + q, k = qk_norm(q), qk_norm(k) - # Step 3: atom encoder - a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder( - ref_pos=ref_pos, - atom_attention_mask=ref_mask, - ref_space_uid=ref_space_uid, - ref_charge=ref_charge, - ref_element=ref_element, - ref_atom_name_chars=ref_atom_name_chars, - atom_to_token=tok_idx, - r_l=r_noisy, - s_i=s_trunk, - num_diffusion_samples=num_diffusion_samples, - return_intermediates=return_atom_repr, - inference_cache=inference_cache, - ) + q = apply_rotary_emb_3d(q, cos, sin) + k = apply_rotary_emb_3d(k, cos, sin) - # Step 4: add conditioned s - a = a + self.s_to_token(self.s_step_norm(s)) + input_dtype = q.dtype + if q.dtype not in (torch.float16, torch.bfloat16): + q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16() - # Step 5: token transformer - a, _ = self.token_transformer( - a, - s, - z, - beta=0.0, - attention_mask=token_attention_mask, - num_diffusion_samples=num_diffusion_samples, + # ESMFold2 does not advertise FlashAttention. Keep this atom path on + # PyTorch. Models that advertise FlashAttention dispatch through the + # precompiled Hugging Face kernels interface in fastplms.attention. + q_t = q.transpose(1, 2) + k_t = k.transpose(1, 2) + v_t = v.transpose(1, 2) + attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale + attn = F.softmax(attn, dim=-1) + out = torch.matmul(attn, v_t).transpose(1, 2) + + out = out.to(input_dtype).reshape( # type: ignore[union-attr] + batch_size, n_atoms, -1 ) + out = out * torch.sigmoid(self.gate_proj(x_input)) + return self.out_proj(out) - # Step 6: token norm - a = self.token_norm(a) - # Step 7: atom decoder - r_update, dec_intermediates = self.atom_decoder( - a_i=a, - q_l=q_skip, - c_l=c_skip, - p_lm=p_skip, - atom_to_token=tok_idx, - atom_attention_mask=ref_mask, - num_diffusion_samples=num_diffusion_samples, - return_intermediates=return_atom_repr, - ) +# =========================================================================== +# SWAAtomBlock, SWAAtomTransformer +# =========================================================================== - # Step 8: compute denoised output - sigma2 = sigma * sigma - t2 = t * t - out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy - out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update - # Collect atom intermediates from encoder + decoder - atom_intermediates: Tensor | None = None - if return_atom_repr: - all_ints = enc_intermediates + dec_intermediates - if all_ints: - atom_intermediates = torch.stack(all_ints, dim=2) +def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift - return { - "x_denoised": out, - "token_repr": a if return_token_repr else None, - "atom_intermediates": atom_intermediates, - } +def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor: + return x + gate * y -# =========================================================================== -# DiffusionStructureHead -# =========================================================================== +class SWAAtomBlock(nn.Module): + """adaLN-Zero + SWA attention + SwiGLU FFN. -class DiffusionStructureHead(nn.Module): - """Wrapper around DiffusionModule with diffusion sampling.""" + Creates adaln_modulation = Sequential(SiLU(), Linear) -> keys like adaln_modulation.1.weight + """ - def __init__(self, config: ESMFold2Config) -> None: + def __init__( + self, + d_atom: int, + n_heads: int, + half_window: int = 64, + expansion_ratio: int = 2, + use_compile_fusions: bool = False, + ) -> None: super().__init__() - dm = config.structure_head.diffusion_module - swa_cfg = config.inputs.atom_encoder - sh = config.structure_head + self.attn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) - self.diffusion_module = DiffusionModule( - c_atom=dm.c_atom, - c_token=dm.c_token, - c_z=dm.c_z, - c_s_inputs=dm.c_s_inputs, - sigma_data=dm.sigma_data, - fourier_dim=dm.fourier_dim, - atom_num_blocks=dm.atom_num_blocks, - atom_num_heads=dm.atom_num_heads, - token_num_blocks=dm.token_num_blocks, - token_num_heads=dm.token_num_heads, - transition_multiplier=dm.transition_multiplier, - swa_window_size=swa_cfg.swa_window_size, - spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, - uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False) + nn.init.zeros_(adaln_linear.weight) + self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear) + + self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window) + self.ffn = SwiGLUFFN(d_atom, expansion_ratio) + + self._rms_adaln = torch.compile(_rms_adaln_raw) if use_compile_fusions else _rms_adaln_raw + self._gated_residual = ( + torch.compile(_gated_residual_raw) if use_compile_fusions else _gated_residual_raw ) - # Sampling hyperparameters - self.sigma_data = dm.sigma_data - self.gamma_0 = sh.gamma_0 - self.gamma_min = sh.gamma_min - self.noise_scale = sh.noise_scale - self.step_scale = sh.step_scale - self.inference_s_max = sh.inference_s_max - self.inference_s_min = sh.inference_s_min - self.inference_p = sh.inference_p - self.inference_num_steps = sh.inference_num_steps + def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor: + mod = self.adaln_modulation(c_l) + if mod.dim() == 2: + mod = mod.unsqueeze(1) + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1) - def set_kernel_backend(self, backend: str | None) -> None: - self.diffusion_module.set_kernel_backend(backend) + attn_input = self._rms_adaln(x, scale_a, shift_a) + attn_out = self.attn(attn_input, attention_params) + x = self._gated_residual(x, gate_a, attn_out) - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ + ffn_input = self._rms_adaln(x, scale_f, shift_f) + ffn_out = self.ffn(ffn_input) + x = self._gated_residual(x, gate_f, ffn_out) + return x - def inference_noise_schedule( - self, num_steps: int | None = None, device: torch.device | None = None - ) -> Tensor: - """Karras power-law noise schedule.""" - steps = self.inference_num_steps if num_steps is None else int(num_steps) - if steps == 1: - return torch.tensor( - [self.inference_s_max * self.sigma_data, 0.0], - device=device, - dtype=torch.float32, - ) - p = float(self.inference_p) - inv_p = 1.0 / p - k = torch.arange(steps, device=device, dtype=torch.float32) - base = self.inference_s_max**inv_p + (k / (steps - 1)) * ( - self.inference_s_min**inv_p - self.inference_s_max**inv_p - ) - schedule = self.sigma_data * base.pow(p) - return F.pad(schedule, (0, 1), value=0.0) - - @staticmethod - def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor: - q = torch.randn((n, 4), dtype=dtype, device=device) - scale = torch.sqrt((q * q).sum(dim=1)) - signs = torch.where(q[:, 0] < 0, -scale, scale) - q = q / signs[:, None] - r, i, j, k = torch.unbind(q, dim=-1) - two_s = 2.0 / (q * q).sum(dim=-1) - return torch.stack( - ( - 1 - two_s * (j * j + k * k), - two_s * (i * j - k * r), - two_s * (i * k + j * r), - two_s * (i * j + k * r), - 1 - two_s * (i * i + k * k), - two_s * (j * k - i * r), - two_s * (i * k - j * r), - two_s * (j * k + i * r), - 1 - two_s * (i * i + j * j), - ), - dim=-1, - ).reshape(n, 3, 3) - - def _center_random_augmentation( - self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None - ) -> tuple[Tensor, Tensor | None]: - """Algorithm 19: center + random rotation + translation.""" - bsz = x.shape[0] - mask = atom_mask.unsqueeze(-1) # [B, A, 1] - denom = mask.sum(dim=1, keepdim=True).clamp(min=1) - mean = (x * mask).sum(dim=1, keepdim=True) / denom - x = x - mean - if second_coords is not None: - second_coords = second_coords - mean - - r = self._random_rotations(bsz, x.dtype, x.device) - x = torch.einsum("bmd,bds->bms", x, r) - if second_coords is not None: - second_coords = torch.einsum("bmd,bds->bms", second_coords, r) - - t = torch.randn_like(x[:, 0:1, :]) - x = x + t - if second_coords is not None: - second_coords = second_coords + t - return x, second_coords - @staticmethod - def _weighted_rigid_align( - x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor - ) -> Tensor: - """Kabsch alignment: align x to x_gt with weights w.""" - w = (mask * w).unsqueeze(-1) # [B, N, 1] - denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8) - mu = (x * w).sum(dim=-2, keepdim=True) / denom - mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom - x_c = x - mu - xgt_c = x_gt - mu_gt - H = torch.einsum("bni,bnj->bij", w * xgt_c, x_c) - H32 = H.float() - U, _, Vh = torch.linalg.svd(H32, driver="gesvd" if H32.is_cuda else None) - det = torch.linalg.det(U @ Vh) - ones = torch.ones_like(det) - R = (U @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ Vh).to( - H.dtype - ) - return x_c @ R.transpose(-1, -2) + mu_gt - - # ------------------------------------------------------------------ - # Sampling - # ------------------------------------------------------------------ +class SWAAtomTransformer(nn.Module): + """Stack of SWAAtomBlocks.""" - @torch.inference_mode() - def sample( + def __init__( self, - z_trunk: Tensor, - s_inputs: Tensor, - s_trunk: Tensor | None, - relative_position_encoding: Tensor, - ref_pos: Tensor, - ref_charge: Tensor, - ref_mask: Tensor, - ref_element: Tensor, - ref_atom_name_chars: Tensor, - ref_space_uid: Tensor, - tok_idx: Tensor, - asym_id: Tensor, - residue_index: Tensor, - entity_id: Tensor, - token_index: Tensor, - sym_id: Tensor, - token_attention_mask: Tensor | None = None, - num_diffusion_samples: int = 1, - num_sampling_steps: int | None = None, - max_inference_sigma: float | None = 256.0, - noise_scale: float | None = None, - step_scale: float | None = None, - return_atom_repr: bool = False, - use_inference_cache: bool = True, - denoising_early_exit_rmsd: float | None = None, - ) -> dict[str, Tensor | None]: - """Diffusion sampling (Algorithm 18). - - ``num_sampling_steps`` is the number of denoising steps actually run. - When ``max_inference_sigma`` is set, the Karras schedule built with - ``num_sampling_steps`` entries would lose its high-σ tail to the cap, - so we inflate the underlying schedule length here to land back at the - requested step count post-truncation. - """ - n_atoms = tok_idx.shape[1] - device = s_inputs.device - target_batch = s_inputs.shape[0] * num_diffusion_samples - - inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None - - steps = ( - self.inference_num_steps - if num_sampling_steps is None - else int(num_sampling_steps) - ) - - schedule = self.inference_noise_schedule(steps, device) - if max_inference_sigma is not None: - schedule = schedule[schedule <= float(max_inference_sigma)] - schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma)) - - lam = self.noise_scale if noise_scale is None else float(noise_scale) - eta = self.step_scale if step_scale is None else float(step_scale) + d_atom: int = 128, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.swa_window_size = swa_window_size + self.head_dim = d_atom // n_heads + self.spatial_rope_base_frequency = spatial_rope_base_frequency + self.n_spatial_rope_pairs_per_axis = n_spatial_rope_pairs_per_axis + self.n_uid_rope_pairs = n_uid_rope_pairs + self.uid_rope_base_frequency = uid_rope_base_frequency - x = schedule[0] * torch.randn( - target_batch, n_atoms, 3, device=device, dtype=torch.float32 + self.blocks = nn.ModuleList( + [ + SWAAtomBlock( + d_atom=d_atom, + n_heads=n_heads, + half_window=swa_window_size // 2, + expansion_ratio=expansion_ratio, + ) + for _ in range(n_blocks) + ] ) - atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float() - gammas = torch.where( - schedule > self.gamma_min, - torch.full_like(schedule, self.gamma_0), - torch.zeros_like(schedule), + def _build_3d_rope(self, ref_pos: Tensor, ref_space_uid: Tensor) -> tuple[Tensor, Tensor]: + return build_3d_rope( + ref_pos=ref_pos, + ref_space_uid=ref_space_uid, + head_dim=self.head_dim, + n_spatial_per_axis=self.n_spatial_rope_pairs_per_axis, + n_uid_pairs=self.n_uid_rope_pairs, + spatial_base_freq=self.spatial_rope_base_frequency, + uid_base_freq=self.uid_rope_base_frequency, ) - x_denoised_prev: Tensor | None = None - token_repr: Tensor | None = None - diff_atom_intermediates: Tensor | None = None - - step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:])) - num_steps = len(step_pairs) - - for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs): - x, x_denoised_prev = self._center_random_augmentation( - x, atom_mask, second_coords=x_denoised_prev - ) - - sigma_tm_val = float(sigma_tm.item()) - t_hat_val = sigma_tm_val * (1.0 + float(gamma.item())) - eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5 - x_noisy = x + eps_std * torch.randn_like(x) - - is_last_step = step_idx == num_steps - 1 - request_atom_repr = return_atom_repr and ( - is_last_step or denoising_early_exit_rmsd is not None - ) - - dm_out = self.diffusion_module( - x_noisy=x_noisy, - t_hat=torch.full( - (target_batch,), t_hat_val, device=device, dtype=torch.float32 - ), - ref_pos=ref_pos, - ref_charge=ref_charge, - ref_mask=ref_mask, - ref_element=ref_element, - ref_atom_name_chars=ref_atom_name_chars, - ref_space_uid=ref_space_uid, - tok_idx=tok_idx, - s_inputs=s_inputs, - s_trunk=s_trunk, - z_trunk=z_trunk, - relative_position_encoding=relative_position_encoding, - asym_id=asym_id, - residue_index=residue_index, - entity_id=entity_id, - token_index=token_index, - sym_id=sym_id, - token_attention_mask=token_attention_mask, - num_diffusion_samples=num_diffusion_samples, - return_token_repr=True, - return_atom_repr=request_atom_repr, - inference_cache=inference_cache, - ) - - x_denoised = dm_out["x_denoised"] - token_repr = dm_out["token_repr"] - if request_atom_repr: - diff_atom_intermediates = dm_out.get("atom_intermediates") - - # Reverse diffusion alignment (Kabsch) - with torch.autocast(device_type="cuda", enabled=False): - x_noisy = self._weighted_rigid_align( - x_noisy.float(), x_denoised.float(), atom_mask, atom_mask - ) - x_noisy = x_noisy.to(dtype=x_denoised.dtype) - - # ODE/SDE step - sigma_t_val = float(sigma_t.item()) - denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val - x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma - - # Denoising early-exit: stop when consecutive predictions converge - if ( - denoising_early_exit_rmsd is not None - and x_denoised_prev is not None - and step_idx >= 1 - ): - with torch.autocast(device_type="cuda", enabled=False): - aligned = self._weighted_rigid_align( - x_denoised_prev.float(), - x_denoised.float(), - atom_mask, - atom_mask, - ) - diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1) - per_sample_rmsd = ( - diff.pow(2).sum(dim=(-1, -2)) / atom_mask.sum(dim=-1).clamp(min=1) - ).sqrt() - if per_sample_rmsd.max().item() < denoising_early_exit_rmsd: - x = x_denoised - x_denoised_prev = x_denoised - break - - x_denoised_prev = x_denoised - - result: dict[str, Tensor | None] = { - "sample_atom_coords": x, - "diff_token_repr": token_repr, - } - if return_atom_repr: - result["diff_atom_intermediates"] = diff_atom_intermediates - return result + def forward( + self, + q_l: Tensor, + c_l: Tensor, + attention_params: tuple, + return_intermediates: bool = False, + ) -> Tensor | tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + for block in self.blocks: + q_l = block(q_l, c_l, attention_params) + if return_intermediates: + intermediates.append(q_l) + if return_intermediates: + return q_l, intermediates + return q_l # =========================================================================== -# RowAttentionPooling +# ESMFold2AtomEncoder (for both inputs_embedder and diffusion_module) # =========================================================================== -class RowAttentionPooling(nn.Module): - """Row-wise attention pooling: attn_proj, out_proj.""" - - def __init__(self, d_pair: int, d_single: int) -> None: - super().__init__() - self.attn_proj = nn.Linear(d_pair, 1, bias=False) - self.out_proj = nn.Linear(d_pair, d_single, bias=False) - - def forward(self, z: Tensor, mask: Tensor) -> Tensor: - scores = self.attn_proj(z).squeeze(-1) - mask_bias = torch.where( - mask[:, None, :].bool(), - torch.zeros_like(scores), - torch.full_like(scores, -1e9), - ) - scores = scores + mask_bias - weights = F.softmax(scores, dim=-1) - pooled = torch.einsum("bnm,bnmd->bnd", weights, z) - return self.out_proj(pooled) - +class ESMFold2AtomEncoder(nn.Module): + """Encode atom inputs with normalization and sliding-window attention. -# =========================================================================== -# InputsEmbedder -# =========================================================================== + Args: + d_atom: atom hidden dim + d_token: token dim for atom_to_token aggregation + n_blocks, n_heads, swa_window_size, expansion_ratio: transformer params + structure_prediction: if True, creates coords_linear and uses full d_token + spatial_rope_base_frequency, n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs, uid_rope_base_frequency: 3D RoPE config + """ + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + structure_prediction: bool = True, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.d_atom = d_atom + self.d_token = d_token + self.structure_prediction = structure_prediction -class InputsEmbedder(nn.Module): - """Embeds input features including atom-level encoding via SWA attention.""" + self.atom_linear = nn.Linear(ATOM_FEATURE_DIM, d_atom, bias=False) + self.atom_norm = nn.LayerNorm(d_atom) - def __init__(self, config: ESMFold2Config) -> None: - super().__init__() - swa_cfg = config.inputs.atom_encoder + if structure_prediction: + self.coords_linear = nn.Linear(6, d_atom, bias=False) - self.atom_attention_encoder = ESMFold2AtomEncoder( - d_atom=swa_cfg.d_atom, - d_token=swa_cfg.d_token, - n_blocks=swa_cfg.n_blocks, - n_heads=swa_cfg.n_heads, - swa_window_size=swa_cfg.swa_window_size, - expansion_ratio=swa_cfg.expansion_ratio, - structure_prediction=False, # no coords_linear - spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, - n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, - n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, - uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, ) + # Output aggregation: d_token for structure prediction, d_token//2 for inputs + out_dim = d_token if structure_prediction else d_token // 2 + self.atom_to_token_linear = nn.Linear(d_atom, out_dim, bias=False) + def forward( self, - aatype: Tensor, - profile: Tensor, - deletion_mean: Tensor, ref_pos: Tensor, atom_attention_mask: Tensor, ref_space_uid: Tensor, @@ -1963,828 +1643,1057 @@ def forward( ref_element: Tensor, ref_atom_name_chars: Tensor, atom_to_token: Tensor, - ) -> Tensor: - """Embed inputs into per-token features. + r_l: Tensor | None = None, + pred_r1: Tensor | None = None, + s_i: Tensor | None = None, + z_ij: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + inference_cache: dict | None = None, + ) -> tuple[Tensor, Tensor, Tensor, tuple, list[Tensor]]: + """Returns (a, q, c, attention_params, intermediates). - Returns: - [B, L, d_inputs] concatenation of atom encoding, aatype, profile, - and deletion_mean. + ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE, + attention indices, n_tokens) across diffusion steps. """ - a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder( - ref_pos=ref_pos, - atom_attention_mask=atom_attention_mask, - ref_space_uid=ref_space_uid, - ref_charge=ref_charge, - ref_element=ref_element, - ref_atom_name_chars=ref_atom_name_chars, - atom_to_token=atom_to_token, - ) - return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1) - - -# =========================================================================== -# ResIdxAsymIdSymIdEntityIdEncoding (trunk relative position) -# =========================================================================== - - -class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module): - """embed.weight [d_pair, n_features] where n_features = 2*(2*r_bins+2) + 1 + (2*c_bins+2). - - For default r_bins=32, c_bins=2: 2*66 + 1 + 6 = 139. - """ - - def __init__( - self, - n_relative_residx_bins: int = 32, - n_relative_chain_bins: int = 2, - d_pair: int = 256, - ) -> None: - super().__init__() - self.n_relative_residx_bins = n_relative_residx_bins - self.n_relative_chain_bins = n_relative_chain_bins - self.d_pair = d_pair + batch_size, n_atoms = ref_pos.shape[:2] - n_feats_residue = 2 * n_relative_residx_bins + 2 - n_feats_token = 2 * n_relative_residx_bins + 2 - n_feats_chain = 2 * n_relative_chain_bins + 2 - n_feats_same_entity = 1 - total_feats = ( - n_feats_residue + n_feats_token + n_feats_chain + n_feats_same_entity - ) - self.embed = nn.Linear(total_feats, d_pair, bias=False) + layer_cache = None + if inference_cache is not None: + layer_cache = inference_cache.setdefault("atomencoder", {}) - def forward( - self, - residue_index: Tensor, - asym_id: Tensor, - sym_id: Tensor, - entity_id: Tensor, - token_index: Tensor, - ) -> Tensor: - bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1) - bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1) - bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1) + if layer_cache is None or len(layer_cache) == 0: + atom_feats = torch.cat( + [ + ref_pos, + ref_charge.unsqueeze(-1), + atom_attention_mask.unsqueeze(-1), + ref_element, + ref_atom_name_chars.reshape(batch_size, n_atoms, MAX_CHARS * CHAR_VOCAB_SIZE), + ], + dim=-1, + ) + c_base = self.atom_norm(self.atom_linear(atom_feats)) + cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid) + cos = cos.repeat_interleave(num_diffusion_samples, 0) + sin = sin.repeat_interleave(num_diffusion_samples, 0) + mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0) + seqlens = mask_exp.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten() + max_seqlen = int(seqlens.max().item()) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + attention_params = (cos, sin, indices, cu_seqlens, max_seqlen) + n_tokens = int(atom_to_token.max().item()) + 1 + if layer_cache is not None: + layer_cache["c_base"] = c_base + layer_cache["attention_params"] = attention_params + layer_cache["mask_exp"] = mask_exp + layer_cache["n_tokens"] = n_tokens + layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave( + num_diffusion_samples, 0 + ) + else: + c_base = layer_cache["c_base"] + attention_params = layer_cache["attention_params"] + mask_exp = layer_cache["mask_exp"] + n_tokens = layer_cache["n_tokens"] - dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1) - dij_residue = torch.clip( - dij_residue + self.n_relative_residx_bins, - 0, - 2 * self.n_relative_residx_bins, - ) - dij_residue = torch.where( - bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1 - ) - aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2) + c = c_base - dij_token = torch.clip( - token_index.unsqueeze(2) - - token_index.unsqueeze(1) - + self.n_relative_residx_bins, - 0, - 2 * self.n_relative_residx_bins, - ) - dij_token = torch.where( - bij_same_chain & bij_same_residue, - dij_token, - 2 * self.n_relative_residx_bins + 1, - ) - aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2) + q = c - dij_chain = torch.clip( - sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins, - 0, - 2 * self.n_relative_chain_bins, - ) - dij_chain = torch.where( - bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain - ) - aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2) + if self.structure_prediction and r_l is not None: + q = q.repeat_interleave(num_diffusion_samples, 0) + if pred_r1 is None: + pred_r1 = torch.zeros_like(r_l) + r_input = torch.cat([r_l, pred_r1], dim=-1) + r_to_q = self.coords_linear(r_input) + q = q + r_to_q - feats = torch.cat( - [ - aij_rel_pos.float(), - aij_rel_token.float(), - bij_same_entity.float().unsqueeze(-1), - aij_rel_chain.float(), - ], - dim=-1, + c = c.repeat_interleave(num_diffusion_samples, 0) + + result = self.atom_transformer( + q_l=q, + c_l=c, + attention_params=attention_params, + return_intermediates=return_intermediates, ) + if return_intermediates: + q, intermediates = result + else: + q = result + intermediates = [] - return self.embed(feats) + q_to_a = F.relu(self.atom_to_token_linear(q)) + if layer_cache is not None and "atom_to_token_exp" in layer_cache: + atom_to_token_exp = layer_cache["atom_to_token_exp"] + else: + atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) + a = scatter_atom_to_token(q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool()) + + return a, q, c, attention_params, intermediates # =========================================================================== -# SingleToPair (for LanguageModelShim) +# ESMFold2AtomDecoder # =========================================================================== -class SingleToPair(nn.Module): - """downproject, output_mlp (Sequential of Linear, GELU, Linear).""" +class ESMFold2AtomDecoder(nn.Module): + """SWA atom decoder with token_to_atom_linear, atom_transformer, norm, output_linear.""" - def __init__(self, input_dim: int, downproject_dim: int, output_dim: int) -> None: + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: super().__init__() - self.downproject = nn.Linear(input_dim, downproject_dim) - self.output_mlp = nn.Sequential( - nn.Linear(2 * downproject_dim, output_dim), - nn.GELU(), - nn.Linear(output_dim, output_dim), + self.token_to_atom_linear = nn.Linear(d_token, d_atom, bias=False) + + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, ) - def forward(self, x: Tensor) -> Tensor: - x = self.downproject(x) - x = torch.cat( - [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))], - dim=3, + self.norm = nn.LayerNorm(d_atom) + self.output_linear = nn.Linear(d_atom, XYZ_DIMS, bias=False) + + def forward( + self, + a_i: Tensor, + q_l: Tensor, + c_l: Tensor, + p_lm: tuple, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + """Returns (r_update, intermediates).""" + atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) + a_to_q = self.token_to_atom_linear(a_i) + a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp) + q_l = q_l + a_to_q + + result = self.atom_transformer( + q_l=q_l, + c_l=c_l, + attention_params=p_lm, + return_intermediates=return_intermediates, ) - return self.output_mlp(x) + if return_intermediates: + q_l, intermediates = result + else: + q_l = result + intermediates = [] + + r_l = self.output_linear(self.norm(q_l)) + return r_l, intermediates # =========================================================================== -# LanguageModelShim +# AttentionPairBias (DiffusionTransformer attention block) # =========================================================================== -class LanguageModelShim(nn.Module): - """Shim holding the trainable projection weights for LM integration. - - Contains: - - base_z_combine: nn.Parameter [num_layers+1] - - base_z_linear: Sequential(LayerNorm(d_model), Linear(d_model, d_z, bias=False)) - - base_z_mlp: Sequential(SingleToPair(d_z, d_z, d_z), LayerNorm(d_z)) - """ +class AttentionPairBias(nn.Module): + """Gated multi-head attention with pair bias conditioning.""" def __init__( - self, d_z: int = 256, d_model: int = 2560, num_layers: int = 80 + self, + d_model: int, + d_pair: int, + num_heads: int, + d_cond: int | None = None, + use_conditioning: bool = True, ) -> None: super().__init__() + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + d_cond = d_cond or d_model - self.base_z_mlp = nn.Sequential(SingleToPair(d_z, d_z, d_z), nn.LayerNorm(d_z)) - self.base_z_linear = nn.Sequential( - nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False) + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.out_gate = nn.Linear(d_cond, d_model, bias=True) + # adaln init: weight=0, bias=-2 + nn.init.zeros_(self.out_gate.weight) + nn.init.constant_(self.out_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False) + self.g_proj = nn.Linear(d_model, d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + + if d_pair > 0: + self.pair_norm = nn.LayerNorm(d_pair, eps=1e-5) + self.pair_bias_proj = nn.Linear(d_pair, num_heads, bias=False) + + self._kernel_backend: str | None = None + + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + self._kernel_backend = backend + + def _is_zero_beta(self, beta: Tensor | float) -> bool: + if isinstance(beta, (int, float)): + return beta == 0.0 + return bool((beta == 0).all()) + + def _can_use_fused_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool: + return ( + _fused_active(self, z) + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + and hasattr(self, "pair_norm") ) - self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1)) - def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor: - """Project pre-computed ESMC hidden states to pair representation. + def _can_use_cueq_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool: + return ( + _cueq_active(self) + and n_queries > 750 + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + ) - Args: - hidden_states: [B, L, num_layers+1, d_model] from ESMC 6B. - lm_dropout: Dropout probability applied to the pair - representation after ``base_z_mlp``. + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + ) -> Tensor: + bsz, n_queries, d_model = a.shape - Returns: - [B, L, L, d_pair] pair representation. - """ - lm_z = self.base_z_linear(hidden_states) # [B, L, 81, d_z] - weights = self.base_z_combine.softmax(0) # [81] - lm_z = (weights @ lm_z).squeeze(-2) # [B, L, d_z] - lm_z = self.base_z_mlp(lm_z) # [B, L, L, d_z] - if lm_dropout > 0: - lm_z = F.dropout(lm_z, p=lm_dropout, training=True) - return lm_z + x = self.adaln(a, s) if s is not None else self.pre_norm(a) + + n_keys = x.shape[1] + q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim) + kv = self.kv_proj(x) + k, v = kv.chunk(2, dim=-1) + k = k.view(bsz, n_keys, self.num_heads, self.head_dim) + v = v.view(bsz, n_keys, self.num_heads, self.head_dim) + + # Expand z for num_diffusion_samples + if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1: + z = z.repeat_interleave(num_diffusion_samples, dim=0) + if ( + attention_mask is not None + and attention_mask.shape[0] != bsz + and num_diffusion_samples > 1 + ): + attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0) + + if self._can_use_fused_pair_bias(z, n_queries, beta): + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + pair_norm_w = self.pair_norm.weight + pair_norm_b = ( + self.pair_norm.bias + if self.pair_norm.bias is not None + else torch.zeros_like(pair_norm_w) + ) + z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16) + bias = _fused_pair_bias( # type: ignore[misc] + z_bf, + kernel_mask, + self.pair_bias_proj.weight, + num_heads=self.num_heads, + pair_norm_w=pair_norm_w, + pair_norm_b=pair_norm_b, + ) # A has shape (b, h, q, k). + q_bhqd = q.transpose(1, 2) + k_bhqd = k.transpose(1, 2) + v_bhqd = v.transpose(1, 2) + attn_out = F.scaled_dot_product_attention( + q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype) + ) + g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim) + ctx = g * attn_out.transpose(1, 2) + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out + + if self._can_use_cueq_pair_bias(z, n_queries, beta): + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + out, _ = _cue_attn_pair_bias( # type: ignore[misc] + s=x, + q=q.transpose(1, 2), + k=k.transpose(1, 2), + v=v.transpose(1, 2), + z=z, + mask=kernel_mask, + num_heads=self.num_heads, + w_proj_z=self.pair_bias_proj.weight, + w_proj_g=self.g_proj.weight, + w_proj_o=self.out_proj.weight, + w_ln_z=self.pair_norm.weight, + b_ln_z=self.pair_norm.bias, + return_z_proj=False, + is_cached_z_proj=False, + ) + else: + # Standard attention with pair bias + g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim) + logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale -# =========================================================================== -# Reproducibility helper (mirrors evolutionaryscale.utils.reproducibility) -# =========================================================================== + pair_bias = self.pair_bias_proj(self.pair_norm(z)) if z.dim() == 4 else z.unsqueeze(-1) + logits = logits + pair_bias.to(dtype=logits.dtype) + if attention_mask is not None: + min_val = torch.finfo(logits.dtype).min + mask_bias = torch.where(attention_mask.bool()[:, None, :, None], 0.0, min_val) + logits = logits + mask_bias.to(dtype=logits.dtype) -@contextmanager -def _seed_context(seed: int | None, *, cuda: bool = True): - """Temporarily seed Python, NumPy, and PyTorch RNGs.""" - if seed is None: - yield - return - py_state = random.getstate() - np_state = np.random.get_state() - torch_state = torch.get_rng_state() - cuda_states = ( - torch.cuda.get_rng_state_all() if cuda and torch.cuda.is_available() else None - ) - seed = int(seed) % (2**32) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if cuda and torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - try: - yield - finally: - random.setstate(py_state) - np.random.set_state(np_state) - torch.set_rng_state(torch_state) - if cuda_states is not None: - torch.cuda.set_rng_state_all(cuda_states) + attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype) + ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v) + ctx = g * ctx + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out # =========================================================================== -# ESMFold2ExperimentalModel — the top-level PreTrainedModel +# ConditionedTransitionBlock # =========================================================================== -def compute_lm_hidden_states( - esmc: nn.Module, - input_ids: Tensor, - asym_id: Tensor, - residue_index: Tensor, - mol_type: Tensor, - token_mask: Tensor, - pad_to_multiple: int | None = None, - lm_mask_pct: float = 0.0, - mask_token_id: int = 32, -) -> Tensor: - """Run ESMC with BOS/EOS wrapping, return hidden states [B, L, N, D] with N=81 layers. +class ConditionedTransitionBlock(nn.Module): + """Conditioned SwiGLU transition with adaptive layer norm.""" - Atom-tokenized modified residues (HYP, MSE, ACE, NH2, ...) span multiple - structure tokens but share a single ``(asym_id, residue_index)`` key — - collapse them to one LM token per residue before running the LM (the LM - was trained on per-residue inputs, not per-atom), then scatter the - hidden states back to the per-token layout. - """ - B, L = input_ids.shape - device = input_ids.device - protein_mask = (mol_type == 0) & token_mask + def __init__( + self, + d_model: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + hidden = transition_multiplier * d_model - lm_input_list = [] - lm_lengths = [] - # Per-batch maps from (original protein-token index) to (LM input position). - expand_maps: list[Tensor] = [] - for b in range(B): - mask_b = protein_mask[b] - ids_b = input_ids[b][mask_b] - asym_b = asym_id[b][mask_b] - res_b = residue_index[b][mask_b] + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.output_gate = nn.Linear(d_cond, d_model, bias=True) + nn.init.zeros_(self.output_gate.weight) + nn.init.constant_(self.output_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) - # Collapse: keep first token per (asym_id, residue_index) key, in - # input order. ``inverse`` maps each original protein-token to its - # collapsed residue index. - keys = torch.stack((asym_b, res_b), dim=1) - unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True) - n_unique = unique_keys.size(0) - token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long) - first_pos = torch.full( - (n_unique,), keys.size(0), device=device, dtype=torch.long - ) - first_pos.scatter_reduce_( - 0, inverse, token_positions, reduce="amin", include_self=True - ) - ordered = torch.argsort(first_pos) - first_pos_ordered = first_pos[ordered] - ids_collapsed = ids_b[first_pos_ordered] - asym_collapsed = asym_b[first_pos_ordered] - remap = torch.empty_like(ordered) - remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long) - inverse_ordered = remap[inverse] + self.lin_swish = nn.Linear(d_model, 2 * hidden, bias=False) + self.lin_out = nn.Linear(hidden, d_model, bias=False) - chain_ids = asym_collapsed.unique(sorted=True) - # [BOS] chain1 [EOS BOS] chain2 ... [EOS] - parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)] - # Per-chain LM positions accumulate; track them for the expand map. - per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long) - cursor = 1 # position 0 is the leading BOS - for i, cid in enumerate(chain_ids): - in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0] - parts.append(ids_collapsed[in_chain]) - per_token_lm_pos[in_chain] = torch.arange( - cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long - ) - cursor += in_chain.shape[0] - if i < len(chain_ids) - 1: - parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype)) - cursor += 2 # EOS + BOS - parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype)) - lm_seq = torch.cat(parts) - lm_input_list.append(lm_seq) - lm_lengths.append(lm_seq.shape[0]) + def forward(self, a: Tensor, s: Tensor | None) -> Tensor: + x = self.adaln(a, s) if s is not None else self.pre_norm(a) - # Original protein-token position → LM input position. - prot_pos_b = mask_b.nonzero(as_tuple=True)[0] - expand_map = torch.full((L,), -1, device=device, dtype=torch.long) - expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered] - expand_maps.append(expand_map) + swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1) + b = F.silu(swish_a) * swish_b + out = self.lin_out(b) - # Pad to longest LM input; round to ``pad_to_multiple`` when fp8 is on - # (TE fp8 kernels assert prod(shape[:-1]) % 8 == 0). - max_len = max(lm_lengths) - if pad_to_multiple is not None and pad_to_multiple > 1: - max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple - lm_input_ids = torch.full( - (B, max_len), - 1, - device=device, - dtype=input_ids.dtype, # PAD=1 - ) - for b in range(B): - lm_input_ids[b, : lm_lengths[b]] = lm_input_list[b] + if s is not None: + out = torch.sigmoid(self.output_gate(s)) * out + return out - # sequence_id for chain-aware attention; PAD tokens get -1 (no attention). - sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1 # BOS=0 - sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1) # PAD=1 - if lm_mask_pct > 0.0: - special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2) - do_mask = ( - torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct - ) & ~special - lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id) +# =========================================================================== +# DiffusionTransformer (token transformer) +# =========================================================================== - with torch.inference_mode(): - esmc_out = esmc( - input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True + +class DiffusionTransformer(nn.Module): + """Diffusion denoising transformer with attention pair bias.""" + + def __init__( + self, + d_model: int, + d_pair: int, + num_heads: int, + num_blocks: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + + self.attn_blocks = nn.ModuleList( + [ + AttentionPairBias( + d_model=d_model, + d_pair=d_pair, + num_heads=num_heads, + d_cond=d_cond, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] + ) + self.transition_blocks = nn.ModuleList( + [ + ConditionedTransitionBlock( + d_model=d_model, + d_cond=d_cond, + transition_multiplier=transition_multiplier, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] ) - hs = esmc_out.hidden_states # [n_layers+1, B, max_len, D] - n_layers_plus_1, _, _, D = hs.shape - result = torch.zeros(B, L, n_layers_plus_1, D, device=device, dtype=hs.dtype) - for b in range(B): - mb = protein_mask[b] - em = expand_maps[b][mb] # [n_protein_tokens] LM positions - # hs[:, b, em, :] -> [n_layers+1, n_protein_tokens, D] - gathered = hs[:, b, em, :].permute(1, 0, 2) - result[b, mb.nonzero(as_tuple=True)[0]] = gathered + def set_kernel_backend(self, backend: str | None) -> None: + for attn in self.attn_blocks: + cast(AttentionPairBias, attn).set_kernel_backend(backend) - return result.detach() + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + x = a + for attn, transition in zip(self.attn_blocks, self.transition_blocks, strict=True): + x = x + attn( + x, + s, + z, + beta, + attention_mask=attention_mask, + num_diffusion_samples=num_diffusion_samples, + ) + x = x + transition(x, s) + if return_intermediates: + intermediates.append(x) + return x, intermediates # =========================================================================== -# TriangleMultiplicativeUpdate +# DiffusionConditioning # =========================================================================== -class TriangleMultiplicativeBlock(nn.Module): - """Triangle multiplicative update block with gated signal routing.""" - _FLOW_TO_EINSUM = {"outgoing": "bikd,bjkd->bijd", "incoming": "bkid,bkjd->bijd"} - _VALID_FLOWS = ("outgoing", "incoming") - def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None: +class DiffusionConditioning(nn.Module): + """Conditions pair and single representations on noise timestep.""" + + def __init__( + self, + c_z: int = 256, + c_s: int = 768, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + transition_multiplier: int = 2, + layer_norm_eps: float = 1e-5, + ) -> None: super().__init__() - if flow not in self._FLOW_TO_EINSUM: - raise ValueError( - f"Invalid flow={flow!r}. Expected one of {self._VALID_FLOWS}." - ) + self.sigma_data = float(sigma_data) + self.c_z = c_z + self.c_s = c_s + self.c_s_inputs = c_s_inputs - self.input_channels = input_channels - self.latent_channels = latent_channels - self.flow = flow - self._einsum_equation = self._FLOW_TO_EINSUM[flow] - self.norm_start = nn.LayerNorm(self.input_channels, eps=_EPS) - self.norm_mix = nn.LayerNorm(self.latent_channels, eps=_EPS) - self.proj_bundle = nn.Linear( - self.input_channels, 4 * self.latent_channels, bias=False - ) - self.proj_emit = nn.Linear( - self.latent_channels, self.input_channels, bias=False + self.z_input_norm = nn.LayerNorm(2 * c_z, eps=layer_norm_eps) + self.z_proj = nn.Linear(2 * c_z, c_z, bias=False) + self.z_transitions = nn.ModuleList( + [TransitionLayer(c_z, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)] ) - self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False) - - self._use_kernels: bool = False - # Default chunked for memory on long sequences; tests override with - # ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16 - # parity checks. - self._chunk_size: int | None = 64 - - def set_chunk_size(self, chunk_size: int | None) -> None: - self._chunk_size = chunk_size - def split_kernel_weights(self) -> tuple[Tensor, Tensor]: - return ( - self.proj_bundle.weight[: 2 * self.latent_channels, :], - self.proj_bundle.weight[2 * self.latent_channels :, :], + self.s_input_norm = nn.LayerNorm(c_s_inputs, eps=layer_norm_eps) + self.s_proj = nn.Linear(c_s_inputs, c_s, bias=False) + self.fourier = FourierEmbedding(fourier_dim) + self.noise_norm = nn.LayerNorm(fourier_dim, eps=layer_norm_eps) + self.noise_proj = nn.Linear(fourier_dim, c_s, bias=False) + self.s_transitions = nn.ModuleList( + [TransitionLayer(c_s, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)] ) - def _kernel_flow_direction(self) -> str: - return self.flow + def forward( + self, + t_hat: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + sigma_data: float | None = None, + num_diffusion_samples: int = 1, + inference_cache: dict[str, Tensor] | None = None, + ) -> tuple[Tensor, Tensor]: + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + base_batch = z_trunk.shape[0] + target_batch = base_batch * num_diffusion_samples + + # z conditioning (cached across diffusion steps: independent of t_hat) + if inference_cache is not None and "z" in inference_cache: + z = inference_cache["z"] + else: + z_rel = relative_position_encoding.to(dtype=torch.float32) + z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1) + z = self.z_proj(self.z_input_norm(z)) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for block in self.z_transitions: + z = z + block(z) + if inference_cache is not None: + inference_cache["z"] = z - def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor: - return torch.einsum(self._einsum_equation, left_stream, right_stream) + # s conditioning + s_inputs_eff = s_inputs + if s_inputs_eff.shape[0] != target_batch: + s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0) - def _triangular_contract_chunked( - self, left_stream: Tensor, right_stream: Tensor, chunk_size: int - ) -> Tensor: - """Compute the triangular einsum in chunks along the output i-dimension.""" - L = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2] - chunks = [] - for start in range(0, L, chunk_size): - end = min(start + chunk_size, L) - if self.flow == "outgoing": - chunk = torch.einsum( - self._einsum_equation, left_stream[:, start:end], right_stream - ) - else: - chunk = torch.einsum( - self._einsum_equation, left_stream[:, :, start:end], right_stream - ) - chunks.append(chunk) - return torch.cat(chunks, dim=1) + s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32))) - def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor: - if visibility is None: - visibility = pair_grid.new_ones(pair_grid.shape[:-1]) + # Noise embedding + t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1) + if t.numel() == 1: + t = t.expand(target_batch) + elif t.shape[0] != target_batch: + t = t.repeat_interleave(num_diffusion_samples, 0) + t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20)) + n = self.fourier(t_noise) + n = self.noise_proj(self.noise_norm(n)) + s = s + n.unsqueeze(1) - if self._use_kernels: - p_in_weight, g_in_weight = self.split_kernel_weights() + for block in self.s_transitions: + s = s + block(s) - try: - return _cue_tri_mul( # type: ignore[misc] - pair_grid, - direction=self._kernel_flow_direction(), - mask=visibility, - norm_in_weight=self.norm_start.weight, - norm_in_bias=self.norm_start.bias, - p_in_weight=p_in_weight, - g_in_weight=g_in_weight, - norm_out_weight=self.norm_mix.weight, - norm_out_bias=self.norm_mix.bias, - p_out_weight=self.proj_emit.weight, - g_out_weight=self.proj_gate.weight, - eps=_EPS, - ) - except Exception as e: - import logging as _logging - - _logging.getLogger(__name__).warning( - "cuequivariance triangle_multiplicative_update kernel failed " - "(flow=%s, shape=%s, dtype=%s); falling back to chunked einsum. " - "Error: %s", - self.flow, - tuple(pair_grid.shape), - pair_grid.dtype, - e, - ) + return s, z - normalized_grid = self.norm_start(pair_grid) - bundled = self.proj_bundle(normalized_grid) - signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1) - routed = signal * torch.sigmoid(gate_logits) - routed = routed * visibility.unsqueeze(-1) - left_stream, right_stream = routed.float().chunk(2, dim=-1) - if self._chunk_size is not None: - contracted = self._triangular_contract_chunked( - left_stream, right_stream, self._chunk_size - ) - else: - contracted = self._triangular_contract(left_stream, right_stream) - mixed = self.proj_emit(self.norm_mix(contracted)) - output_gate = torch.sigmoid(self.proj_gate(normalized_grid)) - return mixed * output_gate +# =========================================================================== +# DiffusionModule +# =========================================================================== -class TriangleMultiplicativeUpdate(nn.Module): - """Thin wrapper exposing the triangular mixer with explicit orientation (v3).""" +class DiffusionModule(nn.Module): + """Diffusion denoising module for structure prediction.""" - def __init__(self, dim: int = 128, _outgoing: bool = True) -> None: + def __init__( + self, + c_atom: int = 128, + c_token: int = 768, + c_z: int = 256, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + atom_num_blocks: int = 3, + atom_num_heads: int = 4, + token_num_blocks: int = 12, + token_num_heads: int = 16, + transition_multiplier: int = 2, + swa_window_size: int = 128, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: super().__init__() - flow = "outgoing" if _outgoing else "incoming" - self._engine = TriangleMultiplicativeBlock( - input_channels=dim, latent_channels=dim, flow=flow - ) + self.sigma_data = float(sigma_data) - def set_kernel_backend(self, backend: str | None) -> None: - # Engine uses cueq when backend=="cuequivariance"; the "fused" backend - # routes through the parent PairUpdateBlock's fused path (bypassing this). - self._engine._use_kernels = backend == BACKEND_CUEQ - if backend == BACKEND_CUEQ and not CUE_AVAILABLE: - raise RuntimeError( - "backend='cuequivariance' but cuequivariance_torch is not installed." - ) + self.conditioning = DiffusionConditioning( + c_z=c_z, + c_s=c_token, # conditioning s output is c_token + c_s_inputs=c_s_inputs, + sigma_data=sigma_data, + fourier_dim=fourier_dim, + transition_multiplier=transition_multiplier, + ) - def set_chunk_size(self, chunk_size: int | None) -> None: - self._engine.set_chunk_size(chunk_size) + # Atom encoder (structure_prediction=True, with coords_linear) + self.atom_encoder = ESMFold2AtomEncoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + structure_prediction=True, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) - def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor: - return self._engine(z, visibility=mask) + # Atom decoder + self.atom_decoder = ESMFold2AtomDecoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + self.s_to_token = nn.Linear(c_token, c_token, bias=False) + nn.init.zeros_(self.s_to_token.weight) -# =========================================================================== -# FoldingTrunk: Transition, PairUpdateBlock, FoldingTrunk -# =========================================================================== + # Token transformer (DiffusionTransformer with pair bias) + self.token_transformer = DiffusionTransformer( + d_model=c_token, + d_pair=c_z, + num_heads=token_num_heads, + num_blocks=token_num_blocks, + d_cond=c_token, + transition_multiplier=transition_multiplier, + use_conditioning=True, + ) + self.s_step_norm = nn.LayerNorm(c_token) + self.token_norm = nn.LayerNorm(c_token) -class Transition(nn.Module): - """LN + SwiGLU FFN with addmm-fused residual; optional Triton LN+w12+SwiGLU kernel.""" + def set_kernel_backend(self, backend: str | None) -> None: + self.token_transformer.set_kernel_backend(backend) - def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: - super().__init__() - self.norm = nn.LayerNorm(d_model) - self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) - # Default chunked; set_chunk_size(None) disables for bit-exact parity tests. - self._chunk_size: int | None = 64 - self._fused_swiglu: nn.Module | None = None - self._kernel_backend: str | None = None + def forward( + self, + x_noisy: Tensor, + t_hat: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + sigma_data: float | None = None, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_token_repr: bool = False, + return_atom_repr: bool = False, + inference_cache: dict[str, Tensor] | None = None, + ) -> dict[str, Tensor | None]: + bsz = x_noisy.shape[0] + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1) + if t.numel() == 1: + t = t.expand(bsz) - def set_chunk_size(self, chunk_size: int | None) -> None: - self._chunk_size = chunk_size + # Step 1: conditioning (pair z is cached across diffusion steps) + s, z = self.conditioning( + t_hat=t, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + sigma_data=sigma, + num_diffusion_samples=num_diffusion_samples, + inference_cache=inference_cache, + ) - def set_kernel_backend(self, backend: str | None) -> None: - """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent).""" - if backend not in _VALID_BACKENDS: - raise ValueError( - f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" - ) - self._kernel_backend = backend - if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE: - assert _FusedLNLinearSwiGLU is not None - d_model = self.norm.normalized_shape[0] - d_inner = self.ffn.hidden_features - has_ln_bias = self.norm.bias is not None - device = self.ffn.w12.weight.device - dtype = self.ffn.w12.weight.dtype - fused = _FusedLNLinearSwiGLU( - d_model=d_model, - d_inner=d_inner, - has_ln_bias=has_ln_bias, - device=device, - dtype=dtype, - ) - with torch.no_grad(): - fused.LN_W.copy_(self.norm.weight) - if has_ln_bias: - fused.LN_B.copy_(self.norm.bias) # type: ignore[union-attr] - # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once. - fused.W12.copy_(self.ffn.w12.weight.t().contiguous()) - self._fused_swiglu = fused.eval().requires_grad_(False) - else: - self._fused_swiglu = None + # Step 2: normalize noisy coords + denom = torch.sqrt(t * t + sigma * sigma) + r_noisy = x_noisy / denom[:, None, None] - def _can_use_fused_path(self, x: Tensor) -> bool: - return ( - _fused_active(self, x) - and self._fused_swiglu is not None - and x.dtype == torch.bfloat16 + # Step 3: atom encoder + a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder( + ref_pos=ref_pos, + atom_attention_mask=ref_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=tok_idx, + r_l=r_noisy, + s_i=s_trunk, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + inference_cache=inference_cache, ) - def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor: - """SwiGLU through silu(x1)*x2, before the final w3.""" - ffn = self.ffn - x12 = ffn.w12(x_normed) - x1, x2 = x12.split(ffn.hidden_features, dim=-1) - return F.silu(x1) * x2 + # Step 4: add conditioned s + a = a + self.s_to_token(self.s_step_norm(s)) - def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor: - """x + w3(hidden) via single cuBLAS addmm — avoids transition-output allocation.""" - ffn = self.ffn - x_shape = x.shape - out = torch.addmm( - x.contiguous().view(-1, x_shape[-1]), - hidden.view(-1, hidden.shape[-1]), - ffn.w3.weight.t(), + # Step 5: token transformer + a, _ = self.token_transformer( + a, + s, + z, + beta=0.0, + attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, ) - return out.view(x_shape) - def forward(self, x: Tensor) -> Tensor: - # Inference-only fast path (addmm-fused residual + pre-alloc out) - # — diverges bit-exactly from ``x + ffn(norm(x))`` so we only use - # it when grad is disabled (binder-design / bit-exact tests run - # with grad on and need the reference path). - if not torch.is_grad_enabled() and self._can_use_fused_path(x): - fused = self._fused_swiglu - assert fused is not None - pre_w3 = fused - if self._chunk_size is None or x.shape[1] <= self._chunk_size: - hidden = pre_w3(x) - return self._addmm_residual(x, hidden) - out = torch.empty_like(x) - for s in range(0, x.shape[1], self._chunk_size): - e = min(s + self._chunk_size, x.shape[1]) - sl = x[:, s:e] - hidden = pre_w3(sl) - out[:, s:e] = self._addmm_residual(sl, hidden) - return out - # Reference path — bit-exact with main: x + ffn(norm(x)). - if self._chunk_size is None or x.shape[1] <= self._chunk_size: - return x + self.ffn(self.norm(x)) - out_list: list[Tensor] = [] - for s in range(0, x.shape[1], self._chunk_size): - e = min(s + self._chunk_size, x.shape[1]) - sl = x[:, s:e] - out_list.append(sl + self.ffn(self.norm(sl))) - return torch.cat(out_list, dim=1) + # Step 6: token norm + a = self.token_norm(a) + # Step 7: atom decoder + r_update, dec_intermediates = self.atom_decoder( + a_i=a, + q_l=q_skip, + c_l=c_skip, + p_lm=p_skip, + atom_to_token=tok_idx, + atom_attention_mask=ref_mask, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + ) -class PairUpdateBlock(nn.Module): - """tri_mul_out, tri_mul_in, pair_transition.""" + # Step 8: compute denoised output + sigma2 = sigma * sigma + t2 = t * t + out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy + out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update - def __init__(self, d_pair: int = 256, expansion_ratio: int = 4) -> None: - super().__init__() - self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) - self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) - self.pair_transition = Transition(d_pair, expansion_ratio=expansion_ratio) - self._kernel_backend: str | None = None - # Row-shared dropout-residual; r=0 for inference (HF model is inference-only). - # backend='fused' swaps in the FusedDropoutResidual Triton kernel. - self.row_drop = DropoutResidual(0.0, batch_dim=1, use_fused_kernels=False) + # Collect atom intermediates from encoder + decoder + atom_intermediates: Tensor | None = None + if return_atom_repr: + all_ints = enc_intermediates + dec_intermediates + if all_ints: + atom_intermediates = torch.stack(all_ints, dim=2) - def set_kernel_backend(self, backend: str | None) -> None: - if backend not in _VALID_BACKENDS: - raise ValueError( - f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" - ) - self.tri_mul_out.set_kernel_backend(backend) - self.tri_mul_in.set_kernel_backend(backend) - self.pair_transition.set_kernel_backend(backend) - self._kernel_backend = backend - self.row_drop = DropoutResidual( - 0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED) - ) + return { + "x_denoised": out, + "token_repr": a if return_token_repr else None, + "atom_intermediates": atom_intermediates, + } - def set_chunk_size(self, chunk_size: int | None) -> None: - self.tri_mul_out.set_chunk_size(chunk_size) - self.tri_mul_in.set_chunk_size(chunk_size) - self.pair_transition.set_chunk_size(chunk_size) - def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool: - return _fused_active(self, pair) and pair.dtype == torch.bfloat16 +# =========================================================================== +# DiffusionStructureHead +# =========================================================================== - def _fused_trimul_with_residual( - self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None - ) -> Tensor: - """Fused TriMul+residual call; weights from the corresponding engine.""" - tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in - engine: TriangleMultiplicativeBlock = tri._engine # type: ignore[assignment] - p_in_weight, g_in_weight = engine.split_kernel_weights() - def _bf16(t: Tensor) -> Tensor: - return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16) +class DiffusionStructureHead(nn.Module): + """Wrapper around DiffusionModule with diffusion sampling.""" - return _fused_trimul_with_residual( # type: ignore[misc] - pair, - direction, - residual=pair, - drop_mask=None, # inference: no dropout, matches internal's eval path - norm_in_weight=_bf16(engine.norm_start.weight), - norm_in_bias=_bf16(engine.norm_start.bias), - p_in_weight=_bf16(p_in_weight), - g_in_weight=_bf16(g_in_weight), - norm_out_weight=_bf16(engine.norm_mix.weight), - norm_out_bias=_bf16(engine.norm_mix.bias), - p_out_weight=_bf16(engine.proj_emit.weight), - g_out_weight=_bf16(engine.proj_gate.weight), - mask=pair_attention_mask, - eps=_EPS, + def __init__(self, config: ESMFold2Config) -> None: + super().__init__() + dm = config.structure_head.diffusion_module + swa_cfg = config.inputs.atom_encoder + sh = config.structure_head + + self.diffusion_module = DiffusionModule( + c_atom=dm.c_atom, + c_token=dm.c_token, + c_z=dm.c_z, + c_s_inputs=dm.c_s_inputs, + sigma_data=dm.sigma_data, + fourier_dim=dm.fourier_dim, + atom_num_blocks=dm.atom_num_blocks, + atom_num_heads=dm.atom_num_heads, + token_num_blocks=dm.token_num_blocks, + token_num_heads=dm.token_num_heads, + transition_multiplier=dm.transition_multiplier, + swa_window_size=swa_cfg.swa_window_size, + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, ) - def forward( - self, pair: Tensor, pair_attention_mask: Tensor | None = None + # Sampling hyperparameters + self.sigma_data = dm.sigma_data + self.gamma_0 = sh.gamma_0 + self.gamma_min = sh.gamma_min + self.noise_scale = sh.noise_scale + self.step_scale = sh.step_scale + self.inference_s_max = sh.inference_s_max + self.inference_s_min = sh.inference_s_min + self.inference_p = sh.inference_p + self.inference_num_steps = sh.inference_num_steps + + def set_kernel_backend(self, backend: str | None) -> None: + self.diffusion_module.set_kernel_backend(backend) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def inference_noise_schedule( + self, num_steps: int | None = None, device: torch.device | None = None ) -> Tensor: - if self._can_use_fused_trimul_with_residual(pair): - pair = self._fused_trimul_with_residual( - pair, "outgoing", pair_attention_mask - ) - pair = self._fused_trimul_with_residual( - pair, "incoming", pair_attention_mask + """Karras power-law noise schedule.""" + steps = self.inference_num_steps if num_steps is None else int(num_steps) + if steps == 1: + return torch.tensor( + [self.inference_s_max * self.sigma_data, 0.0], + device=device, + dtype=torch.float32, ) - else: - pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask)) - pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask)) - pair = self.pair_transition(pair) - return pair + p = float(self.inference_p) + inv_p = 1.0 / p + k = torch.arange(steps, device=device, dtype=torch.float32) + base = self.inference_s_max**inv_p + (k / (steps - 1)) * ( + self.inference_s_min**inv_p - self.inference_s_max**inv_p + ) + schedule = self.sigma_data * base.pow(p) + return F.pad(schedule, (0, 1), value=0.0) + + @staticmethod + def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor: + q = torch.randn((n, 4), dtype=dtype, device=device) + scale = torch.sqrt((q * q).sum(dim=1)) + signs = torch.where(q[:, 0] < 0, -scale, scale) + q = q / signs[:, None] + r, i, j, k = torch.unbind(q, dim=-1) + two_s = 2.0 / (q * q).sum(dim=-1) + return torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + dim=-1, + ).reshape(n, 3, 3) + + def _center_random_augmentation( + self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None + ) -> tuple[Tensor, Tensor | None]: + """Algorithm 19: center + random rotation + translation.""" + bsz = x.shape[0] + mask = atom_mask.unsqueeze(-1) # M has shape (b, a, 1). + denom = mask.sum(dim=1, keepdim=True).clamp(min=1) + mean = (x * mask).sum(dim=1, keepdim=True) / denom + x = x - mean + if second_coords is not None: + second_coords = second_coords - mean + r = self._random_rotations(bsz, x.dtype, x.device) + x = torch.einsum("bmd,bds->bms", x, r) + if second_coords is not None: + second_coords = torch.einsum("bmd,bds->bms", second_coords, r) -class FoldingTrunk(nn.Module): - """ModuleList of PairUpdateBlocks.""" + t = torch.randn_like(x[:, 0:1, :]) + x = x + t + if second_coords is not None: + second_coords = second_coords + t + return x, second_coords - def __init__( - self, n_layers: int = 24, d_pair: int = 256, expansion_ratio: int = 4 - ) -> None: - super().__init__() - self.blocks = nn.ModuleList( - [ - PairUpdateBlock(d_pair=d_pair, expansion_ratio=expansion_ratio) - for _ in range(n_layers) - ] + @staticmethod + def _weighted_rigid_align(x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor) -> Tensor: + """Kabsch alignment: align x to x_gt with weights w.""" + w = (mask * w).unsqueeze(-1) # W has shape (b, n, 1). + denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8) + mu = (x * w).sum(dim=-2, keepdim=True) / denom + mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom + x_c = x - mu + xgt_c = x_gt - mu_gt + covariance = torch.einsum("bni,bnj->bij", w * xgt_c, x_c) + covariance_f32 = covariance.float() + u, _, vh = torch.linalg.svd( + covariance_f32, driver="gesvd" if covariance_f32.is_cuda else None + ) + det = torch.linalg.det(u @ vh) + ones = torch.ones_like(det) + rotation = (u @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ vh).to( + covariance.dtype ) + return x_c @ rotation.transpose(-1, -2) + mu_gt - def set_kernel_backend(self, backend: str | None) -> None: - for block in self.blocks: - cast(PairUpdateBlock, block).set_kernel_backend(backend) + # ------------------------------------------------------------------ + # Sampling + # ------------------------------------------------------------------ - def set_chunk_size(self, chunk_size: int | None) -> None: - for block in self.blocks: - cast(PairUpdateBlock, block).set_chunk_size(chunk_size) + @torch.inference_mode() + def sample( + self, + z_trunk: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + relative_position_encoding: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + num_sampling_steps: int | None = None, + max_inference_sigma: float | None = 256.0, + noise_scale: float | None = None, + step_scale: float | None = None, + return_atom_repr: bool = False, + use_inference_cache: bool = True, + denoising_early_exit_rmsd: float | None = None, + ) -> dict[str, Tensor | None]: + """Diffusion sampling (Algorithm 18). - def forward( - self, pair: Tensor, pair_attention_mask: Tensor | None = None - ) -> Tensor: - # Cast pair → bf16 internally when the fused trimul backend is enabled - # (its bwd kernel requires bf16). Other backends keep the input dtype. - orig_dtype = pair.dtype - fused_on = ( - len(self.blocks) > 0 - and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED - ) - if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16: - pair = pair.to(torch.bfloat16) - for block in self.blocks: - fn = partial(block, pair_attention_mask=pair_attention_mask) - if torch.is_grad_enabled(): - pair = checkpoint(fn, pair, use_reentrant=False) # pyright: ignore - else: - pair = fn(pair) - if pair.dtype != orig_dtype: - pair = pair.to(orig_dtype) - return pair + ``num_sampling_steps`` is the number of denoising steps actually run. + When ``max_inference_sigma`` is set, the Karras schedule built with + ``num_sampling_steps`` entries would lose its high-sigma tail to the cap, + so we inflate the underlying schedule length here to land back at the + requested step count post-truncation. + """ + n_atoms = tok_idx.shape[1] + device = s_inputs.device + target_batch = s_inputs.shape[0] * num_diffusion_samples + inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None -# =========================================================================== -# MSA Encoder -# =========================================================================== + steps = self.inference_num_steps if num_sampling_steps is None else int(num_sampling_steps) + schedule = self.inference_noise_schedule(steps, device) + if max_inference_sigma is not None: + schedule = schedule[schedule <= float(max_inference_sigma)] + schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma)) -class OuterProductMean(nn.Module): - """Outer-product mean: maps an MSA representation into a pair update. + lam = self.noise_scale if noise_scale is None else float(noise_scale) + eta = self.step_scale if step_scale is None else float(step_scale) - The order of the ``/ n_valid`` divide vs. the ``Wout`` projection is - selectable via ``divide_outer_before_proj`` because different ESMFold2 - checkpoints were trained with different orderings: + x = schedule[0] * torch.randn(target_batch, n_atoms, 3, device=device, dtype=torch.float32) + atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float() - * ``False`` (default): ``Wout(outer) / n_valid`` — the projection bias - is scaled by 1/n_valid alongside the outer product. - * ``True``: ``Wout(outer / n_valid)`` — the projection bias is added - unscaled, post-divide. - """ + gammas = torch.where( + schedule > self.gamma_min, + torch.full_like(schedule, self.gamma_0), + torch.zeros_like(schedule), + ) - def __init__( - self, - d_msa: int, - d_hidden: int, - d_pair: int, - divide_outer_before_proj: bool = False, - ) -> None: - super().__init__() - self.d_hidden = d_hidden - self.divide_outer_before_proj = divide_outer_before_proj - self.norm = nn.LayerNorm(d_msa) - self.W = nn.Linear(d_msa, 2 * d_hidden, bias=False) - self.Wout = nn.Linear(d_hidden * d_hidden, d_pair, bias=True) - # Off for bit-exact bf16; ``set_chunk_size(64)`` for long sequences. - self._chunk_size: int | None = None + x_denoised_prev: Tensor | None = None + token_repr: Tensor | None = None + diff_atom_intermediates: Tensor | None = None - def set_chunk_size(self, chunk_size: int | None) -> None: - self._chunk_size = chunk_size + step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:], strict=True)) + num_steps = len(step_pairs) - def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor: - m_norm = self.norm(m) - x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype) - a, b = x.chunk(2, dim=-1) - mask_f = msa_attention_mask.to(a.dtype) - n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0) - if self._chunk_size is None: - outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2) - if self.divide_outer_before_proj: - return self.Wout(outer / n_valid) - return self.Wout(outer) / n_valid - # Chunk along the left (i) axis so the peak einsum intermediate is - # [B, chunk, L, c, d] instead of [B, L, L, c, d]. - L = a.shape[1] - out_chunks: list[Tensor] = [] - for s in range(0, L, self._chunk_size): - e = min(s + self._chunk_size, L) - outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, s:e], b).flatten(-2) - if self.divide_outer_before_proj: - out_chunks.append(self.Wout(outer_chunk / n_valid[:, s:e])) - else: - out_chunks.append(self.Wout(outer_chunk) / n_valid[:, s:e]) - return torch.cat(out_chunks, dim=1) + for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs): + x, x_denoised_prev = self._center_random_augmentation( + x, atom_mask, second_coords=x_denoised_prev + ) + sigma_tm_val = float(sigma_tm.item()) + t_hat_val = sigma_tm_val * (1.0 + float(gamma.item())) + eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5 + x_noisy = x + eps_std * torch.randn_like(x) -class MSAPairWeightedAveraging(nn.Module): - """Pair-biased MSA row update (AF3 Supplement Algorithm 10).""" + is_last_step = step_idx == num_steps - 1 + request_atom_repr = return_atom_repr and ( + is_last_step or denoising_early_exit_rmsd is not None + ) - def __init__( - self, d_msa: int, d_pair: int, n_heads: int = 8, head_width: int = 32 - ) -> None: - super().__init__() - self.n_heads = n_heads - self.head_width = head_width - self.norm_single = nn.LayerNorm(d_msa) - self.compute_bias = nn.Sequential( - nn.LayerNorm(d_pair), nn.Linear(d_pair, n_heads, bias=False) - ) - self.Wv = nn.Linear(d_msa, n_heads * head_width, bias=False) - self.Wgate = nn.Linear(d_msa, n_heads * head_width, bias=False) - self.Wout = nn.Linear(n_heads * head_width, d_msa, bias=False) + dm_out = self.diffusion_module( + x_noisy=x_noisy, + t_hat=torch.full((target_batch,), t_hat_val, device=device, dtype=torch.float32), + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=ref_mask, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + ref_space_uid=ref_space_uid, + tok_idx=tok_idx, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, + return_token_repr=True, + return_atom_repr=request_atom_repr, + inference_cache=inference_cache, + ) - def forward( - self, msa_repr: Tensor, pair_repr: Tensor, pair_attention_mask: Tensor - ) -> Tensor: - """ - Args: - msa_repr: [B, L, M, d_msa] - pair_repr: [B, L, L, d_pair] - pair_attention_mask:[B, L, L] - Returns: - [B, L, M, d_msa] - """ - B, L, M, _ = msa_repr.shape - h, dh = self.n_heads, self.head_width + x_denoised = dm_out["x_denoised"] + token_repr = dm_out["token_repr"] + if request_atom_repr: + diff_atom_intermediates = dm_out.get("atom_intermediates") - msa_normed = self.norm_single(msa_repr) - bias = self.compute_bias(pair_repr) # [B, L, L, n_heads] - bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5) - attn = torch.softmax(bias, dim=-2) # softmax over j + # Reverse diffusion alignment (Kabsch) + with torch.autocast(device_type="cuda", enabled=False): + x_noisy = self._weighted_rigid_align( + x_noisy.float(), x_denoised.float(), atom_mask, atom_mask + ) + x_noisy = x_noisy.to(dtype=x_denoised.dtype) + + # ODE/SDE step + sigma_t_val = float(sigma_t.item()) + denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val + x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma + + # Denoising early-exit: stop when consecutive predictions converge + if ( + denoising_early_exit_rmsd is not None + and x_denoised_prev is not None + and step_idx >= 1 + ): + with torch.autocast(device_type="cuda", enabled=False): + aligned = self._weighted_rigid_align( + x_denoised_prev.float(), + x_denoised.float(), + atom_mask, + atom_mask, + ) + diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1) + per_sample_rmsd = ( + diff.pow(2).sum(dim=(-1, -2)) / atom_mask.sum(dim=-1).clamp(min=1) + ).sqrt() + if per_sample_rmsd.max().item() < denoising_early_exit_rmsd: + x = x_denoised + x_denoised_prev = x_denoised + break - v = self.Wv(msa_normed).reshape(B, L, M, h, dh) - gate = torch.sigmoid(self.Wgate(msa_normed)).reshape(B, L, M, h, dh) + x_denoised_prev = x_denoised - output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate) - return self.Wout(output.reshape(B, L, M, h * dh)) + result: dict[str, Tensor | None] = { + "sample_atom_coords": x, + "diff_token_repr": token_repr, + } + if return_atom_repr: + result["diff_atom_intermediates"] = diff_atom_intermediates + return result diff --git a/fastplms/esmfold2/modeling_esmfold2_experimental.py b/src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py similarity index 79% rename from fastplms/esmfold2/modeling_esmfold2_experimental.py rename to src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py index ad03830..d09f783 100644 --- a/fastplms/esmfold2/modeling_esmfold2_experimental.py +++ b/src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py @@ -8,9 +8,9 @@ from __future__ import annotations -import math +import gc from pathlib import Path -from typing import Any, cast +from typing import Any, ClassVar, cast import torch import torch.nn as nn @@ -18,14 +18,25 @@ from torch import Tensor from transformers.modeling_utils import PreTrainedModel +from .attention import ESMFold2AttentionMixin from .configuration_esmfold2 import ESMFold2Config +from .embedding import ESMFold2EmbeddingMixin from .modeling_esmfold2 import ( - _load_fastplms_esmplusplus_for_esmfold2, + ESMCPrecision, + ESMCPrecisionStatus, + ESMFold2Output, + _drop_transient_esmc_state, + _finalize_structure_output, + _install_esmc_backbone, _lm_precision_context, + _reload_esmc_bf16_for_gradients, + _resolve_structure_output_controls, + _transformer_engine_version, ) from .modeling_esmfold2_common import ( CHAR_VOCAB_SIZE, MAX_ATOMIC_NUMBER, + MSA_CONDITIONING_INPUT_NAMES, NUM_RES_TYPES, DiffusionModule, DiffusionStructureHead, @@ -46,6 +57,8 @@ compute_lm_hidden_states, gather_rep_atom_coords, gather_token_to_atom, + validate_kernel_backend, + validate_msa_conditioning_inputs, ) _EPS = 1e-5 @@ -78,14 +91,10 @@ def __init__(self, config: ESMFold2Config) -> None: self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False) self.s_inputs_norm = nn.LayerNorm(d_inputs) self.z_norm = nn.LayerNorm(d_pair) - self.row_attention_pooling = RowAttentionPooling( - d_pair=d_pair, d_single=d_single - ) + self.row_attention_pooling = RowAttentionPooling(d_pair=d_pair, d_single=d_single) pf = ch.folding_trunk - self.folding_trunk = FoldingTrunk( - n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4 - ) + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) self.plddt_ln = nn.LayerNorm(d_single) max_atoms_per_token = 23 @@ -95,6 +104,7 @@ def __init__(self, config: ESMFold2Config) -> None: self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False) def set_kernel_backend(self, backend: str | None) -> None: + validate_kernel_backend(backend) self.folding_trunk.set_kernel_backend(backend) def set_chunk_size(self, chunk_size: int | None) -> None: @@ -153,9 +163,7 @@ def forward( rep_distances = torch.cdist( rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" ) - distogram_bins = ( - (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() - ) + distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() pair = pair + self.dist_bin_pairwise_embed(distogram_bins) pair_mask = mask[:, :, None].float() * mask[:, None, :].float() @@ -190,13 +198,11 @@ def forward( expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) is_ligand = (expanded_type == _NONPOLYMER_ID).float() - inter_chain = ( - expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) - ).float() + inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float() near_contact = (rep_distances < 8).float() - interface_per_token = ( - near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1) - ).amax(dim=-1) + interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax( + dim=-1 + ) iplddt_weight = torch.where( is_ligand.bool(), torch.full_like(interface_per_token, 2.0), @@ -216,9 +222,7 @@ def forward( n_bins = pae_logits.shape[-1] bin_width = 32.0 / n_bins - bin_centers = torch.arange( - 0.5 * bin_width, 32.0, bin_width, device=pae_logits.device - ) + bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device) mask_f = mask.float() n_res = mask_f.sum(dim=-1, keepdim=True) d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8 @@ -227,9 +231,7 @@ def forward( tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1) pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) - ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / ( - pair_mask_2d.sum(dim=-1) + _EPS - ) + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS) ptm = ptm_per_row.max(dim=-1).values inter_chain_mask = ( @@ -257,9 +259,7 @@ def forward( chain_c2 = (expanded_asym == c2).float() * mask_f pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) denom = pair_m.sum(dim=(-1, -2)) + _EPS - pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum( - dim=(-1, -2) - ) / denom + pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom return { "plddt_logits": plddt_logits, @@ -330,9 +330,7 @@ def forward( pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None - msa_update = self.msa_pair_weighted_averaging( - msa_repr, pair_repr, pair_attention_mask - ) + msa_update = self.msa_pair_weighted_averaging(msa_repr, pair_repr, pair_attention_mask) if mask4d is not None: msa_update = msa_update * mask4d msa_repr = msa_repr + msa_update @@ -413,9 +411,7 @@ def forward( if depth > 1: msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2)) else: - msa_track_mask = torch.zeros( - batch_size, dtype=torch.bool, device=x_pair.device - ) + msa_track_mask = torch.zeros(batch_size, dtype=torch.bool, device=x_pair.device) tok_mask = msa_attention_mask[:, :, 0] pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1) for block in self.blocks: @@ -429,11 +425,11 @@ def forward( return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype) -class ESMFold2ExperimentalModel(PreTrainedModel): +class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin, PreTrainedModel): """Experimental ESMFold2 architecture used by binder-design checkpoints.""" config_class = ESMFold2Config - _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"\._extra_state$"] def __init__(self, config: ESMFold2Config) -> None: super().__init__(config) @@ -454,21 +450,32 @@ def __init__(self, config: ESMFold2Config) -> None: ) self._esmc: nn.Module | None = None self._esmc_fp8 = False + self._esmc_fp8_module_paths: tuple[str, ...] = () + self._esmc_source: str = config.esmc_id + self._esmc_source_revision: str | None = None + self._esmc_source_files: dict[str, str] = {} + self._esmc_local_files_only = False + self._esmc_precision_policy: str = str(getattr(config, "esmc_precision", "auto")) + self._esmc_precision_status = ESMCPrecisionStatus( + requested=self._esmc_precision_policy, + resolved="unloaded", + reason="ESMC has not been loaded.", + device=str(self.device), + transformer_engine_version=_transformer_engine_version(), + ) + self._ttt_lm_head: nn.Module | None = None self._esmfold2_input_builder: Any | None = None + self._kernel_backend: str | None = None pf = config.folding_trunk - self.folding_trunk = FoldingTrunk( - n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4 - ) + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) self.pair_loop_proj = nn.Sequential( nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False) ) nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight) self.structure_head = DiffusionStructureHead(config) - self.distogram_head = nn.Linear( - d_pair, config.structure_head.distogram_bins, bias=True - ) + self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True) self.confidence_head: ConfidenceHead | None = ( ConfidenceHead(config) if config.confidence_head.enabled else None ) @@ -487,16 +494,19 @@ def __init__(self, config: ESMFold2Config) -> None: ) self.post_init() + self._register_state_dict_hook(_drop_transient_esmc_state) @property def device(self) -> torch.device: return next(self.parameters()).device def set_kernel_backend(self, backend: str | None) -> None: + validate_kernel_backend(backend) self.folding_trunk.set_kernel_backend(backend) if self.confidence_head is not None: self.confidence_head.set_kernel_backend(backend) self.structure_head.set_kernel_backend(backend) + self._kernel_backend = backend def set_chunk_size(self, chunk_size: int | None) -> None: self.folding_trunk.set_chunk_size(chunk_size) @@ -512,41 +522,57 @@ def configure_lm_dropout( force_lm_dropout_during_inference: bool = True, ) -> None: self.config.lm_dropout = lm_dropout - self.config.force_lm_dropout_during_inference = ( - force_lm_dropout_during_inference - ) + self.config.force_lm_dropout_during_inference = force_lm_dropout_during_inference - def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None: - dtype_map = { - "bf16": torch.bfloat16, - "fp32": torch.float32, - } - if precision not in dtype_map: - if precision == "fp8": - raise RuntimeError( - "esmc_precision='fp8' is supported only by the standard " - "released ESMFold2 model. The experimental binder-design " - "model keeps the FastPLMs ESM++ backbone in bf16 or fp32." - ) - raise ValueError(f"precision must be one of {list(dtype_map)}, got {precision!r}") - esmc = _load_fastplms_esmplusplus_for_esmfold2( - esmc_model_path=esmc_model_path, - attn_backend=self.config.esmc_attn_backend, - device=self.device, - dtype=dtype_map[precision], - ) - assert esmc.config.hidden_size == self.config.lm_d_model, ( - f"ESMFold2 expected lm_d_model={self.config.lm_d_model}, " - f"but loaded ESM++ hidden_size={esmc.config.hidden_size}." - ) - assert esmc.config.num_hidden_layers == self.config.lm_num_layers, ( - f"ESMFold2 expected lm_num_layers={self.config.lm_num_layers}, " - f"but loaded ESM++ num_hidden_layers={esmc.config.num_hidden_layers}." + @property + def esmc_precision_status(self) -> ESMCPrecisionStatus: + return self._esmc_precision_status + + def load_esmc( + self, + esmc_model_path: str, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool = False, + ) -> None: + """Load ESMC with the same precision policy as released checkpoints.""" + + _install_esmc_backbone( + self, + esmc_model_path, + precision=precision, + device=device, + local_files_only=local_files_only, ) - for parameter in esmc.parameters(): - parameter.requires_grad_(False) + + def reload_esmc( + self, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool | None = None, + ) -> None: + """Reload canonical ESMC weights and discard runtime quantization.""" + + source = self._esmc_source or self.config.esmc_id + old_esmc = self._esmc + self._esmc = None self._esmc_fp8 = False - self._esmc = esmc + self._esmc_fp8_module_paths = () + self._ttt_lm_head = None + del old_esmc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.load_esmc( + source, + precision=precision, + device=device, + local_files_only=( + self._esmc_local_files_only + if local_files_only is None + else local_files_only + ), + ) @classmethod def from_pretrained( @@ -560,21 +586,23 @@ def from_pretrained( kwargs["config"] = ESMFold2Config.from_pretrained( pretrained_model_name_or_path, **kwargs ) - esmc_precision = kwargs.pop("esmc_precision", "bf16") - model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + esmc_precision = kwargs.pop("esmc_precision", None) + local_files_only = bool(kwargs.get("local_files_only", False)) + output_loading_info = bool(kwargs.get("output_loading_info", False)) + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + if output_loading_info: + model, loading_info = loaded + else: + model = loaded if load_esmc: - model.load_esmc(model.config.esmc_id, precision=esmc_precision) - return model - - def apply_torch_compile( - self, mode: str = "fixed_seqlen", dynamic: bool | None = None - ) -> None: - import torch._dynamo - - torch._dynamo.config.cache_size_limit = 512 - torch._dynamo.config.accumulated_cache_size_limit = 512 - torch._dynamo.config.capture_scalar_outputs = True + model.load_esmc( + model.config.esmc_id, + precision=esmc_precision or model.config.esmc_precision, + local_files_only=local_files_only, + ) + return (model, loading_info) if output_loading_info else model + def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None: if dynamic is None: dynamic = mode == "dynamic_seqlen" compile_kwargs: dict[str, bool] = {"dynamic": dynamic} @@ -599,9 +627,18 @@ def _compute_lm_hidden_states( mol_type: Tensor, tok_mask: Tensor, ) -> Tensor: - assert self._esmc is not None - pad_to = 8 if self._esmc_fp8 else None - with _lm_precision_context(self._esmc_fp8): + if self._esmc_fp8 and torch.is_grad_enabled(): + _reload_esmc_bf16_for_gradients( + self, + reason=( + "Gradient-enabled ESMC execution requires BF16; the persisted " + "serving policy is unchanged." + ), + ) + if self._esmc is None: + raise RuntimeError("ESMFold2 language-model features require load_esmc=True.") + pad_to = 16 if self._esmc_fp8 else None + with _lm_precision_context(self._esmc_precision_status.resolved, self.device): return compute_lm_hidden_states( self._esmc, input_ids, @@ -648,9 +685,25 @@ def forward( provide_soft_sequence_to_msa_and_profile: bool = True, noise_scale: float | None = None, step_scale: float | None = None, - max_inference_sigma: int | None = None, - ) -> dict[str, Tensor]: - del noise_scale, step_scale, max_inference_sigma + max_inference_sigma: float | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> ESMFold2Output | tuple[Any, ...]: + output_hidden_states, return_dict = _resolve_structure_output_controls( + self.config, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + validate_msa_conditioning_inputs( + self.config, + msa=msa, + msa_attention_mask=msa_attention_mask, + has_deletion=has_deletion, + deletion_value=deletion_value, + deletion_mean=deletion_mean, + ) tok_mask = token_attention_mask atm_mask = atom_attention_mask n_loops = num_loops if num_loops is not None else self.config.num_loops @@ -680,10 +733,7 @@ def forward( if res_type_soft is not None: res_type_oh = res_type_soft.float() - if ( - not self.config.disable_msa_features - and provide_soft_sequence_to_msa_and_profile - ): + if not self.config.disable_msa_features and provide_soft_sequence_to_msa_and_profile: profile = res_type_oh msa = res_type_oh.unsqueeze(1) msa_attention_mask = tok_mask.unsqueeze(1) @@ -696,24 +746,17 @@ def forward( profile = torch.zeros_like(profile) deletion_mean = torch.zeros_like(deletion_mean) - ref_element_oh = F.one_hot( - ref_element.long(), num_classes=MAX_ATOMIC_NUMBER - ).float() + ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float() ref_atom_name_chars_oh = F.one_hot( ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE ).float() atm_mask_f = atm_mask.float() ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1) - ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze( - -1 - ).unsqueeze(-1) + ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1) atom_to_token = atom_to_token * atm_mask.long() use_amp = ref_pos.device.type == "cuda" - with ( - torch.set_grad_enabled(res_type_soft is not None), - torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16), - ): + with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16): x_inputs = self.inputs_embedder( aatype=res_type_oh, profile=profile.float(), @@ -727,9 +770,7 @@ def forward( atom_to_token=atom_to_token, ) - z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( - x_inputs - ).unsqueeze(1) + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1) relative_position_encoding = self.rel_pos( residue_index=residue_index, asym_id=asym_id, @@ -740,11 +781,7 @@ def forward( token_bonds_encoding = self.token_bonds(token_bonds.float()) z_init = z_init + relative_position_encoding + token_bonds_encoding - if ( - lm_hidden_states is None - and input_ids is not None - and self._esmc is not None - ): + if lm_hidden_states is None and input_ids is not None and self._esmc is not None: lm_hidden_states = self._compute_lm_hidden_states( input_ids, asym_id, residue_index, mol_type, tok_mask ) @@ -754,9 +791,7 @@ def forward( if self.config.force_lm_dropout_during_inference or self.training else 0.0 ) - lm_z = self.language_model( - lm_hidden_states.detach(), lm_dropout=lm_dropout - ) + lm_z = self.language_model(lm_hidden_states.detach(), lm_dropout=lm_dropout) z_init = z_init + lm_z.to(z_init.dtype) msa_kwargs: dict[str, Tensor] | None = None @@ -806,9 +841,9 @@ def forward( if early_exit and loop_num < n_loops: l2_converged = False if prev_pair is not None and loop_num > 0: - rel_l2 = (z.float() - prev_pair.float()).norm() / prev_pair.float().norm().clamp( - min=1e-8 - ) + rel_l2 = ( + z.float() - prev_pair.float() + ).norm() / prev_pair.float().norm().clamp(min=1e-8) l2_converged = rel_l2.item() < 0.25 prev_pair = z.detach().clone() sym_z = z.float() + z.float().transpose(-2, -3) @@ -816,10 +851,7 @@ def forward( if prev_disto_probs is not None and loop_num > 0: kl_per_pair = ( cur_probs - * ( - cur_probs.clamp(min=1e-8) - / prev_disto_probs.clamp(min=1e-8) - ).log() + * (cur_probs.clamp(min=1e-8) / prev_disto_probs.clamp(min=1e-8)).log() ).sum(-1) kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2 if l2_converged or kl.item() < 0.05: @@ -849,11 +881,15 @@ def forward( token_attention_mask=tok_mask, num_diffusion_samples=n_samples, num_sampling_steps=num_sampling_steps, + max_inference_sigma=max_inference_sigma, + noise_scale=noise_scale, + step_scale=step_scale, return_atom_repr=False, denoising_early_exit_rmsd=(0.10 if early_exit else None), ) sample_coords = structure_output["sample_atom_coords"] - assert sample_coords is not None + if sample_coords is None: + raise RuntimeError("ESMFold2 structure sampling did not return coordinates.") if sample_coords.ndim == 4: batch, sample_count, atom_count, coord_dim = sample_coords.shape sample_coords_for_gather = sample_coords.reshape( @@ -891,12 +927,16 @@ def forward( token_bonds_encoding=token_bonds_encoding.detach(), ) output.update(confidence_output) - output["atom_pad_mask"] = ( - atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask - ) + output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask output["residue_index"] = residue_index output["entity_id"] = entity_id - return output + return _finalize_structure_output( + output, + token_input_state=x_inputs, + pair_state=z, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) @property def input_builder(self): @@ -913,15 +953,27 @@ def input_types(self): return esmfold2_types def prepare_structure_input(self, input, seed: int | None = None): - return self.input_builder.prepare_input(input, seed=seed, device=self.device) + return self.input_builder.prepare_model_input( + self, + input, + seed=seed, + device=self.device, + ) @torch.no_grad() - def infer_protein(self, seq: str, **forward_kwargs) -> dict[str, Tensor]: + def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output: from .protein_utils import prepare_protein_features + if forward_kwargs.pop("return_dict", True) is not True: + raise ValueError( + "infer_protein always returns a mapping; return_dict=False is invalid." + ) features = prepare_protein_features(seq) + if not self.config.msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) features = {name: tensor.to(self.device) for name, tensor in features.items()} - output = self(**features, **forward_kwargs) + output = self(**features, **forward_kwargs, return_dict=True) for name in ( "res_type", "atom_to_token", @@ -974,9 +1026,7 @@ def fold_protein( ): from .esmfold2_types import ProteinInput, StructurePredictionInput - input = StructurePredictionInput( - sequences=[ProteinInput(id=chain_id, sequence=sequence)] - ) + input = StructurePredictionInput(sequences=[ProteinInput(id=chain_id, sequence=sequence)]) return self.fold( input, num_loops=num_loops, @@ -988,12 +1038,14 @@ def fold_protein( @staticmethod def result_to_cif(result) -> str: - assert not isinstance(result, list), "Pass one MolecularComplexResult at a time." + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") return result.complex.to_mmcif() @staticmethod def result_to_pdb(result) -> str: - assert not isinstance(result, list), "Pass one MolecularComplexResult at a time." + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") return result.complex.to_protein_complex().to_pdb_string() def save_as_cif(self, result, output_path: str | Path) -> None: @@ -1011,7 +1063,8 @@ def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str: __all__ = [ "ConfidenceHead", + "ESMFold2ExperimentalModel", + "ESMFold2Output", "MSAEncoder", "MSAEncoderBlock", - "ESMFold2ExperimentalModel", ] diff --git a/src/fastplms/models/esmfold2/protein_reference_geometry.json b/src/fastplms/models/esmfold2/protein_reference_geometry.json new file mode 100644 index 0000000..5cab64b --- /dev/null +++ b/src/fastplms/models/esmfold2/protein_reference_geometry.json @@ -0,0 +1 @@ +{"dtype":"float32","provenance":{"contract":"biohub_esmfold2_input_v1","manifest_family":"esmfold2"},"residues":{"ALA":{"C":[1.2127548456192017,0.4737588167190552,0.19521640241146088],"CA":[-0.04190138354897499,0.17447763681411743,-0.5729365348815918],"CB":[-1.276943325996399,0.4288230538368225,0.29937705397605896],"N":[-0.01003183238208294,-1.2073018550872803,-1.0555061101913452],"O":[1.9390329122543335,1.4484562873840332,-0.13759790360927582]},"ARG":{"C":[-3.469440460205078,-1.0612813234329224,-0.2755832374095917],"CA":[-2.0503084659576416,-0.5735036730766296,-0.4097220301628113],"CB":[-1.4193516969680786,-0.3735991418361664,0.9852858781814575],"CD":[0.6643245816230774,1.0068185329437256,0.3963329493999481],"CG":[0.11878877878189087,-0.3112654983997345,0.963895857334137],"CZ":[3.098905324935913,0.3215920031070709,-0.09047172218561172],"N":[-2.0170421600341797,0.6717798113822937,-1.1794233322143555],"NE":[2.1090238094329834,1.0977025032043457,0.6120952367782593],"NH1":[4.461230278015137,0.3844667971134186,0.34141138195991516],"NH2":[2.7856509685516357,-0.4166366159915924,-1.1148239374160767],"O":[-3.8218462467193604,-2.1369943618774414,-0.8294969797134399]},"ASN":{"C":[-1.9211044311523438,-0.6982439160346985,-0.42196929454803467],"CA":[-0.76087886095047,0.23876343667507172,-0.23573364317417145],"CB":[0.5504899024963379,-0.5078350305557251,-0.5390339493751526],"CG":[1.7250099182128906,0.4264017939567566,-0.5778228640556335],"N":[-0.7595629096031189,0.7503494620323181,1.1369825601577759],"ND2":[2.57365345954895,0.5730618834495544,0.5608599781990051],"O":[-2.677666187286377,-0.5753439664840698,-1.4223182201385498],"OD1":[1.9470350742340088,1.1086392402648926,-1.613560438156128]},"ASP":{"C":[-0.9431572556495667,1.0356197357177734,0.18555717170238495],"CA":[-0.6379959583282471,-0.41974392533302307,0.41681644320487976],"CB":[0.48594576120376587,-0.8970447778701782,-0.5209363698959351],"CG":[1.780342936515808,-0.19918935000896454,-0.2310730367898941],"N":[-1.8452696800231934,-1.2169504165649414,0.19437327980995178],"O":[-1.5183608531951904,1.4045922756195068,-0.8739855885505676],"OD1":[2.5202910900115967,-0.6044584512710571,0.7049641013145447],"OD2":[2.1454880237579346,0.9208861589431763,-0.9712985157966614]},"CYS":{"C":[-1.2652032375335693,-0.6832379698753357,-0.3594406247138977],"CA":[0.11344368755817413,-0.09400428831577301,-0.45952197909355164],"CB":[0.6919880509376526,0.09034398198127747,0.952482283115387],"N":[0.0469963513314724,1.190075159072876,-1.1607273817062378],"O":[-1.4631439447402954,-1.8851220607757568,-0.6826791763305664],"SG":[2.4619927406311035,0.5235707759857178,0.9020372629165649]},"GLN":{"C":[-1.7545503377914429,0.7091967463493347,0.8433493971824646],"CA":[-1.370002269744873,-0.6000258922576904,0.2103111445903778],"CB":[0.02040259726345539,-0.5004461407661438,-0.44764479994773865],"CD":[2.4745187759399414,-0.24800164997577667,-0.09364881366491318],"CG":[1.1377512216567993,-0.28680720925331116,0.582992434501648],"N":[-2.370004653930664,-0.9637529850006104,-0.7942749261856079],"NE2":[2.947425603866577,0.9601329565048218,-0.6888364553451538],"O":[-1.8520662784576416,0.7999289631843567,2.0964975357055664],"OE1":[3.1685523986816406,-1.2966246604919434,-0.1717153936624527]},"GLU":{"C":[-1.7741456031799316,0.9664392471313477,0.09259600937366486],"CA":[-1.0560977458953857,0.027459044009447098,1.0306966304779053],"CB":[0.4706551432609558,0.048803869634866714,0.8114414811134338],"CD":[2.398822069168091,-0.3097084164619446,-0.7210537791252136],"CG":[0.9133604764938354,-0.4219329059123993,-0.5830985307693481],"N":[-1.5850872993469238,-1.337684154510498,0.9490851163864136],"O":[-1.9012441635131836,2.181349992752075,0.402479350566864],"OE1":[3.1389315128326416,-1.274524450302124,-0.39029765129089355],"OE2":[2.9647817611694336,0.8781346082687378,-1.1732689142227173]},"GLY":{"C":[0.9440054893493652,-0.10314033925533295,0.19859643280506134],"CA":[-0.39974430203437805,0.5488945245742798,0.15242962539196014],"N":[-1.3942985534667969,-0.39875128865242004,-0.3370324671268463],"O":[1.3352899551391602,-0.669218122959137,1.2541258335113525]},"HIS":{"C":[-2.675257921218872,0.6571555733680725,-0.30441102385520935],"CA":[-1.3396095037460327,0.24797579646110535,0.24960045516490936],"CB":[-0.3041955828666687,0.21721023321151733,-0.8885309100151062],"CD2":[1.780855417251587,-1.1011489629745483,-0.3814258575439453],"CE1":[2.9566943645477295,0.4924798905849457,0.6477115750312805],"CG":[1.0887513160705566,0.028941065073013306,-0.36419469118118286],"N":[-1.4532867670059204,-1.0689626932144165,0.881072461605072],"ND1":[1.840459942817688,1.0411773920059204,0.29804590344429016],"NE2":[3.0280203819274902,-0.8751969337463379,0.26084381341934204],"O":[-3.1311378479003906,1.8079776763916016,-0.06785715371370316]},"ILE":{"C":[-1.3896740674972534,0.8142145276069641,-1.1164065599441528],"CA":[-1.0636085271835327,-0.35169270634651184,-0.21393552422523499],"CB":[0.061667006462812424,0.01599610224366188,0.8057394623756409],"CD1":[1.7929610013961792,0.899773120880127,-0.8863027691841125],"CG1":[1.502519965171814,-0.08899776637554169,0.24154816567897797],"CG2":[-0.053174979984760284,-0.8521055579185486,2.0702083110809326],"N":[-0.7167549729347229,-1.5426139831542969,-0.9983330368995667],"O":[-1.2377792596817017,0.7302915453910828,-2.3656840324401855]},"LEU":{"C":[1.9905058145523071,0.24182087182998657,0.7879968285560608],"CA":[1.3077669143676758,-0.6677430868148804,-0.19492436945438385],"CB":[-0.20306941866874695,-0.8093230128288269,0.11243502795696259],"CD1":[-2.4228057861328125,0.29949337244033813,0.573042094707489],"CD2":[-1.0282856225967407,1.1250264644622803,-1.346014380455017],"CG":[-0.9916267395019531,0.5234957337379456,0.06723011285066605],"N":[1.9657520055770874,-1.9763224124908447,-0.18391533195972443],"O":[2.06896710395813,-0.07880014181137085,2.0048046112060547]},"LYS":{"C":[2.7168593406677246,1.595757246017456,-0.20924785733222961],"CA":[2.0314927101135254,0.2786507308483124,-0.4298512041568756],"CB":[0.5018402934074402,0.4873858690261841,-0.49062973260879517],"CD":[-1.769762635231018,-0.5552700161933899,-1.040329933166504],"CE":[-2.576533555984497,-1.0221366882324219,0.18493641912937164],"CG":[-0.25062066316604614,-0.7894009947776794,-0.9055535793304443],"N":[2.4221372604370117,-0.6473312377929688,0.6370573043823242],"NZ":[-2.269151210784912,-0.24293844401836395,1.3849012851715088],"O":[3.397681713104248,2.116427421569824,-1.1332510709762573]},"MET":{"C":[2.30391001701355,0.8367712497711182,-0.7254616618156433],"CA":[1.2630571126937866,-0.24417810142040253,-0.7626462578773499],"CB":[0.10567972809076309,0.10861825942993164,0.19741646945476532],"CE":[-3.265165090560913,0.7033554911613464,-0.11588376015424728],"CG":[-1.0658042430877686,-0.8736631274223328,0.08811883628368378],"N":[1.8903918266296387,-1.5252995491027832,-0.42638593912124634],"O":[2.465414524078369,1.5928632020950317,-1.7207728624343872],"SD":[-2.4557132720947266,-0.3332225978374481,1.1461700201034546]},"PHE":{"C":[-1.8900631666183472,0.45833414793014526,1.0232222080230713],"CA":[-1.591969609260559,-0.8545162677764893,0.35214468836784363],"CB":[-0.760358452796936,-0.6342853307723999,-0.9257160425186157],"CD1":[0.8468314409255981,1.2480632066726685,-0.7146694660186768],"CD2":[1.6827683448791504,-0.9758077263832092,-0.1423054188489914],"CE1":[2.1801748275756836,1.7875733375549316,-0.3744623064994812],"CE2":[2.888307809829712,-0.48277512192726135,0.16804970800876617],"CG":[0.604112982749939,-0.07200468331575394,-0.6148118376731873],"CZ":[3.149812936782837,0.9656873941421509,0.04440271109342575],"N":[-2.8484435081481934,-1.525790810585022,0.01789816841483116],"O":[-1.3424992561340332,0.74432373046875,2.121629476547241]},"PRO":{"C":[1.6121541261672974,-1.1711241006851196,0.31082412600517273],"CA":[0.32722190022468567,-0.6164458394050598,-0.25072571635246277],"CB":[0.3248198926448822,0.9028244018554688,-0.33368146419525146],"CD":[-1.8495968580245972,0.026575811207294464,0.2681289613246918],"CG":[-1.1425083875656128,1.2730128765106201,-0.2590600252151489],"N":[-0.836250364780426,-0.9899801015853882,0.5561304688453674],"O":[1.6127740144729614,-2.2771971225738525,0.9156193733215332]},"SER":{"C":[0.9941009879112244,-0.5374617576599121,0.73505038022995],"CA":[0.00013792862591799349,0.4966467022895813,0.28510504961013794],"CB":[-1.1279288530349731,-0.1659376323223114,-0.5160963535308838],"N":[0.674650251865387,1.5018702745437622,-0.5367295145988464],"O":[1.0545241832733154,-0.8683545589447021,1.9495396614074707],"OG":[-1.8135979175567627,-1.085249662399292,0.28947514295578003]},"THR":{"C":[-1.294381856918335,0.7077372074127197,-0.5549946427345276],"CA":[-0.5433306097984314,-0.16364754736423492,0.41697052121162415],"CB":[0.853203296661377,-0.5363803505897522,-0.14109353721141815],"CG2":[1.7225933074951172,0.7054727077484131,-0.3651331067085266],"N":[-1.325830340385437,-1.3728225231170654,0.6882233023643494],"O":[-1.6939635276794434,0.23654410243034363,-1.6540418863296509],"OG1":[1.5220820903778076,-1.379003643989563,0.7635167837142944]},"TRP":{"C":[2.1113572120666504,-0.6121063232421875,-0.7733646035194397],"CA":[2.384092092514038,0.09079249948263168,0.5325262546539307],"CB":[1.281521201133728,1.1139036417007446,0.8559791445732117],"CD1":[-0.42329534888267517,-0.15470874309539795,2.2227554321289062],"CD2":[-1.1023900508880615,0.2158389836549759,0.11529432237148285],"CE2":[-2.045644998550415,-0.4881173074245453,0.710669219493866],"CE3":[-1.2173502445220947,0.6102271676063538,-1.300106406211853],"CG":[-0.04292375594377518,0.44645074009895325,1.0942792892456055],"CH2":[-3.3817875385284424,-0.5677337646484375,-1.3032053709030151],"CZ2":[-3.256009340286255,-0.9164394736289978,-0.00984987337142229],"CZ3":[-2.315925121307373,0.2306906282901764,-1.9776310920715332],"N":[3.686030864715576,0.7599999904632568,0.496155709028244],"NE1":[-1.7030320167541504,-0.7665823101997375,2.0595016479492188],"O":[1.796526312828064,-1.8323148488998413,-0.7775964140892029]},"TYR":{"C":[-3.347280740737915,0.3588399887084961,-0.09830684959888458],"CA":[-1.913882851600647,0.23552845418453217,0.330669641494751],"CB":[-1.0093992948532104,0.0004731413209810853,-0.8981552124023438],"CD1":[1.0992432832717896,1.1877919435501099,-0.3579142987728119],"CD2":[1.1803174018859863,-1.253401279449463,-0.31122180819511414],"CE1":[2.5253450870513916,1.1990256309509277,0.029804613441228867],"CE2":[2.471151113510132,-1.240687608718872,0.043534230440855026],"CG":[0.4520410895347595,0.021162061020731926,-0.5305932760238647],"CZ":[3.180687665939331,0.04672492295503616,0.2214856892824173],"N":[-1.7900604009628296,-0.8409399390220642,1.3180142641067505],"O":[-3.967811346054077,-0.6449354290962219,-0.5423302054405212],"OH":[4.523719787597656,0.0671030730009079,0.5877485871315002]},"UNK":{"C":[0.0,0.0,0.0],"CA":[0.0,0.0,0.0],"N":[0.0,0.0,0.0],"O":[0.0,0.0,0.0]},"VAL":{"C":[1.8391697406768799,0.4067850410938263,0.06351757049560547],"CA":[0.6014357209205627,-0.10503966361284256,-0.6336286664009094],"CB":[-0.694736897945404,0.4259096384048462,0.03581475466489792],"CG1":[-1.9276031255722046,0.09515828639268875,-0.8172357082366943],"CG2":[-0.8938426971435547,-0.08640842139720917,1.472349762916565],"N":[0.5987519025802612,-1.569443702697754,-0.7379124760627747],"O":[2.3952062129974365,-0.2666190266609192,0.9731166958808899]}},"schema":"fastplms.esmfold2.reference_geometry.v1"} diff --git a/src/fastplms/models/esmfold2/protein_utils.py b/src/fastplms/models/esmfold2/protein_utils.py new file mode 100644 index 0000000..1d7d0ea --- /dev/null +++ b/src/fastplms/models/esmfold2/protein_utils.py @@ -0,0 +1,178 @@ +"""Protein-only ESMFold2 featurization without the Biohub runtime package. + +Input is one amino-acid sequence. The transformation expands each residue into +the checkpoint atom schema, pads atoms to a multiple of 32, and emits batched +token, atom, and single-sequence MSA tensors. Reference coordinates are loaded +lazily from a provenance-bearing declarative package asset. +""" + +from __future__ import annotations + +import json +from functools import cache +from importlib.resources import files +from typing import Any + +import torch +from torch import Tensor + +from .esmfold2_constants import ( + CHARGED_ATOMS, + ELEMENT_TO_ATOMIC_NUM, + ESM_PROTEIN_VOCAB, + MOL_TYPE_PROTEIN, + PROTEIN_1TO3, + PROTEIN_HEAVY_ATOMS, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, +) + +_GEOMETRY_ASSET = "protein_reference_geometry.json" +_GEOMETRY_SCHEMA = "fastplms.esmfold2.reference_geometry.v1" + + +@cache +def _reference_geometry() -> dict[str, dict[str, tuple[float, float, float]]]: + resource = files(__package__).joinpath(_GEOMETRY_ASSET) + with resource.open(mode="r", encoding="utf-8") as handle: + payload = json.load(handle) + if ( + payload.get("schema") != _GEOMETRY_SCHEMA + or payload.get("dtype") != "float32" + or payload.get("provenance", {}).get("manifest_family") != "esmfold2" + ): + raise RuntimeError("The ESMFold2 reference-geometry asset has invalid provenance.") + + raw_residues = payload.get("residues") + if not isinstance(raw_residues, dict): + raise RuntimeError("The ESMFold2 reference-geometry asset has no residue table.") + geometry: dict[str, dict[str, tuple[float, float, float]]] = {} + for residue, atom_positions in raw_residues.items(): + if not isinstance(residue, str) or not isinstance(atom_positions, dict): + raise RuntimeError("The ESMFold2 reference-geometry residue table is malformed.") + geometry[residue] = {} + for atom_name, position in atom_positions.items(): + if ( + not isinstance(atom_name, str) + or not isinstance(position, list) + or len(position) != 3 + ): + raise RuntimeError("The ESMFold2 reference-geometry atom table is malformed.") + geometry[residue][atom_name] = tuple(float(value) for value in position) + + expected_residues = set(PROTEIN_HEAVY_ATOMS) - {"MSE"} + if set(geometry) != expected_residues: + raise RuntimeError("The ESMFold2 reference-geometry residue set is incomplete.") + for residue, atom_names in PROTEIN_HEAVY_ATOMS.items(): + if residue == "MSE": + continue + if set(geometry[residue]) != set(atom_names): + raise RuntimeError(f"Reference geometry differs from the atom schema for {residue}.") + return geometry + + +def _encode_atom_name(atom_name: str) -> tuple[int, int, int, int]: + padded = atom_name.ljust(4)[:4] + return tuple(ord(character) - 32 if character != " " else 0 for character in padded) + + +def _padded_atom_count(actual_count: int) -> int: + return max(32, ((actual_count + 31) // 32) * 32) + + +def _residue_records(sequence: str) -> tuple[list[dict[str, Any]], list[int], list[int], list[int]]: + geometry = _reference_geometry() + atoms: list[dict[str, Any]] = [] + residue_types: list[int] = [] + input_ids: list[int] = [] + representative_atoms: list[int] = [] + + for token_index, residue_letter in enumerate(sequence): + residue_name = PROTEIN_1TO3.get(residue_letter, "UNK") + atom_names = PROTEIN_HEAVY_ATOMS[residue_name] + atom_start = len(atoms) + for atom_name in atom_names: + atoms.append( + { + "token_index": token_index, + "name": atom_name, + "element": atom_name[0], + "charge": CHARGED_ATOMS.get((residue_name, atom_name), 0), + "position": geometry[residue_name][atom_name], + } + ) + + representative_name = "CB" if "CB" in atom_names else "CA" + representative_atoms.append(atom_start + atom_names.index(representative_name)) + residue_types.append(PROTEIN_RESIDUE_TO_RES_TYPE.get(residue_name, PROTEIN_UNK_RES_TYPE)) + input_ids.append(ESM_PROTEIN_VOCAB.get(residue_letter, ESM_PROTEIN_VOCAB["X"])) + + return atoms, residue_types, input_ids, representative_atoms + + +def prepare_protein_features(sequence: str) -> dict[str, Tensor]: + """Build the protein-only feature mapping consumed by ESMFold2. + + Every tensor includes a leading batch dimension. Biological tokens have + length ``l``; atom tensors have length ``n_atoms``, where ``n_atoms`` is the + smallest multiple of 32 covering all heavy atoms. + """ + + if not sequence: + raise ValueError("sequence must be non-empty") + + atoms, residue_types, input_ids, representative_atoms = _residue_records(sequence) + sequence_length = len(sequence) + n_atoms = _padded_atom_count(len(atoms)) + + ref_pos = torch.zeros((n_atoms, 3), dtype=torch.float32) + ref_element = torch.zeros(n_atoms, dtype=torch.int64) + ref_charge = torch.zeros(n_atoms, dtype=torch.int8) + ref_atom_name_chars = torch.zeros((n_atoms, 4), dtype=torch.int64) + ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64) + atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool) + atom_to_token = torch.zeros(n_atoms, dtype=torch.int64) + + for atom_index, atom in enumerate(atoms): + token_index = atom["token_index"] + ref_pos[atom_index] = torch.tensor(atom["position"], dtype=torch.float32) + ref_element[atom_index] = ELEMENT_TO_ATOMIC_NUM[atom["element"]] + ref_charge[atom_index] = atom["charge"] + ref_atom_name_chars[atom_index] = torch.tensor( + _encode_atom_name(atom["name"]), dtype=torch.int64 + ) + ref_space_uid[atom_index] = token_index + atom_attention_mask[atom_index] = True + atom_to_token[atom_index] = token_index + + residue_type_tensor = torch.tensor(residue_types, dtype=torch.int64) + msa = residue_type_tensor.unsqueeze(0) + features = { + "token_index": torch.arange(sequence_length, dtype=torch.int64), + "residue_index": torch.arange(sequence_length, dtype=torch.int64), + "asym_id": torch.zeros(sequence_length, dtype=torch.int64), + "sym_id": torch.zeros(sequence_length, dtype=torch.int64), + "entity_id": torch.ones(sequence_length, dtype=torch.int64), + "mol_type": torch.full((sequence_length,), MOL_TYPE_PROTEIN, dtype=torch.int64), + "res_type": residue_type_tensor, + "input_ids": torch.tensor(input_ids, dtype=torch.int64), + "token_bonds": torch.zeros((sequence_length, sequence_length, 1), dtype=torch.float32), + "token_attention_mask": torch.ones(sequence_length, dtype=torch.bool), + "ref_pos": ref_pos, + "ref_element": ref_element, + "ref_charge": ref_charge, + "ref_atom_name_chars": ref_atom_name_chars, + "ref_space_uid": ref_space_uid, + "atom_attention_mask": atom_attention_mask, + "atom_to_token": atom_to_token, + "distogram_atom_idx": torch.tensor(representative_atoms, dtype=torch.int64), + "msa": msa, + "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool), + "has_deletion": torch.zeros_like(msa, dtype=torch.bool), + "deletion_value": torch.zeros_like(msa, dtype=torch.float32), + "deletion_mean": torch.zeros(sequence_length, dtype=torch.float32), + } + return {name: tensor.unsqueeze(0) for name, tensor in features.items()} + + +__all__ = ["prepare_protein_features"] diff --git a/src/fastplms/models/esmfold2/reproducibility.py b/src/fastplms/models/esmfold2/reproducibility.py new file mode 100644 index 0000000..92194ce --- /dev/null +++ b/src/fastplms/models/esmfold2/reproducibility.py @@ -0,0 +1,64 @@ +"""Dependency-light RNG scoping for ESMFold2 workflows.""" + +from __future__ import annotations + +import random +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from torch import Tensor + + +@dataclass(frozen=True) +class _RandomState: + python: object + numpy: tuple[Any, ...] + torch_cpu: Tensor + torch_cuda: list[Tensor] | None + + +def _capture_random_state() -> _RandomState: + cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + return _RandomState( + python=random.getstate(), + numpy=np.random.get_state(), + torch_cpu=torch.random.get_rng_state(), + torch_cuda=cuda_state, + ) + + +def _restore_random_state(state: _RandomState) -> None: + random.setstate(state.python) + np.random.set_state(state.numpy) + torch.random.set_rng_state(state.torch_cpu) + if state.torch_cuda is not None: + torch.cuda.set_rng_state_all(state.torch_cuda) + + +@contextmanager +def seed_context(seed: int | None) -> Iterator[None]: + """Seed Python, NumPy, and Torch temporarily, then restore every stream.""" + + if seed is None: + yield + return + if isinstance(seed, bool) or not isinstance(seed, int): + raise TypeError("seed must be None or an integer (excluding bool).") + state = _capture_random_state() + normalized_seed = seed % (2**32) + random.seed(normalized_seed) + np.random.seed(normalized_seed) + torch.manual_seed(normalized_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(normalized_seed) + try: + yield + finally: + _restore_random_state(state) + + +__all__ = ["seed_context"] diff --git a/src/fastplms/models/ttt.py b/src/fastplms/models/ttt.py new file mode 100644 index 0000000..56b4dba --- /dev/null +++ b/src/fastplms/models/ttt.py @@ -0,0 +1,882 @@ +from __future__ import annotations + +import contextlib +import math +import numbers +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections.abc import Iterator, Mapping +from dataclasses import asdict, dataclass, fields +from typing import Any + + +_STANDARD_AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY" +_TTT_SERIALIZATION_VERSION = 1 + + +@dataclass +class TTTConfig: + lr: float = 4e-4 + steps: int = 30 + ags: int = 16 + batch_size: int = 2 + mask_ratio: float = 0.15 + crop_size: int = 1024 + bert_leave_prob: float = 0.1 + bert_replace_prob: float = 0.1 + optimizer: str = "sgd" + momentum: float = 0.0 + weight_decay: float = 0.0 + seed: int | None = 0 + lora_rank: int = 8 + lora_alpha: float = 32.0 + lora_target_replace_module: str | None = None + lora_target_modules: tuple[str, ...] | None = None + initial_state_reset: bool = True + automatic_best_state_reset: bool = False + eval_each_step: bool = False + gradient_clip: bool = False + gradient_clip_max_norm: float = 1.0 + + def __post_init__(self) -> None: + self.verify() + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> TTTConfig: + valid_names = {field.name for field in fields(cls)} + unknown_names = set(kwargs) - valid_names + if unknown_names: + raise ValueError(f"Unknown TTTConfig fields: {sorted(unknown_names)}") + # JSON has no tuple type. Normalize the serialized representation while + # keeping the public constructor and runtime overrides type-strict. + if isinstance(kwargs.get("lora_target_modules"), list): + kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"]) + return cls(**kwargs) + + def merged(self, overrides: Mapping[str, Any] | TTTConfig | None) -> TTTConfig: + if overrides is None: + return self + if isinstance(overrides, TTTConfig): + return overrides + values = {field.name: self.__dict__[field.name] for field in fields(self)} + for name, value in overrides.items(): + if name not in values: + raise ValueError(f"Unknown TTTConfig field: {name}") + values[name] = value + return TTTConfig(**values) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def verify(self) -> None: + numeric_fields = { + "lr": self.lr, + "mask_ratio": self.mask_ratio, + "lora_alpha": self.lora_alpha, + "bert_leave_prob": self.bert_leave_prob, + "bert_replace_prob": self.bert_replace_prob, + "gradient_clip_max_norm": self.gradient_clip_max_norm, + "momentum": self.momentum, + "weight_decay": self.weight_decay, + } + for name, value in numeric_fields.items(): + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise TypeError(f"TTT {name} must be a real number.") + if not math.isfinite(float(value)): + raise ValueError(f"TTT {name} must be finite.") + + integer_fields = { + "steps": self.steps, + "ags": self.ags, + "batch_size": self.batch_size, + "crop_size": self.crop_size, + "lora_rank": self.lora_rank, + } + for name, value in integer_fields.items(): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"TTT {name} must be an integer.") + + if self.seed is not None and ( + isinstance(self.seed, bool) or not isinstance(self.seed, int) + ): + raise TypeError("TTT seed must be None or an integer.") + + boolean_fields = { + "initial_state_reset": self.initial_state_reset, + "automatic_best_state_reset": self.automatic_best_state_reset, + "eval_each_step": self.eval_each_step, + "gradient_clip": self.gradient_clip, + } + for name, value in boolean_fields.items(): + if type(value) is not bool: + raise TypeError(f"TTT {name} must be a boolean.") + + if self.lr <= 0.0: + raise ValueError("TTT learning rate must be positive.") + if self.steps < 1: + raise ValueError("TTT steps must be >= 1.") + if self.ags < 1: + raise ValueError("TTT gradient accumulation steps must be >= 1.") + if self.batch_size < 1: + raise ValueError("TTT batch_size must be >= 1.") + if not 0.0 < self.mask_ratio <= 1.0: + raise ValueError("TTT mask_ratio must be in (0, 1].") + if self.crop_size < 1: + raise ValueError("TTT crop_size must be >= 1.") + if self.lora_rank < 1: + raise ValueError("TTT v1 is LoRA-only, so lora_rank must be >= 1.") + if self.lora_alpha <= 0.0: + raise ValueError("TTT lora_alpha must be positive.") + if not isinstance(self.optimizer, str): + raise TypeError("TTT optimizer must be a string.") + if self.optimizer not in {"adamw", "sgd"}: + raise ValueError("TTT optimizer must be 'adamw' or 'sgd'.") + if self.momentum < 0.0: + raise ValueError("TTT momentum must be non-negative.") + if self.weight_decay < 0.0: + raise ValueError("TTT weight_decay must be non-negative.") + if not 0.0 <= self.bert_leave_prob <= 1.0: + raise ValueError("bert_leave_prob must be in [0, 1].") + if not 0.0 <= self.bert_replace_prob <= 1.0: + raise ValueError("bert_replace_prob must be in [0, 1].") + if self.bert_leave_prob + self.bert_replace_prob > 1.0: + raise ValueError("bert_leave_prob + bert_replace_prob must be <= 1.") + if self.gradient_clip and self.gradient_clip_max_norm <= 0.0: + raise ValueError("gradient_clip_max_norm must be positive.") + if self.lora_target_replace_module is not None: + if not isinstance(self.lora_target_replace_module, str): + raise TypeError("lora_target_replace_module must be None or a string.") + if not self.lora_target_replace_module.strip(): + raise ValueError("lora_target_replace_module must not be empty.") + if self.lora_target_modules is not None: + if not isinstance(self.lora_target_modules, tuple): + raise TypeError("lora_target_modules must be None or a tuple of strings.") + if not self.lora_target_modules: + raise ValueError("lora_target_modules must not be empty.") + if any(not isinstance(name, str) for name in self.lora_target_modules): + raise TypeError("lora_target_modules must contain only strings.") + if any(not name.strip() for name in self.lora_target_modules): + raise ValueError( + "lora_target_modules must contain only non-empty strings." + ) + if len(set(self.lora_target_modules)) != len(self.lora_target_modules): + raise ValueError("lora_target_modules must not contain duplicates.") + + +class LoraInjectedLinear(nn.Module): + """ProteinTTT-compatible low-rank adapter. + + ``alpha`` is the direct adapter-output multiplier used by the pinned + ProteinTTT ``inject_trainable_lora(..., scale=lora_alpha)`` contract. It + is intentionally not divided by ``rank`` as it would be in the common + PEFT LoRA convention. + """ + + def __init__( + self, + linear: nn.Module, + rank: int, + alpha: float, + generator: torch.Generator | None = None, + ) -> None: + super().__init__() + weight = linear._parameters.get("weight") + if not isinstance(weight, torch.Tensor): + raise TypeError("LoRA targets must expose a tensor weight parameter.") + if weight.ndim != 2: + raise ValueError("LoRA can only wrap 2D linear weights.") + self.linear = linear + self.linear.requires_grad_(False) + self.rank = rank + # ProteinTTT names this setting ``lora_alpha`` but passes it directly + # to cloneofsimo/lora's ``scale`` argument. Preserve that numerical + # contract for parity and for saved FastPLMs TTT configurations. + self.scale = alpha + in_features = weight.shape[1] + out_features = weight.shape[0] + # ``nn.Linear`` initializes from the process-global CPU generator. Preserve + # that state when TTT supplies its own generator so lazy adapter injection + # is reproducible without perturbing the caller's RNG stream. + with torch.random.fork_rng(devices=[], enabled=generator is not None): + self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32) + self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32) + nn.init.normal_(self.lora_down.weight, std=1.0 / rank, generator=generator) + nn.init.zeros_(self.lora_up.weight) + self.lora_down.to(device=weight.device) + self.lora_up.to(device=weight.device) + self.register_buffer( + "_ttt_initial_lora_down", + self.lora_down.weight.detach().clone(), + persistent=True, + ) + self.register_buffer( + "_ttt_initial_lora_up", + self.lora_up.weight.detach().clone(), + persistent=True, + ) + + @property + def weight(self) -> torch.Tensor: + return self.linear._parameters["weight"] + + @property + def bias(self) -> torch.Tensor | None: + return self.linear._parameters["bias"] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (..., d_in) + base = self.linear(x) # (..., d_out) + delta = ( # (..., d_out) + self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale + ) + return base + delta.to(dtype=base.dtype) # (..., d_out) + + def reset_lora_parameters(self) -> None: + with torch.no_grad(): + self.lora_down.weight.copy_(self._ttt_initial_lora_down) + self.lora_up.weight.copy_(self._ttt_initial_lora_up) + + +class FastPLMTestTimeTrainingMixin: + def init_ttt(self, ttt_config: TTTConfig | Mapping[str, Any] | None = None) -> None: + base_config = self.__dict__.get("_ttt_cfg") + if base_config is None: + base_config = TTTConfig() + if not isinstance(base_config, TTTConfig): + raise TypeError("Existing TTT configuration must be a TTTConfig instance.") + configured = base_config.merged(ttt_config) + serialized = getattr(getattr(self, "config", None), "fastplms_ttt", None) + serialized_initialized = False + if serialized is not None: + if not isinstance(serialized, Mapping): + raise ValueError("config.fastplms_ttt must be a mapping.") + version = serialized.get("version") + if version != _TTT_SERIALIZATION_VERSION: + raise ValueError( + "Unsupported FastPLMs TTT serialization version " + f"{version!r}; expected {_TTT_SERIALIZATION_VERSION}." + ) + serialized_config = serialized.get("config") + if not isinstance(serialized_config, Mapping): + raise ValueError("Serialized FastPLMs TTT state is missing its config mapping.") + configured = TTTConfig.from_kwargs(**dict(serialized_config)) + initialized_value = serialized.get("initialized", False) + if type(initialized_value) is not bool: + raise ValueError("Serialized FastPLMs TTT initialized flag must be a boolean.") + serialized_initialized = initialized_value + + self._ttt_cfg = configured + self._ttt_cfg.verify() + self._ttt_initialized = False + if serialized_initialized: + self._ttt_inject_lora() + self._ttt_initialized = True + + @property + def ttt_config(self) -> TTTConfig: + if "_ttt_cfg" not in self.__dict__: + self.init_ttt() + return self._ttt_cfg + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self] + + def _ttt_get_frozen_modules(self) -> list[nn.Module]: + return [] + + def _ttt_tokenize( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | dict[str, torch.Tensor]: + del kwargs + if input_ids is not None: + return input_ids # (b, l) + if seq is None: + raise ValueError("Pass either seq or input_ids for TTT.") + tokenized = self.tokenizer(seq, return_tensors="pt", padding=True) + return tokenized["input_ids"] # (b, l) + + def _ttt_mask_token(self) -> int: + return int(self.tokenizer.mask_token_id) + + def _ttt_padding_token(self) -> int: + return int(self.tokenizer.pad_token_id) + + def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + # input_ids: (b, l) + tokenizer = self.tokenizer + special_ids = set(tokenizer.all_special_ids) + vocab_size = int(self.config.vocab_size) + unknown_id = getattr(tokenizer, "unk_token_id", None) + if unknown_id is not None: + special_ids.add(int(unknown_id)) + + vocab: Mapping[str, Any] = {} + get_vocab = getattr(tokenizer, "get_vocab", None) + if callable(get_vocab): + vocab = get_vocab() + elif isinstance(getattr(tokenizer, "vocab", None), Mapping): + vocab = tokenizer.vocab + elif isinstance(getattr(tokenizer, "_token_to_id", None), Mapping): + vocab = tokenizer._token_to_id + + ids: list[int] = [] + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + for amino_acid in _STANDARD_AMINO_ACIDS: + token_id = convert(amino_acid) if callable(convert) else vocab.get(amino_acid) + if ( + isinstance(token_id, int) + and 0 <= token_id < vocab_size + and token_id not in special_ids + and token_id not in ids + ): + ids.append(token_id) + if not ids: + raise ValueError( + "TTT could not resolve any canonical amino-acid token IDs from the tokenizer; " + "refusing to sample arbitrary or reserved vocabulary entries." + ) + return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) # (c_aa,) + + def _ttt_predict_logits( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + if isinstance(batch, dict): + output = self(**batch) + return output.logits # (b, l, c) + attention_mask = batch.ne(self._ttt_padding_token()) # (b, l) + output = self(input_ids=batch, attention_mask=attention_mask) + return output.logits # (b, l, c) + + def _ttt_eval_step( + self, + step: int, + loss: float, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> tuple[dict[str, Any], float | None]: + del step, loss, seq, input_ids, kwargs + return {}, None + + def _ttt_is_lora_target( + self, + name: str, + full_name: str, + module: nn.Module, + active: bool, + target_modules: tuple[str, ...] | None, + ) -> bool: + if not active: + return False + if isinstance(module, LoraInjectedLinear): + return False + if ( + target_modules is not None + and name not in target_modules + and full_name not in target_modules + ): + return False + if isinstance(module, nn.Linear): + return True + if "weight" not in module._parameters: + return False + weight = module._parameters["weight"] + if weight is None or weight.ndim != 2: + return False + return "Linear" in module.__class__.__name__ + + def _ttt_inject_lora(self) -> int: + cfg = self.ttt_config + cfg.verify() + target_class = cfg.lora_target_replace_module + target_modules = cfg.lora_target_modules + wrapped = 0 + generator = None + if cfg.seed is not None: + generator = torch.Generator(device="cpu") + generator.manual_seed(cfg.seed) + + def inject(module: nn.Module, prefix: str, active: bool) -> None: + nonlocal wrapped + for name, child in list(module.named_children()): + full_name = f"{prefix}.{name}" if prefix else name + child_active = active + if target_class is not None: + child_active = active or child.__class__.__name__ == target_class + if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules): + setattr( + module, + name, + LoraInjectedLinear( + child, + rank=cfg.lora_rank, + alpha=cfg.lora_alpha, + generator=generator, + ), + ) + wrapped += 1 + continue + inject(child, full_name, child_active) + + for trainable_module in self._ttt_get_trainable_modules(): + inject(trainable_module, "", target_class is None) + if wrapped == 0: + raise ValueError("TTT LoRA injection did not find any target modules.") + return wrapped + + def _ttt_lora_modules(self) -> list[LoraInjectedLinear]: + return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)] + + def _ttt_lora_parameters(self) -> list[nn.Parameter]: + params: list[nn.Parameter] = [] + for module in self._ttt_lora_modules(): + params.extend(module.lora_down.parameters()) + params.extend(module.lora_up.parameters()) + if not params: + raise RuntimeError("TTT has no LoRA parameters.") + return params + + def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]: + snapshot = [] + for module in self._ttt_lora_modules(): + snapshot.append( + { + "lora_down.weight": module.lora_down.weight.detach().clone(), # (r, d_in) + "lora_up.weight": module.lora_up.weight.detach().clone(), # (d_out, r) + } + ) + if not snapshot: + raise RuntimeError("TTT has no LoRA state to snapshot.") + return snapshot + + def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None: + modules = self._ttt_lora_modules() + if len(modules) != len(state): + raise RuntimeError("TTT LoRA state/module count mismatch.") + with torch.no_grad(): + for module, module_state in zip(modules, state, strict=True): + module.lora_down.weight.copy_(module_state["lora_down.weight"]) + module.lora_up.weight.copy_(module_state["lora_up.weight"]) + + def _ttt_ensure_initialized(self) -> None: + if "_ttt_cfg" not in self.__dict__: + self.init_ttt() + if self._ttt_initialized: + return + self._ttt_inject_lora() + self._ttt_initialized = True + + def ttt_reset(self) -> None: + self._ttt_ensure_initialized() + for module in self._ttt_lora_modules(): + module.reset_lora_parameters() + + def _ttt_serialized_contract(self) -> dict[str, Any]: + return { + "version": _TTT_SERIALIZATION_VERSION, + "initialized": bool(self._ttt_initialized), + "config": self.ttt_config.to_dict(), + } + + def save_pretrained(self, save_directory: Any, *args: Any, **kwargs: Any) -> Any: + """Save initialized adapters, their reset baseline, and the TTT config. + + Adapter injection changes the module tree, so the serialized config must + reconstruct that tree before Transformers loads the state dict. Models + whose own state-dict hooks omit their trainable TTT modules fail closed + instead of producing an artifact that cannot restore the adaptation. + """ + + if self._ttt_initialized: + state_keys = set(self.state_dict()) + missing_adapter_keys = [ + name + for name, _ in self.named_parameters() + if ".lora_" in name and name not in state_keys + ] + if missing_adapter_keys: + raise RuntimeError( + "This model attaches TTT adapters to transient modules that its " + "checkpoint excludes, so save_pretrained cannot persist the adapted " + "state safely. Reset the model or use a model-specific adapter export." + ) + self.config.fastplms_ttt = self._ttt_serialized_contract() + return super().save_pretrained(save_directory, *args, **kwargs) + + def _ttt_make_optimizer(self) -> torch.optim.Optimizer: + cfg = self.ttt_config + params = self._ttt_lora_parameters() + if cfg.optimizer == "sgd": + return torch.optim.SGD( + params, + lr=cfg.lr, + momentum=cfg.momentum, + weight_decay=cfg.weight_decay, + ) + return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay) + + def _ttt_to_device( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + device: torch.device, + ) -> torch.Tensor | dict[str, torch.Tensor]: + if isinstance(batch, dict): + return {name: tensor.to(device) for name, tensor in batch.items()} # unchanged shapes + return batch.to(device) # unchanged shape + + def _ttt_input_ids_from_batch( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + ) -> torch.Tensor: + if isinstance(batch, dict): + return batch["input_ids"] # (b, l) + return batch # (b, l) + + def _ttt_set_input_ids( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + input_ids: torch.Tensor, + ) -> torch.Tensor | dict[str, torch.Tensor]: + if isinstance(batch, dict): + updated = dict(batch) + updated["input_ids"] = input_ids # (b, l) + return updated + return input_ids # (b, l) + + def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: + # input_ids: (b, l) + residue_ids = self._ttt_replacement_tokens(input_ids) # (c_aa,) + return torch.isin(input_ids, residue_ids) # (b, l) + + def _ttt_validate_tokenized_batch( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + ) -> None: + input_ids = self._ttt_input_ids_from_batch(batch) + if input_ids.ndim != 2 or input_ids.shape[0] == 0 or input_ids.shape[1] == 0: + raise ValueError( + "TTT input_ids must have non-empty shape (batch, sequence); got " + f"{tuple(input_ids.shape)}." + ) + + if str(getattr(self.config, "model_type", "")) == "dplm2": + tokenizer = self.tokenizer + token_to_id = getattr(tokenizer, "_token_to_id", {}) + struct_cls_token = getattr(tokenizer, "struct_cls_token", None) + struct_boundary = token_to_id.get(struct_cls_token) + if struct_boundary is None: + raise ValueError( + "DPLM2 TTT could not resolve the structure-token boundary safely." + ) + pad_token = self._ttt_padding_token() + generic_aa_special_ids = torch.tensor( # (4,) + [int(self.config.vocab_size) + offset for offset in range(4)], + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_structure = input_ids.ge(int(struct_boundary)) & input_ids.ne( # (b, l) + pad_token + ) + is_structure &= ~torch.isin(input_ids, generic_aa_special_ids) + if bool(is_structure.any()): + raise ValueError( + "DPLM2 TTT currently supports amino-acid-only inputs. Packed or " + "structure-token inputs require a modality-specific corruption objective." + ) + + if isinstance(batch, dict) and "type_ids" in batch: + type_ids = batch["type_ids"] # (b, l) + attention_mask = batch.get( # (b, l) + "attention_mask", + input_ids.ne(pad_token), + ).bool() + if bool(((type_ids == int(self.config.struct_type)) & attention_mask).any()): + raise ValueError( + "DPLM2 TTT currently supports amino-acid-only inputs; structure " + "type_ids are not accepted." + ) + + if not bool(self._ttt_non_special_mask(input_ids).any()): + raise ValueError( + "TTT input contains no trainable biological residue tokens after excluding " + "padding, boundary, mask, and reserved tokens." + ) + + def _ttt_sample_crop( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + generator: torch.Generator, + ) -> torch.Tensor | dict[str, torch.Tensor]: + input_ids = self._ttt_input_ids_from_batch(batch) + cfg = self.ttt_config + if input_ids.shape[1] <= cfg.crop_size: + return batch + position_has_residue = ( # (l,) + self._ttt_non_special_mask(input_ids).any(dim=0).to(torch.int64) + ) + prefix = F.pad(position_has_residue.cumsum(dim=0), (1, 0)) # (l + 1,) + window_counts = prefix[cfg.crop_size :] - prefix[: -cfg.crop_size] # (l-crop+1,) + valid_starts = torch.where(window_counts > 0)[0] # (n_valid,) + if valid_starts.numel() == 0: + raise ValueError("TTT could not find a crop containing a biological residue token.") + selected = torch.randint( # (1,) + valid_starts.numel(), + (1,), + generator=generator, + device=input_ids.device, + ) + start = int(valid_starts[selected].item()) + end = start + cfg.crop_size + if isinstance(batch, dict): + cropped = {} + for name, tensor in batch.items(): + if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]: + cropped[name] = tensor[:, start:end] # (b, crop_size, ...) + else: + cropped[name] = tensor + return cropped + return input_ids[:, start:end] # (b, crop_size) + + def _ttt_sample_batch( + self, + tokenized: torch.Tensor | dict[str, torch.Tensor], + generator: torch.Generator, + ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]: + cfg = self.ttt_config + batch = self._ttt_sample_crop(tokenized, generator) + input_ids = self._ttt_input_ids_from_batch(batch) # (b, l) + row_has_residue = self._ttt_non_special_mask(input_ids).any(dim=1) # (b,) + eligible_rows = torch.where(row_has_residue)[0] # (n_eligible,) + if eligible_rows.numel() == 0: + raise ValueError( + "TTT sampled batch contains no trainable biological residue tokens." + ) + sampled_row_indices = torch.randint( # (b_sample,) + eligible_rows.numel(), + (cfg.batch_size,), + generator=generator, + device=input_ids.device, + ) + rows = eligible_rows[sampled_row_indices] # (b_sample,) + if isinstance(batch, dict): + sampled: torch.Tensor | dict[str, torch.Tensor] = {} + for name, tensor in batch.items(): + if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]: + sampled[name] = tensor.index_select(0, rows) # (b_sample, ...) + else: + sampled[name] = tensor + else: + sampled = input_ids.index_select(0, rows) # (b_sample, l) + + sampled_ids = self._ttt_input_ids_from_batch(sampled) # (b_sample, l) + labels = sampled_ids.clone() # (b_sample, l) + non_special = self._ttt_non_special_mask(sampled_ids) # (b_sample, l) + label_mask = torch.zeros_like(non_special) # (b_sample, l) + for row_idx in range(sampled_ids.shape[0]): + candidate_positions = torch.where(non_special[row_idx])[0] # (n_candidates,) + if candidate_positions.numel() == 0: + continue + num_mask = max(1, round(candidate_positions.numel() * cfg.mask_ratio)) + order = torch.randperm( # (n_candidates,) + candidate_positions.numel(), + generator=generator, + device=sampled_ids.device, + ) + chosen = candidate_positions[order[:num_mask]] # (n_mask,) + label_mask[row_idx, chosen] = True + labels = labels.masked_fill(~label_mask, -100) # (b_sample, l) + + masked_ids = sampled_ids.clone() # (b_sample, l) + chosen_positions = torch.where(label_mask) # two (n_chosen,) tensors + if chosen_positions[0].numel() > 0: + random_values = torch.rand( # (n_chosen,) + chosen_positions[0].shape, + generator=generator, + device=sampled_ids.device, + ) + leave = random_values < cfg.bert_leave_prob # (n_chosen,) + replace = (random_values >= cfg.bert_leave_prob) & ( # (n_chosen,) + random_values < cfg.bert_leave_prob + cfg.bert_replace_prob + ) + mask = ~(leave | replace) # (n_chosen,) + if mask.any(): + masked_ids[ + chosen_positions[0][mask], + chosen_positions[1][mask], + ] = self._ttt_mask_token() + if replace.any(): + replacement_tokens = self._ttt_replacement_tokens(sampled_ids) # (c_aa,) + replacement_idx = torch.randint( # (n_replace,) + replacement_tokens.shape[0], + (int(replace.sum().item()),), + generator=generator, + device=sampled_ids.device, + ) + masked_ids[ + chosen_positions[0][replace], + chosen_positions[1][replace], + ] = replacement_tokens[replacement_idx] + + return self._ttt_set_input_ids(sampled, masked_ids), labels # batch, (b_sample, l) + + @contextlib.contextmanager + def _ttt_seed_scope(self, seed: int | None) -> Iterator[None]: + if seed is None: + yield + return + cuda_devices = sorted( + { + parameter.device.index + for parameter in self.parameters() + if parameter.device.type == "cuda" and parameter.device.index is not None + } + ) + with torch.random.fork_rng(devices=cuda_devices): + torch.random.default_generator.manual_seed(seed) + for device_index in cuda_devices: + with torch.cuda.device(device_index): + torch.cuda.manual_seed(seed) + yield + + def ttt( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + ttt_config: TTTConfig | Mapping[str, Any] | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + if ttt_config is not None: + if "_ttt_initialized" in self.__dict__ and self._ttt_initialized: + next_cfg = self.ttt_config.merged(ttt_config) + current_cfg = self.ttt_config + if next_cfg.lora_rank != current_cfg.lora_rank: + raise ValueError( + "Changing lora_rank after TTT initialization is not supported." + ) + if next_cfg.lora_alpha != current_cfg.lora_alpha: + raise ValueError( + "Changing lora_alpha after TTT initialization is not supported." + ) + if ( + next_cfg.lora_target_replace_module + != current_cfg.lora_target_replace_module + ): + raise ValueError( + "Changing LoRA target class after TTT initialization is not supported." + ) + if next_cfg.lora_target_modules != current_cfg.lora_target_modules: + raise ValueError( + "Changing LoRA target modules after TTT initialization is not supported." + ) + self._ttt_cfg = next_cfg + else: + # Family constructors preconfigure the attention class that may + # receive LoRA adapters. A first-call mapping changes only the + # requested fields; rebuilding from TTTConfig defaults here + # would erase that family target immediately before injection. + self._ttt_cfg = self.ttt_config.merged(ttt_config) + self._ttt_cfg.verify() + + cfg = self.ttt_config + device = next(self.parameters()).device + tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs) + tokenized = self._ttt_to_device(tokenized, device) + self._ttt_validate_tokenized_batch(tokenized) + self._ttt_ensure_initialized() + if cfg.initial_state_reset: + self.ttt_reset() + + generator_device = device if device.type == "cuda" else torch.device("cpu") + generator = torch.Generator(device=generator_device) + if cfg.seed is not None: + generator.manual_seed(cfg.seed) + + module_modes = {module: module.training for module in self.modules()} + requires_grad = {param: param.requires_grad for param in self.parameters()} + losses: list[float] = [] + step_metrics: list[dict[str, Any]] = [] + best_state: list[dict[str, torch.Tensor]] | None = None + best_metric: float | None = None + best_step = 0 + + with self._ttt_seed_scope(cfg.seed): + try: + self.train() + for param in self.parameters(): + param.requires_grad_(False) + for param in self._ttt_lora_parameters(): + param.requires_grad_(True) + + optimizer = self._ttt_make_optimizer() + optimizer.zero_grad(set_to_none=True) + total_micro_steps = cfg.steps * cfg.ags + for micro_step in range(total_micro_steps): + batch, labels = self._ttt_sample_batch( # batch, (b_sample, l) + tokenized, + generator, + ) + if not bool(labels.ne(-100).any()): + raise RuntimeError( + "TTT produced an all-ignored label batch; refusing a NaN update." + ) + logits = self._ttt_predict_logits(batch, **kwargs) # (b_sample, l, c) + labels = labels.to(device=logits.device) # (b_sample, l) + loss = F.cross_entropy( # () + logits.reshape(-1, logits.shape[-1]), + labels.reshape(-1), + ignore_index=-100, + ) + if not bool(torch.isfinite(loss)): + raise FloatingPointError( + f"TTT loss is non-finite at micro-step {micro_step + 1}." + ) + (loss / cfg.ags).backward() + if (micro_step + 1) % cfg.ags != 0: + continue + + if cfg.gradient_clip: + torch.nn.utils.clip_grad_norm_( + self._ttt_lora_parameters(), + cfg.gradient_clip_max_norm, + ) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + step = (micro_step + 1) // cfg.ags + loss_value = float(loss.detach().item()) + losses.append(loss_value) + if cfg.eval_each_step: + metrics, metric = self._ttt_eval_step( + step=step, + loss=loss_value, + seq=seq, + input_ids=input_ids, + **kwargs, + ) + if len(metrics) > 0: + step_metrics.append(metrics) + if metric is not None and (best_metric is None or metric > best_metric): + best_metric = metric + best_step = step + best_state = self._ttt_snapshot_lora_state() + + if cfg.automatic_best_state_reset and best_state is not None: + self._ttt_restore_lora_state(best_state) + finally: + for param, value in requires_grad.items(): + param.requires_grad_(value) + for module, training in module_modes.items(): + module.train(training) + + return { + "losses": losses, + "step_metrics": step_metrics, + "best_step": best_step, + "best_metric": best_metric, + } diff --git a/src/fastplms/registry.py b/src/fastplms/registry.py new file mode 100644 index 0000000..6d2e86f --- /dev/null +++ b/src/fastplms/registry.py @@ -0,0 +1,1487 @@ +"""Typed access to the FastPLMs model and provenance manifest. + +The registry is intentionally independent of Torch and Transformers. Tooling can +therefore inspect supported checkpoints, licenses, and reference sources without +initializing a model runtime or downloading any files. +""" + +from __future__ import annotations + +import re +import tomllib +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from functools import lru_cache +from importlib import resources +from pathlib import Path, PurePosixPath, PureWindowsPath +from types import MappingProxyType +from typing import Any, Literal, cast +from urllib.parse import urlparse + + +_HEX_RE = re.compile(r"^[0-9a-f]+$") +_IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$") +_HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+") +_WINDOWS_INVALID_PATH_CHARACTERS = frozenset('<>:"|?*') +_WINDOWS_RESERVED_PATH_NAMES = frozenset( + {"AUX", "CON", "NUL", "PRN"} + | {f"COM{index}" for index in range(1, 10)} + | {f"LPT{index}" for index in range(1, 10)} +) +_REPOSITORY_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$") +_REFERENCE_CONTAINER_RE = re.compile(r"^reference-[a-z0-9]+(?:-[a-z0-9]+)*$") +_REFERENCE_ADAPTER_RE = re.compile( + r"^tests\.parity\.support\.reference_adapters\.[a-z_][a-z0-9_]*$" +) +_DOCUMENTATION_FRAGMENT_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") +_ALLOWED_ATTENTION = frozenset( + {"eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"} +) +_ALLOWED_DTYPES = frozenset({"float32", "bfloat16"}) +_ALLOWED_PRECISIONS = frozenset({"default", "auto", "fp32", "bf16", "fp8"}) +_ALLOWED_BF16_EXECUTIONS = frozenset({"static_parameters", "fp32_parameters_autocast"}) +HUB_LICENSE_IDENTIFIERS = frozenset({"mit", "apache-2.0", "cc-by-nc-sa-4.0", "other"}) +_ALLOWED_TOKENIZER_MODES = frozenset({"tokenizer", "sequence", "structure"}) +_ALLOWED_SIZE_CATEGORIES = frozenset({"small", "medium", "large", "xlarge", "structure"}) +RuntimeExtra = Literal["core", "structure"] +TestTier = Literal["check", "compliance", "structure", "feature", "artifact", "benchmark"] +VramTier = Literal["sequence", "large-sequence", "structure", "structure-6b"] +GenerationContract = Literal["not_applicable", "required", "official_unavailable"] +RuntimeAssetTrustKind = Literal["hash_pinned_pickle"] +Bf16Execution = Literal["static_parameters", "fp32_parameters_autocast"] +DtypeName = Literal["float32", "bfloat16"] +_ALLOWED_EXTRAS = frozenset({"core", "structure"}) +_ALLOWED_TEST_TIERS = frozenset( + {"check", "compliance", "structure", "feature", "artifact", "benchmark"} +) +_ALLOWED_VRAM_TIERS = frozenset({"sequence", "large-sequence", "structure", "structure-6b"}) +_ALLOWED_GENERATION_CONTRACTS = frozenset({"not_applicable", "required", "official_unavailable"}) +_ALLOWED_RUNTIME_ASSET_TRUST_KINDS = frozenset({"hash_pinned_pickle"}) +_ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS = frozenset({"requires_cached_verified_file"}) +_ALLOWED_AUTO_CLASSES = frozenset( + { + "AutoConfig", + "AutoModel", + "AutoModelForMaskedLM", + "AutoModelForProteinFolding", + "AutoModelForSequenceClassification", + "AutoModelForSeq2SeqLM", + "AutoModelForTokenClassification", + } +) +_WEIGHT_SUFFIXES = (".bin", ".ckpt", ".pt", ".pth", ".safetensors") +_ALLOWED_ORACLE_ASSET_ROLES = frozenset({"weights", "contact_regression"}) +_FAIR_ESM_ASSET_HOST = "dl.fbaipublicfiles.com" +_ROOT_FIELDS = frozenset( + { + "schema_version", + "legal_files", + "attention_kernels", + "upstreams", + "families", + "models", + "runtime_assets", + } +) +_UPSTREAM_FIELDS = frozenset( + { + "id", + "path", + "url", + "revision", + "license", + "license_files", + "license_digests", + "distribution_files", + } +) +_FAMILY_FIELDS = frozenset( + { + "architecture", + "upstreams", + "tokenizer_mode", + "public_input", + "extra", + "reference_container", + "reference_adapter", + "attention", + "dtypes", + "bf16_execution", + "precisions", + "experimental_precisions", + "vram_tier", + "checkpoint_license", + "hub_license", + "hub_license_name", + "hub_license_link", + "state_transform", + "conversion_provenance", + "representative", + "documentation", + "test_tiers", + "runtime_paths", + "requires_complete_weight_publication", + "weights_publication_allowed", + "auto_map", + "tokenizer_class", + "backbone_model", + } +) +_MODEL_FIELDS = frozenset( + { + "id", + "family", + "size_category", + "generation_contract", + "fast_repo", + "fast_revision", + "fast_files", + "fast_unresolved_files", + "official_repo", + "official_revision", + "official_files", + "official_unresolved_files", + "oracle_assets", + "official_golden", + "artifact_source", + "canonical_state_sha256", + "tokenizer_source", + "auto_map", + "notes", + "msa_conditioning", + } +) +_RUNTIME_ASSET_FIELDS = frozenset( + { + "id", + "repository", + "revision", + "path", + "sha256", + "size", + "consumer_family", + "trust_kind", + "license", + "offline_behavior", + } +) + + +class RegistryError(ValueError): + """Raised when the model manifest is incomplete or internally inconsistent.""" + + +def _portable_relative_path(value: str, context: str) -> PurePosixPath: + """Return one normalized cross-platform relative path or fail closed.""" + + posix = PurePosixPath(value) + windows = PureWindowsPath(value) + unsafe_windows_part = any( + part.rstrip(" .") != part + or part.split(".", maxsplit=1)[0].upper() in _WINDOWS_RESERVED_PATH_NAMES + or any( + ord(character) < 32 or character in _WINDOWS_INVALID_PATH_CHARACTERS + for character in part + ) + for part in posix.parts + ) + if ( + not value + or not posix.parts + or posix == PurePosixPath(".") + or posix.is_absolute() + or windows.is_absolute() + or windows.drive + or "\\" in value + or "." in posix.parts + or ".." in posix.parts + or value != posix.as_posix() + or any( + part.lower() in {".git", ".cache", "__pycache__"} + for part in posix.parts + ) + or unsafe_windows_part + ): + raise RegistryError(f"{context} is not portable: {value!r}") + return posix + + +@dataclass(frozen=True, slots=True) +class FileDigest: + """Expected content identity for one pinned file.""" + + path: str + algorithm: str + digest: str + + @classmethod + def parse(cls, value: str) -> FileDigest: + try: + path, encoded_digest = value.split("=", maxsplit=1) + algorithm, digest = encoded_digest.split(":", maxsplit=1) + except ValueError as error: + raise RegistryError("File digests must use '=:'.") from error + + _portable_relative_path(path, "Checkpoint file path") + + expected_length = {"git-sha1": 40, "sha256": 64}.get(algorithm) + if expected_length is None: + raise RegistryError(f"Unsupported file digest algorithm: {algorithm!r}") + if len(digest) != expected_length or _HEX_RE.fullmatch(digest) is None: + raise RegistryError(f"Invalid {algorithm} digest for {path!r}: {digest!r}") + return cls(path=path, algorithm=algorithm, digest=digest) + + @property + def encoded(self) -> str: + return f"{self.algorithm}:{self.digest}" + + +@dataclass(frozen=True, slots=True) +class CheckpointSource: + """One immutable Hugging Face repository snapshot.""" + + repo_id: str + revision: str + files: tuple[FileDigest, ...] + unresolved_files: tuple[str, ...] = () + + @property + def file_map(self) -> Mapping[str, FileDigest]: + return MappingProxyType({item.path: item for item in self.files}) + + +@dataclass(frozen=True, slots=True) +class OracleAsset: + """Hash-pinned external file required by a native parity oracle.""" + + role: str + path: str + url: str + sha256: str + size: int + + +@dataclass(frozen=True, slots=True) +class RuntimeAsset: + """Immutable runtime data with an explicit deserialization trust boundary.""" + + id: str + repository: str + revision: str + path: str + sha256: str + size: int + consumer_family: str + trust_kind: RuntimeAssetTrustKind + license_expression: str + offline_behavior: str + + +@dataclass(frozen=True, slots=True) +class OfficialGolden: + """Hash-pinned official output bundle required by the check tier.""" + + metadata: FileDigest + tensors: FileDigest + + +@dataclass(frozen=True, slots=True) +class UpstreamSource: + """Pinned official implementation used as a parity oracle.""" + + id: str + path: str + url: str + revision: str + license_expression: str + license_files: tuple[str, ...] + license_digests: tuple[FileDigest, ...] = () + distribution_files: tuple[FileDigest, ...] = () + + +@dataclass(frozen=True, slots=True) +class AttentionKernelSpec: + """Immutable Hugging Face kernel used by one attention backend.""" + + implementation: str + repository: str + revision: str + version: int + expected_variant: str + dtypes: tuple[DtypeName, ...] + + +@dataclass(frozen=True, slots=True) +class ModelFamily: + """Shared runtime and compliance contract for one architecture family.""" + + id: str + architecture: str + upstreams: tuple[str, ...] + tokenizer_mode: str + public_input: str + extra: RuntimeExtra + reference_container: str + reference_adapter: str + attention: tuple[str, ...] + dtypes: tuple[DtypeName, ...] + bf16_execution: Bf16Execution + precisions: tuple[str, ...] + vram_tier: VramTier + checkpoint_license: str + hub_license: str + state_transform: str + representative: str + documentation: str + test_tiers: tuple[TestTier, ...] + runtime_paths: tuple[str, ...] + auto_map_items: tuple[tuple[str, str], ...] + requires_complete_weight_publication: bool = False + weights_publication_allowed: bool = False + experimental_precisions: tuple[str, ...] = () + tokenizer_class: str | None = None + hub_license_name: str | None = None + hub_license_link: str | None = None + conversion_provenance: str = "" + backbone_model: str | None = None + + @property + def auto_map(self) -> Mapping[str, str]: + return MappingProxyType(dict(self.auto_map_items)) + + @property + def hub_license_metadata(self) -> Mapping[str, str]: + """Return valid Hugging Face model-card license fields.""" + + metadata = {"license": self.hub_license} + if self.hub_license_name is not None: + # Hugging Face validates custom license names as lowercase slugs, + # while the manifest retains the reader-facing display name used + # in generated prose. + metadata["license_name"] = _HUB_LICENSE_NAME_RE.sub( + "-", + self.hub_license_name.lower(), + ).strip("-.") + if self.hub_license_link is not None: + metadata["license_link"] = self.hub_license_link + return MappingProxyType(metadata) + + @property + def stable_precisions(self) -> tuple[str, ...]: + """Return precision policies covered by the release contract.""" + + experimental = set(self.experimental_precisions) + return tuple(precision for precision in self.precisions if precision not in experimental) + + +@dataclass(frozen=True, slots=True) +class ModelSpec: + """Complete immutable source and runtime contract for one checkpoint.""" + + id: str + family: ModelFamily + fast: CheckpointSource + official: CheckpointSource + size_category: str + generation_contract: GenerationContract = "not_applicable" + oracle_assets: tuple[OracleAsset, ...] = () + official_golden: OfficialGolden | None = None + artifact_source: str = "fast" + canonical_state_sha256: str | None = None + tokenizer_source_id: str | None = None + auto_map_items: tuple[tuple[str, str], ...] = () + notes: str = "" + msa_conditioning: bool | None = None + + @property + def is_deep_reference(self) -> bool: + return self.id == self.family.representative + + @property + def auto_map(self) -> Mapping[str, str]: + if self.auto_map_items: + return MappingProxyType(dict(self.auto_map_items)) + return self.family.auto_map + + @property + def artifact_checkpoint(self) -> CheckpointSource: + """Return the checkpoint selected for local artifact construction.""" + + return self.fast if self.artifact_source == "fast" else self.official + + @property + def oracle_asset_map(self) -> Mapping[str, OracleAsset]: + """Return native oracle assets keyed by their declared role.""" + + return MappingProxyType({asset.role: asset for asset in self.oracle_assets}) + + +class ModelRegistry(Mapping[str, ModelSpec]): + """Validated mapping of model IDs to typed model specifications.""" + + def __init__( + self, + *, + schema_version: int, + upstreams: Mapping[str, UpstreamSource], + families: Mapping[str, ModelFamily], + models: Mapping[str, ModelSpec], + runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}), + attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}), + legal_files: tuple[FileDigest, ...] = (), + ) -> None: + self.schema_version = schema_version + self.upstreams = MappingProxyType(dict(upstreams)) + self.attention_kernels = MappingProxyType(dict(attention_kernels)) + self.families = MappingProxyType(dict(families)) + self._models = MappingProxyType(dict(models)) + self.runtime_assets = MappingProxyType(dict(runtime_assets)) + self.legal_files = legal_files + + def __getitem__(self, key: str) -> ModelSpec: + return self._models[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._models) + + def __len__(self) -> int: + return len(self._models) + + def by_family(self, family_id: str) -> tuple[ModelSpec, ...]: + if family_id not in self.families: + raise KeyError(family_id) + return tuple(model for model in self._models.values() if model.family.id == family_id) + + def supported_attention_dtypes( + self, + family_id: str, + implementation: str, + ) -> tuple[DtypeName, ...]: + """Return manifest-supported dtypes for one family/backend pair.""" + + family = self.families[family_id] + if implementation not in family.attention: + raise KeyError( + f"Family {family_id!r} does not advertise attention backend " + f"{implementation!r}." + ) + kernel = self.attention_kernels.get(implementation) + if kernel is None: + return family.dtypes + return tuple(dtype for dtype in family.dtypes if dtype in kernel.dtypes) + + def require_resolved(self, model_id: str | None = None) -> None: + """Fail release validation when required file identities remain unresolved.""" + + selected = self._models.values() if model_id is None else (self._models[model_id],) + unresolved: list[str] = [] + for model in selected: + for label, checkpoint in (("fast", model.fast), ("official", model.official)): + for path in checkpoint.unresolved_files: + unresolved.append(f"{model.id}.{label}:{path}") + if unresolved: + detail = ", ".join(unresolved) + raise RegistryError(f"Release provenance is unresolved: {detail}") + + +def _reject_unknown_fields( + table: Mapping[str, Any], + allowed: frozenset[str], + context: str, +) -> None: + unknown = sorted(set(table).difference(allowed)) + if unknown: + raise RegistryError(f"{context} contains unknown fields: {unknown}.") + + +def _require_str(table: Mapping[str, Any], key: str, context: str) -> str: + value = table.get(key) + if not isinstance(value, str) or not value.strip(): + raise RegistryError(f"{context}.{key} must be a non-empty string.") + return value + + +def _require_enum( + table: Mapping[str, Any], + key: str, + context: str, + allowed: frozenset[str], +) -> str: + value = _require_str(table, key, context) + if value not in allowed: + raise RegistryError( + f"{context}.{key} must be one of {sorted(allowed)}; received {value!r}." + ) + return value + + +def _parse_reference_container(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "reference_container", context) + if _REFERENCE_CONTAINER_RE.fullmatch(value) is None: + raise RegistryError( + f"{context}.reference_container must be a portable 'reference-' target." + ) + return value + + +def _parse_reference_adapter(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "reference_adapter", context) + if _REFERENCE_ADAPTER_RE.fullmatch(value) is None: + raise RegistryError( + f"{context}.reference_adapter must name one module under " + "tests.parity.support.reference_adapters." + ) + return value + + +def _parse_documentation_path(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "documentation", context) + if value.count("#") > 1 or "\\" in value: + raise RegistryError(f"{context}.documentation must be a portable documentation path.") + raw_path, separator, fragment = value.partition("#") + path = PurePosixPath(raw_path) + if ( + path.is_absolute() + or ".." in path.parts + or len(path.parts) < 2 + or path.parts[0] != "docs" + or path.suffix != ".md" + or path.as_posix() != raw_path + ): + raise RegistryError( + f"{context}.documentation must reference a normalized Markdown file under docs/." + ) + if separator and _DOCUMENTATION_FRAGMENT_RE.fullmatch(fragment) is None: + raise RegistryError(f"{context}.documentation has an invalid heading fragment.") + return value + + +def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: + value = table.get(key) + if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value): + raise RegistryError(f"{context}.{key} must be a non-empty string array.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: + value = table.get(key, []) + if not isinstance(value, list) or any(not isinstance(item, str) for item in value): + raise RegistryError(f"{context}.{key} must be a string array.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None: + value = table.get(key) + if value is None: + return None + if ( + not isinstance(value, str) + or not value.strip() + or value != value.strip() + or "\n" in value + or "\r" in value + ): + raise RegistryError(f"{context}.{key} must be a non-empty single-line string.") + return value + + +def _parse_hub_license( + table: Mapping[str, Any], + *, + checkpoint_license: str, + context: str, +) -> tuple[str, str | None, str | None]: + expected_fields = {"hub_license", "hub_license_name", "hub_license_link"} + unknown_fields = sorted( + key for key in table if key.startswith("hub_") and key not in expected_fields + ) + if unknown_fields: + raise RegistryError(f"{context} contains unsupported Hub license fields: {unknown_fields}.") + identifier = _require_str(table, "hub_license", context) + if identifier not in HUB_LICENSE_IDENTIFIERS: + raise RegistryError( + f"{context}.hub_license must be a supported Hugging Face license identifier." + ) + expected_identifier: str | None = None + for prefix, candidate in ( + ("MIT", "mit"), + ("Apache-2.0", "apache-2.0"), + ("CC-BY-NC-SA-4.0", "cc-by-nc-sa-4.0"), + ("Profluent-E1-Agreement", "other"), + ("Unresolved", "other"), + ): + if checkpoint_license.startswith(prefix): + expected_identifier = candidate + break + if expected_identifier is None: + raise RegistryError( + f"{context}.checkpoint_license has no declared Hugging Face identifier mapping." + ) + if identifier != expected_identifier: + raise RegistryError( + f"{context}.hub_license must be {expected_identifier!r} for " + f"checkpoint terms {checkpoint_license!r}." + ) + + name = _optional_str(table, "hub_license_name", context) + link = _optional_str(table, "hub_license_link", context) + if identifier != "other": + if name is not None or link is not None: + raise RegistryError( + f"{context} may define hub_license_name and hub_license_link only " + "when hub_license='other'." + ) + return identifier, None, None + if name is None or link is None: + raise RegistryError( + f"{context} must define hub_license_name and hub_license_link when hub_license='other'." + ) + parsed_link = urlparse(link) + if ( + parsed_link.scheme != "https" + or not parsed_link.netloc + or not parsed_link.path + or parsed_link.username is not None + or parsed_link.password is not None + ): + raise RegistryError(f"{context}.hub_license_link must be an absolute HTTPS URL.") + return identifier, name, link + + +def _require_digest_list( + table: Mapping[str, Any], key: str, context: str +) -> tuple[FileDigest, ...]: + encoded = _require_str_list(table, key, context) + digests = tuple(FileDigest.parse(value) for value in encoded) + paths = [item.path for item in digests] + if len(paths) != len(set(paths)): + raise RegistryError(f"{context}.{key} contains duplicate paths.") + return digests + + +def _validate_revision(revision: str, context: str) -> None: + if len(revision) != 40 or _HEX_RE.fullmatch(revision) is None: + raise RegistryError(f"{context} must be an immutable 40-character commit revision.") + + +def _parse_checkpoint(table: Mapping[str, Any], prefix: str, context: str) -> CheckpointSource: + repo_id = _require_str(table, f"{prefix}_repo", context) + if _REPOSITORY_ID_RE.fullmatch(repo_id) is None: + raise RegistryError(f"{context}.{prefix}_repo must be a Hugging Face repository ID.") + revision = _require_str(table, f"{prefix}_revision", context) + _validate_revision(revision, f"{context}.{prefix}_revision") + encoded_files = _require_str_list(table, f"{prefix}_files", context) + files = tuple(FileDigest.parse(value) for value in encoded_files) + paths = [item.path for item in files] + if len(paths) != len(set(paths)): + raise RegistryError(f"{context}.{prefix}_files contains duplicate paths.") + if not any(item.path.endswith(_WEIGHT_SUFFIXES) for item in files): + raise RegistryError(f"{context}.{prefix}_files does not identify a weight file.") + unresolved_files = _optional_str_list(table, f"{prefix}_unresolved_files", context) + for unresolved_path in unresolved_files: + _portable_relative_path(unresolved_path, "Unresolved checkpoint path") + if unresolved_path in paths: + raise RegistryError( + f"{context}.{prefix} marks {unresolved_path!r} both resolved and unresolved." + ) + return CheckpointSource( + repo_id=repo_id, + revision=revision, + files=files, + unresolved_files=unresolved_files, + ) + + +def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[OracleAsset, ...]: + raw = table.get("oracle_assets", []) + if not isinstance(raw, list): + raise RegistryError(f"{context}.oracle_assets must be an array of tables.") + assets: list[OracleAsset] = [] + for index, value in enumerate(raw): + asset_context = f"{context}.oracle_assets[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{asset_context} must be a table.") + expected_fields = {"role", "path", "url", "sha256", "size"} + if set(value) != expected_fields: + raise RegistryError(f"{asset_context} must contain exactly {sorted(expected_fields)}.") + role = _require_str(value, "role", asset_context) + if role not in _ALLOWED_ORACLE_ASSET_ROLES: + raise RegistryError(f"Unsupported oracle asset role: {role!r}.") + path = _require_str(value, "path", asset_context) + try: + normalized_path = _portable_relative_path(path, "Oracle asset path") + except RegistryError as error: + raise RegistryError(f"Invalid oracle asset path: {path!r}.") from error + if normalized_path.suffix != ".pt": + raise RegistryError(f"Invalid oracle asset path: {path!r}.") + url = _require_str(value, "url", asset_context) + parsed_url = urlparse(url) + if ( + parsed_url.scheme != "https" + or parsed_url.hostname != _FAIR_ESM_ASSET_HOST + or parsed_url.path != f"/fair-esm/{path}" + or parsed_url.params + or parsed_url.query + or parsed_url.fragment + ): + raise RegistryError(f"Invalid fair-esm oracle asset URL: {url!r}.") + sha256 = _require_str(value, "sha256", asset_context) + if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: + raise RegistryError(f"Invalid oracle asset SHA-256 for {path!r}.") + size = value.get("size") + if isinstance(size, bool) or not isinstance(size, int) or size <= 0: + raise RegistryError(f"{asset_context}.size must be a positive byte count.") + assets.append( + OracleAsset( + role=role, + path=path, + url=url, + sha256=sha256, + size=size, + ) + ) + roles = [asset.role for asset in assets] + paths = [asset.path for asset in assets] + urls = [asset.url for asset in assets] + if ( + len(roles) != len(set(roles)) + or len(paths) != len(set(paths)) + or len(urls) != len(set(urls)) + ): + raise RegistryError(f"{context}.oracle_assets contains duplicate identities.") + return tuple(assets) + + +def _parse_official_golden( + table: Mapping[str, Any], + model_id: str, + context: str, +) -> OfficialGolden | None: + raw = table.get("official_golden") + if raw is None: + return None + if not isinstance(raw, dict) or set(raw) != {"metadata", "tensors"}: + raise RegistryError( + f"{context}.official_golden must contain exactly 'metadata' and 'tensors'." + ) + parsed: dict[str, FileDigest] = {} + for role in ("metadata", "tensors"): + value = raw[role] + if not isinstance(value, str): + raise RegistryError(f"{context}.official_golden.{role} must be a file digest.") + digest = FileDigest.parse(value) + if digest.algorithm != "sha256": + raise RegistryError( + f"{context}.official_golden.{role} must use an immutable SHA-256 digest." + ) + expected = f"tests/goldens/{model_id}.{'json' if role == 'metadata' else 'safetensors'}" + if digest.path != expected: + raise RegistryError(f"{context}.official_golden.{role} must use path {expected!r}.") + parsed[role] = digest + return OfficialGolden(metadata=parsed["metadata"], tensors=parsed["tensors"]) + + +def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain [[attention_kernels]] entries.") + kernels: dict[str, AttentionKernelSpec] = {} + expected_variants = { + "flash_attention_2": "flash_attn2", + "flash_attention_3": "flash_attn3", + } + for index, value in enumerate(raw): + context = f"attention_kernels[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + expected_fields = frozenset( + { + "implementation", + "repository", + "revision", + "version", + "expected_variant", + "dtypes", + } + ) + _reject_unknown_fields(value, expected_fields, context) + implementation = _require_str(value, "implementation", context) + if implementation not in expected_variants: + raise RegistryError(f"Unsupported attention kernel {implementation!r}.") + if implementation in kernels: + raise RegistryError(f"Duplicate attention kernel {implementation!r}.") + repository = _require_str(value, "repository", context) + if _REPOSITORY_ID_RE.fullmatch(repository) is None: + raise RegistryError(f"Invalid attention-kernel repository {repository!r}.") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + kernel_version = value.get("version") + if ( + isinstance(kernel_version, bool) + or not isinstance(kernel_version, int) + or kernel_version <= 0 + ): + raise RegistryError(f"{context}.version must be a positive integer.") + expected_variant = _require_str(value, "expected_variant", context) + if expected_variant != expected_variants[implementation]: + raise RegistryError( + f"{context}.expected_variant must be {expected_variants[implementation]!r}." + ) + dtypes = _require_str_list(value, "dtypes", context) + if not set(dtypes).issubset(_ALLOWED_DTYPES): + raise RegistryError(f"{context}.dtypes contains unsupported dtypes.") + kernels[implementation] = AttentionKernelSpec( + implementation=implementation, + repository=repository, + revision=revision, + version=kernel_version, + expected_variant=expected_variant, + dtypes=cast(tuple[DtypeName, ...], dtypes), + ) + if set(kernels) != set(expected_variants): + raise RegistryError("The manifest must pin both FlashAttention kernel versions.") + return kernels + + +def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[upstreams]] entry.") + upstreams: dict[str, UpstreamSource] = {} + paths: set[str] = set() + for index, value in enumerate(raw): + context = f"upstreams[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _UPSTREAM_FIELDS, context) + source_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(source_id) is None: + raise RegistryError(f"Invalid upstream ID: {source_id!r}") + if source_id in upstreams: + raise RegistryError(f"Duplicate upstream ID: {source_id!r}") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + path = _require_str(value, "path", context) + try: + normalized_path = _portable_relative_path(path, f"{context}.path") + except RegistryError as error: + raise RegistryError( + f"{context}.path must be a normalized directory directly under " + "'vendor/upstream/'." + ) from error + if ( + normalized_path.parts[:2] != ("vendor", "upstream") + or len(normalized_path.parts) != 3 + ): + raise RegistryError( + f"{context}.path must be a normalized directory directly under " + "'vendor/upstream/'." + ) + if path in paths: + raise RegistryError(f"Duplicate upstream path: {path!r}") + paths.add(path) + url = _require_str(value, "url", context) + if not url.startswith("https://github.com/") or not url.endswith(".git"): + raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.") + license_files = _require_str_list(value, "license_files", context) + license_digests = _require_digest_list(value, "license_digests", context) + if tuple(item.path for item in license_digests) != license_files: + raise RegistryError( + f"{context}.license_digests must cover license_files in the same order." + ) + distribution_files = _require_digest_list(value, "distribution_files", context) + distribution_map = {item.path: item for item in distribution_files} + for canonical in license_digests: + distributed = distribution_map.get(canonical.path) + if distributed is None or distributed.encoded != canonical.encoded: + raise RegistryError( + f"{context}.distribution_files must include an exact copy of " + f"{canonical.path!r}." + ) + if source_id == "e1": + required_e1 = { + "LICENSE", + "ATTRIBUTION", + "NOTICE", + "Apache-2.0.txt", + "BSD-3-Clause.txt", + "MODIFICATIONS.md", + } + missing_e1 = sorted(required_e1.difference(distribution_map)) + if missing_e1: + raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}") + upstreams[source_id] = UpstreamSource( + id=source_id, + path=path, + url=url, + revision=revision, + license_expression=_require_str(value, "license", context), + license_files=license_files, + license_digests=license_digests, + distribution_files=distribution_files, + ) + return upstreams + + +def _parse_families( + raw: object, + upstreams: Mapping[str, UpstreamSource], +) -> dict[str, ModelFamily]: + if not isinstance(raw, dict) or not raw: + raise RegistryError("The manifest must contain [families.] tables.") + families: dict[str, ModelFamily] = {} + for family_id, value in raw.items(): + context = f"families.{family_id}" + if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict): + raise RegistryError(f"Invalid family table: {family_id!r}") + checkpoint_license = _require_str(value, "checkpoint_license", context) + hub_license, hub_license_name, hub_license_link = _parse_hub_license( + value, + checkpoint_license=checkpoint_license, + context=context, + ) + _reject_unknown_fields(value, _FAMILY_FIELDS, context) + source_ids = _require_str_list(value, "upstreams", context) + unknown_sources = sorted(set(source_ids).difference(upstreams)) + if unknown_sources: + raise RegistryError(f"{context} references unknown upstreams: {unknown_sources}") + tokenizer_mode = _require_str(value, "tokenizer_mode", context) + if tokenizer_mode not in _ALLOWED_TOKENIZER_MODES: + raise RegistryError(f"Unsupported tokenizer mode in {context}: {tokenizer_mode!r}") + public_input = _require_str(value, "public_input", context) + attention = _require_str_list(value, "attention", context) + if not set(attention).issubset(_ALLOWED_ATTENTION): + raise RegistryError(f"Unsupported attention implementation in {context}.") + dtypes = _require_str_list(value, "dtypes", context) + if not set(dtypes).issubset(_ALLOWED_DTYPES): + raise RegistryError(f"Unsupported dtype in {context}.") + bf16_execution = cast( + Bf16Execution, + _require_enum( + value, + "bf16_execution", + context, + _ALLOWED_BF16_EXECUTIONS, + ), + ) + precisions = _require_str_list(value, "precisions", context) + if not set(precisions).issubset(_ALLOWED_PRECISIONS): + raise RegistryError(f"Unsupported precision policy in {context}.") + experimental_precisions = _optional_str_list( + value, + "experimental_precisions", + context, + ) + unknown_experimental_precisions = sorted( + set(experimental_precisions).difference(precisions) + ) + if unknown_experimental_precisions: + raise RegistryError( + f"{context}.experimental_precisions must be a subset of precisions; " + f"unknown values: {unknown_experimental_precisions}." + ) + extra = cast(RuntimeExtra, _require_enum(value, "extra", context, _ALLOWED_EXTRAS)) + vram_tier = cast( + VramTier, + _require_enum(value, "vram_tier", context, _ALLOWED_VRAM_TIERS), + ) + test_tiers_raw = _require_str_list(value, "test_tiers", context) + unknown_test_tiers = sorted(set(test_tiers_raw).difference(_ALLOWED_TEST_TIERS)) + if unknown_test_tiers: + raise RegistryError( + f"{context}.test_tiers contains unsupported tiers: {unknown_test_tiers}." + ) + test_tiers = cast(tuple[TestTier, ...], test_tiers_raw) + reference_container = _parse_reference_container(value, context) + reference_adapter = _parse_reference_adapter(value, context) + documentation = _parse_documentation_path(value, context) + runtime_paths = _require_str_list(value, "runtime_paths", context) + if len(runtime_paths) != len(set(runtime_paths)): + raise RegistryError(f"{context}.runtime_paths must not contain duplicates.") + for runtime_path in runtime_paths: + try: + _portable_relative_path(runtime_path, f"{context}.runtime_paths entry") + except RegistryError as error: + raise RegistryError( + f"Unsafe runtime path in {context}: {runtime_path!r}" + ) from error + if runtime_path.startswith("vendor/"): + raise RegistryError(f"Unsafe runtime path in {context}: {runtime_path!r}") + requires_complete_weight_publication = value.get( + "requires_complete_weight_publication", + False, + ) + if not isinstance(requires_complete_weight_publication, bool): + raise RegistryError( + f"{context}.requires_complete_weight_publication must be a boolean." + ) + if "weights_publication_allowed" not in value: + raise RegistryError( + f"{context}.weights_publication_allowed must be declared explicitly." + ) + weights_publication_allowed = value["weights_publication_allowed"] + if not isinstance(weights_publication_allowed, bool): + raise RegistryError(f"{context}.weights_publication_allowed must be a boolean.") + raw_auto_map = value.get("auto_map") + if not isinstance(raw_auto_map, dict) or not raw_auto_map: + raise RegistryError(f"{context}.auto_map must be a non-empty table.") + auto_map: list[tuple[str, str]] = [] + for auto_class, class_path in raw_auto_map.items(): + if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): + raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") + if not class_path.startswith("fastplms.") or class_path.count(".") < 2: + raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") + auto_map.append((auto_class, class_path)) + tokenizer_class = value.get("tokenizer_class") + if tokenizer_class is not None: + if tokenizer_mode != "tokenizer": + raise RegistryError( + f"{context}.tokenizer_class requires tokenizer_mode='tokenizer'." + ) + if ( + not isinstance(tokenizer_class, str) + or not tokenizer_class.startswith("fastplms.") + or tokenizer_class.count(".") < 2 + ): + raise RegistryError( + f"Invalid tokenizer class path in {context}: {tokenizer_class!r}" + ) + backbone_model = value.get("backbone_model") + if backbone_model is not None and ( + not isinstance(backbone_model, str) + or _IDENTIFIER_RE.fullmatch(backbone_model) is None + ): + raise RegistryError( + f"{context}.backbone_model must be a valid manifest model ID." + ) + state_transform = _require_str(value, "state_transform", context) + conversion_provenance = _require_str(value, "conversion_provenance", context) + required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:") + missing_sections = [ + section for section in required_sections if section not in conversion_provenance + ] + if missing_sections or state_transform not in conversion_provenance: + raise RegistryError( + f"{context}.conversion_provenance must identify {state_transform!r} and " + f"contain mechanism-first sections; missing {missing_sections}." + ) + families[family_id] = ModelFamily( + id=family_id, + architecture=_require_str(value, "architecture", context), + upstreams=source_ids, + tokenizer_mode=tokenizer_mode, + public_input=public_input, + extra=extra, + reference_container=reference_container, + reference_adapter=reference_adapter, + attention=attention, + dtypes=cast(tuple[DtypeName, ...], dtypes), + bf16_execution=bf16_execution, + precisions=precisions, + experimental_precisions=experimental_precisions, + vram_tier=vram_tier, + checkpoint_license=checkpoint_license, + hub_license=hub_license, + state_transform=state_transform, + representative=_require_str(value, "representative", context), + documentation=documentation, + test_tiers=test_tiers, + runtime_paths=runtime_paths, + auto_map_items=tuple(auto_map), + requires_complete_weight_publication=requires_complete_weight_publication, + weights_publication_allowed=weights_publication_allowed, + tokenizer_class=tokenizer_class, + hub_license_name=hub_license_name, + hub_license_link=hub_license_link, + conversion_provenance=conversion_provenance, + backbone_model=backbone_model, + ) + return families + + +def _parse_runtime_assets( + raw: object, + families: Mapping[str, ModelFamily], +) -> dict[str, RuntimeAsset]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.") + runtime_assets: dict[str, RuntimeAsset] = {} + identities: set[tuple[str, str, str]] = set() + for index, value in enumerate(raw): + context = f"runtime_assets[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _RUNTIME_ASSET_FIELDS, context) + asset_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(asset_id) is None: + raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}") + if asset_id in runtime_assets: + raise RegistryError(f"Duplicate runtime asset ID: {asset_id!r}") + repository = _require_str(value, "repository", context) + if _REPOSITORY_ID_RE.fullmatch(repository) is None: + raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + path = _require_str(value, "path", context) + try: + normalized_path = _portable_relative_path(path, "Runtime asset path") + except RegistryError as error: + raise RegistryError(f"Runtime asset path is not portable: {path!r}") from error + sha256 = _require_str(value, "sha256", context) + if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: + raise RegistryError(f"Invalid runtime asset SHA-256 for {path!r}.") + size = value.get("size") + if isinstance(size, bool) or not isinstance(size, int) or size <= 0: + raise RegistryError(f"{context}.size must be a positive byte count.") + consumer_family = _require_str(value, "consumer_family", context) + if consumer_family not in families: + raise RegistryError( + f"{context}.consumer_family references unknown family {consumer_family!r}." + ) + trust_kind = cast( + RuntimeAssetTrustKind, + _require_enum( + value, + "trust_kind", + context, + _ALLOWED_RUNTIME_ASSET_TRUST_KINDS, + ), + ) + license_expression = _require_str(value, "license", context) + offline_behavior = _require_str(value, "offline_behavior", context) + if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS: + raise RegistryError( + f"{context}.offline_behavior is unsupported: {offline_behavior!r}." + ) + if trust_kind == "hash_pinned_pickle" and normalized_path.suffix != ".pkl": + raise RegistryError( + f"{context}.path must end in '.pkl' for trust_kind='hash_pinned_pickle'." + ) + identity = (repository, revision, path) + if identity in identities: + raise RegistryError(f"Duplicate runtime asset identity: {identity!r}") + identities.add(identity) + runtime_assets[asset_id] = RuntimeAsset( + id=asset_id, + repository=repository, + revision=revision, + path=path, + sha256=sha256, + size=size, + consumer_family=consumer_family, + trust_kind=trust_kind, + license_expression=license_expression, + offline_behavior=offline_behavior, + ) + return runtime_assets + + +def _parse_models( + raw: object, + families: Mapping[str, ModelFamily], +) -> dict[str, ModelSpec]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[models]] entry.") + models: dict[str, ModelSpec] = {} + fast_repositories: set[str] = set() + for index, value in enumerate(raw): + context = f"models[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _MODEL_FIELDS, context) + model_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(model_id) is None: + raise RegistryError(f"Invalid model ID: {model_id!r}") + if model_id in models: + raise RegistryError(f"Duplicate model ID: {model_id!r}") + family_id = _require_str(value, "family", context) + if family_id not in families: + raise RegistryError(f"{context} references unknown family {family_id!r}.") + fast = _parse_checkpoint(value, "fast", context) + official = _parse_checkpoint(value, "official", context) + if fast.repo_id in fast_repositories: + raise RegistryError(f"Duplicate FastPLMs repository ID: {fast.repo_id!r}") + fast_repositories.add(fast.repo_id) + family = families[family_id] + oracle_assets = _parse_oracle_assets(value, context) + official_golden = _parse_official_golden(value, model_id, context) + size_category = _require_str(value, "size_category", context) + if size_category not in _ALLOWED_SIZE_CATEGORIES: + raise RegistryError(f"Unsupported size category in {context}: {size_category!r}") + generation_contract = cast( + GenerationContract, + _require_enum( + value, + "generation_contract", + context, + _ALLOWED_GENERATION_CONTRACTS, + ), + ) + if family.tokenizer_mode == "structure" and size_category != "structure": + raise RegistryError( + f"Structure checkpoint {model_id!r} must use size_category='structure'." + ) + artifact_source = value.get("artifact_source", "fast") + if artifact_source not in {"fast", "official"}: + raise RegistryError(f"{context}.artifact_source must be 'fast' or 'official'.") + canonical_state_sha256 = value.get("canonical_state_sha256") + if artifact_source == "official": + if ( + not isinstance(canonical_state_sha256, str) + or len(canonical_state_sha256) != 64 + or _HEX_RE.fullmatch(canonical_state_sha256) is None + ): + raise RegistryError( + f"{context}.canonical_state_sha256 must be a SHA-256 commitment " + "for an official-source artifact." + ) + elif canonical_state_sha256 is not None: + raise RegistryError( + f"{context}.canonical_state_sha256 is restricted to official-source artifacts." + ) + if family.tokenizer_mode == "tokenizer" and not any( + "tokenizer" in item.path or "vocab" in item.path for item in fast.files + ): + raise RegistryError(f"{context} does not pin a tokenizer asset.") + tokenizer_source_id = value.get("tokenizer_source") + if tokenizer_source_id is not None and ( + family.tokenizer_mode != "tokenizer" + or not isinstance(tokenizer_source_id, str) + or _IDENTIFIER_RE.fullmatch(tokenizer_source_id) is None + ): + raise RegistryError(f"{context}.tokenizer_source is invalid.") + notes = value.get("notes", "") + if not isinstance(notes, str): + raise RegistryError(f"{context}.notes must be a string.") + msa_conditioning = value.get("msa_conditioning") + if family_id == "esmfold2": + if not isinstance(msa_conditioning, bool): + raise RegistryError( + f"{context}.msa_conditioning must be an explicit boolean for " + "ESMFold2 checkpoints." + ) + elif "msa_conditioning" in value: + raise RegistryError( + f"{context}.msa_conditioning is only valid for ESMFold2 checkpoints." + ) + raw_auto_map = value.get("auto_map") + auto_map: list[tuple[str, str]] = [] + if raw_auto_map is not None: + if not isinstance(raw_auto_map, dict) or not raw_auto_map: + raise RegistryError(f"{context}.auto_map must be a non-empty table.") + for auto_class, class_path in raw_auto_map.items(): + if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): + raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") + if not class_path.startswith("fastplms.") or class_path.count(".") < 2: + raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") + auto_map.append((auto_class, class_path)) + models[model_id] = ModelSpec( + id=model_id, + family=family, + fast=fast, + official=official, + size_category=size_category, + generation_contract=generation_contract, + oracle_assets=oracle_assets, + official_golden=official_golden, + artifact_source=artifact_source, + canonical_state_sha256=canonical_state_sha256, + tokenizer_source_id=tokenizer_source_id, + auto_map_items=tuple(auto_map), + notes=notes, + msa_conditioning=msa_conditioning, + ) + return models + + +def _validate_registry( + upstreams: Mapping[str, UpstreamSource], + attention_kernels: Mapping[str, AttentionKernelSpec], + families: Mapping[str, ModelFamily], + models: Mapping[str, ModelSpec], +) -> None: + for spec in models.values(): + if spec.tokenizer_source_id is None: + continue + source = models.get(spec.tokenizer_source_id) + if source is None: + raise RegistryError( + f"Model {spec.id!r} references unknown tokenizer source " + f"{spec.tokenizer_source_id!r}." + ) + if not any( + PurePosixPath(item.path).name + in { + "added_tokens.json", + "merges.txt", + "sentencepiece.bpe.model", + "special_tokens_map.json", + "spiece.model", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "vocab.txt", + } + for item in source.official.files + ): + raise RegistryError( + f"Tokenizer source {source.id!r} has no official tokenizer assets." + ) + expected_esmfold2 = { + "esmfold2": ("Synthyra/ESMFold2", "biohub/ESMFold2"), + "esmfold2_fast": ("Synthyra/ESMFold2-Fast", "biohub/ESMFold2-Fast"), + "esmfold2_experimental_cutoff2025": ( + "Synthyra/ESMFold2-Experimental-Cutoff2025", + "biohub/ESMFold2-Experimental-Cutoff2025", + ), + "esmfold2_experimental_fast_cutoff2025": ( + "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025", + "biohub/ESMFold2-Experimental-Fast-Cutoff2025", + ), + } + actual_esmfold2 = { + model.id: (model.fast.repo_id, model.official.repo_id) + for model in models.values() + if model.family.id == "esmfold2" + } + if actual_esmfold2 != expected_esmfold2: + raise RegistryError( + "ESMFold2 support must contain exactly the four approved model IDs and " + "official/Synthyra repositories." + ) + + golden_paths: list[str] = [] + for model in models.values(): + if model.official_golden is not None: + golden_paths.extend( + ( + model.official_golden.metadata.path, + model.official_golden.tensors.path, + ) + ) + if len(golden_paths) != len(set(golden_paths)): + raise RegistryError("Official golden paths must be unique across model declarations.") + unused_upstreams = sorted( + set(upstreams).difference( + source for family in families.values() for source in family.upstreams + ) + ) + if unused_upstreams: + raise RegistryError( + f"Upstream sources are not connected to a model family: {unused_upstreams}" + ) + advertised_flash = { + implementation + for family in families.values() + for implementation in family.attention + if implementation.startswith("flash_attention_") + } + missing_kernels = sorted(advertised_flash.difference(attention_kernels)) + if missing_kernels: + raise RegistryError( + f"Advertised FlashAttention backends lack kernel specs: {missing_kernels}." + ) + for family in families.values(): + for implementation in family.attention: + kernel = attention_kernels.get(implementation) + if kernel is not None and not set(family.dtypes).intersection(kernel.dtypes): + raise RegistryError( + f"Family {family.id!r} and attention kernel {implementation!r} " + "have no supported dtype in common." + ) + family_models = [model for model in models.values() if model.family.id == family.id] + if not family_models: + raise RegistryError(f"Family {family.id!r} has no checkpoints.") + representative = models.get(family.representative) + if representative is None or representative.family.id != family.id: + raise RegistryError( + f"Family {family.id!r} has invalid representative {family.representative!r}." + ) + if family.backbone_model is not None and family.backbone_model not in models: + raise RegistryError( + f"Family {family.id!r} references unknown backbone model " + f"{family.backbone_model!r}." + ) + + +def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry: + try: + manifest = tomllib.loads(raw_bytes.decode("utf-8")) + except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error: + raise RegistryError(f"Unable to parse model manifest: {error}") from error + _reject_unknown_fields(manifest, _ROOT_FIELDS, "manifest") + if manifest.get("schema_version") != 1: + raise RegistryError("Unsupported model manifest schema_version; expected 1.") + legal_files = _require_digest_list(manifest, "legal_files", "manifest") + required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"} + if {item.path for item in legal_files} != required_legal_paths: + raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.") + attention_kernels = _parse_attention_kernels(manifest.get("attention_kernels")) + upstreams = _parse_upstreams(manifest.get("upstreams")) + families = _parse_families(manifest.get("families"), upstreams) + runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families) + models = _parse_models(manifest.get("models"), families) + _validate_registry(upstreams, attention_kernels, families, models) + return ModelRegistry( + schema_version=1, + upstreams=upstreams, + attention_kernels=attention_kernels, + families=families, + models=models, + runtime_assets=runtime_assets, + legal_files=legal_files, + ) + + +def load_model_registry(path: str | Path | None = None) -> ModelRegistry: + """Load and validate a model manifest without importing model code.""" + + if path is None: + manifest = resources.files("fastplms").joinpath("models.toml") + return _load_manifest_bytes(manifest.read_bytes()) + return _load_manifest_bytes(Path(path).read_bytes()) + + +@lru_cache(maxsize=1) +def get_model_registry() -> ModelRegistry: + """Return the validated package registry, cached after its first read.""" + + return load_model_registry() + + +def get_model_spec(model_id: str) -> ModelSpec: + """Return one model specification by its stable manifest ID.""" + + try: + return get_model_registry()[model_id] + except KeyError as error: + supported = ", ".join(get_model_registry()) + raise KeyError( + f"Unknown FastPLMs model ID {model_id!r}. Supported IDs: {supported}" + ) from error + + +__all__ = [ + "HUB_LICENSE_IDENTIFIERS", + "CheckpointSource", + "FileDigest", + "GenerationContract", + "ModelFamily", + "ModelRegistry", + "ModelSpec", + "OracleAsset", + "RegistryError", + "RuntimeAsset", + "RuntimeAssetTrustKind", + "RuntimeExtra", + "TestTier", + "UpstreamSource", + "VramTier", + "get_model_registry", + "get_model_spec", + "load_model_registry", +] diff --git a/src/fastplms/runtime.py b/src/fastplms/runtime.py new file mode 100644 index 0000000..49e60db --- /dev/null +++ b/src/fastplms/runtime.py @@ -0,0 +1,69 @@ +"""Explicit, reversible Torch runtime configuration. + +Importing FastPLMs does not change global Torch settings. Callers that want a +runtime profile opt in with :func:`runtime_profile` and receive their previous +settings back when the context exits. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +MatmulPrecision = Literal["highest", "high", "medium"] + + +@dataclass(frozen=True, slots=True) +class RuntimeProfile: + """Requested Torch settings for a bounded inference or training block.""" + + float32_matmul_precision: MatmulPrecision = "highest" + allow_tf32: bool | None = None + + +@contextmanager +def runtime_profile(profile: RuntimeProfile | None = None) -> Iterator[None]: + """Apply a Torch runtime profile and restore the previous global settings. + + The default profile requests the highest float32 matrix-multiplication + precision and leaves TF32 policy unchanged. Torch is imported only when the + context is entered. + """ + + import torch + + selected = profile or RuntimeProfile() + previous_matmul_precision = torch.get_float32_matmul_precision() + matmul_backend = getattr(getattr(torch.backends, "cuda", None), "matmul", None) + cudnn_backend = getattr(torch.backends, "cudnn", None) + previous_matmul_tf32 = ( + getattr(matmul_backend, "allow_tf32", None) if matmul_backend is not None else None + ) + previous_cudnn_tf32 = ( + getattr(cudnn_backend, "allow_tf32", None) if cudnn_backend is not None else None + ) + + torch.set_float32_matmul_precision(selected.float32_matmul_precision) + if selected.allow_tf32 is not None: + if matmul_backend is not None and hasattr(matmul_backend, "allow_tf32"): + matmul_backend.allow_tf32 = selected.allow_tf32 + if cudnn_backend is not None and hasattr(cudnn_backend, "allow_tf32"): + cudnn_backend.allow_tf32 = selected.allow_tf32 + try: + yield + finally: + torch.set_float32_matmul_precision(previous_matmul_precision) + if selected.allow_tf32 is not None: + if matmul_backend is not None and previous_matmul_tf32 is not None: + matmul_backend.allow_tf32 = previous_matmul_tf32 + if cudnn_backend is not None and previous_cudnn_tf32 is not None: + cudnn_backend.allow_tf32 = previous_cudnn_tf32 + + +__all__ = ["MatmulPrecision", "RuntimeProfile", "runtime_profile"] diff --git a/testing/compliance.py b/testing/compliance.py deleted file mode 100644 index 49d256f..0000000 --- a/testing/compliance.py +++ /dev/null @@ -1,212 +0,0 @@ -import entrypoint_setup - -import torch -import random -from typing import List, Tuple -from torch.nn.functional import mse_loss -from tqdm import tqdm -from collections import defaultdict -from transformers import AutoModelForMaskedLM - -from fastplms.esm2.modeling_fastesm import FastEsmForMaskedLM -from fastplms.esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM -from fastplms.e1.modeling_e1 import E1ForMaskedLM -from testing.official.esm2 import load_official_model as load_official_esm2_model -from testing.official.esm_plusplus import load_official_model as load_official_esmc_model -from testing.official.e1 import load_official_model as load_official_e1_model - -from fastplms.weight_parity_utils import assert_state_dict_equal - - -class ComplianceChecker: - def __init__( - self, - test_number_batches: int = 25, - batch_size: int = 8, - min_sequence_length: int = 16, - max_sequence_length: int = 128, - ) -> None: - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - self.test_number_batches = test_number_batches - self.batch_size = batch_size - self.min_sequence_length = min_sequence_length - self.max_sequence_length = max_sequence_length - self.canonical_amino_acids = "ACDEFGHIKLMNPQRSTVWY" - - def _load_esmc(self, from_auto_model: bool = False, force_download: bool = False) -> Tuple[torch.nn.Module, torch.nn.Module, object]: - official_model_path = "biohub/ESMC-300M" - fast_model_path = "Synthyra/ESMplusplus_small" - official_model, tokenizer = load_official_esmc_model( - reference_repo_id=official_model_path, - device=self.device, - dtype=torch.bfloat16, - ) - load_class = AutoModelForMaskedLM if from_auto_model else ESMplusplusForMaskedLM - fast_model = load_class.from_pretrained( - fast_model_path, - dtype=torch.bfloat16, - device_map=self.device, - force_download=force_download, - trust_remote_code=True, - ).eval() - return official_model, fast_model, tokenizer - - def _load_esm2(self, from_auto_model: bool = False, force_download: bool = False) -> Tuple[torch.nn.Module, torch.nn.Module, object]: - official_model_path = "facebook/esm2_t6_8M_UR50D" - fast_model_path = "Synthyra/ESM2-8M" - official_model, tokenizer = load_official_esm2_model( - reference_repo_id=official_model_path, - device=self.device, - dtype=torch.bfloat16, - ) - load_class = AutoModelForMaskedLM if from_auto_model else FastEsmForMaskedLM - fast_model = load_class.from_pretrained( - fast_model_path, - dtype=torch.bfloat16, - device_map=self.device, - force_download=force_download, - trust_remote_code=True, - ).eval() - return official_model, fast_model, tokenizer - - def _load_e1(self, from_auto_model: bool = False, force_download: bool = False) -> Tuple[torch.nn.Module, torch.nn.Module, object]: - official_model_path = "Profluent-Bio/E1-150m" - fast_model_path = "Synthyra/Profluent-E1-150M" - official_model, tokenizer = load_official_e1_model( - reference_repo_id=official_model_path, - device=self.device, - dtype=torch.bfloat16, - ) - load_class = AutoModelForMaskedLM if from_auto_model else E1ForMaskedLM - fast_model = load_class.from_pretrained( - fast_model_path, - dtype=torch.bfloat16, - device_map=self.device, - force_download=force_download, - trust_remote_code=True, - ).eval() - return official_model, fast_model, tokenizer - - def _generate_random_sequence(self, length: int) -> str: - return 'M' + "".join(random.choices(self.canonical_amino_acids, k=length)) - - def _generate_random_batch(self, batch_size: int, min_length: int, max_length: int) -> List[str]: - return [self._generate_random_sequence(random.randint(min_length, max_length)) for _ in range(batch_size)] - - def _weight_compliance(self, official_model: torch.nn.Module, fast_model: torch.nn.Module) -> None: - for (official_name, official_param), (fast_name, fast_param) in zip(official_model.model.state_dict().items(), fast_model.state_dict().items()): - if official_name == fast_name: - diff = mse_loss(official_param, fast_param).item() - if diff > 0.0: - print(f"{official_name}: {diff}") - assert diff < 1e-3, f"Parameter {official_name} has a large difference: {diff}" - else: - print(f"Name mismatch: {official_name} != {fast_name}") - - @torch.inference_mode() - def _foward_compliance(self, model_type: str, official_model: torch.nn.Module, fast_model: torch.nn.Module, tokenizer: object, only_non_pad_tokens: bool = False) -> None: - cumulative_logits_mse = 0 - cumulative_preds_accuracy = 0 - hidden_state_diff_dict = defaultdict(int) - - for _ in tqdm(range(self.test_number_batches)): - batch = self._generate_random_batch(self.batch_size, self.min_sequence_length, self.max_sequence_length) - if model_type == "E1": - tokenized = tokenizer.get_batch_kwargs(batch, device=self.device) - tokenized = { - "input_ids": tokenized["input_ids"], - "within_seq_position_ids": tokenized["within_seq_position_ids"], - "global_position_ids": tokenized["global_position_ids"], - "sequence_ids": tokenized["sequence_ids"], - "attention_mask": (tokenized["sequence_ids"] != -1).long(), - } - else: - tokenized = tokenizer(batch, return_tensors="pt", padding=True) - tokenized = {k: v.to(self.device) for k, v in tokenized.items()} - - attention_mask = tokenized['attention_mask'].cpu().bool() - model_inputs = tokenized.copy() - if model_type == "ESMC": - model_inputs["sequence_id"] = model_inputs["attention_mask"].to(dtype=torch.bool) - - official_output = official_model(**model_inputs, output_hidden_states=True) - official_hidden_states = official_output.hidden_states - official_logits = official_output.logits.cpu() - if only_non_pad_tokens: - official_logits = official_logits[attention_mask] - official_preds = official_logits.argmax(dim=-1) - - fast_output = fast_model(**model_inputs, output_hidden_states=True) - fast_hidden_states = fast_output.hidden_states - fast_logits = fast_output.logits.cpu() - if only_non_pad_tokens: - fast_logits = fast_logits[attention_mask] - fast_preds = fast_logits.argmax(dim=-1) - - cumulative_logits_mse += mse_loss(official_logits, fast_logits) - cumulative_preds_accuracy += (official_preds == fast_preds).float().mean() - - for i in range(len(official_hidden_states)): - official_state, fast_state = official_hidden_states[i], fast_hidden_states[i] - if only_non_pad_tokens: - official_state, fast_state = official_state[attention_mask], fast_state[attention_mask] - hidden_state_diff_dict[i] += mse_loss(official_state, fast_state).item() - - avg_logits_mse = cumulative_logits_mse / self.test_number_batches - avg_preds_accuracy = cumulative_preds_accuracy / self.test_number_batches - print(f"Average logits MSE: {avg_logits_mse}") - print(f"Average preds accuracy: {avg_preds_accuracy}") - - if avg_logits_mse > 1e-3 or avg_preds_accuracy < 0.95: - print("Differences were too large, printing hidden state differences for debugging...") - for k, v in hidden_state_diff_dict.items(): - print(f"Hidden state {k} Avg MSE: {v / self.test_number_batches}") - - - def __call__( - self, - model_type: str = "ESMC", - force_download: bool = False, - from_auto_model: bool = False, - only_non_pad_tokens: bool = False, - ) -> None: - if model_type == "ESMC": - official_model, fast_model, tokenizer = self._load_esmc(from_auto_model, force_download) - elif model_type == "ESM2": - official_model, fast_model, tokenizer = self._load_esm2(from_auto_model, force_download) - elif model_type == "E1": - official_model, fast_model, tokenizer = self._load_e1(from_auto_model, force_download) - else: - raise ValueError(f"Unsupported model type: {model_type}. Supported: ESMC, ESM2, E1") - assert_state_dict_equal( - reference_state_dict=official_model.model.state_dict(), - candidate_state_dict=fast_model.state_dict(), - context=f"{model_type} weight parity", - ) - self._weight_compliance(official_model, fast_model) - self._foward_compliance(model_type, official_model, fast_model, tokenizer, only_non_pad_tokens) - - -if __name__ == "__main__": - import argparse - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument("--only_non_pad_tokens", action="store_true") - parser.add_argument("--force_download", action="store_true") - parser.add_argument("--from_auto_model", action="store_true") - parser.add_argument("--model_types", nargs="+", default=["ESMC", "ESM2", "E1"]) - args = parser.parse_args() - - if args.hf_token is not None: - from huggingface_hub import login - login(token=args.hf_token) - - checker = ComplianceChecker() - for model_type in args.model_types: - print(f"Checking {model_type}...") - checker( - model_type=model_type, - from_auto_model=args.from_auto_model, - only_non_pad_tokens=args.only_non_pad_tokens, - force_download=args.force_download - ) diff --git a/testing/conftest.py b/testing/conftest.py deleted file mode 100644 index dcb1376..0000000 --- a/testing/conftest.py +++ /dev/null @@ -1,420 +0,0 @@ -import contextlib -import os -import random -from typing import Dict, List, Tuple - -import pytest -import torch - - -def pytest_configure(config): - config.addinivalue_line("markers", "gpu: requires CUDA GPU") - config.addinivalue_line("markers", "slow: loads two models simultaneously (compliance tests)") - config.addinivalue_line("markers", "large: requires 24+ GB VRAM (3B parameter models)") - config.addinivalue_line("markers", "structure: structure prediction models (Boltz2, ESMFold, ESMFold2)") - - -# Standalone scripts that are not pytest tests -collect_ignore = [ - os.path.join(os.path.dirname(__file__), "test_contact_maps.py"), - os.path.join(os.path.dirname(__file__), "compliance.py"), - os.path.join(os.path.dirname(__file__), "throughput.py"), - os.path.join(os.path.dirname(__file__), "run_boltz2_compliance.py"), -] - -CANONICAL_AAS = "ACDEFGHIKLMNPQRSTVWY" -SEED = 42 -DEFAULT_BATCH_SIZE = 4 -MAX_EMBED_LEN = 128 - - -@contextlib.contextmanager -def strict_fp32_matmul(): - """Temporarily disable TF32 for fp32 numerical parity checks.""" - try: - old_fp32_precision = torch.backends.fp32_precision - old_matmul_precision = torch.backends.cuda.matmul.fp32_precision - old_cudnn_precision = torch.backends.cudnn.fp32_precision - except AttributeError: - old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 - old_cudnn_tf32 = torch.backends.cudnn.allow_tf32 - old_matmul_precision = torch.get_float32_matmul_precision() - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.set_float32_matmul_precision("highest") - try: - yield - finally: - torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32 - torch.backends.cudnn.allow_tf32 = old_cudnn_tf32 - torch.set_float32_matmul_precision(old_matmul_precision) - return - - torch.backends.fp32_precision = "ieee" - torch.backends.cuda.matmul.fp32_precision = "ieee" - torch.backends.cudnn.fp32_precision = "ieee" - try: - yield - finally: - torch.backends.fp32_precision = old_fp32_precision - torch.backends.cuda.matmul.fp32_precision = old_matmul_precision - torch.backends.cudnn.fp32_precision = old_cudnn_precision - -# Default registry: one small model per family for fast CI -MODEL_REGISTRY: Dict[str, Dict] = { - "esm2": { - "fast_path": "Synthyra/ESM2-8M", - "official_path": "facebook/esm2_t6_8M_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - }, - "esmc": { - "fast_path": "Synthyra/ESMplusplus_small", - "official_path": "biohub/ESMC-300M", - "load_official": "testing.official.esm_plusplus", - "model_type": "ESMC", - "uses_tokenizer": True, - }, - "esm3": { - "fast_path": "Synthyra/ESM3_small", - "official_path": "esm3-sm-open-v1", - "load_official": "testing.official.esm3", - "model_type": "ESM3", - "uses_tokenizer": True, - }, - "e1": { - "fast_path": "Synthyra/Profluent-E1-150M", - "official_path": "Profluent-Bio/E1-150m", - "load_official": "testing.official.e1", - "model_type": "E1", - "uses_tokenizer": False, - }, - "dplm": { - "fast_path": "Synthyra/DPLM-150M", - "official_path": "airkingbd/dplm_150m", - "load_official": "testing.official.dplm", - "model_type": "DPLM", - "uses_tokenizer": True, - }, - "dplm2": { - "fast_path": "Synthyra/DPLM2-150M", - "official_path": "airkingbd/dplm2_150m", - "load_official": "testing.official.dplm2", - "model_type": "DPLM2", - "uses_tokenizer": True, - }, - "ankh": { - "fast_path": "Synthyra/ANKH_base", - "official_path": "ElnaggarLab/ankh-base", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - }, -} - -# Full registry: every checkpoint across all model families -FULL_MODEL_REGISTRY: Dict[str, Dict] = { - # ESM2 family - "esm2_8m": { - "fast_path": "Synthyra/ESM2-8M", - "official_path": "facebook/esm2_t6_8M_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - "size_category": "small", - }, - "esm2_35m": { - "fast_path": "Synthyra/ESM2-35M", - "official_path": "facebook/esm2_t12_35M_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - "size_category": "small", - }, - "esm2_150m": { - "fast_path": "Synthyra/ESM2-150M", - "official_path": "facebook/esm2_t30_150M_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - "size_category": "medium", - }, - "esm2_650m": { - "fast_path": "Synthyra/ESM2-650M", - "official_path": "facebook/esm2_t33_650M_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - "size_category": "large", - }, - "esm2_3b": { - "fast_path": "Synthyra/ESM2-3B", - "official_path": "facebook/esm2_t36_3B_UR50D", - "load_official": "testing.official.esm2", - "model_type": "ESM2", - "uses_tokenizer": True, - "size_category": "xlarge", - }, - # ESM++ family - "esmc_small": { - "fast_path": "Synthyra/ESMplusplus_small", - "official_path": "biohub/ESMC-300M", - "load_official": "testing.official.esm_plusplus", - "model_type": "ESMC", - "uses_tokenizer": True, - "size_category": "medium", - }, - "esmc_large": { - "fast_path": "Synthyra/ESMplusplus_large", - "official_path": "biohub/ESMC-600M", - "load_official": "testing.official.esm_plusplus", - "model_type": "ESMC", - "uses_tokenizer": True, - "size_category": "large", - }, - "esmc_6b": { - "fast_path": "Synthyra/ESMplusplus_6B", - "official_path": "biohub/ESMC-6B", - "load_official": "testing.official.esm_plusplus", - "model_type": "ESMC", - "uses_tokenizer": True, - "size_category": "xlarge", - }, - "esm3_small": { - "fast_path": "Synthyra/ESM3_small", - "official_path": "esm3-sm-open-v1", - "load_official": "testing.official.esm3", - "model_type": "ESM3", - "uses_tokenizer": True, - "size_category": "large", - }, - # E1 family - "e1_150m": { - "fast_path": "Synthyra/Profluent-E1-150M", - "official_path": "Profluent-Bio/E1-150m", - "load_official": "testing.official.e1", - "model_type": "E1", - "uses_tokenizer": False, - "size_category": "small", - }, - "e1_300m": { - "fast_path": "Synthyra/Profluent-E1-300M", - "official_path": "Profluent-Bio/E1-300m", - "load_official": "testing.official.e1", - "model_type": "E1", - "uses_tokenizer": False, - "size_category": "medium", - }, - "e1_600m": { - "fast_path": "Synthyra/Profluent-E1-600M", - "official_path": "Profluent-Bio/E1-600m", - "load_official": "testing.official.e1", - "model_type": "E1", - "uses_tokenizer": False, - "size_category": "large", - }, - # DPLM family - "dplm_150m": { - "fast_path": "Synthyra/DPLM-150M", - "official_path": "airkingbd/dplm_150m", - "load_official": "testing.official.dplm", - "model_type": "DPLM", - "uses_tokenizer": True, - "size_category": "small", - }, - "dplm_650m": { - "fast_path": "Synthyra/DPLM-650M", - "official_path": "airkingbd/dplm_650m", - "load_official": "testing.official.dplm", - "model_type": "DPLM", - "uses_tokenizer": True, - "size_category": "large", - }, - "dplm_3b": { - "fast_path": "Synthyra/DPLM-3B", - "official_path": "airkingbd/dplm_3b", - "load_official": "testing.official.dplm", - "model_type": "DPLM", - "uses_tokenizer": True, - "size_category": "xlarge", - }, - # DPLM2 family - "dplm2_150m": { - "fast_path": "Synthyra/DPLM2-150M", - "official_path": "airkingbd/dplm2_150m", - "load_official": "testing.official.dplm2", - "model_type": "DPLM2", - "uses_tokenizer": True, - "size_category": "small", - }, - "dplm2_650m": { - "fast_path": "Synthyra/DPLM2-650M", - "official_path": "airkingbd/dplm2_650m", - "load_official": "testing.official.dplm2", - "model_type": "DPLM2", - "uses_tokenizer": True, - "size_category": "large", - }, - "dplm2_3b": { - "fast_path": "Synthyra/DPLM2-3B", - "official_path": "airkingbd/dplm2_3b", - "load_official": "testing.official.dplm2", - "model_type": "DPLM2", - "uses_tokenizer": True, - "size_category": "xlarge", - }, - # ANKH family - "ankh_base": { - "fast_path": "Synthyra/ANKH_base", - "official_path": "ElnaggarLab/ankh-base", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - "size_category": "medium", - }, - "ankh_large": { - "fast_path": "Synthyra/ANKH_large", - "official_path": "ElnaggarLab/ankh-large", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - "size_category": "large", - }, - "ankh2_large": { - "fast_path": "Synthyra/ANKH2_large", - "official_path": "ElnaggarLab/ankh2-ext2", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - "size_category": "large", - }, - "ankh3_large": { - "fast_path": "Synthyra/ANKH3_large", - "official_path": "ElnaggarLab/ankh3-large", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - "size_category": "large", - }, - "ankh3_xl": { - "fast_path": "Synthyra/ANKH3_xl", - "official_path": "ElnaggarLab/ankh3-xl", - "load_official": "testing.official.ankh", - "model_type": "ANKH", - "uses_tokenizer": True, - "size_category": "xlarge", - }, -} - -# Structure prediction models (separate API, not MaskedLM) -STRUCTURE_MODEL_REGISTRY: Dict[str, Dict] = { - "boltz2": { - "fast_path": "Synthyra/Boltz2", - "model_type": "Boltz2", - "size_category": "structure", - }, - "esmfold": { - "fast_path": "Synthyra/FastESMFold", - "model_type": "ESMFold", - "size_category": "structure", - }, - "esmfold2": { - "fast_path": "Synthyra/ESMFold2", - "official_path": "biohub/ESMFold2", - "model_type": "ESMFold2", - "size_category": "structure", - }, - "esmfold2_fast": { - "fast_path": "Synthyra/ESMFold2-Fast", - "official_path": "biohub/ESMFold2-Fast", - "model_type": "ESMFold2", - "size_category": "structure", - }, -} - -BACKENDS = ("sdpa", "flex", "kernels_flash") - - -def get_models_by_size(*categories: str) -> Dict[str, Dict]: - return {k: v for k, v in FULL_MODEL_REGISTRY.items() if v["size_category"] in categories} - - -# Pre-built key lists by size category -SMALL_MODEL_KEYS = list(get_models_by_size("small").keys()) -MEDIUM_MODEL_KEYS = list(get_models_by_size("small", "medium").keys()) -LARGE_MODEL_KEYS = list(get_models_by_size("large").keys()) -XLARGE_MODEL_KEYS = list(get_models_by_size("xlarge").keys()) -ALL_FULL_MODEL_KEYS = list(FULL_MODEL_REGISTRY.keys()) -SEQUENCE_MODEL_KEYS = [k for k in ALL_FULL_MODEL_KEYS if FULL_MODEL_REGISTRY[k]["size_category"] != "structure"] -STRUCTURE_MODEL_KEYS = list(STRUCTURE_MODEL_REGISTRY.keys()) - - -def mark_by_size(keys: List[str], registry: Dict[str, Dict], extra_marks: List = None) -> List: - """Return pytest.param list with appropriate markers based on size_category.""" - params = [] - for k in keys: - marks = list(extra_marks or []) - if registry[k]["size_category"] == "xlarge": - marks.append(pytest.mark.large) - elif registry[k]["size_category"] in ("large", "medium"): - marks.append(pytest.mark.slow) - params.append(pytest.param(k, marks=marks)) - return params - - -def tokenize_batch( - model, - model_key: str, - sequences: List[str], - device: torch.device, - registry: Dict[str, Dict] = None, -) -> Dict[str, torch.Tensor]: - """Tokenize a batch of sequences, handling E1's sequence mode. - - Shared helper used across multiple test files to avoid duplication. - """ - if registry is None: - registry = FULL_MODEL_REGISTRY - config = registry[model_key] if model_key in registry else MODEL_REGISTRY[model_key] - - if config["model_type"] == "E1": - batch = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) - return { - "input_ids": batch["input_ids"], - "within_seq_position_ids": batch["within_seq_position_ids"], - "global_position_ids": batch["global_position_ids"], - "sequence_ids": batch["sequence_ids"], - "attention_mask": (batch["sequence_ids"] != -1).long(), - } - tokenizer = model.tokenizer - tokenized = tokenizer(sequences, return_tensors="pt", padding=True) - return {k: v.to(device) for k, v in tokenized.items()} - - -def add_model_specific_inputs( - model_inputs: Dict[str, torch.Tensor], - model_type: str, -) -> Dict[str, torch.Tensor]: - """Add model-specific extra inputs (e.g. sequence_id for ESMC).""" - if model_type == "ESMC": - model_inputs["sequence_id"] = model_inputs["attention_mask"].to(dtype=torch.bool) - return model_inputs - - -def random_sequences(n: int, min_len: int = 8, max_len: int = 64) -> List[str]: - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=random.randint(min_len, max_len))) - for _ in range(n) - ] - - -def random_sequences_fixed_len(n: int, length: int = 64) -> List[str]: - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=length - 1)) - for _ in range(n) - ] - - -def get_device() -> torch.device: - return torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/testing/count_parameters.py b/testing/count_parameters.py deleted file mode 100644 index 6ace4a4..0000000 --- a/testing/count_parameters.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Count parameters for all models in the registry. - -Loads each model from HuggingFace and prints total / trainable parameter counts. -Run on a machine with enough VRAM or use --cpu to load on CPU (slow for large models). - -Usage: - py -m testing.count_parameters - py -m testing.count_parameters --families esm2 ankh - py -m testing.count_parameters --cpu - py -m testing.count_parameters --include-structure -""" - -import argparse -import json -import sys - -import torch -from transformers import AutoModel, AutoModelForMaskedLM - -SEQUENCE_MODELS = { - # ESM2 - "esm2_8m": "Synthyra/ESM2-8M", - "esm2_35m": "Synthyra/ESM2-35M", - "esm2_150m": "Synthyra/ESM2-150M", - "esm2_650m": "Synthyra/ESM2-650M", - "esm2_3b": "Synthyra/ESM2-3B", - # ESM++ - "esmc_small": "Synthyra/ESMplusplus_small", - "esmc_large": "Synthyra/ESMplusplus_large", - "esmc_6b": "Synthyra/ESMplusplus_6B", - # E1 - "e1_150m": "Synthyra/Profluent-E1-150M", - "e1_300m": "Synthyra/Profluent-E1-300M", - "e1_600m": "Synthyra/Profluent-E1-600M", - # DPLM - "dplm_150m": "Synthyra/DPLM-150M", - "dplm_650m": "Synthyra/DPLM-650M", - "dplm_3b": "Synthyra/DPLM-3B", - # DPLM2 - "dplm2_150m": "Synthyra/DPLM2-150M", - "dplm2_650m": "Synthyra/DPLM2-650M", - "dplm2_3b": "Synthyra/DPLM2-3B", - # ANKH - "ankh_base": "Synthyra/ANKH_base", - "ankh_large": "Synthyra/ANKH_large", - "ankh2_large": "Synthyra/ANKH2_large", - "ankh3_large": "Synthyra/ANKH3_large", - "ankh3_xl": "Synthyra/ANKH3_xl", -} - -STRUCTURE_MODELS = { - "boltz2": "Synthyra/Boltz2", - "esmfold": "Synthyra/FastESMFold", -} - -FAMILY_PREFIXES = { - "esm2": "esm2_", - "esmc": "esmc_", - "esm++": "esmc_", - "e1": "e1_", - "dplm2": "dplm2_", - "dplm": "dplm_", - "ankh": "ankh", - "boltz": "boltz", - "esmfold": "esmfold", -} - - -def _format_params(n: int) -> str: - if n >= 1_000_000_000: - return f"{n / 1_000_000_000:.2f}B" - if n >= 1_000_000: - return f"{n / 1_000_000:.1f}M" - if n >= 1_000: - return f"{n / 1_000:.1f}K" - return str(n) - - -def count_params(repo_id: str, device: str) -> dict: - model = AutoModelForMaskedLM.from_pretrained( - repo_id, - trust_remote_code=True, - dtype=torch.float32, - ).to(device) - - total = sum(p.numel() for p in model.parameters()) - trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) - - # Deduplicate tied parameters (e.g. shared embedding/lm_head) - seen = set() - unique_total = 0 - for p in model.parameters(): - ptr = p.data_ptr() - if ptr not in seen: - seen.add(ptr) - unique_total += p.numel() - - del model - if device != "cpu": - torch.cuda.empty_cache() - - return { - "total": total, - "trainable": trainable, - "unique": unique_total, - "total_str": _format_params(total), - "unique_str": _format_params(unique_total), - } - - -def count_structure_params(repo_id: str, device: str) -> dict: - model = AutoModel.from_pretrained( - repo_id, - trust_remote_code=True, - dtype=torch.float32, - ).to(device) - - total = sum(p.numel() for p in model.parameters()) - trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) - - seen = set() - unique_total = 0 - for p in model.parameters(): - ptr = p.data_ptr() - if ptr not in seen: - seen.add(ptr) - unique_total += p.numel() - - del model - if device != "cpu": - torch.cuda.empty_cache() - - return { - "total": total, - "trainable": trainable, - "unique": unique_total, - "total_str": _format_params(total), - "unique_str": _format_params(unique_total), - } - - -def main(): - parser = argparse.ArgumentParser(description="Count parameters for FastPLMs models") - parser.add_argument( - "--families", - nargs="*", - help="Filter to specific families (e.g. esm2 ankh dplm)", - ) - parser.add_argument( - "--cpu", - action="store_true", - help="Load models on CPU instead of CUDA", - ) - parser.add_argument( - "--include-structure", - action="store_true", - help="Also count structure models (Boltz2, ESMFold)", - ) - parser.add_argument( - "--json", - action="store_true", - help="Output as JSON", - ) - args = parser.parse_args() - - device = "cpu" if args.cpu else "cuda" - - models_to_count = {} - if args.families: - for family in args.families: - prefix = FAMILY_PREFIXES[family.lower()] - for key, repo in SEQUENCE_MODELS.items(): - if key.startswith(prefix): - models_to_count[key] = repo - if args.include_structure: - for key, repo in STRUCTURE_MODELS.items(): - if key.startswith(prefix): - models_to_count[key] = repo - else: - models_to_count.update(SEQUENCE_MODELS) - if args.include_structure: - models_to_count.update(STRUCTURE_MODELS) - - results = {} - structure_keys = set(STRUCTURE_MODELS.keys()) - - print(f"{'Model Key':<20} {'Repo ID':<35} {'Total':>12} {'Unique':>12}") - print("-" * 83) - - for key, repo in models_to_count.items(): - try: - if key in structure_keys: - info = count_structure_params(repo, device) - else: - info = count_params(repo, device) - results[key] = {"repo_id": repo, **info} - print(f"{key:<20} {repo:<35} {info['total_str']:>12} {info['unique_str']:>12}") - except Exception as e: - print(f"{key:<20} {repo:<35} {'ERROR':>12} {str(e)[:30]}", file=sys.stderr) - results[key] = {"repo_id": repo, "error": str(e)} - - if args.json: - print("\n" + json.dumps(results, indent=2)) - - -if __name__ == "__main__": - main() diff --git a/testing/debug_scripts/README.md b/testing/debug_scripts/README.md deleted file mode 100644 index d937dfd..0000000 --- a/testing/debug_scripts/README.md +++ /dev/null @@ -1,5 +0,0 @@ -# Parity debug scripts - -One-off investigation scripts written during parity fixes. Not picked up by pytest. - -Keep for reproducing the original investigation if a related regression shows up. If you're adding a new family or investigating a new parity bug, write a similar script here rather than polluting `testing/` with non-test files. diff --git a/testing/debug_scripts/investigate_backend_cosine.py b/testing/debug_scripts/investigate_backend_cosine.py deleted file mode 100644 index 4f76ddf..0000000 --- a/testing/debug_scripts/investigate_backend_cosine.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Diagnostic for backend-consistency regression. - -Compares sdpa vs kernels_flash vs flex on ESM2-8M in bf16. -Measures per-position cosine, pooled cosine, maxabs, and argmax agreement. -""" -from __future__ import annotations - -import random -import torch -import torch.nn.functional as F -from transformers import AutoModelForMaskedLM - -from testing.conftest import CANONICAL_AAS, SEED - - -def gen_sequences(seed: int = SEED): - rng = random.Random(seed) - lengths = [16, 32, 48, 64, 80, 96, 112, 128] - return ["M" + "".join(rng.choices(CANONICAL_AAS, k=L - 1)) for L in lengths] - - -def run(model, sequences, device): - tok = model.tokenizer - enc = tok(sequences, return_tensors="pt", padding=True) - enc = {k: v.to(device) for k, v in enc.items()} - with torch.inference_mode(): - out = model(**enc) - return out.last_hidden_state, out.logits, enc["attention_mask"] - - -def main() -> int: - device = torch.device("cuda") - sequences = gen_sequences() - - print("Loading ESM2-8M in bf16 with sdpa ...") - m_sdpa = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-8M", trust_remote_code=True, - dtype=torch.bfloat16, device_map=device, - ).eval() - m_sdpa.attn_backend = "sdpa" - sdpa_last, sdpa_logits, mask = run(m_sdpa, sequences, device) - del m_sdpa - torch.cuda.empty_cache() - - for backend in ("flex", "kernels_flash"): - print(f"\nLoading ESM2-8M in bf16 with {backend} ...") - m_alt = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-8M", trust_remote_code=True, - dtype=torch.bfloat16, device_map=device, - ).eval() - try: - m_alt.attn_backend = backend - except (AssertionError, RuntimeError) as e: - print(f" SKIP {backend}: {e}") - del m_alt - torch.cuda.empty_cache() - continue - alt_last, alt_logits, _ = run(m_alt, sequences, device) - - # Per-position cosine vs sdpa. - mask_b = mask.bool() - sdpa_valid = sdpa_last[mask_b].float() - alt_valid = alt_last[mask_b].float() - per_pos_cos = F.cosine_similarity(sdpa_valid, alt_valid, dim=-1) - print(f" per-position cosine: min={per_pos_cos.min().item():.4f} mean={per_pos_cos.mean().item():.4f} max={per_pos_cos.max().item():.4f}") - - # Pooled cosine per sequence. - m = mask.bool().unsqueeze(-1).float() - sdpa_pooled = (sdpa_last.float() * m).sum(dim=1) / m.sum(dim=1).clamp_min(1.0) - alt_pooled = (alt_last.float() * m).sum(dim=1) / m.sum(dim=1).clamp_min(1.0) - pooled_cos = F.cosine_similarity(sdpa_pooled, alt_pooled, dim=-1) - print(f" per-seq pooled cosine: {[f'{c.item():.4f}' for c in pooled_cos]}") - - # Raw diffs. - diff = (alt_last.float() - sdpa_last.float())[mask_b] - print(f" raw diff: mse={(diff ** 2).mean().item():.3e} maxabs={diff.abs().max().item():.3e}") - print(f" sdpa magnitudes: mean_norm={sdpa_valid.norm(dim=-1).mean().item():.3f} maxabs={sdpa_valid.abs().max().item():.3e}") - print(f" alt magnitudes: mean_norm={alt_valid.norm(dim=-1).mean().item():.3f} maxabs={alt_valid.abs().max().item():.3e}") - # NaN / Inf check. - print(f" alt has NaN={torch.isnan(alt_last).any().item()} Inf={torch.isinf(alt_last).any().item()}") - - # Logit argmax agreement. - sdpa_argmax = sdpa_logits.float()[mask_b].argmax(dim=-1) - alt_argmax = alt_logits.float()[mask_b].argmax(dim=-1) - agreement = (sdpa_argmax == alt_argmax).float().mean().item() - print(f" argmax agreement: {agreement:.4f}") - - # First divergent position. - nonmatching = (sdpa_argmax != alt_argmax).nonzero(as_tuple=False).flatten() - if len(nonmatching) > 0: - print(f" first {min(5, len(nonmatching))} mismatched positions: {nonmatching[:5].tolist()}") - - del m_alt, alt_last, alt_logits - torch.cuda.empty_cache() - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/investigate_dplm2_tokenize.py b/testing/debug_scripts/investigate_dplm2_tokenize.py deleted file mode 100644 index d6054b6..0000000 --- a/testing/debug_scripts/investigate_dplm2_tokenize.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Figure out why native DPLM2 OOBs on tokens produced by the native tokenizer.""" -from __future__ import annotations - -import torch -from transformers import AutoModelForMaskedLM, EsmForMaskedLM, EsmTokenizer, AutoConfig - - -def main() -> int: - device = torch.device("cuda") - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/DPLM2-150M", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - native = EsmForMaskedLM.from_pretrained( - "airkingbd/dplm2_150m", dtype=torch.float32, device_map=device, - ).eval() - native_tok = EsmTokenizer.from_pretrained("airkingbd/dplm2_150m") - fast_tok = fast.tokenizer - native_cfg = AutoConfig.from_pretrained("airkingbd/dplm2_150m") - - print(f"fast config.vocab_size = {fast.config.vocab_size}") - print(f"native config.vocab_size = {native_cfg.vocab_size}") - print(f"fast word_emb shape: {fast.state_dict()['esm.embeddings.word_embeddings.weight'].shape}") - print(f"native word_emb shape: {native.state_dict()['esm.embeddings.word_embeddings.weight'].shape}") - - print(f"\nfast tok vocab size: {len(fast_tok.get_vocab())}") - print(f"native tok vocab size: {len(native_tok.get_vocab())}") - print(f"fast tok pad_token_id: {fast_tok.pad_token_id} cls: {fast_tok.cls_token_id} eos: {fast_tok.eos_token_id} mask: {fast_tok.mask_token_id}") - print(f"native tok pad_token_id: {native_tok.pad_token_id} cls: {native_tok.cls_token_id} eos: {native_tok.eos_token_id} mask: {native_tok.mask_token_id}") - - seqs = ["MALW", "MKTIIALSY"] - fast_enc = fast_tok(seqs, return_tensors="pt", padding=True) - native_enc = native_tok(seqs, return_tensors="pt", padding=True) - print(f"\nfast_enc input_ids[0] = {fast_enc['input_ids'][0].tolist()}") - print(f"native_enc input_ids[0] = {native_enc['input_ids'][0].tolist()}") - print(f"max fast_enc id = {fast_enc['input_ids'].max().item()}") - print(f"max native_enc id = {native_enc['input_ids'].max().item()}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/investigate_dplm2_weights.py b/testing/debug_scripts/investigate_dplm2_weights.py deleted file mode 100644 index 05333f9..0000000 --- a/testing/debug_scripts/investigate_dplm2_weights.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Check whether DPLM2 fast lm_head matches native DPLM2 lm_head or embedding.""" -from __future__ import annotations - -import torch -from transformers import AutoModelForMaskedLM, EsmForMaskedLM, AutoConfig - - -def main() -> int: - device = torch.device("cuda") - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/DPLM2-150M", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - native = EsmForMaskedLM.from_pretrained( - "airkingbd/dplm2_150m", dtype=torch.float32, device_map=device, - ).eval() - native_cfg = AutoConfig.from_pretrained("airkingbd/dplm2_150m") - - print(f"native tie_word_embeddings: {getattr(native_cfg, 'tie_word_embeddings', None)}") - print(f"fast tie_word_embeddings: {fast.config.tie_word_embeddings}") - - fsd = fast.state_dict() - nsd = native.state_dict() - - f_lm = fsd["lm_head.decoder.weight"] - n_lm = nsd["lm_head.decoder.weight"] - n_emb = nsd["esm.embeddings.word_embeddings.weight"] - f_emb = fsd["esm.embeddings.word_embeddings.weight"] - - print(f"fast lm_head == native lm_head? max|d|={(f_lm - n_lm).abs().max().item():.3e}") - print(f"native lm_head == native word_emb? max|d|={(n_lm - n_emb).abs().max().item():.3e}") - print(f"fast lm_head == native word_emb? max|d|={(f_lm - n_emb).abs().max().item():.3e}") - print(f"fast word_emb == native word_emb? max|d|={(f_emb - n_emb).abs().max().item():.3e}") - print(f"fast word_emb == fast lm_head? max|d|={(f_emb - f_lm).abs().max().item():.3e}") - - # Check biases if present. - for k in ("lm_head.decoder.bias", "lm_head.bias"): - if k in fsd and k in nsd: - print(f"fast[{k}] vs native[{k}] max|d|={(fsd[k] - nsd[k]).abs().max().item():.3e}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_ankh.py b/testing/debug_scripts/parity_debug_ankh.py deleted file mode 100644 index 947048c..0000000 --- a/testing/debug_scripts/parity_debug_ankh.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Trace where FastAnkh diverges from native HuggingFace T5EncoderModel. - -Strategy: - 1. Load both with identical weights, fp32, same input. - 2. Confirm embedding outputs match. - 3. Hook both encoders to capture per-layer hidden states. - 4. Compare layer by layer with both absolute and relative metrics. - 5. Once we find the first diverging layer, dig into that layer's submodules - (layer_norm output, q/k/v projections, attention output, FFN output) - by hooking inside the block. - -Run inside fastplms-ankh image: - docker run --rm --gpus all --ipc=host -v $(pwd):/workspace fastplms-ankh \ - python -m testing.parity_debug_ankh -""" -from __future__ import annotations - -import random -from typing import Dict, List, Tuple - -import torch -import torch.nn as nn - -from testing.conftest import CANONICAL_AAS, SEED - - -SEQUENCES = ["MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVK"] - - -def diff_metrics(fast: torch.Tensor, native: torch.Tensor, tag: str) -> str: - """Both absolute and relative diff metrics.""" - f = fast.float() - n = native.float() - d = f - n - abs_mse = (d ** 2).mean().item() - abs_max = d.abs().max().item() - n_std = n.std().item() - rel = d.std().item() / max(n_std, 1e-12) - return f"{tag:50s} abs_mse={abs_mse:.3e} abs_max={abs_max:.3e} rel_std={rel:.3e} | native_mean={n.mean().item():.3e} native_std={n_std:.3e}" - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - from transformers import AutoModelForMaskedLM, T5EncoderModel, AutoTokenizer - print("Loading fast model...") - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ANKH_base", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - print("Loading native T5EncoderModel...") - native = T5EncoderModel.from_pretrained( - "ElnaggarLab/ankh-base", - device_map=device, - dtype=torch.float32, - ).eval() - tokenizer = AutoTokenizer.from_pretrained("ElnaggarLab/ankh-base") - - # ---- Inputs ---- - enc = tokenizer(SEQUENCES, return_tensors="pt", padding=True) - enc = {k: v.to(device) for k, v in enc.items()} - input_ids = enc["input_ids"] - attention_mask = enc["attention_mask"] - print(f"\ninput_ids shape={tuple(input_ids.shape)}, mask sum={attention_mask.sum().item()}") - - # ---- Locate fast encoder + native encoder ---- - fast_enc = fast.encoder - native_enc = native.encoder if hasattr(native, "encoder") else native - - n_blocks_fast = len(fast_enc.block) - n_blocks_native = len(native_enc.block) - print(f"fast blocks={n_blocks_fast}, native blocks={n_blocks_native}") - assert n_blocks_fast == n_blocks_native, "block count mismatch" - - # ---- Hook all hidden states (output of each block + final norm) ---- - fast_hs: List[torch.Tensor] = [] - native_hs: List[torch.Tensor] = [] - - def make_block_hook(store: List[torch.Tensor]): - def hook(_mod, _inp, out): - # block forward returns (hidden_states, attn_weights, position_bias) for fast - # and (hidden_states, ...) for native T5Block. Take the first element. - if isinstance(out, tuple): - store.append(out[0].detach().clone()) - else: - store.append(out.detach().clone()) - return hook - - handles = [] - for blk in fast_enc.block: - handles.append(blk.register_forward_hook(make_block_hook(fast_hs))) - for blk in native_enc.block: - handles.append(blk.register_forward_hook(make_block_hook(native_hs))) - - # ---- Forward both ---- - print("\nRunning fast forward...") - with torch.no_grad(): - fout = fast(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) - print(f"fast last_hidden_state shape={tuple(fout.last_hidden_state.shape)}") - - print("Running native forward...") - with torch.no_grad(): - nout = native(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) - print(f"native last_hidden_state shape={tuple(nout.last_hidden_state.shape)}") - - for h in handles: - h.remove() - - # ---- Per-block diff ---- - print("\n=== Per-block hidden-state diff (output of each block, BEFORE final norm) ===") - print(f"fast hooked={len(fast_hs)}, native hooked={len(native_hs)}") - for i in range(min(len(fast_hs), len(native_hs))): - print(diff_metrics(fast_hs[i], native_hs[i], f" block {i:02d} output")) - - # ---- Full hidden states tuple from outputs ---- - print("\n=== Per-layer hidden_states from .hidden_states tuple ===") - fh = fout.hidden_states - nh = nout.hidden_states - print(f"fast hs tuple len={len(fh)}, native hs tuple len={len(nh)}") - for i in range(min(len(fh), len(nh))): - print(diff_metrics(fh[i], nh[i], f" hidden_states[{i:02d}]")) - - # ---- Final last_hidden_state ---- - print("\n=== Final last_hidden_state (after final_layer_norm) ===") - print(diff_metrics(fout.last_hidden_state, nout.last_hidden_state, " last_hidden_state")) - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_ankh_mask.py b/testing/debug_scripts/parity_debug_ankh_mask.py deleted file mode 100644 index e8b2ac3..0000000 --- a/testing/debug_scripts/parity_debug_ankh_mask.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Quick test: what does bool_tensor.masked_fill(bool_mask, float('-inf')) do?""" -import torch - -am_4d = torch.tensor([[True, True, True, False, False]])[None, None] -print(f"am_4d dtype={am_4d.dtype} shape={am_4d.shape}") -print(f"am_4d: {am_4d}") - -result = am_4d.masked_fill(am_4d.logical_not(), float("-inf")) -print(f"\nresult dtype={result.dtype}") -print(f"result: {result}") - -# What we actually want -zero_mask = torch.zeros(am_4d.shape, dtype=torch.float32) -zero_mask = zero_mask.masked_fill(am_4d.logical_not(), float("-inf")) -print(f"\nzero_mask: {zero_mask}") - -# Test with SDPA -torch.manual_seed(0) -B, H, Q, K, D = 2, 2, 4, 8, 8 -q = torch.randn(B, H, Q, D, device="cuda", dtype=torch.float32) -k = torch.randn(B, H, K, D, device="cuda", dtype=torch.float32) -v = torch.randn(B, H, K, D, device="cuda", dtype=torch.float32) - -# Mask: first 4 keys valid, last 4 padded for batch 0; all valid for batch 1 -mask_2d = torch.tensor([[1,1,1,1,0,0,0,0],[1,1,1,1,1,1,1,1]], dtype=torch.bool, device="cuda") -mask_4d = mask_2d[:, None, None, :] # (B, 1, 1, K) bool - -# Approach 1: bool masked_fill (current code) -m1 = mask_4d.masked_fill(mask_4d.logical_not(), float("-inf")) -print(f"\nm1 dtype={m1.dtype} m1[0]={m1[0]}") - -# Approach 2: zeros + masked_fill (proper additive mask) -m2 = torch.zeros(mask_4d.shape, dtype=torch.float32, device="cuda") -m2 = m2.masked_fill(mask_4d.logical_not(), float("-inf")) -print(f"m2 dtype={m2.dtype} m2[0]={m2[0]}") - -import torch.nn.functional as F -out1 = F.scaled_dot_product_attention(q, k, v, attn_mask=m1.float() if m1.dtype == torch.bool else m1, scale=1.0) -out2 = F.scaled_dot_product_attention(q, k, v, attn_mask=m2, scale=1.0) -print(f"\nout1 vs out2 (batch=0): mse={((out1[0]-out2[0])**2).mean().item():.3e}") -print(f"out1 vs out2 (batch=1): mse={((out1[1]-out2[1])**2).mean().item():.3e}") - -# Now compare with single-batch (B=1) using just batch 0's data -q_single = q[0:1] -k_single = k[0:1, :, :4, :] -v_single = v[0:1, :, :4, :] -out_single = F.scaled_dot_product_attention(q_single, k_single, v_single, scale=1.0) -print(f"\nout_single shape={out_single.shape}, out1[0:1] shape={out1[0:1].shape}") -print(f"out_single vs out1[0]: mse={((out_single[0]-out1[0])**2).mean().item():.3e} maxabs={(out_single[0]-out1[0]).abs().max().item():.3e}") -print(f"out_single vs out2[0]: mse={((out_single[0]-out2[0])**2).mean().item():.3e} maxabs={(out_single[0]-out2[0]).abs().max().item():.3e}") diff --git a/testing/debug_scripts/parity_debug_ankh_padding.py b/testing/debug_scripts/parity_debug_ankh_padding.py deleted file mode 100644 index 18a9056..0000000 --- a/testing/debug_scripts/parity_debug_ankh_padding.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Trace where ANKH padding-isolation breaks: single seq alone vs batched padded. - -The forward parity tests pass with no padding but fail when padding is introduced. -Compares hidden_states at every layer for the short sequence in two settings: - - alone: forward([short]), output[0, :16] - - padded: forward([short, long_]), output[0, :16] - -If they differ, padding is bleeding into valid attention. -""" -from __future__ import annotations - -import random - -import torch - -from testing.conftest import CANONICAL_AAS, SEED - - -def diff(label, a, b): - a = a.float() - b = b.float() - d = (a - b) - print(f" {label:50s} mse={(d**2).mean().item():.3e} maxabs={d.abs().max().item():.3e}") - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - rng = random.Random(SEED) - short = "M" + "".join(rng.choices(CANONICAL_AAS, k=15)) - long_ = "M" + "".join(rng.choices(CANONICAL_AAS, k=127)) - print(f"short len={len(short)} long len={len(long_)}") - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ANKH_base", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - - tok = fast.tokenizer - enc_alone = tok([short], return_tensors="pt", padding=True) - enc_alone = {k: v.to(device) for k, v in enc_alone.items()} - enc_padded = tok([short, long_], return_tensors="pt", padding=True) - enc_padded = {k: v.to(device) for k, v in enc_padded.items()} - print(f"alone: input_ids shape={tuple(enc_alone['input_ids'].shape)}, mask sum={enc_alone['attention_mask'].sum().item()}") - print(f"padded: input_ids shape={tuple(enc_padded['input_ids'].shape)}, mask sum row0={enc_padded['attention_mask'][0].sum().item()}") - - valid_len = int(enc_alone["attention_mask"].sum().item()) - print(f"valid_len={valid_len}\n") - - with torch.no_grad(): - out_alone = fast(input_ids=enc_alone["input_ids"], attention_mask=enc_alone["attention_mask"], output_hidden_states=True) - out_padded = fast(input_ids=enc_padded["input_ids"], attention_mask=enc_padded["attention_mask"], output_hidden_states=True) - - print("=== Per-layer hidden_states diff: alone[0, :v] vs padded[0, :v] ===") - for i in range(len(out_alone.hidden_states)): - ha = out_alone.hidden_states[i][0, :valid_len] - hp = out_padded.hidden_states[i][0, :valid_len] - diff(f"hidden_states[{i:02d}]", ha, hp) - - print("\n=== Now repeat with manual-attn path to rule out SDPA -inf issues ===") - # Patch in _manual_attn for layer 0 to inspect raw attention scores - from fastplms.ankh.modeling_ankh import AnkhSelfAttention, AttentionBackend - layer0_attn = fast.encoder.block[0].layer[0].SelfAttention - - captured = {} - orig_compute_bias = layer0_attn.compute_bias - - def cb_hook(*args, **kw): - out = orig_compute_bias(*args, **kw) - captured.setdefault("bias", []).append(out.clone()) - return out - - layer0_attn.compute_bias = cb_hook # type: ignore - - with torch.no_grad(): - _ = fast(input_ids=enc_alone["input_ids"], attention_mask=enc_alone["attention_mask"], output_hidden_states=False) - _ = fast(input_ids=enc_padded["input_ids"], attention_mask=enc_padded["attention_mask"], output_hidden_states=False) - layer0_attn.compute_bias = orig_compute_bias # type: ignore - - bias_alone, bias_padded = captured["bias"] - print(f" bias_alone shape = {tuple(bias_alone.shape)}") - print(f" bias_padded shape = {tuple(bias_padded.shape)}") - if bias_padded.shape[-1] >= valid_len and bias_padded.shape[-2] >= valid_len: - diff("bias[upper_left v x v] alone vs padded", bias_alone[:, :, :valid_len, :valid_len], bias_padded[:, :, :valid_len, :valid_len]) - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_block_internals.py b/testing/debug_scripts/parity_debug_block_internals.py deleted file mode 100644 index 13dbf71..0000000 --- a/testing/debug_scripts/parity_debug_block_internals.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Instrument each sub-operation of block 0 and block 1 to pinpoint where -FastPLMs and native ESMC diverge in fp32. - -Methodology: grab intermediate tensors from both implementations via forward -hooks on corresponding submodules, then compare element-wise. -""" -from __future__ import annotations - -import random -import torch -from torch.nn.attention import SDPBackend, sdpa_kernel - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -def diff_stats(label: str, a: torch.Tensor, b: torch.Tensor) -> None: - a = a.float() - b = b.float() - d = a - b - print(f" {label:40s} mse={((d**2).mean().item()):.3e} maxabs={d.abs().max().item():.3e} shape={tuple(a.shape)}") - - -class Capture: - """Hook that stores the module's input and output.""" - def __init__(self): - self.inputs = None - self.outputs = None - def __call__(self, module, inputs, outputs): - self.inputs = tuple(x.detach().clone() if isinstance(x, torch.Tensor) else x for x in inputs) - self.outputs = outputs.detach().clone() if isinstance(outputs, torch.Tensor) else outputs - - -def attach(mod, cap): - return mod.register_forward_hook(cap) - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - seq = "M" + "".join(random.choices(CANONICAL_AAS, k=63)) - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - native, _ = load_native_esmc(reference_repo_id="esmc-300", device=device, dtype=torch.float32) - - enc = fast.tokenizer([seq], return_tensors="pt", padding=False) - enc = {k: v.to(device) for k, v in enc.items()} - - fast_caps = {} - native_caps = {} - hooks = [] - for i in (0, 1): - fast_block = fast.transformer.blocks[i] - native_block = native.model.transformer.blocks[i] - - for name, fmod, nmod in [ - ("attn.layernorm_qkv", fast_block.attn.layernorm_qkv, native_block.attn.layernorm_qkv), - ("attn.q_ln", fast_block.attn.q_ln, native_block.attn.q_ln), - ("attn.k_ln", fast_block.attn.k_ln, native_block.attn.k_ln), - ("attn.rotary", fast_block.attn.rotary, native_block.attn.rotary), - ("attn.out_proj", fast_block.attn.out_proj, native_block.attn.out_proj), - ("attn", fast_block.attn, native_block.attn), - ("ffn", fast_block.ffn, native_block.ffn), - ]: - fcap = Capture() - ncap = Capture() - fast_caps[(i, name)] = fcap - native_caps[(i, name)] = ncap - hooks.append(attach(fmod, fcap)) - hooks.append(attach(nmod, ncap)) - - with torch.no_grad(), sdpa_kernel(SDPBackend.MATH): - fo = fast(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"], output_hidden_states=True) - no = native(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) - - for h in hooks: - h.remove() - - for i in (0, 1): - print(f"\n== block {i} ==") - diff_stats(f"block {i} INPUT (x)", fast_caps[(i, "attn")].inputs[0], native_caps[(i, "attn")].inputs[0]) - diff_stats(f" attn.layernorm_qkv OUTPUT", fast_caps[(i, "attn.layernorm_qkv")].outputs, native_caps[(i, "attn.layernorm_qkv")].outputs) - diff_stats(f" attn.q_ln OUTPUT", fast_caps[(i, "attn.q_ln")].outputs, native_caps[(i, "attn.q_ln")].outputs) - diff_stats(f" attn.k_ln OUTPUT", fast_caps[(i, "attn.k_ln")].outputs, native_caps[(i, "attn.k_ln")].outputs) - - f_rot_in_q = fast_caps[(i, "attn.rotary")].inputs[0] - n_rot_in_q = native_caps[(i, "attn.rotary")].inputs[0] - diff_stats(f" attn.rotary INPUT q", f_rot_in_q, n_rot_in_q) - f_rot_out_q = fast_caps[(i, "attn.rotary")].outputs[0] - n_rot_out_q = native_caps[(i, "attn.rotary")].outputs[0] - diff_stats(f" attn.rotary OUTPUT q", f_rot_out_q, n_rot_out_q) - - diff_stats(f" attn.out_proj INPUT", fast_caps[(i, "attn.out_proj")].inputs[0], native_caps[(i, "attn.out_proj")].inputs[0]) - diff_stats(f" attn.out_proj OUTPUT", fast_caps[(i, "attn.out_proj")].outputs, native_caps[(i, "attn.out_proj")].outputs) - diff_stats(f" attn OUTPUT (full)", fast_caps[(i, "attn")].outputs[0], native_caps[(i, "attn")].outputs) - - diff_stats(f" ffn INPUT", fast_caps[(i, "ffn")].inputs[0], native_caps[(i, "ffn")].inputs[0]) - diff_stats(f" ffn OUTPUT", fast_caps[(i, "ffn")].outputs, native_caps[(i, "ffn")].outputs) - diff_stats(f"block {i} OUTPUT", fo.hidden_states[i], no.hidden_states[i]) - - print("\nRotary cos/sin caches at block 0 vs block 1 within FastPLMs (sanity):") - f0 = fast.transformer.blocks[0].attn.rotary - f1 = fast.transformer.blocks[1].attn.rotary - n0 = native.model.transformer.blocks[0].attn.rotary - n1 = native.model.transformer.blocks[1].attn.rotary - diff_stats("fast b0 cos vs native b0 cos", f0._cos_cached, n0._cos_cached) - diff_stats("fast b1 cos vs native b1 cos", f1._cos_cached, n1._cos_cached) - diff_stats("fast b0 cos vs fast b1 cos", f0._cos_cached, f1._cos_cached) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_diff_structure.py b/testing/debug_scripts/parity_debug_diff_structure.py deleted file mode 100644 index 363ae8a..0000000 --- a/testing/debug_scripts/parity_debug_diff_structure.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Decompose the diff between FastPLMs and native ESMC hidden states. - -Test what LayerNorm actually absorbs by APPLYING the native norm manually -on top of fast_29 and native_29 and seeing what happens. -""" -from __future__ import annotations - -import random -import torch - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -def per_position_affine_fit(fast: torch.Tensor, native: torch.Tensor): - mean_n = native.mean(dim=-1, keepdim=True) - mean_f = fast.mean(dim=-1, keepdim=True) - n_c = native - mean_n - f_c = fast - mean_f - cov = (f_c * n_c).mean(dim=-1, keepdim=True) - var = (n_c * n_c).mean(dim=-1, keepdim=True).clamp_min(1e-12) - beta = cov / var - alpha = mean_f - beta * mean_n - fitted = alpha + beta * native - residual = fast - fitted - return alpha.squeeze(-1), beta.squeeze(-1), residual - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - seq = "M" + "".join(random.choices(CANONICAL_AAS, k=63)) - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - - native, _ = load_native_esmc(reference_repo_id="esmc-300", device=device, dtype=torch.float32) - - enc = fast.tokenizer([seq], return_tensors="pt", padding=False) - enc = {k: v.to(device) for k, v in enc.items()} - - with torch.no_grad(): - fout = fast(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"], output_hidden_states=True) - nout = native(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) - - f29 = fout.hidden_states[29].float().clone() - n29 = nout.hidden_states[29].float().clone() - f30 = fout.hidden_states[30].float().clone() - n30 = nout.hidden_states[30].float().clone() - - print(f"f29 vs n29: mse={((f29-n29)**2).mean().item():.3e} maxabs={(f29-n29).abs().max().item():.3e}") - print(f"f30 vs n30: mse={((f30-n30)**2).mean().item():.3e} maxabs={(f30-n30).abs().max().item():.3e}") - - fast_norm = fast.transformer.norm - native_norm = native.model.transformer.norm - - print(f"\n fast norm weight max abs: {fast_norm.weight.abs().max().item():.6f}") - print(f" native norm weight max abs: {native_norm.weight.abs().max().item():.6f}") - print(f" weights equal: {torch.equal(fast_norm.weight, native_norm.weight)}") - print(f" fast norm bias None: {fast_norm.bias is None} native: {native_norm.bias is None}") - - manual_fast30 = fast_norm(f29) - manual_native30 = native_norm(n29) - print(f"\n manual fast_norm(f29) vs fout.hidden_states[30]:" - f" mse={((manual_fast30 - f30)**2).mean().item():.3e}") - print(f" manual native_norm(n29) vs nout.hidden_states[30]:" - f" mse={((manual_native30 - n30)**2).mean().item():.3e}") - print(f" manual fast_norm(f29) vs manual native_norm(n29):" - f" mse={((manual_fast30 - manual_native30)**2).mean().item():.3e}") - - cross_fn = fast_norm(n29) - cross_nf = native_norm(f29) - print(f" cross fast_norm(n29) vs nout.hidden_states[30]:" - f" mse={((cross_fn - n30)**2).mean().item():.3e}") - print(f" cross native_norm(f29) vs fout.hidden_states[30]:" - f" mse={((cross_nf - f30)**2).mean().item():.3e}") - - print("\n alpha/beta at a few positions (from affine fit):") - alpha, beta, residual = per_position_affine_fit(f29, n29) - print(f" beta per-position: min={beta.min().item():.6f} max={beta.max().item():.6f} mean={beta.mean().item():.6f}") - print(f" alpha per-position abs max: {alpha.abs().max().item():.6f}") - print(f" residual mse: {(residual**2).mean().item():.3e} ({((residual**2).mean().item() / ((f29-n29)**2).mean().item()) * 100:.1f}% of total)") - - print("\n check: std of diff per position vs std of native per position:") - diff = f29 - n29 - diff_std = diff.std(dim=-1).squeeze(0) - native_std = n29.std(dim=-1).squeeze(0) - ratio = diff_std / native_std - print(f" ratio diff_std / native_std: min={ratio.min().item():.4f} max={ratio.max().item():.4f} mean={ratio.mean().item():.4f}") - - print("\n compare f29 and n29 row-wise (per position):") - for pos in [0, 1, 10, 32, 64]: - if pos >= f29.shape[1]: - continue - fp = f29[0, pos] - np = n29[0, pos] - print(f" pos {pos}: mean_f={fp.mean().item():+.4f} std_f={fp.std().item():.4f} " - f"mean_n={np.mean().item():+.4f} std_n={np.std().item():.4f} " - f"mean_diff={(fp - np).mean().item():+.4e} corr={torch.corrcoef(torch.stack([fp, np]))[0,1].item():.6f}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_esm2.py b/testing/debug_scripts/parity_debug_esm2.py deleted file mode 100644 index be50299..0000000 --- a/testing/debug_scripts/parity_debug_esm2.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Sanity-check parity on ESM2 to confirm large absolute per-layer diffs are -architectural (huge activation magnitudes + LN), not an ESMC-specific bug.""" -from __future__ import annotations - -import random -import torch - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm2 import load_official_model as load_native_esm2 - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - seq = "M" + "".join(random.choices(CANONICAL_AAS, k=63)) - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESM2-8M", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - - native, tokenizer = load_native_esm2( - reference_repo_id="facebook/esm2_t6_8M_UR50D", - device=device, dtype=torch.float32, - ) - - enc = tokenizer([seq], return_tensors="pt", padding=False) - enc = {k: v.to(device) for k, v in enc.items()} - - with torch.no_grad(): - fout = fast(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"], output_hidden_states=True) - nout = native(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"], output_hidden_states=True) - - print(f"fast #hiddens={len(fout.hidden_states)} native #hiddens={len(nout.hidden_states)}") - n = min(len(fout.hidden_states), len(nout.hidden_states)) - for i in range(n): - f = fout.hidden_states[i].float() - nt = nout.hidden_states[i].float() - diff = f - nt - mse = (diff ** 2).mean().item() - maxabs = diff.abs().max().item() - native_std = nt.std(dim=-1).mean().item() - diff_std = diff.std(dim=-1).mean().item() - rel = diff_std / max(native_std, 1e-12) - print(f" layer {i:2d}: mse={mse:.3e} maxabs={maxabs:.3e} " - f"native_std={native_std:8.3f} rel_diff={rel:.3e}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_esmc.py b/testing/debug_scripts/parity_debug_esmc.py deleted file mode 100644 index d4d2f14..0000000 --- a/testing/debug_scripts/parity_debug_esmc.py +++ /dev/null @@ -1,414 +0,0 @@ -""" -Standalone ESMC parity investigation — NOT a pytest test. - -Run inside the fastplms Docker image: - python /app/testing/parity_debug_esmc.py - -Goal: surface where FastPLMs ESMC diverges from native Biohub ESMC. - -What this script checks, in order, printing a clear PASS/FAIL per check: - 1. Tokenizer vocab parity (size, every token mapping, special token IDs). - 2. Tokenization of fixed sequences produces identical input_ids. - 3. State dict parity in fp32 (per-parameter MSE, no aggregate). - 4. Forward parity in fp32 per layer, SDPA backend, with/without padding. - 5. Forward parity in bf16 per layer, SDPA backend. - 6. Forward parity across attention backends {sdpa, kernels_flash, flex} (fp32). - 7. `embed_dataset()` pipeline output vs a manual native-wrapped pipeline. - -The point is: if everything prints PASS and downstream parity is still off on a -real task, the issue is not in the encoder — look at how embeddings are consumed. -If anything prints FAIL, the message names which layer / which parameter / which -input diverged and by how much. -""" - -from __future__ import annotations - -import random -import sys -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import torch -import torch.nn as nn -from torch.nn.functional import mse_loss - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -FAST_PATH = "Synthyra/ESMplusplus_small" -NATIVE_REF = "esmc-300" - -PAD_TOK = "" - -FP32_PARAM_MSE_TOL = 0.0 -FP32_HIDDEN_MSE_TOL = 5e-8 -FP32_HIDDEN_MAXABS_TOL = 5e-4 -BF16_HIDDEN_MSE_TOL = 1e-3 -BF16_HIDDEN_MAXABS_TOL = 5e-2 -BACKEND_MSE_TOL = 1e-5 -BACKEND_MAXABS_TOL = 5e-3 - - -def banner(msg: str) -> None: - print("=" * 78) - print(msg) - print("=" * 78) - - -def check(name: str, ok: bool, detail: str = "") -> bool: - tag = "PASS" if ok else "FAIL" - print(f" [{tag}] {name}{(': ' + detail) if detail else ''}") - return ok - - -def gen_sequences(seed: int, n: int, lengths: List[int]) -> List[str]: - rng = random.Random(seed) - assert len(lengths) == n, f"{len(lengths)} != {n}" - out: List[str] = [] - for L in lengths: - out.append("M" + "".join(rng.choices(CANONICAL_AAS, k=L - 1))) - return out - - -@dataclass -class ForwardOutputs: - last_hidden_state: torch.Tensor - hidden_states: Tuple[torch.Tensor, ...] - logits: Optional[torch.Tensor] - - -def load_fast(dtype: torch.dtype, device: torch.device, attn_backend: str = "sdpa") -> nn.Module: - from fastplms.esm_plusplus.modeling_esm_plusplus import ESMplusplusConfig - from transformers import AutoModelForMaskedLM - - ESMplusplusConfig.attn_backend = attn_backend - model = AutoModelForMaskedLM.from_pretrained( - FAST_PATH, - trust_remote_code=True, - dtype=dtype, - device_map=device, - ).eval() - if hasattr(model, "transformer"): - trans = model.transformer - elif hasattr(model, "esm") and hasattr(model.esm, "transformer"): - trans = model.esm.transformer - else: - trans = None - if trans is not None and hasattr(trans, "set_attention_backend"): - trans.set_attention_backend(attn_backend) - return model - - -def fast_forward(model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> ForwardOutputs: - sequence_id = attention_mask.to(dtype=torch.bool) - out = model( - input_ids=input_ids, - attention_mask=attention_mask, - sequence_id=sequence_id, - output_hidden_states=True, - ) - last = out.last_hidden_state if out.last_hidden_state is not None else out.hidden_states[-1] - return ForwardOutputs( - last_hidden_state=last, - hidden_states=tuple(out.hidden_states), - logits=getattr(out, "logits", None), - ) - - -def native_forward(model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> ForwardOutputs: - out = model(input_ids=input_ids, attention_mask=attention_mask) - return ForwardOutputs( - last_hidden_state=out.last_hidden_state, - hidden_states=tuple(out.hidden_states), - logits=out.logits, - ) - - -def hidden_mse_maxabs(a: torch.Tensor, b: torch.Tensor, mask: torch.Tensor) -> Tuple[float, float]: - a_valid = a[mask] - b_valid = b[mask] - mse = mse_loss(a_valid.float(), b_valid.float()).item() - maxabs = (a_valid.float() - b_valid.float()).abs().max().item() - return mse, maxabs - - -def check_tokenizer_parity(fast_model: nn.Module, native_tokenizer) -> bool: - banner("1. Tokenizer vocab parity") - ft = fast_model.tokenizer - all_pass = True - - fast_vocab = ft.get_vocab() - native_vocab = native_tokenizer.get_vocab() - all_pass &= check( - "vocab size equal", - len(fast_vocab) == len(native_vocab), - f"fast={len(fast_vocab)} native={len(native_vocab)}", - ) - - for tok, nid in sorted(native_vocab.items(), key=lambda kv: kv[1]): - fid = fast_vocab.get(tok, None) - if fid != nid: - check(f"token {tok!r} id match", False, f"native={nid} fast={fid}") - all_pass = False - if all_pass: - check("every token id matches", True, f"{len(native_vocab)} tokens") - - for name in ("pad_token_id", "cls_token_id", "eos_token_id", "mask_token_id", "unk_token_id", "bos_token_id"): - fast_v = getattr(ft, name, None) - native_v = getattr(native_tokenizer, name, None) - all_pass &= check(f"{name} match", fast_v == native_v, f"fast={fast_v} native={native_v}") - return all_pass - - -def tokenize_fast(tokenizer, sequences: List[str], device: torch.device) -> Dict[str, torch.Tensor]: - enc = tokenizer(sequences, return_tensors="pt", padding=True) - return {k: v.to(device) for k, v in enc.items()} - - -def check_tokenization(fast_model: nn.Module, native_tokenizer, sequences: List[str], device: torch.device) -> Tuple[bool, Dict[str, torch.Tensor]]: - banner("2. Tokenization produces identical input_ids") - fast_enc = tokenize_fast(fast_model.tokenizer, sequences, device) - native_enc = tokenize_fast(native_tokenizer, sequences, device) - ok_ids = torch.equal(fast_enc["input_ids"], native_enc["input_ids"]) - ok_mask = torch.equal(fast_enc["attention_mask"], native_enc["attention_mask"]) - ok = check("input_ids exact", ok_ids, f"shape={tuple(fast_enc['input_ids'].shape)}") - ok &= check("attention_mask exact", ok_mask) - if not ok_ids: - diff = (fast_enc["input_ids"] != native_enc["input_ids"]).nonzero(as_tuple=False) - print(f" first 5 mismatches: {diff[:5].tolist()}") - return ok, fast_enc - - -def check_weight_parity(fast_sd: Dict[str, torch.Tensor], native_sd: Dict[str, torch.Tensor]) -> bool: - banner("3. State dict parity (fp32, per-parameter)") - all_pass = True - - fast_keys = set(fast_sd.keys()) - native_keys = set(native_sd.keys()) - if fast_keys != native_keys: - only_fast = sorted(fast_keys - native_keys) - only_native = sorted(native_keys - fast_keys) - check("key sets match", False, f"only_fast={only_fast[:3]} only_native={only_native[:3]}") - all_pass = False - else: - check("key sets match", True, f"{len(fast_keys)} parameters") - - failed: List[str] = [] - for name in sorted(fast_keys & native_keys): - a = fast_sd[name].float() - b = native_sd[name].float() - if a.shape != b.shape: - failed.append(f"{name}: shape {tuple(a.shape)} vs {tuple(b.shape)}") - continue - diff = (a - b).abs().max().item() - if diff > FP32_PARAM_MSE_TOL: - mse = mse_loss(a, b).item() - failed.append(f"{name}: max|Δ|={diff:.3e} mse={mse:.3e}") - all_pass &= check( - f"every parameter matches (tol {FP32_PARAM_MSE_TOL})", - not failed, - f"{len(failed)} divergent; first: {failed[:3]}", - ) - return all_pass - - -def check_forward_parity( - fast_model: nn.Module, - native_model: nn.Module, - enc: Dict[str, torch.Tensor], - dtype: torch.dtype, - mse_tol: float, - maxabs_tol: float, - tag: str, -) -> bool: - banner(f"{tag}") - mask = enc["attention_mask"].bool() - - with torch.inference_mode(): - fo = fast_forward(fast_model, enc["input_ids"], enc["attention_mask"]) - no = native_forward(native_model, enc["input_ids"], enc["attention_mask"]) - - all_pass = True - all_pass &= check( - "hidden_states tuple length matches", - len(fo.hidden_states) == len(no.hidden_states), - f"fast={len(fo.hidden_states)} native={len(no.hidden_states)}", - ) - - n = min(len(fo.hidden_states), len(no.hidden_states)) - failures: List[str] = [] - for i in range(n): - mse, maxabs = hidden_mse_maxabs(fo.hidden_states[i], no.hidden_states[i], mask) - if mse > mse_tol or maxabs > maxabs_tol: - failures.append(f"layer {i}: mse={mse:.3e} maxabs={maxabs:.3e}") - print(f" layer {i:2d}: mse={mse:.3e} maxabs={maxabs:.3e}") - all_pass &= check( - f"every layer within tol (mse<{mse_tol}, maxabs<{maxabs_tol})", - not failures, - f"{len(failures)} divergent layers", - ) - - mse, maxabs = hidden_mse_maxabs(fo.last_hidden_state, no.last_hidden_state, mask) - all_pass &= check( - "last_hidden_state within tol", - mse <= mse_tol and maxabs <= maxabs_tol, - f"mse={mse:.3e} maxabs={maxabs:.3e}", - ) - - if fo.logits is not None and no.logits is not None: - mse, maxabs = hidden_mse_maxabs(fo.logits, no.logits, mask) - print(f" logits: mse={mse:.3e} maxabs={maxabs:.3e}") - return all_pass - - -def check_backend_consistency( - fast_model_sdpa: nn.Module, - enc: Dict[str, torch.Tensor], - dtype: torch.dtype, - device: torch.device, -) -> bool: - banner(f"6. Attention backend consistency (dtype={dtype})") - mask = enc["attention_mask"].bool() - - with torch.inference_mode(): - fo_sdpa = fast_forward(fast_model_sdpa, enc["input_ids"], enc["attention_mask"]) - - all_pass = True - for backend in ("kernels_flash", "flex"): - try: - fm_alt = load_fast(dtype=dtype, device=device, attn_backend=backend) - with torch.inference_mode(): - fo_alt = fast_forward(fm_alt, enc["input_ids"], enc["attention_mask"]) - except Exception as e: - check(f"backend {backend} loadable+runnable", False, str(e)[:120]) - all_pass = False - continue - - mse, maxabs = hidden_mse_maxabs(fo_alt.last_hidden_state, fo_sdpa.last_hidden_state, mask) - all_pass &= check( - f"backend {backend} ≈ sdpa (last_hidden_state)", - mse <= BACKEND_MSE_TOL and maxabs <= BACKEND_MAXABS_TOL, - f"mse={mse:.3e} maxabs={maxabs:.3e}", - ) - del fm_alt - torch.cuda.empty_cache() - return all_pass - - -def check_embed_dataset_pipeline( - fast_model: nn.Module, - native_model: nn.Module, - sequences: List[str], - dtype: torch.dtype, - device: torch.device, -) -> bool: - banner(f"7. embed_dataset() pipeline vs manual native (dtype={dtype})") - fast_embeddings = fast_model.embed_dataset( - sequences=sequences, - tokenizer=fast_model.tokenizer, - batch_size=4, - max_len=max(len(s) for s in sequences) + 2, - truncate=True, - full_embeddings=False, - embed_dtype=torch.float32, - pooling_types=["mean"], - num_workers=0, - sql=False, - save=False, - padding="longest", - ) - assert fast_embeddings is not None - - native_tokenizer = native_model.tokenizer - native_embeddings: Dict[str, torch.Tensor] = {} - for seq in sequences: - enc = tokenize_fast(native_tokenizer, [seq], device) - with torch.inference_mode(): - out = native_forward(native_model, enc["input_ids"], enc["attention_mask"]) - m = enc["attention_mask"].bool().unsqueeze(-1).float() - pooled = (out.last_hidden_state.float() * m).sum(dim=1) / m.sum(dim=1).clamp_min(1.0) - native_embeddings[seq] = pooled.squeeze(0).cpu() - - failures: List[str] = [] - for seq in sequences: - f = fast_embeddings[seq].cpu().float() - n = native_embeddings[seq].cpu().float() - if f.shape != n.shape: - failures.append(f"seq_len={len(seq)}: shape {tuple(f.shape)} vs {tuple(n.shape)}") - continue - diff = (f - n).abs().max().item() - mse = mse_loss(f, n).item() - print(f" seq_len={len(seq):3d}: mse={mse:.3e} maxabs={diff:.3e}") - if mse > BF16_HIDDEN_MSE_TOL or diff > BF16_HIDDEN_MAXABS_TOL: - failures.append(f"seq_len={len(seq)}: mse={mse:.3e} maxabs={diff:.3e}") - - return check( - f"mean-pool embedding parity (mse<{BF16_HIDDEN_MSE_TOL})", - not failures, - f"{len(failures)} divergent sequences", - ) - - -def main() -> int: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - assert device.type == "cuda", "This script requires a CUDA GPU." - random.seed(SEED) - torch.manual_seed(SEED) - - lengths = [16, 32, 48, 64, 80, 96, 112, 128] - sequences = gen_sequences(SEED, len(lengths), lengths) - - native_fp32, native_tok = load_native_esmc( - reference_repo_id=NATIVE_REF, device=device, dtype=torch.float32, - ) - fast_fp32 = load_fast(dtype=torch.float32, device=device, attn_backend="sdpa") - - overall: List[bool] = [] - - overall.append(check_tokenizer_parity(fast_fp32, native_tok)) - ok, enc = check_tokenization(fast_fp32, native_tok, sequences, device) - overall.append(ok) - overall.append(check_weight_parity(fast_fp32.state_dict(), native_fp32.model.state_dict())) - overall.append(check_forward_parity( - fast_fp32, native_fp32, enc, torch.float32, - mse_tol=FP32_HIDDEN_MSE_TOL, maxabs_tol=FP32_HIDDEN_MAXABS_TOL, - tag="4. Forward parity (fp32, sdpa)", - )) - - del fast_fp32, native_fp32 - torch.cuda.empty_cache() - - fast_bf16 = load_fast(dtype=torch.bfloat16, device=device, attn_backend="sdpa") - native_bf16, _ = load_native_esmc( - reference_repo_id=NATIVE_REF, device=device, dtype=torch.bfloat16, - ) - overall.append(check_forward_parity( - fast_bf16, native_bf16, enc, torch.bfloat16, - mse_tol=BF16_HIDDEN_MSE_TOL, maxabs_tol=BF16_HIDDEN_MAXABS_TOL, - tag="5. Forward parity (bf16, sdpa)", - )) - - overall.append(check_backend_consistency(fast_bf16, enc, torch.bfloat16, device)) - - overall.append(check_embed_dataset_pipeline( - fast_bf16, native_bf16, sequences, torch.bfloat16, device, - )) - - banner("SUMMARY") - labels = [ - "1. tokenizer vocab", - "2. tokenization ids", - "3. weight parity (fp32)", - "4. forward parity (fp32, sdpa)", - "5. forward parity (bf16, sdpa)", - "6. backend consistency", - "7. embed_dataset pipeline", - ] - for lbl, ok in zip(labels, overall): - print(f" [{'PASS' if ok else 'FAIL'}] {lbl}") - return 0 if all(overall) else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/testing/debug_scripts/parity_debug_esmc_minimal.py b/testing/debug_scripts/parity_debug_esmc_minimal.py deleted file mode 100644 index 3f2870a..0000000 --- a/testing/debug_scripts/parity_debug_esmc_minimal.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Minimal ESMC parity — single sequence, no padding, fp32, sdpa. - -Answers: does the per-layer divergence go away without padding? -""" -from __future__ import annotations - -import random -import torch -from torch.nn.functional import mse_loss - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -def fast_hidden_states(model, input_ids, attention_mask): - out = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) - return tuple(out.hidden_states), out.last_hidden_state - - -def native_hidden_states(model, input_ids, attention_mask): - out = model(input_ids=input_ids, attention_mask=attention_mask) - return tuple(out.hidden_states), out.last_hidden_state - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - seq = "M" + "".join(random.choices(CANONICAL_AAS, k=63)) - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - - native, _ = load_native_esmc(reference_repo_id="esmc-300", device=device, dtype=torch.float32) - - enc = fast.tokenizer([seq], return_tensors="pt", padding=False) - enc = {k: v.to(device) for k, v in enc.items()} - - print(f"seq_len={enc['input_ids'].shape[1]} (no padding)") - - with torch.inference_mode(): - fh, flast = fast_hidden_states(fast, enc["input_ids"], enc["attention_mask"]) - nh, nlast = native_hidden_states(native, enc["input_ids"], enc["attention_mask"]) - - print(f"len(fast_hs)={len(fh)} len(native_hs)={len(nh)}") - n = min(len(fh), len(nh)) - for i in range(n): - diff = (fh[i] - nh[i]).abs() - mse = mse_loss(fh[i].float(), nh[i].float()).item() - print(f" layer {i:2d}: mse={mse:.3e} maxabs={diff.max().item():.3e}") - diff = (flast - nlast).abs() - mse = mse_loss(flast.float(), nlast.float()).item() - print(f" last: mse={mse:.3e} maxabs={diff.max().item():.3e}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_kernel.py b/testing/debug_scripts/parity_debug_kernel.py deleted file mode 100644 index cc8d2f8..0000000 --- a/testing/debug_scripts/parity_debug_kernel.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Pinpoint the root cause of FastPLMs vs native ESMC per-layer drift. - -Hypotheses to test, each with its own column: - A. Baseline: FastPLMs sdpa + default PyTorch kernel dispatch. - B. Force MATH kernel on FastPLMs (bit-deterministic reference): if native - also uses MATH, A-vs-native will match B-vs-native only if dispatch is - identical. If B matches native much better than A, dispatch differs. - C. Force MATH kernel on BOTH fast and native: if this makes per-layer - near-zero, the drift is pure SDPA-kernel-dispatch numerics (no bug). - D. Force MATH kernel on fast, no mask at all: rules out mask-shape bias. - -Run: - docker run --gpus all --ipc=host --rm -v $(pwd):/workspace \ - fastplms-esm_plusplus python /workspace/testing/parity_debug_kernel.py -""" -from __future__ import annotations - -import random -from contextlib import nullcontext - -import torch -from torch.nn.attention import SDPBackend, sdpa_kernel - -from testing.conftest import CANONICAL_AAS, SEED -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -def per_layer_mse(fh, nh): - return [((fh[i].float() - nh[i].float()) ** 2).mean().item() for i in range(min(len(fh), len(nh)))] - - -def run_pair(fast, native, enc, fast_backend, native_backend, label): - fast_ctx = sdpa_kernel(fast_backend) if fast_backend is not None else nullcontext() - native_ctx = sdpa_kernel(native_backend) if native_backend is not None else nullcontext() - try: - with torch.no_grad(): - with fast_ctx: - fo = fast( - input_ids=enc["input_ids"], - attention_mask=enc.get("attention_mask"), - output_hidden_states=True, - ) - with native_ctx: - no = native(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) - mses = per_layer_mse(fo.hidden_states, no.hidden_states) - print(f" {label:50s} L0={mses[0]:.2e} L1={mses[1]:.2e} L15={mses[15]:.2e} L29={mses[29]:.2e} L30={mses[30]:.2e}") - return mses - except Exception as e: - print(f" {label:50s} skipped: {type(e).__name__}: {str(e)[:80]}") - return None - - -def main() -> int: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - seq = "M" + "".join(random.choices(CANONICAL_AAS, k=63)) - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - - native, _ = load_native_esmc(reference_repo_id="esmc-300", device=device, dtype=torch.float32) - - enc = fast.tokenizer([seq], return_tensors="pt", padding=False) - enc = {k: v.to(device) for k, v in enc.items()} - enc_no_mask = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]} - - print("Per-layer MSE between FastPLMs and native ESMC at selected layers (fp32, no padding):\n") - - run_pair(fast, native, enc, None, None, "A. default / default") - run_pair(fast, native, enc, SDPBackend.MATH, None, "B. fast=MATH / native=default") - run_pair(fast, native, enc, None, SDPBackend.MATH, "C. fast=default / native=MATH") - run_pair(fast, native, enc, SDPBackend.MATH, SDPBackend.MATH, "D. both=MATH (deterministic)") - run_pair(fast, native, enc, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, "E. both=EFFICIENT") - run_pair(fast, native, enc, SDPBackend.FLASH_ATTENTION, SDPBackend.FLASH_ATTENTION, "F. both=FLASH") - - print("\nWhich kernels get dispatched by default?") - - def probe_fast(): - x = fast.embed(enc["input_ids"]) - block = fast.transformer.blocks[0] - q = torch.randn(1, 15, enc["input_ids"].shape[1], 64, device=device) - k = torch.randn_like(q) - v = torch.randn_like(q) - attn_mask = enc["attention_mask"][:, None, None, :].bool() - for backend in (SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH, SDPBackend.CUDNN_ATTENTION): - try: - with sdpa_kernel(backend): - _ = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) - print(f" FastPLMs mask=(B,1,1,L) attn_mask=bool: {backend.name}: OK") - except Exception as e: - print(f" FastPLMs mask=(B,1,1,L) attn_mask=bool: {backend.name}: {type(e).__name__}: {str(e)[:80]}") - attn_mask2 = enc["attention_mask"][:, None, :, None].bool() & enc["attention_mask"][:, None, None, :].bool() - for backend in (SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH, SDPBackend.CUDNN_ATTENTION): - try: - with sdpa_kernel(backend): - _ = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask2) - print(f" native mask=(B,1,L,L) attn_mask=bool: {backend.name}: OK") - except Exception as e: - print(f" native mask=(B,1,L,L) attn_mask=bool: {backend.name}: {type(e).__name__}: {str(e)[:80]}") - - probe_fast() - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/parity_debug_rotary.py b/testing/debug_scripts/parity_debug_rotary.py deleted file mode 100644 index 72e3084..0000000 --- a/testing/debug_scripts/parity_debug_rotary.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Verify the rotary inv_freq difference between FastPLMs and native ESMC.""" -from __future__ import annotations - -import torch - -from testing.official.esm_plusplus import load_official_model as load_native_esmc - - -def diff(label, a, b): - a = a.float() - b = b.float() - print(f" {label:55s} mse={((a-b)**2).mean().item():.3e} maxabs={(a-b).abs().max().item():.3e}") - - -def main() -> int: - device = torch.device("cuda") - - from transformers import AutoModelForMaskedLM - fast = AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", trust_remote_code=True, - dtype=torch.float32, device_map=device, - ).eval() - native, _ = load_native_esmc(reference_repo_id="esmc-300", device=device, dtype=torch.float32) - - f_rot = fast.transformer.blocks[0].attn.rotary - n_rot = native.model.transformer.blocks[0].attn.rotary - - print("Before any forward call:") - diff("fast.inv_freq vs native.inv_freq", f_rot.inv_freq, n_rot.inv_freq) - - print("\nManual recompute on GPU vs on CPU-then-moved:") - cpu_inv_freq = 1.0 / (10000.0 ** (torch.arange(0, 64, 2, device="cpu", dtype=torch.float32) / 64)) - gpu_inv_freq = 1.0 / (10000.0 ** (torch.arange(0, 64, 2, device=device, dtype=torch.float32) / 64)) - diff("cpu_computed.to(cuda) vs gpu_computed", cpu_inv_freq.to(device), gpu_inv_freq) - diff("fast.inv_freq vs cpu_computed.to(cuda)", f_rot.inv_freq, cpu_inv_freq.to(device)) - diff("fast.inv_freq vs gpu_computed", f_rot.inv_freq, gpu_inv_freq) - diff("native.inv_freq vs cpu_computed.to(cuda)", n_rot.inv_freq, cpu_inv_freq.to(device)) - diff("native.inv_freq vs gpu_computed", n_rot.inv_freq, gpu_inv_freq) - - print(f"\n fast inv_freq first 3 entries: {f_rot.inv_freq[:3].tolist()}") - print(f" native inv_freq first 3 entries: {n_rot.inv_freq[:3].tolist()}") - print(f" cpu_computed first 3 entries: {cpu_inv_freq[:3].tolist()}") - print(f" gpu_computed first 3 entries: {gpu_inv_freq[:3].tolist()}") - - f_inv_before = f_rot.inv_freq.clone() - _ = torch.randn(1, 5, 15, 64, device=device) - f_rot._update_cos_sin_cache(5, device=device, dtype=torch.float32) - print("\nAfter FastPLMs _update_cos_sin_cache:") - diff("fast.inv_freq before vs after cache update", f_inv_before, f_rot.inv_freq) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/debug_scripts/patch_hub_cache.py b/testing/debug_scripts/patch_hub_cache.py deleted file mode 100644 index 03286df..0000000 --- a/testing/debug_scripts/patch_hub_cache.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Overwrite Hub-cached composite modeling files with a freshly-built local composite. - -Purpose: validate in-flight changes to fastplms/ against the parity suite WITHOUT -having to push to HF Hub first. The parity tests load via trust_remote_code, which -caches the Hub `modeling_*.py` at -`~/.cache/huggingface/modules/transformers_modules///modeling_*.py`. -This script rebuilds those files from the current checkout so the next pytest run -exercises local code. - -Usage: - python testing/debug_scripts/patch_hub_cache.py # patch every Synthyra family - python testing/debug_scripts/patch_hub_cache.py dplm esm2 # patch only listed families -""" -from __future__ import annotations - -import os -import sys -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(REPO_ROOT)) - -from update_HF import MODEL_REGISTRY, build_composite # noqa: E402 - - -def repo_slug(repo_id: str) -> str: - """HF transformers_modules slug is the repo id with '-' replaced by '_hyphen_'.""" - return repo_id.replace("-", "_hyphen_") - - -def cache_dir_for(repo_id: str) -> Path: - base = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) - return base / "modules" / "transformers_modules" / repo_slug(repo_id) - - -def patch_entry(entry: dict) -> int: - if not entry.get("composite", False): - return 0 - composite_body = build_composite( - entry["modeling_src"], - include_embedding_mixin=entry.get("include_embedding_mixin", True), - ) - dest_filename = entry["modeling_dest"] - patched = 0 - for repo_id in entry["repo_ids"]: - cache_root = cache_dir_for(repo_id) - if not cache_root.exists(): - print(f" [skip] {repo_id}: no cache at {cache_root}") - continue - # Every revision subdir gets patched. - revs = [p for p in cache_root.iterdir() if p.is_dir()] - if not revs: - print(f" [skip] {repo_id}: no revisions under {cache_root}") - continue - for rev in revs: - dest = rev / dest_filename - if not dest.exists(): - print(f" [skip] {repo_id}@{rev.name}: no {dest_filename}") - continue - dest.write_text(composite_body, encoding="utf-8") - print(f" [patch] {repo_id}@{rev.name}: {dest}") - patched += 1 - return patched - - -def main() -> int: - families_filter = set(sys.argv[1:]) - total = 0 - for entry in MODEL_REGISTRY: - if families_filter and entry["family"] not in families_filter: - continue - print(f"== {entry['family']} ==") - total += patch_entry(entry) - print(f"\nPatched {total} composite file(s).") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/official/__init__.py b/testing/official/__init__.py deleted file mode 100644 index c956502..0000000 --- a/testing/official/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Helpers for loading official model implementations from submodules.""" -import os -import sys - -_ESM_SUBMODULE = os.path.join(os.path.dirname(__file__), "..", "..", "official", "esm") - - -def use_esm_submodule(): - """Load esm from official/esm submodule instead of pip. - - The Biohub esm package uses the same top-level `esm` import as fair-esm. - """ - if _ESM_SUBMODULE not in sys.path: - sys.path.insert(0, _ESM_SUBMODULE) diff --git a/testing/official/ankh.py b/testing/official/ankh.py deleted file mode 100644 index 81f0baf..0000000 --- a/testing/official/ankh.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Load official ANKH model (T5EncoderModel) from HuggingFace for comparison.""" - -import torch -import torch.nn as nn -from typing import Tuple - -from transformers import T5EncoderModel, AutoTokenizer - - -class _AnkhComplianceOutput: - """Mimics HuggingFace model output so the test suite can access .logits and .hidden_states.""" - - def __init__(self, logits: torch.Tensor, last_hidden_state: torch.Tensor, hidden_states: Tuple[torch.Tensor, ...]) -> None: - self.logits = logits - self.last_hidden_state = last_hidden_state - self.hidden_states = hidden_states - - -class _OfficialAnkhForwardWrapper(nn.Module): - def __init__(self, model: T5EncoderModel, tokenizer) -> None: - super().__init__() - self.model = model - self.tokenizer = tokenizer - self.lm_head = nn.Linear(model.config.d_model, model.config.vocab_size, bias=False) - self.lm_head.weight = model.shared.weight - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs): - outputs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - logits = self.lm_head(outputs.last_hidden_state) - return _AnkhComplianceOutput( - logits=logits, - last_hidden_state=outputs.last_hidden_state, - hidden_states=outputs.hidden_states, - ) - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, object]: - """Load the official ANKH model as a T5EncoderModel. - - The official ElnaggarLab repos are T5ForConditionalGeneration but - T5EncoderModel.from_pretrained extracts just the encoder. - - Returns (wrapped_model, tokenizer). - """ - model = T5EncoderModel.from_pretrained( - reference_repo_id, - device_map=device, - dtype=dtype, - ).eval() - tokenizer = AutoTokenizer.from_pretrained(reference_repo_id) - wrapped = _OfficialAnkhForwardWrapper(model, tokenizer).to(device=device, dtype=dtype).eval() - return wrapped, tokenizer - - -if __name__ == "__main__": - model, tokenizer = load_official_model( - reference_repo_id="ElnaggarLab/ankh-base", - device=torch.device("cuda"), - dtype=torch.float32, - ) - print(model) - print(tokenizer) diff --git a/testing/official/dplm.py b/testing/official/dplm.py deleted file mode 100644 index ad4970a..0000000 --- a/testing/official/dplm.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Load official DPLM model from HuggingFace transformers for comparison. - -DPLM uses the ESM2 architecture internally, so the official weights load -directly via EsmForMaskedLM from HuggingFace transformers. -""" - -import torch -import torch.nn as nn -from typing import Tuple - -from transformers import EsmForMaskedLM, EsmTokenizer - - -class _OfficialDPLMForwardWrapper(nn.Module): - def __init__(self, model: EsmForMaskedLM) -> None: - super().__init__() - self.model = model - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs): - return self.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, EsmTokenizer]: - """Load the official DPLM model (ESM2 architecture) from HuggingFace. - - Returns (wrapped_model, tokenizer). - """ - model = EsmForMaskedLM.from_pretrained( - reference_repo_id, - device_map=device, - dtype=dtype, - ).eval() - tokenizer = EsmTokenizer.from_pretrained(reference_repo_id) - wrapped = _OfficialDPLMForwardWrapper(model) - return wrapped, tokenizer diff --git a/testing/official/dplm2.py b/testing/official/dplm2.py deleted file mode 100644 index e260c0a..0000000 --- a/testing/official/dplm2.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Load official DPLM2 model from HuggingFace transformers for comparison. - -DPLM2 uses the ESM2 architecture internally, so the official weights load -directly via EsmForMaskedLM from HuggingFace transformers. - -The DPLM2 tokenizer emits special-token IDs at positions [vocab_size .. -vocab_size+3] (cls, eos, pad, mask) that are OUT OF RANGE of the size- -vocab_size embedding table. The real DPLM2 forward remaps those to in-range -AA-side IDs before embedding lookup. Our wrapper does the same remap so -feeding tokenizer output straight to native EsmForMaskedLM doesn't OOB. -""" - -import torch -import torch.nn as nn -from typing import Tuple - -from transformers import EsmForMaskedLM, EsmTokenizer - - -def _normalize_dplm2_input_ids(input_ids: torch.Tensor, vocab_size: int) -> torch.Tensor: - """Remap DPLM2's high-ID generic special tokens to the AA-side special IDs. - - Must match fastplms.dplm2.modeling_dplm2._normalize_dplm2_input_ids exactly; - the DPLM2 forward applies this normalization before the embedding lookup so - the shared ESM backbone sees in-range token IDs. - """ - if input_ids.numel() == 0: - return input_ids - normalized = input_ids.clone() - generic_to_aa_special_ids = { - vocab_size: 2, # eos-generic -> AA eos - vocab_size + 1: 3, # unk-generic -> AA unk - vocab_size + 2: 0, # pad-generic -> AA pad - vocab_size + 3: 32, # mask-generic -> AA mask - } - for generic_id, aa_id in generic_to_aa_special_ids.items(): - normalized[input_ids == generic_id] = aa_id - # The DPLM2 tokenizer emits cls at vocab_size + 2 for one line of the - # cls/eos/pad/mask family, depending on vocabulary file; the mapping above - # is what fastplms.dplm2.modeling_dplm2 normalizes. If any ID survives - # above the embedding table, remap the whole row of "generic specials" - # defensively so feeding the tokenizer output into native does not OOB. - in_range = normalized.lt(vocab_size) & normalized.ge(0) - if not bool(in_range.all()): - # Everything outside range collapses to 0 (pad) on the native side. - # The fp32 hidden-state parity with fast is still meaningful because - # fast applies the same mapping on its side. - normalized = torch.where(in_range, normalized, torch.zeros_like(normalized)) - return normalized - - -class _OfficialDPLM2ForwardWrapper(nn.Module): - def __init__(self, model: EsmForMaskedLM) -> None: - super().__init__() - self.model = model - self.vocab_size = int(model.config.vocab_size) - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs): - input_ids = _normalize_dplm2_input_ids(input_ids, self.vocab_size) - return self.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, EsmTokenizer]: - """Load the official DPLM2 model (ESM2 architecture) from HuggingFace. - - Returns (wrapped_model, tokenizer). - """ - model = EsmForMaskedLM.from_pretrained( - reference_repo_id, - device_map=device, - dtype=dtype, - ).eval() - tokenizer = EsmTokenizer.from_pretrained(reference_repo_id) - wrapped = _OfficialDPLM2ForwardWrapper(model) - return wrapped, tokenizer diff --git a/testing/official/esm2.py b/testing/official/esm2.py deleted file mode 100644 index a3f8dcd..0000000 --- a/testing/official/esm2.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Load official ESM2 model from HuggingFace transformers for comparison.""" - -import torch -import torch.nn as nn -from typing import Tuple - -from transformers import EsmForMaskedLM, EsmTokenizer - - -class _OfficialESM2ForwardWrapper(nn.Module): - def __init__(self, model: EsmForMaskedLM) -> None: - super().__init__() - self.model = model - self.tokenizer = EsmTokenizer.from_pretrained(model.config._name_or_path) - - def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs): - outputs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - return outputs - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, EsmTokenizer]: - """Load the official HuggingFace ESM2 model. - - Returns (wrapped_model, tokenizer). - The wrapped model's forward returns standard HF outputs with hidden_states. - """ - model = EsmForMaskedLM.from_pretrained( - reference_repo_id, - device_map=device, - dtype=dtype, - attn_implementation="sdpa", - position_embedding_type="rotary", - ).eval() - tokenizer = EsmTokenizer.from_pretrained(reference_repo_id) - wrapped = _OfficialESM2ForwardWrapper(model) - return wrapped, tokenizer - - -if __name__ == "__main__": - model, tokenizer = load_official_model(reference_repo_id="facebook/esm2_t6_8M_UR50D", device=torch.device("cuda"), dtype=torch.float32) - print(model) - print(tokenizer) diff --git a/testing/official/esm3.py b/testing/official/esm3.py deleted file mode 100644 index 5a21d22..0000000 --- a/testing/official/esm3.py +++ /dev/null @@ -1,370 +0,0 @@ -"""Load official ESM3 from the official/esm submodule for comparison.""" -import math -from typing import Optional, Tuple - -import einops -import torch -import torch.nn as nn -from torch.nn import functional as F - -from testing.official import use_esm_submodule - -use_esm_submodule() - - -def _patch_official_geom_attention_dtype_cast() -> None: - from esm.layers.geom_attention import GeometricReasoningOriginalImpl - - if getattr(GeometricReasoningOriginalImpl, "_fastplms_dtype_patch", False): - return - - def forward(self, s, affine, affine_mask, sequence_id, chain_id): - if sequence_id is None: - sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64) - attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2) - attn_bias = attn_bias.unsqueeze(1).float() - attn_bias = attn_bias.masked_fill( - ~affine_mask[:, None, None, :], - torch.finfo(attn_bias.dtype).min, - ) - chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2) - attn_bias = attn_bias.masked_fill( - chain_id_mask.unsqueeze(1), - torch.finfo(s.dtype).min, - ) - - ns = self.s_norm(s) - vec_rot, vec_dist = self.proj(ns).split( - [ - self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages, - self.v_heads * 2 * 3, - ], - dim=-1, - ) - query_rot, key_rot, value = ( - affine.rot[..., None] - .apply(einops.rearrange(vec_rot, "... (h c) -> ... h c", c=3)) - .split( - [self.v_heads, self.v_heads, self.v_heads * self.num_vector_messages], - dim=-2, - ) - ) - query_dist, key_dist = ( - affine[..., None] - .apply(einops.rearrange(vec_dist, "... (h c) -> ... h c", c=3)) - .chunk(2, dim=-2) - ) - - query_dist = einops.rearrange(query_dist, "b s h d -> b h s 1 d") - key_dist = einops.rearrange(key_dist, "b s h d -> b h 1 s d") - query_rot = einops.rearrange(query_rot, "b s h d -> b h s d") - key_rot = einops.rearrange(key_rot, "b s h d -> b h d s") - value = einops.rearrange( - value, - "b s (h m) d -> b h s (m d)", - m=self.num_vector_messages, - ) - - distance_term = (query_dist - key_dist).norm(dim=-1) / math.sqrt(3) - rotation_term = query_rot.matmul(key_rot) / math.sqrt(3) - distance_term_weight = einops.rearrange( - F.softplus(self.distance_scale_per_head), - "h -> h 1 1", - ) - rotation_term_weight = einops.rearrange( - F.softplus(self.rotation_scale_per_head), - "h -> h 1 1", - ) - attn_weight = ( - rotation_term * rotation_term_weight - - distance_term * distance_term_weight - ) - - if attn_bias is not None: - seq_q = attn_weight.size(2) - seq_k = attn_weight.size(3) - bias_q = max(0, attn_bias.size(2) - seq_q) - bias_k = max(0, attn_bias.size(3) - seq_k) - attn_bias = attn_bias[:, :, bias_q:, bias_k:] - attn_weight = attn_weight + attn_bias - - attn_weight = torch.softmax(attn_weight, dim=-1) - attn_out = attn_weight.matmul(value) - attn_out = ( - affine.rot[..., None] - .invert() - .apply( - einops.rearrange( - attn_out, - "b h s (m d) -> b s (h m) d", - m=self.num_vector_messages, - ) - ) - ) - attn_out = einops.rearrange( - attn_out, - "b s (h m) d -> b s (h m d)", - m=self.num_vector_messages, - ) - if self.mask_and_zero_frameless: - attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0) - attn_out = attn_out.to(self.out_proj.weight.dtype) - return self.out_proj(attn_out) - - GeometricReasoningOriginalImpl.forward = forward - GeometricReasoningOriginalImpl._fastplms_dtype_patch = True - - -_patch_official_geom_attention_dtype_cast() - - -class _ESM3ComplianceOutput: - def __init__( - self, - logits: torch.Tensor, - last_hidden_state: torch.Tensor, - hidden_states: Tuple[torch.Tensor, ...], - sequence_logits: torch.Tensor, - structure_logits: torch.Tensor, - function_logits: torch.Tensor, - residue_logits: torch.Tensor, - ) -> None: - self.logits = logits - self.last_hidden_state = last_hidden_state - self.hidden_states = hidden_states - self.sequence_logits = sequence_logits - self.structure_logits = structure_logits - self.function_logits = function_logits - self.residue_logits = residue_logits - - -class _ESM3StateDictRoot(nn.Module): - def __init__(self, model: nn.Module) -> None: - super().__init__() - self.esm3 = model - - -class _OfficialESM3ForwardWrapper(nn.Module): - def __init__(self, model: nn.Module) -> None: - super().__init__() - self.model = _ESM3StateDictRoot(model) - self.tokenizer = model.tokenizers.sequence - - @property - def esm3(self) -> nn.Module: - return self.model.esm3 - - def _encode_inputs( - self, - sequence_tokens: torch.Tensor, - structure_tokens: torch.Tensor, - average_plddt: torch.Tensor, - per_res_plddt: torch.Tensor, - ss8_tokens: torch.Tensor, - sasa_tokens: torch.Tensor, - function_tokens: torch.Tensor, - residue_annotation_tokens: torch.Tensor, - ) -> torch.Tensor: - from esm.utils.misc import rbf - - encoder = self.esm3.encoder - sequence_embed = encoder.sequence_embed(sequence_tokens) - rbf_16_fn = lambda x: rbf(x, v_min=0.0, v_max=1.0, n_bins=16) - plddt_embed = encoder.plddt_projection( - rbf_16_fn(average_plddt).to(encoder.plddt_projection.weight.dtype) - ) - structure_per_res_plddt = encoder.structure_per_res_plddt_projection( - rbf_16_fn(per_res_plddt).to( - encoder.structure_per_res_plddt_projection.weight.dtype - ) - ) - structure_embed = encoder.structure_tokens_embed(structure_tokens) - ss8_embed = encoder.ss8_embed(ss8_tokens) - sasa_embed = encoder.sasa_embed(sasa_tokens) - function_embed = torch.cat( - [ - embed_fn(funcs) - for embed_fn, funcs in zip( - encoder.function_embed, - function_tokens.unbind(-1), - ) - ], - -1, - ) - - batch_size, seq_len, num_annotations = residue_annotation_tokens.shape - residue_embed = encoder.residue_embed( - einops.rearrange( - residue_annotation_tokens, - "b l n -> (b l) n", - b=batch_size, - l=seq_len, - n=num_annotations, - ) - ) - residue_embed = einops.rearrange( - residue_embed, - "(b l) d -> b l d", - b=batch_size, - l=seq_len, - ) - - return ( - sequence_embed - + plddt_embed - + structure_per_res_plddt - + structure_embed - + ss8_embed - + sasa_embed - + function_embed - + residue_embed - ) - - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - sequence_tokens: Optional[torch.Tensor] = None, - structure_tokens: Optional[torch.Tensor] = None, - ss8_tokens: Optional[torch.Tensor] = None, - sasa_tokens: Optional[torch.Tensor] = None, - function_tokens: Optional[torch.Tensor] = None, - residue_annotation_tokens: Optional[torch.Tensor] = None, - average_plddt: Optional[torch.Tensor] = None, - per_res_plddt: Optional[torch.Tensor] = None, - structure_coords: Optional[torch.Tensor] = None, - chain_id: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - **kwargs, - ) -> _ESM3ComplianceOutput: - from esm.utils.constants import esm3 as C - from esm.utils.structure.affine3d import build_affine3d_from_coordinates - - del output_hidden_states, kwargs - output_attentions = bool(output_attentions) - if sequence_tokens is None: - sequence_tokens = input_ids - if sequence_id is None and attention_mask is not None: - sequence_id = attention_mask.to(dtype=torch.bool) - - present_inputs = [ - sequence_tokens, - structure_tokens, - ss8_tokens, - sasa_tokens, - structure_coords, - function_tokens, - residue_annotation_tokens, - ] - try: - seq_len, device = next( - (x.shape[1], x.device) for x in present_inputs if x is not None - ) - except StopIteration: - raise ValueError("At least one of the inputs must be non-None") - - def defaults(x: Optional[torch.Tensor], token: int) -> torch.Tensor: - if x is None: - return torch.full( - (1, seq_len), - token, - dtype=torch.long, - device=device, - ) - return x - - sequence_tokens = defaults(sequence_tokens, self.tokenizer.mask_token_id) - ss8_tokens = defaults(ss8_tokens, C.SS8_PAD_TOKEN) - sasa_tokens = defaults(sasa_tokens, C.SASA_PAD_TOKEN) - average_plddt = defaults(average_plddt, 1).float() - per_res_plddt = defaults(per_res_plddt, 0).float() - chain_id = defaults(chain_id, 0) - - if residue_annotation_tokens is None: - residue_annotation_tokens = torch.full( - (1, seq_len, C.MAX_RESIDUE_ANNOTATIONS), - C.RESIDUE_PAD_TOKEN, - dtype=torch.long, - device=device, - ) - if function_tokens is None: - function_tokens = torch.full( - (1, seq_len, C.FUNCTION_TOKENS_DEPTH), - C.INTERPRO_PAD_TOKEN, - dtype=torch.long, - device=device, - ) - if structure_coords is None: - structure_coords = torch.full( - (1, seq_len, 3, 3), - float("nan"), - dtype=torch.float, - device=device, - ) - - structure_coords = structure_coords[..., :3, :] - affine, affine_mask = build_affine3d_from_coordinates(structure_coords) - structure_tokens = defaults(structure_tokens, C.STRUCTURE_MASK_TOKEN) - structure_tokens = ( - structure_tokens.masked_fill(structure_tokens == -1, C.STRUCTURE_MASK_TOKEN) - .masked_fill(sequence_tokens == C.SEQUENCE_BOS_TOKEN, C.STRUCTURE_BOS_TOKEN) - .masked_fill(sequence_tokens == C.SEQUENCE_PAD_TOKEN, C.STRUCTURE_PAD_TOKEN) - .masked_fill(sequence_tokens == C.SEQUENCE_EOS_TOKEN, C.STRUCTURE_EOS_TOKEN) - .masked_fill( - sequence_tokens == C.SEQUENCE_CHAINBREAK_TOKEN, - C.STRUCTURE_CHAINBREAK_TOKEN, - ) - ) - - x = self._encode_inputs( - sequence_tokens, - structure_tokens, - average_plddt, - per_res_plddt, - ss8_tokens, - sasa_tokens, - function_tokens, - residue_annotation_tokens, - ) - x, embedding, hidden_states, attentions = self.esm3.transformer( - x, - sequence_id, - affine, - affine_mask, - chain_id, - output_attentions=output_attentions, - ) - output = self.esm3.output_heads(x, embedding, attentions=attentions) - - return _ESM3ComplianceOutput( - logits=output.sequence_logits, - last_hidden_state=output.embeddings, - hidden_states=hidden_states, - sequence_logits=output.sequence_logits, - structure_logits=output.structure_logits, - function_logits=output.function_logits, - residue_logits=output.residue_logits, - ) - - -def _normalize_reference_repo_id(reference_repo_id: str) -> str: - if reference_repo_id == "biohub/esm3-sm-open-v1": - return "esm3-sm-open-v1" - return reference_repo_id - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, object]: - from esm.pretrained import load_local_model - from esm.utils.constants.models import normalize_model_name - - model_name = normalize_model_name(_normalize_reference_repo_id(reference_repo_id)) - model = load_local_model(model_name, device=device).eval() - model = model.to(device=device, dtype=dtype).eval() - wrapped = _OfficialESM3ForwardWrapper(model).to(device=device, dtype=dtype).eval() - return wrapped, wrapped.tokenizer diff --git a/testing/official/esm_plusplus.py b/testing/official/esm_plusplus.py deleted file mode 100644 index e6f7cba..0000000 --- a/testing/official/esm_plusplus.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Load official ESMC model from the official/esm submodule for comparison.""" -from typing import Optional, Tuple - -import torch -import torch.nn as nn - -from testing.official import use_esm_submodule - -use_esm_submodule() - - -class _ESMCComplianceOutput: - """Mimics HuggingFace model output so the test suite can access .logits and .hidden_states.""" - def __init__(self, logits: torch.Tensor, last_hidden_state: torch.Tensor, hidden_states: Tuple[torch.Tensor, ...]) -> None: - self.logits = logits - self.last_hidden_state = last_hidden_state - self.hidden_states = hidden_states - - -class _OfficialESMCForwardWrapper(nn.Module): - """Wraps official ESMC model to produce outputs compatible with our test suite.""" - def __init__(self, model: nn.Module, tokenizer: object) -> None: - super().__init__() - self.model = model - self.tokenizer = tokenizer - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - sequence_id: Optional[torch.Tensor] = None, - **kwargs, - ): - esmc_output = self.model(sequence_tokens=input_ids) - # ESMC returns: sequence_logits, embeddings, hidden_states (stacked [n_layers, B, L, D]) - logits = esmc_output.sequence_logits - embeddings = esmc_output.embeddings - raw_hiddens = esmc_output.hidden_states - # Convert stacked tensor to tuple for compatibility with hidden_states[-1] - if raw_hiddens is not None: - hidden_states = tuple(raw_hiddens[i] for i in range(raw_hiddens.shape[0])) - hidden_states = hidden_states + (embeddings,) - else: - hidden_states = (embeddings,) - return _ESMCComplianceOutput( - logits=logits, - last_hidden_state=embeddings, - hidden_states=hidden_states, - ) - - -def load_official_model( - reference_repo_id: str, - device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, object]: - """Load the official ESMC model from the esm submodule. - - Args: - reference_repo_id: e.g. "biohub/ESMC-300M" or "esmc-300" - device: target device - dtype: target dtype (should be float32 for comparison) - - Returns (wrapped_model, tokenizer). - """ - from esm.models.esmc import ESMC - from esm.tokenization import get_esmc_model_tokenizers - from fastplms.esm_plusplus.modeling_esm_plusplus import ( - _ESMC_CHECKPOINT_SPECS, - _load_safetensors_state_dict, - _resolve_esmc_checkpoint_key, - get_esmc_checkpoint_path, - ) - - key = _resolve_esmc_checkpoint_key(reference_repo_id) - spec = _ESMC_CHECKPOINT_SPECS[key] - with torch.device(device): - official_model = ESMC( - d_model=spec["hidden_size"], - n_heads=spec["num_attention_heads"], - n_layers=spec["num_hidden_layers"], - tokenizer=get_esmc_model_tokenizers(), - use_flash_attn=False, - ).eval() - _load_safetensors_state_dict( - model_obj=official_model, - checkpoint_path=get_esmc_checkpoint_path(reference_repo_id), - device=device, - ) - - official_model = official_model.to(device=device, dtype=dtype).eval() - tokenizer = official_model.tokenizer - wrapped = _OfficialESMCForwardWrapper(official_model, tokenizer).to(device=device, dtype=dtype).eval() - return wrapped, tokenizer - - -if __name__ == "__main__": - model, tokenizer = load_official_model(reference_repo_id="biohub/ESMC-300M", device=torch.device("cuda"), dtype=torch.float32) - print(model) - print(tokenizer) diff --git a/testing/run_boltz2_compliance.py b/testing/run_boltz2_compliance.py deleted file mode 100644 index e82cd3f..0000000 --- a/testing/run_boltz2_compliance.py +++ /dev/null @@ -1,928 +0,0 @@ -import argparse -import json -import os -import random -import re -import subprocess -import sys -import tempfile -import time -import urllib.request - -import biotite.structure as struc -import matplotlib -import numpy as np -import torch -from tqdm.auto import tqdm - -from pathlib import Path -from typing import Any -from typing import Dict -from typing import List -from typing import Optional -from typing import Tuple -from transformers import AutoModel - -from fastplms.boltz.cif_writer import write_cif -from fastplms.boltz.get_boltz2_weights import BOLTZ2_CKPT_URL -from fastplms.boltz.minimal_featurizer import build_boltz2_features -from fastplms.boltz.minimal_structures import ProteinStructureTemplate -from testing.common import autocast_context -from testing.common import build_output_dir -from testing.common import login_if_needed -from testing.common import resolve_device -from testing.common import resolve_dtype -from testing.reporting import write_csv -from testing.reporting import write_json -from testing.reporting import write_summary - - -matplotlib.use("Agg") -import matplotlib.pyplot as plt - - -assert "tm_score" in dir(struc), ( - "biotite.structure.tm_score is unavailable. Install biotite>=1.5.0 in the target environment." -) - -TM_SCORE_FN = struc.tm_score -BOLTZ2_FIXED_RECYCLING_STEPS = 3 -BOLTZ2_FIXED_SAMPLING_STEPS = 200 -BOLTZ2_FIXED_DIFFUSION_SAMPLES = 20 -MIN_SEED_VALUE = int(np.iinfo(np.uint32).min) -MAX_SEED_VALUE = int(np.iinfo(np.uint32).max) - -SEQUENCE_OPTIONS = [ - "MDDADPEERNYDNMLKMLSDLNKDLEKLLEEMEKISVQATWMAYDMVVMRTNPTLAESMRRLEDAFVNCKEEMEKNWQELLHETKQRL", - "MASLGHILVFCVGLLTMAKAESPKEHDPFTYDYQSLQIGGLVIAGILFILGILIVLSRRCRCKFNQQQRTGEPDEEEGTFRSSIRRLSTRRR", - "MAVESRVTQEEIKKEPEKPIDREKTCPLLLRVFTTNNGRHHRMDEFSRGNVPSSELQIYTWMDATLKELTSLVKEVYPEARKKGTHFNFAIVFTDVKRPGYRVKEIGSTMSGRKGTDDSMTLQSQKFQIGDYLDIAITPPNRAPPPSGRMRPY", -] - - -def _enforce_determinism() -> None: - if "CUBLAS_WORKSPACE_CONFIG" not in os.environ: - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.allow_tf32 = False - if torch.cuda.is_available(): - torch.backends.cuda.matmul.allow_tf32 = False - torch.use_deterministic_algorithms(True) - - -def _seed_everything(seed: Optional[int] = None, workers: bool = False) -> int: - if seed is None: - env_seed = os.environ.get("PL_GLOBAL_SEED") - if env_seed is None: - seed = 0 - else: - seed = int(env_seed) - elif isinstance(seed, int) is False: - seed = int(seed) - - if not (MIN_SEED_VALUE <= seed <= MAX_SEED_VALUE): - raise ValueError(f"{seed} is not in bounds, numpy accepts from {MIN_SEED_VALUE} to {MAX_SEED_VALUE}") - - os.environ["PL_GLOBAL_SEED"] = str(seed) - os.environ["PL_SEED_WORKERS"] = f"{int(workers)}" - os.environ["PYTHONHASHSEED"] = str(seed) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - return seed - - -def _download_checkpoint_if_needed(checkpoint_path: Path) -> Path: - checkpoint_path.parent.mkdir(parents=True, exist_ok=True) - if not checkpoint_path.exists(): - urllib.request.urlretrieve(BOLTZ2_CKPT_URL, str(checkpoint_path)) # noqa: S310 - return checkpoint_path - - -def _detect_no_kernels_support() -> bool: - command = [sys.executable, "-m", "boltz.main", "predict", "--help"] - completed = subprocess.run(command, capture_output=True, text=True, check=False) - combined_output = f"{completed.stdout}\n{completed.stderr}" - return "--no_kernels" in combined_output - - -def _set_sequence_seed(seed: int, sequence_index: int) -> None: - _seed_everything(seed=seed + sequence_index, workers=False) - - -def _to_device(feats: Dict[str, torch.Tensor], device: torch.device, dtype: torch.dtype) -> Dict[str, torch.Tensor]: - output: Dict[str, torch.Tensor] = {} - for key in feats: - value = feats[key] - if value.is_floating_point(): - output[key] = value.to(device=device, dtype=dtype) - else: - output[key] = value.to(device=device) - return output - - -def _clone_feats(feats: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - output: Dict[str, torch.Tensor] = {} - for key in feats: - output[key] = feats[key].clone() - return output - - -def _summary_metric(value: torch.Tensor) -> torch.Tensor: - if value.ndim == 0: - return value.reshape(1) - if value.ndim == 1: - return value - return value.reshape(value.shape[0], -1)[:, 0] - - -def _extract_primary_plddt_vector(output: Dict[str, torch.Tensor], feats: Dict[str, torch.Tensor]) -> torch.Tensor: - assert "plddt" in output, "Missing pLDDT in model output." - plddt = output["plddt"].detach().cpu() - if plddt.ndim == 0: - return plddt.reshape(1).float() - if plddt.ndim >= 2: - plddt = plddt[0] - plddt = plddt.reshape(-1).float() - - token_mask = feats["token_pad_mask"][0].detach().cpu().reshape(-1) > 0 - atom_mask = feats["atom_pad_mask"][0].detach().cpu().reshape(-1) > 0 - if plddt.numel() == token_mask.numel(): - plddt = plddt[token_mask] - elif plddt.numel() == atom_mask.numel(): - plddt = plddt[atom_mask] - return plddt - - -def _compute_confidence_score(ptm: torch.Tensor, iptm: torch.Tensor, complex_plddt: torch.Tensor) -> torch.Tensor: - if torch.allclose(iptm, torch.zeros_like(iptm)): - return (4 * complex_plddt + ptm) / 5 - return (4 * complex_plddt + iptm) / 5 - - -def _run_ours_forward( - model, - feats_ours: Dict[str, torch.Tensor], - args: argparse.Namespace, - device: torch.device, - dtype: torch.dtype, - sequence_index: int, -) -> Dict[str, torch.Tensor]: - with torch.no_grad(), autocast_context(device=device, dtype=dtype): - _set_sequence_seed(args.seed, sequence_index) - return model.forward( - feats=feats_ours, - recycling_steps=BOLTZ2_FIXED_RECYCLING_STEPS, - num_sampling_steps=BOLTZ2_FIXED_SAMPLING_STEPS, - diffusion_samples=BOLTZ2_FIXED_DIFFUSION_SAMPLES, - run_confidence_sequentially=args.run_confidence_sequentially, - ) - - -def _vector_metrics(lhs: torch.Tensor, rhs: torch.Tensor) -> Tuple[float, float, float]: - delta = lhs.float() - rhs.float() - abs_delta = torch.abs(delta) - mae = float(abs_delta.mean().item()) - rmse = float(torch.sqrt(torch.mean(delta * delta)).item()) - max_abs = float(abs_delta.max().item()) - return mae, rmse, max_abs - - -def _kabsch_align_mobile_to_target(mobile: torch.Tensor, target: torch.Tensor) -> torch.Tensor: - assert mobile.ndim == 2 and target.ndim == 2, "Expected coordinate tensors with shape [N, 3]." - assert mobile.shape == target.shape, "Coordinate tensors must have matching shapes." - assert mobile.shape[1] == 3, "Coordinate tensors must have last dimension size 3." - assert mobile.shape[0] > 0, "Expected at least one shared atom for alignment." - - mobile_32 = mobile.float() - target_32 = target.float() - if mobile_32.shape[0] < 3: - mobile_centroid = mobile_32.mean(dim=0, keepdim=True) - target_centroid = target_32.mean(dim=0, keepdim=True) - return mobile_32 - mobile_centroid + target_centroid - - mobile_centroid = mobile_32.mean(dim=0, keepdim=True) - target_centroid = target_32.mean(dim=0, keepdim=True) - mobile_centered = mobile_32 - mobile_centroid - target_centered = target_32 - target_centroid - - covariance = mobile_centered.transpose(0, 1).matmul(target_centered) - u_mat, _, vh_mat = torch.linalg.svd(covariance, full_matrices=False) - correction = torch.eye(3, dtype=mobile_32.dtype, device=mobile_32.device) - det_sign = torch.det(vh_mat.transpose(0, 1).matmul(u_mat.transpose(0, 1))).item() - if det_sign < 0: - correction[2, 2] = -1.0 - rotation = vh_mat.transpose(0, 1).matmul(correction).matmul(u_mat.transpose(0, 1)) - return mobile_centered.matmul(rotation) + target_centroid - - -def _pairwise_distance_mae(lhs: torch.Tensor, rhs: torch.Tensor) -> float: - assert lhs.ndim == 2 and rhs.ndim == 2, "Expected coordinate tensors with shape [N, 3]." - assert lhs.shape == rhs.shape, "Coordinate tensors must have matching shapes." - assert lhs.shape[1] == 3, "Coordinate tensors must have last dimension size 3." - lhs_dist = torch.cdist(lhs.float(), lhs.float()) - rhs_dist = torch.cdist(rhs.float(), rhs.float()) - return float(torch.mean(torch.abs(lhs_dist - rhs_dist)).item()) - - -def _write_single_chain_fasta(sequence: str, path: Path) -> None: - text = f">A|protein|empty\n{sequence}\n" - path.write_text(text, encoding="utf-8") - - -def _parse_pdb_atom_map(path: Path) -> Dict[Tuple[str, int, str], torch.Tensor]: - atom_map: Dict[Tuple[str, int, str], torch.Tensor] = {} - for line in path.read_text(encoding="utf-8").splitlines(): - if not (line.startswith("ATOM") or line.startswith("HETATM")): - continue - atom_name = line[12:16].strip() - chain_id = line[21:22].strip() - residue_index = int(line[22:26]) - x_val = float(line[30:38]) - y_val = float(line[38:46]) - z_val = float(line[46:54]) - atom_map[(chain_id, residue_index, atom_name)] = torch.tensor([x_val, y_val, z_val], dtype=torch.float32) - assert len(atom_map) > 0, f"No atoms parsed from PDB: {path}" - return atom_map - - -def _extract_model_id_from_name(filename: str) -> int: - match = re.search(r"_model_(\d+)\.", filename) - assert match is not None, f"Could not parse model id from filename: {filename}" - return int(match.group(1)) - - -def _map_paths_by_model(paths: List[Path]) -> Dict[int, Path]: - path_map: Dict[int, Path] = {} - for path in paths: - model_id = _extract_model_id_from_name(path.name) - assert model_id not in path_map, f"Found duplicate artifacts for model id {model_id}: {path}" - path_map[model_id] = path - return path_map - - -def _build_ours_atom_maps( - sample_coords: torch.Tensor, - atom_mask: torch.Tensor, - atom_names: List[str], - atom_residue_index: List[int], - atom_chain_id: List[str], -) -> List[Dict[Tuple[str, int, str], torch.Tensor]]: - coords = sample_coords.detach().cpu() - if coords.ndim == 4: - assert coords.shape[0] == 1, "Expected singleton batch dimension for sample coordinates." - coords = coords[0] - if coords.ndim == 2: - coords = coords.unsqueeze(0) - assert coords.ndim == 3, f"Expected sample_atom_coords with 3 dimensions, got shape {coords.shape}." - assert coords.shape[0] >= BOLTZ2_FIXED_DIFFUSION_SAMPLES, ( - f"Expected at least {BOLTZ2_FIXED_DIFFUSION_SAMPLES} samples, got {coords.shape[0]}." - ) - - atom_mask_bool = atom_mask.detach().cpu() > 0 - output: List[Dict[Tuple[str, int, str], torch.Tensor]] = [] - for sample_index in range(BOLTZ2_FIXED_DIFFUSION_SAMPLES): - valid_coords = coords[sample_index][atom_mask_bool] - assert valid_coords.shape[0] >= len(atom_names), ( - "Our model returned fewer valid atom coordinates than template atoms." - ) - atom_map: Dict[Tuple[str, int, str], torch.Tensor] = {} - for atom_idx in range(len(atom_names)): - key = ( - atom_chain_id[atom_idx], - atom_residue_index[atom_idx] + 1, - atom_names[atom_idx], - ) - atom_map[key] = valid_coords[atom_idx].float().cpu() - output.append(atom_map) - return output - - -def _build_reference_cif_tensors( - template: ProteinStructureTemplate, - atom_pad_mask: torch.Tensor, - ref_atom_map: Dict[Tuple[str, int, str], torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - assert atom_pad_mask.ndim == 1, "Expected atom pad mask with shape [atoms]." - atom_slots = atom_pad_mask.shape[0] - coords = torch.zeros((1, atom_slots, 3), dtype=torch.float32) - ref_mask = atom_pad_mask.detach().cpu().float().clone() - for atom_idx in range(template.num_atoms): - key = ( - template.atom_chain_id[atom_idx], - template.atom_residue_index[atom_idx] + 1, - template.atom_names[atom_idx], - ) - if key in ref_atom_map: - coords[0, atom_idx] = ref_atom_map[key].float().cpu() - else: - ref_mask[atom_idx] = 0.0 - return coords, ref_mask - - -def _run_boltz_cli_reference( - sequence: str, - sequence_index: int, - checkpoint_path: Path, - args: argparse.Namespace, - device: torch.device, - supports_no_kernels: bool, -) -> Tuple[List[Dict[Tuple[str, int, str], torch.Tensor]], List[torch.Tensor], List[Dict[str, float]]]: - sequence_seed = args.seed + sequence_index - with tempfile.TemporaryDirectory(prefix=f"boltz2_ref_{sequence_index}_") as tmp_dir_str: - tmp_dir = Path(tmp_dir_str) - fasta_path = tmp_dir / f"seq_{sequence_index}.fasta" - out_root = tmp_dir / "ref_out" - _write_single_chain_fasta(sequence=sequence, path=fasta_path) - - command = [ - sys.executable, - "-m", - "boltz.main", - "predict", - str(fasta_path), - "--out_dir", - str(out_root), - "--model", - "boltz2", - "--checkpoint", - str(checkpoint_path), - "--recycling_steps", - str(BOLTZ2_FIXED_RECYCLING_STEPS), - "--sampling_steps", - str(BOLTZ2_FIXED_SAMPLING_STEPS), - "--diffusion_samples", - str(BOLTZ2_FIXED_DIFFUSION_SAMPLES), - "--seed", - str(sequence_seed), - "--output_format", - "pdb", - ] - if supports_no_kernels: - command.append("--no_kernels") - - env = os.environ.copy() - env["PL_GLOBAL_SEED"] = str(sequence_seed) - env["PL_SEED_WORKERS"] = "0" - env["PYTHONHASHSEED"] = str(sequence_seed) - completed = subprocess.run(command, capture_output=True, text=True, check=False, env=env) - if completed.returncode != 0: - stderr = completed.stderr[-4000:] - stdout = completed.stdout[-4000:] - raise RuntimeError( - "pip boltz CLI prediction failed.\n" - f"Command: {' '.join(command)}\n" - f"STDOUT tail:\n{stdout}\n" - f"STDERR tail:\n{stderr}" - ) - - results_root = out_root / f"boltz_results_{fasta_path.stem}" / "predictions" - assert results_root.exists(), f"Reference predictions directory not found: {results_root}" - - pdb_candidates = sorted(results_root.rglob("*_model_*.pdb")) - plddt_candidates = sorted(results_root.rglob("plddt_*_model_*.npz")) - confidence_candidates = sorted(results_root.rglob("confidence_*_model_*.json")) - - assert len(pdb_candidates) > 0, f"No reference PDB artifacts found under {results_root}" - assert len(plddt_candidates) > 0, f"No reference pLDDT npz artifacts found under {results_root}" - assert len(confidence_candidates) > 0, f"No reference confidence json artifacts found under {results_root}" - - pdb_by_model = _map_paths_by_model(pdb_candidates) - plddt_by_model = _map_paths_by_model(plddt_candidates) - confidence_by_model = _map_paths_by_model(confidence_candidates) - - expected_model_ids = list(range(BOLTZ2_FIXED_DIFFUSION_SAMPLES)) - for model_id in expected_model_ids: - assert model_id in pdb_by_model, f"Missing reference PDB for model {model_id}" - assert model_id in plddt_by_model, f"Missing reference pLDDT for model {model_id}" - assert model_id in confidence_by_model, f"Missing reference confidence JSON for model {model_id}" - - atom_maps: List[Dict[Tuple[str, int, str], torch.Tensor]] = [] - plddt_samples: List[torch.Tensor] = [] - confidence_summaries: List[Dict[str, float]] = [] - for model_id in expected_model_ids: - atom_maps.append(_parse_pdb_atom_map(pdb_by_model[model_id])) - - with np.load(plddt_by_model[model_id]) as handle: - assert "plddt" in handle.files, f"Missing 'plddt' array in {plddt_by_model[model_id]}" - plddt_samples.append(torch.tensor(handle["plddt"], dtype=torch.float32)) - - confidence_summary = json.loads(confidence_by_model[model_id].read_text(encoding="utf-8")) - for key in ["ptm", "iptm", "complex_plddt", "confidence_score"]: - assert key in confidence_summary, ( - f"Reference confidence summary missing key '{key}' in {confidence_by_model[model_id]}" - ) - confidence_summaries.append(confidence_summary) - - return atom_maps, plddt_samples, confidence_summaries - - -def _shared_ca_key_order( - ours_atom_maps: List[Dict[Tuple[str, int, str], torch.Tensor]], - ref_atom_maps: List[Dict[Tuple[str, int, str], torch.Tensor]], -) -> List[Tuple[str, int, str]]: - shared_keys = {key for key in ours_atom_maps[0] if key[2] == "CA"} - for atom_map in ours_atom_maps: - ca_keys = {key for key in atom_map if key[2] == "CA"} - shared_keys = shared_keys.intersection(ca_keys) - for atom_map in ref_atom_maps: - ca_keys = {key for key in atom_map if key[2] == "CA"} - shared_keys = shared_keys.intersection(ca_keys) - assert len(shared_keys) > 0, "No shared CA atoms found across all samples." - ordered_keys = list(shared_keys) - ordered_keys.sort() - return ordered_keys - - -def _stack_coords_for_keys( - atom_maps: List[Dict[Tuple[str, int, str], torch.Tensor]], - ordered_keys: List[Tuple[str, int, str]], -) -> torch.Tensor: - stacked_samples: List[torch.Tensor] = [] - for atom_map in atom_maps: - coords: List[torch.Tensor] = [] - for key in ordered_keys: - assert key in atom_map, f"Missing key {key} in atom map." - coords.append(atom_map[key].float()) - stacked_samples.append(torch.stack(coords, dim=0)) - return torch.stack(stacked_samples, dim=0) - - -def _coords_to_ca_atom_array(coords: np.ndarray) -> struc.AtomArray: - assert coords.ndim == 2, "Expected coordinate array with shape [N, 3]." - assert coords.shape[1] == 3, "Expected coordinate array with shape [N, 3]." - assert coords.shape[0] > 0, "Expected at least one CA atom for TM-score." - assert np.all(np.isfinite(coords)), "Coordinate array for TM-score contains non-finite values." - - atom_count = coords.shape[0] - array = struc.AtomArray(atom_count) - array.coord = coords.astype(np.float32, copy=False) - array.atom_name = np.full(atom_count, "CA") - array.res_name = np.full(atom_count, "GLY") - array.chain_id = np.full(atom_count, "A") - array.res_id = np.arange(1, atom_count + 1, dtype=np.int32) - array.element = np.full(atom_count, "C") - return array - - -def _tm_score_from_coords(reference_coords: torch.Tensor, subject_coords: torch.Tensor) -> float: - aligned_subject = _kabsch_align_mobile_to_target(subject_coords, reference_coords) - reference_np = reference_coords.detach().cpu().numpy().astype(np.float64) - aligned_np = aligned_subject.detach().cpu().numpy().astype(np.float64) - reference_atom_array = _coords_to_ca_atom_array(reference_np) - subject_atom_array = _coords_to_ca_atom_array(aligned_np) - index_array = np.arange(reference_np.shape[0], dtype=np.int32) - tm_value = float( - TM_SCORE_FN( - reference=reference_atom_array, - subject=subject_atom_array, - reference_indices=index_array, - subject_indices=index_array, - reference_length="shorter", - ) - ) - assert np.isfinite(tm_value), "TM-score computation produced non-finite value." - return tm_value - - -def _build_tm_matrix(reference_stack: torch.Tensor, subject_stack: torch.Tensor, symmetric: bool) -> np.ndarray: - assert reference_stack.ndim == 3 and subject_stack.ndim == 3, "Expected stacks with shape [S, N, 3]." - assert reference_stack.shape[1:] == subject_stack.shape[1:], "Reference and subject stacks must share atom layout." - matrix = np.zeros((reference_stack.shape[0], subject_stack.shape[0]), dtype=np.float32) - if symmetric: - assert reference_stack.shape[0] == subject_stack.shape[0], "Symmetric matrix requires same sample count." - for row_idx in range(reference_stack.shape[0]): - for col_idx in range(row_idx, subject_stack.shape[0]): - tm_value = _tm_score_from_coords(reference_stack[row_idx], subject_stack[col_idx]) - matrix[row_idx, col_idx] = tm_value - matrix[col_idx, row_idx] = tm_value - return matrix - - for row_idx in range(reference_stack.shape[0]): - for col_idx in range(subject_stack.shape[0]): - matrix[row_idx, col_idx] = _tm_score_from_coords(reference_stack[row_idx], subject_stack[col_idx]) - return matrix - - -def _write_tm_matrix_heatmap(path: Path, matrix: np.ndarray, title: str) -> None: - fig, axis = plt.subplots(figsize=(7, 6)) - image = axis.imshow(matrix, cmap="viridis", vmin=0.0, vmax=1.0, aspect="auto") - axis.set_title(title) - axis.set_xlabel("Column sample index") - axis.set_ylabel("Row sample index") - fig.colorbar(image, ax=axis, fraction=0.046, pad=0.04) - fig.tight_layout() - fig.savefig(path, dpi=300) - plt.close(fig) - - -def _write_tm_matrix_artifacts( - matrix_dir: Path, - matrix_name: str, - title: str, - matrix: np.ndarray, -) -> Tuple[str, str, str]: - csv_path = matrix_dir / f"{matrix_name}.csv" - npy_path = matrix_dir / f"{matrix_name}.npy" - png_path = matrix_dir / f"{matrix_name}.png" - np.savetxt(csv_path, matrix, delimiter=",", fmt="%.6f") - np.save(npy_path, matrix) - _write_tm_matrix_heatmap(path=png_path, matrix=matrix, title=title) - return str(csv_path), str(npy_path), str(png_path) - - -def _matrix_stats(matrix: np.ndarray) -> Dict[str, float]: - flattened = matrix.reshape(-1) - return { - "mean": float(np.mean(flattened)), - "median": float(np.median(flattened)), - "min": float(np.min(flattened)), - "max": float(np.max(flattened)), - } - - -def run_boltz2_compliance_suite(args: argparse.Namespace) -> int: - if args.enforce_determinism: - _enforce_determinism() - login_if_needed(args.token) - device = resolve_device(args.device) - dtype = resolve_dtype(args.dtype, device) - _seed_everything(seed=args.seed, workers=False) - output_dir = build_output_dir(args.output_dir, "boltz2_compliance") - checkpoint_path = _download_checkpoint_if_needed(Path(args.checkpoint_path)) - sequences = SEQUENCE_OPTIONS - supports_no_kernels = _detect_no_kernels_support() - - model = AutoModel.from_pretrained(args.repo_id, trust_remote_code=True) - model = model.to(device=device, dtype=torch.float32).eval() - - rows: List[Dict[str, object]] = [] - overall_pass = True - - for sequence_index, sequence in tqdm(list(enumerate(sequences)), desc="Boltz2 sequences", unit="seq"): - started = time.perf_counter() - row: Dict[str, object] = { - "sequence_index": sequence_index, - "sequence": sequence, - "sequence_seed": args.seed + sequence_index, - "ours_dtype_effective": str(dtype), - "num_ours_samples": 0, - "num_ref_samples": 0, - "shared_atoms": 0, - "shared_ca_atoms": 0, - "coord_mae": float("nan"), - "coord_rmse": float("nan"), - "coord_max_abs": float("nan"), - "coord_mae_aligned": float("nan"), - "coord_rmse_aligned": float("nan"), - "coord_max_abs_aligned": float("nan"), - "pairwise_dist_mae": float("nan"), - "plddt_mae": float("nan"), - "ptm_abs_diff": float("nan"), - "iptm_abs_diff": float("nan"), - "complex_plddt_abs_diff": float("nan"), - "confidence_score_abs_diff": float("nan"), - "tm_cross_median": float("nan"), - "tm_cross_mean": float("nan"), - "tm_cross_min": float("nan"), - "tm_cross_max": float("nan"), - "tm_ref_within_median": float("nan"), - "tm_ref_within_mean": float("nan"), - "tm_ref_within_min": float("nan"), - "tm_ref_within_max": float("nan"), - "tm_ours_within_median": float("nan"), - "tm_ours_within_mean": float("nan"), - "tm_ours_within_min": float("nan"), - "tm_ours_within_max": float("nan"), - "tm_official_vs_ours_csv": "", - "tm_official_vs_ours_npy": "", - "tm_official_vs_ours_png": "", - "tm_official_vs_official_csv": "", - "tm_official_vs_official_npy": "", - "tm_official_vs_official_png": "", - "tm_ours_vs_ours_csv": "", - "tm_ours_vs_ours_npy": "", - "tm_ours_vs_ours_png": "", - "ours_cif_path": "", - "ref_cif_path": "", - "pass": False, - "seconds": 0.0, - "error": "", - } - - try: - feats, template = build_boltz2_features( - amino_acid_sequence=sequence, - num_bins=model.config.num_bins, - atoms_per_window_queries=model.core.input_embedder.atom_encoder.atoms_per_window_queries, - ) - feats_ours = _to_device(_clone_feats(feats), device=device, dtype=torch.float32) - - try: - out_ours = _run_ours_forward( - model=model, - feats_ours=feats_ours, - args=args, - device=device, - dtype=dtype, - sequence_index=sequence_index, - ) - except RuntimeError as exc: - error_text = str(exc) - bf16_mismatch = "expected scalar type Float but found BFloat16" in error_text - fp16_mismatch = "expected scalar type Float but found Half" in error_text - if bf16_mismatch or fp16_mismatch: - out_ours = _run_ours_forward( - model=model, - feats_ours=feats_ours, - args=args, - device=device, - dtype=torch.float32, - sequence_index=sequence_index, - ) - row["ours_dtype_effective"] = str(torch.float32) - else: - raise - - ours_atom_maps = _build_ours_atom_maps( - sample_coords=out_ours["sample_atom_coords"], - atom_mask=feats_ours["atom_pad_mask"][0], - atom_names=template.atom_names, - atom_residue_index=template.atom_residue_index, - atom_chain_id=template.atom_chain_id, - ) - ref_atom_maps, ref_plddt_samples, ref_confidence_samples = _run_boltz_cli_reference( - sequence=sequence, - sequence_index=sequence_index, - checkpoint_path=checkpoint_path, - args=args, - device=device, - supports_no_kernels=supports_no_kernels, - ) - - row["num_ours_samples"] = len(ours_atom_maps) - row["num_ref_samples"] = len(ref_atom_maps) - - ours_atom_map_primary = ours_atom_maps[0] - ref_atom_map_primary = ref_atom_maps[0] - ref_plddt_primary = ref_plddt_samples[0].float().cpu().reshape(-1) - ref_confidence_primary = ref_confidence_samples[0] - - shared_keys = [] - for key in ours_atom_map_primary: - if key in ref_atom_map_primary: - shared_keys.append(key) - shared_keys.sort() - assert len(shared_keys) > 0, "No overlapping atom keys between our output and pip boltz CLI output." - row["shared_atoms"] = len(shared_keys) - - shared_ca_keys = _shared_ca_key_order(ours_atom_maps=ours_atom_maps, ref_atom_maps=ref_atom_maps) - row["shared_ca_atoms"] = len(shared_ca_keys) - - ours_coords_stack = torch.stack([ours_atom_map_primary[key] for key in shared_keys], dim=0) - ref_coords_stack = torch.stack([ref_atom_map_primary[key] for key in shared_keys], dim=0) - coord_mae, coord_rmse, coord_max_abs = _vector_metrics(ours_coords_stack, ref_coords_stack) - row["coord_mae"] = coord_mae - row["coord_rmse"] = coord_rmse - row["coord_max_abs"] = coord_max_abs - ours_coords_aligned = _kabsch_align_mobile_to_target(ours_coords_stack, ref_coords_stack) - coord_mae_aligned, coord_rmse_aligned, coord_max_abs_aligned = _vector_metrics( - ours_coords_aligned, - ref_coords_stack, - ) - row["coord_mae_aligned"] = coord_mae_aligned - row["coord_rmse_aligned"] = coord_rmse_aligned - row["coord_max_abs_aligned"] = coord_max_abs_aligned - row["pairwise_dist_mae"] = _pairwise_distance_mae(ours_coords_stack, ref_coords_stack) - - ours_plddt = _extract_primary_plddt_vector(out_ours, feats_ours) - assert ours_plddt.numel() == ref_plddt_primary.numel(), ( - f"pLDDT size mismatch (ours={ours_plddt.numel()}, ref={ref_plddt_primary.numel()})." - ) - row["plddt_mae"] = float(torch.mean(torch.abs(ours_plddt - ref_plddt_primary)).item()) - - ours_ptm = _summary_metric(out_ours["ptm"]).float().cpu() - ours_iptm = _summary_metric(out_ours["iptm"]).float().cpu() - ours_complex_plddt = _summary_metric(out_ours["complex_plddt"]).float().cpu() - ours_confidence_score = _compute_confidence_score( - ptm=ours_ptm, - iptm=ours_iptm, - complex_plddt=ours_complex_plddt, - ) - - ref_ptm = torch.tensor([float(ref_confidence_primary["ptm"])], dtype=torch.float32) - ref_iptm = torch.tensor([float(ref_confidence_primary["iptm"])], dtype=torch.float32) - ref_complex_plddt = torch.tensor([float(ref_confidence_primary["complex_plddt"])], dtype=torch.float32) - ref_confidence_score = torch.tensor([float(ref_confidence_primary["confidence_score"])], dtype=torch.float32) - - row["ptm_abs_diff"] = float(torch.mean(torch.abs(ours_ptm - ref_ptm)).item()) - row["iptm_abs_diff"] = float(torch.mean(torch.abs(ours_iptm - ref_iptm)).item()) - row["complex_plddt_abs_diff"] = float( - torch.mean(torch.abs(ours_complex_plddt - ref_complex_plddt)).item() - ) - row["confidence_score_abs_diff"] = float( - torch.mean(torch.abs(ours_confidence_score - ref_confidence_score)).item() - ) - - if args.write_cif_artifacts: - structure_dir = output_dir / "structures" / f"seq_{sequence_index}" - structure_dir.mkdir(parents=True, exist_ok=True) - - ours_cif_path = structure_dir / f"ours_seq{sequence_index}.cif" - write_cif( - structure_template=template, - atom_coords=out_ours["sample_atom_coords"].detach().cpu(), - atom_mask=feats_ours["atom_pad_mask"][0].detach().cpu(), - output_path=str(ours_cif_path), - plddt=out_ours["plddt"].detach().cpu() if "plddt" in out_ours else None, - sample_index=0, - ) - row["ours_cif_path"] = str(ours_cif_path) - - ref_coords_cif, ref_atom_mask_cif = _build_reference_cif_tensors( - template=template, - atom_pad_mask=feats_ours["atom_pad_mask"][0].detach().cpu(), - ref_atom_map=ref_atom_map_primary, - ) - ref_cif_path = structure_dir / f"ref_seq{sequence_index}.cif" - write_cif( - structure_template=template, - atom_coords=ref_coords_cif, - atom_mask=ref_atom_mask_cif, - output_path=str(ref_cif_path), - plddt=ref_plddt_primary, - sample_index=0, - ) - row["ref_cif_path"] = str(ref_cif_path) - - ours_ca_stack = _stack_coords_for_keys(atom_maps=ours_atom_maps, ordered_keys=shared_ca_keys) - ref_ca_stack = _stack_coords_for_keys(atom_maps=ref_atom_maps, ordered_keys=shared_ca_keys) - - tm_official_vs_ours = _build_tm_matrix( - reference_stack=ref_ca_stack, - subject_stack=ours_ca_stack, - symmetric=False, - ) - tm_official_vs_official = _build_tm_matrix( - reference_stack=ref_ca_stack, - subject_stack=ref_ca_stack, - symmetric=True, - ) - tm_ours_vs_ours = _build_tm_matrix( - reference_stack=ours_ca_stack, - subject_stack=ours_ca_stack, - symmetric=True, - ) - - matrix_dir = output_dir / "tm_matrices" / f"seq_{sequence_index}" - matrix_dir.mkdir(parents=True, exist_ok=True) - csv_path, npy_path, png_path = _write_tm_matrix_artifacts( - matrix_dir=matrix_dir, - matrix_name="official_vs_ours", - title=f"Sequence {sequence_index}: official vs ours TM-score", - matrix=tm_official_vs_ours, - ) - row["tm_official_vs_ours_csv"] = csv_path - row["tm_official_vs_ours_npy"] = npy_path - row["tm_official_vs_ours_png"] = png_path - - csv_path, npy_path, png_path = _write_tm_matrix_artifacts( - matrix_dir=matrix_dir, - matrix_name="official_vs_official", - title=f"Sequence {sequence_index}: official vs official TM-score", - matrix=tm_official_vs_official, - ) - row["tm_official_vs_official_csv"] = csv_path - row["tm_official_vs_official_npy"] = npy_path - row["tm_official_vs_official_png"] = png_path - - csv_path, npy_path, png_path = _write_tm_matrix_artifacts( - matrix_dir=matrix_dir, - matrix_name="ours_vs_ours", - title=f"Sequence {sequence_index}: ours vs ours TM-score", - matrix=tm_ours_vs_ours, - ) - row["tm_ours_vs_ours_csv"] = csv_path - row["tm_ours_vs_ours_npy"] = npy_path - row["tm_ours_vs_ours_png"] = png_path - - cross_stats = _matrix_stats(tm_official_vs_ours) - row["tm_cross_mean"] = cross_stats["mean"] - row["tm_cross_median"] = cross_stats["median"] - row["tm_cross_min"] = cross_stats["min"] - row["tm_cross_max"] = cross_stats["max"] - - ref_within_stats = _matrix_stats(tm_official_vs_official) - row["tm_ref_within_mean"] = ref_within_stats["mean"] - row["tm_ref_within_median"] = ref_within_stats["median"] - row["tm_ref_within_min"] = ref_within_stats["min"] - row["tm_ref_within_max"] = ref_within_stats["max"] - - ours_within_stats = _matrix_stats(tm_ours_vs_ours) - row["tm_ours_within_mean"] = ours_within_stats["mean"] - row["tm_ours_within_median"] = ours_within_stats["median"] - row["tm_ours_within_min"] = ours_within_stats["min"] - row["tm_ours_within_max"] = ours_within_stats["max"] - - row["pass"] = bool(row["tm_cross_median"] >= args.tm_pass_threshold) - if row["pass"] is False: - overall_pass = False - except Exception as exc: - row["error"] = str(exc) - overall_pass = False - finally: - row["seconds"] = round(time.perf_counter() - started, 4) - rows.append(row) - - payload: Dict[str, object] = { - "suite": "boltz2_compliance", - "all_passed": overall_pass, - "repo_id": args.repo_id, - "checkpoint_path": str(checkpoint_path), - "device": str(device), - "dtype": str(dtype), - "seed": args.seed, - "enforce_determinism": args.enforce_determinism, - "write_cif_artifacts": args.write_cif_artifacts, - "num_sequences": len(sequences), - "recycling_steps": BOLTZ2_FIXED_RECYCLING_STEPS, - "num_sampling_steps": BOLTZ2_FIXED_SAMPLING_STEPS, - "diffusion_samples": BOLTZ2_FIXED_DIFFUSION_SAMPLES, - "tm_pass_threshold": args.tm_pass_threshold, - "supports_no_kernels": supports_no_kernels, - "rows": rows, - } - write_json(output_dir / "metrics.json", payload) - write_csv(output_dir / "metrics.csv", rows) - - passed_count = 0 - for row in rows: - if bool(row["pass"]): - passed_count += 1 - summary_lines = [ - "Suite: boltz2_compliance", - f"Sequences tested: {len(rows)}", - f"Sequences passed: {passed_count}", - f"Sequences failed: {len(rows) - passed_count}", - f"Output directory: {output_dir}", - f"Device: {device}", - f"Dtype: {dtype}", - f"Recycling steps (fixed): {BOLTZ2_FIXED_RECYCLING_STEPS}", - f"Sampling steps (fixed): {BOLTZ2_FIXED_SAMPLING_STEPS}", - f"Diffusion samples (fixed): {BOLTZ2_FIXED_DIFFUSION_SAMPLES}", - f"TM pass threshold: {args.tm_pass_threshold}", - f"Write CIF artifacts: {args.write_cif_artifacts}", - f"Reference CLI supports --no_kernels: {supports_no_kernels}", - ] - for row in rows: - status = "PASS" if bool(row["pass"]) else "FAIL" - summary_lines.append( - f"{status} | idx={row['sequence_index']} | seed={row['sequence_seed']} | " - f"ours_dtype={row['ours_dtype_effective']} | shared_atoms={row['shared_atoms']} | " - f"shared_ca={row['shared_ca_atoms']} | tm_cross_median={row['tm_cross_median']} | " - f"tm_ref_within_median={row['tm_ref_within_median']} | " - f"tm_ours_within_median={row['tm_ours_within_median']} | " - f"coord_aln_mae={row['coord_mae_aligned']} | plddt_mae={row['plddt_mae']} | " - f"official_vs_ours_csv={row['tm_official_vs_ours_csv']} | " - f"official_vs_official_csv={row['tm_official_vs_official_csv']} | " - f"ours_vs_ours_csv={row['tm_ours_vs_ours_csv']} | " - f"ours_cif={row['ours_cif_path']} | ref_cif={row['ref_cif_path']} | error={row['error']}" - ) - write_summary(output_dir / "summary.txt", summary_lines) - print("\n".join(summary_lines)) - - if overall_pass: - return 0 - return 1 - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description="Run Boltz2 compliance test against pip boltz CLI outputs.") - parser.add_argument("--token", type=str, default=None) - parser.add_argument("--repo-id", type=str, default="Synthyra/Boltz2") - parser.add_argument("--checkpoint-path", type=str, default="fastplms/boltz/weights/boltz2_conf.ckpt") - parser.add_argument("--device", type=str, default="auto") - parser.add_argument("--dtype", type=str, default="float32", choices=["auto", "float32", "float16", "bfloat16"]) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--enforce-determinism", action=argparse.BooleanOptionalAction, default=True) - parser.add_argument("--write-cif-artifacts", action=argparse.BooleanOptionalAction, default=True) - parser.add_argument("--pass-coord-metric", type=str, default="aligned", choices=["raw", "aligned"]) - parser.add_argument("--run-confidence-sequentially", action="store_true") - parser.add_argument("--coord-mae-threshold", type=float, default=5e-3) - parser.add_argument("--coord-rmse-threshold", type=float, default=5e-3) - parser.add_argument("--coord-max-abs-threshold", type=float, default=5e-2) - parser.add_argument("--plddt-mae-threshold", type=float, default=5e-3) - parser.add_argument("--summary-metric-abs-threshold", type=float, default=5e-3) - parser.add_argument("--tm-pass-threshold", type=float, default=0.60) - parser.add_argument("--output-dir", type=str, default=None) - return parser - - -def main(argv: Optional[List[str]] = None) -> int: - parser = build_parser() - args = parser.parse_args(argv) - return run_boltz2_compliance_suite(args) - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/testing/test_automodel.py b/testing/test_automodel.py deleted file mode 100644 index 5da0fb4..0000000 --- a/testing/test_automodel.py +++ /dev/null @@ -1,139 +0,0 @@ -"""Test that all FastPLM models load correctly via AutoModelForMaskedLM.""" - -import random - -import pytest -import torch -from transformers import AutoModelForMaskedLM - -from testing.conftest import ( - CANONICAL_AAS, FULL_MODEL_REGISTRY, MODEL_REGISTRY, SEED, - add_model_specific_inputs, mark_by_size, tokenize_batch, -) - - -MODEL_KEYS = list(MODEL_REGISTRY.keys()) -FULL_KEYS = list(FULL_MODEL_REGISTRY.keys()) - - -def _tokenize_single(model, model_key: str, sequence: str, device: torch.device, registry=None): - """Tokenize a single sequence, handling E1's sequence-mode separately.""" - return tokenize_batch(model, model_key, [sequence], device, registry=registry) - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", MODEL_KEYS) -def test_automodel_loads(model_key: str) -> None: - """Model loads via AutoModelForMaskedLM with trust_remote_code=True.""" - config = MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - assert model is not None - # E1 uses sequence-mode (tokenizer lives on model.model.prep_tokens, not model.tokenizer) - if config["uses_tokenizer"]: - assert hasattr(model, "tokenizer") - assert hasattr(model, "attn_backend") - assert model.attn_backend == "sdpa" - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", MODEL_KEYS) -def test_automodel_forward_pass(model_key: str) -> None: - """Single forward pass produces valid logits with no NaN.""" - random.seed(SEED) - config = MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - sequence = "M" + "".join(random.choices(CANONICAL_AAS, k=31)) - inputs = _tokenize_single(model, model_key, sequence, device) - - with torch.inference_mode(): - output = model(**inputs) - - assert hasattr(output, "logits") - logits = output.logits - assert logits.ndim == 3 - assert logits.shape[0] == 1 - assert not torch.isnan(logits).any(), f"NaN in logits for {model_key}" - assert not torch.isinf(logits).any(), f"Inf in logits for {model_key}" - - del model - torch.cuda.empty_cache() - - -# --------------------------------------------------------------------------- -# Full model registry tests: all checkpoints across all families -# --------------------------------------------------------------------------- - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", mark_by_size(FULL_KEYS, FULL_MODEL_REGISTRY)) -def test_full_automodel_loads(model_key: str) -> None: - """Every checkpoint loads via AutoModelForMaskedLM and exposes expected attributes.""" - config = FULL_MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - assert model is not None - if config["uses_tokenizer"]: - assert hasattr(model, "tokenizer") - assert hasattr(model, "attn_backend") - assert model.attn_backend == "sdpa" - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", mark_by_size(FULL_KEYS, FULL_MODEL_REGISTRY)) -def test_full_automodel_forward(model_key: str) -> None: - """Every checkpoint produces valid logits on a single forward pass.""" - random.seed(SEED) - config = FULL_MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - sequence = "M" + "".join(random.choices(CANONICAL_AAS, k=31)) - inputs = _tokenize_single(model, model_key, sequence, device, registry=FULL_MODEL_REGISTRY) - inputs = add_model_specific_inputs(inputs, config["model_type"]) - - with torch.inference_mode(): - output = model(**inputs) - - assert hasattr(output, "logits") - logits = output.logits - assert logits.ndim == 3 - assert logits.shape[0] == 1 - assert not torch.isnan(logits).any(), f"NaN in logits for {model_key}" - assert not torch.isinf(logits).any(), f"Inf in logits for {model_key}" - - del model - torch.cuda.empty_cache() diff --git a/testing/test_backend_consistency.py b/testing/test_backend_consistency.py deleted file mode 100644 index 5058e70..0000000 --- a/testing/test_backend_consistency.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Test that all attention backends produce consistent outputs for each model. - -Loads each model once in float32, runs forward passes with SDPA, Flex, and -Flash backends, and verifies that logits are within tolerance. -""" - -import random -from typing import Dict, List - -import pytest -import torch -from transformers import AutoModelForMaskedLM - -from testing.conftest import ( - BACKENDS, CANONICAL_AAS, FULL_MODEL_REGISTRY, MODEL_REGISTRY, SEED, - add_model_specific_inputs, mark_by_size, tokenize_batch, -) - - -MODEL_KEYS = list(MODEL_REGISTRY.keys()) -FULL_KEYS = list(FULL_MODEL_REGISTRY.keys()) -# bfloat16 has ~3 decimal digits precision; backends differ in tiling/accumulation order. -# SDPA vs Flex can show max abs diffs up to ~0.5 in logit space at bfloat16. -# We check that predictions (argmax) agree rather than raw logit values. -PRED_AGREEMENT_THRESHOLDS = { - "default": 0.95, - # ESMC's 30-layer stack accumulates bfloat16 backend rounding in the LM head. - # The stricter parity suite checks native parity and pooled cosine separately. - "ESMC": 0.90, -} -NUM_SEQUENCES = 4 -SEQ_LEN = 64 - - -def _generate_sequences(model_key: str) -> List[str]: - """Generate test sequences (fixed-length for reproducibility).""" - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=SEQ_LEN - 1)) - for _ in range(NUM_SEQUENCES) - ] - - -def _tokenize_batch( - model, - model_key: str, - sequences: List[str], - device: torch.device, - registry: Dict[str, Dict] = None, -) -> Dict[str, torch.Tensor]: - """Tokenize sequences, handling E1's sequence mode.""" - if registry is None: - registry = MODEL_REGISTRY - config = registry[model_key] if model_key in registry else MODEL_REGISTRY[model_key] - if config["model_type"] == "E1": - batch = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) - return { - "input_ids": batch["input_ids"], - "within_seq_position_ids": batch["within_seq_position_ids"], - "global_position_ids": batch["global_position_ids"], - "sequence_ids": batch["sequence_ids"], - "attention_mask": (batch["sequence_ids"] != -1).long(), - } - tokenizer = model.tokenizer - tokenized = tokenizer( - sequences, - return_tensors="pt", - padding="max_length", - max_length=SEQ_LEN + 2, # account for special tokens - truncation=True, - ) - return {k: v.to(device) for k, v in tokenized.items()} - - -def _run_backend_consistency(model_key: str, registry: Dict[str, Dict]) -> None: - """Core backend consistency logic shared by default and full-registry tests.""" - random.seed(SEED) - config = registry[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - sequences = _generate_sequences(model_key) - inputs = _tokenize_batch(model, model_key, sequences, device, registry=registry) - - model_inputs = inputs.copy() - model_inputs = add_model_specific_inputs(model_inputs, config["model_type"]) - - backend_logits: Dict[str, torch.Tensor] = {} - - for backend in BACKENDS: - try: - model.attn_backend = backend - except (AssertionError, RuntimeError) as e: - print(f"Skipping backend '{backend}' for {model_key}: {e}") - continue - - try: - with torch.inference_mode(): - output = model(**model_inputs) - except (AssertionError, RuntimeError) as e: - print(f"Backend '{backend}' failed at runtime for {model_key}: {e}") - continue - - backend_logits[backend] = output.logits.cpu() - - assert len(backend_logits) >= 1, f"No backends available for {model_key}" - - if "sdpa" not in backend_logits: - pytest.skip(f"SDPA backend not available for {model_key}, cannot compare") - - reference = backend_logits["sdpa"] - attention_mask = inputs["attention_mask"].cpu().bool() - ref_masked = reference[attention_mask] - ref_preds = ref_masked.argmax(dim=-1) - - for backend, logits in backend_logits.items(): - if backend == "sdpa": - continue - cand_masked = logits[attention_mask] - cand_preds = cand_masked.argmax(dim=-1) - agreement = (ref_preds == cand_preds).float().mean().item() - model_type = config["model_type"] - if model_type in PRED_AGREEMENT_THRESHOLDS: - threshold = PRED_AGREEMENT_THRESHOLDS[model_type] - else: - threshold = PRED_AGREEMENT_THRESHOLDS["default"] - assert agreement >= threshold, ( - f"{model_key}: SDPA vs {backend} prediction agreement = {agreement:.4f} " - f"(threshold: {threshold})" - ) - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", MODEL_KEYS) -def test_backend_consistency(model_key: str) -> None: - """All available backends produce equivalent logits (default registry).""" - _run_backend_consistency(model_key, MODEL_REGISTRY) - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", mark_by_size(FULL_KEYS, FULL_MODEL_REGISTRY)) -def test_full_backend_consistency(model_key: str) -> None: - """All available backends produce equivalent logits (all checkpoints).""" - _run_backend_consistency(model_key, FULL_MODEL_REGISTRY) diff --git a/testing/test_binder_design_fastplms.py b/testing/test_binder_design_fastplms.py deleted file mode 100644 index efb1fbf..0000000 --- a/testing/test_binder_design_fastplms.py +++ /dev/null @@ -1,552 +0,0 @@ -"""Tests for the FastPLMs binder design tutorial.""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import pytest -import torch - -import cookbook.tutorials.binder_design_fastplms as binder - - -@dataclass -class TinyTransformerOutput: - last_hidden_state: torch.Tensor - - -class TinyTransformer(torch.nn.Module): - def forward( - self, - x: torch.Tensor, - attention_mask: torch.Tensor | None = None, - output_hidden_states: bool = False, - output_attentions: bool = False, - ) -> TinyTransformerOutput: - del attention_mask, output_hidden_states, output_attentions - return TinyTransformerOutput(last_hidden_state=x) - - -class TinyTokenizer: - cls_token_id = 0 - eos_token_id = 1 - mask_token_id = 2 - - -class TinyConfig: - vocab_size = 24 - - -class TinyLM(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.config = TinyConfig() - self.tokenizer = TinyTokenizer() - self.embed = torch.nn.Embedding(self.config.vocab_size, 8) - self.transformer = TinyTransformer() - self.sequence_head = torch.nn.Linear(8, self.config.vocab_size) - - -class FakeInputBuilder: - def decode( - self, - output: dict[str, torch.Tensor], - inputs: dict[str, torch.Tensor], - chain_infos: list[Any], - num_diffusion_samples: int, - complex_id: str, - ) -> dict[str, Any]: - del output, inputs, chain_infos, num_diffusion_samples - return {"complex_id": complex_id} - - -class FakeCritic: - input_builder = FakeInputBuilder() - - def result_to_cif(self, complex_result: dict[str, Any]) -> str: - return f"data_{complex_result['complex_id']}\n" - - def result_to_pdb(self, complex_result: dict[str, Any]) -> str: - return "HEADER FASTPLMS TEST\nEND\n" - - -class FakeScalingCritic(FakeCritic): - def __init__(self) -> None: - self.device_moves: list[str] = [] - - def to(self, device: torch.device | str) -> "FakeScalingCritic": - self.device_moves.append(str(device)) - return self - - -@dataclass -class FakeProteinInput: - id: str - sequence: str - msa: Any - - -@dataclass -class FakeStructurePredictionInput: - sequences: list[FakeProteinInput] - - -class FakeInputTypes: - ProteinInput = FakeProteinInput - StructurePredictionInput = FakeStructurePredictionInput - - -class FakeFoldModel(torch.nn.Module): - input_types = FakeInputTypes() - device = torch.device("cpu") - - def __init__(self) -> None: - super().__init__() - self.saw_model_input_type = False - - def prepare_structure_input( - self, input_data: FakeStructurePredictionInput, seed: int | None = None - ) -> tuple[dict[str, torch.Tensor], list[Any]]: - del seed - assert isinstance(input_data, FakeStructurePredictionInput) - self.saw_model_input_type = True - features = { - "dummy": torch.zeros(1, 1), - "pocket_feature": torch.ones(1, 1), - } - return features, [] - - def forward( - self, dummy: torch.Tensor, res_type_soft: torch.Tensor - ) -> dict[str, torch.Tensor]: - del dummy - batch, total_length, _ = res_type_soft.shape - distogram_logits = torch.zeros(batch, total_length, total_length, 128) - return {"distogram_logits": distogram_logits} - - -def test_prompt_sampling_is_reproducible() -> None: - factory = binder.BINDER_PROMPT_FACTORIES["trastuzumab_framework_vhvl"] - - first = factory.sample(seed=17) - second = factory.sample(seed=17) - third = factory.sample(seed=18) - - assert first == second - assert first != third - assert binder.MUTABLE_TOKEN in first - assert 160 <= len(first) <= 260 - - -def test_fold_uses_model_input_types() -> None: - model = FakeFoldModel() - target_one_hot = binder.sequence_to_one_hot("ACD", device="cpu") - logits = torch.randn(1, 3, 20) - design = torch.softmax(logits, dim=-1) - - result = binder.fold_and_get_distogram( - model=model, - target_seq="ACD", - target_one_hot=target_one_hot, - design=design, - num_sampling_steps=1, - ) - - assert model.saw_model_input_type - assert result["distogram_logits"].shape == (1, 6, 6, 128) - - -def test_design_input_validation_rejects_ambiguous_targets() -> None: - with pytest.raises(AssertionError, match="Provide either target name"): - binder.design_binder( - inversion_models={}, - critic_models={}, - lm_model=object(), - target_name="pd-l1", - target_sequence="ACDE", - binder_name="minibinder", - binder_sequence=None, - is_antibody=None, - seed=0, - steps=1, - ) - - -def test_design_input_validation_rejects_ambiguous_binders() -> None: - with pytest.raises(AssertionError, match="Provide either binder name"): - binder.design_binder( - inversion_models={}, - critic_models={}, - lm_model=object(), - target_name="pd-l1", - target_sequence=None, - binder_name="minibinder", - binder_sequence="####", - is_antibody=None, - seed=0, - steps=1, - ) - - -def test_fastplms_pseudoperplexity_nll_is_differentiable() -> None: - torch.manual_seed(0) - lm_model = TinyLM() - logits = torch.randn(1, 5, 20, requires_grad=True) - design = torch.softmax(logits, dim=-1) - score_mask = torch.ones(1, 5, dtype=torch.bool) - - loss = binder.compute_fastplms_pseudoperplexity_nll( - lm_model=lm_model, - binder_design=design, - score_mask=score_mask, - batch_size=1, - n_passes=2, - mask_fraction=0.5, - ) - loss.sum().backward() - - assert loss.shape == (1,) - assert torch.isfinite(loss).all() - assert logits.grad is not None - assert torch.isfinite(logits.grad).all() - assert logits.grad.abs().sum().item() > 0 - - -def test_antibody_proxy_uses_supplied_cdr_indices( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def fail_cdr_lookup(binder_sequence: str) -> list[int]: - del binder_sequence - raise AssertionError("CDR lookup should not be called") - - monkeypatch.setattr(binder, "_cdr_indices", fail_cdr_lookup) - distogram_logits = torch.zeros(1, 6, 6, 128) - - scores = binder.compute_distogram_iptm_proxy( - distogram_logits=distogram_logits, - target_length=3, - binder_sequence="AAA", - is_antibody=True, - cdr_indices=[0, 2], - ) - - assert 0.0 <= scores["distogram_iptm_proxy"] <= 1.0 - assert 0.0 <= scores["cdr_distogram_iptm_proxy"] <= 1.0 - - -def test_official_selection_uses_mean_hero_iptm_and_scaling_proxy() -> None: - rows = [ - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.9, - "distogram_iptm_proxy": 0.1, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Cutoff2025", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.5, - "distogram_iptm_proxy": 0.2, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Fast-base300M-step250k", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": None, - "distogram_iptm_proxy": 0.8, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|EEEEEE", - "is_antibody": False, - "iptm": 0.95, - "distogram_iptm_proxy": 0.1, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 1, - }, - { - "critic_name": "ESMFold2-Experimental-Cutoff2025", - "designed_sequence": "AAA|EEEEEE", - "is_antibody": False, - "iptm": 0.95, - "distogram_iptm_proxy": 0.2, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 1, - }, - { - "critic_name": "ESMFold2-Experimental-Fast-base300M-step250k", - "designed_sequence": "AAA|EEEEEE", - "is_antibody": False, - "iptm": None, - "distogram_iptm_proxy": 0.2, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 1, - }, - ] - - selection = binder.select_official_designs(rows) - - assert selection.iloc[0]["designed_sequence"] == "AAA|DEDEDE" - assert selection.iloc[0]["iptm_score"] == pytest.approx(0.7) - assert selection.iloc[0]["iptm_proxy_score"] == pytest.approx(0.8) - assert selection.iloc[0]["selection_score"] == pytest.approx(0.75) - assert selection.iloc[0]["hero_iptm_min"] == pytest.approx(0.5) - assert not bool(selection.iloc[0]["all_hero_critics_pass"]) - - -def test_official_selection_filters_high_pi_minibinders() -> None: - rows = [ - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|KKKKKK", - "is_antibody": False, - "iptm": 0.99, - "distogram_iptm_proxy": 0.99, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.5, - "distogram_iptm_proxy": 0.5, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 1, - }, - ] - - selection = binder.select_official_designs(rows) - - assert selection["designed_sequence"].tolist() == ["AAA|DEDEDE"] - - -def test_official_selection_empty_after_pi_filter_keeps_schema() -> None: - rows = [ - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|KKKKKK", - "is_antibody": False, - "iptm": 0.99, - "distogram_iptm_proxy": 0.99, - "cdr_distogram_iptm_proxy": float("nan"), - } - ] - - selection = binder.select_official_designs(rows) - - assert selection.empty - assert "selection_score" in selection.columns - assert "hero_iptm_min" in selection.columns - assert "all_hero_critics_pass" in selection.columns - - -def test_official_selection_uses_cdr_proxy_for_scaling_antibodies() -> None: - rows = [ - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|EVQLVESGGG", - "is_antibody": True, - "iptm": 0.6, - "distogram_iptm_proxy": 0.1, - "cdr_distogram_iptm_proxy": 0.2, - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Fast-base300M-step250k", - "designed_sequence": "AAA|EVQLVESGGG", - "is_antibody": True, - "iptm": None, - "distogram_iptm_proxy": 0.1, - "cdr_distogram_iptm_proxy": 0.9, - "batch_idx": 0, - }, - ] - - selection = binder.select_official_designs(rows) - - assert selection.iloc[0]["iptm_score"] == pytest.approx(0.6) - assert selection.iloc[0]["iptm_proxy_score"] == pytest.approx(0.9) - assert selection.iloc[0]["selection_score"] == pytest.approx(0.75) - - -def test_official_selection_flags_all_hero_critics_pass() -> None: - rows = [ - { - "critic_name": "ESMFold2-Experimental-Fast", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.91, - "distogram_iptm_proxy": 0.4, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Fast-Cutoff2025", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.92, - "distogram_iptm_proxy": 0.4, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.93, - "distogram_iptm_proxy": 0.4, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - { - "critic_name": "ESMFold2-Experimental-Cutoff2025", - "designed_sequence": "AAA|DEDEDE", - "is_antibody": False, - "iptm": 0.94, - "distogram_iptm_proxy": 0.4, - "cdr_distogram_iptm_proxy": float("nan"), - "batch_idx": 0, - }, - ] - - selection = binder.select_official_designs(rows) - - assert bool(selection.iloc[0]["all_hero_critics_pass"]) - assert selection.iloc[0]["consensus_iptm_threshold"] == pytest.approx(0.9) - - -@pytest.mark.gpu -@pytest.mark.slow -def test_fastplms_esmplusplus_small_pseudoperplexity_smoke() -> None: - from transformers import AutoModelForMaskedLM - - model = ( - AutoModelForMaskedLM.from_pretrained( - "Synthyra/ESMplusplus_small", - trust_remote_code=True, - dtype=torch.float32, - ) - .cuda() - .eval() - .requires_grad_(False) - ) - logits = torch.randn(1, 4, 20, device="cuda", requires_grad=True) - design = torch.softmax(logits, dim=-1) - score_mask = torch.ones(1, 4, dtype=torch.bool, device="cuda") - - loss = binder.compute_fastplms_pseudoperplexity_nll( - lm_model=model, - binder_design=design, - score_mask=score_mask, - batch_size=1, - n_passes=1, - mask_fraction=0.5, - ) - loss.sum().backward() - - assert loss.shape == (1,) - assert torch.isfinite(loss).all() - assert logits.grad is not None - assert torch.isfinite(logits.grad).all() - assert logits.grad.abs().sum().item() > 0 - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -def test_tiny_design_dry_run_writes_outputs( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - def fake_fold_and_get_distogram( - model: Any, - target_seq: str, - target_one_hot: torch.Tensor, - design: torch.Tensor, - num_loops: int = 0, - num_sampling_steps: int = 1, - calculate_confidence: bool = False, - seed: int | None = None, - ) -> dict[str, Any]: - del model, num_loops, num_sampling_steps, calculate_confidence, seed - batch, binder_length, aa_dim = design.shape - target_length = target_one_hot.size(1) - total_length = target_length + binder_length - aa_weight = torch.linspace(-1.0, 1.0, aa_dim, device=design.device) - binder_signal = (design * aa_weight).sum(dim=-1) - token_signal = torch.cat( - (torch.zeros(batch, target_length, device=design.device), binder_signal), - dim=1, - ) - pair_signal = token_signal[:, :, None] + token_signal[:, None, :] - bin_basis = torch.linspace(-1.0, 1.0, 128, device=design.device) - distogram_logits = pair_signal.unsqueeze(-1) * bin_basis - seq_list = [target_seq + "|" + "A" * binder_length for _ in range(batch)] - return { - "distogram_logits": distogram_logits, - "inputs": {}, - "chain_info_list": [[] for _ in range(batch)], - "output": {"distogram_logits": distogram_logits}, - "seq_list": seq_list, - "iptm": torch.ones(batch, device=design.device), - "plddt": torch.ones(batch, 1, device=design.device), - } - - def fake_pseudoperplexity_nll( - lm_model: Any, - binder_design: torch.Tensor, - score_mask: torch.Tensor, - batch_size: int = 4, - n_passes: int = 4, - mask_fraction: float = binder.DEFAULT_ESMC_MASK_FRACTION, - ) -> torch.Tensor: - del lm_model, score_mask, batch_size, n_passes, mask_fraction - return binder_design.square().mean(dim=(1, 2)) - - monkeypatch.setattr(binder, "fold_and_get_distogram", fake_fold_and_get_distogram) - monkeypatch.setattr( - binder, "compute_fastplms_pseudoperplexity_nll", fake_pseudoperplexity_nll - ) - - scaling_critic = FakeScalingCritic() - best_sequences, trajectory, rows = binder.design_binder( - inversion_models={"fake_inversion": object()}, - critic_models={ - "fake_critic": FakeCritic(), - "ESMFold2-Experimental-Fast-base300M-step250k": scaling_critic, - }, - lm_model=object(), - target_name=None, - target_sequence="ACD", - binder_name=None, - binder_sequence="###", - is_antibody=False, - seed=0, - batch_size=1, - steps=1, - output_dir=tmp_path, - ) - - assert best_sequences == ["ACD|AAA"] - assert list(trajectory) == [0] - assert len(rows) == 2 - assert rows[0]["binder_sequence"] == "AAA" - assert rows[0]["binder_length"] == 3 - assert rows[0]["target_length"] == 3 - assert rows[0]["mean_plddt"] == 1.0 - assert scaling_critic.device_moves == ["cuda", "cpu"] - assert (tmp_path / "trajectory.jsonl").exists() - assert (tmp_path / "best_sequences.fasta").exists() - assert (tmp_path / "results.parquet").exists() - assert (tmp_path / "selection.parquet").exists() diff --git a/testing/test_compliance.py b/testing/test_compliance.py deleted file mode 100644 index 3dfcf92..0000000 --- a/testing/test_compliance.py +++ /dev/null @@ -1,447 +0,0 @@ -"""Weight and forward-pass compliance tests against original implementations. - -Tests that FastPLM weights are bit-exact with the originals and that forward -pass outputs (logits, hidden states) are numerically equivalent. - -Marked as `slow` because each test loads two models simultaneously. -""" - -import importlib -import random -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import pytest -import torch -from transformers import AutoModelForMaskedLM - -from testing.conftest import ( - CANONICAL_AAS, FULL_MODEL_REGISTRY, MODEL_REGISTRY, SEED, - add_model_specific_inputs, mark_by_size, strict_fp32_matmul, -) -from fastplms.weight_parity_utils import assert_state_dict_equal - - -MODEL_KEYS = list(MODEL_REGISTRY.keys()) - -TEST_NUM_BATCHES = 25 -BATCH_SIZE = 8 -MIN_SEQ_LEN = 16 -MAX_SEQ_LEN = 128 - -FORWARD_DTYPE = torch.float32 - - -@dataclass(frozen=True) -class ForwardComplianceTolerances: - hidden_mse: Optional[float] = 1e-8 - hidden_maxabs: Optional[float] = 5e-4 - hidden_rel_std: float = 5e-3 - hidden_rel_maxabs: float = 5e-2 - last_hidden_mse: float = 1e-8 - last_hidden_maxabs: float = 5e-4 - last_hidden_rel_maxabs: float = 5e-4 - logits_mse: float = 1e-8 - logits_maxabs: float = 5e-4 - - -FORWARD_COMPLIANCE_TOLERANCES: Dict[str, ForwardComplianceTolerances] = { - "ESM2": ForwardComplianceTolerances( - hidden_mse=1e-12, - hidden_maxabs=1e-5, - hidden_rel_std=1e-6, - hidden_rel_maxabs=1e-5, - last_hidden_mse=1e-12, - last_hidden_maxabs=1e-5, - last_hidden_rel_maxabs=1e-5, - logits_mse=1e-12, - logits_maxabs=1e-5, - ), - "ESMC": ForwardComplianceTolerances( - # ESMC intermediate states are high-magnitude pre-norm streams; use - # tight relative guards there and absolute guards after final norm. - hidden_mse=None, - hidden_maxabs=None, - hidden_rel_std=1e-4, - hidden_rel_maxabs=1e-4, - last_hidden_mse=1e-8, - last_hidden_maxabs=1e-3, - last_hidden_rel_maxabs=1e-3, - logits_mse=1e-4, - logits_maxabs=5e-3, - ), - "ESM3": ForwardComplianceTolerances( - hidden_mse=1e-8, - hidden_maxabs=1e-3, - hidden_rel_std=5e-3, - hidden_rel_maxabs=5e-2, - last_hidden_mse=1e-8, - last_hidden_maxabs=1e-3, - last_hidden_rel_maxabs=1e-3, - logits_mse=1e-4, - logits_maxabs=5e-3, - ), - "E1": ForwardComplianceTolerances( - hidden_mse=5e-7, - hidden_maxabs=2e-2, - hidden_rel_std=1e-2, - hidden_rel_maxabs=2e-2, - last_hidden_mse=5e-7, - last_hidden_maxabs=2e-2, - last_hidden_rel_maxabs=2e-3, - logits_mse=1e-4, - logits_maxabs=5e-2, - ), - "DPLM": ForwardComplianceTolerances(), - "DPLM2": ForwardComplianceTolerances( - hidden_mse=1e-12, - hidden_maxabs=1e-5, - hidden_rel_std=1e-5, - hidden_rel_maxabs=1e-5, - last_hidden_mse=1e-12, - last_hidden_maxabs=1e-5, - last_hidden_rel_maxabs=1e-5, - logits_mse=1e-4, - logits_maxabs=5e-3, - ), - "ANKH": ForwardComplianceTolerances( - hidden_mse=1e-12, - hidden_maxabs=1e-5, - hidden_rel_std=1e-5, - hidden_rel_maxabs=1e-5, - last_hidden_mse=1e-12, - last_hidden_maxabs=1e-5, - last_hidden_rel_maxabs=1e-5, - logits_mse=1e-6, - logits_maxabs=5e-3, - ), -} - - -def _generate_random_batch(batch_size: int, min_len: int, max_len: int) -> List[str]: - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=random.randint(min_len, max_len))) - for _ in range(batch_size) - ] - - -def _load_models( - model_key: str, - device: torch.device, - dtype: torch.dtype = torch.bfloat16, - registry: Dict[str, Dict] = None, -) -> Tuple[torch.nn.Module, torch.nn.Module, object]: - """Load official and fast models for a given family. - - Returns (official_wrapped, fast_model, tokenizer). - """ - if registry is None: - registry = MODEL_REGISTRY - config = registry[model_key] - - # Load official - module = importlib.import_module(config["load_official"]) - official_model, tokenizer = module.load_official_model( - reference_repo_id=config["official_path"], - device=device, - dtype=dtype, - ) - - # Load fast - fast_model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=dtype, - device_map=device, - ).eval() - - return official_model, fast_model, tokenizer - - -def _tokenize_batch( - model_key: str, - tokenizer: object, - batch: List[str], - device: torch.device, - registry: Dict[str, Dict] = None, -) -> Dict[str, torch.Tensor]: - """Tokenize a batch, handling E1's sequence mode.""" - if registry is None: - registry = MODEL_REGISTRY - config = registry[model_key] - if config["model_type"] == "E1": - tokenized = tokenizer.get_batch_kwargs(batch, device=device) - return { - "input_ids": tokenized["input_ids"], - "within_seq_position_ids": tokenized["within_seq_position_ids"], - "global_position_ids": tokenized["global_position_ids"], - "sequence_ids": tokenized["sequence_ids"], - "attention_mask": (tokenized["sequence_ids"] != -1).long(), - } - tokenized = tokenizer(batch, return_tensors="pt", padding=True) - return {k: v.to(device) for k, v in tokenized.items()} - - -def _masked_metrics( - candidate: torch.Tensor, - reference: torch.Tensor, - attention_mask: torch.Tensor, -) -> Dict[str, float]: - mask = attention_mask.bool() - cand = candidate[mask].float() - ref = reference[mask].float() - diff = cand - ref - diff_std = diff.std().item() - ref_std = ref.std().item() - diff_maxabs = diff.abs().max().item() - ref_maxabs = ref.abs().max().item() - rel_std = diff_std / ref_std if ref_std > 1e-12 else 0.0 - rel_maxabs = diff_maxabs / ref_maxabs if ref_maxabs > 1e-12 else 0.0 - return { - "mse": (diff ** 2).mean().item(), - "maxabs": diff_maxabs, - "rel_std": rel_std, - "rel_maxabs": rel_maxabs, - } - - -def _record_worst( - worst: Dict[str, Dict[str, float]], - label: str, - metrics: Dict[str, float], -) -> None: - if label not in worst: - worst[label] = metrics - return - for metric_name, value in metrics.items(): - if value > worst[label][metric_name]: - worst[label][metric_name] = value - - -def _render_worst(worst: Dict[str, Dict[str, float]]) -> str: - lines = [] - for label in sorted(worst): - metrics = worst[label] - lines.append( - f"{label}: mse={metrics['mse']:.3e}, maxabs={metrics['maxabs']:.3e}, " - f"rel_std={metrics['rel_std']:.3e}, rel_maxabs={metrics['rel_maxabs']:.3e}" - ) - return "\n".join(lines) - - -def _exceeds(value: float, tolerance: Optional[float]) -> bool: - return tolerance is not None and value > tolerance - - -def _format_tolerance(tolerance: Optional[float]) -> str: - if tolerance is None: - return "relative-only" - return f"{tolerance:.3e}" - - -def _select_final_hidden_state( - output: object, - hidden_states: Tuple[torch.Tensor, ...], -) -> Tuple[str, torch.Tensor]: - if isinstance(output, dict) and "last_hidden_state" in output: - return "last_hidden_state", output["last_hidden_state"] - try: - last_hidden_state = output.last_hidden_state - except AttributeError: - return "hidden_states[-1]", hidden_states[-1] - if last_hidden_state is not None: - return "last_hidden_state", last_hidden_state - return "hidden_states[-1]", hidden_states[-1] - - -def _run_weight_compliance(model_key: str, registry: Dict[str, Dict]) -> None: - """Core weight compliance logic shared by default and full-registry tests.""" - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - try: - official_model, fast_model, _ = _load_models(model_key, device, dtype=torch.float32, registry=registry) - except ModuleNotFoundError as e: - pytest.skip(f"Dependency not installed for {model_key}: {e}") - - assert_state_dict_equal( - reference_state_dict=official_model.model.state_dict(), - candidate_state_dict=fast_model.state_dict(), - context=f"{model_key} weight parity", - ) - - del official_model, fast_model - torch.cuda.empty_cache() - - -def _run_forward_compliance(model_key: str, registry: Dict[str, Dict]) -> None: - """Core forward compliance logic shared by default and full-registry tests.""" - random.seed(SEED) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - config = registry[model_key] - model_type = config["model_type"] - assert model_type in FORWARD_COMPLIANCE_TOLERANCES, ( - f"{model_key}: missing forward compliance tolerances for model_type={model_type}" - ) - tol = FORWARD_COMPLIANCE_TOLERANCES[model_type] - - try: - official_model, fast_model, tokenizer = _load_models( - model_key, device, dtype=FORWARD_DTYPE, registry=registry, - ) - except ModuleNotFoundError as e: - pytest.skip(f"Dependency not installed for {model_key}: {e}") - - failures: List[str] = [] - worst_metrics: Dict[str, Dict[str, float]] = {} - - with torch.inference_mode(), strict_fp32_matmul(): - for _ in range(TEST_NUM_BATCHES): - batch = _generate_random_batch(BATCH_SIZE, MIN_SEQ_LEN, MAX_SEQ_LEN) - tokenized = _tokenize_batch(model_key, tokenizer, batch, device, registry=registry) - attention_mask = tokenized["attention_mask"].bool() - - model_inputs = tokenized.copy() - model_inputs = add_model_specific_inputs(model_inputs, model_type) - - official_output = official_model(**model_inputs, output_hidden_states=True) - fast_output = fast_model(**model_inputs, output_hidden_states=True) - - official_logits = official_output.logits - fast_logits = fast_output.logits - assert official_logits is not None, f"{model_key}: official output has no logits" - assert fast_logits is not None, f"{model_key}: fast output has no logits" - logits_metrics = _masked_metrics( - fast_logits, - official_logits, - attention_mask, - ) - _record_worst(worst_metrics, "logits", logits_metrics) - if logits_metrics["mse"] > tol.logits_mse or logits_metrics["maxabs"] > tol.logits_maxabs: - failures.append( - f"logits: mse={logits_metrics['mse']:.3e} (tol={tol.logits_mse:.3e}), " - f"maxabs={logits_metrics['maxabs']:.3e} (tol={tol.logits_maxabs:.3e})" - ) - - official_hidden = official_output.hidden_states - fast_hidden = fast_output.hidden_states - assert len(official_hidden) == len(fast_hidden), ( - f"{model_key}: hidden_states tuple length mismatch " - f"fast={len(fast_hidden)} official={len(official_hidden)}" - ) - assert len(fast_hidden) > 0, f"{model_key}: no hidden states returned" - for i, (fast_h, official_h) in enumerate(zip(fast_hidden, official_hidden)): - hidden_metrics = _masked_metrics(fast_h, official_h, attention_mask) - label = f"hidden_states[{i}]" - _record_worst(worst_metrics, label, hidden_metrics) - if ( - _exceeds(hidden_metrics["mse"], tol.hidden_mse) - or _exceeds(hidden_metrics["maxabs"], tol.hidden_maxabs) - or hidden_metrics["rel_std"] > tol.hidden_rel_std - or hidden_metrics["rel_maxabs"] > tol.hidden_rel_maxabs - ): - failures.append( - f"{label}: mse={hidden_metrics['mse']:.3e} " - f"(tol={_format_tolerance(tol.hidden_mse)}), " - f"maxabs={hidden_metrics['maxabs']:.3e} " - f"(tol={_format_tolerance(tol.hidden_maxabs)}), " - f"rel_std={hidden_metrics['rel_std']:.3e} (tol={tol.hidden_rel_std:.3e}), " - f"rel_maxabs={hidden_metrics['rel_maxabs']:.3e} " - f"(tol={tol.hidden_rel_maxabs:.3e})" - ) - - official_last_label, official_last = _select_final_hidden_state( - official_output, - official_hidden, - ) - fast_last_label, fast_last = _select_final_hidden_state( - fast_output, - fast_hidden, - ) - final_label = official_last_label - if official_last_label != fast_last_label: - final_label = f"fast {fast_last_label} vs official {official_last_label}" - last_metrics = _masked_metrics(fast_last, official_last, attention_mask) - _record_worst(worst_metrics, final_label, last_metrics) - if ( - last_metrics["mse"] > tol.last_hidden_mse - or last_metrics["maxabs"] > tol.last_hidden_maxabs - or last_metrics["rel_maxabs"] > tol.last_hidden_rel_maxabs - ): - failures.append( - f"{final_label}: mse={last_metrics['mse']:.3e} " - f"(tol={tol.last_hidden_mse:.3e}), " - f"maxabs={last_metrics['maxabs']:.3e} " - f"(tol={tol.last_hidden_maxabs:.3e}), " - f"rel_maxabs={last_metrics['rel_maxabs']:.3e} " - f"(tol={tol.last_hidden_rel_maxabs:.3e})" - ) - - if failures: - rendered_failures = "\n".join(failures[:20]) - rendered_worst = _render_worst(worst_metrics) - pytest.fail( - f"{model_key} forward compliance failed under fp32 strict matmul:\n" - f"{rendered_failures}\n" - f"Worst observed metrics:\n{rendered_worst}" - ) - - del official_model, fast_model - torch.cuda.empty_cache() - - -# --------------------------------------------------------------------------- -# Default registry tests (small models, fast CI) -# --------------------------------------------------------------------------- - -# DPLM2 original has an extra contact_head not present in the FastPLM version, -# so positional state_dict comparison fails. Skip weight compliance for DPLM2. -WEIGHT_COMPLIANCE_KEYS = [k for k in MODEL_KEYS if k != "dplm2"] - -# DPLM2 original has structural differences (contact head, vocab mapping) that -# cause CUDA assertion failures when running through the ESM2 forward wrapper. -FORWARD_COMPLIANCE_KEYS = [k for k in MODEL_KEYS if k != "dplm2"] - - -@pytest.mark.slow -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", WEIGHT_COMPLIANCE_KEYS) -def test_weight_compliance(model_key: str) -> None: - """FastPLM weights are bit-exact with the original implementation.""" - _run_weight_compliance(model_key, MODEL_REGISTRY) - - -@pytest.mark.slow -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", FORWARD_COMPLIANCE_KEYS) -def test_forward_compliance(model_key: str) -> None: - """FastPLM forward pass outputs match the original within tolerance.""" - _run_forward_compliance(model_key, MODEL_REGISTRY) - - -# --------------------------------------------------------------------------- -# Full registry tests (all checkpoints across all families) -# --------------------------------------------------------------------------- - -FULL_WEIGHT_KEYS = [k for k in FULL_MODEL_REGISTRY if not k.startswith("dplm2")] -FULL_FORWARD_KEYS = [k for k in FULL_MODEL_REGISTRY if not k.startswith("dplm2")] - - -@pytest.mark.slow -@pytest.mark.gpu -@pytest.mark.parametrize( - "model_key", - mark_by_size(FULL_WEIGHT_KEYS, FULL_MODEL_REGISTRY, extra_marks=[pytest.mark.slow]), -) -def test_full_weight_compliance(model_key: str) -> None: - """Every checkpoint's weights are bit-exact with the original implementation.""" - _run_weight_compliance(model_key, FULL_MODEL_REGISTRY) - - -@pytest.mark.slow -@pytest.mark.gpu -@pytest.mark.parametrize( - "model_key", - mark_by_size(FULL_FORWARD_KEYS, FULL_MODEL_REGISTRY, extra_marks=[pytest.mark.slow]), -) -def test_full_forward_compliance(model_key: str) -> None: - """Every checkpoint's forward pass matches the original within tolerance.""" - _run_forward_compliance(model_key, FULL_MODEL_REGISTRY) diff --git a/testing/test_contact_maps.py b/testing/test_contact_maps.py deleted file mode 100644 index 007feb5..0000000 --- a/testing/test_contact_maps.py +++ /dev/null @@ -1,209 +0,0 @@ -import argparse -import os -import random -import requests -import tempfile -import sys -from typing import Tuple - -import matplotlib.pyplot as plt -import numpy as np -import torch -from Bio.PDB import PDBParser, PPBuilder - -from fastplms.esm2.modeling_fastesm import FastEsmModel - - -def download_random_pdb() -> str: - """ - Download a random protein chain PDB file. - - Returns: - str: Path to the downloaded PDB file. - """ - example_pdbs = ["1AKE"] - - # Select a random PDB ID - pdb_id = random.choice(example_pdbs) - print(f"Selected random PDB ID: {pdb_id}") - - # Create a temporary file to store the PDB - temp_file = tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) - temp_file_path = temp_file.name - temp_file.close() - - # Download the PDB file - url = f"https://files.rcsb.org/download/{pdb_id}.pdb" - response = requests.get(url) - - if response.status_code == 200: - with open(temp_file_path, 'wb') as f: - f.write(response.content) - print(f"Downloaded PDB file to: {temp_file_path}") - return temp_file_path - else: - raise Exception(f"Failed to download PDB file: {response.status_code}") - - -def parse_pdb(pdb_file: str) -> Tuple[str, np.ndarray]: - """ - Parse a PDB file and extract the protein sequence and CA atom coordinates. - - Parameters: - pdb_file (str): Path to the PDB file. - - Returns: - tuple: (sequence (str), coords (np.ndarray of shape (L, 3))) - """ - parser = PDBParser(QUIET=True) - structure = parser.get_structure("protein", pdb_file) - ppb = PPBuilder() - - # Assume a single protein chain; take the first polypeptide found. - for pp in ppb.build_peptides(structure): - sequence = str(pp.get_sequence()) - coords = [] - for residue in pp: - # Only add the CA atom if available. - if 'CA' in residue: - coords.append(residue['CA'].get_coord()) - if len(coords) == 0: - raise ValueError("No CA atoms found in the polypeptide.") - return sequence, np.array(coords) - - raise ValueError("No polypeptide chains were found in the PDB file.") - - -def compute_distance_matrix(coords: np.ndarray) -> np.ndarray: - """ - Compute the pairwise Euclidean distance matrix from a set of coordinates. - - Parameters: - coords (np.ndarray): Array of shape (L, 3) where L is the number of residues. - - Returns: - np.ndarray: A matrix of shape (L, L) containing distances. - """ - diff = coords[:, None, :] - coords[None, :, :] - dist_matrix = np.sqrt(np.sum(diff**2, axis=-1)) - - return dist_matrix - - -def get_esm_contact_map(sequence: str) -> np.ndarray: - """ - Use the ESM model to predict a contact map for the given protein sequence. - - Parameters: - sequence (str): Amino acid sequence. - - Returns: - np.ndarray: A 2D array (L x L) with contact probabilities. - """ - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model_path = "Synthyra/ESM2-650M" - model = FastEsmModel.from_pretrained(model_path).eval().to(device) - tokenizer = model.tokenizer - - inputs = tokenizer(sequence, return_tensors="pt") - inputs = {key: value.to(device) for key, value in inputs.items()} - with torch.no_grad(): - contact_map = model.predict_contacts(inputs["input_ids"], inputs["attention_mask"]) - print(contact_map.shape) - contact_map = contact_map.squeeze().cpu().numpy() - print(contact_map.shape) - return contact_map - - -def plot_maps(true_contact_map: np.ndarray, predicted_contact_map: np.ndarray, pdb_file: str) -> None: - """ - Generate two subplots: - 1. ESM predicted contact map. - 2. True contact map from the PDB (binary, thresholded). - - Parameters: - true_contact_map (np.ndarray): Binary (0/1) contact map from PDB. - predicted_contact_map (np.ndarray): Predicted contact probabilities from ESM. - pdb_file (str): Path to the PDB file, used to generate output filename. - """ - fig, axs = plt.subplots(1, 2, figsize=(12, 6)) - - # Plot the ESM-predicted contact map. - im0 = axs[0].imshow(predicted_contact_map, cmap='RdYlBu_r', aspect='equal') - axs[0].set_title("Predicted contact probabilities") - axs[0].set_xlabel("Residue index") - axs[0].set_ylabel("Residue index") - fig.colorbar(im0, ax=axs[0], fraction=0.046, pad=0.04) - - # Plot the true contact map (binary contacts). - im1 = axs[1].imshow(true_contact_map, cmap='RdYlBu_r', aspect='equal') - axs[1].set_title("True contacts (PDB, threshold = 8 Å)") - axs[1].set_xlabel("Residue index") - axs[1].set_ylabel("Residue index") - fig.colorbar(im1, ax=axs[1], fraction=0.046, pad=0.04) - - plt.tight_layout() - - # Generate output filename from PDB filename - pdb_name = os.path.splitext(os.path.basename(pdb_file))[0] - output_file = f"contact_maps_{pdb_name}.png" - plt.savefig(output_file, dpi=300, bbox_inches='tight') - plt.close() - - -def main() -> None: - # py tests/test_contact_maps.py - parser = argparse.ArgumentParser( - description="Extract protein sequence and compute contact maps from a PDB file using ESM predictions." - ) - parser.add_argument("--pdb_file", type=str, help="Path to the PDB file of the protein. If not provided, a random PDB will be downloaded.", default=None) - parser.add_argument( - "--threshold", - type=float, - default=8.0, - help="Distance threshold (in Å) for defining true contacts (default: 8.0 Å)." - ) - args = parser.parse_args() - - # If no PDB file is provided, download a random one - if args.pdb_file is None: - pdb_file = download_random_pdb() - else: - pdb_file = args.pdb_file - - try: - # Parse the PDB file. - sequence, coords = parse_pdb(pdb_file) - print("Extracted Protein Sequence:") - print(sequence) - - # Compute the pairwise distance matrix. - dist_matrix = compute_distance_matrix(coords) - - # Create a binary contact map from the distance matrix using the threshold. - true_contact_map = (dist_matrix < args.threshold).astype(float) - - # Get the predicted contact map from the ESM model. - predicted_contact_map = get_esm_contact_map(sequence) - - # Check that the dimensions agree. - if predicted_contact_map.shape[0] != true_contact_map.shape[0]: - print("Warning: The predicted contact map and true contact map have different dimensions.") - - # Plot the maps. - plot_maps(true_contact_map, predicted_contact_map, pdb_file) - - print(f"Contact maps saved to: contact_maps_{os.path.splitext(os.path.basename(pdb_file))[0]}.png") - - finally: - # Clean up the temporary file if we downloaded a random PDB - if args.pdb_file is None and os.path.exists(pdb_file): - os.remove(pdb_file) - print(f"Removed temporary PDB file: {pdb_file}") - - -if __name__ == '__main__': - main() - - - diff --git a/testing/test_e1_rag.py b/testing/test_e1_rag.py deleted file mode 100644 index 9a59c48..0000000 --- a/testing/test_e1_rag.py +++ /dev/null @@ -1,317 +0,0 @@ -from __future__ import annotations - -import io -import sys -import tarfile -from pathlib import Path - -import pytest -import torch - -from fastplms.e1.modeling_e1 import ( - ColabFoldSearcher, - ContextCache, - ContextSpecification, - E1Config, - E1ForMaskedLM, - HomologueSearcher, - _safe_extract_tar, - get_msa_for_sequence, - get_query_from_a3m, - load_msa_dir, - parse_msa, - sample_context, - sample_multiple_contexts, -) - - -def _write_tiny_a3m(path) -> None: - path.write_text( - ">query\n" - "ACDEFG\n" - ">near\n" - "ACDEYG\n" - ">far\n" - "TTTTTT\n", - encoding="utf-8", - ) - - -def _write_parity_a3m(path) -> None: - path.write_text( - ">query\n" - "ACDEFGHI\n" - ">near\n" - "ACDEYGH-\n" - ">gapped\n" - "AC-EFGHI\n" - ">mid\n" - "TCD-FGHI\n" - ">far\n" - "TTTTTTTT\n", - encoding="utf-8", - ) - - -def _load_official_msa_sampling(): - official_src = Path(__file__).resolve().parents[1] / "official" / "e1" / "src" - sys.path.insert(0, str(official_src)) - try: - return pytest.importorskip("E1.msa_sampling") - finally: - sys.path.remove(str(official_src)) - - -def _tiny_e1_model(device: torch.device) -> E1ForMaskedLM: - config = E1Config( - hidden_size=32, - intermediate_size=64, - num_hidden_layers=1, - num_attention_heads=4, - num_key_value_heads=4, - max_num_sequences=8, - max_num_positions_within_seq=64, - max_num_positions_global=256, - dtype="float32", - ) - return E1ForMaskedLM(config=config).eval().to(device) - - -def test_a3m_parsing_query_lookup_and_context_sampling(tmp_path) -> None: - a3m_path = tmp_path / "query.a3m" - _write_tiny_a3m(a3m_path) - - records = parse_msa(str(a3m_path)) - assert [record.id for record in records] == ["query", "near", "far"] - assert get_query_from_a3m(str(a3m_path)) == "ACDEFG" - - msa_lookup = load_msa_dir(str(tmp_path)) - assert msa_lookup["ACDEFG"] == str(a3m_path) - assert get_msa_for_sequence("ACDEYG", msa_lookup, min_identity=0.80) == str(a3m_path) - - context, ids = sample_context( - msa_path=str(a3m_path), - max_num_samples=1, - max_token_length=32, - max_query_similarity=0.1, - min_query_similarity=0.0, - seed=0, - device=torch.device("cpu"), - ) - assert context == "TTTTTT" - assert ids == ["far"] - - -def test_context_sampling_matches_official_e1(tmp_path) -> None: - official_msa_sampling = _load_official_msa_sampling() - a3m_path = tmp_path / "parity.a3m" - _write_parity_a3m(a3m_path) - - kwargs = { - "msa_path": str(a3m_path), - "max_num_samples": 3, - "max_token_length": 32, - "max_query_similarity": 0.99, - "min_query_similarity": 0.0, - "neighbor_similarity_lower_bound": 0.8, - "device": torch.device("cpu"), - } - for seed in (0, 3, 11): - context, ids = sample_context(seed=seed, **kwargs) - official_context, official_ids = official_msa_sampling.sample_context(seed=seed, **kwargs) - assert context == official_context - assert ids == official_ids - - specs = [ - ContextSpecification( - max_num_samples=3, - max_token_length=16, - max_query_similarity=0.99, - min_query_similarity=0.0, - neighbor_similarity_lower_bound=0.8, - ), - ContextSpecification( - max_num_samples=4, - max_token_length=32, - max_query_similarity=1.0, - min_query_similarity=0.2, - neighbor_similarity_lower_bound=0.8, - ), - ] - official_specs = [ - official_msa_sampling.ContextSpecification( - max_num_samples=spec.max_num_samples, - max_token_length=spec.max_token_length, - max_query_similarity=spec.max_query_similarity, - min_query_similarity=spec.min_query_similarity, - neighbor_similarity_lower_bound=spec.neighbor_similarity_lower_bound, - ) - for spec in specs - ] - - contexts, ids = sample_multiple_contexts( - msa_path=str(a3m_path), - context_specifications=specs, - seed=7, - device=torch.device("cpu"), - ) - official_contexts, official_ids = official_msa_sampling.sample_multiple_contexts( - msa_path=str(a3m_path), - context_specifications=official_specs, - seed=7, - device=torch.device("cpu"), - ) - assert contexts == official_contexts - assert ids == official_ids - - -def test_context_cache_round_trip(tmp_path) -> None: - cache = ContextCache(str(tmp_path), specs_hash="abc123", seed=7) - assert cache.load("msa") is None - cache.store("msa", {"ctx": "ACDEFG"}) - assert cache.load("msa") == {"ctx": "ACDEFG"} - - -def test_safe_tar_extraction_blocks_traversal(tmp_path) -> None: - tar_path = tmp_path / "bad.tar" - payload = b"bad" - with tarfile.open(tar_path, "w") as tar: - info = tarfile.TarInfo("../escape.a3m") - info.size = len(payload) - tar.addfile(info, io.BytesIO(payload)) - - with tarfile.open(tar_path) as tar: - with pytest.raises(ValueError): - _safe_extract_tar(tar, str(tmp_path / "out")) - - -def test_public_e1_rag_methods_exist() -> None: - model = _tiny_e1_model(torch.device("cpu")) - methods = [ - model.search_homologues, - model.batch_search_homologues, - model.sample_msa_contexts, - model.score_ppll, - model.embed_with_msa, - model.embed_dataset_with_msa, - ] - for method in methods: - assert callable(method) - assert callable(model.embed_dataset) - - -def test_mmseqs_searcher_subprocess_path_is_mockable(tmp_path, monkeypatch) -> None: - monkeypatch.chdir(tmp_path) - searcher = HomologueSearcher(target_db="target_db") - calls = [] - - def fake_run(cmd, **kwargs): - calls.append(cmd) - return None - - monkeypatch.setattr(searcher, "_ensure_docker_image", lambda: None) - monkeypatch.setattr(searcher, "_run_docker_command", fake_run) - - a3m_path = searcher.search("ACDEFG", output_dir="msas", seq_id="query") - - assert a3m_path == "msas/query/query.a3m" - assert any("createdb" in call for call in calls) - assert any("search" in call for call in calls) - assert any("result2msa" in call for call in calls) - - -def test_colabfold_searcher_http_path_is_mockable(tmp_path, monkeypatch) -> None: - searcher = ColabFoldSearcher(inter_request_delay=(0.0, 0.0)) - - def fake_download(ticket_id: str, output_path: str) -> None: - payload = b">query\nACDEFG\n" - with tarfile.open(output_path, "w:gz") as tar: - info = tarfile.TarInfo("uniref.a3m") - info.size = len(payload) - tar.addfile(info, io.BytesIO(payload)) - - monkeypatch.setattr(searcher, "_submit", lambda sequence: {"status": "RUNNING", "id": "ticket"}) - monkeypatch.setattr(searcher, "_poll", lambda ticket_id: {"status": "COMPLETE"}) - monkeypatch.setattr(searcher, "_download", fake_download) - - a3m_path = searcher.search("ACDEFG", str(tmp_path), seq_id="query") - - assert a3m_path.endswith("query.a3m") - assert get_query_from_a3m(a3m_path) == "ACDEFG" - - -@pytest.mark.gpu -def test_e1_score_ppll_with_tiny_synthetic_msa(tmp_path) -> None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = _tiny_e1_model(device) - a3m_path = tmp_path / "query.a3m" - _write_tiny_a3m(a3m_path) - - scores = model.score_ppll( - ["ACDEFG"], - a3m_path=str(a3m_path), - max_context_tokens=[64], - similarity_thresholds=[1.0], - min_query_similarity=0.0, - progress=False, - ) - - assert len(scores) == 1 - assert 0.0 <= scores[0] <= 1.0 - - per_context_scores = model.score_ppll( - ["ACDEFG"], - a3m_path=str(a3m_path), - ensemble=False, - max_context_tokens=[64, 128], - similarity_thresholds=[1.0], - min_query_similarity=0.0, - progress=False, - ) - assert len(per_context_scores) == 1 - assert len(per_context_scores[0]) == 2 - for score in per_context_scores[0]: - assert 0.0 <= score <= 1.0 - - -@pytest.mark.gpu -def test_e1_embed_with_msa_shapes(tmp_path) -> None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = _tiny_e1_model(device) - a3m_path = tmp_path / "query.a3m" - _write_tiny_a3m(a3m_path) - - pooled = model.embed_with_msa( - ["ACDEFG"], - a3m_path=str(a3m_path), - pooling_types=["mean"], - progress=False, - ) - matrix = model.embed_with_msa( - ["ACDEFG"], - a3m_path=str(a3m_path), - matrix_embed=True, - progress=False, - ) - - assert pooled.shape == (1, model.config.hidden_size) - assert len(matrix) == 1 - assert matrix[0].shape == (6, model.config.hidden_size) - - -@pytest.mark.gpu -def test_e1_embed_dataset_with_msa_falls_back_without_msa() -> None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = _tiny_e1_model(device) - - embeddings = model.embed_dataset_with_msa( - ["ACDEFG"], - msa_lookup={}, - batch_size=1, - max_len=16, - pooling_types=["mean"], - progress=False, - ) - - assert set(embeddings) == {"ACDEFG"} - assert embeddings["ACDEFG"].shape == (model.config.hidden_size,) diff --git a/testing/test_embedding_mixin.py b/testing/test_embedding_mixin.py deleted file mode 100644 index b6dcc3c..0000000 --- a/testing/test_embedding_mixin.py +++ /dev/null @@ -1,598 +0,0 @@ -"""Embedding mixin tests: NaN stability, batch-vs-single match, FASTA parsing, DPLM2 utilities.""" - -import os -import random -import sqlite3 -import tempfile -from typing import Dict, List - -import pytest -import torch - -from fastplms.embedding_mixin import ( - EmbeddingMixin, - load_pooled_embeddings_from_db, - load_pooled_embeddings_from_pth, - pool_embeddings, - tensor_to_embedding_blob, -) -from testing.conftest import ( - CANONICAL_AAS, FULL_MODEL_REGISTRY, MODEL_REGISTRY, SEED, - mark_by_size, strict_fp32_matmul, -) - - -BATCH_SIZE = 4 -MAX_EMBED_LEN = 128 -EMBED_MATCH_TOL = { - "default": {"maxabs": 6e-3, "rel_maxabs": None}, - # ESM3 exposes the pre-final-norm residual stream as embeddings. Absolute - # fp32 batch-shape noise can slightly exceed 1e-2 while relative error stays tiny. - "esm3": {"maxabs": 1.5e-2, "rel_maxabs": 5e-6}, -} - - -# Models that use tokenizer mode (not E1) -TOKENIZER_MODEL_KEYS = [k for k, v in MODEL_REGISTRY.items() if v["uses_tokenizer"]] -ALL_MODEL_KEYS = list(MODEL_REGISTRY.keys()) -ALL_FULL_MODEL_KEYS = list(FULL_MODEL_REGISTRY.keys()) -FULL_TOKENIZER_KEYS = [k for k, v in FULL_MODEL_REGISTRY.items() if v["uses_tokenizer"]] - - -class DummyEmbeddingConfig: - model_type = "E1" - hidden_size = 2 - - -class DummyHiddenStateModel(torch.nn.Module, EmbeddingMixin): - def __init__(self) -> None: - super().__init__() - self.config = DummyEmbeddingConfig() - self._parameter = torch.nn.Parameter(torch.zeros(1)) - - def _embed( - self, - sequences: List[str], - return_attention_mask: bool = False, - hidden_state_index: int = -1, - store_all_hidden_states: bool = False, - **kwargs, - ) -> torch.Tensor: - del kwargs - max_len = max(len(sequence) for sequence in sequences) - attention_mask = torch.zeros(len(sequences), max_len, dtype=torch.long) - hidden_states = [] - for layer_idx in range(3): - layer = torch.zeros(len(sequences), max_len, 2) - for batch_idx, sequence in enumerate(sequences): - seq_len = len(sequence) - attention_mask[batch_idx, :seq_len] = 1 - positions = torch.arange(seq_len, dtype=torch.float32) - layer[batch_idx, :seq_len, 0] = layer_idx - layer[batch_idx, :seq_len, 1] = positions - hidden_states.append(layer) - if store_all_hidden_states: - embeddings = torch.stack(hidden_states, dim=1) - elif hidden_state_index == -1: - embeddings = hidden_states[-1] - else: - embeddings = hidden_states[hidden_state_index] - if return_attention_mask: - return embeddings, attention_mask - return embeddings - - -class FixedLengthTokenizer: - """Wraps a tokenizer so every call pads to exactly MAX_EMBED_LEN tokens. - - Both batch=1 and batch=N therefore receive tensors of the same shape, - keeping max_seqlen_in_batch identical and eliminating floating-point - variability from different softmax vector lengths / flash-attention tile sizes. - """ - def __init__(self, tokenizer: object, max_length: int = MAX_EMBED_LEN) -> None: - self._tok = tokenizer - self.max_length = max_length - - def __call__(self, sequences: List[str], **kwargs) -> Dict[str, torch.Tensor]: - return self._tok( - sequences, - return_tensors="pt", - padding="max_length", - max_length=self.max_length, - truncation=True, - ) - - -def _random_sequences(n: int, min_len: int = 8, max_len: int = 64) -> List[str]: - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=random.randint(min_len, max_len))) - for _ in range(n) - ] - - -def _random_sequences_fixed_len(n: int, length: int = 64) -> List[str]: - return [ - "M" + "".join(random.choices(CANONICAL_AAS, k=length - 1)) - for _ in range(n) - ] - - -def _assert_no_nan(embeddings: Dict[str, torch.Tensor], label: str) -> None: - for seq, emb in embeddings.items(): - assert not torch.isnan(emb).any(), ( - f"[{label}] NaN found in embedding for sequence '{seq[:20]}...'" - ) - - -def _assert_embeddings_match( - a: Dict[str, torch.Tensor], - b: Dict[str, torch.Tensor], - label: str, -) -> None: - assert set(a) == set(b), f"[{label}] Key sets differ between batch and single runs" - if label.startswith("esm3"): - tol = EMBED_MATCH_TOL["esm3"] - else: - tol = EMBED_MATCH_TOL["default"] - for seq in a: - ea, eb = a[seq].float(), b[seq].float() - assert ea.shape == eb.shape, ( - f"[{label}] Shape mismatch for '{seq[:20]}': {ea.shape} vs {eb.shape}" - ) - max_diff = (ea - eb).abs().max().item() - base_maxabs = max(ea.abs().max().item(), eb.abs().max().item()) - rel_maxabs = max_diff / base_maxabs if base_maxabs > 1e-12 else 0.0 - rel_tol = tol["rel_maxabs"] - rel_ok = rel_tol is None or rel_maxabs <= rel_tol - assert max_diff <= tol["maxabs"] and rel_ok, ( - f"[{label}] Max abs diff {max_diff:.5f} > {tol['maxabs']} " - f"or rel maxabs {rel_maxabs:.3e} > {rel_tol} for '{seq[:20]}'" - ) - - -@pytest.fixture -def disable_tf32_for_batch_single_match(): - # TF32 kernels can be batch-shape-dependent on Hopper/GH200, which defeats - # this test's batch-vs-single equality check. - with strict_fp32_matmul(): - yield - - -# --- CPU-only utility tests --- - -def test_parse_fasta() -> None: - from fastplms.embedding_mixin import parse_fasta - - fasta_content = ( - ">seq1 a simple protein\n" - "MKTLLLTLVVVTIVCLDLGYT\n" - ">seq2 multi-line sequence\n" - "ACDEFGHIKL\n" - "MNPQRSTVWY\n" - ">seq3 another entry\n" - "MALWMRLLPLLALL\n" - ) - expected = [ - "MKTLLLTLVVVTIVCLDLGYT", - "ACDEFGHIKLMNPQRSTVWY", - "MALWMRLLPLLALL", - ] - with tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False) as f: - f.write(fasta_content) - tmp_path = f.name - parsed = parse_fasta(tmp_path) - os.unlink(tmp_path) - assert parsed == expected - - -def test_pool_embeddings_selects_layer_from_all_hidden_states() -> None: - all_layers = torch.tensor( - [ - [[0.0, 0.0], [0.0, 2.0]], - [[1.0, 0.0], [1.0, 2.0]], - [[2.0, 0.0], [2.0, 2.0]], - ] - ) - pooled = pool_embeddings( - {"AA": all_layers}, - pooling_types=["mean"], - hidden_state_index=1, - ) - assert torch.equal(pooled["AA"], torch.tensor([1.0, 1.0])) - - -def test_load_pooled_embeddings_from_pth_and_db() -> None: - all_layers = torch.tensor( - [ - [[0.0, 0.0], [0.0, 2.0]], - [[1.0, 0.0], [1.0, 2.0]], - [[2.0, 0.0], [2.0, 2.0]], - ] - ) - with tempfile.TemporaryDirectory() as tmp_dir: - pth_path = os.path.join(tmp_dir, "embeddings.pth") - db_path = os.path.join(tmp_dir, "embeddings.db") - torch.save({"AA": all_layers}, pth_path) - - with sqlite3.connect(db_path) as conn: - cursor = conn.cursor() - cursor.execute( - "CREATE TABLE embeddings (sequence TEXT PRIMARY KEY, embedding BLOB NOT NULL)" - ) - cursor.execute( - "INSERT INTO embeddings VALUES (?, ?)", - ("AA", tensor_to_embedding_blob(all_layers)), - ) - conn.commit() - - pth_pooled = load_pooled_embeddings_from_pth( - pth_path, - pooling_types=["mean"], - hidden_state_index=1, - ) - db_pooled = load_pooled_embeddings_from_db( - db_path, - pooling_types=["mean"], - hidden_state_index=1, - ) - - assert torch.equal(pth_pooled["AA"], torch.tensor([1.0, 1.0])) - assert torch.equal(db_pooled["AA"], torch.tensor([1.0, 1.0])) - - -def test_embed_dataset_hidden_state_index_sequence_mode() -> None: - model = DummyHiddenStateModel() - sequences = ["ACD", "M"] - with tempfile.TemporaryDirectory() as tmp_dir: - save_path = os.path.join(tmp_dir, "embeddings.pth") - default_embeddings = model.embed_dataset( - sequences=sequences, - tokenizer=None, - batch_size=2, - pooling_types=["mean"], - save=False, - save_path=save_path, - padding="longest", - ) - selected_embeddings = model.embed_dataset( - sequences=sequences, - tokenizer=None, - batch_size=2, - pooling_types=["mean"], - hidden_state_index=1, - save=False, - save_path=save_path, - padding="longest", - ) - - assert torch.equal(default_embeddings["ACD"], torch.tensor([2.0, 1.0])) - assert torch.equal(selected_embeddings["ACD"], torch.tensor([1.0, 1.0])) - assert torch.equal(selected_embeddings["M"], torch.tensor([1.0, 0.0])) - - -def test_embed_dataset_store_all_hidden_states_sequence_mode() -> None: - model = DummyHiddenStateModel() - with tempfile.TemporaryDirectory() as tmp_dir: - save_path = os.path.join(tmp_dir, "embeddings.pth") - embeddings = model.embed_dataset( - sequences=["ACD", "M"], - tokenizer=None, - batch_size=2, - full_embeddings=True, - store_all_hidden_states=True, - save=False, - save_path=save_path, - padding="longest", - ) - - assert embeddings["ACD"].shape == (3, 3, 2) - assert embeddings["M"].shape == (3, 1, 2) - assert torch.equal(embeddings["ACD"][1].mean(dim=0), torch.tensor([1.0, 1.0])) - - -def test_embed_dataset_store_all_hidden_states_sql_roundtrip() -> None: - model = DummyHiddenStateModel() - with tempfile.TemporaryDirectory() as tmp_dir: - db_path = os.path.join(tmp_dir, "embeddings.db") - model.embed_dataset( - sequences=["ACD", "M"], - tokenizer=None, - batch_size=2, - full_embeddings=True, - store_all_hidden_states=True, - sql=True, - sql_db_path=db_path, - save=False, - padding="longest", - ) - pooled = model.load_pooled_embeddings_from_db( - db_path, - pooling_types=["mean"], - hidden_state_index=1, - ) - - assert torch.equal(pooled["ACD"], torch.tensor([1.0, 1.0])) - assert torch.equal(pooled["M"], torch.tensor([1.0, 0.0])) - - -def test_embed_dataset_store_all_hidden_states_requires_full_embeddings() -> None: - model = DummyHiddenStateModel() - with pytest.raises(AssertionError, match="requires full_embeddings"): - model.embed_dataset( - sequences=["ACD"], - tokenizer=None, - store_all_hidden_states=True, - save=False, - padding="longest", - ) - - -@pytest.mark.gpu -def test_dplm2_multimodal_layout_guard() -> None: - from fastplms.dplm2.modeling_dplm2 import _has_packed_multimodal_layout - - plain = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 2], [1, 1, 1, 1, 1, 0, 2, 2]]) - packed = torch.tensor([[1, 1, 1, 2, 0, 0, 0, 2], [1, 1, 2, 2, 0, 0, 2, 2]]) - mismatched = torch.tensor([[1, 1, 1, 2, 0, 0, 2, 2]]) - - assert not _has_packed_multimodal_layout(plain, aa_type=1, struct_type=0, pad_type=2) - assert _has_packed_multimodal_layout(packed, aa_type=1, struct_type=0, pad_type=2) - assert not _has_packed_multimodal_layout(mismatched, aa_type=1, struct_type=0, pad_type=2) - - -@pytest.mark.gpu -def test_dplm2_special_token_normalization() -> None: - from fastplms.dplm2.modeling_dplm2 import _normalize_dplm2_input_ids - - input_ids = torch.tensor([[8231, 5, 23, 13, 8229, 1, 8232, -100]]) - normalized = _normalize_dplm2_input_ids(input_ids, vocab_size=8229) - expected = torch.tensor([[0, 5, 23, 13, 2, 1, 32, -100]]) - assert torch.equal(normalized, expected), ( - f"DPLM2 normalization mismatch: got {normalized.tolist()}, expected {expected.tolist()}" - ) - - -# --- GPU model tests --- - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", ALL_MODEL_KEYS) -def test_nan_stability(model_key: str) -> None: - """Batched embed_dataset produces no NaN in real-token rows.""" - from transformers import AutoModelForMaskedLM - - random.seed(SEED) - config = MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - uses_tokenizer = config["uses_tokenizer"] - if uses_tokenizer: - tokenizer = FixedLengthTokenizer(model.tokenizer) - sequences = _random_sequences(n=8) - else: - tokenizer = None - sequences = _random_sequences_fixed_len(n=8) - - embs = model.embed_dataset( - sequences=sequences, - batch_size=BATCH_SIZE, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.bfloat16, - save=False, - ) - _assert_no_nan(embs, f"{model_key} NaN check batch_size={BATCH_SIZE}") - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", TOKENIZER_MODEL_KEYS) -def test_batch_single_match(model_key: str, disable_tf32_for_batch_single_match) -> None: - """Batched and single-item embedding produce matching results (tokenizer models only). - - E1 is excluded: flash varlen is not bit-deterministic across different batch sizes. - For SDPA models we cast to float32 to avoid bfloat16 CUBLAS algorithm selection differences. - """ - from transformers import AutoModelForMaskedLM - - random.seed(SEED) - config = MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.float32, - device_map=device, - ).eval() - - tokenizer = FixedLengthTokenizer(model.tokenizer) - sequences = _random_sequences(n=8) - - with strict_fp32_matmul(): - batch_embs = model.embed_dataset( - sequences=sequences, - batch_size=BATCH_SIZE, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.float32, - save=False, - ) - single_embs = model.embed_dataset( - sequences=sequences, - batch_size=1, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.float32, - save=False, - ) - _assert_no_nan(batch_embs, f"{model_key} match test batch_size={BATCH_SIZE}") - _assert_no_nan(single_embs, f"{model_key} match test batch_size=1") - _assert_embeddings_match(batch_embs, single_embs, model_key) - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -def test_tokenizer_model_embed_dataset_uses_default_tokenizer() -> None: - from transformers import AutoModelForMaskedLM - - config = MODEL_REGISTRY["esmc"] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - embeddings = model.embed_dataset( - sequences=["MKTAYIAKQ", "GGGG"], - batch_size=2, - max_len=16, - pooling_types=["mean"], - save=False, - ) - - assert set(embeddings) == {"MKTAYIAKQ", "GGGG"} - assert embeddings["MKTAYIAKQ"].shape == (model.config.hidden_size,) - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", ALL_MODEL_KEYS) -def test_hidden_state_index_embed_dataset_smoke(model_key: str) -> None: - from transformers import AutoModelForMaskedLM - - config = MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - if config["uses_tokenizer"]: - tokenizer = FixedLengthTokenizer(model.tokenizer, max_length=32) - sequences = ["MKTAYIAKQ", "GGGG"] - else: - tokenizer = None - sequences = ["MKTAYIAKQ", "GGGG"] - - embeddings = model.embed_dataset( - sequences=sequences, - tokenizer=tokenizer, - batch_size=2, - max_len=32, - pooling_types=["mean"], - hidden_state_index=0, - save=False, - ) - - assert set(embeddings) == set(sequences) - for embedding in embeddings.values(): - assert embedding.shape == (model.config.hidden_size,) - assert not torch.isnan(embedding).any() - - del model - torch.cuda.empty_cache() - - -# --------------------------------------------------------------------------- -# Full model registry tests: NaN stability across all checkpoints -# --------------------------------------------------------------------------- - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", mark_by_size(ALL_FULL_MODEL_KEYS, FULL_MODEL_REGISTRY)) -def test_full_nan_stability(model_key: str) -> None: - """Every checkpoint's embed_dataset produces no NaN in real-token rows.""" - from transformers import AutoModelForMaskedLM - - random.seed(SEED) - config = FULL_MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.bfloat16, - device_map=device, - ).eval() - - uses_tokenizer = config["uses_tokenizer"] - if uses_tokenizer: - tokenizer = FixedLengthTokenizer(model.tokenizer) - sequences = _random_sequences(n=8) - else: - tokenizer = None - sequences = _random_sequences_fixed_len(n=8) - - embs = model.embed_dataset( - sequences=sequences, - batch_size=BATCH_SIZE, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.bfloat16, - save=False, - ) - _assert_no_nan(embs, f"{model_key} NaN check batch_size={BATCH_SIZE}") - - del model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", mark_by_size(FULL_TOKENIZER_KEYS, FULL_MODEL_REGISTRY)) -def test_full_batch_single_match(model_key: str, disable_tf32_for_batch_single_match) -> None: - """Every tokenizer-mode checkpoint matches batch vs single-item embedding.""" - from transformers import AutoModelForMaskedLM - - random.seed(SEED) - config = FULL_MODEL_REGISTRY[model_key] - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], - trust_remote_code=True, - dtype=torch.float32, - device_map=device, - ).eval() - - tokenizer = FixedLengthTokenizer(model.tokenizer) - sequences = _random_sequences(n=8) - - with strict_fp32_matmul(): - batch_embs = model.embed_dataset( - sequences=sequences, - batch_size=BATCH_SIZE, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.float32, - save=False, - ) - single_embs = model.embed_dataset( - sequences=sequences, - batch_size=1, - tokenizer=tokenizer, - full_embeddings=True, - embed_dtype=torch.float32, - save=False, - ) - _assert_no_nan(batch_embs, f"{model_key} match test batch_size={BATCH_SIZE}") - _assert_no_nan(single_embs, f"{model_key} match test batch_size=1") - _assert_embeddings_match(batch_embs, single_embs, model_key) - - del model - torch.cuda.empty_cache() diff --git a/testing/test_esm3.py b/testing/test_esm3.py deleted file mode 100644 index 146dd3a..0000000 --- a/testing/test_esm3.py +++ /dev/null @@ -1,121 +0,0 @@ -import shutil -from pathlib import Path - -import pytest -import torch -from transformers import AutoModel - -from fastplms.esm3 import modeling_esm3 -from fastplms.esm3.modeling_esm3 import FastESM3Config, FastESM3Model -from testing.conftest import strict_fp32_matmul - - -HUB_AUTO_MAP = { - "AutoConfig": "modeling_esm3.FastESM3Config", - "AutoModel": "modeling_esm3.FastESM3Model", - "AutoModelForMaskedLM": "modeling_esm3.FastESM3Model", -} - - -def _small_config() -> FastESM3Config: - config = FastESM3Config( - hidden_size=64, - num_attention_heads=4, - num_vector_heads=8, - num_hidden_layers=2, - ) - config.architectures = ["FastESM3Model"] - config.auto_map = HUB_AUTO_MAP - return config - - -def _small_model() -> FastESM3Model: - try: - model = FastESM3Model(_small_config()).eval() - except ModuleNotFoundError as exc: - pytest.skip(f"Biohub ESM3 runtime dependency is unavailable: {exc}") - return model - - -def test_esm3_sequence_only_forward() -> None: - model = _small_model() - batch = model.tokenize_sequences(["MKTAYIAKQ", "GGGG"], device=model.device) - - with torch.inference_mode(): - output = model(**batch) - - assert output.logits is not None - assert output.function_logits is not None - assert output.residue_logits is not None - assert output.logits.shape[:2] == batch["input_ids"].shape - assert output.logits.shape[-1] == 64 - assert output.structure_logits.shape[-1] == 4096 - assert output.function_logits.shape[-2:] == (8, 260) - assert output.residue_logits.shape[-1] == 1478 - assert not torch.isnan(output.logits).any() - - -def test_esm3_accepts_function_tokens_argument() -> None: - model = _small_model() - batch = model.tokenize_sequences(["MKTAYIAKQ"], device=model.device) - function_tokens = batch["input_ids"].new_zeros((*batch["input_ids"].shape, 8)) - - with torch.inference_mode(): - output = model(**batch, function_tokens=function_tokens) - - assert output.logits is not None - assert output.logits.shape[:2] == batch["input_ids"].shape - - -def test_esm3_loads_with_automodel(tmp_path: Path) -> None: - model = _small_model() - model.save_pretrained(tmp_path) - shutil.copyfile(Path(modeling_esm3.__file__), tmp_path / "modeling_esm3.py") - - loaded = AutoModel.from_pretrained(tmp_path, trust_remote_code=True).eval() - batch = loaded.tokenize_sequences(["MKTAYIAKQ"], device=loaded.device) - - with torch.inference_mode(): - output = loaded(**batch) - - assert output.logits is not None - assert output.logits.shape[:2] == batch["input_ids"].shape - - -def test_esm3_embed_dataset(tmp_path: Path) -> None: - model = _small_model() - save_path = tmp_path / "embeddings.pth" - - embeddings = model.embed_dataset( - sequences=["MKTAYIAKQ", "GGGG"], - batch_size=2, - max_len=16, - pooling_types=["mean", "cls"], - save=True, - save_path=str(save_path), - ) - - assert set(embeddings) == {"MKTAYIAKQ", "GGGG"} - assert embeddings["MKTAYIAKQ"].shape == (128,) - assert save_path.exists() - - -def test_esm3_flex_matches_sdpa() -> None: - if not torch.cuda.is_available(): - pytest.skip("Flex attention ESM3 equivalence is validated on CUDA.") - model = _small_model().to(torch.device("cuda")) - batch = model.tokenize_sequences(["MKTAYIAKQ", "GGGG"], device=model.device) - - with torch.inference_mode(), strict_fp32_matmul(): - model.attn_backend = "sdpa" - sdpa_output = model(**batch).last_hidden_state - try: - model.attn_backend = "flex" - except AssertionError as exc: - pytest.skip(f"Flex attention is unavailable: {exc}") - flex_output = model(**batch).last_hidden_state - - max_abs = (sdpa_output - flex_output).float().abs().max().item() - mse = ((sdpa_output - flex_output).float() ** 2).mean().item() - assert max_abs < 1e-4 - assert mse < 1e-8 diff --git a/testing/test_esmfold2_experimental.py b/testing/test_esmfold2_experimental.py deleted file mode 100644 index fc62d81..0000000 --- a/testing/test_esmfold2_experimental.py +++ /dev/null @@ -1,231 +0,0 @@ -"""ESMFold2 experimental model tests.""" - -from __future__ import annotations - -import importlib -import os -import subprocess -import sys -import tempfile -from pathlib import Path - -import pytest -import torch -import torch.nn.functional as F -from transformers import AutoModel - -from fastplms.esmfold2.configuration_esmfold2 import ESMFold2Config -from fastplms.esmfold2.get_weights import EXPERIMENTAL_AUTO_MAP -from fastplms.esmfold2.modeling_esmfold2_common import NUM_RES_TYPES -from fastplms.esmfold2.modeling_esmfold2_experimental import ( - ESMFold2ExperimentalModel, -) -from fastplms.esmfold2.protein_utils import prepare_protein_features - - -TEST_SEQUENCE = "MSTNPKPQRKTKRNT" -OFFICIAL_REPO = "biohub/ESMFold2-Experimental-Fast" -FAST_REPO = "Synthyra/ESMFold2-Experimental-Fast" -OUTPUT_TOLERANCES = { - "distogram_logits": 0.0, - "plddt": 1e-6, - "pae": 0.0, - "ptm": 0.0, - "iptm": 0.0, -} - - -def _enable_deterministic_forward() -> None: - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = True - torch.use_deterministic_algorithms(True) - - -def _load_official_model() -> torch.nn.Module: - module = pytest.importorskip( - "transformers.models.esmfold2.modeling_esmfold2_experimental" - ) - official_cls = module.ESMFold2ExperimentalModel - return ( - official_cls.from_pretrained( - OFFICIAL_REPO, - load_esmc=False, - dtype=torch.float32, - ) - .eval() - .cuda() - ) - - -def _load_fast_model() -> ESMFold2ExperimentalModel: - return ( - ESMFold2ExperimentalModel.from_pretrained( - OFFICIAL_REPO, - load_esmc=False, - dtype=torch.float32, - ) - .eval() - .cuda() - ) - - -def _run_short_fold(model: torch.nn.Module) -> dict[str, torch.Tensor]: - common_module_name = ( - model.__class__.__module__.rsplit(".", 1)[0] + ".modeling_esmfold2_common" - ) - common_module = importlib.import_module(common_module_name) - with common_module._seed_context(0), torch.no_grad(): - return model.infer_protein( - TEST_SEQUENCE, - num_loops=1, - num_sampling_steps=2, - num_diffusion_samples=1, - calculate_confidence=True, - seed=0, - ) - - -def _assert_forward_parity() -> None: - _enable_deterministic_forward() - official_model = _load_official_model() - fast_model = _load_fast_model() - - official_output = _run_short_fold(official_model) - fast_output = _run_short_fold(fast_model) - - for key, atol in OUTPUT_TOLERANCES.items(): - torch.testing.assert_close( - fast_output[key], - official_output[key], - rtol=0.0, - atol=atol, - msg=f"ESMFold2 experimental output mismatch: {key}", - ) - - del official_model, fast_model, official_output, fast_output - torch.cuda.empty_cache() - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -def test_esmfold2_experimental_weight_parity() -> None: - official_model = _load_official_model() - fast_model = _load_fast_model() - - official_state = official_model.state_dict() - fast_state = fast_model.state_dict() - assert official_state.keys() == fast_state.keys() - - for name, official_tensor in official_state.items(): - torch.testing.assert_close( - fast_state[name], - official_tensor, - rtol=0.0, - atol=0.0, - msg=f"ESMFold2 experimental parameter mismatch: {name}", - ) - - del official_model, fast_model - torch.cuda.empty_cache() - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -def test_esmfold2_experimental_forward_parity() -> None: - env = os.environ.copy() - with tempfile.TemporaryDirectory() as module_cache: - env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - env["HF_MODULES_CACHE"] = module_cache - result = subprocess.run( - [ - sys.executable, - __file__, - "--esmfold2-experimental-forward-parity", - ], - capture_output=True, - text=True, - check=False, - env=env, - ) - if result.returncode != 0 and "Skipped:" in result.stderr: - pytest.skip(result.stderr.split("Skipped:", 1)[1].strip()) - assert result.returncode == 0, result.stdout + result.stderr - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -def test_esmfold2_experimental_res_type_soft_gradients() -> None: - model = _load_fast_model() - features = { - name: tensor.cuda() for name, tensor in prepare_protein_features(TEST_SEQUENCE).items() - } - res_type_soft = F.one_hot( - features["res_type"].long(), num_classes=NUM_RES_TYPES - ).float() - res_type_soft.requires_grad_(True) - - output = model( - **features, - res_type_soft=res_type_soft, - num_loops=0, - num_sampling_steps=1, - num_diffusion_samples=1, - calculate_confidence=False, - seed=0, - ) - loss = output["distogram_logits"].float().mean() - loss.backward() - - assert "representative_atom_coords" in output - assert output["representative_atom_coords"].shape[-1] == 3 - assert output["representative_atom_coords"].shape[-2] == features["res_type"].shape[1] - assert res_type_soft.grad is not None - assert torch.isfinite(res_type_soft.grad).all() - assert res_type_soft.grad.abs().sum().item() > 0 - - del model, output, features - torch.cuda.empty_cache() - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -def test_esmfold2_experimental_automodel_loads() -> None: - try: - model = AutoModel.from_pretrained( - FAST_REPO, - trust_remote_code=True, - load_esmc=False, - dtype=torch.float32, - ) - except OSError as exc: - pytest.skip(f"{FAST_REPO} is not available yet: {exc}") - model = model.eval().cuda() - - assert callable(model.infer_protein_as_pdb) - assert callable(model.fold) - assert callable(model.prepare_structure_input) - - del model - torch.cuda.empty_cache() - - -def test_esmfold2_experimental_export_config(tmp_path: Path) -> None: - config = ESMFold2Config(type="experimental") - config.auto_map = EXPERIMENTAL_AUTO_MAP - config.architectures = ["ESMFold2ExperimentalModel"] - config.save_pretrained(tmp_path) - - loaded = ESMFold2Config.from_pretrained(tmp_path) - assert loaded.auto_map == EXPERIMENTAL_AUTO_MAP - assert loaded.architectures == ["ESMFold2ExperimentalModel"] - - -if __name__ == "__main__": - assert len(sys.argv) == 2 - assert sys.argv[1] == "--esmfold2-experimental-forward-parity" - _assert_forward_parity() diff --git a/testing/test_parity.py b/testing/test_parity.py deleted file mode 100644 index 09862c0..0000000 --- a/testing/test_parity.py +++ /dev/null @@ -1,1053 +0,0 @@ -"""Rigorous parity tests between FastPLMs and native implementations. - -The suite favors small, specific asserts so a failure identifies the tensor, -metric, and tolerance that diverged. fp32 checks stay tight; bf16 checks are -documented per family because depth and architecture change accumulated -rounding. Intermediate hidden states use relative metrics so large pre-norm -residual streams cannot make absolute MSE look worse than it is. - -Each family runs in the same file, but a family-specific Docker image can skip -families whose native reference dependencies are not installed. - -Run (per family image; see Dockerfile.): - docker run --gpus all --ipc=host --rm -v $(pwd):/workspace \ - fastplms-esm_plusplus python -m pytest /workspace/testing/test_parity.py -k esmc -v -s -""" -from __future__ import annotations - -import contextlib -import importlib -import random -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import pytest -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import AutoModelForMaskedLM - -from testing.conftest import ( - CANONICAL_AAS, MODEL_REGISTRY, SEED, strict_fp32_matmul, tokenize_batch, -) - - -@dataclass -class ParityTolerances: - """Per-family / per-dtype numerical tolerances. - - Tolerances are intentionally tight in fp32 and documented per-family in bf16. - There are three complementary metrics so that no single collapsed scalar can - hide a localized regression: - - - `*_last_hidden_mse` / `*_last_hidden_maxabs`: absolute errors at the final - (post-final-norm) residue-level representation. - - `*_last_hidden_rel_maxabs`: the absolute maxabs divided by native's maxabs. - Catches cases where native activations are also large, and cases where a - constant bias is added while the absolute maxabs still looks small. - - `*_hidden_rel_std`: per-layer std-of-diff / std-of-native. Captures overall - distribution agreement at each intermediate layer. - - `*_hidden_rel_maxabs`: per-layer (maxabs of diff) / (maxabs of native). - Captures sharp per-head or per-position regressions that std-of-diff - would average out. - """ - fp32_last_hidden_mse: float = 1e-8 - fp32_last_hidden_maxabs: float = 5e-4 - fp32_last_hidden_rel_maxabs: float = 5e-4 - fp32_logits_mse: float = 1e-4 - fp32_hidden_rel_std: float = 5e-3 - fp32_hidden_rel_maxabs: float = 5e-2 - bf16_last_hidden_mse: float = 1e-5 - bf16_last_hidden_maxabs: float = 2e-2 - bf16_last_hidden_rel_maxabs: float = 5e-2 - bf16_logits_mse: float = 5e-2 - bf16_hidden_rel_std: float = 5e-2 - bf16_hidden_rel_maxabs: float = 1e-1 - - -FAMILY_TOLERANCES: Dict[str, ParityTolerances] = { - "esm2": ParityTolerances( - fp32_last_hidden_mse=1e-12, fp32_last_hidden_maxabs=1e-5, fp32_last_hidden_rel_maxabs=1e-5, - fp32_logits_mse=1e-12, fp32_hidden_rel_std=1e-6, fp32_hidden_rel_maxabs=1e-5, - bf16_last_hidden_mse=1e-5, bf16_last_hidden_maxabs=2e-2, bf16_last_hidden_rel_maxabs=3e-2, - bf16_logits_mse=5e-2, bf16_hidden_rel_std=1e-2, bf16_hidden_rel_maxabs=5e-2, - ), - "esmc": ParityTolerances( - fp32_last_hidden_mse=1e-8, fp32_last_hidden_maxabs=1e-3, fp32_last_hidden_rel_maxabs=1e-3, - fp32_logits_mse=1e-4, fp32_hidden_rel_std=5e-3, fp32_hidden_rel_maxabs=5e-2, - bf16_last_hidden_mse=1e-5, bf16_last_hidden_maxabs=5e-2, bf16_last_hidden_rel_maxabs=5e-2, - bf16_logits_mse=5e-2, bf16_hidden_rel_std=5e-2, bf16_hidden_rel_maxabs=1e-1, - ), - "esm3": ParityTolerances( - fp32_last_hidden_mse=1e-8, fp32_last_hidden_maxabs=1e-3, fp32_last_hidden_rel_maxabs=1e-3, - fp32_logits_mse=1e-4, fp32_hidden_rel_std=5e-3, fp32_hidden_rel_maxabs=5e-2, - # ESM3 exposes the pre-final-norm embedding stream as `last_hidden_state`; - # bf16 absolute error scales with that large residual magnitude, so the - # relative checks are the meaningful guardrails here. - bf16_last_hidden_mse=2e1, bf16_last_hidden_maxabs=2.6e2, bf16_last_hidden_rel_maxabs=5e-2, - bf16_logits_mse=5e-2, bf16_hidden_rel_std=7e-2, bf16_hidden_rel_maxabs=1.5e-1, - ), - "e1": ParityTolerances( - fp32_last_hidden_mse=5e-7, fp32_last_hidden_maxabs=2e-2, fp32_last_hidden_rel_maxabs=2e-3, - fp32_hidden_rel_std=1e-2, fp32_hidden_rel_maxabs=2e-2, - # Grouped-query attention plus block-causal global layers can produce - # a few bf16-bucket absolute outliers; MSE and relative maxabs stay tight. - bf16_last_hidden_maxabs=1.5e-1, bf16_last_hidden_rel_maxabs=5e-2, - bf16_hidden_rel_std=1e-1, bf16_hidden_rel_maxabs=1e-1, - ), - "dplm": ParityTolerances(), - "dplm2": ParityTolerances( - # DPLM2 encoder is bit-identical to native's ESM backbone on pure AA - # input (same weights; ModifiedRotaryEmbedding falls through to vanilla - # rotary when no packed multimodal layout is detected). Logits differ - # by construction (head is separately learned); logits parity is - # skipped via _family_has_head_mismatch. - fp32_last_hidden_mse=1e-12, fp32_last_hidden_maxabs=1e-5, fp32_last_hidden_rel_maxabs=1e-5, - fp32_hidden_rel_std=1e-5, fp32_hidden_rel_maxabs=1e-5, - # fp32_logits_mse is unused (skipped) but kept for dataclass defaults. - # DPLM2-150M has 30 layers (vs ESM2-8M's 6), so bf16 accumulation gives - # slightly higher MSE and maxabs than ESM2 at the post-final-norm output. - bf16_last_hidden_mse=5e-5, bf16_last_hidden_maxabs=1.5e-1, bf16_last_hidden_rel_maxabs=3e-2, - bf16_hidden_rel_std=3e-2, bf16_hidden_rel_maxabs=5e-2, - ), - "ankh": ParityTolerances( - fp32_last_hidden_mse=1e-12, fp32_last_hidden_maxabs=1e-5, fp32_last_hidden_rel_maxabs=1e-5, - fp32_logits_mse=1e-4, - fp32_hidden_rel_std=1e-5, fp32_hidden_rel_maxabs=1e-5, - # ANKH has no per-block norm (T5-style residual accumulation across 48+ blocks), so - # bf16 rounding compounds into larger absolute values at the pre-norm residual stream. - # Use relative maxabs so a biased activation stream cannot hide behind - # ANKH's large native bf16 residual magnitudes. - bf16_last_hidden_mse=5e-4, bf16_last_hidden_maxabs=2e-1, bf16_last_hidden_rel_maxabs=4e-2, - bf16_logits_mse=5e-2, bf16_hidden_rel_std=5e-2, bf16_hidden_rel_maxabs=1e-1, - ), -} - -EXPECTED_WEIGHT_EXTRAS: Dict[str, set] = { - # fast has these keys, native does not - "ankh": {"lm_head.weight"}, -} - -EXPECTED_NATIVE_EXTRAS: Dict[str, set] = { - # native has these keys, fast does not - # DPLM2: native preserves the pretrained contact-prediction head that the - # DPLM2 authors shipped on top of the ESM2 backbone; the FastPLMs variant - # strips it because FastPLMs DPLM2 is an MLM-only model. - "dplm2": { - "esm.contact_head.regression.weight", - "esm.contact_head.regression.bias", - }, -} - -EXPECTED_VALUE_MISMATCHES: Dict[str, set] = { - # Shared key names whose values are expected to differ for structural - # (non-bug) reasons. Logits parity is also skipped for these families. - # DPLM2: native has `tie_word_embeddings=True` so `lm_head.decoder.weight` - # is an alias for `esm.embeddings.word_embeddings.weight`. FastPLMs DPLM2 - # has `tie_word_embeddings=False` and stores a separately-learned head. - # Word embeddings themselves match exactly; only the head value differs. - "dplm2": {"lm_head.decoder.weight"}, -} - - -def _family_has_head_mismatch(model_key: str) -> bool: - """Return True if fast and native have known lm_head / output-head differences, - so logits parity should not be asserted. - """ - return bool( - EXPECTED_WEIGHT_EXTRAS.get(model_key) - or EXPECTED_NATIVE_EXTRAS.get(model_key) - or EXPECTED_VALUE_MISMATCHES.get(model_key) - ) - - -FIXED_SEQUENCE_LENGTHS = [16, 32, 48, 64, 80, 96, 112, 128] - -# Tokenizer-mode batches used to stress padding behavior. "single" exercises -# no padding; "uniform" exercises mild padding (all lengths within ~50%); -# "skewed" exercises extreme padding (one short, one near-max), which is -# where mask-handling bugs typically surface. -PADDING_SCENARIOS: Dict[str, List[int]] = { - "single": [128], - "uniform": [16, 32, 48, 64, 80, 96, 112, 128], - "skewed": [16, 16, 16, 128, 128], -} - - -def generate_fixed_sequences(seed: int = SEED, lengths: Optional[List[int]] = None) -> List[str]: - if lengths is None: - lengths = FIXED_SEQUENCE_LENGTHS - rng = random.Random(seed) - return [ - "M" + "".join(rng.choices(CANONICAL_AAS, k=L - 1)) - for L in lengths - ] - - -def try_load_native(model_key: str, device: torch.device, dtype: torch.dtype): - config = MODEL_REGISTRY[model_key] - try: - module = importlib.import_module(config["load_official"]) - return module.load_official_model( - reference_repo_id=config["official_path"], - device=device, - dtype=dtype, - ) - except (ImportError, ModuleNotFoundError, FileNotFoundError) as e: - pytest.skip(f"Native deps not installed for {model_key}: {e}") - - -def load_fast(model_key: str, device: torch.device, dtype: torch.dtype) -> nn.Module: - config = MODEL_REGISTRY[model_key] - model = AutoModelForMaskedLM.from_pretrained( - config["fast_path"], trust_remote_code=True, - dtype=dtype, device_map=device, - ).eval() - return model - - -def fast_forward( - model: nn.Module, - model_key: str, - sequences: List[str], - device: torch.device, - output_hidden_states: bool = True, -): - config = MODEL_REGISTRY[model_key] - if config["model_type"] == "E1": - batch = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) - attention_mask = (batch["sequence_ids"] != -1).long() - out = model( - input_ids=batch["input_ids"], - within_seq_position_ids=batch["within_seq_position_ids"], - global_position_ids=batch["global_position_ids"], - sequence_ids=batch["sequence_ids"], - output_hidden_states=output_hidden_states, - ) - return out, attention_mask - batch = tokenize_batch(model, model_key, sequences, device) - kwargs = dict( - input_ids=batch["input_ids"], - attention_mask=batch["attention_mask"], - output_hidden_states=output_hidden_states, - ) - if config["model_type"] == "ESMC": - kwargs["sequence_id"] = batch["attention_mask"].to(torch.bool) - out = model(**kwargs) - return out, batch["attention_mask"] - - -def native_forward( - model: nn.Module, - model_key: str, - sequences: List[str], - device: torch.device, - native_tokenizer, -): - config = MODEL_REGISTRY[model_key] - if config["model_type"] == "E1": - batch = native_tokenizer.get_batch_kwargs(sequences, device=device) - attention_mask = (batch["sequence_ids"] != -1).long() - out = model(**batch, attention_mask=attention_mask) - return out, attention_mask - enc = native_tokenizer(sequences, return_tensors="pt", padding=True) - enc = {k: v.to(device) for k, v in enc.items()} - out = model(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) - return out, enc["attention_mask"] - - -def _masked(tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - return tensor[mask.bool()] - - -def relative_hidden_std(fast: torch.Tensor, native: torch.Tensor, mask: torch.Tensor) -> float: - mask_b = mask.bool() - f = fast[mask_b].float() - n = native[mask_b].float() - diff_std = (f - n).std().item() - native_std = n.std().item() - if native_std < 1e-12: - return 0.0 - return diff_std / native_std - - -def relative_hidden_maxabs(fast: torch.Tensor, native: torch.Tensor, mask: torch.Tensor) -> float: - """Worst-case localized relative error at this layer: maxabs(diff) / maxabs(native). - - Complement to relative_hidden_std. The std metric collapses every position x every - hidden dim into a single scalar, which can hide a single-dimension or single-position - regression. This metric asks "at the worst element, how large is the error relative - to the worst element of native?" - """ - mask_b = mask.bool() - f = fast[mask_b].float() - n = native[mask_b].float() - diff_maxabs = (f - n).abs().max().item() - native_maxabs = n.abs().max().item() - if native_maxabs < 1e-12: - return 0.0 if diff_maxabs < 1e-12 else float("inf") - return diff_maxabs / native_maxabs - - -# Tokenizer parity checks the token contract separately from encoder numerics. - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", [k for k in MODEL_REGISTRY if MODEL_REGISTRY[k]["uses_tokenizer"]]) -def test_tokenizer_parity(model_key: str) -> None: - device = torch.device("cuda") - fast = load_fast(model_key, device, torch.float32) - native_model, native_tok = try_load_native(model_key, device, torch.float32) - - fast_tok = fast.tokenizer - fast_vocab = fast_tok.get_vocab() - native_vocab = native_tok.get_vocab() - assert len(fast_vocab) == len(native_vocab), ( - f"{model_key}: vocab size mismatch fast={len(fast_vocab)} native={len(native_vocab)}" - ) - missing_in_fast = [t for t in native_vocab if t not in fast_vocab] - assert not missing_in_fast, f"{model_key}: tokens missing from fast tokenizer: {missing_in_fast[:5]}" - id_mismatches = [ - (t, native_vocab[t], fast_vocab[t]) - for t in native_vocab - if native_vocab[t] != fast_vocab[t] - ] - assert not id_mismatches, f"{model_key}: token id mismatches: {id_mismatches[:5]}" - - for attr in ("pad_token_id", "cls_token_id", "eos_token_id", "mask_token_id", "unk_token_id"): - f_id = getattr(fast_tok, attr, None) - n_id = getattr(native_tok, attr, None) - assert f_id == n_id, f"{model_key}: {attr} mismatch fast={f_id} native={n_id}" - - del fast, native_model - torch.cuda.empty_cache() - - -# Weight parity is bit exact in fp32 except for documented structural extras. - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY.keys())) -def test_weight_parity_fp32(model_key: str) -> None: - device = torch.device("cuda") - fast = load_fast(model_key, device, torch.float32) - native_model, _ = try_load_native(model_key, device, torch.float32) - - fast_sd = fast.state_dict() - native_sd = native_model.model.state_dict() if hasattr(native_model, "model") else native_model.state_dict() - - expected_fast_extras = EXPECTED_WEIGHT_EXTRAS.get(model_key, set()) - expected_native_extras = EXPECTED_NATIVE_EXTRAS.get(model_key, set()) - expected_value_mismatches = EXPECTED_VALUE_MISMATCHES.get(model_key, set()) - - fast_keys = set(fast_sd.keys()) - expected_fast_extras - native_keys = set(native_sd.keys()) - expected_native_extras - assert fast_keys == native_keys, ( - f"{model_key}: state_dict key sets differ\n" - f" (allowlisted fast extras: {expected_fast_extras})\n" - f" (allowlisted native extras: {expected_native_extras})\n" - f" only_fast (unexpected): {sorted(fast_keys - native_keys)[:5]}\n" - f" only_native (unexpected): {sorted(native_keys - fast_keys)[:5]}" - ) - - shape_mismatches: List[str] = [] - value_mismatches: List[str] = [] - unexpected_value_matches: List[str] = [] - for name in sorted(fast_keys & native_keys): - f = fast_sd[name] - n = native_sd[name] - if f.shape != n.shape: - shape_mismatches.append(f"{name}: {tuple(f.shape)} vs {tuple(n.shape)}") - continue - values_equal = torch.equal(f.float(), n.float()) - if name in expected_value_mismatches: - # Allowlisted mismatches should remain real mismatches; if they - # start matching, the allowlist is stale. - if values_equal: - unexpected_value_matches.append(name) - else: - if not values_equal: - max_abs = (f.float() - n.float()).abs().max().item() - value_mismatches.append(f"{name}: max|Δ|={max_abs:.3e}") - assert not shape_mismatches, f"{model_key}: shape mismatches:\n" + "\n".join(shape_mismatches[:10]) - assert not value_mismatches, f"{model_key}: value mismatches:\n" + "\n".join(value_mismatches[:10]) - assert not unexpected_value_matches, ( - f"{model_key}: the following keys were allowlisted in EXPECTED_VALUE_MISMATCHES " - f"but actually match natively. Remove them from the allowlist:\n" - + "\n".join(unexpected_value_matches) - ) - - del fast, native_model - torch.cuda.empty_cache() - - -# fp32 forward parity keeps the strictest tolerances. - -def _run_forward_parity(model_key: str, dtype: torch.dtype, tol: ParityTolerances, dtype_label: str, scenario: str = "uniform") -> None: - device = torch.device("cuda") - random.seed(SEED) - torch.manual_seed(SEED) - - fast = load_fast(model_key, device, dtype) - native_model, native_tok = try_load_native(model_key, device, dtype) - - sequences = generate_fixed_sequences(lengths=PADDING_SCENARIOS[scenario]) - - parity_context = strict_fp32_matmul() if dtype == torch.float32 else contextlib.nullcontext() - with torch.no_grad(), parity_context: - fout, fmask = fast_forward(fast, model_key, sequences, device, output_hidden_states=True) - nout, nmask = native_forward(native_model, model_key, sequences, device, native_tok) - - assert torch.equal(fmask, nmask), f"{model_key}: attention_mask mismatch between fast and native tokenization" - - fh: Tuple[torch.Tensor, ...] = tuple(fout.hidden_states) - nh: Tuple[torch.Tensor, ...] = tuple(nout.hidden_states) - assert len(fh) == len(nh), f"{model_key}: hidden_states tuple length mismatch fast={len(fh)} native={len(nh)}" - - flast_attr = getattr(fout, "last_hidden_state", None) - nlast_attr = getattr(nout, "last_hidden_state", None) - flast = flast_attr if flast_attr is not None else fh[-1] - nlast = nlast_attr if nlast_attr is not None else nh[-1] - mask_b = fmask.bool() - - last_diff = (flast - nlast).float()[mask_b] - last_native = nlast.float()[mask_b] - last_mse = (last_diff ** 2).mean().item() - last_maxabs = last_diff.abs().max().item() - last_native_maxabs = last_native.abs().max().item() - last_rel_maxabs = last_maxabs / last_native_maxabs if last_native_maxabs > 1e-12 else 0.0 - if dtype == torch.float32: - last_mse_tol = tol.fp32_last_hidden_mse - last_maxabs_tol = tol.fp32_last_hidden_maxabs - last_rel_maxabs_tol = tol.fp32_last_hidden_rel_maxabs - else: - last_mse_tol = tol.bf16_last_hidden_mse - last_maxabs_tol = tol.bf16_last_hidden_maxabs - last_rel_maxabs_tol = tol.bf16_last_hidden_rel_maxabs - assert last_mse <= last_mse_tol, ( - f"{model_key} ({dtype_label}): last_hidden_state MSE={last_mse:.3e} > tol={last_mse_tol:.3e} " - f"(maxabs={last_maxabs:.3e}, rel_maxabs={last_rel_maxabs:.3e})" - ) - assert last_maxabs <= last_maxabs_tol, ( - f"{model_key} ({dtype_label}): last_hidden_state maxabs={last_maxabs:.3e} > tol={last_maxabs_tol:.3e} " - f"(mse={last_mse:.3e}, rel_maxabs={last_rel_maxabs:.3e})" - ) - assert last_rel_maxabs <= last_rel_maxabs_tol, ( - f"{model_key} ({dtype_label}): last_hidden_state rel_maxabs={last_rel_maxabs:.3e} > tol={last_rel_maxabs_tol:.3e} " - f"(maxabs_diff={last_maxabs:.3e}, maxabs_native={last_native_maxabs:.3e}). " - f"A systematic bias may have been introduced even though absolute error looks small." - ) - - # Skip logits parity when fast and native have a known head difference. - # - ANKH: fast is ForMaskedLM with its own head; native T5EncoderModel has - # no head and testing/official/ankh.py bolts on a fresh tied-weight head. - # - DPLM2: native ties lm_head to word embeddings; FastPLMs stores a - # separately-learned head (see EXPECTED_VALUE_MISMATCHES). Encoder hidden - # states match; logits by construction do not. - has_head_mismatch = _family_has_head_mismatch(model_key) - if not has_head_mismatch and hasattr(fout, "logits") and hasattr(nout, "logits") and fout.logits is not None and nout.logits is not None: - logits_diff = (fout.logits - nout.logits).float()[mask_b] - logits_mse = (logits_diff ** 2).mean().item() - logits_mse_tol = tol.fp32_logits_mse if dtype == torch.float32 else tol.bf16_logits_mse - assert logits_mse <= logits_mse_tol, ( - f"{model_key} ({dtype_label}): logits MSE={logits_mse:.3e} > tol={logits_mse_tol:.3e}" - ) - - if dtype == torch.float32: - rel_std_tol = tol.fp32_hidden_rel_std - rel_maxabs_tol = tol.fp32_hidden_rel_maxabs - else: - rel_std_tol = tol.bf16_hidden_rel_std - rel_maxabs_tol = tol.bf16_hidden_rel_maxabs - per_layer: List[Tuple[int, float, float]] = [] - for i in range(len(fh)): - rel_std = relative_hidden_std(fh[i], nh[i], fmask) - rel_maxabs = relative_hidden_maxabs(fh[i], nh[i], fmask) - per_layer.append((i, rel_std, rel_maxabs)) - std_violations = [(i, s, m) for i, s, m in per_layer if s > rel_std_tol] - maxabs_violations = [(i, s, m) for i, s, m in per_layer if m > rel_maxabs_tol] - if std_violations or maxabs_violations: - rendered = "\n".join( - f" layer {i}: rel_diff_std={s:.3e} rel_diff_maxabs={m:.3e}" - for i, s, m in per_layer - ) - reason_parts: List[str] = [] - if std_violations: - reason_parts.append(f"rel_diff_std > tol={rel_std_tol:.3e}") - if maxabs_violations: - reason_parts.append(f"rel_diff_maxabs > tol={rel_maxabs_tol:.3e}") - reason = " and ".join(reason_parts) - pytest.fail( - f"{model_key} ({dtype_label}): per-layer hidden-state divergence ({reason}):\n{rendered}" - ) - - del fast, native_model - torch.cuda.empty_cache() - - -@pytest.mark.gpu -@pytest.mark.parametrize("scenario", list(PADDING_SCENARIOS.keys())) -@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY.keys())) -def test_forward_parity_fp32(model_key: str, scenario: str) -> None: - tol = FAMILY_TOLERANCES[model_key] - _run_forward_parity(model_key, torch.float32, tol, "fp32", scenario=scenario) - - -@pytest.mark.gpu -@pytest.mark.parametrize("scenario", list(PADDING_SCENARIOS.keys())) -@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY.keys())) -def test_forward_parity_bf16(model_key: str, scenario: str) -> None: - tol = FAMILY_TOLERANCES[model_key] - _run_forward_parity(model_key, torch.bfloat16, tol, "bf16", scenario=scenario) - - -# Padding isolation catches backend-specific mask bugs on valid positions. - -# Backends to exercise for the padding-isolation test. kernels_flash is excluded: -# its unpad/pad helpers strip padding entirely, so a batch-shape-dependent -# regression would surface as "doesn't even load" long before this test. -PADDING_BACKENDS: Tuple[str, ...] = ("sdpa", "flex") -PADDING_ISOLATION_TOL: Dict[str, Dict[str, float]] = { - "default": {"maxabs": 1e-3, "mse": 1e-7, "rel_maxabs": float("inf")}, - # ESM3's 48-layer stack exposes a high-magnitude pre-final residual stream. - # TF32 is disabled here; remaining absolute batch-shape outliers are small - # relative to the residual magnitude and logits stay stable. - "esm3": {"maxabs": 1e-2, "mse": 1e-7, "rel_maxabs": 5e-6}, -} - - -@pytest.mark.gpu -@pytest.mark.parametrize("backend", PADDING_BACKENDS) -@pytest.mark.parametrize("model_key", [k for k in MODEL_REGISTRY if MODEL_REGISTRY[k]["uses_tokenizer"]]) -def test_padding_does_not_pollute_valid_positions_fp32(model_key: str, backend: str) -> None: - """A padded batch's valid-position `last_hidden_state` must match the same - sequence run unpadded. - - We only check `last_hidden_state` (not intermediate hidden states) because - `F.scaled_dot_product_attention` is not bit-deterministic across batch - shapes because kernel dispatch and reduction order can differ between batch=1 - and batch=N runs, producing tiny per-layer diffs (~1e-5 maxabs at - intermediate layers, decaying to ~1e-6 after the final norm). Those - diffs are PyTorch SDPA noise, not a parity bug. What WOULD be a bug: - padded keys bleeding into valid-query attention through a broken mask, - which would produce a much larger and persistent diff at - `last_hidden_state`; that is exactly what this test catches. - - Parametrized over backends so a FLEX-specific block-mask bug or an ANKH - flex score_mod that forgets to honor the block mask is caught independently - of the SDPA path. - """ - device = torch.device("cuda") - random.seed(SEED) - - fast = load_fast(model_key, device, torch.float32) - resolved = _apply_backend(fast, backend, model_key) - if resolved is None: - pytest.skip(f"{model_key}: backend {backend} unavailable or fell back") - - short = generate_fixed_sequences(lengths=[16])[0] - long_ = generate_fixed_sequences(lengths=[128])[0] - - with torch.no_grad(), strict_fp32_matmul(): - out_alone, mask_alone = fast_forward(fast, model_key, [short], device, output_hidden_states=True) - out_padded, mask_padded = fast_forward(fast, model_key, [short, long_], device, output_hidden_states=True) - - valid_len = mask_alone.sum().item() - la = getattr(out_alone, "last_hidden_state", None) - lp = getattr(out_padded, "last_hidden_state", None) - last_alone = (la if la is not None else out_alone.hidden_states[-1])[0, :valid_len].float() - last_padded = (lp if lp is not None else out_padded.hidden_states[-1])[0, :valid_len].float() - - diff = (last_alone - last_padded).abs() - diff_max = diff.max().item() - diff_mse = (diff ** 2).mean().item() - base_maxabs = max(last_alone.abs().max().item(), last_padded.abs().max().item()) - diff_rel_maxabs = diff_max / base_maxabs if base_maxabs > 1e-12 else 0.0 - tol = PADDING_ISOLATION_TOL["esm3"] if model_key == "esm3" else PADDING_ISOLATION_TOL["default"] - assert ( - diff_max < tol["maxabs"] - and diff_mse < tol["mse"] - and diff_rel_maxabs < tol["rel_maxabs"] - ), ( - f"{model_key} ({backend}): padding appears to be polluting valid-position outputs (fp32). " - f"At `last_hidden_state`, valid-position diff vs unpadded run is " - f"max|Δ|={diff_max:.3e}, mse={diff_mse:.3e}, " - f"rel_maxabs={diff_rel_maxabs:.3e} " - f"(expected max<{tol['maxabs']:.1e}, mse<{tol['mse']:.1e}, " - f"rel_maxabs<{tol['rel_maxabs']:.1e}). " - f"Failing the combined absolute, MSE, and relative guards indicates an " - f"attention-mask bug: padded keys are likely bleeding into valid query attention." - ) - - del fast - torch.cuda.empty_cache() - - -# Backend consistency is fast-only. fp32 compares SDPA against fp32-capable -# alternatives; bf16 also includes kernels_flash, whose kernels reject fp32. -# ANKH may resolve kernels_flash to flex because T5 relative position bias -# cannot be passed to the flash kernels. - -BACKEND_CONSISTENCY_FP32_MATRIX: Dict[str, Tuple[str, ...]] = { - "esm2": ("flex",), - "esmc": ("flex",), - "esm3": ("flex",), - "e1": ("flex",), - "dplm": ("flex",), - "dplm2": ("flex",), - "ankh": ("flex",), -} - -BACKEND_CONSISTENCY_BF16_MATRIX: Dict[str, Tuple[str, ...]] = { - "esm2": ("kernels_flash", "flex"), - "esmc": ("kernels_flash", "flex"), - "esm3": ("flex",), - "e1": ("kernels_flash", "flex"), - "dplm": ("kernels_flash", "flex"), - "dplm2": ("kernels_flash", "flex"), - "ankh": ("flex",), -} - - -def _apply_backend(model: nn.Module, backend: str, model_key: str) -> Optional[str]: - """Set `model.attn_backend = backend` using the per-family property setter - that every FastPLMs sequence model exposes. Returns the *resolved* backend - as a string, or None if the backend is unavailable on this GPU / image. - - Why this exists: earlier versions of the test suite tried to switch backends - by setting class attributes on the Config subclass, which is silently a - no-op (every config's `__init__` overwrites the class attr). This helper - uses the correct mechanism and verifies the switch actually took effect. - """ - try: - model.attn_backend = backend - except AssertionError as e: - # Backend resolution asserts when the requested implementation is not - # installed in this image or is unsupported by this torch build. - print(f"{model_key}: backend {backend} unavailable: {e}") - return None - except Exception as e: # noqa: BLE001 - backend assertion failures should skip. - print(f"{model_key}: backend {backend} failed to apply: {e}") - return None - return _get_resolved_backend(model, model_key) - - -def _get_resolved_backend(model: nn.Module, model_key: str) -> str: - """Introspect the resolved backend from a known attention submodule. - - This matters for ANKH, which silently falls back kernels_flash -> flex at - the encoder level. If we asked for kernels_flash and got flex back, we want - the test to know. - """ - # Walk modules looking for an attention sub-module with an attn_backend enum. - for module in model.modules(): - attn_backend = getattr(module, "attn_backend", None) - if attn_backend is None: - continue - if hasattr(attn_backend, "value"): - return attn_backend.value - # Fall back to the config. - return model.config.attn_backend - - -# Backend tolerances use two regimes. fp32-capable backends should stay close to -# SDPA at the hidden-state level. bf16 kernels can legitimately drift in raw -# values because tiling and reductions differ, so those checks assert downstream -# agreement: mean-pooled representation cosine and masked-LM top-1 agreement. -# Depth still matters, so tolerances are per family and per backend. -BACKEND_TOL_FP32: Dict[str, Dict[str, Dict[str, float]]] = { - "esm2": {"flex": {"mse": 1e-6, "maxabs": 5e-3, "rel_maxabs": 5e-3}}, - "esmc": {"flex": {"mse": 1e-6, "maxabs": 1e-2, "rel_maxabs": 5e-3}}, - # ESM3 exposes the pre-final-norm stream as last_hidden_state, with fp32 - # activations around 1e4 on the fixed parity batch. With TF32 disabled, - # SDPA vs manual/Flex attention stays at sub-1e-6 relative drift. - "esm3": {"flex": {"mse": 1e-6, "maxabs": 1e-2, "rel_maxabs": 5e-6}}, - "e1": {"flex": {"mse": 1e-6, "maxabs": 1e-2, "rel_maxabs": 5e-3}}, - "dplm": {"flex": {"mse": 1e-6, "maxabs": 5e-2, "rel_maxabs": 5e-3}}, - "dplm2": {"flex": {"mse": 1e-6, "maxabs": 5e-2, "rel_maxabs": 5e-3}}, - "ankh": {"flex": {"mse": 1e-6, "maxabs": 5e-2, "rel_maxabs": 1e-2}}, -} -# ESM3's `last_hidden_state` is the pre-final-norm residual stream. Its absolute -# scale is intentionally large, so fp32 backend consistency is gated by MSE and -# relative maxabs; absolute maxabs is still reported if another guard fails. -BACKEND_FP32_RELATIVE_DOMINANT = {"esm3"} - -# bf16 backend-consistency thresholds are per-family per-backend. Two physics- -# driven metrics with different behaviors across families: -# -# - min_pooled_cosine: per-sequence mean-pool(last_hidden_state) cosine vs sdpa, -# min across sequences. Measures "does the representation direction agree?" -# - min_argmax_agreement: fraction of positions whose top-1 LM logit matches -# sdpa's. Measures "does the downstream MLM prediction agree?" -# -# Empirical behavior (see testing/debug_scripts/investigate_backend_cosine.py): -# ESM2-8M (6 layers): -# flex: pooled_cosine ~ 1.0000, argmax ~ 1.0000 -# kernels_flash: pooled_cosine 0.70-0.86, argmax ~ 0.98 -# ESMC-300M (30 layers): -# flex: pooled_cosine ~ 1.0000, argmax ~ 0.94 -# kernels_flash: pooled_cosine ~ 1.0000, argmax ~ 0.95 -# -# kernels_flash's online softmax + different tile reduction order drifts the -# residual stream in direction at short depth (ESM2-8M) even though the argmax -# is stable. At 30 layers the residual stabilizes direction-wise because each -# layer's LayerNorm re-anchors magnitude, and the argmax disagreement is from -# bf16 rounding on the final LM head softmax, not attention kernel drift. -# -# None for min_cosine means "informational only, don't assert". Use that when -# we know the metric is dominated by a known architectural phenomenon rather -# than by bugs we want to catch. -BACKEND_TOL_BF16_DOWNSTREAM: Dict[str, Dict[str, Dict[str, Optional[float]]]] = { - "esm2": { - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.98}, - "kernels_flash": {"min_cosine": None, "min_argmax_agreement": 0.95}, - }, - "esmc": { - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.90}, - "kernels_flash": {"min_cosine": 0.995, "min_argmax_agreement": 0.90}, - }, - "esm3": { - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.90}, - }, - "e1": { - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.90}, - "kernels_flash": {"min_cosine": None, "min_argmax_agreement": 0.90}, - }, - "dplm": { - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.95}, - "kernels_flash": {"min_cosine": None, "min_argmax_agreement": 0.95}, - }, - "dplm2": { - # DPLM2 has a separately learned lm_head. Native ties this weight, fast - # does not, so backend agreement is checked on the fast head itself. - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.95}, - "kernels_flash": {"min_cosine": None, "min_argmax_agreement": 0.95}, - }, - "ankh": { - # ANKH supports only flex at the user-facing level (kernels_flash falls back). - "flex": {"min_cosine": 0.995, "min_argmax_agreement": 0.90}, - }, -} - - -def _run_backend_consistency_fp32( - model_key: str, - backends: Tuple[str, ...], - backend_tol: Dict[str, Dict[str, float]], -) -> None: - """backend_tol is the per-backend tol dict FOR THIS family (already indexed by caller).""" - device = torch.device("cuda") - random.seed(SEED) - sequences = generate_fixed_sequences() - - baseline = load_fast(model_key, device, torch.float32) - base_resolved = _apply_backend(baseline, "sdpa", model_key) - assert base_resolved == "sdpa", f"{model_key}: sdpa baseline resolved to {base_resolved} (expected 'sdpa')" - - with torch.no_grad(), strict_fp32_matmul(): - base_out, mask = fast_forward(baseline, model_key, sequences, device, output_hidden_states=False) - base_last_attr = getattr(base_out, "last_hidden_state", None) - base_last = base_last_attr if base_last_attr is not None else base_out.hidden_states[-1] - mask_b = mask.bool() - base_valid = base_last[mask_b].float() - base_maxabs = base_valid.abs().max().item() - del baseline, base_out - torch.cuda.empty_cache() - - failures: List[str] = [] - for backend in backends: - alt = load_fast(model_key, device, torch.float32) - resolved = _apply_backend(alt, backend, model_key) - if resolved is None: - del alt - torch.cuda.empty_cache() - pytest.skip(f"{model_key} (fp32): backend {backend} not available") - if resolved != backend: - failures.append( - f"{backend}: requested backend was silently resolved to '{resolved}'. " - f"If this is an intentional fallback (e.g. ANKH kernels_flash -> flex), " - f"remove {backend!r} from the consistency matrix for {model_key!r}." - ) - del alt - torch.cuda.empty_cache() - continue - with torch.no_grad(), strict_fp32_matmul(): - alt_out, _ = fast_forward(alt, model_key, sequences, device, output_hidden_states=False) - alt_last_attr = getattr(alt_out, "last_hidden_state", None) - alt_last = alt_last_attr if alt_last_attr is not None else alt_out.hidden_states[-1] - diff = (alt_last[mask_b].float() - base_valid) - mse = (diff ** 2).mean().item() - maxabs = diff.abs().max().item() - rel_maxabs = maxabs / base_maxabs if base_maxabs > 1e-12 else 0.0 - tol = backend_tol[backend] - maxabs_failed = maxabs > tol["maxabs"] - if model_key in BACKEND_FP32_RELATIVE_DOMINANT: - maxabs_failed = False - if mse > tol["mse"] or maxabs_failed or rel_maxabs > tol["rel_maxabs"]: - failures.append( - f"{backend}: mse={mse:.3e} (tol={tol['mse']:.3e}), " - f"maxabs={maxabs:.3e} (tol={tol['maxabs']:.3e}), " - f"rel_maxabs={rel_maxabs:.3e} (tol={tol['rel_maxabs']:.3e})" - ) - del alt, alt_out - torch.cuda.empty_cache() - - assert not failures, f"{model_key} (fp32): backend consistency failed:\n" + "\n".join(failures) - - -def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """Mean-pool over valid positions per sequence. Returns (batch, hidden).""" - m = attention_mask.bool().unsqueeze(-1).float() - summed = (last_hidden_state.float() * m).sum(dim=1) - counts = m.sum(dim=1).clamp_min(1.0) - return summed / counts - - -def _run_backend_consistency_bf16_downstream( - model_key: str, - backends: Tuple[str, ...], - backend_tol: Dict[str, Dict[str, Optional[float]]], -) -> None: - device = torch.device("cuda") - random.seed(SEED) - sequences = generate_fixed_sequences() - - baseline = load_fast(model_key, device, torch.bfloat16) - base_resolved = _apply_backend(baseline, "sdpa", model_key) - assert base_resolved == "sdpa", f"{model_key}: sdpa baseline resolved to {base_resolved} (expected 'sdpa')" - - with torch.no_grad(): - base_out, mask = fast_forward(baseline, model_key, sequences, device, output_hidden_states=False) - base_last_attr = getattr(base_out, "last_hidden_state", None) - base_last = base_last_attr if base_last_attr is not None else base_out.hidden_states[-1] - base_pooled = _mean_pool(base_last, mask) # (B, D) in fp32 - mask_b = mask.bool() - - base_logits = getattr(base_out, "logits", None) - base_argmax = None - if base_logits is not None: - base_argmax = base_logits.float()[mask_b].argmax(dim=-1) - - del baseline, base_out - torch.cuda.empty_cache() - - failures: List[str] = [] - for backend in backends: - alt = load_fast(model_key, device, torch.bfloat16) - resolved = _apply_backend(alt, backend, model_key) - if resolved is None: - del alt - torch.cuda.empty_cache() - pytest.skip(f"{model_key} (bf16): backend {backend} not available") - if resolved != backend: - failures.append( - f"{backend}: requested backend was silently resolved to '{resolved}'. " - f"Remove {backend!r} from the consistency matrix for {model_key!r} if this is intentional." - ) - del alt - torch.cuda.empty_cache() - continue - with torch.no_grad(): - alt_out, _ = fast_forward(alt, model_key, sequences, device, output_hidden_states=False) - alt_last_attr = getattr(alt_out, "last_hidden_state", None) - alt_last = alt_last_attr if alt_last_attr is not None else alt_out.hidden_states[-1] - alt_pooled = _mean_pool(alt_last, mask) - # Per-sequence cosine similarity, then min across sequences. - cos_per_seq = F.cosine_similarity(base_pooled, alt_pooled, dim=-1) # (B,) - min_cos = cos_per_seq.min().item() - tol = backend_tol[backend] - # Always print diagnostics; backend cosine positions are useful context. - print(f" {backend}: min_pooled_cosine={min_cos:.4f}") - if tol["min_cosine"] is not None and min_cos < tol["min_cosine"]: - failures.append( - f"{backend}: min per-sequence pooled cosine = {min_cos:.4f} " - f"(tol >= {tol['min_cosine']:.4f})" - ) - # Argmax agreement (if logits are available). - alt_logits = getattr(alt_out, "logits", None) - if base_argmax is not None and alt_logits is not None: - alt_argmax = alt_logits.float()[mask_b].argmax(dim=-1) - agreement = (base_argmax == alt_argmax).float().mean().item() - print(f" {backend}: argmax_agreement={agreement:.4f}") - if agreement < tol["min_argmax_agreement"]: - failures.append( - f"{backend}: logits argmax agreement vs sdpa = {agreement:.4f} " - f"(tol >= {tol['min_argmax_agreement']:.4f})" - ) - del alt, alt_out - torch.cuda.empty_cache() - - assert not failures, f"{model_key} (bf16 downstream): backend consistency failed:\n" + "\n".join(failures) - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", list(BACKEND_CONSISTENCY_FP32_MATRIX.keys())) -def test_backend_consistency_fp32(model_key: str) -> None: - """fp32 backend parity: sdpa vs flex. Strict raw-value agreement. - - kernels_flash is excluded: its ops reject fp32 at the kernel level. - """ - _run_backend_consistency_fp32( - model_key, - BACKEND_CONSISTENCY_FP32_MATRIX[model_key], - BACKEND_TOL_FP32[model_key], - ) - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", list(BACKEND_CONSISTENCY_BF16_MATRIX.keys())) -def test_backend_consistency_bf16_downstream(model_key: str) -> None: - """bf16 backend parity: sdpa vs flex vs kernels_flash at the downstream level. - - In bf16, different attention kernels produce meaningfully different raw - hidden states (tiling / reduction order; well-documented in the legacy - compliance test). What matters for downstream users is whether: - - - mean-pooled embeddings (classification / regression use) agree in - direction (cosine similarity) - - top-1 argmax predictions (MLM use) agree - - Both are checked here with per-family per-backend thresholds. - """ - _run_backend_consistency_bf16_downstream( - model_key, - BACKEND_CONSISTENCY_BF16_MATRIX[model_key], - BACKEND_TOL_BF16_DOWNSTREAM[model_key], - ) - - -# `embed_dataset(...)` must match native forward plus mean pooling, because this -# is the representation path downstream users actually call. - -# Per-family absolute/max-abs tolerances for the mean-pooled embedding. -EMBED_DATASET_TOL: Dict[str, Dict[str, float]] = { - "esm2": {"mse": 5e-8, "maxabs": 5e-3}, - "esmc": {"mse": 5e-8, "maxabs": 5e-3}, - "esm3": {"mse": 5e-8, "maxabs": 5e-3}, - "dplm": {"mse": 5e-8, "maxabs": 5e-3}, - # DPLM2: encoder output identical to ESM backbone on AA input; mean-pool parity tight. - "dplm2": {"mse": 5e-8, "maxabs": 5e-3}, - "e1": {"mse": 5e-6, "maxabs": 2e-2}, # GQA + block-causal global layers accumulate rounding. - # ANKH-base post-final-RMSNorm activations are modest (~O(1)) but the - # mean-pool aggregates over the full sequence; 5e-3 maxabs is plenty tight. - "ankh": {"mse": 5e-8, "maxabs": 5e-3}, -} - -PIPELINE_MODEL_KEYS = [k for k in EMBED_DATASET_TOL if k in MODEL_REGISTRY] - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", PIPELINE_MODEL_KEYS) -def test_embed_dataset_pipeline_parity(model_key: str) -> None: - device = torch.device("cuda") - random.seed(SEED) - - config = MODEL_REGISTRY[model_key] - sequences = generate_fixed_sequences() - fast = load_fast(model_key, device, torch.float32) - native_model, native_tok = try_load_native(model_key, device, torch.float32) - - # Tokenizer mode vs sequence mode: E1 has no tokenizer. - tokenizer_mode = config["uses_tokenizer"] - with strict_fp32_matmul(): - fast_embeddings = fast.embed_dataset( - sequences=sequences, - tokenizer=fast.tokenizer if tokenizer_mode else None, - batch_size=4, max_len=256, truncate=True, - full_embeddings=False, - embed_dtype=torch.float32, - pooling_types=["mean"], - num_workers=0, sql=False, save=False, - padding="max_length" if tokenizer_mode else "longest", - ) - assert fast_embeddings is not None - - tol = EMBED_DATASET_TOL[model_key] - with torch.no_grad(), strict_fp32_matmul(): - failures: List[str] = [] - for seq in sequences: - # Produce a native mean-pooled embedding for this single sequence. - if tokenizer_mode: - enc = native_tok([seq], return_tensors="pt", padding=True) - enc = {k: v.to(device) for k, v in enc.items()} - out = native_model( - input_ids=enc["input_ids"], - attention_mask=enc["attention_mask"], - output_hidden_states=True, - ) - last_attr = getattr(out, "last_hidden_state", None) - last = (last_attr if last_attr is not None else out.hidden_states[-1]).float() - m = enc["attention_mask"].bool().unsqueeze(-1).float() - else: - # E1: native_tok is the E1BatchPreparer. - batch = native_tok.get_batch_kwargs([seq], device=device) - attention_mask = (batch["sequence_ids"] != -1).long() - out = native_model(**batch, attention_mask=attention_mask) - last_attr = getattr(out, "last_hidden_state", None) - last = (last_attr if last_attr is not None else out.hidden_states[-1]).float() - m = attention_mask.bool().unsqueeze(-1).float() - pooled = (last * m).sum(dim=1) / m.sum(dim=1).clamp_min(1.0) - pooled = pooled.squeeze(0).cpu() - fast_vec = fast_embeddings[seq].cpu().float() - assert fast_vec.shape == pooled.shape, ( - f"{model_key} seq_len={len(seq)}: shape mismatch fast={tuple(fast_vec.shape)} " - f"native={tuple(pooled.shape)}" - ) - mse = ((fast_vec - pooled) ** 2).mean().item() - maxabs = (fast_vec - pooled).abs().max().item() - if mse > tol["mse"] or maxabs > tol["maxabs"]: - failures.append( - f"seq_len={len(seq)}: mse={mse:.3e} (tol={tol['mse']:.3e}) " - f"maxabs={maxabs:.3e} (tol={tol['maxabs']:.3e})" - ) - - assert not failures, ( - f"{model_key}: embed_dataset pipeline parity failed:\n" + "\n".join(failures) - ) - del fast, native_model - torch.cuda.empty_cache() - - -# Backend setter semantics: post-load switching must propagate to every -# attention layer, or backend-parametrized tests silently become no-ops. - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY.keys())) -def test_attn_backend_setter_propagates(model_key: str) -> None: - """Structural check: does `model.attn_backend = X` propagate to every attention submodule? - - Unlike the other parity tests, this does not require the native package for - the family; it is a FastPLMs-only invariant. Run in any image. - """ - device = torch.device("cuda") - fast = load_fast(model_key, device, torch.float32) - - # Start from SDPA, switch to flex, verify every attention submodule flipped. - fast.attn_backend = "sdpa" - assert _get_resolved_backend(fast, model_key) == "sdpa", ( - f"{model_key}: attn_backend setter did not propagate 'sdpa' to attention modules" - ) - - try: - fast.attn_backend = "flex" - except (AssertionError, ValueError) as e: - pytest.skip(f"{model_key}: flex backend unavailable: {e}") - resolved = _get_resolved_backend(fast, model_key) - assert resolved == "flex", ( - f"{model_key}: after setting attn_backend='flex', attention modules report '{resolved}'. " - f"The setter is not propagating. Every other backend-parametrized test becomes a no-op." - ) - - # Every attention-like submodule should report 'flex' now (not just one). - mismatches = [] - for name, module in fast.named_modules(): - ab = getattr(module, "attn_backend", None) - if ab is None or not hasattr(ab, "value"): - continue - if ab.value != "flex": - mismatches.append(f"{name}: {ab.value}") - assert not mismatches, ( - f"{model_key}: after setting attn_backend='flex', these submodules are still on a " - f"different backend:\n" + "\n".join(mismatches[:10]) - ) - - del fast - torch.cuda.empty_cache() diff --git a/testing/test_throughput.py b/testing/test_throughput.py deleted file mode 100644 index 40abfbe..0000000 --- a/testing/test_throughput.py +++ /dev/null @@ -1,90 +0,0 @@ -"""Pytest-based throughput benchmark across models, backends, batch sizes, and sequence lengths. - -Wraps ThroughputChecker from throughput.py and saves structured results as JSON, CSV, -and a PNG plot to /workspace/ (or the current directory if /workspace/ does not exist). - -Marked as `slow` + `gpu` because it runs compiled inference across many configurations. -""" - -import json -from pathlib import Path -from typing import Dict, List - -import pytest -import torch - -from testing.conftest import BACKENDS, FULL_MODEL_REGISTRY -from testing.throughput import ThroughputChecker, plot_results, save_structured_results - - -# Models to benchmark: one representative per family that uses a tokenizer. -# E1 is excluded because throughput.py uses model.tokenizer (tokenizer mode only). -THROUGHPUT_MODELS: Dict[str, str] = { - "ESM2-8M": "Synthyra/ESM2-8M", - "ESMplusplus_small": "Synthyra/ESMplusplus_small", - "DPLM-150M": "Synthyra/DPLM-150M", - "DPLM2-150M": "Synthyra/DPLM2-150M", -} - -BATCH_SIZES = [2, 4, 8] -SEQUENCE_LENGTHS = [64, 128, 256, 512, 1024] -MIN_LENGTH = 32 - -# Fewer timed batches than standalone for faster pytest runs -WARMUP_BATCHES = 5 -TIMED_BATCHES = 25 - - -def _get_output_dir() -> Path: - workspace = Path("/workspace") - if workspace.is_dir(): - return workspace - return Path(".") - - -@pytest.mark.gpu -@pytest.mark.slow -def test_throughput_benchmark() -> None: - """Benchmark tokens/sec across backends, batch sizes, and sequence lengths. - - Saves results as JSON, CSV, and PNG to the output directory. - """ - assert torch.cuda.is_available(), "Throughput benchmark requires CUDA" - - checker = ThroughputChecker( - warmup_batches=WARMUP_BATCHES, - timed_batches=TIMED_BATCHES, - ) - - all_results: Dict[str, Dict] = {} - - for model_name, model_path in THROUGHPUT_MODELS.items(): - print(f"\n--- Benchmarking {model_name} ({model_path}) ---") - results = checker.evaluate( - model_path, - BATCH_SIZES, - min_length=MIN_LENGTH, - sequence_lengths=SEQUENCE_LENGTHS, - backends=list(BACKENDS), - ) - all_results[model_path] = results - - output_dir = _get_output_dir() - save_structured_results(all_results, str(output_dir)) - plot_results(all_results, str(output_dir / "throughput_comparison.png")) - - # Verify that results were collected - json_path = output_dir / "throughput_results.json" - assert json_path.exists(), "throughput_results.json was not created" - with open(json_path) as f: - rows = json.load(f) - assert len(rows) > 0, "No throughput results collected" - - # Sanity check: at least one backend produced positive throughput - max_throughput = max(r["tokens_per_sec"] for r in rows) - assert max_throughput > 0, "All throughput measurements are zero" - - print(f"\nResults saved to {output_dir}") - print(f" JSON: {output_dir / 'throughput_results.json'}") - print(f" CSV: {output_dir / 'throughput_results.csv'}") - print(f" PNG: {output_dir / 'throughput_comparison.png'}") diff --git a/testing/test_tokenizer_contract.py b/testing/test_tokenizer_contract.py deleted file mode 100644 index 3229d3a..0000000 --- a/testing/test_tokenizer_contract.py +++ /dev/null @@ -1,347 +0,0 @@ -"""Tokenizer contract tests for all FastPLMs sequence checkpoints.""" - -from __future__ import annotations - -import shutil -from pathlib import Path -from typing import Dict -from unittest.mock import patch - -import pytest -import torch -from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer, EsmTokenizer - -from fastplms.ankh.modeling_ankh import FAST_ANKH_ENCODER, FastAnkhConfig -from fastplms.dplm2.modeling_dplm2 import _normalize_dplm2_input_ids -from fastplms.e1.modeling_e1 import E1BatchPreparer, E1Config, E1ForMaskedLM, get_tokenizer -from fastplms.esm3.modeling_esm3 import ( - SEQUENCE_VOCAB as ESM3_SEQUENCE_VOCAB, - EsmSequenceTokenizer as ESM3SequenceTokenizer, -) -from fastplms.esm_plusplus.modeling_esm_plusplus import EsmSequenceTokenizer -from testing.conftest import CANONICAL_AAS, FULL_MODEL_REGISTRY, mark_by_size - - -TOKENIZER_REFERENCE_KEYS = [ - key - for key, value in FULL_MODEL_REGISTRY.items() - if value["uses_tokenizer"] and value["model_type"] != "ESM3" -] -ESM3_MODEL_KEYS = [ - key - for key, value in FULL_MODEL_REGISTRY.items() - if value["model_type"] == "ESM3" -] -DPLM2_MODEL_KEYS = [ - key - for key, value in FULL_MODEL_REGISTRY.items() - if value["model_type"] == "DPLM2" -] -CANONICAL_SEQUENCES = [ - "M" + CANONICAL_AAS, - "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSH", - "MXXBZUOACDEFGHIKLMNPQRSTVWY", -] - - -def _repo_root() -> Path: - return Path(__file__).resolve().parents[1] - - -def _e1_tokenizer_json() -> Path: - return _repo_root() / "fastplms" / "e1" / "tokenizer.json" - - -def _tiny_ankh_config(name_or_path: str = "") -> FastAnkhConfig: - config = FastAnkhConfig( - vocab_size=4, - d_model=8, - d_kv=4, - d_ff=16, - num_heads=2, - num_layers=1, - ) - config._name_or_path = name_or_path - return config - - -def _fast_tokenizer(config: Dict): - if config["model_type"] == "ANKH": - encoder = FAST_ANKH_ENCODER(_tiny_ankh_config(config["fast_path"])) - return encoder.tokenizer - if config["model_type"] == "ESMC": - return EsmSequenceTokenizer() - if config["model_type"] in ("ESM2", "DPLM", "DPLM2"): - return EsmTokenizer.from_pretrained(config["fast_path"]) - return AutoTokenizer.from_pretrained( - config["fast_path"], - trust_remote_code=True, - ) - - -def _reference_tokenizer(config: Dict): - if config["model_type"] == "ESMC": - return EsmSequenceTokenizer() - if config["model_type"] in ("ESM2", "DPLM", "DPLM2"): - return EsmTokenizer.from_pretrained(config["official_path"]) - return AutoTokenizer.from_pretrained( - config["official_path"], - trust_remote_code=True, - ) - - -def _token_ids(tokenizer, sequence: str) -> torch.Tensor: - encoded = tokenizer( - sequence, - return_tensors="pt", - ) - return encoded["input_ids"] - - -def _special_token_ids(tokenizer) -> Dict[str, int | None]: - return { - "pad_token_id": tokenizer.pad_token_id, - "cls_token_id": tokenizer.cls_token_id, - "eos_token_id": tokenizer.eos_token_id, - "mask_token_id": tokenizer.mask_token_id, - "unk_token_id": tokenizer.unk_token_id, - } - - -@pytest.mark.parametrize( - "model_key", - mark_by_size(TOKENIZER_REFERENCE_KEYS, FULL_MODEL_REGISTRY), -) -def test_sequence_tokenizer_matches_reference(model_key: str) -> None: - config = FULL_MODEL_REGISTRY[model_key] - fast_tok = _fast_tokenizer(config) - reference_tok = _reference_tokenizer(config) - - fast_vocab = fast_tok.get_vocab() - reference_vocab = reference_tok.get_vocab() - assert len(fast_vocab) == len(reference_vocab), ( - f"{model_key}: vocab size mismatch fast={len(fast_vocab)} " - f"reference={len(reference_vocab)}" - ) - - missing_in_fast = [ - token - for token in reference_vocab - if token not in fast_vocab - ] - assert not missing_in_fast, ( - f"{model_key}: tokens missing from fast tokenizer: {missing_in_fast[:5]}" - ) - - id_mismatches = [ - (token, reference_vocab[token], fast_vocab[token]) - for token in reference_vocab - if reference_vocab[token] != fast_vocab[token] - ] - assert not id_mismatches, ( - f"{model_key}: token id mismatches: {id_mismatches[:5]}" - ) - - assert _special_token_ids(fast_tok) == _special_token_ids(reference_tok), ( - f"{model_key}: special token ids differ" - ) - - for sequence in CANONICAL_SEQUENCES: - fast_ids = _token_ids(fast_tok, sequence) - reference_ids = _token_ids(reference_tok, sequence) - assert torch.equal(fast_ids, reference_ids), ( - f"{model_key}: encoded ids differ for {sequence[:16]} " - f"fast={fast_ids[0, :8].tolist()} " - f"reference={reference_ids[0, :8].tolist()}" - ) - - -@pytest.mark.parametrize( - "model_key", - mark_by_size(ESM3_MODEL_KEYS, FULL_MODEL_REGISTRY), -) -def test_esm3_sequence_tokenizer_contract(model_key: str) -> None: - tokenizer = ESM3SequenceTokenizer() - expected_vocab = { - token: token_id - for token_id, token in enumerate(ESM3_SEQUENCE_VOCAB) - } - - assert tokenizer.get_vocab() == expected_vocab, ( - f"{model_key}: ESM3 sequence vocabulary changed" - ) - assert _special_token_ids(tokenizer) == { - "pad_token_id": 1, - "cls_token_id": 0, - "eos_token_id": 2, - "mask_token_id": 32, - "unk_token_id": 3, - } - - for sequence in CANONICAL_SEQUENCES: - encoded = _token_ids(tokenizer, sequence) - expected_ids = [0] + [expected_vocab[token] for token in sequence] + [2] - assert encoded[0].tolist() == expected_ids, ( - f"{model_key}: encoded ids differ for {sequence[:16]}" - ) - - -def test_ankh_tokenizer_loader_falls_back_for_bare_config() -> None: - encoder = FAST_ANKH_ENCODER(_tiny_ankh_config()) - fast_tok = encoder.tokenizer - reference_tok = AutoTokenizer.from_pretrained("ElnaggarLab/ankh-base") - - assert len(fast_tok.get_vocab()) == len(reference_tok.get_vocab()) - assert _special_token_ids(fast_tok) == _special_token_ids(reference_tok) - assert torch.equal( - _token_ids(fast_tok, CANONICAL_SEQUENCES[0]), - _token_ids(reference_tok, CANONICAL_SEQUENCES[0]), - ) - - -@pytest.mark.parametrize( - "model_key", - mark_by_size(DPLM2_MODEL_KEYS, FULL_MODEL_REGISTRY), -) -def test_dplm2_tokenizer_special_ids_normalize_in_range(model_key: str) -> None: - config = FULL_MODEL_REGISTRY[model_key] - fast_config = AutoConfig.from_pretrained( - config["fast_path"], - trust_remote_code=True, - ) - tokenizer = EsmTokenizer.from_pretrained(config["fast_path"]) - - generic_special_ids = torch.tensor([[ - fast_config.vocab_size, - fast_config.vocab_size + 1, - fast_config.vocab_size + 2, - fast_config.vocab_size + 3, - -100, - ]]) - expected = torch.tensor([[2, 3, 0, 32, -100]]) - normalized_special_ids = _normalize_dplm2_input_ids( - generic_special_ids, - vocab_size=fast_config.vocab_size, - ) - assert torch.equal(normalized_special_ids, expected) - - encoded = tokenizer( - CANONICAL_SEQUENCES, - return_tensors="pt", - padding=True, - ) - normalized_input_ids = _normalize_dplm2_input_ids( - encoded["input_ids"], - vocab_size=fast_config.vocab_size, - ) - valid_ids = normalized_input_ids[normalized_input_ids.ge(0)] - assert bool(valid_ids.lt(fast_config.vocab_size).all()) - - -def test_e1_config_uses_static_token_constants() -> None: - with patch( - "fastplms.e1.modeling_e1.get_tokenizer", - side_effect=AssertionError("E1Config should not load tokenizer.json"), - ): - config = E1Config() - - assert config.vocab_size == 34 - assert config.pad_token_id == 0 - assert config.bos_token_id == 1 - assert config.eos_token_id == 2 - - -def test_e1_get_tokenizer_prefers_local_model_dir(tmp_path: Path) -> None: - shutil.copyfile(_e1_tokenizer_json(), tmp_path / "tokenizer.json") - - with patch( - "huggingface_hub.hf_hub_download", - side_effect=AssertionError("local tokenizer load should not call Hub download"), - ) as hf_hub_download: - tokenizer = get_tokenizer(tmp_path, local_files_only=True) - - assert not hf_hub_download.called - assert tokenizer.token_to_id("") == 0 - assert tokenizer.get_vocab_size() == 34 - - -def test_e1_get_tokenizer_local_files_only_missing_local_source_raises( - tmp_path: Path, -) -> None: - with patch("fastplms.e1.modeling_e1.os.path.isfile", return_value=False): - with patch( - "huggingface_hub.hf_hub_download", - side_effect=AssertionError("missing local tokenizer should not call Hub download"), - ) as hf_hub_download: - with pytest.raises(FileNotFoundError): - get_tokenizer(tmp_path, local_files_only=True) - - assert not hf_hub_download.called - - -def test_e1_automodel_local_files_only_uses_local_tokenizer(tmp_path: Path) -> None: - config = E1Config( - hidden_size=8, - intermediate_size=16, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - max_num_sequences=4, - max_num_positions_within_seq=64, - max_num_positions_global=128, - ) - config.auto_map = { - "AutoConfig": "modeling_e1.E1Config", - "AutoModelForMaskedLM": "modeling_e1.E1ForMaskedLM", - } - model = E1ForMaskedLM(config) - model.save_pretrained(tmp_path) - shutil.copyfile(_e1_tokenizer_json(), tmp_path / "tokenizer.json") - shutil.copyfile( - _repo_root() / "fastplms" / "e1" / "modeling_e1.py", - tmp_path / "modeling_e1.py", - ) - - with patch( - "huggingface_hub.hf_hub_download", - side_effect=AssertionError("local AutoModel load should not call Hub download"), - ) as hf_hub_download: - loaded = AutoModelForMaskedLM.from_pretrained( - tmp_path, - trust_remote_code=True, - local_files_only=True, - ) - - assert not hf_hub_download.called - assert loaded.prep_tokens.tokenizer.token_to_id("") == 0 - - -def test_e1_sequence_mode_tokenizer_contract() -> None: - tokenizer = get_tokenizer() - preparer = E1BatchPreparer(tokenizer=tokenizer) - sequences = [ - "M" + CANONICAL_AAS, - "M" + CANONICAL_AAS[::-1], - ] - - assert tokenizer.token_to_id("") == 0 - for token in ("", "", "1", "2", "?", "X"): - token_id = tokenizer.token_to_id(token) - assert token_id is not None, f"E1 token missing from tokenizer: {token}" - - batch = preparer.get_batch_kwargs( - sequences, - device=torch.device("cpu"), - ) - input_ids = batch["input_ids"] - sequence_ids = batch["sequence_ids"] - within_seq_position_ids = batch["within_seq_position_ids"] - global_position_ids = batch["global_position_ids"] - - assert input_ids.shape == sequence_ids.shape - assert input_ids.shape == within_seq_position_ids.shape - assert input_ids.shape == global_position_ids.shape - assert input_ids.shape[0] == len(sequences) - assert bool((sequence_ids == -1).eq(input_ids == tokenizer.token_to_id("")).all()) - assert bool((within_seq_position_ids[sequence_ids != -1] >= 0).all()) - assert bool((global_position_ids[sequence_ids != -1] >= 0).all()) diff --git a/testing/test_ttt.py b/testing/test_ttt.py deleted file mode 100644 index fd87404..0000000 --- a/testing/test_ttt.py +++ /dev/null @@ -1,238 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -import torch.nn as nn - -from fastplms.ankh.modeling_ankh import FastAnkhForMaskedLM -from fastplms.dplm.modeling_dplm import DPLMForMaskedLM -from fastplms.dplm2.modeling_dplm2 import DPLM2ForMaskedLM -from fastplms.e1.modeling_e1 import E1ForMaskedLM -from fastplms.esm2.modeling_fastesm import FastEsmForMaskedLM -from fastplms.esm3.modeling_esm3 import FastESM3Model -from fastplms.esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM -from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model -from fastplms.test_time_training import ( - FastPLMTestTimeTrainingMixin, - LoraInjectedLinear, -) -from testing.conftest import MODEL_REGISTRY, STRUCTURE_MODEL_REGISTRY - - -TEST_SEQUENCE = "MSTNPKPQRKTKRNT" -LOCAL_MODEL_CLASSES = { - "esm2": FastEsmForMaskedLM, - "esmc": ESMplusplusForMaskedLM, - "esm3": FastESM3Model, - "e1": E1ForMaskedLM, - "dplm": DPLMForMaskedLM, - "dplm2": DPLM2ForMaskedLM, - "ankh": FastAnkhForMaskedLM, -} - - -class DummyConfig: - vocab_size = 8 - - -class DummyTokenizer: - pad_token_id = 0 - cls_token_id = 1 - eos_token_id = 2 - mask_token_id = 3 - all_special_ids = [0, 1, 2, 3] - - def __init__(self) -> None: - self.vocab = { - "A": 4, - "C": 5, - "D": 6, - "E": 7, - } - - def __call__( - self, - seq: str | list[str], - return_tensors: str = "pt", - padding: bool = True, - ) -> dict[str, torch.Tensor]: - del return_tensors, padding - sequences = [seq] if isinstance(seq, str) else seq - encoded = [] - for sequence in sequences: - encoded.append( - [self.cls_token_id] - + [self.vocab[aa] for aa in sequence] - + [self.eos_token_id] - ) - max_len = max(len(ids) for ids in encoded) - input_ids = torch.full((len(encoded), max_len), self.pad_token_id) - for row, ids in enumerate(encoded): - input_ids[row, : len(ids)] = torch.tensor(ids) - return {"input_ids": input_ids.long()} - - -class DummyTTTModel(FastPLMTestTimeTrainingMixin, nn.Module): - def __init__(self) -> None: - nn.Module.__init__(self) - self.config = DummyConfig() - self.tokenizer = DummyTokenizer() - self.embed = nn.Embedding(self.config.vocab_size, 8) - self.backbone = nn.Sequential( - nn.Linear(8, 8), - nn.GELU(), - nn.Linear(8, 8), - ) - self.lm_head = nn.Linear(8, self.config.vocab_size) - self.init_ttt( - { - "steps": 1, - "ags": 1, - "batch_size": 1, - "mask_ratio": 1.0, - "bert_leave_prob": 0.0, - "bert_replace_prob": 0.0, - "lora_rank": 2, - "lora_alpha": 1.0, - } - ) - - def _ttt_get_trainable_modules(self) -> list[nn.Module]: - return [self.backbone] - - def forward( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor | None = None, - ): - del attention_mask - hidden = self.backbone(self.embed(input_ids)) - return SimpleNamespace(logits=self.lm_head(hidden)) - - -def test_ttt_masking_masks_only_residue_tokens() -> None: - model = DummyTTTModel() - tokenized = model._ttt_tokenize(seq="ACDE") - generator = torch.Generator() - generator.manual_seed(0) - - batch, labels = model._ttt_sample_batch(tokenized, generator) - - assert isinstance(batch, torch.Tensor) - assert labels[0, 0].item() == -100 - assert labels[0, -1].item() == -100 - assert torch.all(batch[labels != -100] == model.tokenizer.mask_token_id) - - -def test_ttt_lora_injection_is_lazy_and_backbone_scoped() -> None: - model = DummyTTTModel() - - assert all("lora_" not in name for name in model.state_dict()) - - model._ttt_ensure_initialized() - - assert any(isinstance(module, LoraInjectedLinear) for module in model.backbone.modules()) - assert not any(isinstance(module, LoraInjectedLinear) for module in model.lm_head.modules()) - - -def test_ttt_only_lora_params_change_and_reset_restores_adapter() -> None: - model = DummyTTTModel() - model._ttt_ensure_initialized() - initial = { - name: parameter.detach().clone() - for name, parameter in model.named_parameters() - } - - metrics = model.ttt(seq="ACDE") - - changed = [ - name - for name, parameter in model.named_parameters() - if not torch.equal(parameter.detach(), initial[name]) - ] - assert len(metrics["losses"]) == 1 - assert len(changed) > 0 - assert all("lora_" in name for name in changed) - - model.ttt_reset() - for name, parameter in model.named_parameters(): - torch.testing.assert_close(parameter.detach(), initial[name]) - - -@pytest.mark.gpu -@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY)) -def test_sequence_model_ttt_smoke(model_key: str) -> None: - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for FastPLMs model smoke tests.") - config = MODEL_REGISTRY[model_key] - model_cls = LOCAL_MODEL_CLASSES[model_key] - model = ( - model_cls.from_pretrained( - config["fast_path"], - dtype=torch.float32, - ) - .eval() - .cuda() - ) - metrics = model.ttt( - seq=TEST_SEQUENCE, - ttt_config={ - "steps": 1, - "ags": 1, - "batch_size": 1, - "crop_size": 64, - "lora_rank": 2, - "lora_alpha": 1.0, - }, - ) - - assert len(metrics["losses"]) == 1 - assert callable(model.ttt_reset) - model.ttt_reset() - del model - torch.cuda.empty_cache() - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -def test_esmfold2_ttt_smoke() -> None: - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for ESMFold2 TTT smoke tests.") - config = STRUCTURE_MODEL_REGISTRY["esmfold2_fast"] - model = ( - ESMFold2Model.from_pretrained( - config["fast_path"], - load_esmc=True, - dtype=torch.float32, - ) - .eval() - .cuda() - ) - - result = model.fold_protein( - TEST_SEQUENCE, - num_loops=1, - num_sampling_steps=1, - num_diffusion_samples=1, - seed=0, - ttt=True, - ttt_config={ - "steps": 1, - "ags": 1, - "batch_size": 1, - "crop_size": 64, - "lora_rank": 2, - "lora_alpha": 1.0, - }, - ) - - assert result.ttt_metrics is not None - assert len(result.ttt_metrics["losses"]) == 1 - assert len(result.ttt_metrics["step_plddts"]) == 2 - assert result.ttt_metrics["best_step"] in {0, 1} - - del model, result - torch.cuda.empty_cache() diff --git a/testing/throughput.py b/testing/throughput.py deleted file mode 100644 index 1eff110..0000000 --- a/testing/throughput.py +++ /dev/null @@ -1,342 +0,0 @@ -import entrypoint_setup - -import argparse -import copy -import random -import time -from pathlib import Path -from typing import Dict, List, Tuple - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns -import torch -from tqdm.auto import tqdm -from transformers import AutoModelForMaskedLM - - -def set_seed(seed: int) -> None: - random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - - -SUPPORTED_BACKENDS = ("sdpa", "flex", "kernels_flash") - - -class ThroughputChecker: - def __init__( - self, - warmup_batches: int = 10, - timed_batches: int = 100, - ) -> None: - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - self.warmup_batches = warmup_batches - self.timed_batches = timed_batches - self.canonical_amino_acids = "ACDEFGHIKLMNPQRSTVWY" - - def _load_model(self, model_path: str) -> torch.nn.Module: - model = AutoModelForMaskedLM.from_pretrained( - model_path, - dtype=torch.bfloat16, - device_map=self.device, - trust_remote_code=True, - ).eval() - return model - - def _generate_random_sequence(self, length: int) -> str: - return "M" + "".join(random.choices(self.canonical_amino_acids, k=length - 1)) - - def _generate_random_batch(self, batch_size: int, min_length: int, max_length: int) -> List[str]: - max_length_example = self._generate_random_sequence(max_length) - return [max_length_example] + [ - self._generate_random_sequence(random.randint(min_length, max_length)) - for _ in range(batch_size - 1) - ] - - def _make_batch(self, tokenizer: object, batch_size: int, min_length: int, max_length: int) -> Tuple[Dict[str, torch.Tensor], int]: - """Generate and tokenize one batch, returning (tokenized_on_device, nonpad_token_count).""" - batch = self._generate_random_batch(batch_size, min_length, max_length) - tokenized = tokenizer( - batch, - return_tensors="pt", - padding="max_length", - max_length=max_length, - truncation=True, - add_special_tokens=True, - ) - if "attention_mask" in tokenized: - nonpad_tokens = tokenized["attention_mask"].sum().item() - else: - pad_token_id = tokenizer.pad_token_id - if pad_token_id is not None: - nonpad_tokens = (tokenized["input_ids"] != pad_token_id).sum().item() - else: - nonpad_tokens = tokenized["input_ids"].numel() - tokenized = {k: v.to(self.device) for k, v in tokenized.items()} - return tokenized, nonpad_tokens - - @torch.inference_mode() - def _time(self, model: torch.nn.Module, tokenizer: object, batch_size: int, min_length: int, max_length: int) -> Tuple[float, int]: - model = model.to(self.device).eval() - set_seed(42) - min_dynamic_warmup_batches = self.warmup_batches - max_dynamic_warmup_batches = self.warmup_batches * 10 - stability_window = 3 - relative_stability_tolerance = 0.10 - - def synchronize(): - if self.device.type == "cuda": - torch.cuda.synchronize() - - def run_one_batch() -> int: - tokenized, nonpad_tokens = self._make_batch(tokenizer, batch_size, min_length, max_length) - _ = model(**tokenized) - return nonpad_tokens - - def time_batches(num_batches: int, message: str): - processed_tokens = 0 - synchronize() - start_time = time.time() - for _ in tqdm(range(num_batches), desc=message, leave=False): - processed_tokens += run_one_batch() - synchronize() - end_time = time.time() - return end_time - start_time, processed_tokens - - # Two torch.compile compatibility fixes for flex attention: - # - # 1. create_block_mask must NOT run inside torch.compile (per PyTorch docs). - # Wrap get_attention_mask and E1's block mask helpers with torch.compiler.disable - # so they execute eagerly even when the outer model is compiled. - # - # 2. The modeling files internally compile flex_attention via _get_flex_attention_fn(), - # caching the result in a module-level _compiled_flex_attention global. When we - # also wrap the full model in torch.compile(), the inductor hits double-compilation. - # Set the debug flag so _get_flex_attention_fn returns raw flex_attention, and - # clear any already-cached compiled version. - import sys - for mod in list(sys.modules.values()): - for fn_name in ("get_attention_mask", "create_block_causal_mask_optimized", "create_within_seq_block_mask"): - fn = getattr(mod, fn_name, None) - if fn is not None and callable(fn): - setattr(mod, fn_name, torch.compiler.disable(fn)) - if hasattr(mod, "_compiled_flex_attention"): - mod._compiled_flex_attention = None - flex_mod = getattr(torch.nn.attention, "flex_attention", None) - if flex_mod is not None: - flex_mod._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True - - # Eager pass to warm internal caches (rotary _seq_len_cached, flex block masks) - # before compile. Without this, dynamo recompiles when cache state changes. - _ = run_one_batch() - synchronize() - - torch._dynamo.reset() - model = torch.compile(model) - - warmup_latencies = [] - for warmup_idx in tqdm(range(max_dynamic_warmup_batches), desc="Warmup", leave=False): - synchronize() - warmup_start = time.time() - _ = run_one_batch() - synchronize() - warmup_latency = time.time() - warmup_start - warmup_latencies.append(warmup_latency) - - if warmup_idx + 1 < min_dynamic_warmup_batches: - continue - if len(warmup_latencies) < 2 * stability_window: - continue - - previous_window = warmup_latencies[-2 * stability_window:-stability_window] - current_window = warmup_latencies[-stability_window:] - previous_mean = sum(previous_window) / stability_window - current_mean = sum(current_window) / stability_window - assert previous_mean > 0.0, "Warmup latency mean should be positive." - relative_change = abs(current_mean - previous_mean) / previous_mean - if relative_change <= relative_stability_tolerance: - break - - time_taken, timed_tokens_sum = time_batches(self.timed_batches, "Timed") - if self.device.type == "cuda": - torch.cuda.empty_cache() - return time_taken, timed_tokens_sum - - def evaluate(self, model_path: str, batch_sizes: List[int], min_length: int, sequence_lengths: List[int], backends: List[str]) -> Dict[str, Dict[Tuple[int, int], Dict[str, float]]]: - results = {backend: {} for backend in backends} - - original_model = self._load_model(model_path) - tokenizer = original_model.tokenizer - - for backend in backends: - print(f"Benchmarking {model_path} with backend={backend}") - try: - backend_model = copy.deepcopy(original_model) - backend_model.attn_backend = backend - except AssertionError as error: - print(f"Skipping backend '{backend}' for {model_path}: {error}") - continue - - for bs in batch_sizes: - for max_length in sequence_lengths: - model_copy = copy.deepcopy(backend_model) - time_taken, tokens = self._time( - model_copy, - tokenizer, - bs, - min_length, - max_length, - ) - results[backend][(bs, max_length)] = {"time": time_taken, "tokens": tokens} - - original_model.cpu() - del original_model - if self.device.type == "cuda": - torch.cuda.empty_cache() - return results - - -def save_structured_results(all_results: Dict[str, Dict], output_dir: str) -> None: - """Save throughput results as JSON and CSV files. - - Each row contains: model, backend, batch_size, seq_len, tokens_per_sec, - total_tokens, elapsed_sec. - """ - import csv - import json - - rows = [] - for model_path, results in all_results.items(): - model_name = Path(model_path).name - for backend in sorted(results.keys()): - for (bs, max_length), entry in results[backend].items(): - time_taken = entry["time"] - nonpad_tokens = entry["tokens"] - tokens_per_sec = nonpad_tokens / time_taken if time_taken > 0 else 0.0 - rows.append({ - "model": model_name, - "model_path": model_path, - "backend": backend, - "batch_size": bs, - "seq_len": max_length, - "tokens_per_sec": round(tokens_per_sec, 2), - "total_tokens": nonpad_tokens, - "elapsed_sec": round(time_taken, 4), - }) - - output_path = Path(output_dir) - - json_path = output_path / "throughput_results.json" - with open(json_path, "w") as f: - json.dump(rows, f, indent=2) - print(f"JSON results saved to {json_path}") - - if rows: - csv_path = output_path / "throughput_results.csv" - with open(csv_path, "w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=rows[0].keys()) - writer.writeheader() - writer.writerows(rows) - print(f"CSV results saved to {csv_path}") - - -def plot_results(all_results: Dict[str, Dict], output_path: str) -> None: - sns.set_theme(style="whitegrid") - plot_data = [] - - for model_path, results in all_results.items(): - model_name = Path(model_path).name - for backend in sorted(results.keys()): - for (bs, max_length), entry in results[backend].items(): - time_taken = entry["time"] - nonpad_tokens = entry["tokens"] - tokens_per_sec = nonpad_tokens / time_taken if time_taken > 0 else 0.0 - plot_data.append( - { - "Model": model_name, - "Backend": backend, - "Batch": bs, - "SeqLen": max_length, - "TokensPerSec": tokens_per_sec, - "NonPadTokens": nonpad_tokens, - "Seconds": time_taken, - } - ) - - if not plot_data: - return - - plot_df = pd.DataFrame(plot_data) - sequence_lengths = sorted(plot_df["SeqLen"].dropna().unique().tolist()) - - plot = sns.relplot( - data=plot_df, - x="SeqLen", - y="TokensPerSec", - hue="Backend", - style="Batch", - kind="line", - marker="o", - dashes=False, - col="Model", - col_wrap=1, - height=4.5, - aspect=1.5, - facet_kws={"sharey": False}, - ) - plot.set_titles("{col_name}") - plot.set(xticks=sequence_lengths) - plot.set_axis_labels("Sequence length", "Non-pad tokens/s") - plot.figure.suptitle("Throughput comparison by model") - plot.tight_layout() - plot.figure.subplots_adjust(top=0.93, right=0.95, bottom=0.06) - plot.add_legend(title="Backend / Batch") - plt.savefig(output_path, dpi=300) - print(f"Results saved to {output_path}") - - -if __name__ == "__main__": - # On Windows, use "%cd%" instead of "${PWD}" to get the current working directory: - # docker run --gpus all -v "%cd%":/workspace fastplms python -m testing.throughput - # On Linux/macOS, keep using ${PWD}: - # docker run --gpus all -v ${PWD}:/workspace fastplms python -m testing.throughput - parser = argparse.ArgumentParser() - parser.add_argument("--hf_token", type=str, default=None) - parser.add_argument( - "--model_paths", - nargs="+", - default=["Synthyra/ESM2-8M", "Synthyra/ESMplusplus_small"], - ) - parser.add_argument("--batch_sizes", nargs="+", type=int, default=[2, 4, 8]) - parser.add_argument("--sequence_lengths", nargs="+", type=int, default=[64, 128, 256, 512, 1024, 2048]) - parser.add_argument("--backends", nargs="+", choices=SUPPORTED_BACKENDS, default=list(SUPPORTED_BACKENDS)) - parser.add_argument("--min_length", type=int, default=32) - parser.add_argument("--warmup_batches", type=int, default=10) - parser.add_argument("--timed_batches", type=int, default=100) - parser.add_argument("--output_path", type=str, default="throughput_comparison.png") - args = parser.parse_args() - - if args.hf_token: - from huggingface_hub import login - - login(token=args.hf_token) - - checker = ThroughputChecker(warmup_batches=args.warmup_batches, timed_batches=args.timed_batches) - - all_results = {} - for model_path in args.model_paths: - all_results[model_path] = checker.evaluate( - model_path, - args.batch_sizes, - min_length=args.min_length, - sequence_lengths=args.sequence_lengths, - backends=args.backends, - ) - - output_dir = str(Path(args.output_path).parent) - save_structured_results(all_results, output_dir) - plot_results(all_results, args.output_path) diff --git a/testing/__init__.py b/tests/__init__.py similarity index 100% rename from testing/__init__.py rename to tests/__init__.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..59e0694 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,225 @@ +import builtins +import contextlib +import importlib.util +import random +import sys +from pathlib import Path + + +def _cpu_contract_requested() -> bool: + cpu_root = (Path(__file__).resolve().parent / "cpu").resolve() + for raw_argument in sys.argv[1:]: + if raw_argument.startswith("-"): + continue + candidate_text = raw_argument.split("::", maxsplit=1)[0] + try: + candidate = Path(candidate_text).resolve() + except (OSError, ValueError): + continue + if candidate == cpu_root or cpu_root in candidate.parents: + return True + return False + + +def _bootstrap_cpu_contract() -> None: + if not _cpu_contract_requested() or getattr( + builtins, + "_fastplms_cpu_process_bootstrapped", + False, + ): + return + bootstrap = Path(__file__).resolve().parent / "cpu" / "bootstrap" / "sitecustomize.py" + spec = importlib.util.spec_from_file_location("_fastplms_cpu_sitecustomize", bootstrap) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load CPU contract bootstrap: {bootstrap}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + +_bootstrap_cpu_contract() + +# The hermetic CPU bootstrap must run before these import-time initializers. +import pytest # noqa: E402 +import torch # noqa: E402 + +from fastplms.registry import ModelSpec, get_model_registry # noqa: E402 + + +def pytest_configure(config): + config.addinivalue_line("markers", "gpu: requires CUDA GPU") + config.addinivalue_line("markers", "slow: loads two models simultaneously (compliance tests)") + config.addinivalue_line("markers", "large: requires 24+ GB VRAM (3B parameter models)") + config.addinivalue_line( + "markers", "structure: structure prediction models (Boltz2, ESMFold, ESMFold2)" + ) + + +CANONICAL_AAS = "ACDEFGHIKLMNPQRSTVWY" +SEED = 42 +DEFAULT_BATCH_SIZE = 4 +MAX_EMBED_LEN = 128 + + +@contextlib.contextmanager +def strict_fp32_matmul(): + """Temporarily disable TF32 for fp32 numerical parity checks.""" + try: + old_fp32_precision = torch.backends.fp32_precision + old_matmul_precision = torch.backends.cuda.matmul.fp32_precision + old_cudnn_precision = torch.backends.cudnn.fp32_precision + except AttributeError: + old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + old_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + old_matmul_precision = torch.get_float32_matmul_precision() + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32 + torch.backends.cudnn.allow_tf32 = old_cudnn_tf32 + torch.set_float32_matmul_precision(old_matmul_precision) + return + + torch.backends.fp32_precision = "ieee" + torch.backends.cuda.matmul.fp32_precision = "ieee" + torch.backends.cudnn.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.fp32_precision = old_fp32_precision + torch.backends.cuda.matmul.fp32_precision = old_matmul_precision + torch.backends.cudnn.fp32_precision = old_cudnn_precision + + +# The package manifest is the sole checkpoint registry. These dictionaries keep +# the legacy helper surface while deriving every value from typed ModelSpec data. +def _legacy_entry(model: ModelSpec) -> dict: + return { + "fast_path": model.fast.repo_id, + "fast_revision": model.fast.revision, + "official_path": model.official.repo_id, + "official_revision": model.official.revision, + "load_official": model.family.reference_adapter, + "model_type": model.family.architecture, + "uses_tokenizer": model.family.tokenizer_mode == "tokenizer", + "size_category": model.size_category, + "attention": model.family.attention, + "dtypes": model.family.dtypes, + "precisions": model.family.precisions, + "state_transform": model.family.state_transform, + "reference_container": model.family.reference_container, + } + + +_TYPED_REGISTRY = get_model_registry() +FULL_MODEL_REGISTRY: dict[str, dict] = { + model.id: _legacy_entry(model) + for model in _TYPED_REGISTRY.values() + if model.family.tokenizer_mode != "structure" +} +STRUCTURE_MODEL_REGISTRY: dict[str, dict] = { + model.id: _legacy_entry(model) + for model in _TYPED_REGISTRY.values() + if model.family.tokenizer_mode == "structure" +} +MODEL_REGISTRY: dict[str, dict] = { + model.family.architecture.lower(): _legacy_entry(model) + for model in _TYPED_REGISTRY.values() + if model.is_deep_reference and model.family.tokenizer_mode != "structure" +} +BACKENDS = ( + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", +) + + +def get_models_by_size(*categories: str) -> dict[str, dict]: + return {k: v for k, v in FULL_MODEL_REGISTRY.items() if v["size_category"] in categories} + + +# Pre-built key lists by size category +SMALL_MODEL_KEYS = list(get_models_by_size("small").keys()) +MEDIUM_MODEL_KEYS = list(get_models_by_size("small", "medium").keys()) +LARGE_MODEL_KEYS = list(get_models_by_size("large").keys()) +XLARGE_MODEL_KEYS = list(get_models_by_size("xlarge").keys()) +ALL_FULL_MODEL_KEYS = list(FULL_MODEL_REGISTRY.keys()) +SEQUENCE_MODEL_KEYS = [ + k for k in ALL_FULL_MODEL_KEYS if FULL_MODEL_REGISTRY[k]["size_category"] != "structure" +] +STRUCTURE_MODEL_KEYS = list(STRUCTURE_MODEL_REGISTRY.keys()) + + +def mark_by_size( + keys: list[str], + registry: dict[str, dict], + extra_marks: list | None = None, +) -> list: + """Return pytest.param list with appropriate markers based on size_category.""" + params = [] + for k in keys: + marks = list(extra_marks or []) + if registry[k]["size_category"] == "xlarge": + marks.append(pytest.mark.large) + elif registry[k]["size_category"] in ("large", "medium"): + marks.append(pytest.mark.slow) + params.append(pytest.param(k, marks=marks)) + return params + + +def tokenize_batch( + model, + model_key: str, + sequences: list[str], + device: torch.device, + registry: dict[str, dict] | None = None, +) -> dict[str, torch.Tensor]: + """Tokenize a batch of sequences, handling E1's sequence mode. + + Shared helper used across multiple test files to avoid duplication. + """ + if registry is None: + registry = FULL_MODEL_REGISTRY + config = registry[model_key] if model_key in registry else MODEL_REGISTRY[model_key] + + if config["model_type"] == "E1": + batch = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) + return { + "input_ids": batch["input_ids"], + "within_seq_position_ids": batch["within_seq_position_ids"], + "global_position_ids": batch["global_position_ids"], + "sequence_ids": batch["sequence_ids"], + "attention_mask": (batch["sequence_ids"] != -1).long(), + } + tokenizer = model.tokenizer + tokenized = tokenizer(sequences, return_tensors="pt", padding=True) + return {k: v.to(device) for k, v in tokenized.items()} + + +def add_model_specific_inputs( + model_inputs: dict[str, torch.Tensor], + model_type: str, +) -> dict[str, torch.Tensor]: + """Add model-specific extra inputs (e.g. sequence_id for ESMC).""" + if model_type == "ESMC": + # model_inputs["sequence_id"]: (b, l) + model_inputs["sequence_id"] = model_inputs["attention_mask"].to(dtype=torch.bool) + return model_inputs + + +def random_sequences(n: int, min_len: int = 8, max_len: int = 64) -> list[str]: + return [ + "M" + "".join(random.choices(CANONICAL_AAS, k=random.randint(min_len, max_len))) + for _ in range(n) + ] + + +def random_sequences_fixed_len(n: int, length: int = 64) -> list[str]: + return ["M" + "".join(random.choices(CANONICAL_AAS, k=length - 1)) for _ in range(n)] + + +def get_device() -> torch.device: + return torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/cpu/bootstrap/sitecustomize.py b/tests/cpu/bootstrap/sitecustomize.py new file mode 100644 index 0000000..e6c5e5e --- /dev/null +++ b/tests/cpu/bootstrap/sitecustomize.py @@ -0,0 +1,309 @@ +"""Python-startup policy for the hermetic CPU contract lane. + +This module is intentionally standard-library-only. The CPU workflow adds its +directory to ``PYTHONPATH``, so Python imports it before pytest can import the +repository-level ``tests/conftest.py`` (and therefore before Torch or the model +registry). Child Python processes inherit the same startup policy. +""" + +from __future__ import annotations + +import builtins +import io +import os +import shlex +import socket +import subprocess +import tempfile +from pathlib import Path +from typing import Any + + +_CACHE_ENVIRONMENT = { + "HF_HOME": "huggingface", + "HF_HUB_CACHE": "huggingface/hub", + "HUGGINGFACE_HUB_CACHE": "huggingface/hub", + "TRANSFORMERS_CACHE": "huggingface/transformers", + "HF_DATASETS_CACHE": "huggingface/datasets", + "TORCH_HOME": "torch", + "TORCH_EXTENSIONS_DIR": "torch-extensions", + "TORCHINDUCTOR_CACHE_DIR": "torch-inductor", + "TRITON_CACHE_DIR": "triton", + "XDG_CACHE_HOME": "xdg", +} +_FIXED_ENVIRONMENT = { + "CUDA_VISIBLE_DEVICES": "", + "DO_NOT_TRACK": "1", + "HF_DATASETS_OFFLINE": "1", + "HF_HUB_DISABLE_TELEMETRY": "1", + "HF_HUB_OFFLINE": "1", + "MKL_NUM_THREADS": "1", + "NUMEXPR_NUM_THREADS": "1", + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "PYTEST_XDIST_AUTO_NUM_WORKERS": "4", + "TOKENIZERS_PARALLELISM": "false", + "TRANSFORMERS_OFFLINE": "1", + "VECLIB_MAXIMUM_THREADS": "1", +} +_WORKSPACE = Path(__file__).resolve().parents[3] +_FORBIDDEN_READ_ROOTS = tuple( + path.resolve() + for path in ( + _WORKSPACE / "vendor" / "upstream", + _WORKSPACE / ".git" / "modules", + _WORKSPACE / "official", + ) +) +_CHECKPOINT_SUFFIXES = frozenset({".bin", ".ckpt", ".pt", ".pth", ".safetensors"}) +_FORBIDDEN_CONTAINER_EXECUTABLES = frozenset( + {"buildx", "docker", "docker-compose", "podman"} +) +_SHELL_EXECUTABLES = frozenset({"bash", "cmd", "dash", "powershell", "pwsh", "sh", "zsh"}) +_COMMAND_WRAPPERS = frozenset({"command", "env", "nohup", "sudo"}) + + +def _network_blocked(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("Network access is forbidden in tests/cpu") + + +def _dir_fd_path(file: object, dir_fd: int | None) -> object: + if dir_fd is None or not isinstance(file, (str, bytes, os.PathLike)): + return file + decoded = os.fsdecode(file) + if os.path.isabs(decoded): + return decoded + for descriptor_root in ("/proc/self/fd", "/dev/fd"): + try: + directory = Path(os.readlink(f"{descriptor_root}/{dir_fd}")) + except OSError: + continue + return directory / decoded + raise RuntimeError(f"CPU contracts could not resolve directory descriptor {dir_fd}") + + +def _assert_portable_path(file: object, *, dir_fd: int | None = None) -> None: + if not isinstance(file, (str, bytes, os.PathLike)): + return + try: + resolved = Path(_dir_fd_path(file, dir_fd)).resolve() + except (OSError, TypeError, ValueError): + return + if any(resolved == root or root in resolved.parents for root in _FORBIDDEN_READ_ROOTS): + raise RuntimeError(f"CPU contracts may not access submodule/reference path: {resolved}") + if ( + (resolved == _WORKSPACE or _WORKSPACE in resolved.parents) + and ( + resolved.suffix.lower() in _CHECKPOINT_SUFFIXES + or resolved.name.endswith(".safetensors.index.json") + ) + ): + raise RuntimeError(f"CPU contracts may not access checkpoint path: {resolved}") + + +def _executable_basename(value: object) -> str: + if not isinstance(value, (str, bytes, os.PathLike)): + return "" + name = Path(os.fsdecode(value)).name.lower() + return name.removesuffix(".exe") + + +def _command_tokens(command: object) -> list[str] | None: + if isinstance(command, (str, bytes, os.PathLike)): + try: + return shlex.split(os.fsdecode(command), posix=os.name != "nt") + except ValueError: + return None + if isinstance(command, (list, tuple)): + try: + return [os.fsdecode(value) for value in command] + except TypeError: + return None + return None + + +def _command_starts(tokens: list[str]) -> list[list[str]]: + commands: list[list[str]] = [] + current: list[str] = [] + for token in tokens: + if token in {"&", "&&", ";", "|", "||"}: + if current: + commands.append(current) + current = [] + else: + current.append(token) + if current: + commands.append(current) + return commands + + +def _unwrap_command(tokens: list[str]) -> list[str]: + remaining = list(tokens) + while remaining and _executable_basename(remaining[0]) in _COMMAND_WRAPPERS: + wrapper = _executable_basename(remaining.pop(0)) + while remaining and ( + remaining[0].startswith("-") + or (wrapper == "env" and "=" in remaining[0]) + ): + remaining.pop(0) + return remaining + + +def _forbidden_container_command(command: object, *, executable: object = None) -> str | None: + explicit = _executable_basename(executable) + if explicit in _FORBIDDEN_CONTAINER_EXECUTABLES: + return explicit + tokens = _command_tokens(command) + if tokens is None: + return None + for candidate in _command_starts(tokens): + unwrapped = _unwrap_command(candidate) + if not unwrapped: + continue + executable_name = _executable_basename(unwrapped[0]) + if executable_name in _FORBIDDEN_CONTAINER_EXECUTABLES: + return executable_name + if executable_name in _SHELL_EXECUTABLES: + lowered = [token.lower() for token in unwrapped] + for flag in ("-c", "-ec", "-lc", "/c", "-command"): + try: + index = lowered.index(flag) + except ValueError: + continue + if index + 1 < len(unwrapped): + nested = _forbidden_container_command(unwrapped[index + 1]) + if nested is not None: + return nested + return None + + +def _assert_portable_spawn(command: object, *, executable: object = None) -> None: + forbidden = _forbidden_container_command(command, executable=executable) + if forbidden is not None: + raise RuntimeError( + f"Container execution is forbidden in tests/cpu: {forbidden}" + ) + + +def _install_spawn_guards() -> None: + if getattr(builtins, "_fastplms_cpu_spawn_guard", False): + return + original_popen = subprocess.Popen + original_system = os.system + + class GuardedPopen(original_popen): # type: ignore[misc, valid-type] + def __init__(self, args: Any, *popen_args: Any, **popen_kwargs: Any) -> None: + _assert_portable_spawn(args, executable=popen_kwargs.get("executable")) + super().__init__(args, *popen_args, **popen_kwargs) + + def guarded_system(command: object) -> int: + _assert_portable_spawn(command) + return original_system(command) + + subprocess.Popen = GuardedPopen # type: ignore[assignment] + os.system = guarded_system # type: ignore[assignment] + + for name in ( + "spawnl", + "spawnle", + "spawnlp", + "spawnlpe", + "spawnv", + "spawnve", + "spawnvp", + "spawnvpe", + ): + original = getattr(os, name, None) + if original is None: + continue + + def guarded_spawn(*args: Any, _original: Any = original, **kwargs: Any) -> Any: + command = args[1] if len(args) > 1 else kwargs.get("file") + _assert_portable_spawn(command, executable=command) + return _original(*args, **kwargs) + + setattr(os, name, guarded_spawn) + builtins.__dict__["_fastplms_cpu_spawn_guard"] = True + + +def _install_open_guards() -> None: + if getattr(builtins, "_fastplms_cpu_open_guard", False): + return + original_builtin_open = builtins.open + original_io_open = io.open + original_os_open = os.open + + def guarded_builtin_open(file: object, *args: Any, **kwargs: Any) -> Any: + _assert_portable_path(file) + return original_builtin_open(file, *args, **kwargs) + + def guarded_io_open(file: object, *args: Any, **kwargs: Any) -> Any: + _assert_portable_path(file) + return original_io_open(file, *args, **kwargs) + + def guarded_os_open(file: object, *args: Any, **kwargs: Any) -> int: + _assert_portable_path(file, dir_fd=kwargs.get("dir_fd")) + return original_os_open(file, *args, **kwargs) + + builtins.open = guarded_builtin_open + io.open = guarded_io_open + os.open = guarded_os_open # type: ignore[assignment] + builtins.__dict__["_fastplms_cpu_open_guard"] = True + + +def _install_hub_guards() -> None: + # Hub is present in the required environment. Keeping this import optional + # lets the startup bootstrap remain harmless while an environment is built. + try: + import huggingface_hub + import huggingface_hub._snapshot_download + import huggingface_hub.file_download + except ImportError: + return + huggingface_hub.hf_hub_download = _network_blocked # type: ignore[assignment] + huggingface_hub.snapshot_download = _network_blocked # type: ignore[assignment] + huggingface_hub.file_download.hf_hub_download = ( # type: ignore[assignment] + _network_blocked + ) + huggingface_hub._snapshot_download.snapshot_download = ( # type: ignore[assignment] + _network_blocked + ) + huggingface_hub.file_download.http_get = _network_blocked # type: ignore[assignment] + + +def _install() -> None: + bootstrap_root = str(Path(__file__).resolve().parent) + python_path = os.environ.get("PYTHONPATH", "") + python_path_entries = [entry for entry in python_path.split(os.pathsep) if entry] + if bootstrap_root not in python_path_entries: + os.environ["PYTHONPATH"] = os.pathsep.join((bootstrap_root, *python_path_entries)) + cache_root_value = os.environ.get("FASTPLMS_CPU_CACHE_ROOT") + if cache_root_value is None: + cache_root = Path(tempfile.mkdtemp(prefix="fastplms-cpu-contract-cache-")).resolve() + os.environ["FASTPLMS_CPU_CACHE_ROOT"] = str(cache_root) + os.environ["FASTPLMS_CPU_CACHE_STARTED_EMPTY"] = "1" + else: + cache_root = Path(cache_root_value).resolve() + cache_root.mkdir(parents=True, exist_ok=True) + for name, relative in _CACHE_ENVIRONMENT.items(): + path = cache_root / relative + path.mkdir(parents=True, exist_ok=True) + os.environ[name] = str(path) + os.environ.update(_FIXED_ENVIRONMENT) + os.environ["FASTPLMS_CPU_BOOTSTRAPPED"] = "1" + + socket.create_connection = _network_blocked # type: ignore[assignment] + socket.getaddrinfo = _network_blocked # type: ignore[assignment] + socket.socket.connect = _network_blocked # type: ignore[assignment] + socket.socket.connect_ex = _network_blocked # type: ignore[assignment] + socket.socket.sendto = _network_blocked # type: ignore[assignment] + if hasattr(socket.socket, "sendmsg"): + socket.socket.sendmsg = _network_blocked # type: ignore[assignment] + _install_open_guards() + _install_spawn_guards() + _install_hub_guards() + builtins.__dict__["_fastplms_cpu_assert_portable_path"] = _assert_portable_path + builtins.__dict__["_fastplms_cpu_process_bootstrapped"] = True + + +_install() diff --git a/tests/cpu/conftest.py b/tests/cpu/conftest.py new file mode 100644 index 0000000..d06001e --- /dev/null +++ b/tests/cpu/conftest.py @@ -0,0 +1,512 @@ +"""Hermetic policy for the mandatory CPU contract lane.""" + +from __future__ import annotations + +import builtins +import hashlib +import json +import os +import platform +import random +import signal +import time +import warnings +from collections.abc import Iterator +from importlib.metadata import version +from pathlib import Path +from typing import Any + +from tests.cpu.resource_telemetry import ( + ConcurrentProcessTreeSampler, + MemoryEvidenceError, + aggregate_process_memory, + capture_process_memory, + select_concurrent_memory_gate, +) + + +if os.environ.get("FASTPLMS_CPU_BOOTSTRAPPED") != "1" or not getattr( + builtins, + "_fastplms_cpu_process_bootstrapped", + False, +): + raise RuntimeError("tests/cpu requires the Python-startup hermetic bootstrap") +_CPU_CACHE_ROOT = Path(os.environ["FASTPLMS_CPU_CACHE_ROOT"]).resolve() +_CACHE_NAMES = tuple( + sorted( + { + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "TRANSFORMERS_CACHE", + "HF_DATASETS_CACHE", + "TORCH_HOME", + "TORCH_EXTENSIONS_DIR", + "TORCHINDUCTOR_CACHE_DIR", + "TRITON_CACHE_DIR", + "XDG_CACHE_HOME", + } + ) +) + +import numpy as np # noqa: E402 +import pytest # noqa: E402 +import torch # noqa: E402 + + +def _install_checkpoint_loader_guards() -> None: + """Block native checkpoint loaders that can bypass Python's open wrappers.""" + + if getattr(builtins, "_fastplms_cpu_checkpoint_loader_guard", False): + return + assert_path = getattr(builtins, "_fastplms_cpu_assert_portable_path", None) + if not callable(assert_path): + raise RuntimeError("CPU checkpoint path guard was not installed at Python startup") + original_torch_load = torch.load + + def guarded_torch_load(file: object, *args: Any, **kwargs: Any) -> Any: + assert_path(file) + return original_torch_load(file, *args, **kwargs) + + torch.load = guarded_torch_load # type: ignore[assignment] + try: + import safetensors + import safetensors.torch + except ImportError: + pass + else: + original_safe_open = safetensors.safe_open + original_load_file = safetensors.torch.load_file + + def guarded_safe_open(filename: object, *args: Any, **kwargs: Any) -> Any: + assert_path(filename) + return original_safe_open(filename, *args, **kwargs) + + def guarded_load_file(filename: object, *args: Any, **kwargs: Any) -> Any: + assert_path(filename) + return original_load_file(filename, *args, **kwargs) + + safetensors.safe_open = guarded_safe_open # type: ignore[assignment] + safetensors.torch.safe_open = guarded_safe_open # type: ignore[assignment] + safetensors.torch.load_file = guarded_load_file # type: ignore[assignment] + builtins.__dict__["_fastplms_cpu_checkpoint_loader_guard"] = True + + +_install_checkpoint_loader_guards() +torch.set_num_interop_threads(1) + +_FORBIDDEN_MARKERS = { + "artifact", + "benchmark", + "checkpoint", + "compliance", + "gpu", + "large", + "network", + "packaging", + "reference", + "slow", +} +_MAX_TEST_SECONDS = 10.0 +_MAX_SUITE_SECONDS = 300.0 +_MAX_SUITE_RSS_BYTES = 4 * 1024**3 +_MEMORY_SAMPLE_INTERVAL_SECONDS = 0.05 +_WORKSPACE = Path(__file__).resolve().parents[2] +_TELEMETRY_PATH = Path( + os.environ.get( + "FASTPLMS_CPU_TELEMETRY_PATH", + str(_WORKSPACE / "artifacts" / "telemetry" / "cpu-contract.json"), + ) +).resolve() + + +class CpuContractTimeoutError(TimeoutError): + """Raised while a CPU contract still owns the per-test wall-clock budget.""" + + +def _timeout_test(_signum: int, _frame: object) -> None: + raise CpuContractTimeoutError( + f"CPU contract exceeded its {_MAX_TEST_SECONDS:.0f}s execution budget." + ) + + +def _fail_session(session: pytest.Session, message: str) -> None: + failures = getattr(session.config, "_fastplms_gate_failures", None) + if failures is None: + failures = [] + session.config._fastplms_gate_failures = failures + if message not in failures: + failures.append(message) + session.exitstatus = pytest.ExitCode.TESTS_FAILED + warnings.warn(pytest.PytestWarning(message), stacklevel=2) + + +def _enforce_concurrent_memory_budget( + session: pytest.Session, + memory_gate: dict[str, object], +) -> None: + selected_peak = memory_gate.get("peak_bytes") + if not isinstance(selected_peak, int) or isinstance(selected_peak, bool): + _fail_session(session, "CPU contract concurrent memory peak is unavailable") + return + if selected_peak > _MAX_SUITE_RSS_BYTES: + _fail_session( + session, + "CPU contract concurrent memory budget exceeded " + f"({memory_gate.get('metric', 'unavailable')}): " + f"{selected_peak / 1024**3:.2f} GiB > 4 GiB", + ) + + +def _cache_snapshot() -> dict[str, object]: + records: list[dict[str, object]] = [] + seen: set[Path] = set() + for name in _CACHE_NAMES: + root = Path(os.environ[name]).resolve() + if root in seen: + continue + seen.add(root) + files = sorted(path for path in root.rglob("*") if path.is_file()) + digest = hashlib.sha256() + total_bytes = 0 + for path in files: + size = path.stat().st_size + total_bytes += size + relative = path.relative_to(_CPU_CACHE_ROOT).as_posix().encode("utf-8") + digest.update(len(relative).to_bytes(8, "big")) + digest.update(relative) + digest.update(size.to_bytes(8, "big")) + records.append( + { + "root": root.relative_to(_CPU_CACHE_ROOT).as_posix(), + "files": len(files), + "bytes": total_bytes, + "inventory_sha256": digest.hexdigest(), + } + ) + return {"roots": sorted(records, key=lambda record: str(record["root"]))} + + +def _atomic_json(path: Path, payload: dict[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def pytest_sessionstart(session: pytest.Session) -> None: + session.config._fastplms_cpu_started = time.monotonic() # type: ignore[attr-defined] + session.config._fastplms_cpu_durations = [] # type: ignore[attr-defined] + session.config._fastplms_cpu_phase_durations = {} # type: ignore[attr-defined] + session.config._fastplms_cpu_outcomes = {} # type: ignore[attr-defined] + session.config._fastplms_gate_failures = [] # type: ignore[attr-defined] + session.config._fastplms_cache_before = _cache_snapshot() # type: ignore[attr-defined] + if not hasattr(session.config, "workerinput"): + session.config._fastplms_worker_memory = {} # type: ignore[attr-defined] + session.config._fastplms_worker_memory_errors = [] # type: ignore[attr-defined] + if ( + not hasattr(session.config, "workerinput") + and os.environ.get("FASTPLMS_CPU_CACHE_STARTED_EMPTY") != "1" + ): + raise pytest.UsageError( + "CPU contract bootstrap did not attest an empty fresh cache root" + ) + if not hasattr(session.config, "workerinput"): + sampler = ConcurrentProcessTreeSampler( + root_pids=(os.getpid(),), + sample_interval_seconds=_MEMORY_SAMPLE_INTERVAL_SECONDS, + ) + try: + sampler.start() + except MemoryEvidenceError as error: + raise pytest.UsageError( + f"CPU contract concurrent memory sampler is unavailable: {error}" + ) from error + session.config._fastplms_concurrent_memory_sampler = sampler # type: ignore[attr-defined] + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + del exitstatus + started = session.config._fastplms_cpu_started # type: ignore[attr-defined] + elapsed = time.monotonic() - started + is_worker = hasattr(session.config, "workerinput") + worker_output = getattr(session.config, "workeroutput", None) + if isinstance(worker_output, dict): + worker_input = session.config.workerinput # type: ignore[attr-defined] + worker_id = str(worker_input.get("workerid", "")) + try: + worker_output["fastplms_memory_evidence"] = capture_process_memory( + worker_id, + role="worker", + ) + except MemoryEvidenceError as error: + worker_output["fastplms_memory_error"] = str(error) + _fail_session(session, f"CPU contract memory evidence unavailable: {error}") + worker_output["fastplms_test_durations"] = list( + session.config._fastplms_cpu_durations # type: ignore[attr-defined] + ) + if not is_worker and elapsed > _MAX_SUITE_SECONDS: + _fail_session( + session, + f"CPU contract budget exceeded: {elapsed:.2f}s > {_MAX_SUITE_SECONDS:.0f}s", + ) + if not is_worker: + worker_memory = getattr(session.config, "_fastplms_worker_memory", {}) + diagnostic_errors = list( + getattr(session.config, "_fastplms_worker_memory_errors", []) + ) + expected_workers = _expected_xdist_workers(session.config) + if expected_workers != len(worker_memory): + diagnostic_errors.append( + "CPU contract expected " + f"{expected_workers} xdist worker memory records, received {len(worker_memory)}" + ) + try: + controller_memory = capture_process_memory("controller", role="controller") + temporal_upper_bound = aggregate_process_memory( + controller_memory, + worker_memory.values(), + ) + except MemoryEvidenceError as error: + diagnostic_errors.append(str(error)) + temporal_upper_bound = { + "available": False, + "errors": sorted(set(diagnostic_errors)), + } + if diagnostic_errors: + temporal_upper_bound["available"] = False + temporal_upper_bound["errors"] = sorted(set(diagnostic_errors)) + warnings.warn( + pytest.PytestWarning( + "CPU contract temporal-upper-bound diagnostic is incomplete: " + + "; ".join(sorted(set(diagnostic_errors))) + ), + stacklevel=2, + ) + + concurrent_errors: list[str] = [] + try: + sampler = ( # type: ignore[attr-defined] + session.config._fastplms_concurrent_memory_sampler + ) + concurrent_memory = sampler.stop() + memory_gate = select_concurrent_memory_gate(concurrent_memory) + except (AttributeError, MemoryEvidenceError) as error: + concurrent_errors.append(str(error)) + concurrent_memory = { + "available": False, + "errors": sorted(set(concurrent_errors)), + } + memory_gate = { + "available": False, + "errors": sorted(set(concurrent_errors)), + } + memory_evidence = { + "available": not concurrent_errors, + "gate": memory_gate, + "concurrent_process_tree": concurrent_memory, + "temporal_upper_bound": temporal_upper_bound, + } + if concurrent_errors: + _fail_session( + session, + "CPU contract concurrent memory evidence unavailable: " + + "; ".join(concurrent_errors), + ) + selected_peak = memory_gate.get("peak_bytes") + if memory_gate.get("fallback_used") is True: + warnings.warn( + pytest.PytestWarning( + "CPU contract PSS evidence was incomplete for one or more " + "processes; enforcing the 4 GiB budget with the per-process " + "PSS/RSS hybrid: " + + "; ".join(memory_gate.get("fallback_reasons", [])) + ), + stacklevel=2, + ) + _enforce_concurrent_memory_budget(session, memory_gate) + durations = list( + getattr(session.config, "_fastplms_cpu_durations", []) + ) + list(getattr(session.config, "_fastplms_worker_durations", [])) + durations.sort(key=lambda record: str(record["nodeid"])) + cache_before = session.config._fastplms_cache_before # type: ignore[attr-defined] + _atomic_json( + _TELEMETRY_PATH, + { + "schema_version": 3, + "report": "fastplms-cpu-contract", + "source_revision": os.environ.get("GITHUB_SHA", "unbound-local-source"), + "runtime": { + "python": platform.python_version(), + "torch": version("torch"), + "transformers": version("transformers"), + "platform": platform.platform(), + }, + "budgets": { + "suite_seconds": _MAX_SUITE_SECONDS, + "test_seconds": _MAX_TEST_SECONDS, + "concurrent_physical_memory_bytes": _MAX_SUITE_RSS_BYTES, + }, + "observed": { + "suite_seconds": round(elapsed, 6), + "concurrent_peak_memory_bytes": selected_peak, + "concurrent_peak_memory_metric": memory_gate.get("metric"), + "peak_concurrent_rss_bytes": concurrent_memory.get( + "peak_concurrent_rss_bytes" + ), + "peak_concurrent_hybrid_bytes": concurrent_memory.get( + "peak_concurrent_hybrid_bytes" + ), + "peak_concurrent_pss_bytes": concurrent_memory.get( + "peak_concurrent_pss_bytes" + ), + "temporal_upper_bound_rss_bytes": temporal_upper_bound.get( + "temporal_upper_bound_rss_bytes" + ), + "memory": memory_evidence, + "tests": durations, + }, + "cache": { + "started_empty": True, + "before_collection": cache_before, + "after_session": _cache_snapshot(), + }, + "gate_failures": list( + getattr(session.config, "_fastplms_gate_failures", []) + ), + }, + ) + if getattr(session.config, "_fastplms_gate_failures", []): + # This assignment runs try-last so no normal session-finish hook can + # turn a resource-gate failure back into a successful process exit. + session.exitstatus = pytest.ExitCode.TESTS_FAILED + + +@pytest.hookimpl(optionalhook=True) +def pytest_testnodedown(node: Any, error: object) -> None: + worker_output = getattr(node, "workeroutput", {}) + worker_memory = getattr(node.config, "_fastplms_worker_memory", None) + memory_errors = getattr(node.config, "_fastplms_worker_memory_errors", None) + if worker_memory is None or memory_errors is None: + return + memory_record = worker_output.get("fastplms_memory_evidence") + if isinstance(memory_record, dict): + worker_id = str(memory_record.get("process_id", "")) + if not worker_id or worker_id in worker_memory: + memory_errors.append(f"duplicate or missing worker memory identity: {worker_id!r}") + else: + worker_memory[worker_id] = memory_record + else: + detail = worker_output.get("fastplms_memory_error") or error or "missing record" + memory_errors.append(f"xdist worker memory evidence unavailable: {detail}") + durations = getattr(node.config, "_fastplms_worker_durations", None) + if durations is None: + durations = [] + node.config._fastplms_worker_durations = durations + durations.extend(worker_output.get("fastplms_test_durations", [])) + + +def _expected_xdist_workers(config: pytest.Config) -> int: + raw_value = config.getoption("numprocesses", default=0) + if raw_value in (None, 0, "0"): + return 0 + if raw_value == "auto": + raw_value = os.environ.get("PYTEST_XDIST_AUTO_NUM_WORKERS", "") + try: + count = int(raw_value) + except (TypeError, ValueError) as error: + raise pytest.UsageError( + f"Cannot determine configured xdist worker count from {raw_value!r}" + ) from error + if count < 0: + raise pytest.UsageError(f"Invalid xdist worker count: {count}") + return count + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Make directory ownership positive and reject expensive resource markers.""" + + for item in items: + item.add_marker(pytest.mark.cpu_contract) + forbidden = sorted( + mark.name for mark in item.iter_markers() if mark.name in _FORBIDDEN_MARKERS + ) + if forbidden: + raise pytest.UsageError( + f"CPU contract {item.nodeid} carries forbidden markers: {forbidden}" + ) + if item.get_closest_marker("skip") or item.get_closest_marker("skipif"): + raise pytest.UsageError(f"CPU contract {item.nodeid} may not be skipped") + if item.get_closest_marker("xfail"): + raise pytest.UsageError(f"CPU contract {item.nodeid} may not be xfailed") + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Iterator[None]: + """Bound the complete setup/call/teardown protocol for one contract.""" + + del item, nextitem + if ( + not hasattr(signal, "SIGALRM") + or not hasattr(signal, "ITIMER_REAL") + or not hasattr(signal, "setitimer") + ): + yield + return + previous_handler = signal.getsignal(signal.SIGALRM) + signal.signal(signal.SIGALRM, _timeout_test) + previous_timer = signal.setitimer(signal.ITIMER_REAL, _MAX_TEST_SECONDS) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0.0) + signal.signal(signal.SIGALRM, previous_handler) + if previous_timer[0] > 0.0: + signal.setitimer(signal.ITIMER_REAL, *previous_timer) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item: pytest.Item, call: pytest.CallInfo[object]) -> Iterator[None]: + """Turn dynamic import/runtime skips into hard failures.""" + + outcome = yield + report = outcome.get_result() + if report.skipped: + report.outcome = "failed" + report.longrepr = "Mandatory CPU contracts may not skip at runtime." + phase_durations = item.config._fastplms_cpu_phase_durations # type: ignore[attr-defined] + total_duration = float(phase_durations.get(report.nodeid, 0.0)) + report.duration + phase_durations[report.nodeid] = total_duration + outcomes = item.config._fastplms_cpu_outcomes # type: ignore[attr-defined] + if report.outcome != "passed": + outcomes[report.nodeid] = report.outcome + if total_duration > _MAX_TEST_SECONDS: + report.outcome = "failed" + outcomes[report.nodeid] = "failed" + report.longrepr = ( + f"CPU contract exceeded its {_MAX_TEST_SECONDS:.0f}s budget: " + f"{total_duration:.2f}s across setup/call/teardown" + ) + if report.when == "teardown": + item.config._fastplms_cpu_durations.append( # type: ignore[attr-defined] + { + "nodeid": report.nodeid, + "seconds": round(total_duration, 6), + "outcome": outcomes.get(report.nodeid, "passed"), + } + ) + + +@pytest.fixture(autouse=True) +def _hermetic_cpu(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + """Deny network access and keep tiny numerical tests deterministic.""" + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + try: + yield + finally: + torch.set_num_threads(previous_threads) diff --git a/tests/cpu/resource_telemetry.py b/tests/cpu/resource_telemetry.py new file mode 100644 index 0000000..c9c34c7 --- /dev/null +++ b/tests/cpu/resource_telemetry.py @@ -0,0 +1,667 @@ +"""Linux process-tree memory accounting for the mandatory CPU gate.""" + +from __future__ import annotations + +import errno +import os +import threading +import time +from collections.abc import Iterable, Mapping +from pathlib import Path +from typing import Any + + +class MemoryEvidenceError(RuntimeError): + """Raised when process memory evidence is absent or internally inconsistent.""" + + +_TRANSIENT_PROC_ERRNOS = frozenset({errno.ENOENT, errno.ESRCH}) +_KIBIBYTE = 1024 + + +def _validated_pid(value: object) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise MemoryEvidenceError(f"invalid process identifier: {value!r}") + return value + + +def _read_proc_text(path: Path) -> str | None: + """Read one procfs file, treating a vanished process as a sampling race.""" + + try: + return path.read_text(encoding="utf-8") + except (FileNotFoundError, ProcessLookupError): + return None + except PermissionError as error: + raise MemoryEvidenceError(f"permission denied while reading {path}") from error + except OSError as error: + if error.errno in _TRANSIENT_PROC_ERRNOS: + return None + raise MemoryEvidenceError(f"cannot read procfs memory evidence from {path}") from error + + +def _direct_child_pids(proc_root: Path, process_id: int) -> set[int]: + """Return children created by any thread in one Linux process.""" + + task_root = proc_root / str(process_id) / "task" + try: + task_directories = tuple(task_root.iterdir()) + except (FileNotFoundError, ProcessLookupError): + return set() + except PermissionError as error: + raise MemoryEvidenceError(f"permission denied while reading {task_root}") from error + except OSError as error: + if error.errno in _TRANSIENT_PROC_ERRNOS: + return set() + raise MemoryEvidenceError(f"cannot enumerate process tasks in {task_root}") from error + + children: set[int] = set() + for task_directory in task_directories: + if not task_directory.name.isdigit(): + continue + child_text = _read_proc_text(task_directory / "children") + if child_text is None: + continue + for raw_child in child_text.split(): + try: + child_id = int(raw_child) + except ValueError as error: + raise MemoryEvidenceError( + f"invalid child process identifier {raw_child!r} in {task_directory}" + ) from error + children.add(_validated_pid(child_id)) + return children + + +def collect_process_tree_pids( + proc_root: Path, + root_pids: Iterable[int], +) -> tuple[int, ...]: + """Collect overlapping Linux process roots and descendants exactly once.""" + + roots = tuple(sorted({_validated_pid(process_id) for process_id in root_pids})) + if not roots: + raise MemoryEvidenceError("process-tree sampling requires at least one root PID") + observed: set[int] = set() + pending = list(reversed(roots)) + while pending: + process_id = pending.pop() + if process_id in observed: + continue + observed.add(process_id) + pending.extend( + child_id + for child_id in sorted( + _direct_child_pids(proc_root, process_id), + reverse=True, + ) + if child_id not in observed + ) + return tuple(sorted(observed)) + + +def _parse_kib_field(text: str, field: str, *, path: Path) -> int | None: + for line in text.splitlines(): + name, separator, raw_value = line.partition(":") + if not separator or name != field: + continue + parts = raw_value.split() + if len(parts) != 2 or parts[1] != "kB": + raise MemoryEvidenceError(f"invalid {field} value in {path}: {raw_value!r}") + try: + kibibytes = int(parts[0]) + except ValueError as error: + raise MemoryEvidenceError( + f"non-integer {field} value in {path}: {parts[0]!r}" + ) from error + if kibibytes < 0: + raise MemoryEvidenceError(f"negative {field} value in {path}: {kibibytes}") + return kibibytes * _KIBIBYTE + return None + + +def _read_rss_bytes(proc_root: Path, process_id: int) -> int | None: + process_root = proc_root / str(process_id) + status_path = process_root / "status" + status_text = _read_proc_text(status_path) + if status_text is None: + return None + rss_bytes = _parse_kib_field(status_text, "VmRSS", path=status_path) + if rss_bytes is not None: + return rss_bytes + + # A zombie can omit VmRSS. statm still provides an exact zero-resident + # record and avoids incorrectly treating a zero-footprint process as an + # unavailable PSS sample. + statm_path = process_root / "statm" + statm_text = _read_proc_text(statm_path) + if statm_text is None: + return None + fields = statm_text.split() + if len(fields) < 2: + raise MemoryEvidenceError(f"invalid statm record in {statm_path}") + try: + resident_pages = int(fields[1]) + page_size = int(os.sysconf("SC_PAGE_SIZE")) + except (OSError, TypeError, ValueError) as error: + raise MemoryEvidenceError(f"cannot decode resident pages in {statm_path}") from error + if resident_pages < 0 or page_size <= 0: + raise MemoryEvidenceError(f"invalid resident-page accounting in {statm_path}") + return resident_pages * page_size + + +def _read_pss_bytes(proc_root: Path, process_id: int) -> tuple[int | None, str | None]: + path = proc_root / str(process_id) / "smaps_rollup" + try: + text = path.read_text(encoding="utf-8") + except (FileNotFoundError, ProcessLookupError): + return None, "smaps_rollup disappeared during sampling" + except PermissionError: + return None, "smaps_rollup permission denied" + except OSError as error: + if error.errno in _TRANSIENT_PROC_ERRNOS: + return None, "smaps_rollup disappeared during sampling" + return None, f"smaps_rollup read failed with errno {error.errno}" + pss_bytes = _parse_kib_field(text, "Pss", path=path) + if pss_bytes is None: + return None, "smaps_rollup contained no Pss field" + return pss_bytes, None + + +def _read_process_identity( + proc_root: Path, + process_id: int, +) -> tuple[int, str] | None: + """Return Linux process start time and state for PID-reuse-safe sampling.""" + + path = proc_root / str(process_id) / "stat" + text = _read_proc_text(path) + if text is None: + return None + closing_parenthesis = text.rfind(")") + if closing_parenthesis < 0: + raise MemoryEvidenceError(f"invalid process identity record in {path}") + raw_pid = text[:closing_parenthesis].partition("(")[0].strip() + fields = text[closing_parenthesis + 1 :].split() + if raw_pid != str(process_id) or len(fields) < 20: + raise MemoryEvidenceError(f"invalid process identity record in {path}") + state = fields[0] + if len(state) != 1: + raise MemoryEvidenceError(f"invalid process state in {path}: {state!r}") + try: + start_time = int(fields[19]) + except ValueError as error: + raise MemoryEvidenceError(f"invalid process start time in {path}") from error + if start_time < 0: + raise MemoryEvidenceError(f"negative process start time in {path}") + return start_time, state + + +def sample_process_tree_memory( + *, + root_pids: Iterable[int], + proc_root: Path = Path("/proc"), +) -> dict[str, object]: + """Measure one concurrent, PID-deduplicated Linux process-tree snapshot.""" + + proc_root = proc_root.resolve() + roots = tuple(sorted({_validated_pid(process_id) for process_id in root_pids})) + process_ids = collect_process_tree_pids(proc_root, roots) + processes: list[dict[str, object]] = [] + fallback_reasons: set[str] = set() + transient_process_events: set[str] = set() + for process_id in process_ids: + identity_before = _read_process_identity(proc_root, process_id) + if identity_before is None: + transient_process_events.add( + f"pid {process_id}: disappeared before memory sampling" + ) + continue + if identity_before[1] == "Z": + processes.append( + { + "pid": process_id, + "rss_bytes": 0, + "pss_bytes": 0, + "accounted_bytes": 0, + "accounting_metric": "zero-resident-zombie", + "state": "Z", + } + ) + continue + rss_bytes = _read_rss_bytes(proc_root, process_id) + if rss_bytes is None: + transient_process_events.add( + f"pid {process_id}: disappeared before RSS sampling completed" + ) + continue + pss_bytes, pss_error = _read_pss_bytes(proc_root, process_id) + identity_after = _read_process_identity(proc_root, process_id) + if identity_after is not None and identity_after[0] != identity_before[0]: + transient_process_events.add( + f"pid {process_id}: PID was reused during memory sampling" + ) + pss_bytes = None + pss_error = "PID identity changed during PSS sampling" + if pss_bytes is None: + if identity_after is None: + transient_process_events.add( + f"pid {process_id}: exited before PSS sampling completed" + ) + elif identity_after[1] == "Z": + transient_process_events.add( + f"pid {process_id}: became a zombie during PSS sampling" + ) + elif identity_after is None: + # Both resident measurements were completed while this PID was live. + # Retaining them describes that observed instant without inventing a + # fallback; record the subsequent exit explicitly. + transient_process_events.add( + f"pid {process_id}: exited after complete RSS/PSS sampling" + ) + if pss_error is not None: + fallback_reasons.add(f"pid {process_id}: {pss_error}") + accounted_bytes = pss_bytes if pss_bytes is not None else rss_bytes + processes.append( + { + "pid": process_id, + "rss_bytes": rss_bytes, + "pss_bytes": pss_bytes, + "accounted_bytes": accounted_bytes, + "accounting_metric": ( + "proportional-set-size" + if pss_bytes is not None + else "resident-set-size-fallback" + ), + "state": identity_after[1] if identity_after is not None else "exited", + } + ) + + observed_pids = [int(process["pid"]) for process in processes] + if not observed_pids or not set(roots).intersection(observed_pids): + raise MemoryEvidenceError( + "procfs sampling did not capture any requested process-tree root" + ) + aggregate_rss = sum(int(process["rss_bytes"]) for process in processes) + pss_complete = all(isinstance(process["pss_bytes"], int) for process in processes) + aggregate_pss = ( + sum(int(process["pss_bytes"]) for process in processes) + if pss_complete + else None + ) + aggregate_hybrid = sum(int(process["accounted_bytes"]) for process in processes) + rss_fallback_process_count = sum( + process["pss_bytes"] is None for process in processes + ) + return { + "root_pids": list(roots), + "process_ids": observed_pids, + "process_count": len(observed_pids), + "pid_accounting": "unique-live-process-id", + "aggregate_rss_bytes": aggregate_rss, + "aggregate_pss_bytes": aggregate_pss, + "aggregate_hybrid_bytes": aggregate_hybrid, + "hybrid_accounting": ( + "sum each live process PSS when available, otherwise that process RSS" + ), + "rss_fallback_process_count": rss_fallback_process_count, + "pss_complete": pss_complete, + "pss_fallback_reasons": sorted(fallback_reasons), + "transient_process_event_count": len(transient_process_events), + "transient_process_events": sorted(transient_process_events), + "processes": processes, + } + + +class ConcurrentProcessTreeSampler: + """Sample a Linux controller process tree at a fixed high frequency.""" + + def __init__( + self, + *, + root_pids: Iterable[int] | None = None, + sample_interval_seconds: float = 0.05, + proc_root: Path = Path("/proc"), + ) -> None: + roots = (os.getpid(),) if root_pids is None else tuple(root_pids) + self.root_pids = tuple(sorted({_validated_pid(process_id) for process_id in roots})) + if not self.root_pids: + raise MemoryEvidenceError("concurrent memory sampling requires a root PID") + if sample_interval_seconds <= 0: + raise MemoryEvidenceError("memory sample interval must be positive") + self.sample_interval_seconds = float(sample_interval_seconds) + self.proc_root = proc_root.resolve() + self._lock = threading.Lock() + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._started = False + self._stopped = False + self._errors: list[str] = [] + self._sample_count = 0 + self._pss_complete_sample_count = 0 + self._pss_complete_for_all_samples = True + self._pss_fallback_reasons: set[str] = set() + self._transient_process_events: set[str] = set() + self._transient_process_event_count = 0 + self._observed_process_ids: set[int] = set() + self._max_process_count = 0 + self._peak_concurrent_rss_bytes = 0 + self._peak_concurrent_pss_bytes = 0 + self._peak_concurrent_hybrid_bytes = 0 + self._peak_rss_snapshot: dict[str, object] | None = None + self._peak_pss_snapshot: dict[str, object] | None = None + self._peak_hybrid_snapshot: dict[str, object] | None = None + self._last_sample_started: float | None = None + self._max_sample_gap_seconds = 0.0 + + def _record_snapshot( + self, + snapshot: dict[str, object], + *, + sample_started: float, + ) -> None: + with self._lock: + if self._last_sample_started is not None: + self._max_sample_gap_seconds = max( + self._max_sample_gap_seconds, + sample_started - self._last_sample_started, + ) + self._last_sample_started = sample_started + self._sample_count += 1 + process_ids = {int(process_id) for process_id in snapshot["process_ids"]} + self._observed_process_ids.update(process_ids) + self._max_process_count = max(self._max_process_count, len(process_ids)) + transient_events = snapshot["transient_process_events"] + if not isinstance(transient_events, list): + raise MemoryEvidenceError( + "process-tree snapshot has invalid transient lifecycle evidence" + ) + self._transient_process_events.update(str(event) for event in transient_events) + self._transient_process_event_count += int( + snapshot["transient_process_event_count"] + ) + aggregate_rss = int(snapshot["aggregate_rss_bytes"]) + if aggregate_rss > self._peak_concurrent_rss_bytes: + self._peak_concurrent_rss_bytes = aggregate_rss + self._peak_rss_snapshot = snapshot + aggregate_hybrid = int(snapshot["aggregate_hybrid_bytes"]) + if aggregate_hybrid > self._peak_concurrent_hybrid_bytes: + self._peak_concurrent_hybrid_bytes = aggregate_hybrid + self._peak_hybrid_snapshot = snapshot + if snapshot["pss_complete"] is True: + self._pss_complete_sample_count += 1 + aggregate_pss = int(snapshot["aggregate_pss_bytes"]) + if aggregate_pss > self._peak_concurrent_pss_bytes: + self._peak_concurrent_pss_bytes = aggregate_pss + self._peak_pss_snapshot = snapshot + else: + self._pss_complete_for_all_samples = False + self._pss_fallback_reasons.update(snapshot["pss_fallback_reasons"]) + + def sample_now(self) -> dict[str, object]: + if not self._started or self._stopped: + raise MemoryEvidenceError("concurrent memory sampler is not running") + sample_started = time.monotonic() + snapshot = sample_process_tree_memory( + root_pids=self.root_pids, + proc_root=self.proc_root, + ) + self._record_snapshot(snapshot, sample_started=sample_started) + return snapshot + + def _run(self) -> None: + while not self._stop_event.wait(self.sample_interval_seconds): + try: + self.sample_now() + except Exception as error: # pragma: no cover - exercised through stop() + with self._lock: + self._errors.append(f"{type(error).__name__}: {error}") + self._stop_event.set() + return + + def start(self) -> None: + if self._started: + raise MemoryEvidenceError("concurrent memory sampler was already started") + self._started = True + try: + self.sample_now() + except Exception: + self._started = False + raise + self._thread = threading.Thread( + target=self._run, + name="fastplms-cpu-memory-sampler", + daemon=True, + ) + self._thread.start() + + def stop(self) -> dict[str, object]: + if not self._started: + raise MemoryEvidenceError("concurrent memory sampler was never started") + if self._stopped: + raise MemoryEvidenceError("concurrent memory sampler was already stopped") + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=max(1.0, 10 * self.sample_interval_seconds)) + if self._thread.is_alive(): + with self._lock: + self._errors.append("sampler thread did not stop") + if not self._errors: + try: + self.sample_now() + except Exception as error: + with self._lock: + self._errors.append(f"{type(error).__name__}: {error}") + self._stopped = True + with self._lock: + if self._errors: + raise MemoryEvidenceError( + "concurrent process-tree sampling failed: " + "; ".join(self._errors) + ) + if ( + self._sample_count <= 0 + or self._peak_rss_snapshot is None + or self._peak_hybrid_snapshot is None + ): + raise MemoryEvidenceError("concurrent process-tree sampler produced no evidence") + return { + "available": True, + "measurement": "linux-procfs-concurrent-process-tree", + "root_pids": list(self.root_pids), + "sample_interval_seconds": self.sample_interval_seconds, + "sample_count": self._sample_count, + "max_sample_gap_seconds": round(self._max_sample_gap_seconds, 6), + "observed_process_ids": sorted(self._observed_process_ids), + "max_concurrent_process_count": self._max_process_count, + "pid_accounting": "each live descendant PID is counted once per sample", + "peak_concurrent_rss_bytes": self._peak_concurrent_rss_bytes, + "peak_concurrent_hybrid_bytes": self._peak_concurrent_hybrid_bytes, + "peak_concurrent_pss_bytes": ( + self._peak_concurrent_pss_bytes + if self._pss_complete_sample_count + else None + ), + "pss_complete_sample_count": self._pss_complete_sample_count, + "pss_complete_for_all_samples": self._pss_complete_for_all_samples, + "pss_fallback_reasons": sorted(self._pss_fallback_reasons), + "transient_process_event_count": self._transient_process_event_count, + "transient_process_events": sorted(self._transient_process_events), + "peak_rss_snapshot": self._peak_rss_snapshot, + "peak_pss_snapshot": self._peak_pss_snapshot, + "peak_hybrid_snapshot": self._peak_hybrid_snapshot, + } + + +def select_concurrent_memory_gate(evidence: Mapping[str, Any]) -> dict[str, object]: + """Gate on the peak per-process PSS/RSS hybrid across all samples.""" + + sample_count = evidence.get("sample_count") + peak_rss = evidence.get("peak_concurrent_rss_bytes") + peak_hybrid = evidence.get("peak_concurrent_hybrid_bytes") + if not isinstance(sample_count, int) or isinstance(sample_count, bool) or sample_count <= 0: + raise MemoryEvidenceError("concurrent memory evidence has no samples") + if not isinstance(peak_rss, int) or isinstance(peak_rss, bool) or peak_rss <= 0: + raise MemoryEvidenceError("concurrent memory evidence has no positive RSS peak") + if ( + not isinstance(peak_hybrid, int) + or isinstance(peak_hybrid, bool) + or peak_hybrid <= 0 + ): + raise MemoryEvidenceError("concurrent memory evidence has no positive hybrid peak") + if evidence.get("pss_complete_for_all_samples") is True: + peak_pss = evidence.get("peak_concurrent_pss_bytes") + if not isinstance(peak_pss, int) or isinstance(peak_pss, bool) or peak_pss <= 0: + raise MemoryEvidenceError( + "PSS evidence is marked complete but has no positive concurrent peak" + ) + return { + "metric": "proportional-set-size", + "source": "/proc//smaps_rollup:Pss", + "peak_bytes": peak_hybrid, + "fallback_used": False, + "fallback_is_conservative": False, + "fallback_reasons": [], + } + + reasons = evidence.get("pss_fallback_reasons") + normalized_reasons = ( + [str(reason) for reason in reasons] + if isinstance(reasons, list) and reasons + else ["one or more concurrent samples lacked complete PSS evidence"] + ) + return { + "metric": "per-process-pss-rss-hybrid", + "source": ( + "/proc//smaps_rollup:Pss with per-process " + "/proc//status:VmRSS fallback" + ), + "peak_bytes": peak_hybrid, + "fallback_used": True, + "fallback_is_conservative": True, + "fallback_reasons": normalized_reasons, + } + + +def _peak_rss_bytes(scope: int) -> int: + """Return Linux ``getrusage`` peak RSS in bytes for one resource scope.""" + + try: + import resource + + peak = int(resource.getrusage(scope).ru_maxrss) + except (ImportError, OSError, ValueError) as error: + raise MemoryEvidenceError("getrusage peak RSS is unavailable") from error + if peak < 0: + raise MemoryEvidenceError(f"getrusage returned a negative peak RSS: {peak}") + return peak * 1024 if os.name != "nt" else peak + + +def capture_process_memory(process_id: str, *, role: str) -> dict[str, object]: + """Capture one process and the largest child it has already waited for.""" + + try: + import resource + except ImportError as error: + raise MemoryEvidenceError("getrusage is required by the CPU contract") from error + if not process_id: + raise MemoryEvidenceError("process memory evidence requires a process identifier") + if role not in {"controller", "worker"}: + raise MemoryEvidenceError(f"unsupported process memory role: {role!r}") + process_peak = _peak_rss_bytes(resource.RUSAGE_SELF) + children_peak = _peak_rss_bytes(resource.RUSAGE_CHILDREN) + return { + "process_id": process_id, + "pid": os.getpid(), + "role": role, + "process_peak_rss_bytes": process_peak, + "waited_children_peak_rss_bytes": children_peak, + } + + +def _validated_record(record: Mapping[str, Any], *, role: str) -> dict[str, object]: + process_id = record.get("process_id") + if not isinstance(process_id, str) or not process_id: + raise MemoryEvidenceError("process memory evidence has no process_id") + if record.get("role") != role: + raise MemoryEvidenceError( + f"process {process_id!r} reported role {record.get('role')!r}, expected {role!r}" + ) + process_pid = _validated_pid(record.get("pid")) + values: dict[str, int] = {} + for name in ("process_peak_rss_bytes", "waited_children_peak_rss_bytes"): + value = record.get(name) + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise MemoryEvidenceError( + f"process {process_id!r} has invalid {name}: {value!r}" + ) + if name == "process_peak_rss_bytes" and value == 0: + raise MemoryEvidenceError( + f"process {process_id!r} reported no process peak RSS" + ) + values[name] = value + return { + "process_id": process_id, + "pid": process_pid, + "role": role, + **values, + } + + +def aggregate_process_memory( + controller: Mapping[str, Any], + workers: Iterable[Mapping[str, Any]], +) -> dict[str, object]: + """Build a temporal upper bound from nonconcurrent getrusage maxima. + + Under xdist, the controller's ``RUSAGE_CHILDREN`` contains the same workers + that report their own usage, so it is evidence only and is not added. Each + worker's waited-child peak is added because isolated AutoClass probes are + separate processes and are absent from that worker's ``RUSAGE_SELF``. + In serial mode, the controller's waited children are added directly. + """ + + controller_record = _validated_record(controller, role="controller") + worker_records = [_validated_record(record, role="worker") for record in workers] + worker_records.sort(key=lambda record: str(record["process_id"])) + worker_ids = [str(record["process_id"]) for record in worker_records] + if len(worker_ids) != len(set(worker_ids)): + raise MemoryEvidenceError("duplicate xdist worker memory evidence") + worker_pids = [int(record["pid"]) for record in worker_records] + if len(worker_pids) != len(set(worker_pids)): + raise MemoryEvidenceError("duplicate xdist worker operating-system PIDs") + + controller_self = int(controller_record["process_peak_rss_bytes"]) + if worker_records: + accounted_controller_children = 0 + accounted_workers = sum( + int(record["process_peak_rss_bytes"]) + + int(record["waited_children_peak_rss_bytes"]) + for record in worker_records + ) + accounting_mode = "xdist-conservative-process-tree-upper-bound" + else: + accounted_controller_children = int( + controller_record["waited_children_peak_rss_bytes"] + ) + accounted_workers = 0 + accounting_mode = "serial-conservative-process-tree-upper-bound" + aggregate = controller_self + accounted_controller_children + accounted_workers + return { + "available": True, + "measurement": "getrusage.ru_maxrss", + "units": "bytes", + "accounting_mode": accounting_mode, + "budget_enforced": False, + "interpretation": ( + "Diagnostic temporal upper bound only: component maxima can occur at " + "different times and must not enforce the concurrent memory budget." + ), + "double_counting_policy": ( + "Exclude controller RUSAGE_CHILDREN when workers report themselves; " + "include each worker RUSAGE_SELF and RUSAGE_CHILDREN exactly once." + ), + "aggregate_peak_rss_bytes": aggregate, + "temporal_upper_bound_rss_bytes": aggregate, + "controller": controller_record, + "workers": worker_records, + } diff --git a/tests/cpu/test_ankh_contracts.py b/tests/cpu/test_ankh_contracts.py new file mode 100644 index 0000000..f9046e1 --- /dev/null +++ b/tests/cpu/test_ankh_contracts.py @@ -0,0 +1,270 @@ +"""Mandatory ANKH encoder/decoder and Hugging Face contracts. + +The detailed assertions live beside the unit tests so they are also useful to +developers running a focused file. This module is the positive, offline CPU +allowlist used by the required status check. +""" + +import pytest +import torch +from pathlib import Path +from transformers.utils import ModelOutput + +from tests.unit import test_ankh_cpu_contract as contracts + + +def _assert_nested_output_close(actual: object, expected: object) -> None: + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, (tuple, list)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_value, expected_value in zip(actual, expected, strict=True): + _assert_nested_output_close(actual_value, expected_value) + return + assert actual == expected + + +def _encoder_labels( + model_class: ( + type[contracts.FastAnkhModel] + | type[contracts.FastAnkhForMaskedLMExtension] + | type[contracts.FastAnkhForSequenceClassification] + | type[contracts.FastAnkhForTokenClassification] + ), + input_ids: torch.Tensor, + attention_mask: torch.Tensor, +) -> torch.Tensor | None: + if model_class is contracts.FastAnkhForMaskedLMExtension: + return input_ids.masked_fill(~attention_mask, -100) + if model_class is contracts.FastAnkhForSequenceClassification: + return torch.tensor([1, 2]) + if model_class is contracts.FastAnkhForTokenClassification: + return input_ids.remainder(3).masked_fill(~attention_mask, -100) + return None + + +test_complete_t5_checkpoint_loads_clean_encoder_and_seq2seq_views = ( + contracts.test_complete_t5_checkpoint_loads_clean_encoder_and_seq2seq_views +) +test_decoder_default_attention_mask_preserves_t5_start_and_masks_padding = ( + contracts.test_decoder_default_attention_mask_preserves_t5_start_and_masks_padding +) +test_decoder_embedding_batch_masks_start_eos_padding_and_sentinels = ( + contracts.test_decoder_embedding_batch_masks_start_eos_padding_and_sentinels +) +test_decoder_embed_dataset_slices_aligned_inputs_and_records_provenance = ( + contracts.test_decoder_embed_dataset_slices_aligned_inputs_and_records_provenance +) +test_decoder_embeddings_require_explicit_aligned_inputs = ( + contracts.test_decoder_embeddings_require_explicit_aligned_inputs +) +test_encoder_auto_classes_honor_tuple_and_dict_outputs = ( + contracts.test_encoder_auto_classes_honor_tuple_and_dict_outputs +) +test_encoder_auto_classes_resize_shared_input_embeddings = ( + contracts.test_encoder_auto_classes_resize_shared_input_embeddings +) +test_encoder_only_view_rejects_decoder_hidden_states = ( + contracts.test_encoder_only_view_rejects_decoder_hidden_states +) +test_encoder_task_heads_produce_finite_loss_and_gradients = ( + contracts.test_encoder_task_heads_produce_finite_loss_and_gradients +) +test_offline_auto_tokenizer_flags_and_seq2seq_generation_config = ( + contracts.test_offline_auto_tokenizer_flags_and_seq2seq_generation_config +) +test_ankh_decoder_embedding_inputs_use_tight_sentinel_tokenization = ( + contracts.test_ankh_decoder_embedding_inputs_use_tight_sentinel_tokenization +) +test_ankh_explicit_and_model_owned_tokenizers_share_the_raw_sequence_contract = ( + contracts.test_ankh_explicit_and_model_owned_tokenizers_share_the_raw_sequence_contract +) +test_ankh_tokenization_normalizes_raw_sequences_and_tight_sentinel_prompts = ( + contracts.test_ankh_tokenization_normalizes_raw_sequences_and_tight_sentinel_prompts +) +test_ankh_tokenization_rejects_empty_inputs_and_real_slow_tokenizers = ( + contracts.test_ankh_tokenization_rejects_empty_inputs_and_real_slow_tokenizers +) +test_ankh_ttt_uses_raw_residue_tokenization = contracts.test_ankh_ttt_uses_raw_residue_tokenization +test_sdpa_output_attentions_fallback_keeps_padding_mask_and_backend = ( + contracts.test_sdpa_output_attentions_fallback_keeps_padding_mask_and_backend +) +test_seq2seq_embedding_selects_encoder_and_explicit_decoder_layers = ( + contracts.test_seq2seq_embedding_selects_encoder_and_explicit_decoder_layers +) +test_seq2seq_head_produces_finite_loss_and_gradients = ( + contracts.test_seq2seq_head_produces_finite_loss_and_gradients +) +test_seq2seq_view_forces_eager_without_changing_encoder_backend_contract = ( + contracts.test_seq2seq_view_forces_eager_without_changing_encoder_backend_contract +) +test_tokenizer_load_context_is_per_instance_and_offline_scoped = ( + contracts.test_tokenizer_load_context_is_per_instance_and_offline_scoped +) +test_tokenizer_load_context_is_isolated_during_concurrent_first_access = ( + contracts.test_tokenizer_load_context_is_isolated_during_concurrent_first_access +) + + +@pytest.mark.parametrize( + "model_class", + ( + contracts.FastAnkhModel, + contracts.FastAnkhForMaskedLMExtension, + contracts.FastAnkhForSequenceClassification, + contracts.FastAnkhForTokenClassification, + ), +) +def test_ankh_encoder_views_backward_and_save_reload( + model_class: ( + type[contracts.FastAnkhModel] + | type[contracts.FastAnkhForMaskedLMExtension] + | type[contracts.FastAnkhForSequenceClassification] + | type[contracts.FastAnkhForTokenClassification] + ), + tmp_path: Path, +) -> None: + model = model_class(contracts._config(num_labels=3)).eval() + input_ids = torch.tensor([[2, 3, 1], [4, 1, 0]]) # (b=2, l=3) + attention_mask = input_ids.ne(0) # (b, l) + labels = _encoder_labels(model_class, input_ids, attention_mask) # (b, ...) or None + model_arguments = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "output_attentions": True, + "output_hidden_states": True, + } + if labels is not None: + model_arguments["labels"] = labels + output = model( + **model_arguments, + return_dict=True, + ) + tuple_output = model( + **model_arguments, + return_dict=False, + ) + + assert isinstance(output, ModelOutput) + assert output.hidden_states is not None + assert output.attentions is not None + _assert_nested_output_close(tuple_output, output.to_tuple()) + if labels is not None: + torch.testing.assert_close(tuple_output[0], output.loss) + torch.testing.assert_close(tuple_output[1], output.logits) + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model( + input_ids=input_ids, + attention_mask=attention_mask, + unexpected_cpu_contract=True, + ) + + tensor = ( # (b, l, d) for the base model; otherwise task-head logits + output.last_hidden_state if model_class is contracts.FastAnkhModel else output.logits + ) + loss = output.loss if labels is not None else tensor.square().mean() # () + assert loss is not None and torch.isfinite(loss) + loss.backward() + assert model.shared.weight.grad is not None + + save_dir = tmp_path / model_class.__name__ + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained(save_dir, local_files_only=True).eval() + with torch.inference_mode(): + observed = reloaded( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + ) + observed_tensor = ( + observed.last_hidden_state if model_class is contracts.FastAnkhModel else observed.logits + ) + torch.testing.assert_close(observed_tensor, tensor.detach(), rtol=0.0, atol=0.0) + + +def test_ankh_seq2seq_view_honors_tuple_output_and_resize() -> None: + model = contracts.FastAnkhForConditionalGeneration(contracts._config()).eval() + input_ids = torch.tensor([[2, 3, 1, 0], [4, 1, 0, 0]]) # (b=2, l_enc=4) + attention_mask = input_ids.ne(0) # (b, l_enc) + decoder_input_ids = torch.tensor([[0, 5, 1, 0], [0, 6, 7, 1]]) # (b, l_dec=4) + decoder_attention_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]]) # (b, l_dec) + structured = model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + output_attentions=True, + output_hidden_states=True, + use_cache=False, + return_dict=True, + ) + tuple_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + output_attentions=True, + output_hidden_states=True, + use_cache=False, + return_dict=False, + ) + + assert isinstance(structured, ModelOutput) + assert isinstance(tuple_output, tuple) + assert structured.decoder_hidden_states is not None + assert structured.encoder_hidden_states is not None + assert structured.decoder_attentions is not None + assert structured.cross_attentions is not None + assert structured.encoder_attentions is not None + causal_future = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1) # (l_dec, l_dec) + for decoder_attention in structured.decoder_attentions: + assert torch.count_nonzero(decoder_attention.masked_select(causal_future)) == 0 + assert torch.count_nonzero(decoder_attention[0, :, :, 3]) == 0 + for cross_attention in structured.cross_attentions: + assert torch.count_nonzero(cross_attention[0, :, :, 3]) == 0 + assert torch.count_nonzero(cross_attention[1, :, :, 2:]) == 0 + _assert_nested_output_close(tuple_output, structured.to_tuple()) + torch.testing.assert_close(tuple_output[0], structured.logits) + + labels = torch.tensor([[5, 1, -100, -100], [6, 7, 1, -100]]) # (b, l_dec) + loss_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + use_cache=False, + return_dict=True, + ) + loss_tuple = model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + use_cache=False, + return_dict=False, + ) + _assert_nested_output_close(loss_tuple, loss_output.to_tuple()) + torch.testing.assert_close(loss_tuple[0], loss_output.loss) + torch.testing.assert_close(loss_tuple[1], loss_output.logits) + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + unexpected_cpu_contract=True, + ) + + model.resize_token_embeddings(19) + assert model.get_input_embeddings().num_embeddings == 19 + assert model.get_output_embeddings().out_features == 19 + assert model.encoder.embed_tokens is model.shared + assert model.decoder.embed_tokens is model.shared diff --git a/tests/cpu/test_attention_contracts.py b/tests/cpu/test_attention_contracts.py new file mode 100644 index 0000000..259310a --- /dev/null +++ b/tests/cpu/test_attention_contracts.py @@ -0,0 +1,449 @@ +"""Mandatory attention dispatch, masking, fallback, and cache contracts.""" + +import pytest +import torch +from collections import OrderedDict + +import fastplms.attention.interfaces as attention_interfaces +import fastplms.models.esm_plusplus.modeling_esm_plusplus as esmpp_module +from fastplms.attention import _core as attention_core +from fastplms.models.esm2.modeling_fastesm import FastEsmConfig, FastEsmModel +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, +) +from tests.unit import test_attention_interfaces as interface_contracts +from tests.unit import test_attention_regressions as contracts +from tests.unit import test_esmc_diagnostics as esmc_contracts + + +test_flash_backend_loads_only_its_hugging_face_kernel = ( + interface_contracts.test_flash_backend_loads_only_its_hugging_face_kernel +) +test_causal_masked_flash_uses_varlen_and_zeroes_padding = ( + interface_contracts.test_causal_masked_flash_uses_varlen_and_zeroes_padding +) +test_masked_flash_validates_padding_mask_shape_before_kernel_loading = ( + interface_contracts.test_masked_flash_validates_padding_mask_shape_before_kernel_loading +) +test_public_attention_setter_matches_transformers_513_kernel_policy = ( + interface_contracts.test_public_attention_setter_matches_transformers_513_kernel_policy +) +test_model_flash_flags_match_the_manifest = ( + interface_contracts.test_model_flash_flags_match_the_manifest +) +test_ankh_sdpa_never_mutates_process_global_reduction_policy = ( + interface_contracts.test_ankh_sdpa_never_mutates_process_global_reduction_policy +) +test_ankh_concurrent_fallback_and_sdpa_keep_backend_and_global_policy = ( + interface_contracts.test_ankh_concurrent_fallback_and_sdpa_keep_backend_and_global_policy +) +test_compiled_flex_cache_key_covers_execution_not_batch_contents = ( + interface_contracts.test_compiled_flex_cache_key_covers_execution_not_batch_contents +) +test_flex_block_mask_supports_disjoint_valid_spans_and_exact_cache_keys = ( + interface_contracts.test_flex_block_mask_supports_disjoint_valid_spans_and_exact_cache_keys +) +test_flex_block_mask_key_separates_equal_bytes_with_different_pattern_dtypes = ( + interface_contracts.test_flex_block_mask_key_separates_equal_bytes_with_different_pattern_dtypes +) +test_esmplusplus_flex_sequence_masks_share_exact_bounded_cache = ( + interface_contracts.test_esmplusplus_flex_sequence_masks_share_exact_bounded_cache +) +test_flash_kernel_variant_mismatch_fails_closed = ( + interface_contracts.test_flash_kernel_variant_mismatch_fails_closed +) +test_locked_kernel_is_hash_validated_before_import = ( + interface_contracts.test_locked_kernel_is_hash_validated_before_import +) +test_locked_kernel_offline_resolves_sparse_snapshot_without_hub_api = ( + interface_contracts.test_locked_kernel_offline_resolves_sparse_snapshot_without_hub_api +) +test_locked_kernel_offline_rejects_unlocked_cached_variant = ( + interface_contracts.test_locked_kernel_offline_rejects_unlocked_cached_variant +) + +test_attention_masks_require_exact_batch_sequence_shape = ( + contracts.test_attention_masks_require_exact_batch_sequence_shape +) +test_attention_masks_reject_rows_without_valid_keys_before_dispatch = ( + contracts.test_attention_masks_reject_rows_without_valid_keys_before_dispatch +) +test_ankh_output_attentions_eager_fallback_honors_attention_dropout = ( + contracts.test_ankh_output_attentions_eager_fallback_honors_attention_dropout +) +test_ankh_sdpa_receives_training_attention_dropout = ( + contracts.test_ankh_sdpa_receives_training_attention_dropout +) +test_ankh_ttt_missing_input_uses_optimization_safe_validation = ( + contracts.test_ankh_ttt_missing_input_uses_optimization_safe_validation +) +test_bool_to_additive_mask_rejects_non_boolean_input_explicitly = ( + contracts.test_bool_to_additive_mask_rejects_non_boolean_input_explicitly +) +test_dplm_output_attentions_fallback_preserves_padding_mask = ( + contracts.test_dplm_output_attentions_fallback_preserves_padding_mask +) +test_dplm2_output_attentions_fallback_is_call_scoped = ( + contracts.test_dplm2_output_attentions_fallback_is_call_scoped +) +test_dplm_sdpa_output_attentions_fallback_preserves_cross_attention_mask_and_backend = ( + contracts.test_dplm_sdpa_output_attentions_fallback_preserves_cross_attention_mask_and_backend +) +test_e1_output_attentions_fallback_preserves_block_causal_mask_and_backend = ( + contracts.test_e1_output_attentions_fallback_preserves_block_causal_mask_and_backend +) +test_e1_public_fallback_warns_once_masks_padding_and_preserves_flex = ( + contracts.test_e1_public_fallback_warns_once_masks_padding_and_preserves_flex +) +test_esm3_sdpa_fallback_is_call_scoped_and_preserves_padding_mask = ( + contracts.test_esm3_sdpa_fallback_is_call_scoped_and_preserves_padding_mask +) +test_esm3_sequence_id_grouping_combines_with_public_padding_mask = ( + contracts.test_esm3_sequence_id_grouping_combines_with_public_padding_mask +) +test_esm3_rejects_malformed_padding_and_sequence_masks = ( + contracts.test_esm3_rejects_malformed_padding_and_sequence_masks +) +test_esm2_output_attentions_fallback_preserves_padding_mask = ( + contracts.test_esm2_output_attentions_fallback_preserves_padding_mask +) +test_esm2_attention_validates_head_divisibility_without_assertions = ( + contracts.test_esm2_attention_validates_head_divisibility_without_assertions +) +test_esm2_flex_unavailability_raises_explicit_runtime_error = ( + contracts.test_esm2_flex_unavailability_raises_explicit_runtime_error +) +test_esm2_legacy_backend_setter_uses_explicit_validation = ( + contracts.test_esm2_legacy_backend_setter_uses_explicit_validation +) +test_esmfold_flex_training_rejects_unimplemented_attention_dropout = ( + contracts.test_esmfold_flex_training_rejects_unimplemented_attention_dropout +) +test_esmfold_output_attentions_fallback_preserves_padding_mask = ( + contracts.test_esmfold_output_attentions_fallback_preserves_padding_mask +) +test_esmplusplus_rejects_empty_attention_rows_without_fallback_or_mutation = ( + contracts.test_esmplusplus_rejects_empty_attention_rows_without_fallback_or_mutation +) +test_esmplusplus_chain_masks_fail_closed_without_assertions = ( + contracts.test_esmplusplus_chain_masks_fail_closed_without_assertions +) +test_esmplusplus_flex_unavailability_raises_explicit_runtime_error = ( + contracts.test_esmplusplus_flex_unavailability_raises_explicit_runtime_error +) +test_flash_attention_2_dense_and_varlen_use_autograd_wrappers = ( + contracts.test_flash_attention_2_dense_and_varlen_use_autograd_wrappers +) +test_flash_attention_2_dense_and_varlen_preserve_lora_and_input_gradients = ( + contracts.test_flash_attention_2_dense_and_varlen_preserve_lora_and_input_gradients +) +test_flash_attention_2_rejects_low_level_only_kernel_artifact = ( + contracts.test_flash_attention_2_rejects_low_level_only_kernel_artifact +) +test_flash_attention_3_preserves_internal_type_errors = ( + contracts.test_flash_attention_3_preserves_internal_type_errors +) +test_kernel_loader_honors_all_offline_environment_variables = ( + contracts.test_kernel_loader_honors_all_offline_environment_variables +) +test_public_flex_cache_cleanup_is_scoped_and_complete = ( + contracts.test_public_flex_cache_cleanup_is_scoped_and_complete +) + +test_esmc_calibration_contains_no_expected_failures = ( + esmc_contracts.test_esmc_calibration_contains_no_expected_failures +) +test_esmc_catastrophic_disagreement_remains_a_hard_failure = ( + esmc_contracts.test_esmc_catastrophic_disagreement_remains_a_hard_failure +) +test_esmc_supported_backend_deviation_warns_and_writes_complete_metrics = ( + esmc_contracts.test_esmc_supported_backend_deviation_warns_and_writes_complete_metrics +) + + +def _tiny_esm2(attn_backend: str) -> FastEsmModel: + return FastEsmModel( + FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=5, + position_embedding_type="absolute", + attn_backend=attn_backend, + ) + ) + + +def _tiny_esmc(attn_backend: str) -> ESMplusplusModel: + return ESMplusplusModel( + ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + dropout=0.0, + pad_token_id=1, + mask_token_id=5, + attn_backend=attn_backend, + ) + ).eval() + + +def test_esmc_flex_dispatch_is_compiled_once_and_receives_exact_padding_mask( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The supported ESMC Flex path must dispatch without hiding mask semantics.""" + + sentinel_block_mask = object() + observed: dict[str, object] = { + "block_masks_created": 0, + "compiled": 0, + "compile_requests": [], + "dispatches": [], + "mask_mod": None, + } + compiled = None + + def fake_create_block_mask(mask_mod, *shape, **kwargs): + observed["block_masks_created"] = int(observed["block_masks_created"]) + 1 + observed["mask_mod"] = mask_mod + observed["block_mask_shape"] = shape + observed["block_mask_device"] = kwargs["device"] + return sentinel_block_mask + + def fake_get_flex_attention_fn(**kwargs): + nonlocal compiled + cast_requests = observed["compile_requests"] + assert isinstance(cast_requests, list) + cast_requests.append(kwargs) + if compiled is None: + observed["compiled"] = int(observed["compiled"]) + 1 + + def run_flex(query, key, value, **call_kwargs): + cast_dispatches = observed["dispatches"] + assert isinstance(cast_dispatches, list) + cast_dispatches.append( + { + "query_shape": tuple(query.shape), + "key_shape": tuple(key.shape), + "value_shape": tuple(value.shape), + **call_kwargs, + } + ) + return value + + compiled = run_flex + return compiled + + monkeypatch.setattr(attention_core, "_flex_block_masks", OrderedDict()) + monkeypatch.setattr(attention_core, "create_block_mask", fake_create_block_mask) + monkeypatch.setattr(esmpp_module, "flex_attention", object()) + monkeypatch.setattr( + esmpp_module, + "_get_flex_attention_fn", + fake_get_flex_attention_fn, + ) + + model = _tiny_esmc("flex_attention") + input_ids = torch.tensor(((0, 3, 4, 2, 1), (0, 6, 2, 1, 1))) # (b=2, l=5) + attention_mask = input_ids.ne(1) # (b, l) + first = model(input_ids=input_ids, attention_mask=attention_mask) # hidden: (b, l, d=8) + second = model(input_ids=input_ids, attention_mask=attention_mask) # hidden: (b, l, d) + + assert first.last_hidden_state.shape == (2, 5, 8) + assert second.last_hidden_state.shape == first.last_hidden_state.shape + assert torch.isfinite(first.last_hidden_state).all() + assert torch.isfinite(second.last_hidden_state).all() + assert model.attn_backend == "flex_attention" + assert model.config.attn_backend == "flex_attention" + assert model.config._attn_implementation == "flex_attention" + + compile_requests = observed["compile_requests"] + dispatches = observed["dispatches"] + assert isinstance(compile_requests, list) and len(compile_requests) == 2 + assert observed["compiled"] == 1 + assert compile_requests[0] == compile_requests[1] + assert compile_requests[0] == { + "device": torch.device("cpu"), + "dtype": torch.float32, + "shape": (2, 2, 5, 4), + "mask_semantics": "padding", + } + assert observed["block_masks_created"] == 1 + assert len(attention_core._flex_block_masks) == 1 + assert isinstance(dispatches, list) and len(dispatches) == 2 + for dispatch in dispatches: + assert dispatch["query_shape"] == (2, 2, 5, 4) + assert dispatch["key_shape"] == (2, 2, 5, 4) + assert dispatch["value_shape"] == (2, 2, 5, 4) + assert dispatch["block_mask"] is sentinel_block_mask + assert dispatch["scale"] == 0.5 + assert dispatch["kernel_options"] == { + "PRESCALE_QK": True, + "BLOCK_N": 32, + } + + assert observed["block_mask_shape"] == (2, 1, 5, 5) + assert observed["block_mask_device"] == torch.device("cpu") + mask_mod = observed["mask_mod"] + assert callable(mask_mod) + for batch_index in range(2): + for query_index in range(5): + for key_index in range(5): + expected = bool( + attention_mask[batch_index, query_index] + == attention_mask[batch_index, key_index] + ) + assert bool(mask_mod(batch_index, 0, query_index, key_index)) is expected + + +def test_esmc_fake_fa3_dispatch_receives_exact_mask_and_returns_finite_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The supported ESMC FA3 path must preserve public shape and mask contracts.""" + + calls: list[dict[str, object]] = [] + + def fake_flash_attention(**kwargs): + query = kwargs["query_states"] + key = kwargs["key_states"] + value = kwargs["value_states"] + attention_mask = kwargs["attention_mask_2d"] + masked_output = value.masked_fill( # (b, l, h, d_h) + attention_mask[:, :, None, None].logical_not(), + 0.0, + ) + calls.append( + { + "query_shape": tuple(query.shape), + "key_shape": tuple(key.shape), + "value_shape": tuple(value.shape), + "attention_mask": attention_mask.clone(), + "causal": kwargs["causal"], + "implementation": kwargs["implementation"], + "masked_output": masked_output, + } + ) + return masked_output + + monkeypatch.setattr( + esmpp_module, + "kernels_flash_attention_func", + fake_flash_attention, + ) + monkeypatch.setattr(attention_interfaces, "require_kernels_package", lambda: None) + + model = _tiny_esmc("flash_attention_3") + input_ids = torch.tensor(((0, 3, 4, 2, 1), (0, 6, 2, 1, 1))) # (b=2, l=5) + attention_mask = input_ids.ne(1) # (b, l) + output = model(input_ids=input_ids, attention_mask=attention_mask) # hidden: (b, l, d=8) + + assert output.last_hidden_state.shape == (2, 5, 8) + assert torch.isfinite(output.last_hidden_state).all() + assert model.attn_backend == "flash_attention_3" + assert model.config.attn_backend == "flash_attention_3" + assert model.config._attn_implementation == "flash_attention_3" + assert len(calls) == 1 + call = calls[0] + assert call["query_shape"] == (2, 5, 2, 4) + assert call["key_shape"] == (2, 5, 2, 4) + assert call["value_shape"] == (2, 5, 2, 4) + assert call["causal"] is False + assert call["implementation"] == "flash_attention_3" + assert torch.equal(call["attention_mask"], attention_mask) + masked_output = call["masked_output"] + assert isinstance(masked_output, torch.Tensor) + assert torch.count_nonzero(masked_output[attention_mask.logical_not()]) == 0 + + +def test_eager_and_sdpa_match_for_dense_and_mixed_padding_with_gradients() -> None: + torch.manual_seed(17) + eager = _tiny_esm2("eager").eval() + sdpa = _tiny_esm2("sdpa").eval() + sdpa.load_state_dict(eager.state_dict()) + input_ids = torch.tensor([[0, 3, 4, 2, 1], [0, 6, 2, 1, 1]]) # (b=2, l=5) + masks = (torch.ones_like(input_ids), input_ids.ne(1)) # each (b, l) + + for attention_mask in masks: + eager.zero_grad(set_to_none=True) + sdpa.zero_grad(set_to_none=True) + eager_output = eager( + input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + sdpa_output = sdpa( + input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + torch.testing.assert_close( + sdpa_output.last_hidden_state, + eager_output.last_hidden_state, + rtol=2e-5, + atol=2e-6, + ) + assert eager_output.hidden_states is not None + assert sdpa_output.hidden_states is not None + for sdpa_hidden, eager_hidden in zip( + sdpa_output.hidden_states, + eager_output.hidden_states, + strict=True, + ): + torch.testing.assert_close( + sdpa_hidden, + eager_hidden, + rtol=2e-5, + atol=2e-6, + ) + + valid = attention_mask.bool().unsqueeze(-1) # (b, l, 1) + # A squared norm after the final LayerNorm is nearly constant and + # produces cancellation-dominated gradients. Project each hidden + # coordinate onto a distinct fixed coefficient so this parity gate + # compares a well-conditioned, nonzero backward signal. + gradient_probe = torch.linspace( + -1.0, + 1.0, + eager_output.last_hidden_state.shape[-1], + dtype=eager_output.last_hidden_state.dtype, + device=eager_output.last_hidden_state.device, + ).view(1, 1, -1) # (1, 1, d) + (eager_output.last_hidden_state * valid * gradient_probe).sum().backward() + (sdpa_output.last_hidden_state * valid * gradient_probe).sum().backward() + eager_gradients = { + name: parameter.grad + for name, parameter in eager.named_parameters() + if parameter.grad is not None + } + sdpa_gradients = { + name: parameter.grad + for name, parameter in sdpa.named_parameters() + if parameter.grad is not None + } + assert eager_gradients + assert sdpa_gradients.keys() == eager_gradients.keys() + expected_attention_projections = ( + "query", + "key", + "value", + "attention.output.dense", + ) + assert all( + any(projection in name for name in eager_gradients) + for projection in expected_attention_projections + ) + for name, eager_gradient in eager_gradients.items(): + torch.testing.assert_close( + sdpa_gradients[name], + eager_gradient, + rtol=3e-5, + atol=3e-6, + ) diff --git a/tests/cpu/test_autoclass_evidence_matrix.py b/tests/cpu/test_autoclass_evidence_matrix.py new file mode 100644 index 0000000..052313b --- /dev/null +++ b/tests/cpu/test_autoclass_evidence_matrix.py @@ -0,0 +1,417 @@ +"""Explicit runtime evidence for every family-level advertised AutoClass.""" + +from __future__ import annotations + +import importlib +import pytest +from dataclasses import dataclass +from transformers import PretrainedConfig + +from fastplms.registry import get_model_registry + + +@dataclass(frozen=True) +class AutoClassEvidence: + symbol_path: str + capabilities: frozenset[str] + runtime_tests: tuple[str, ...] + case_parameter: str | None = None + limitations: tuple[str, ...] = () + + +_CONFIG_TEST = ( + "tests/cpu/test_autoclass_evidence_matrix.py::" + "test_every_advertised_config_round_trips_offline" +) +_SEQUENCE_TEST = ( + "tests/cpu/test_sequence_autoclass_contracts.py::" + "test_esm2_advertised_models_forward_loss_backward_resize_and_reload" +) +_ESMC_TEST = ( + "tests/cpu/test_sequence_autoclass_contracts.py::" + "test_esmc_public_models_forward_loss_backward_resize_and_reload" +) +_DPLM_TEST = ( + "tests/cpu/test_sequence_autoclass_contracts.py::" + "test_dplm_advertised_models_forward_loss_backward_resize_and_reload" +) +_E1_TEST = ( + "tests/cpu/test_sequence_autoclass_contracts.py::" + "test_e1_advertised_models_forward_loss_backward_resize_and_reload" +) +_ANKH_TEST = ( + "tests/cpu/test_ankh_contracts.py::test_ankh_encoder_views_backward_and_save_reload" +) +_BASE_CAPABILITIES = frozenset( + {"backward", "forward", "resize", "return_dict", "save_reload", "tuple"} +) +_HEAD_CAPABILITIES = _BASE_CAPABILITIES | {"loss"} + + +def _config(symbol_path: str) -> AutoClassEvidence: + return AutoClassEvidence( + symbol_path=symbol_path, + capabilities=frozenset({"construct", "serialize", "reload"}), + runtime_tests=(_CONFIG_TEST,), + case_parameter="symbol_path", + ) + + +def _sequence( + symbol_path: str, + runtime_test: str, + *, + loss: bool, +) -> AutoClassEvidence: + return AutoClassEvidence( + symbol_path=symbol_path, + capabilities=_HEAD_CAPABILITIES if loss else _BASE_CAPABILITIES, + runtime_tests=(runtime_test,), + case_parameter="model_class", + limitations=() if loss else ("supervised loss is not applicable to a base model",), + ) + + +AUTOCLASS_EVIDENCE: dict[tuple[str, str], AutoClassEvidence] = { + ("esm2", "AutoConfig"): _config( + "fastplms.models.esm2.modeling_fastesm.FastEsmConfig" + ), + ("esm2", "AutoModel"): _sequence( + "fastplms.models.esm2.modeling_fastesm.FastEsmModel", + _SEQUENCE_TEST, + loss=False, + ), + ("esm2", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.esm2.modeling_fastesm.FastEsmForMaskedLM", + _SEQUENCE_TEST, + loss=True, + ), + ("esm2", "AutoModelForSequenceClassification"): _sequence( + "fastplms.models.esm2.modeling_fastesm.FastEsmForSequenceClassification", + _SEQUENCE_TEST, + loss=True, + ), + ("esm2", "AutoModelForTokenClassification"): _sequence( + "fastplms.models.esm2.modeling_fastesm.FastEsmForTokenClassification", + _SEQUENCE_TEST, + loss=True, + ), + ("esm_plusplus", "AutoConfig"): _config( + "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig" + ), + ("esm_plusplus", "AutoModel"): _sequence( + "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", + _ESMC_TEST, + loss=False, + ), + ("esm_plusplus", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM", + _ESMC_TEST, + loss=True, + ), + ("esm3", "AutoConfig"): _config( + "fastplms.models.esm3.modeling_esm3.FastESM3Config" + ), + ("esm3", "AutoModel"): AutoClassEvidence( + symbol_path="fastplms.models.esm3.modeling_esm3.FastESM3Model", + capabilities=_BASE_CAPABILITIES | {"multimodal_generation"}, + runtime_tests=( + "tests/cpu/test_generation_contracts.py::" + "test_esm3_uses_hugging_face_initialization_and_only_retains_requested_states", + "tests/cpu/test_generation_contracts.py::test_esm3_advertised_model_logits_backward", + "tests/cpu/test_generation_contracts.py::test_esm3_loads_with_automodel", + "tests/cpu/test_generation_contracts.py::" + "test_esm3_resize_updates_sequence_input_and_output_embeddings", + "tests/cpu/test_generation_contracts.py::test_esm3_sequence_only_forward", + ), + limitations=("supervised loss is not applicable to the ESM3 base model",), + ), + ("e1", "AutoConfig"): _config("fastplms.models.e1.modeling_e1.E1Config"), + ("e1", "AutoModel"): _sequence( + "fastplms.models.e1.modeling_e1.E1Model", + _E1_TEST, + loss=False, + ), + ("e1", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.e1.modeling_e1.E1ForMaskedLM", + _E1_TEST, + loss=True, + ), + ("e1", "AutoModelForSequenceClassification"): _sequence( + "fastplms.models.e1.modeling_e1.E1ForSequenceClassification", + _E1_TEST, + loss=True, + ), + ("e1", "AutoModelForTokenClassification"): _sequence( + "fastplms.models.e1.modeling_e1.E1ForTokenClassification", + _E1_TEST, + loss=True, + ), + ("dplm", "AutoConfig"): _config( + "fastplms.models.dplm.modeling_dplm.DPLMConfig" + ), + ("dplm", "AutoModel"): _sequence( + "fastplms.models.dplm.modeling_dplm.DPLMModel", + _DPLM_TEST, + loss=False, + ), + ("dplm", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.dplm.modeling_dplm.DPLMForMaskedLM", + _DPLM_TEST, + loss=True, + ), + ("dplm", "AutoModelForSequenceClassification"): _sequence( + "fastplms.models.dplm.modeling_dplm.DPLMForSequenceClassification", + _DPLM_TEST, + loss=True, + ), + ("dplm", "AutoModelForTokenClassification"): _sequence( + "fastplms.models.dplm.modeling_dplm.DPLMForTokenClassification", + _DPLM_TEST, + loss=True, + ), + ("dplm2", "AutoConfig"): _config( + "fastplms.models.dplm2.modeling_dplm2.DPLM2Config" + ), + ("dplm2", "AutoModel"): _sequence( + "fastplms.models.dplm2.modeling_dplm2.DPLM2Model", + _DPLM_TEST, + loss=False, + ), + ("dplm2", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.dplm2.modeling_dplm2.DPLM2ForMaskedLM", + _DPLM_TEST, + loss=True, + ), + ("dplm2", "AutoModelForSequenceClassification"): _sequence( + "fastplms.models.dplm2.modeling_dplm2.DPLM2ForSequenceClassification", + _DPLM_TEST, + loss=True, + ), + ("dplm2", "AutoModelForTokenClassification"): _sequence( + "fastplms.models.dplm2.modeling_dplm2.DPLM2ForTokenClassification", + _DPLM_TEST, + loss=True, + ), + ("ankh", "AutoConfig"): _config( + "fastplms.models.ankh.modeling_ankh.FastAnkhConfig" + ), + ("ankh", "AutoModel"): _sequence( + "fastplms.models.ankh.modeling_ankh.FastAnkhModel", + _ANKH_TEST, + loss=False, + ), + ("ankh", "AutoModelForMaskedLM"): _sequence( + "fastplms.models.ankh.modeling_ankh.FastAnkhForMaskedLMExtension", + _ANKH_TEST, + loss=True, + ), + ("ankh", "AutoModelForSeq2SeqLM"): AutoClassEvidence( + symbol_path=( + "fastplms.models.ankh.modeling_ankh.FastAnkhForConditionalGeneration" + ), + capabilities=_HEAD_CAPABILITIES | {"encoder_decoder_state"}, + runtime_tests=( + "tests/cpu/test_ankh_contracts.py::" + "test_complete_t5_checkpoint_loads_clean_encoder_and_seq2seq_views", + "tests/cpu/test_ankh_contracts.py::" + "test_ankh_seq2seq_view_honors_tuple_output_and_resize", + "tests/cpu/test_ankh_contracts.py::" + "test_seq2seq_head_produces_finite_loss_and_gradients", + ), + ), + ("ankh", "AutoModelForSequenceClassification"): _sequence( + "fastplms.models.ankh.modeling_ankh.FastAnkhForSequenceClassification", + _ANKH_TEST, + loss=True, + ), + ("ankh", "AutoModelForTokenClassification"): _sequence( + "fastplms.models.ankh.modeling_ankh.FastAnkhForTokenClassification", + _ANKH_TEST, + loss=True, + ), + ("boltz2", "AutoConfig"): _config( + "fastplms.models.boltz.modeling_boltz2.Boltz2Config" + ), + ("boltz2", "AutoModel"): AutoClassEvidence( + symbol_path="fastplms.models.boltz.modeling_boltz2.Boltz2Model", + capabilities=frozenset( + {"backward", "forward", "output_flags", "return_dict", "save_reload", "tuple"} + ), + runtime_tests=( + "tests/cpu/test_structure_contracts.py::" + "test_boltz_public_forward_honors_output_controls_backward_and_reload", + ), + limitations=( + "token embedding resize and classifier-style supervised loss are not applicable " + "to a structure pipeline", + ), + ), + ("esmfold", "AutoConfig"): _config( + "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig" + ), + ("esmfold", "AutoModel"): AutoClassEvidence( + symbol_path=( + "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding" + ), + capabilities=frozenset( + { + "backward", + "forward", + "multimer_infer", + "output_flags", + "return_dict", + "save_reload", + "tuple", + } + ), + runtime_tests=( + "tests/cpu/test_structure_contracts.py::" + "test_fast_esmfold_public_forward_honors_output_controls_and_backward", + "tests/cpu/test_structure_contracts.py::" + "test_fast_esmfold_tiny_model_saves_and_reloads_exact_state", + "tests/cpu/test_structure_contracts.py::" + "test_esmfold_infer_preserves_official_multimer_contract", + ), + limitations=( + "the CPU contract injects the folding core because a complete ESMFold trunk is " + "a release-candidate checkpoint contract", + "token resize and classifier-style task loss are not applicable", + ), + ), + ("esmfold2", "AutoConfig"): _config( + "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config" + ), + ("esmfold2", "AutoModel"): AutoClassEvidence( + symbol_path="fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", + capabilities=frozenset( + {"backward", "forward", "output_flags", "return_dict", "save_reload", "tuple"} + ), + runtime_tests=( + "tests/cpu/test_structure_contracts.py::" + "test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides", + "tests/cpu/test_structure_contracts.py::" + "test_esmfold2_advertised_models_tiny_init_backward_and_save_reload", + ), + limitations=( + "token resize and classifier-style task loss are not applicable to the " + "structure pipeline", + ), + ), +} + + +def _load_symbol(path: str) -> type: + module_name, separator, symbol_name = path.rpartition(".") + if not separator: + raise AssertionError(f"Invalid symbol path: {path}") + symbol = getattr(importlib.import_module(module_name), symbol_name) + assert isinstance(symbol, type), path + return symbol + + +def _manifest_family_entries() -> dict[tuple[str, str], str]: + registry = get_model_registry() + return { + (family_id, auto_class): symbol_path + for family_id, family in registry.families.items() + for auto_class, symbol_path in family.auto_map.items() + } + + +_CONFIG_CASES = tuple( + (family_id, evidence.symbol_path) + for (family_id, auto_class), evidence in AUTOCLASS_EVIDENCE.items() + if auto_class == "AutoConfig" +) + + +@pytest.mark.parametrize(("family_id", "symbol_path"), _CONFIG_CASES) +def test_every_advertised_config_round_trips_offline( + family_id: str, + symbol_path: str, +) -> None: + config_class = _load_symbol(symbol_path) + assert issubclass(config_class, PretrainedConfig), family_id + config = config_class() + serialized = config.to_dict() + reloaded = config_class.from_dict(serialized) + + assert isinstance(reloaded, config_class) + assert reloaded.to_dict() == serialized + + +def test_autoclass_runtime_evidence_matrix_exactly_matches_all_37_entries() -> None: + manifest_entries = _manifest_family_entries() + assert len(manifest_entries) == 37 + assert set(AUTOCLASS_EVIDENCE) == set(manifest_entries) + for key, evidence in AUTOCLASS_EVIDENCE.items(): + assert evidence.symbol_path == manifest_entries[key] + _load_symbol(evidence.symbol_path) + assert evidence.capabilities + assert evidence.runtime_tests + + +def _collected_base_node_id(item: pytest.Item) -> str: + return item.nodeid.replace("\\", "/").partition("[")[0] + + +def _case_parameter_symbol(item: pytest.Item, parameter: str) -> str | None: + callspec = getattr(item, "callspec", None) + if callspec is None or parameter not in callspec.params: + return None + value = callspec.params[parameter] + if isinstance(value, str): + return value + if isinstance(value, type): + return f"{value.__module__}.{value.__qualname__}" + return None + + +def test_autoclass_runtime_evidence_targets_are_collected_cpu_tests( + request: pytest.FixtureRequest, +) -> None: + collected: dict[str, list[pytest.Item]] = {} + for item in request.session.items: + collected.setdefault(_collected_base_node_id(item), []).append(item) + + for evidence in AUTOCLASS_EVIDENCE.values(): + for node_id in evidence.runtime_tests: + cases = collected.get(node_id, []) + assert cases, f"Runtime evidence target was not collected: {node_id}" + if evidence.case_parameter is None: + continue + observed_symbols = { + symbol + for item in cases + if ( + symbol := _case_parameter_symbol(item, evidence.case_parameter) + ) + is not None + } + assert evidence.symbol_path in observed_symbols, ( + f"{node_id} was collected, but not the {evidence.case_parameter!r} case " + f"for {evidence.symbol_path!r}; observed {sorted(observed_symbols)!r}" + ) + + +def test_model_specific_automap_overrides_have_cpu_runtime_evidence() -> None: + registry = get_model_registry() + experimental_path = ( + "fastplms.models.esmfold2.modeling_esmfold2_experimental." + "ESMFold2ExperimentalModel" + ) + overridden_paths = { + symbol_path + for spec in registry.values() + for auto_class, symbol_path in spec.auto_map.items() + if symbol_path != spec.family.auto_map[auto_class] + } + assert overridden_paths == {experimental_path} + _load_symbol(experimental_path) + module = importlib.import_module("tests.cpu.test_structure_contracts") + for test_name in ( + "test_esmfold2_advertised_models_tiny_init_backward_and_save_reload", + "test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides", + ): + assert callable(getattr(module, test_name)) diff --git a/tests/cpu/test_benchmark_artifact_contracts.py b/tests/cpu/test_benchmark_artifact_contracts.py new file mode 100644 index 0000000..638ce3b --- /dev/null +++ b/tests/cpu/test_benchmark_artifact_contracts.py @@ -0,0 +1,265 @@ +"""Offline CPU contracts for local benchmark artifact identity binding.""" + +from __future__ import annotations + +import json +import xml.etree.ElementTree as ET +import pytest +from pathlib import Path +from types import SimpleNamespace + +import benchmarks.suite as benchmark_suite +from benchmarks.suite import benchmark_cases, bind_local_artifacts +from fastplms.registry import ModelSpec, get_model_registry + + +_RUNTIME_REVISION = "1" * 40 +_SOURCE_SHA256 = "2" * 64 +_RUNTIME_BUNDLE_SHA256 = "3" * 64 + + +def _write_identity_artifact(root: Path, spec: ModelSpec) -> Path: + path = root / spec.fast.repo_id.rsplit("/", maxsplit=1)[1] + path.mkdir(parents=True) + config = { + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "fastplms_checkpoint_revision": spec.artifact_checkpoint.revision, + "fastplms_weights_revision": spec.artifact_checkpoint.revision, + "fastplms_runtime_revision": _RUNTIME_REVISION, + "fastplms_source_tree_sha256": _SOURCE_SHA256, + "fastplms_runtime_bundle_sha256": _RUNTIME_BUNDLE_SHA256, + } + provenance = { + "model_id": spec.id, + "artifact_checkpoint": { + "repo_id": spec.artifact_checkpoint.repo_id, + "revision": spec.artifact_checkpoint.revision, + }, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": _RUNTIME_REVISION, + "source_tree_sha256": _SOURCE_SHA256, + "runtime_bundle_sha256": _RUNTIME_BUNDLE_SHA256, + "canonical_weights": { + "state_digest": { + "schema_version": 1, + "algorithm": "sha256", + "sha256": "4" * 64, + } + }, + } + (path / "config.json").write_text(json.dumps(config), encoding="utf-8") + (path / "provenance.json").write_text(json.dumps(provenance), encoding="utf-8") + (path / "artifact-manifest.json").write_text( + json.dumps({"config.json": "sha256:" + "5" * 64}), + encoding="utf-8", + ) + return path + + +def _stub_complete_validation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(benchmark_suite, "_validate_built_artifact", lambda *_args: None) + monkeypatch.setattr( + benchmark_suite, + "_frozen_runtime_identity", + lambda *_args: (_RUNTIME_REVISION, _SOURCE_SHA256), + ) + + +def test_local_benchmark_artifact_identity_is_path_free_and_registry_stable( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_complete_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + artifact = _write_identity_artifact(tmp_path, spec) + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=False)) + + identities = bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + case = cases[0] + assert case.model == spec.fast.repo_id + assert case.revision == spec.fast.revision + assert case.load_model == artifact.resolve() + assert case.load_revision is None + assert case.local_files_only is True + assert identities[spec.id]["runtime_revision"] == _RUNTIME_REVISION + assert identities[spec.id]["weights_revision"] == spec.artifact_checkpoint.revision + assert str(tmp_path) not in json.dumps(identities, sort_keys=True) + + +@pytest.mark.parametrize( + ("field", "stale_value"), + ( + ("fastplms_model_id", "esm2_35m"), + ("fastplms_runtime_revision", "6" * 40), + ("fastplms_source_tree_sha256", "7" * 64), + ), +) +def test_local_benchmark_artifact_rejects_swapped_or_stale_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + field: str, + stale_value: str, +) -> None: + _stub_complete_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + artifact = _write_identity_artifact(tmp_path, spec) + config_path = artifact / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config[field] = stale_value + config_path.write_text(json.dumps(config), encoding="utf-8") + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=True)) + + with pytest.raises(ValueError, match="registry/frozen source"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + +def test_local_benchmark_artifact_rejects_linked_root_and_child( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_complete_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + actual_root = tmp_path / "actual" + actual_root.mkdir() + _write_identity_artifact(actual_root, spec) + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=True)) + + linked_root = tmp_path / "linked-root" + linked_root.symlink_to(actual_root, target_is_directory=True) + with pytest.raises(ValueError, match="link or junction"): + bind_local_artifacts(cases, linked_root, source_root=tmp_path) + + linked_child_root = tmp_path / "linked-child-root" + linked_child_root.mkdir() + child = linked_child_root / spec.fast.repo_id.rsplit("/", maxsplit=1)[1] + child.symlink_to(actual_root / child.name, target_is_directory=True) + with pytest.raises(ValueError, match="Missing or invalid selected benchmark artifacts"): + bind_local_artifacts(cases, linked_child_root, source_root=tmp_path) + + +def test_local_benchmark_artifact_propagates_complete_validator_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_complete_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + _write_identity_artifact(tmp_path, spec) + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=True)) + monkeypatch.setattr( + benchmark_suite, + "_validate_built_artifact", + lambda *_args: (_ for _ in ()).throw(ValueError("digest mismatch")), + ) + + with pytest.raises(ValueError, match="digest mismatch"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + +def test_prepublication_artifact_root_benchmark_command_is_documented() -> None: + root = Path(__file__).resolve().parents[2] + for relative_name in ("docs/benchmarking.md", "benchmarks/README.md"): + text = (root / relative_name).read_text(encoding="utf-8") + assert "tools.artifacts.build_all" in text + assert "--benchmark-suite" in text + assert "benchmarks.suite" in text + assert "--artifact-root dist/hub" in text + assert "--local-files-only" in text + + +def test_suite_report_records_path_free_artifact_inventory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + output = tmp_path / "report.json" + junit = tmp_path / "junit.xml" + artifact_root = tmp_path / "artifacts" + artifact_root.mkdir() + spec = get_model_registry()["esm2_8m"] + identity = { + "model_id": spec.id, + "registry_repo_id": spec.fast.repo_id, + "registry_revision": spec.fast.revision, + "checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": "7" * 40, + "source_tree_sha256": "8" * 64, + "runtime_bundle_sha256": "9" * 64, + "canonical_state_sha256": "a" * 64, + "artifact_manifest_sha256": "b" * 64, + } + + def fake_bind(cases, root): + assert root == artifact_root + for case in cases: + case.load_model = root / "ESM2-8M" + case.load_revision = None + case.local_files_only = True + case.artifact_identity = identity + case.artifact_dependencies = {} + return {"esm2_8m": identity} + + fake_torch = SimpleNamespace( + cuda=SimpleNamespace(empty_cache=lambda: None), + ) + monkeypatch.setattr(benchmark_suite, "bind_local_artifacts", fake_bind) + monkeypatch.setattr(benchmark_suite, "_require_torch", lambda: fake_torch) + monkeypatch.setattr( + benchmark_suite, + "environment_fingerprint", + lambda _torch: {"gpu": "synthetic", "gpu_capability": [0, 0]}, + ) + monkeypatch.setattr( + benchmark_suite, + "_load_model", + lambda _case, _torch: (object(), 1.0), + ) + monkeypatch.setattr( + benchmark_suite, + "run_case", + lambda case, **_kwargs: { + "model": case.model, + "revision": case.revision, + "blocks": [], + }, + ) + + assert ( + benchmark_suite.main( + [ + "--quick", + "--family", + "esm2", + "--artifact-root", + str(artifact_root), + "--output", + str(output), + "--junit-output", + str(junit), + ] + ) + == 0 + ) + + report = json.loads(output.read_text(encoding="utf-8")) + assert report["schema_version"] == 3 + assert report["status"] == "complete" + assert report["expected_case_count"] == report["completed_case_count"] == 1 + assert report["artifact_load_mode"] == "validated_local_build" + assert report["artifacts"] == {"esm2_8m": identity} + assert report["results"][0]["artifact"] == identity + assert report["results"][0]["model"] == spec.fast.repo_id + assert report["timing_contract"] == { + "cold_compile_field": "results[].compile_ms", + "first_forward_field": "results[].first_forward_ms", + "warmup_field": "results[].warmup_samples_ms", + "warm_throughput_field": "results[].blocks", + "compile_amortized_into_throughput": False, + } + assert report["baseline_promotion_contract"]["requires_exact_environment_match"] is True + assert report["baseline_promotion_contract"]["requires_exact_artifact_inventory_match"] is True + junit_root = ET.parse(junit).getroot() + assert junit_root.attrib["failures"] == "0" + assert junit_root.find("testcase/failure") is None + assert str(tmp_path) not in json.dumps(report, sort_keys=True) diff --git a/tests/cpu/test_conversion_contracts.py b/tests/cpu/test_conversion_contracts.py new file mode 100644 index 0000000..336f18e --- /dev/null +++ b/tests/cpu/test_conversion_contracts.py @@ -0,0 +1,11 @@ +"""Mandatory tiny state-conversion contracts for complete ANKH publication.""" + +from tests.release import test_conversion_tools as conversion_contracts + + +test_ankh_transform_requires_and_preserves_complete_t5_state = ( + conversion_contracts.test_ankh_transform_requires_and_preserves_complete_t5_state +) +test_ankh_transform_rejects_encoder_only_publication_state = ( + conversion_contracts.test_ankh_transform_rejects_encoder_only_publication_state +) diff --git a/tests/cpu/test_documentation_contracts.py b/tests/cpu/test_documentation_contracts.py new file mode 100644 index 0000000..f8c5f7c --- /dev/null +++ b/tests/cpu/test_documentation_contracts.py @@ -0,0 +1,981 @@ +"""Execute curated offline examples and the Python migration snippets.""" + +from __future__ import annotations + +import ast +import json +import sqlite3 +import struct +import sys +import pytest +import torch +import transformers +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +from examples import ( + _runtime, + ankh_embeddings, + artifact_loading, + attention_switching, + e1_rag, + embedding_and_retrieval, + generation, + structure_preparation, + task_heads, + ttt, +) +from fastplms.embeddings import ( + EmbeddingRecord, + EmbeddingResult, + embed_dataset, + load_sqlite_result, + save_safetensors_result, + save_sqlite_result, +) +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esmfold2 import esmfold2_types +from tests.integration import test_ttt as ttt_contracts +from tests.unit import test_ankh_cpu_contract as ankh_contracts +from tests.unit import test_e1_cache_contract as e1_contracts +from tests.unit import test_embeddings_api as embedding_contracts + + +_ROOT = Path(__file__).resolve().parents[2] +_CURATED_EXAMPLE_CPU_CASES: dict[str, tuple[str, ...]] = { + "embedding_and_retrieval.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_embedding_and_retrieval_example_executes_with_ordered_sqlite", + ), + "attention_switching.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_attention_switching_main_executes_optimized_and_masked_fallback", + ), + "ankh_embeddings.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_ankh_embedding_example_executes_encoder_and_decoder_layers", + ), + "generation.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_dplm_branch_offline", + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_dplm2_branch_offline", + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_esm3_trace", + ), + "e1_rag.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_e1_rag_example_executes_local_msa_and_shared_persistence", + ), + "ttt.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_ttt_example_executes_seeded_adapt_save_and_reset", + ), + "structure_preparation.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_structure_preparation_example_executes_each_public_branch", + ), + "artifact_loading.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_artifact_loading_example_executes_local_only_autoconfig", + ), + "task_heads.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_task_head_example_executes_all_advertised_heads_offline", + ), + "fine_tuning.py": ( + "tests/cpu/test_peft_contracts.py::" + "test_fine_tuning_main_wires_both_tasks_without_external_io", + "tests/cpu/test_peft_contracts.py::" + "test_shipped_collators_create_tokenizer_aware_sequence_and_pair_batches", + "tests/cpu/test_peft_contracts.py::" + "test_shipped_initializer_drives_one_peft_step_and_atomic_final_reload", + ), + "binder_design_fastplms.py": ( + "tests/cpu/test_structure_contracts.py::" + "test_public_binder_workflow_pads_heterogeneous_prepared_atoms_without_truncation", + "tests/cpu/test_structure_contracts.py::" + "test_binder_example_main_wires_explicit_offline_cli_arguments", + "tests/cpu/test_structure_contracts.py::" + "test_binder_structure_loss_is_finite_and_differentiable", + ), +} + + +def _patch_transformers_imports( + monkeypatch: pytest.MonkeyPatch, + **replacements: object, +) -> None: + """Patch ``from transformers import ...`` through its lazy module proxy.""" + + proxy = ModuleType("transformers") + proxy.__dict__.update(transformers.__dict__) + proxy.__dict__.update(replacements) + proxy.__getattr__ = lambda name: getattr(transformers, name) # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "transformers", proxy) + + +def _tiny_esm3_model() -> FastESM3Model: + return FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + ).eval() + + +class _TinyTokenizer: + all_special_ids = (0, 1) + vocab_size = 16 + name_or_path = "cpu-doc-tokenizer" + + def __call__( + self, + sequences: list[str], + **_kwargs: object, + ) -> dict[str, torch.Tensor]: + rows = [[2 + (ord(character) % 10) for character in value] + [1] for value in sequences] + width = max(map(len, rows)) + input_ids = torch.tensor( # (b, l) + [row + [0] * (width - len(row)) for row in rows] + ) + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(0).long(), + } + + +def _python_blocks(path: Path) -> list[str]: + text = path.read_text(encoding="utf-8") + blocks: list[str] = [] + for section in text.split("```python")[1:]: + blocks.append(section.split("```", maxsplit=1)[0].strip()) + return blocks + + +def test_embedding_and_retrieval_example_executes_with_ordered_sqlite( + tmp_path: Path, +) -> None: + model = embedding_contracts.SyntheticEmbeddingModel() + model.embed_dataset = lambda inputs, **kwargs: embed_dataset( # type: ignore[attr-defined] + model, inputs, **kwargs + ) + output = tmp_path / "example.sqlite" + result = embedding_and_retrieval.run_embeddings( + model, + _TinyTokenizer(), + {"a": "ACD", "b": "GG"}, + output=output, + output_format="sqlite", + max_length=16, + ) + + assert [record.id for record in result] == ["a", "b"] + selected = load_sqlite_result(output, record_ids=["b", "a", "b"]) + assert [record.id for record in selected] == ["b", "a", "b"] + + +@pytest.mark.parametrize( + "retrieval_arguments", + ( + ("--select-id", "0"), + ( + "--output", + "embeddings.safetensors", + "--format", + "safetensors", + "--select-id", + "0", + ), + ), +) +def test_embedding_retrieval_selection_requires_sqlite_output( + tmp_path: Path, + retrieval_arguments: tuple[str, ...], +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + (artifact / "config.json").write_text("{}\n", encoding="utf-8") + + with pytest.raises( + SystemExit, + match="--select-id requires both --output and --format sqlite", + ): + embedding_and_retrieval.main([str(artifact), "--sequence", "ACD", *retrieval_arguments]) + + +@pytest.mark.parametrize( + ("module", "arguments"), + ( + (embedding_and_retrieval, ["artifact", "--sequence", "ACD"]), + (ankh_embeddings, ["artifact"]), + (generation, ["dplm", "artifact"]), + (e1_rag, ["artifact", "query.a3m"]), + (ttt, ["artifact", "adapted"]), + ), +) +def test_sequence_examples_share_explicit_device_dtype_contract( + module: ModuleType, + arguments: list[str], +) -> None: + parsed = module.build_parser().parse_args([*arguments, "--device", "cpu", "--dtype", "float32"]) + assert parsed.device == "cpu" + assert parsed.dtype == "float32" + device, dtype = _runtime.resolve_execution(parsed.device, parsed.dtype) + assert device == torch.device("cpu") + assert dtype is torch.float32 + + +def test_attention_switching_main_executes_optimized_and_masked_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + import torch.nn.functional as functional + + from fastplms import attention as attention_module + + model = ankh_contracts.FastAnkhModel(ankh_contracts._config(attn_backend="sdpa")).eval() + artifact = tmp_path / "ankh-artifact" + artifact.mkdir() + (artifact / "config.json").write_text("{}\n", encoding="utf-8") + model_loads: list[tuple[object, dict[str, object]]] = [] + tokenizer_loads: list[tuple[object, dict[str, object]]] = [] + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, source, **kwargs): + del cls + model_loads.append((source, kwargs)) + return model + + class FakeAutoTokenizer: + @classmethod + def from_pretrained(cls, source, **kwargs): + del cls + tokenizer_loads.append((source, kwargs)) + return _TinyTokenizer() + + sdpa_calls: list[tuple[tuple[int, ...], tuple[int, ...]]] = [] + original_sdpa = functional.scaled_dot_product_attention + + def tracked_sdpa(query, key, value, *args, **kwargs): + sdpa_calls.append((tuple(query.shape), tuple(key.shape))) + return original_sdpa(query, key, value, *args, **kwargs) + + cleared: list[bool] = [] + _patch_transformers_imports( + monkeypatch, + AutoModel=FakeAutoModel, + AutoTokenizer=FakeAutoTokenizer, + ) + monkeypatch.setattr(functional, "scaled_dot_product_attention", tracked_sdpa) + monkeypatch.setattr( + attention_module, + "clear_flex_attention_caches", + lambda: cleared.append(True), + ) + + result = attention_switching.main([str(artifact), "--backend", "sdpa"]) + output = capsys.readouterr().out + + assert result == 0 + assert sdpa_calls + assert "optimized (2," in output + assert "fallback (2," in output + assert "warning" in output and "sdpa" in output and "eager" in output + assert model.attn_backend == "sdpa" + assert model_loads == [ + ( + artifact.resolve(), + { + "trust_remote_code": True, + "local_files_only": True, + "attn_implementation": "sdpa", + "dtype": torch.float32, + }, + ) + ] + assert tokenizer_loads == [ + ( + artifact.resolve(), + {"trust_remote_code": True, "local_files_only": True}, + ) + ] + assert cleared == [True] + + +def test_attention_example_detects_and_does_not_repair_fallback_mutation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = ankh_contracts.FastAnkhModel(ankh_contracts._config(attn_backend="sdpa")).eval() + artifact = tmp_path / "ankh-artifact" + artifact.mkdir() + (artifact / "config.json").write_text("{}\n", encoding="utf-8") + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, *_args: Any, **_kwargs: Any) -> Any: + return model + + class FakeAutoTokenizer: + @classmethod + def from_pretrained(cls, *_args: Any, **_kwargs: Any) -> Any: + return object() + + output = SimpleNamespace(last_hidden_state=torch.zeros((2, 3, 8))) # (b=2, l=3, d=8) + monkeypatch.setattr( + attention_switching, + "run_optimized_attention_example", + lambda *_args, **_kwargs: output, + ) + + def mutate_during_fallback(model: Any, *_args: Any, **_kwargs: Any) -> Any: + model.attn_backend = "eager" + return output, ["configured=sdpa effective=eager"] + + monkeypatch.setattr( + attention_switching, + "run_attention_example", + mutate_during_fallback, + ) + _patch_transformers_imports( + monkeypatch, + AutoModel=FakeAutoModel, + AutoTokenizer=FakeAutoTokenizer, + ) + + with pytest.raises(RuntimeError, match="eager fallback mutated"): + attention_switching.main([str(artifact), "--backend", "sdpa"]) + assert model.attn_backend == "eager" + + +@pytest.mark.parametrize("backend", ("flash_attention_2", "flash_attention_3")) +def test_attention_switching_flash_contract_fails_before_model_loading( + backend: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + (artifact / "config.json").write_text("{}\n", encoding="utf-8") + + class ForbiddenAutoModel: + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: + del cls, args, kwargs + raise AssertionError("invalid FlashAttention execution reached model loading") + + _patch_transformers_imports(monkeypatch, AutoModel=ForbiddenAutoModel) + + with pytest.raises(SystemExit, match=rf"{backend} requires a CUDA device"): + attention_switching.main( + [ + str(artifact), + "--backend", + backend, + "--device", + "cpu", + "--dtype", + "bfloat16", + ] + ) + with pytest.raises(SystemExit, match=rf"{backend} requires --dtype bfloat16"): + attention_switching.main( + [ + str(artifact), + "--backend", + backend, + "--device", + "cuda:0", + "--dtype", + "float32", + ] + ) + + +def test_task_head_example_executes_all_advertised_heads_offline( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmForSequenceClassification, + FastEsmForTokenClassification, + ) + + artifact = tmp_path / "esm2-artifact" + + def config() -> FastEsmConfig: + return FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + num_labels=2, + pad_token_id=1, + mask_token_id=5, + eos_token_id=2, + position_embedding_type="absolute", + attn_backend="eager", + ) + + FastEsmForMaskedLM(config()).save_pretrained(artifact, safe_serialization=True) + + class TaskTokenizer: + all_special_ids = (0, 1, 2, 5) + mask_token_id = 5 + + def __call__( + self, + sequences: list[str], + **_kwargs: object, + ) -> dict[str, torch.Tensor]: + available = (3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15) + alphabet = dict(zip("ACDEFGHIKLMN", available, strict=True)) + rows = [[0, *(alphabet[residue] for residue in sequence), 2] for sequence in sequences] + width = max(map(len, rows)) + input_ids = torch.tensor( # (b, l) + [row + [1] * (width - len(row)) for row in rows] + ) + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(1).long(), + } + + class FakeMaskedLM: + @classmethod + def from_pretrained(cls, source: Any, **kwargs: Any) -> Any: + return FastEsmForMaskedLM.from_pretrained(source, **kwargs) + + class FakeSequenceClassification: + @classmethod + def from_pretrained(cls, source: Any, **kwargs: Any) -> Any: + return FastEsmForSequenceClassification.from_pretrained(source, **kwargs) + + class FakeTokenClassification: + @classmethod + def from_pretrained(cls, source: Any, **kwargs: Any) -> Any: + return FastEsmForTokenClassification.from_pretrained(source, **kwargs) + + class FakeTokenizer: + @classmethod + def from_pretrained(cls, *_args: Any, **_kwargs: Any) -> Any: + return TaskTokenizer() + + _patch_transformers_imports( + monkeypatch, + AutoModelForMaskedLM=FakeMaskedLM, + AutoModelForSequenceClassification=FakeSequenceClassification, + AutoModelForTokenClassification=FakeTokenClassification, + AutoTokenizer=FakeTokenizer, + ) + + result = task_heads.main( + [ + str(artifact), + "--sequence", + "ACDE", + "--sequence", + "AC", + "--device", + "cpu", + "--dtype", + "float32", + "--attn-backend", + "eager", + ] + ) + summary = json.loads(capsys.readouterr().out) + + assert result == 0 + assert summary["masked_lm"]["status"] == "checkpoint-provided pretrained head" + assert summary["contacts"]["status"] == "checkpoint-provided pretrained head" + assert summary["contacts"]["finite"] is True + assert summary["sequence_classification"]["status"] == ("base weights + untrained task head") + assert summary["token_classification"]["status"] == ("base weights + untrained task head") + + +def test_task_head_example_rejects_missing_or_nonfinite_checkpoint_heads() -> None: + with pytest.raises(RuntimeError, match="complete checkpoint-provided"): + task_heads._require_checkpoint_heads( + {"missing_keys": ["lm_head.decoder.weight", "esm.contact_head.regression.weight"]}, + ("lm_head.", "esm.contact_head."), + ) + with pytest.raises(RuntimeError, match="Contact predictions contained non-finite"): + task_heads._require_finite_tensor( + "Contact predictions", + torch.tensor([float("nan")]), + ) + + +def test_ankh_embedding_example_executes_encoder_and_decoder_layers() -> None: + config = ankh_contracts._config() + tokenizer = ankh_contracts._TinyTokenizer() + encoder = ankh_contracts.FastAnkhModel(config).eval() + seq2seq = ankh_contracts.FastAnkhForConditionalGeneration(ankh_contracts._config()).eval() + encoder.tokenizer = tokenizer + seq2seq.tokenizer = tokenizer + + encoder_final, encoder_all, decoder_final = ankh_embeddings.extract_ankh_layers( + encoder, + seq2seq, + tokenizer, + ["ACD", "EF"], + ["AS", "D"], + ) + + assert encoder_final.metadata["hidden_state_source"] == "encoder" + assert encoder_all[0].load_tensor().shape[0] == config.num_layers + 1 + assert decoder_final.metadata["hidden_state_source"] == "decoder" + assert decoder_final.metadata["decoder_input_fingerprint"] + + tokenizer_calls: list[tuple[str, bool]] = [] + generation_arguments: dict[str, Any] = {} + + class PromptTokenizer: + def __call__( + self, + text: str, + *, + return_tensors: str, + add_special_tokens: bool = True, + ) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + tokenizer_calls.append((text, add_special_tokens)) + ids = [[2, 3, 4, 1]] if add_special_tokens else [[9, 7]] + input_ids = torch.tensor(ids) + return { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + } + + class PromptedSeq2Seq: + config = SimpleNamespace(decoder_start_token_id=0) + device = torch.device("cpu") + + def generate(self, **kwargs: Any) -> torch.Tensor: + generation_arguments.update(kwargs) + return kwargs["decoder_input_ids"] + + generated = ankh_embeddings.generate_ankh_task( + PromptedSeq2Seq(), + PromptTokenizer(), + "ACD", + "M", + max_new_tokens=3, + ) + + assert tokenizer_calls == [("ACD", True), ("M", False)] + assert torch.equal(generated, torch.tensor([[0, 9, 7]])) + assert torch.equal(generation_arguments["input_ids"], torch.tensor([[2, 3, 4, 1]])) + assert generation_arguments["do_sample"] is False + assert generation_arguments["num_beams"] == 1 + assert generation_arguments["use_cache"] is True + assert generation_arguments["max_new_tokens"] == 3 + + +class _GenerationExampleTokenizer: + def __init__(self) -> None: + self.encoded_sequences: list[str] = [] + + def __call__(self, sequence: str, *, return_tensors: str) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + self.encoded_sequences.append(sequence) + return { + "input_ids": torch.tensor( + [[0, *([3] * len(sequence)), 2]], + dtype=torch.long, + ) + } + + def get_vocab(self) -> dict[str, int]: + return { + "": 10, + "": 11, + "": 12, + "": 20, + "": 21, + "": 22, + } + + +class _GenerationExampleModel: + device = torch.device("cpu") + + def __init__(self, *, multimodal: bool) -> None: + self.multimodal = multimodal + self.calls: list[tuple[torch.Tensor, dict[str, Any]]] = [] + + def generate(self, input_ids: torch.Tensor, **kwargs: Any) -> Any: + self.calls.append((input_ids.clone(), dict(kwargs))) + sampled = torch.randint(0, 100, (1, 1), device=input_ids.device) + output = torch.cat((input_ids, sampled), dim=1) + return {"output_tokens": output} if self.multimodal else output + + +def test_generation_example_executes_seeded_dplm_branch_offline() -> None: + generation.configure_offline() + model = _GenerationExampleModel(multimodal=False) + tokenizer = _GenerationExampleTokenizer() + caller_rng = torch.random.get_rng_state() + + first = generation.generate_dplm(model, tokenizer, length=3, steps=2, seed=19) + second = generation.generate_dplm(model, tokenizer, length=3, steps=2, seed=19) + + assert torch.equal(first, second) + assert torch.equal(torch.random.get_rng_state(), caller_rng) + assert tokenizer.encoded_sequences == ["AAA", "AAA"] + assert len(model.calls) == 2 + for input_ids, kwargs in model.calls: + assert torch.equal(input_ids, torch.tensor([[0, 3, 3, 3, 2]])) + assert kwargs == { + "max_iter": 2, + "sampling_strategy": "argmax", + "disable_resample": True, + } + assert generation.os.environ["HF_HUB_OFFLINE"] == "1" + assert generation.os.environ["TRANSFORMERS_OFFLINE"] == "1" + + +def test_generation_example_executes_seeded_dplm2_branch_offline() -> None: + generation.configure_offline() + model = _GenerationExampleModel(multimodal=True) + tokenizer = _GenerationExampleTokenizer() + caller_rng = torch.random.get_rng_state() + + first = generation.generate_dplm2(model, tokenizer, length=3, steps=2, seed=23) + second = generation.generate_dplm2(model, tokenizer, length=3, steps=2, seed=23) + + assert torch.equal(first, second) + assert torch.equal(torch.random.get_rng_state(), caller_rng) + assert len(model.calls) == 2 + expected = torch.tensor([[10, 11, 11, 11, 12, 20, 21, 21, 21, 22]]) + for input_ids, kwargs in model.calls: + assert torch.equal(input_ids, expected) + assert kwargs == { + "max_iter": 2, + "sampling_strategy": "argmax", + "unmasking_strategy": "deterministic", + } + assert generation.os.environ["HF_HUB_OFFLINE"] == "1" + assert generation.os.environ["TRANSFORMERS_OFFLINE"] == "1" + + +def test_generation_example_executes_seeded_esm3_trace() -> None: + model = _tiny_esm3_model() + first = generation.generate_esm3(model, "MK__A", steps=2, seed=19) + second = generation.generate_esm3(model, "MK__A", steps=2, seed=19) + + assert first == second + assert first.startswith("MK") and first.endswith("A") + + +def test_generation_example_preserves_every_esm3_multimodal_track() -> None: + model = _tiny_esm3_model() + request = generation.build_esm3_multimodal_request(model, "M_A") + observed: list[dict[str, torch.Tensor]] = [] + original_forward = model.forward + + def capture_forward(*args: Any, **kwargs: Any) -> Any: + observed.append( + { + name: value.detach().clone() + for name, value in kwargs.items() + if torch.is_tensor(value) + } + ) + return original_forward(*args, **kwargs) + + model.forward = capture_forward + output = generation.generate_esm3(model, request, steps=1, seed=23) + + assert torch.is_tensor(output) + assert len(observed) == 1 + assert set(request).issubset(observed[0]) + for name, expected in request.items(): + torch.testing.assert_close(observed[0][name], expected, equal_nan=True) + + +def test_e1_rag_example_executes_local_msa_and_shared_persistence( + tmp_path: Path, +) -> None: + sequence = "ACDEFG" + a3m_path = tmp_path / "query.a3m" + a3m_path.write_text(">query\nACDEFG\n>near\nACDEYG\n", encoding="utf-8") + output = tmp_path / "e1.sqlite" + model = e1_contracts.E1ForMaskedLM(e1_contracts._tiny_e1_config()).eval() + + result = e1_rag.embed_local_msa( + model, + sequence, + a3m_path, + output=output, + output_format="sqlite", + seed=7, + ) + + assert [(record.id, record.sequence) for record in result] == [ + ("0", sequence), + ("1", sequence), + ] + assert [record.id for record in load_sqlite_result(output)] == ["0", "1"] + + +def test_ttt_example_executes_seeded_adapt_save_and_reset(tmp_path: Path) -> None: + model = ttt_contracts.DummyPretrainedTTTModel(ttt_contracts.DummyPretrainedTTTConfig()) + metrics = ttt.adapt_and_save(model, "ACDE", tmp_path / "adapted", seed=7) + + assert len(metrics["losses"]) == 3 + assert (tmp_path / "adapted" / "config.json").is_file() + assert not list(tmp_path.glob(".adapted-*")) + with pytest.raises(FileExistsError, match="Refusing to overwrite"): + ttt.adapt_and_save(model, "ACDE", tmp_path / "adapted", seed=7) + + +def test_ttt_example_resets_after_a_mutating_adaptation_failure(tmp_path: Path) -> None: + class FailingModel: + def __init__(self) -> None: + self.mutated = False + self.reset_calls = 0 + + def ttt(self, **_kwargs: Any) -> None: + self.mutated = True + raise RuntimeError("adaptation failed after mutation") + + def ttt_reset(self) -> None: + self.mutated = False + self.reset_calls += 1 + + model = FailingModel() + with pytest.raises(RuntimeError, match="adaptation failed after mutation"): + ttt.adapt_and_save(model, "ACDE", tmp_path / "failed", seed=7) + + assert model.mutated is False + assert model.reset_calls == 1 + assert not (tmp_path / "failed").exists() + assert not list(tmp_path.glob(".failed-*")) + + +def test_ttt_example_cleans_staging_even_when_reset_fails(tmp_path: Path) -> None: + class ResetFailingModel: + def ttt(self, **_kwargs: Any) -> None: + raise RuntimeError("adaptation failure") + + def ttt_reset(self) -> None: + raise ValueError("reset failure") + + with pytest.raises(ValueError, match="reset failure"): + ttt.adapt_and_save( + ResetFailingModel(), + "ACDE", + tmp_path / "failed-reset", + seed=7, + ) + + assert not (tmp_path / "failed-reset").exists() + assert not list(tmp_path.glob(".failed-reset-*")) + + +def test_ttt_example_refuses_source_and_existing_destinations(tmp_path: Path) -> None: + artifact = tmp_path / "source" + artifact.mkdir() + (artifact / "config.json").write_text("{}\n", encoding="utf-8") + with pytest.raises(SystemExit, match="must not be the source artifact"): + ttt.main([str(artifact), str(artifact)]) + + destination = tmp_path / "adapted" + destination.mkdir() + with pytest.raises(SystemExit, match="Refusing to overwrite"): + ttt.main([str(artifact), str(destination)]) + + +def test_structure_preparation_example_executes_each_public_branch() -> None: + class FakeBoltz: + def predict_structure(self, **kwargs: Any) -> Any: + return kwargs + + class FakeESMFold: + def fold_protein(self, sequence: str) -> Any: + return {"sequence": sequence} + + boltz = structure_preparation.run_structure_helper(FakeBoltz(), "boltz2", "ACD", 7) + esmfold = structure_preparation.run_structure_helper(FakeESMFold(), "esmfold", "ACD", 7) + + assert boltz["seed"] == 7 + assert esmfold == {"sequence": "ACD"} + + request = structure_preparation.build_esmfold2_conditioned_complex(esmfold2_types) + assert len(request.sequences) == 5 + assert ( + sum(isinstance(sequence, esmfold2_types.ProteinInput) for sequence in request.sequences) + == 2 + ) + assert any(isinstance(sequence, esmfold2_types.RNAInput) for sequence in request.sequences) + assert any(isinstance(sequence, esmfold2_types.DNAInput) for sequence in request.sequences) + assert any(isinstance(sequence, esmfold2_types.LigandInput) for sequence in request.sequences) + protein_with_msa = request.sequences[0] + assert isinstance(protein_with_msa, esmfold2_types.ProteinInput) + assert isinstance(protein_with_msa.msa, esmfold2_types.MSA) + modified_protein = request.sequences[1] + assert isinstance(modified_protein, esmfold2_types.ProteinInput) + assert modified_protein.modifications == [esmfold2_types.Modification(position=0, ccd="MSE")] + assert request.covalent_bonds and len(request.covalent_bonds) == 1 + assert request.distogram_conditioning and len(request.distogram_conditioning) == 1 + + class FakeESMFold2: + input_types = esmfold2_types + + def prepare_structure_input(self, prepared: Any, *, seed: int) -> Any: + assert seed == 11 + if prepared.pocket is not None: + raise NotImplementedError("Pocket conditioning is not implemented.") + return prepared + + rejection = structure_preparation.verify_esmfold2_pocket_rejection( + FakeESMFold2(), + seed=11, + ) + assert "Pocket conditioning" in rejection + + +def test_artifact_loading_example_executes_local_only_autoconfig(tmp_path: Path) -> None: + config = transformers.BertConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + ) + config.save_pretrained(tmp_path) + loaded = artifact_loading.load_local_artifact(tmp_path, "AutoConfig") + + assert loaded.model_type == config.model_type + assert artifact_loading.require_local_artifact(str(tmp_path)) == tmp_path.resolve() + + +def test_migration_python_snippets_execute_against_tiny_offline_objects( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + blocks = _python_blocks(_ROOT / "docs" / "migration.md") + assert len(blocks) == 6 + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: + del cls, args, kwargs + return SimpleNamespace(set_attn_implementation=lambda _backend: None) + + class FakeSeq2Seq(FakeAutoModel): + pass + + _patch_transformers_imports( + monkeypatch, + AutoModel=FakeAutoModel, + AutoModelForSeq2SeqLM=FakeSeq2Seq, + ) + exec(blocks[0], {"model_id": "local", "__builtins__": __builtins__}) + + model = embedding_contracts.SyntheticEmbeddingModel() + model.embed_dataset = lambda inputs, **kwargs: embed_dataset( # type: ignore[attr-defined] + model, inputs, **kwargs + ) + inputs = ["ACD", "GG", "ACD"] + cwd = Path.cwd() + monkeypatch.chdir(tmp_path) + namespace = {"model": model, "inputs": inputs, "__builtins__": __builtins__} + exec(blocks[1], namespace) + result = namespace["result"] + exec(blocks[2], {"result": result, "__builtins__": __builtins__}) + + records = EmbeddingResult( + [ + EmbeddingRecord("a", "AC", torch.tensor([1.0, 2.0])), + EmbeddingRecord("b", "GG", torch.tensor([3.0, 4.0])), + ], + {"complete": True, "run_fingerprint": "migration-example-fixture"}, + ) + (tmp_path / "embeddings.sqlite").unlink() + save_safetensors_result(records, tmp_path / "embeddings") + save_sqlite_result(records, tmp_path / "embeddings.sqlite") + tensor = torch.tensor([1.0, 2.0], dtype=torch.float32) + shape = tuple(tensor.shape) + blob = struct.pack(f" Any: + return {"values": list(values), "kwargs": kwargs} + + exec( + blocks[5], + { + "encoder": FakeEmbeddingView(), + "seq2seq": FakeEmbeddingView(), + "inputs": inputs, + "__builtins__": __builtins__, + }, + ) + monkeypatch.chdir(cwd) + + +def test_capability_evidence_routes_every_curated_example_to_collected_cpu_cases() -> None: + evidence = (_ROOT / "docs" / "generated" / "capability_evidence.md").read_text(encoding="utf-8") + for name, required_cases in _CURATED_EXAMPLE_CPU_CASES.items(): + assert name in evidence + missing_cases: list[str] = [] + for nodeid in required_cases: + assert nodeid in evidence, ( + f"Capability evidence for {name} omits exact CPU node {nodeid!r}" + ) + relative_path, separator, test_name = nodeid.partition("::") + path = (_ROOT / relative_path).resolve() + try: + path.relative_to((_ROOT / "tests" / "cpu").resolve()) + except ValueError: + missing_cases.append(nodeid) + continue + if not separator or not path.is_file(): + missing_cases.append(nodeid) + continue + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + declared_tests: set[str] = set() + for node in tree.body: + if isinstance(node, (ast.AsyncFunctionDef, ast.FunctionDef)): + declared_tests.add(node.name) + elif isinstance(node, ast.Assign): + declared_tests.update( + target.id for target in node.targets if isinstance(target, ast.Name) + ) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + declared_tests.add(node.target.id) + if test_name not in declared_tests: + missing_cases.append(nodeid) + assert not missing_cases, ( + f"Capability evidence for {name} names missing CPU cases: {missing_cases!r}" + ) + assert evidence.count("`cpu_contract`") >= len(_CURATED_EXAMPLE_CPU_CASES) diff --git a/tests/cpu/test_e1_contracts.py b/tests/cpu/test_e1_contracts.py new file mode 100644 index 0000000..33818b5 --- /dev/null +++ b/tests/cpu/test_e1_contracts.py @@ -0,0 +1,116 @@ +"""Mandatory E1 cache, output, resize, and backend contracts.""" + +import pytest +import torch +from pathlib import Path +from typing import Literal + +from fastplms.embeddings import EmbeddingResult, load_sqlite_result +from tests.unit import test_e1_cache_contract as contracts + + +test_e1_cache_hit_does_not_slice_target_outputs_twice = ( + contracts.test_e1_cache_hit_does_not_slice_target_outputs_twice +) +test_e1_cache_miss_slices_every_sequence_aligned_output_alias = ( + contracts.test_e1_cache_miss_slices_every_sequence_aligned_output_alias +) +test_e1_cached_flex_dispatch_keeps_or_discards_context_by_layer = ( + contracts.test_e1_cached_flex_dispatch_keeps_or_discards_context_by_layer +) +test_e1_cached_sdpa_preserves_layer_attention_semantics = ( + contracts.test_e1_cached_sdpa_preserves_layer_attention_semantics +) +test_e1_from_pretrained_tokenizer_context_is_thread_local = ( + contracts.test_e1_from_pretrained_tokenizer_context_is_thread_local +) +test_e1_lazy_tokenizer_uses_resolved_weight_commit_per_instance = ( + contracts.test_e1_lazy_tokenizer_uses_resolved_weight_commit_per_instance +) +test_e1_legacy_backend_setter_rejects_unadvertised_backends = ( + contracts.test_e1_legacy_backend_setter_rejects_unadvertised_backends +) +test_e1_loss_bearing_head_tuples_start_with_loss_then_logits = ( + contracts.test_e1_loss_bearing_head_tuples_start_with_loss_then_logits +) +test_e1_masked_lm_resizes_input_and_output_embeddings_together = ( + contracts.test_e1_masked_lm_resizes_input_and_output_embeddings_together +) +test_e1_public_models_honor_config_output_flags_and_return_dict = ( + contracts.test_e1_public_models_honor_config_output_flags_and_return_dict +) +test_e1_public_forwards_reject_unknown_arguments = ( + contracts.test_e1_public_forwards_reject_unknown_arguments +) +test_e1_base_model_rejects_misaligned_biological_indices = ( + contracts.test_e1_base_model_rejects_misaligned_biological_indices +) +test_e1_config_round_trip_preserves_cache_policy = ( + contracts.test_e1_config_round_trip_preserves_cache_policy +) +test_e1_encoder_embedding_filters_training_only_preparer_fields = ( + contracts.test_e1_encoder_embedding_filters_training_only_preparer_fields +) + + +@pytest.mark.parametrize("pooling", ("mean", "cls")) +def test_e1_msa_embeddings_use_ordered_shared_sqlite_persistence( + tmp_path: Path, + pooling: Literal["mean", "cls"], +) -> None: + model = contracts.E1ForMaskedLM(contracts._tiny_e1_config()).eval() + output = tmp_path / f"e1-msa-{pooling}.sqlite" + + embeddings = model.embed_dataset_with_msa( + ["ACDEFG", "ACDEFG"], + msa_lookup={}, + batch_size=1, + max_len=16, + pooling=pooling, + progress=False, + embed_dtype=torch.float32, + output=output, + format="sqlite", + ) + + assert isinstance(embeddings, EmbeddingResult) + assert [(record.id, record.sequence) for record in embeddings] == [ + ("0", "ACDEFG"), + ("1", "ACDEFG"), + ] + assert all(record.load_tensor().shape == (model.config.hidden_size,) for record in embeddings) + assert embeddings.metadata["descriptor_index"] == "sqlite-records" + assert embeddings.metadata["family_adapter"]["kind"] == "e1-msa-v1" + reopened = load_sqlite_result(output) + assert [record.sequence for record in reopened] == ["ACDEFG", "ACDEFG"] + + if pooling == "cls": + with pytest.raises(ValueError, match=r"does not support pooling operations.*cls"): + model.embed_dataset( + ["ACDEFG"], + batch_size=1, + pooling="cls", + ) + + +def test_e1_mlm_loss_excludes_hf_ignore_and_padding_labels_from_normalization() -> None: + model = contracts.E1ForMaskedLM(contracts._tiny_e1_config()).eval() + batch = contracts._tiny_e1_batch() + labels = torch.tensor( # (b=1, l=4) + [[-100, 5, model.model.padding_idx, -100]], + dtype=torch.long, + ) + + output = model(**batch, labels=labels, return_dict=True) + assert output.loss is not None + valid = labels.ne(-100) & labels.ne(model.model.padding_idx) # (b, l) + expected = torch.nn.functional.cross_entropy(output.logits[valid], labels[valid]) # () + + torch.testing.assert_close(output.loss, expected) + actual_gradient = torch.autograd.grad( # (b, l, vocab) + output.loss, + output.logits, + retain_graph=True, + )[0] + expected_gradient = torch.autograd.grad(expected, output.logits)[0] # (b, l, vocab) + torch.testing.assert_close(actual_gradient, expected_gradient) diff --git a/tests/cpu/test_e1_retrieval_security.py b/tests/cpu/test_e1_retrieval_security.py new file mode 100644 index 0000000..bf24728 --- /dev/null +++ b/tests/cpu/test_e1_retrieval_security.py @@ -0,0 +1,304 @@ +"""Offline security contracts for E1's local MMseqs2 runtime.""" + +from __future__ import annotations + +import json +import subprocess +import pytest +from pathlib import Path + +from fastplms.models.e1 import retrieval + + +def _completed( + command: list[str], + *, + returncode: int = 0, + stdout: str = "", + stderr: str = "", +) -> subprocess.CompletedProcess[str]: + return subprocess.CompletedProcess(command, returncode, stdout=stdout, stderr=stderr) + + +def _inspect_payload( + *, + digest: str = retrieval.MMSEQS2_CPU_MANIFEST_DIGEST, + image_id: str = "sha256:" + "a" * 64, + architecture: str | None = None, + repository: str = retrieval.MMSEQS2_IMAGE_REPOSITORY, +) -> str: + if architecture is None: + architecture = retrieval._docker_architecture() + return json.dumps( + [ + { + "RepoDigests": [f"{repository}@{digest}"], + "Id": image_id, + "Os": "linux", + "Architecture": architecture, + } + ] + ) + + +def _image_identity() -> retrieval._DockerImageIdentity: + return retrieval._DockerImageIdentity( + reference=retrieval.DOCKER_IMAGE, + repository=retrieval.MMSEQS2_IMAGE_REPOSITORY, + version=retrieval.MMSEQS2_VERSION, + manifest_digest=retrieval.MMSEQS2_CPU_MANIFEST_DIGEST, + image_id="sha256:" + "a" * 64, + os="linux", + architecture=retrieval._docker_architecture(), + ) + + +def test_mmseqs2_default_is_cpu_offline_and_immutable() -> None: + assert retrieval.DOCKER_IMAGE == ( + "ghcr.io/soedinglab/mmseqs2:18-8cc5c@" + "sha256:41b12b0d5f41432fa1b9976123da6e2e06e7fab49a34964f3b54ec038e5845d9" + ) + searcher = retrieval.HomologueSearcher(target_db="target") + assert searcher.use_gpu is False + assert searcher.allow_pull is False + assert searcher.allow_network is False + assert searcher._docker_base_cmd()[3:5] == ["--network", "none"] + + with pytest.raises(ValueError, match="immutable @sha256"): + retrieval.HomologueSearcher( + target_db="target", + docker_image="ghcr.io/soedinglab/mmseqs2:18-8cc5c", + ) + with pytest.raises(ValueError, match="CPU-only"): + retrieval.HomologueSearcher(target_db="target", use_gpu=True) + + +@pytest.mark.parametrize( + "kwargs", + ( + {"target_db": ""}, + {"target_db": "target", "sensitivity": float("nan")}, + {"target_db": "target", "max_seqs": True}, + {"target_db": "target", "min_seq_id": -0.1}, + {"target_db": "target", "coverage": 1.1}, + {"target_db": "target", "phase_timeout": float("inf")}, + {"target_db": "target", "allow_pull": 1}, + ), +) +def test_mmseqs2_constructor_rejects_unsafe_runtime_values( + kwargs: dict[str, object], +) -> None: + with pytest.raises((TypeError, ValueError)): + retrieval.HomologueSearcher(**kwargs) + + +def test_mmseqs2_missing_image_fails_without_pull(monkeypatch: pytest.MonkeyPatch) -> None: + searcher = retrieval.HomologueSearcher(target_db="target") + calls: list[list[str]] = [] + + def fake_run(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]: + del kwargs + calls.append(command) + if command[:3] == ["docker", "image", "inspect"]: + return _completed(command, returncode=1, stderr="not found") + if command[:2] == ["docker", "pull"]: + raise AssertionError("allow_pull=False must not pull") + return _completed(command) + + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + with pytest.raises(RuntimeError, match="allow_pull=False"): + searcher._ensure_docker_image() + assert not any(command[:2] == ["docker", "pull"] for command in calls) + + +def test_mmseqs2_explicit_pull_is_reinspected_and_digest_verified( + monkeypatch: pytest.MonkeyPatch, +) -> None: + searcher = retrieval.HomologueSearcher(target_db="target", allow_pull=True) + inspect_count = 0 + calls: list[list[str]] = [] + + def fake_run(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]: + nonlocal inspect_count + del kwargs + calls.append(command) + if command[:3] == ["docker", "image", "inspect"]: + inspect_count += 1 + if inspect_count == 1: + return _completed(command, returncode=1, stderr="not found") + return _completed(command, stdout=_inspect_payload()) + return _completed(command) + + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + identity = searcher._ensure_docker_image() + + assert identity.manifest_digest == retrieval.MMSEQS2_CPU_MANIFEST_DIGEST + assert identity.image_id == "sha256:" + "a" * 64 + assert inspect_count == 2 + assert ["docker", "pull", retrieval.DOCKER_IMAGE] in calls + + +@pytest.mark.parametrize( + ("payload", "cause"), + ( + (_inspect_payload(digest="sha256:" + "b" * 64), "RepoDigests"), + (_inspect_payload(repository="example.invalid/mmseqs2"), "RepoDigests"), + (_inspect_payload(image_id="mutable-image-id"), "image ID"), + (_inspect_payload(architecture="s390x"), "architecture"), + ), +) +def test_mmseqs2_inspect_rejects_wrong_digest_and_image_id( + payload: str, + cause: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + searcher = retrieval.HomologueSearcher(target_db="target") + + def fake_run(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]: + del kwargs + if command[:3] == ["docker", "image", "inspect"]: + return _completed(command, stdout=payload) + return _completed(command) + + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + with pytest.raises(RuntimeError) as raised: + searcher._ensure_docker_image() + assert cause in str(raised.value.__cause__) + + +def test_mmseqs2_inspect_preserves_non_missing_docker_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + searcher = retrieval.HomologueSearcher(target_db="target") + + def fake_run(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]: + del kwargs + if command[:3] == ["docker", "image", "inspect"]: + return _completed(command, returncode=13, stderr="permission denied") + return _completed(command) + + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + with pytest.raises(subprocess.CalledProcessError) as raised: + searcher._ensure_docker_image() + assert raised.value.returncode == 13 + assert raised.value.stderr == "permission denied" + + +def test_mmseqs2_phase_timeout_preserves_subprocess_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + searcher = retrieval.HomologueSearcher(target_db="target", phase_timeout=7.5) + expired = subprocess.TimeoutExpired(["docker", "run"], timeout=7.5) + + def timeout(*args, **kwargs): + del args + assert kwargs["timeout"] == 7.5 + raise expired + + monkeypatch.setattr(retrieval.subprocess, "run", timeout) + with pytest.raises(TimeoutError, match=r"search.*7\.5") as raised: + searcher._run_docker_command(["docker", "run"], phase="search", check=True) + assert raised.value.__cause__ is expired + + +def test_mmseqs2_realpath_validation_rejects_symlink_escape( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + working = tmp_path / "working" + outside = tmp_path / "outside" + working.mkdir() + outside.mkdir() + (working / "escape").symlink_to(outside, target_is_directory=True) + monkeypatch.chdir(working) + searcher = retrieval.HomologueSearcher(target_db="escape/target") + + with pytest.raises(ValueError, match="resolve under"): + searcher._validate_paths_under_cwd("escape/target") + + +@pytest.mark.parametrize( + ("sequence", "seq_id"), + ( + ("ACD\n>injected", "query"), + ("ACDEFG", "query\n>injected"), + ), +) +def test_mmseqs2_rejects_fasta_and_filename_injection_before_docker( + sequence: str, + seq_id: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + searcher = retrieval.HomologueSearcher(target_db="target") + monkeypatch.setattr( + searcher, + "_ensure_docker_image", + lambda: (_ for _ in ()).throw(AssertionError("Docker must not run")), + ) + + with pytest.raises(ValueError): + searcher.search(sequence, "results", seq_id=seq_id) + + +def test_mmseqs2_result_provenance_controls_cache_reuse( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + target_prefix = Path("database/target") + target_prefix.parent.mkdir() + target_dbtype = Path(f"{target_prefix}.dbtype") + target_dbtype.write_bytes(b"db-v1") + searcher = retrieval.HomologueSearcher( + target_db=str(target_prefix), + target_db_identity="uniref30-test-revision", + ) + identity = _image_identity() + monkeypatch.setattr(searcher, "_ensure_docker_image", lambda: identity) + calls: list[list[str]] = [] + + def fake_run(command: list[str], **kwargs) -> subprocess.CompletedProcess[str]: + del kwargs + calls.append(command) + if "result2msa" in command: + index = command.index("result2msa") + output = Path(command[index + 4]) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(">query\nACDEFG\n", encoding="utf-8") + return _completed(command) + + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + result = searcher.search("ACDEFG", "results", seq_id="query") + provenance_path = Path(result).with_name(searcher._PROVENANCE_FILENAME) + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + + assert provenance["runtime"]["image_id"] == identity.image_id + assert provenance["runtime"]["manifest_digest"] == identity.manifest_digest + assert provenance["request"]["target_db"]["identity"] == "uniref30-test-revision" + assert provenance["cache_identity_sha256"] + assert all("--network" in command and "none" in command for command in calls) + + calls.clear() + monkeypatch.setattr( + searcher, + "_ensure_docker_image", + lambda: (_ for _ in ()).throw(AssertionError("valid cache must not inspect Docker")), + ) + assert searcher.search("ACDEFG", "results", seq_id="query") == result + assert calls == [] + + Path(result).write_text(">query\nTAMPERED\n", encoding="utf-8") + with pytest.raises(AssertionError, match="valid cache must not inspect Docker"): + searcher.search("ACDEFG", "results", seq_id="query") + + +def test_mmseqs2_digest_is_disclosed_in_e1_documentation() -> None: + documentation = (Path(__file__).resolve().parents[2] / "docs" / "models.md").read_text( + encoding="utf-8" + ) + assert retrieval.MMSEQS2_VERSION in documentation + assert retrieval.MMSEQS2_CPU_MANIFEST_DIGEST in documentation + assert "allow_pull=False" in documentation + assert "search-provenance.json" in documentation diff --git a/tests/cpu/test_embedding_contracts.py b/tests/cpu/test_embedding_contracts.py new file mode 100644 index 0000000..be8fd0d --- /dev/null +++ b/tests/cpu/test_embedding_contracts.py @@ -0,0 +1,388 @@ +"""Mandatory real-family, ordered, streaming, and persistent embedding contracts.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from typing import Any, ClassVar + +from fastplms.embeddings import embed_dataset, load_sqlite_result +from fastplms.models.ankh.modeling_ankh import FastAnkhModel +from fastplms.models.dplm.modeling_dplm import DPLMConfig, DPLMModel +from fastplms.models.dplm2.modeling_dplm2 import DPLM2Config, DPLM2Model +from fastplms.models.e1.modeling_e1 import E1Model +from fastplms.models.esm2.modeling_fastesm import FastEsmModel +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, +) +from tests.cpu.test_sequence_autoclass_contracts import ( + _dplm_config_values, + _esm2_config, +) +from tests.unit import test_embeddings_api as contracts +from tests.unit.test_ankh_cpu_contract import _config as _ankh_config +from tests.unit.test_e1_cache_contract import _tiny_e1_config + + +test_all_hidden_state_embeddings_trim_token_axis_and_round_trip = ( + contracts.test_all_hidden_state_embeddings_trim_token_axis_and_round_trip +) +test_all_poolers_and_output_slices = contracts.test_all_poolers_and_output_slices +test_bounded_length_bucketing_restores_input_order = ( + contracts.test_bounded_length_bucketing_restores_input_order +) +test_decoder_companions_are_fingerprinted_and_bucket_aligned = ( + contracts.test_decoder_companions_are_fingerprinted_and_bucket_aligned +) +test_decoder_embeddings_require_an_explicit_model_capability = ( + contracts.test_decoder_embeddings_require_an_explicit_model_capability +) +test_failed_safetensors_overwrite_preserves_previous_valid_generation = ( + contracts.test_failed_safetensors_overwrite_preserves_previous_valid_generation +) +test_fasta_parser_streams_without_path_read_text = ( + contracts.test_fasta_parser_streams_without_path_read_text +) +test_fasta_preserves_headers_order_and_duplicates = ( + contracts.test_fasta_preserves_headers_order_and_duplicates +) +test_full_embeddings_contain_biological_residues_only = ( + contracts.test_full_embeddings_contain_biological_residues_only +) +test_interrupted_embedding_overwrite_preserves_previous_generation = ( + contracts.test_interrupted_embedding_overwrite_preserves_previous_generation +) +test_invalid_storage_and_pooling_fail_before_input_consumption = ( + contracts.test_invalid_storage_and_pooling_fail_before_input_consumption +) +test_legacy_sqlite_converter_accepts_compact_blobs_without_pickle = ( + contracts.test_legacy_sqlite_converter_accepts_compact_blobs_without_pickle +) +test_large_streaming_inputs_use_bounded_disk_windows = ( + contracts.test_large_streaming_inputs_use_bounded_disk_windows +) +test_mapping_inputs_embed_values_with_mapping_keys_as_ids = ( + contracts.test_mapping_inputs_embed_values_with_mapping_keys_as_ids +) +test_model_state_fingerprint_rehashes_data_and_storage_alias_mutations = ( + contracts.test_model_state_fingerprint_rehashes_data_and_storage_alias_mutations +) +test_runtime_versions_are_part_of_resume_identity = ( + contracts.test_runtime_versions_are_part_of_resume_identity +) +test_tokenizer_content_changes_run_fingerprint = ( + contracts.test_tokenizer_content_changes_run_fingerprint +) +test_native_sequence_tokenizer_loader_context_is_bound_without_secret_values = ( + contracts.test_native_sequence_tokenizer_loader_context_is_bound_without_secret_values +) +test_local_artifact_identity_fills_embedding_provenance = ( + contracts.test_local_artifact_identity_fills_embedding_provenance +) +test_max_length_counts_biological_residues_not_special_tokens = ( + contracts.test_max_length_counts_biological_residues_not_special_tokens +) +test_truncate_false_rejects_overlength_inputs_before_custom_adapter_inference = ( + contracts.test_truncate_false_rejects_overlength_inputs_before_custom_adapter_inference +) +test_truncate_false_rejects_overlength_inputs_before_raw_adapter_inference = ( + contracts.test_truncate_false_rejects_overlength_inputs_before_raw_adapter_inference +) +test_persistent_resume_metadata_records_true_commit_granularity = ( + contracts.test_persistent_resume_metadata_records_true_commit_granularity +) +test_result_preserves_order_and_duplicates = contracts.test_result_preserves_order_and_duplicates +test_resume_recovers_from_authoritative_manifest_when_index_is_missing = ( + contracts.test_resume_recovers_from_authoritative_manifest_when_index_is_missing +) +test_resume_requires_matching_fingerprint = contracts.test_resume_requires_matching_fingerprint +test_open_safetensors_reader_survives_successful_overwrite = ( + contracts.test_open_safetensors_reader_survives_successful_overwrite +) +test_safetensors_manifest_rejects_shard_path_traversal = ( + contracts.test_safetensors_manifest_rejects_shard_path_traversal +) +test_safetensors_round_trip_is_lazy = contracts.test_safetensors_round_trip_is_lazy +test_safetensors_descriptor_shards_have_a_bounded_record_count = ( + contracts.test_safetensors_descriptor_shards_have_a_bounded_record_count +) +test_safetensors_tensor_corruption_is_detected_on_materialization = ( + contracts.test_safetensors_tensor_corruption_is_detected_on_materialization +) +test_safetensors_streaming_resumes_an_ordered_prefix = ( + contracts.test_safetensors_streaming_resumes_an_ordered_prefix +) +test_safetensors_generation_gc_is_dry_run_and_explicitly_exclusive = ( + contracts.test_safetensors_generation_gc_is_dry_run_and_explicitly_exclusive +) +test_sqlite_filtered_retrieval_preserves_selector_order_and_duplicates = ( + contracts.test_sqlite_filtered_retrieval_preserves_selector_order_and_duplicates +) +test_sqlite_successful_overwrite_becomes_default_and_retains_prior_run = ( + contracts.test_sqlite_successful_overwrite_becomes_default_and_retains_prior_run +) +test_interrupted_sqlite_overwrite_retains_prior_run_and_resumable_prefix = ( + contracts.test_interrupted_sqlite_overwrite_retains_prior_run_and_resumable_prefix +) +test_sqlite_first_batch_publication_is_atomic_and_hidden_run_resumes = ( + contracts.test_sqlite_first_batch_publication_is_atomic_and_hidden_run_resumes +) +test_sqlite_same_run_replacement_is_deferred_until_first_batch_commit = ( + contracts.test_sqlite_same_run_replacement_is_deferred_until_first_batch_commit +) +test_sqlite_prepublication_schema_remains_readable_and_migrates = ( + contracts.test_sqlite_prepublication_schema_remains_readable_and_migrates +) +test_sqlite_loading_and_lazy_tensor_reads_use_read_only_connections = ( + contracts.test_sqlite_loading_and_lazy_tensor_reads_use_read_only_connections +) +test_sqlite_round_trip_is_lazy_and_bf16_lossless = ( + contracts.test_sqlite_round_trip_is_lazy_and_bf16_lossless +) +test_sqlite_tensor_corruption_is_detected_on_materialization = ( + contracts.test_sqlite_tensor_corruption_is_detected_on_materialization +) +test_sqlite_streaming_resumes_an_ordered_prefix = ( + contracts.test_sqlite_streaming_resumes_an_ordered_prefix +) + + +class _TinyProteinTokenizer: + """Deterministic tokenizer for real tiny tokenizer-mode model families.""" + + pad_token_id = 1 + aa_cls_token = "" + aa_eos_token = "" + all_special_ids = (0, 1, 2, 5, 32) + name_or_path = "offline/tiny-protein-tokenizer" + model_max_length = 64 + padding_side = "right" + truncation_side = "right" + special_tokens_map: ClassVar[dict[str, str]] = { + "cls_token": "", + "eos_token": "", + "mask_token": "", + "pad_token": "", + } + _residue_ids: ClassVar[dict[str, int]] = { + residue: token_id + for residue, token_id in zip( + "ACDEFGHIKLMNPQRSTVWYX", + ( + 3, + 4, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + ), + strict=True, + ) + } + + def get_vocab(self) -> dict[str, int]: + return { + "": 0, + "": 1, + "": 2, + "": 5, + "": 32, + **self._residue_ids, + } + + def get_added_vocab(self) -> dict[str, int]: + return {} + + def num_special_tokens_to_add(self, *, pair: bool) -> int: + assert pair is False + return 2 + + def __call__( + self, + sequences: str | list[str], + **kwargs: Any, + ) -> dict[str, torch.Tensor | list[int] | list[list[int]]]: + scalar = isinstance(sequences, str) + sequence_rows = [sequences] if scalar else list(sequences) + rows: list[list[int]] = [] + for sequence in sequence_rows: + if sequence.startswith(self.aa_cls_token): + if not sequence.endswith(self.aa_eos_token): + raise ValueError("DPLM2 sequence is missing its amino-acid EOS token") + sequence = sequence[len(self.aa_cls_token) : -len(self.aa_eos_token)] + row = [0, *(self._residue_ids[residue] for residue in sequence), 2] + if kwargs.get("truncation") and kwargs.get("max_length") is not None: + row = row[: int(kwargs["max_length"])] + rows.append(row) + + width = max(map(len, rows)) + padded = [row + [self.pad_token_id] * (width - len(row)) for row in rows] + if kwargs.get("return_tensors") != "pt": + return {"input_ids": padded[0] if scalar else padded} + input_ids = torch.tensor(padded, dtype=torch.long) + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(self.pad_token_id).long(), + } + + +def _real_family_model( + family: str, +) -> ( + FastEsmModel + | ESMplusplusModel + | DPLMModel + | DPLM2Model + | FastAnkhModel + | FastESM3Model + | E1Model +): + tokenizer = _TinyProteinTokenizer() + if family == "esm2": + model = FastEsmModel(_esm2_config()) + elif family == "esm_plusplus": + model = ESMplusplusModel( + ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + dropout=0.0, + pad_token_id=1, + mask_token_id=5, + attn_backend="eager", + ) + ) + elif family == "dplm": + model = DPLMModel(DPLMConfig(**_dplm_config_values(33))) + elif family == "dplm2": + from tests.cpu.test_sequence_autoclass_contracts import _dplm2_config_values + + model = DPLM2Model(DPLM2Config(**_dplm2_config_values())) + elif family == "ankh": + model = FastAnkhModel(_ankh_config(num_layers=1, num_decoder_layers=1)) + elif family == "esm3": + return FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + ).eval() + elif family == "e1": + return E1Model(_tiny_e1_config()).eval() + else: + raise AssertionError(f"Unknown embedding family: {family}") + model.tokenizer = tokenizer + return model.eval() + + +@pytest.mark.parametrize( + ("family", "persist"), + ( + ("esm2", True), + ("esm_plusplus", False), + ("dplm", False), + ("dplm2", False), + ("ankh", False), + ("esm3", False), + ("e1", False), + ), +) +def test_every_real_sequence_family_uses_ordered_biological_embedding_path( + family: str, + persist: bool, + tmp_path: Path, +) -> None: + model = _real_family_model(family) + sequences = ["ACD", "G", "ACD"] + output = tmp_path / "real-family.sqlite" if persist else None + + result = model.embed_dataset( # record tensors: (l_i, d) + sequences, + batch_size=3, + full_embeddings=True, + output=output, + format="sqlite", + ) + + assert [record.id for record in result] == ["0", "1", "2"] + assert [record.sequence for record in result] == sequences + assert [record.load_tensor().shape[0] for record in result] == [3, 1, 3] + assert all(torch.isfinite(record.load_tensor()).all() for record in result) + torch.testing.assert_close(result[0].load_tensor(), result[2].load_tensor()) + assert result.metadata["residue_mask_policy"] == "biological-residues-only" + + if output is not None: + reopened = load_sqlite_result(output) + assert [record.sequence for record in reopened] == sequences + for source, restored in zip(result, reopened, strict=True): + torch.testing.assert_close( + restored.load_tensor(), + source.load_tensor(), + rtol=0.0, + atol=0.0, + ) + + +def test_generator_inputs_are_consumed_once_and_keep_stable_order() -> None: + consumed: list[str] = [] + + def sequences(): + for sequence in ("A", "CCCC", "GG"): + consumed.append(sequence) + yield sequence + + result = embed_dataset( # pooled record tensors: (d,) + contracts.SyntheticEmbeddingModel(), + sequences(), + batch_size=2, + batch_window_size=3, + ) + + assert consumed == ["A", "CCCC", "GG"] + assert [record.sequence for record in result] == consumed + + +test_strict_embedding_controls_fail_before_consuming_inputs = ( + contracts.test_strict_embedding_controls_fail_before_consuming_inputs +) +test_decoder_input_ids_require_nonempty_2d_int32_or_int64 = ( + contracts.test_decoder_input_ids_require_nonempty_2d_int32_or_int64 +) +test_decoder_attention_masks_are_exact_finite_binary_shapes = ( + contracts.test_decoder_attention_masks_are_exact_finite_binary_shapes +) +test_embedding_batch_adapter_outputs_are_validated = ( + contracts.test_embedding_batch_adapter_outputs_are_validated +) +test_embedding_value_types_fail_closed = contracts.test_embedding_value_types_fail_closed +test_pagerank_controls_require_finite_valid_values = ( + contracts.test_pagerank_controls_require_finite_valid_values +) +test_tensor_sha256_uses_bounded_chunks_and_preserves_legacy_digest = ( + contracts.test_tensor_sha256_uses_bounded_chunks_and_preserves_legacy_digest +) +test_sqlite_lazy_references_are_absolute_after_cwd_change = ( + contracts.test_sqlite_lazy_references_are_absolute_after_cwd_change +) diff --git a/tests/cpu/test_environment_contract.py b/tests/cpu/test_environment_contract.py new file mode 100644 index 0000000..718f7e5 --- /dev/null +++ b/tests/cpu/test_environment_contract.py @@ -0,0 +1,211 @@ +"""Environment invariants for the mandatory CPU confidence lane.""" + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import textwrap +import huggingface_hub +import huggingface_hub.file_download +import pytest +import safetensors.torch +import torch +import transformers +from pathlib import Path + +from tests.cpu.resource_telemetry import aggregate_process_memory + + +def test_locked_cpu_runtime() -> None: + assert sys.version_info[:2] == (3, 12) + assert torch.__version__.split("+", maxsplit=1)[0] == "2.13.0" + assert transformers.__version__ == "5.13.0" + assert not torch.cuda.is_available() + assert os.environ["HF_HUB_OFFLINE"] == "1" + assert os.environ["TRANSFORMERS_OFFLINE"] == "1" + assert os.environ["HF_DATASETS_OFFLINE"] == "1" + assert os.environ["PYTEST_XDIST_AUTO_NUM_WORKERS"] == "4" + assert os.environ["FASTPLMS_CPU_BOOTSTRAPPED"] == "1" + assert os.environ["FASTPLMS_CPU_CACHE_STARTED_EMPTY"] == "1" + assert torch.get_num_threads() == 1 + assert torch.get_num_interop_threads() == 1 + for variable in ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", + ): + assert os.environ[variable] == "1" + + +def test_all_runtime_caches_are_fresh_and_task_scoped() -> None: + variables = ( + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "TRANSFORMERS_CACHE", + "HF_DATASETS_CACHE", + "TORCH_HOME", + "TORCH_EXTENSIONS_DIR", + "TORCHINDUCTOR_CACHE_DIR", + "TRITON_CACHE_DIR", + "XDG_CACHE_HOME", + ) + paths = [Path(os.environ[variable]).resolve() for variable in variables] + assert all(path.is_dir() for path in paths) + common_root = Path(os.path.commonpath(paths)) + assert common_root.name.startswith("fastplms-cpu-contract-cache-") + + +def test_network_guard_is_active() -> None: + with pytest.raises(RuntimeError, match="Network access is forbidden"): + socket.getaddrinfo("huggingface.co", 443) + with ( + socket.socket() as client, + pytest.raises(RuntimeError, match="Network access is forbidden"), + ): + client.connect_ex(("huggingface.co", 443)) + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as client: + with pytest.raises(RuntimeError, match="Network access is forbidden"): + client.sendto(b"blocked", ("127.0.0.1", 9)) + if hasattr(client, "sendmsg"): + with pytest.raises(RuntimeError, match="Network access is forbidden"): + client.sendmsg([b"blocked"], [], 0, ("127.0.0.1", 9)) + with pytest.raises(RuntimeError, match="Network access is forbidden"): + huggingface_hub.hf_hub_download("org/model", "config.json") + with pytest.raises(RuntimeError, match="Network access is forbidden"): + huggingface_hub.file_download.hf_hub_download("org/model", "config.json") + + +def test_optimized_subprocess_inherits_socket_and_hub_guards() -> None: + script = textwrap.dedent( + """ + import os + import socket + + import huggingface_hub + + assert os.environ["FASTPLMS_CPU_BOOTSTRAPPED"] == "1" + for operation in ( + lambda: socket.getaddrinfo("huggingface.co", 443), + lambda: huggingface_hub.hf_hub_download("org/model", "config.json"), + ): + try: + operation() + except RuntimeError as error: + assert "Network access is forbidden" in str(error) + else: + raise AssertionError("Inherited CPU network guard was not active") + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as client: + for operation in ( + lambda: client.sendto(b"blocked", ("127.0.0.1", 9)), + lambda: client.sendmsg([b"blocked"], [], 0, ("127.0.0.1", 9)), + ): + try: + operation() + except RuntimeError as error: + assert "Network access is forbidden" in str(error) + else: + raise AssertionError("Inherited CPU UDP guard was not active") + """ + ) + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[2], + env=dict(os.environ), + check=False, + capture_output=True, + text=True, + timeout=5, + ) + assert completed.returncode == 0, completed.stderr + + +def test_container_process_spawns_are_blocked_at_python_startup() -> None: + operations = ( + lambda: subprocess.run(["docker", "version"], check=False), + lambda: subprocess.Popen(["podman", "info"]), + lambda: subprocess.run("docker compose version", shell=True, check=False), + lambda: subprocess.run(["sh", "-lc", "buildx version"], check=False), + lambda: os.system("sudo docker buildx version"), + ) + for operation in operations: + with pytest.raises(RuntimeError, match="Container execution is forbidden"): + operation() + + +def test_repository_checkpoint_reads_are_structurally_blocked() -> None: + workspace = Path(__file__).resolve().parents[2] + checkpoints = sorted((workspace / "tests/goldens").glob("*.safetensors")) + assert checkpoints, "The release tree must contain pinned golden tensors" + checkpoint = checkpoints[0] + + for operation in ( + lambda: checkpoint.open("rb"), + lambda: torch.load(checkpoint, map_location="cpu"), + lambda: safetensors.torch.load_file(checkpoint, device="cpu"), + ): + with pytest.raises(RuntimeError, match="checkpoint path"): + operation() + + +def test_reference_and_submodule_reads_are_blocked_during_test_execution() -> None: + blocked = Path(__file__).resolve().parents[2] / "vendor" / "upstream" / "README.md" + with pytest.raises(RuntimeError, match="submodule/reference path"): + blocked.open(encoding="utf-8") + + +def test_open_guard_resolves_relative_paths_against_dir_fd(tmp_path: Path) -> None: + harmless = tmp_path / "official" + harmless.mkdir() + parent_fd = os.open(tmp_path, os.O_RDONLY) + try: + harmless_fd = os.open("official", os.O_RDONLY, dir_fd=parent_fd) + os.close(harmless_fd) + finally: + os.close(parent_fd) + + workspace = Path(__file__).resolve().parents[2] + workspace_fd = os.open(workspace, os.O_RDONLY) + try: + with pytest.raises(RuntimeError, match="submodule/reference path"): + os.open("official", os.O_RDONLY, dir_fd=workspace_fd) + finally: + os.close(workspace_fd) + + +def test_resource_accounting_includes_probe_children_without_double_counting_workers() -> None: + evidence = aggregate_process_memory( + { + "process_id": "controller", + "pid": 100, + "role": "controller", + "process_peak_rss_bytes": 100, + # This contains the same workers below and must remain unaccounted. + "waited_children_peak_rss_bytes": 10_000, + }, + ( + { + "process_id": "gw1", + "pid": 102, + "role": "worker", + "process_peak_rss_bytes": 400, + "waited_children_peak_rss_bytes": 500, + }, + { + "process_id": "gw0", + "pid": 101, + "role": "worker", + "process_peak_rss_bytes": 200, + "waited_children_peak_rss_bytes": 300, + }, + ), + ) + + assert evidence["aggregate_peak_rss_bytes"] == 1_500 + assert evidence["temporal_upper_bound_rss_bytes"] == 1_500 + assert evidence["accounting_mode"] == "xdist-conservative-process-tree-upper-bound" + assert evidence["budget_enforced"] is False + assert [record["process_id"] for record in evidence["workers"]] == ["gw0", "gw1"] diff --git a/tests/cpu/test_generation_contracts.py b/tests/cpu/test_generation_contracts.py new file mode 100644 index 0000000..d38eae5 --- /dev/null +++ b/tests/cpu/test_generation_contracts.py @@ -0,0 +1,183 @@ +"""Mandatory tiny DPLM, DPLM2, and ESM3 generation contracts.""" + +from __future__ import annotations + +import pytest +import torch + +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from tests.integration import test_dplm_generation as dplm_contracts +from tests.integration import test_esm3 as esm3_contracts + + +_integration_dplm_config = dplm_contracts._common_config + + +def _cpu_dplm_config(vocab_size: int) -> dict[str, object]: + values = _integration_dplm_config(vocab_size) + values.update( + { + "hidden_size": 8, + "num_attention_heads": 2, + "intermediate_size": 16, + "max_position_embeddings": 16, + # DPLM2 intentionally advertises SDPA only; DPLM retains its + # eager-path coverage in this shared tiny integration fixture. + "attn_backend": "sdpa" if vocab_size == 64 else "eager", + } + ) + return values + + +def _cpu_small_esm3_config() -> FastESM3Config: + return FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + + +def _cpu_small_esm3_model() -> FastESM3Model: + return FastESM3Model(_cpu_small_esm3_config()).eval() + + +# The integration contracts look up these helpers when each test executes. +# Override them in this positive CPU allowlist so the gate never grows into a +# benchmark-sized model while retaining the exact public behavior assertions. +esm3_contracts._small_config = _cpu_small_esm3_config +esm3_contracts._small_model = _cpu_small_esm3_model +dplm_contracts._common_config = _cpu_dplm_config + +test_dplm_argmax_generation_preserves_fixed_positions = ( + dplm_contracts.test_dplm_argmax_generation_preserves_fixed_positions +) +test_dplm_automodel_rejects_decoder_cache_contracts = ( + dplm_contracts.test_dplm_automodel_rejects_decoder_cache_contracts +) +test_dplm_masked_lm_rejects_decoder_and_cross_attention_arguments = ( + dplm_contracts.test_dplm_masked_lm_rejects_decoder_and_cross_attention_arguments +) +test_dplm_task_heads_honor_config_and_explicit_return_dict = ( + dplm_contracts.test_dplm_task_heads_honor_config_and_explicit_return_dict +) +test_dplm2_argmax_generation_preserves_modalities_and_fixed_positions = ( + dplm_contracts.test_dplm2_argmax_generation_preserves_modalities_and_fixed_positions +) +test_dplm2_automodel_infers_official_multimodal_types_and_returns_pooling = ( + dplm_contracts.test_dplm2_automodel_infers_official_multimodal_types_and_returns_pooling +) +test_generation_rejects_invalid_controls = dplm_contracts.test_generation_rejects_invalid_controls +test_masked_lm_resize_updates_input_and_output_projections = ( + dplm_contracts.test_masked_lm_resize_updates_input_and_output_projections +) +test_seeded_stochastic_generation_is_repeatable = ( + dplm_contracts.test_seeded_stochastic_generation_is_repeatable +) + +test_esm3_accepts_function_tokens_argument = ( + esm3_contracts.test_esm3_accepts_function_tokens_argument +) +test_esm3_generation_preserves_every_supported_conditioning_track = ( + esm3_contracts.test_esm3_generation_preserves_every_supported_conditioning_track +) +test_esm3_generation_none_num_steps_uses_mask_count = ( + esm3_contracts.test_esm3_generation_none_num_steps_uses_mask_count +) +test_esm3_generation_rejects_noninteger_num_steps = ( + esm3_contracts.test_esm3_generation_rejects_noninteger_num_steps +) +test_esm3_generation_rejects_nonpositive_num_steps = ( + esm3_contracts.test_esm3_generation_rejects_nonpositive_num_steps +) +test_esm3_generation_rejects_unknown_or_ambiguous_inputs = ( + esm3_contracts.test_esm3_generation_rejects_unknown_or_ambiguous_inputs +) +test_esm3_loads_with_automodel = esm3_contracts.test_esm3_loads_with_automodel +test_esm3_rejects_attention_mask_row_without_a_valid_key = ( + esm3_contracts.test_esm3_rejects_attention_mask_row_without_a_valid_key +) +test_esm3_resize_updates_sequence_input_and_output_embeddings = ( + esm3_contracts.test_esm3_resize_updates_sequence_input_and_output_embeddings +) +test_esm3_repeated_save_removes_stale_runtime_outputs = ( + esm3_contracts.test_esm3_repeated_save_removes_stale_runtime_outputs +) +test_esm3_saved_bridge_rejects_poisoned_archive = ( + esm3_contracts.test_esm3_saved_bridge_rejects_poisoned_archive +) +test_esm3_saved_bridge_rejects_preimported_runtime_mismatch = ( + esm3_contracts.test_esm3_saved_bridge_rejects_preimported_runtime_mismatch +) +test_esm3_saved_bridge_reuses_same_runtime_in_process = ( + esm3_contracts.test_esm3_saved_bridge_reuses_same_runtime_in_process +) +test_esm3_saved_model_loads_without_installed_fastplms = ( + esm3_contracts.test_esm3_saved_model_loads_without_installed_fastplms +) +test_esm3_saved_runtime_archive_is_fixed_bounded_and_deterministic = ( + esm3_contracts.test_esm3_saved_runtime_archive_is_fixed_bounded_and_deterministic +) +test_esm3_saved_runtime_rejects_allowlisted_symlink = ( + esm3_contracts.test_esm3_saved_runtime_rejects_allowlisted_symlink +) +test_esm3_saved_runtime_rejects_missing_allowlisted_file = ( + esm3_contracts.test_esm3_saved_runtime_rejects_missing_allowlisted_file +) +test_esm3_saved_runtime_rejects_noncanonical_allowlist_paths = ( + esm3_contracts.test_esm3_saved_runtime_rejects_noncanonical_allowlist_paths +) +test_esm3_saved_runtime_rejects_oversize_allowlisted_file = ( + esm3_contracts.test_esm3_saved_runtime_rejects_oversize_allowlisted_file +) +test_esm3_saved_runtime_rejects_oversize_total = ( + esm3_contracts.test_esm3_saved_runtime_rejects_oversize_total +) +test_esm3_seeded_generation_is_repeatable_and_preserves_context = ( + esm3_contracts.test_esm3_seeded_generation_is_repeatable_and_preserves_context +) +test_esm3_uses_hugging_face_initialization_and_only_retains_requested_states = ( + esm3_contracts.test_esm3_uses_hugging_face_initialization_and_only_retains_requested_states +) + + +def test_esm3_sequence_only_forward() -> None: + model = _cpu_small_esm3_model() + batch = model.tokenize_sequences( # each token track: (b=2, l) + ["MKTAYIAKQ", "GGGG"], + device=model.device, + ) + + with torch.inference_mode(): + output = model(**batch) # sequence logits: (b, l, vocab) + + assert output.logits is not None + assert output.logits.shape == (*batch["input_ids"].shape, model.config.vocab_size) + assert output.last_hidden_state.shape == ( + *batch["input_ids"].shape, + model.config.hidden_size, + ) + assert output.structure_logits.shape[-1] == 4096 + assert output.function_logits.shape[-2:] == (8, 260) + assert output.residue_logits.shape[-1] == 1478 + assert torch.isfinite(output.logits).all() + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model(**batch, unexpected_cpu_contract=True) + + +def test_esm3_advertised_model_logits_backward() -> None: + model = esm3_contracts._small_model().train() + batch = model.tokenize_sequences(["MKT", "GG"], device=model.device) # tracks: (b=2, l) + output = model(**batch, return_dict=True) # sequence logits: (b, l, vocab) + loss = output.sequence_logits.float().square().mean() # () + + assert torch.isfinite(loss) + loss.backward() + gradients = [ + parameter.grad + for parameter in model.parameters() + if parameter.requires_grad and parameter.grad is not None + ] + assert gradients + assert all(torch.isfinite(gradient).all() for gradient in gradients) diff --git a/tests/cpu/test_hf_base_model_contracts.py b/tests/cpu/test_hf_base_model_contracts.py new file mode 100644 index 0000000..6744396 --- /dev/null +++ b/tests/cpu/test_hf_base_model_contracts.py @@ -0,0 +1,379 @@ +"""Hugging Face backbone, prefix-loading, and encoder-only CPU contracts.""" + +from __future__ import annotations + +import json +import pytest +import torch +from collections.abc import Callable +from pathlib import Path +from transformers import PretrainedConfig, PreTrainedModel + +from fastplms.models.ankh.modeling_ankh import FastAnkhConfig, FastAnkhModel +from fastplms.models.dplm.modeling_dplm import ( + FAST_DPLM_ENCODER, + DPLMConfig, + DPLMForMaskedLM, + DPLMForSequenceClassification, + DPLMForTokenClassification, + DPLMModel, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + FAST_DPLM2_ENCODER, + DPLM2Config, + DPLM2ForMaskedLM, + DPLM2ForSequenceClassification, + DPLM2ForTokenClassification, + DPLM2Model, +) +from fastplms.models.esm2.modeling_fastesm import ( + FAST_ESM_ENCODER, + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmForSequenceClassification, + FastEsmForTokenClassification, + FastEsmModel, +) +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, +) + + +def _esm2_config() -> FastEsmConfig: + return FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=5, + num_labels=3, + position_embedding_type="absolute", + add_pooling_layer=False, + attn_backend="eager", + use_cache=False, + ) + + +def _dplm_config() -> DPLMConfig: + return DPLMConfig( + vocab_size=33, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + bos_token_id=0, + eos_token_id=2, + mask_token_id=32, + num_labels=3, + position_embedding_type="rotary", + add_pooling_layer=False, + attn_backend="eager", + use_cache=False, + ) + + +def _dplm2_config() -> DPLM2Config: + return DPLM2Config( + vocab_size=64, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + bos_token_id=0, + eos_token_id=2, + mask_token_id=32, + num_labels=3, + position_embedding_type="rotary", + add_pooling_layer=False, + attn_backend="sdpa", + use_cache=False, + ) + + +def _esmc_config() -> ESMplusplusConfig: + return ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + dropout=0.0, + pad_token_id=1, + mask_token_id=5, + attn_backend="eager", + ) + + +def _ankh_config() -> FastAnkhConfig: + return FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + num_decoder_layers=1, + dropout_rate=0.0, + pad_token_id=0, + eos_token_id=1, + decoder_start_token_id=0, + attn_backend="eager", + use_cache=False, + ) + + +def _esm3_config() -> FastESM3Config: + return FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + + +_BASE_MODEL_CASES: tuple[ + tuple[ + str, + Callable[[], PretrainedConfig], + type[PreTrainedModel], + type[PreTrainedModel], + ], + ..., +] = ( + ("esm2-base", _esm2_config, FAST_ESM_ENCODER, FastEsmModel), + ("esm2-mlm", _esm2_config, FAST_ESM_ENCODER, FastEsmForMaskedLM), + ( + "esm2-sequence", + _esm2_config, + FAST_ESM_ENCODER, + FastEsmForSequenceClassification, + ), + ("esm2-token", _esm2_config, FAST_ESM_ENCODER, FastEsmForTokenClassification), + ("dplm-base", _dplm_config, FAST_DPLM_ENCODER, DPLMModel), + ("dplm-mlm", _dplm_config, FAST_DPLM_ENCODER, DPLMForMaskedLM), + ( + "dplm-sequence", + _dplm_config, + FAST_DPLM_ENCODER, + DPLMForSequenceClassification, + ), + ("dplm-token", _dplm_config, FAST_DPLM_ENCODER, DPLMForTokenClassification), + ("dplm2-base", _dplm2_config, FAST_DPLM2_ENCODER, DPLM2Model), + ("dplm2-mlm", _dplm2_config, FAST_DPLM2_ENCODER, DPLM2ForMaskedLM), + ( + "dplm2-sequence", + _dplm2_config, + FAST_DPLM2_ENCODER, + DPLM2ForSequenceClassification, + ), + ( + "dplm2-token", + _dplm2_config, + FAST_DPLM2_ENCODER, + DPLM2ForTokenClassification, + ), +) + + +def _assert_exact_state(actual: torch.nn.Module, expected: torch.nn.Module) -> None: + actual_state = actual.state_dict() # parameter name -> checkpoint-shaped tensor + expected_state = expected.state_dict() # parameter name -> checkpoint-shaped tensor + assert actual_state.keys() == expected_state.keys() + for name, expected_tensor in expected_state.items(): + torch.testing.assert_close( + actual_state[name], + expected_tensor, + rtol=0.0, + atol=0.0, + msg=lambda message, name=name: f"{name}: {message}", + ) + + +@pytest.mark.parametrize( + ("case_id", "config_factory", "encoder_class", "wrapper_class"), + _BASE_MODEL_CASES, + ids=[case[0] for case in _BASE_MODEL_CASES], +) +def test_advertised_wrappers_expose_and_prefix_load_the_esm_base_model( + case_id: str, + config_factory: Callable[[], PretrainedConfig], + encoder_class: type[PreTrainedModel], + wrapper_class: type[PreTrainedModel], + tmp_path: Path, +) -> None: + del case_id + encoder = encoder_class(config_factory()).eval() + encoder_dir = tmp_path / "encoder" + encoder.save_pretrained(encoder_dir, safe_serialization=True) + + wrapper, loading_info = wrapper_class.from_pretrained( + encoder_dir, + local_files_only=True, + output_loading_info=True, + ) + wrapper.eval() + + assert wrapper.base_model is wrapper.esm + assert wrapper.base_model is not wrapper + _assert_exact_state(wrapper.base_model, encoder) + assert not loading_info["unexpected_keys"] + assert not any(key.startswith("esm.") for key in loading_info["missing_keys"]) + assert not loading_info["mismatched_keys"] + + wrapper_dir = tmp_path / "wrapper" + wrapper.save_pretrained(wrapper_dir, safe_serialization=True) + reloaded = wrapper_class.from_pretrained(wrapper_dir, local_files_only=True).eval() + assert reloaded.base_model is reloaded.esm + assert reloaded.base_model is not reloaded + _assert_exact_state(reloaded, wrapper) + + +_PACKAGE_RESAVE_CASES: tuple[ + tuple[str, Callable[[], PretrainedConfig], type[PreTrainedModel], bool], ... +] = ( + ("esm2", _esm2_config, FastEsmModel, False), + ("dplm", _dplm_config, DPLMModel, False), + ("dplm2", _dplm2_config, DPLM2Model, False), + ("esmc", _esmc_config, ESMplusplusModel, False), + ("ankh", _ankh_config, FastAnkhModel, False), + ("esm3", _esm3_config, FastESM3Model, True), +) + + +@pytest.mark.parametrize( + ("case_id", "config_factory", "model_class", "expected_remote_code"), + _PACKAGE_RESAVE_CASES, + ids=[case[0] for case in _PACKAGE_RESAVE_CASES], +) +def test_package_models_use_hf_local_semantics_across_save_resave( + case_id: str, + config_factory: Callable[[], PretrainedConfig], + model_class: type[PreTrainedModel], + expected_remote_code: bool, + tmp_path: Path, +) -> None: + del case_id + model = model_class(config_factory()).eval() + assert model.is_remote_code() is expected_remote_code + + first_path = tmp_path / "first" + model.save_pretrained(first_path, safe_serialization=True) + reloaded = model_class.from_pretrained(first_path, local_files_only=True).eval() + assert reloaded.is_remote_code() is expected_remote_code + _assert_exact_state(reloaded, model) + + second_path = tmp_path / "second" + reloaded.save_pretrained(second_path, safe_serialization=True) + resaved = model_class.from_pretrained(second_path, local_files_only=True).eval() + assert resaved.is_remote_code() is expected_remote_code + _assert_exact_state(resaved, model) + for checked_model in (model, reloaded, resaved): + assert None not in (getattr(checked_model.config, "auto_map", None) or {}) + + +def test_dplm_rejects_mixed_batch_attention_mask_broadcasting() -> None: + model = DPLMModel(_dplm_config()).eval() + input_ids = torch.tensor([[0, 6, 2, 1], [0, 7, 8, 2]]) # (b=2, l=4) + + for malformed_mask in ( + torch.ones(1, 4, dtype=torch.long), + torch.ones(2, 3, dtype=torch.long), + ): + with pytest.raises( + ValueError, + match=r"attention_mask must have shape \(2, 4\)", + ): + model(input_ids=input_ids, attention_mask=malformed_mask) + + +def test_dplm2_config_rejects_decoder_and_cross_attention_modes() -> None: + for unsupported_flag in ("is_decoder", "add_cross_attention"): + with pytest.raises(ValueError, match="DPLM2 is encoder-only"): + DPLM2Config(**{unsupported_flag: True}) + + +def test_dplm2_legacy_cache_config_warns_once_and_new_artifacts_disable_cache( + tmp_path: Path, +) -> None: + with pytest.warns(UserWarning, match="normalizing use_cache to False") as warning_records: + config = DPLM2Config(use_cache=True) + + assert len(warning_records) == 1 + assert config.use_cache is False + assert config.is_decoder is False + assert config.add_cross_attention is False + + config.save_pretrained(tmp_path) + serialized = json.loads((tmp_path / "config.json").read_text(encoding="utf-8")) + assert serialized["use_cache"] is False + assert serialized["is_decoder"] is False + assert serialized["add_cross_attention"] is False + + reloaded = DPLM2Config.from_pretrained(tmp_path, local_files_only=True) + assert reloaded.use_cache is False + + +def test_dplm2_public_forwards_reject_cache_and_cross_attention_arguments() -> None: + input_ids = torch.tensor([[0, 6, 7, 2]]) # (b=1, l=4) + for model_class in ( + DPLM2Model, + DPLM2ForMaskedLM, + DPLM2ForSequenceClassification, + DPLM2ForTokenClassification, + ): + model = model_class(_dplm2_config()).eval() + for argument, value in ( + ("use_cache", True), + ("past_key_values", ((torch.zeros(1), torch.zeros(1)),)), + ("encoder_hidden_states", torch.zeros(1, 2, 8)), + ): + with pytest.raises(TypeError, match=argument): + model(input_ids=input_ids, **{argument: value}) + + +def test_esmc_sequence_id_is_authoritative_for_chain_and_padding_masks() -> None: + model = ESMplusplusModel(_esmc_config()).eval() + input_ids = torch.tensor([[0, 3, 4, 5, 1, 1]]) # (b=1, l=6) + sequence_id = torch.tensor([[0, 0, 1, 1, -1, -1]]) # (b, l) + + with torch.inference_mode(): + expected = model( # (b, l, d=8) + input_ids=input_ids, + sequence_id=sequence_id, + ).last_hidden_state + actual = model( # (b, l, d) + input_ids=input_ids, + sequence_id=sequence_id, + # The official Biohub contract ignores this mask whenever + # sequence_id is present; sequence_id itself carries padding. + attention_mask=torch.zeros_like(input_ids), + ).last_hidden_state + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + mask_2d, mask_4d, block_mask = model.transformer._sequence_id_attention_masks( + sequence_id, + batch_size=1, + seq_len=6, + device=input_ids.device, + ) # (b, l), (b, 1, l, l), None for eager attention + assert torch.equal(mask_2d, sequence_id.ge(0)) + assert torch.equal( + mask_4d, + sequence_id[:, None, :, None].eq(sequence_id[:, None, None, :]), + ) + assert block_mask is None diff --git a/tests/cpu/test_hopper_hardware_contract.py b/tests/cpu/test_hopper_hardware_contract.py new file mode 100644 index 0000000..669c01b --- /dev/null +++ b/tests/cpu/test_hopper_hardware_contract.py @@ -0,0 +1,16 @@ +"""Positive CPU allowlist aliases for Hopper/SM90 hardware contracts.""" + +from tests.unit.test_hopper_hardware_contract import ( + test_comparisons_require_the_exact_same_hopper_device_fingerprint, + test_golden_comparison_rejects_cross_device_and_honors_new_identity_fields, + test_release_hardware_accepts_named_hopper_sm90_products, + test_release_hardware_rejects_non_hopper_or_incomplete_identity, +) + + +__all__ = [ + "test_comparisons_require_the_exact_same_hopper_device_fingerprint", + "test_golden_comparison_rejects_cross_device_and_honors_new_identity_fields", + "test_release_hardware_accepts_named_hopper_sm90_products", + "test_release_hardware_rejects_non_hopper_or_incomplete_identity", +] diff --git a/tests/cpu/test_manifest_contract.py b/tests/cpu/test_manifest_contract.py new file mode 100644 index 0000000..62d8973 --- /dev/null +++ b/tests/cpu/test_manifest_contract.py @@ -0,0 +1,179 @@ +"""Cheap manifest-wide coverage that never loads checkpoint weights.""" + +from __future__ import annotations + +import importlib + +from fastplms.registry import get_model_registry + + +_CPU_FAMILIES = { + "ankh", + "boltz2", + "dplm", + "dplm2", + "e1", + "esm2", + "esm3", + "esm_plusplus", + "esmfold", + "esmfold2", +} +_CPU_CHECKPOINTS = frozenset( + { + "ankh2_large", + "ankh3_large", + "ankh3_xl", + "ankh_base", + "ankh_large", + "boltz2", + "dplm2_150m", + "dplm2_3b", + "dplm2_650m", + "dplm_150m", + "dplm_3b", + "dplm_650m", + "e1_150m", + "e1_300m", + "e1_600m", + "esm2_150m", + "esm2_35m", + "esm2_3b", + "esm2_650m", + "esm2_8m", + "esm3_small", + "esmc_6b", + "esmc_large", + "esmc_small", + "esmfold", + "esmfold2", + "esmfold2_experimental_cutoff2025", + "esmfold2_experimental_fast_cutoff2025", + "esmfold2_fast", + } +) + +_ANKH_OFFICIAL_FILES = { + "ankh_base": { + "config.json": "git-sha1:abd44a36b5469e9a7cb019e4059b5ac1392d8422", + "pytorch_model.bin": ( + "sha256:9b2a886374f0ff4a893f4e7a989deed76bb2458c8998bd5202ea8e97d92ddcc3" + ), + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0", + }, + "ankh_large": { + "config.json": "git-sha1:1abf33e52ee3d6be67d780ec57d32ac2b27b5306", + "pytorch_model.bin": ( + "sha256:517b6e8b279dedcb477af240b35c46bd6eb3307723eb281e60d4b2c8a87b889b" + ), + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09", + }, + "ankh2_large": { + "config.json": "git-sha1:9286bed4ecbc4f7113024919d16ec9719b0c0748", + "generation_config.json": ( + "git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a" + ), + "pytorch_model.bin": ( + "sha256:2df583f28f111276ee22a7b76007f4297e9a69766d60bccd9c8d7169c06ac606" + ), + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25", + }, + "ankh3_large": { + "config.json": "git-sha1:f5278f77d158cdd8a173df888e3ed365e84a80a3", + "generation_config.json": ( + "git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2" + ), + "pytorch_model.bin": ( + "sha256:26321a345e07a25b21c6c41b651c4db91b420892e52c0dcbc55bd7a8f510f95b" + ), + "special_tokens_map.json": "git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model": ( + "sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0" + ), + "tokenizer.json": "git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json": "git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", + }, + "ankh3_xl": { + "config.json": "git-sha1:f8997040e8913df75fd2eebe71a2a8eb750ed0d0", + "generation_config.json": ( + "git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a" + ), + "pytorch_model-00001-of-00003.bin": ( + "sha256:2c9793cbee16697cd4149debe07d3a27143e280f6e970fa46042aae820fea981" + ), + "pytorch_model-00002-of-00003.bin": ( + "sha256:31c5a860e414513c829ae52affb0970d7cef2c0545df2d6e1338b6806ab7174b" + ), + "pytorch_model-00003-of-00003.bin": ( + "sha256:055a853bdd3623db95a637935aa299427e837cd8ea69fc04708b0262508bec75" + ), + "special_tokens_map.json": "git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model": ( + "sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0" + ), + "tokenizer.json": "git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json": "git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", + }, +} + + +def _load_symbol(path: str) -> type: + module_name, _, symbol_name = path.rpartition(".") + assert module_name and symbol_name, path + return getattr(importlib.import_module(module_name), symbol_name) + + +def test_every_checkpoint_and_family_is_owned_by_the_cpu_matrix() -> None: + registry = get_model_registry() + assert frozenset(registry) == _CPU_CHECKPOINTS + assert {spec.family.id for spec in registry.values()} == _CPU_FAMILIES + for spec in registry.values(): + assert len(spec.fast.revision) == 40 + assert len(spec.official.revision) == 40 + assert not spec.fast.unresolved_files + assert not spec.official.unresolved_files + + +def test_every_advertised_automap_symbol_imports_without_optional_runtime_work() -> None: + registry = get_model_registry() + family_maps = { + family_id: family.auto_map + for family_id, family in registry.families.items() + } + assert sum(len(auto_map) for auto_map in family_maps.values()) == 37 + for family_id, auto_map in sorted(family_maps.items()): + for auto_class, symbol_path in sorted(auto_map.items()): + symbol = _load_symbol(symbol_path) + assert isinstance(symbol, type), (family_id, auto_class, symbol_path) + + for model_id, spec in sorted(registry.items()): + for auto_class, symbol_path in sorted(spec.auto_map.items()): + symbol = _load_symbol(symbol_path) + assert isinstance(symbol, type), (model_id, auto_class, symbol_path) + + +def test_ankh_official_asset_inventory_is_exact_and_complete() -> None: + """Pin every official runtime asset while excluding the obsolete PyTorch index.""" + + registry = get_model_registry() + for model_id, expected in _ANKH_OFFICIAL_FILES.items(): + source = registry[model_id].official + assert {path: item.encoded for path, item in source.file_map.items()} == expected + + for model_id in ("ankh_base", "ankh_large"): + assert "generation_config.json" not in registry[model_id].official.file_map + assert "spiece.model" not in registry[model_id].official.file_map + assert "generation_config.json" in registry["ankh2_large"].official.file_map + for model_id in ("ankh3_large", "ankh3_xl"): + assert "generation_config.json" in registry[model_id].official.file_map + assert "spiece.model" in registry[model_id].official.file_map + + xl = registry["ankh3_xl"] + assert "pytorch_model.bin.index.json" not in xl.official.file_map + assert "official PyTorch shard index is deliberately excluded" in xl.notes diff --git a/tests/cpu/test_optimized_validation_contracts.py b/tests/cpu/test_optimized_validation_contracts.py new file mode 100644 index 0000000..3d4c572 --- /dev/null +++ b/tests/cpu/test_optimized_validation_contracts.py @@ -0,0 +1,542 @@ +"""Public validation must remain active when Python removes ``assert`` statements.""" + +from __future__ import annotations + +import ast +import os +import subprocess +import sys +import textwrap +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace + +from fastplms.models.dplm.modeling_dplm import FAST_DPLM_ENCODER +from fastplms.models.dplm2.modeling_dplm2 import ( + DPLM2ForMaskedLM, + _has_packed_multimodal_layout, + _normalize_dplm2_input_ids, +) +from fastplms.models.e1.attention import _unpad_input +from fastplms.models.e1.cache import DynamicCache, KVCache +from fastplms.models.e1.modeling_e1 import FAST_E1_ENCODER, E1ForMaskedLM +from fastplms.models.e1.preparation import E1BatchPreparer +from fastplms.models.e1.retrieval import ( + IdSequence, + _make_homologue_searcher, + compute_ppll, + convert_to_tensor, + get_query_from_a3m, + read_fasta_sequences, +) +from fastplms.models.esm3.modeling_esm3 import ( + Affine3D, + FastESM3Config, + FastESM3PreTrainedModel, + RotationMatrix, +) +from fastplms.models.esmfold import modeling_fast_esmfold +from fastplms.models.esmfold.modeling_fast_esmfold import ( + EsmSelfAttention as FastEsmFoldSelfAttention, +) +from fastplms.models.esmfold.modeling_fast_esmfold import ( + FastEsmFoldConfig, +) +from fastplms.models.ttt import TTTConfig + + +@pytest.mark.parametrize( + ("values", "error_type"), + ( + ({"lr": 0.0}, ValueError), + ({"steps": 0}, ValueError), + ({"steps": 1.5}, TypeError), + ({"mask_ratio": 1.1}, ValueError), + ({"optimizer": "rmsprop"}, ValueError), + ({"lora_target_modules": "query"}, TypeError), + ), +) +def test_ttt_config_public_validation_uses_explicit_exceptions( + values: dict[str, object], + error_type: type[Exception], +) -> None: + with pytest.raises(error_type): + TTTConfig(**values) + + +def test_sequence_model_public_validations_use_explicit_exceptions( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(ValueError, match="outside the DPLM2 embedding table"): + _normalize_dplm2_input_ids(torch.tensor([[999]]), vocab_size=64) + with pytest.raises(ValueError, match=r"type_ids.*shape"): + _has_packed_multimodal_layout( + torch.zeros(3, dtype=torch.long), + aa_type=0, + struct_type=1, + pad_type=2, + ) + dplm2_stub = SimpleNamespace( + config=SimpleNamespace( + use_return_dict=True, + dplm_type="multimodal", + hidden_size=8, + ), + pad_id=0, + ) + with pytest.raises(ValueError, match="attention_mask is required"): + DPLM2ForMaskedLM.forward( + dplm2_stub, + inputs_embeds=torch.zeros(1, 2, 8), + ) + with pytest.raises(ValueError, match="type_ids is required"): + DPLM2ForMaskedLM.forward( + dplm2_stub, + attention_mask=torch.ones(1, 2), + inputs_embeds=torch.zeros(1, 2, 8), + ) + with pytest.raises(ValueError, match="Pass either seq or input_ids"): + DPLM2ForMaskedLM._ttt_tokenize(SimpleNamespace()) + empty_replacement_stub = SimpleNamespace( + tokenizer=SimpleNamespace( + all_special_ids=(), + struct_cls_token="", + _token_to_id={"": 0}, + ) + ) + with pytest.raises(RuntimeError, match="replacement set is empty"): + DPLM2ForMaskedLM._ttt_replacement_tokens( + empty_replacement_stub, + torch.tensor([[1]]), + ) + with pytest.raises(ValueError, match=r"head_mask\.dim"): + FAST_DPLM_ENCODER._convert_head_mask_to_5d( + SimpleNamespace(dtype=torch.float32), + torch.ones(1, 1, 1), + num_hidden_layers=1, + ) + + DynamicCache().crop(1) + with pytest.raises(ValueError, match="max_length must be positive"): + DynamicCache().crop(0) + with pytest.raises(ValueError, match="use_cache=True"): + KVCache().after_forward( + {"context": ["ctx"], "context_len": [1], "use_cache": False}, + SimpleNamespace(), + ) + with pytest.raises(TypeError, match="Sequence must be a string"): + E1BatchPreparer.validate_sequence( + SimpleNamespace(mask_token=""), + 123, + ) + e1_stub = SimpleNamespace( + config=SimpleNamespace( + hidden_size=8, + max_num_positions_within_seq=8, + max_num_positions_global=8, + max_num_sequences=2, + ) + ) + with pytest.raises(ValueError, match="input_ids must have rank 2"): + FAST_E1_ENCODER._prepare_hidden_states( + e1_stub, + torch.ones(1, 2, 3, dtype=torch.long), + None, + None, + None, + None, + ) + with pytest.raises(ValueError, match="sequence_ids must have shape"): + FAST_E1_ENCODER._prepare_hidden_states( + e1_stub, + torch.ones(1, 2, dtype=torch.long), + None, + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(2, 1, dtype=torch.long), + ) + + malformed_fasta = tmp_path / "malformed.fasta" + malformed_fasta.write_text("ACDE\n", encoding="utf-8") + with pytest.raises(ValueError, match="before header"): + read_fasta_sequences(str(malformed_fasta)) + malformed_a3m = tmp_path / "malformed.a3m" + malformed_a3m.write_text("ACDE\n", encoding="utf-8") + with pytest.raises(ValueError, match="No FASTA header"): + get_query_from_a3m(str(malformed_a3m)) + with pytest.raises(ValueError, match="equal aligned lengths"): + convert_to_tensor([IdSequence("a", "AC"), IdSequence("b", "ACD")]) + with pytest.raises(ValueError, match="empty token sequence"): + compute_ppll(torch.empty(0, 2), torch.empty(0, dtype=torch.long)) + with pytest.raises(ValueError, match="target_db is required"): + _make_homologue_searcher("mmseqs2", None) + + with pytest.raises(ValueError, match="hidden_size must be positive"): + FastESM3Config(hidden_size=0) + with pytest.raises(ValueError, match="divisible"): + FastESM3Config(hidden_size=24, num_attention_heads=5) + with pytest.raises(ValueError, match="currently supports only"): + FastESM3PreTrainedModel.attn_backend.fset( + SimpleNamespace(), + "flash_attention_2", + ) + with pytest.raises( + ValueError, + match="Rotation matrices must have trailing shape", + ): + RotationMatrix(torch.zeros(2, 3, 2)) + rotation = RotationMatrix.identity((2,)) + with pytest.raises( + ValueError, + match="Affine translation and rotation batch shapes must match", + ): + Affine3D(torch.zeros(1, 3), rotation) + + with pytest.raises(ValueError, match="not a multiple"): + FastEsmFoldSelfAttention( + FastEsmFoldConfig( + hidden_size=10, + num_attention_heads=3, + attn_backend="eager", + ) + ) + fold_attention = FastEsmFoldSelfAttention( + FastEsmFoldConfig( + hidden_size=8, + num_attention_heads=2, + attn_backend="eager", + ) + ) + monkeypatch.setattr(modeling_fast_esmfold, "flex_attention", None) + attention_heads = torch.zeros(1, 2, 3, 4) # (b=1, h=2, l=3, d_h=4) + with pytest.raises(RuntimeError, match="Flex attention is not available"): + fold_attention._flex_attn( + attention_heads, + attention_heads, + attention_heads, + ) + + query_layer = torch.zeros(1, 3, 2, 4) # (b=1, l_q=3, h=2, d_h=4) + key_layer = torch.zeros(1, 3, 2, 4) # (b, l_k=3, h, d_h) + value_layer = torch.zeros(1, 3, 2, 4) # (b, l_k, h, d_h) + with pytest.raises( + ValueError, + match="Shape mismatch between query layer and query sequence ids", + ): + _unpad_input( + query_layer, + key_layer, + value_layer, + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(1, 3, dtype=torch.long), + ) + with pytest.raises( + ValueError, + match="key_layer and value_layer must have identical shapes", + ): + _unpad_input( + query_layer, + key_layer, + value_layer[:, :2], + torch.zeros(1, 3, dtype=torch.long), + torch.zeros(1, 3, dtype=torch.long), + ) + + +def test_ttt_and_e1_public_state_validations_use_explicit_exceptions() -> None: + from tests.integration.test_ttt import DummyTTTModel + + model = DummyTTTModel() + with pytest.raises(ValueError, match="Pass either seq or input_ids"): + model._ttt_tokenize() + with pytest.raises(RuntimeError, match="no LoRA parameters"): + model._ttt_lora_parameters() + with pytest.raises(RuntimeError, match="no LoRA state"): + model._ttt_snapshot_lora_state() + + model._ttt_ensure_initialized() + with pytest.raises(ValueError, match="Changing lora_rank"): + model.ttt(seq="AC", ttt_config={"lora_rank": 3}) + with pytest.raises(RuntimeError, match="state/module count mismatch"): + model._ttt_restore_lora_state([]) + + invalid_target = DummyTTTModel() + invalid_target._ttt_cfg = invalid_target.ttt_config.merged( + {"lora_target_modules": ("missing",)} + ) + with pytest.raises(ValueError, match="did not find any target modules"): + invalid_target._ttt_inject_lora() + + with pytest.raises(ValueError, match="E1 token tensors"): + E1ForMaskedLM._ttt_tokenize(SimpleNamespace()) + with pytest.raises(TypeError, match="tensor dictionary"): + E1ForMaskedLM._ttt_predict_logits(SimpleNamespace(), torch.tensor([[1]])) + + class EmptyContexts: + def sample_msa_contexts(self, **_kwargs: object) -> dict[str, object]: + return {} + + with pytest.raises(ValueError, match="sampled MSA context"): + E1ForMaskedLM.score_ppll( + EmptyContexts(), + sequences=["AC"], + a3m_path="unused.a3m", + ) + + +def test_representative_public_validation_survives_python_optimized_mode() -> None: + script = textwrap.dedent( + """ + from types import SimpleNamespace + + import torch + import torch.nn as nn + + from fastplms.models.dplm.modeling_dplm import FAST_DPLM_ENCODER + from fastplms.models.dplm2.modeling_dplm2 import ( + DPLM2ForMaskedLM, + _has_packed_multimodal_layout, + _normalize_dplm2_input_ids, + ) + from fastplms.models.e1.attention import _unpad_input + from fastplms.models.e1.cache import DynamicCache, KVCache + from fastplms.models.e1.modeling_e1 import E1ForMaskedLM, FAST_E1_ENCODER + from fastplms.models.e1.retrieval import ( + _make_homologue_searcher, + compute_ppll, + ) + from fastplms.models.esm3.modeling_esm3 import ( + Affine3D, + FastESM3Config, + FastESM3PreTrainedModel, + RotaryEmbedding, + RotationMatrix, + ) + from fastplms.models.esmfold import modeling_fast_esmfold + from fastplms.models.esmfold.modeling_fast_esmfold import ( + EsmSelfAttention as FastEsmFoldSelfAttention, + FastEsmFoldConfig, + ) + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin, TTTConfig + + class DummyTTTModel(FastPLMTestTimeTrainingMixin, nn.Module): + # Keep the optimized subprocess independent of the broad integration module. + + def __init__(self): + nn.Module.__init__(self) + self.config = SimpleNamespace(vocab_size=8) + self.backbone = nn.Sequential(nn.Linear(8, 8)) + self.init_ttt({"lora_rank": 2, "lora_alpha": 1.0}) + + def _ttt_get_trainable_modules(self): + return [self.backbone] + + def must_raise(error_type, function, *args, **kwargs): + try: + function(*args, **kwargs) + except error_type: + return + raise RuntimeError( + f"{function.__qualname__} did not raise {error_type.__name__} under -O" + ) + + must_raise(ValueError, TTTConfig, steps=0) + must_raise(TypeError, TTTConfig, lora_target_modules="query") + must_raise( + ValueError, + _normalize_dplm2_input_ids, + torch.tensor([[999]]), + 64, + ) + must_raise( + ValueError, + _has_packed_multimodal_layout, + torch.zeros(3, dtype=torch.long), + 0, + 1, + 2, + ) + dplm2_stub = SimpleNamespace( + config=SimpleNamespace( + use_return_dict=True, + dplm_type="multimodal", + hidden_size=8, + ), + pad_id=0, + ) + must_raise( + ValueError, + DPLM2ForMaskedLM.forward, + dplm2_stub, + inputs_embeds=torch.zeros(1, 2, 8), + ) + must_raise( + ValueError, + DPLM2ForMaskedLM.forward, + dplm2_stub, + attention_mask=torch.ones(1, 2), + inputs_embeds=torch.zeros(1, 2, 8), + ) + must_raise(ValueError, DPLM2ForMaskedLM._ttt_tokenize, SimpleNamespace()) + must_raise( + ValueError, + FAST_DPLM_ENCODER._convert_head_mask_to_5d, + SimpleNamespace(dtype=torch.float32), + torch.ones(1, 1, 1), + 1, + ) + must_raise(ValueError, DynamicCache().crop, 0) + must_raise( + ValueError, + KVCache().after_forward, + {"context": ["ctx"], "context_len": [1], "use_cache": False}, + SimpleNamespace(), + ) + e1_stub = SimpleNamespace( + config=SimpleNamespace( + hidden_size=8, + max_num_positions_within_seq=8, + max_num_positions_global=8, + max_num_sequences=2, + ) + ) + must_raise( + ValueError, + FAST_E1_ENCODER._prepare_hidden_states, + e1_stub, + torch.ones(1, 2, dtype=torch.long), + None, + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(2, 1, dtype=torch.long), + ) + must_raise( + ValueError, + compute_ppll, + torch.empty(0, 2), + torch.empty(0, dtype=torch.long), + ) + must_raise(ValueError, _make_homologue_searcher, "mmseqs2", None) + must_raise(ValueError, FastESM3Config, hidden_size=0) + must_raise( + ValueError, + FastESM3PreTrainedModel.attn_backend.fset, + SimpleNamespace(), + "flash_attention_2", + ) + rotary = RotaryEmbedding(4) + rotary._update_cos_sin_cache = lambda *_args, **_kwargs: None + must_raise( + RuntimeError, + rotary.forward, + torch.zeros(1, 2, 1, 4), + torch.zeros(1, 2, 1, 4), + ) + must_raise(ValueError, RotationMatrix, torch.zeros(2, 3, 2)) + rotation = RotationMatrix.identity((2,)) + must_raise( + ValueError, + Affine3D, + torch.zeros(1, 3), + rotation, + ) + must_raise( + ValueError, + FastEsmFoldSelfAttention, + FastEsmFoldConfig( + hidden_size=10, + num_attention_heads=3, + attn_backend="eager", + ), + ) + fold_attention = FastEsmFoldSelfAttention( + FastEsmFoldConfig( + hidden_size=8, + num_attention_heads=2, + attn_backend="eager", + ) + ) + modeling_fast_esmfold.flex_attention = None + attention_heads = torch.zeros(1, 2, 3, 4) + must_raise( + RuntimeError, + fold_attention._flex_attn, + attention_heads, + attention_heads, + attention_heads, + ) + must_raise( + ValueError, + _unpad_input, + torch.zeros(1, 3, 2, 4), + torch.zeros(1, 3, 2, 4), + torch.zeros(1, 3, 2, 4), + torch.zeros(1, 2, dtype=torch.long), + torch.zeros(1, 3, dtype=torch.long), + ) + must_raise( + ValueError, + _unpad_input, + torch.zeros(1, 3, 2, 4), + torch.zeros(1, 3, 2, 4), + torch.zeros(1, 2, 2, 4), + torch.zeros(1, 3, dtype=torch.long), + torch.zeros(1, 3, dtype=torch.long), + ) + + model = DummyTTTModel() + must_raise(ValueError, model._ttt_tokenize) + must_raise(RuntimeError, model._ttt_lora_parameters) + model._ttt_ensure_initialized() + must_raise(ValueError, model.ttt, seq="AC", ttt_config={"lora_rank": 3}) + must_raise(ValueError, E1ForMaskedLM._ttt_tokenize, SimpleNamespace()) + must_raise( + TypeError, + E1ForMaskedLM._ttt_predict_logits, + SimpleNamespace(), + torch.tensor([[1]]), + ) + """ + ) + environment = dict(os.environ) + environment["PYTHONHASHSEED"] = "0" + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[2], + env=environment, + check=False, + capture_output=True, + text=True, + timeout=30, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_esmfold_public_attention_validation_has_no_optimized_away_asserts() -> None: + source_path = ( + Path(__file__).resolve().parents[2] + / "src" + / "fastplms" + / "models" + / "esmfold" + / "modeling_fast_esmfold.py" + ) + module = ast.parse(source_path.read_text(encoding="utf-8")) + attention_class = next( + node + for node in module.body + if isinstance(node, ast.ClassDef) and node.name == "EsmSelfAttention" + ) + public_validation_methods = { + method.name: method + for method in attention_class.body + if isinstance(method, (ast.FunctionDef, ast.AsyncFunctionDef)) + and method.name in {"__init__", "_flex_attn"} + } + + assert set(public_validation_methods) == {"__init__", "_flex_attn"} + for method in public_validation_methods.values(): + assert not any(isinstance(node, ast.Assert) for node in ast.walk(method)) diff --git a/tests/cpu/test_peft_contracts.py b/tests/cpu/test_peft_contracts.py new file mode 100644 index 0000000..dd6029b --- /dev/null +++ b/tests/cpu/test_peft_contracts.py @@ -0,0 +1,669 @@ +"""Mandatory shipped fine-tuning, PEFT, collator, and persistence contracts.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from examples import fine_tuning +from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForSequenceClassification, +) +from tests.unit import test_fine_tuning_example as fine_tuning_contracts + + +test_pair_collator_enforces_longest_first_tokenizer_limit = ( + fine_tuning_contracts.test_pair_collator_enforces_longest_first_tokenizer_limit +) +test_pair_token_budget_includes_special_tokens_at_the_exact_boundary = ( + fine_tuning_contracts.test_pair_token_budget_includes_special_tokens_at_the_exact_boundary +) +test_reporting_is_opt_in_for_the_minimal_training_install = ( + fine_tuning_contracts.test_reporting_is_opt_in_for_the_minimal_training_install +) +test_max_length_contract_is_an_encoded_budget_including_added_tokens = ( + fine_tuning_contracts.test_max_length_contract_is_an_encoded_budget_including_added_tokens +) +test_ordered_training_row_hash_is_content_and_order_sensitive = ( + fine_tuning_contracts.test_ordered_training_row_hash_is_content_and_order_sensitive +) +test_persisted_hash_scope_covers_full_state_or_only_lora_payload = ( + fine_tuning_contracts.test_persisted_hash_scope_covers_full_state_or_only_lora_payload +) +test_atomic_final_artifact_reload_preserves_trainer_and_held_out_logits = ( + fine_tuning_contracts.test_atomic_final_artifact_reload_preserves_trainer_and_held_out_logits +) + + +def test_plot_contract_uses_task_output_paths_and_has_no_interactive_or_overwrite_path() -> None: + contract = fine_tuning_contracts + contract.test_plot_contract_uses_task_output_paths_and_has_no_interactive_or_overwrite_path() + + +def test_training_manifest_records_reproducible_model_data_and_tokenizer_identity( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + contract = fine_tuning_contracts + contract.test_training_manifest_records_reproducible_model_data_and_tokenizer_identity( + monkeypatch, + tmp_path, + ) + + +def test_immutable_sources_reject_moving_refs_pin_only_shipped_defaults_and_detect_drift( + tmp_path: Path, +) -> None: + contract = fine_tuning_contracts + contract.test_immutable_sources_reject_moving_refs_pin_only_shipped_defaults_and_detect_drift( + tmp_path + ) + + +@pytest.mark.parametrize("backend", ("flash_attention_2", "flash_attention_3")) +def test_fine_tuning_example_rejects_flash_without_explicit_bf16_policy( + backend: str, +) -> None: + with pytest.raises(SystemExit): + fine_tuning.build_parser().parse_args(["--attn-backend", backend]) + with pytest.raises(ValueError, match="explicit BF16 CUDA"): + fine_tuning.initialize_model("unused", 2, attn_backend=backend) + + +def test_fine_tuning_main_wires_both_tasks_without_external_io( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Exercise the shipped CLI entry point without loading models or datasets.""" + + regression_calls: list[dict[str, Any]] = [] + classification_calls: list[dict[str, Any]] = [] + monkeypatch.setattr( + fine_tuning, + "train_regression_model", + lambda **kwargs: regression_calls.append(kwargs), + ) + monkeypatch.setattr( + fine_tuning, + "train_classification_model", + lambda **kwargs: classification_calls.append(kwargs), + ) + + result = fine_tuning.main( + [ + "--task", + "both", + "--model_path", + "offline/tiny-model", + "--model-revision", + "a" * 40, + "--classification-dataset-source", + "offline/classification", + "--classification-dataset-revision", + "b" * 40, + "--regression-train-dataset-source", + "offline/regression-train", + "--regression-train-dataset-revision", + "c" * 40, + "--regression-validation-dataset-source", + "offline/regression-validation", + "--regression-validation-dataset-revision", + "d" * 40, + "--regression-test-dataset-source", + "offline/regression-test", + "--regression-test-dataset-revision", + "e" * 40, + "--no-use-lora", + "--batch_size", + "3", + "--lr", + "0.001", + "--epochs", + "2.5", + "--max_length", + "64", + "--attn-backend", + "eager", + "--output-dir", + str(tmp_path / "fine-tuning"), + "--grad_accum", + "2", + "--patience", + "1", + "--seed", + "17", + "--full-determinism", + "--no-plot-results", + ] + ) + + shared = { + "model_name": "offline/tiny-model", + "model_revision": "a" * 40, + "use_lora": False, + "batch_size": 3, + "learning_rate": 0.001, + "num_epochs": 2.5, + "max_length": 64, + "gradient_accumulation_steps": 2, + "patience": 1, + "seed": 17, + "full_determinism": True, + "plot_results": False, + "attn_backend": "eager", + } + assert result == 0 + assert regression_calls == [ + { + **shared, + "train_dataset_source": "offline/regression-train", + "train_dataset_revision": "c" * 40, + "validation_dataset_source": "offline/regression-validation", + "validation_dataset_revision": "d" * 40, + "test_dataset_source": "offline/regression-test", + "test_dataset_revision": "e" * 40, + "output_dir": tmp_path / "fine-tuning" / "regression", + } + ] + assert classification_calls == [ + { + **shared, + "dataset_source": "offline/classification", + "dataset_revision": "b" * 40, + "output_dir": tmp_path / "fine-tuning" / "classification", + } + ] + + +def test_multi_task_preflight_rejects_any_existing_child_before_dispatch( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + output_root = tmp_path / "fine-tuning" + existing = output_root / "classification_lora" + existing.mkdir(parents=True) + calls: list[str] = [] + monkeypatch.setattr( + fine_tuning, + "train_regression_model", + lambda **kwargs: calls.append("regression"), + ) + monkeypatch.setattr( + fine_tuning, + "train_classification_model", + lambda **kwargs: calls.append("classification"), + ) + + with pytest.raises(FileExistsError, match="could mix prior state"): + fine_tuning.main(["--task", "both", "--output-dir", str(output_root)]) + assert calls == [] + + +def test_atomic_output_reservation_rejects_reuse_and_cleans_failed_runs( + tmp_path: Path, +) -> None: + existing = tmp_path / "existing" + existing.mkdir() + with ( + pytest.raises(FileExistsError, match="already exists"), + fine_tuning._reserved_output_directory(existing), + ): + pytest.fail("an existing output must never be entered") + + failed = tmp_path / "failed" + with ( + pytest.raises(RuntimeError, match="training failed"), + fine_tuning._reserved_output_directory(failed) as reserved, + ): + (reserved / "partial-checkpoint.bin").write_bytes(b"partial") + raise RuntimeError("training failed") + assert not failed.exists() + + completed = tmp_path / "completed" + with fine_tuning._reserved_output_directory(completed) as reserved: + (reserved / "result.txt").write_text("complete\n", encoding="utf-8") + assert (completed / "result.txt").is_file() + assert not (completed / fine_tuning._OUTPUT_RESERVATION_FILE).exists() + + +class _ColumnDataset: + def __init__(self, **columns: list[Any]) -> None: + self.columns = columns + self.column_names = list(columns) + + def __getitem__(self, column: str) -> list[Any]: + return self.columns[column] + + def __len__(self) -> int: + return len(next(iter(self.columns.values()), ())) + + +def _classification_data( + *, + train_labels: list[Any] | None = None, + valid_labels: list[Any] | None = None, + test_labels: list[Any] | None = None, +) -> dict[str, _ColumnDataset]: + return { + "train": _ColumnDataset( + seqs=["AC", "DE"], + labels=[0, 1] if train_labels is None else train_labels, + ), + "valid": _ColumnDataset( + seqs=["FG"], + labels=[0] if valid_labels is None else valid_labels, + ), + "test": _ColumnDataset( + seqs=["HI"], + labels=[1] if test_labels is None else test_labels, + ), + } + + +def test_dataset_contracts_reject_noncontiguous_unseen_and_nonfinite_labels() -> None: + assert fine_tuning._validate_classification_dataset_dict(_classification_data()) == 2 + + with pytest.raises(ValueError, match="contiguous zero-based"): + fine_tuning._validate_classification_dataset_dict(_classification_data(train_labels=[0, 2])) + with pytest.raises(ValueError, match="absent from train"): + fine_tuning._validate_classification_dataset_dict(_classification_data(valid_labels=[2])) + with pytest.raises(ValueError, match="must be an integer"): + fine_tuning._validate_classification_dataset_dict(_classification_data(test_labels=[1.0])) + + valid_regression = _ColumnDataset(SeqA=["AC"], SeqB=["DE"], labels=[-8.2]) + fine_tuning._validate_regression_dataset(valid_regression, split="train") + invalid_regression = _ColumnDataset(SeqA=["AC"], SeqB=["DE"], labels=[float("nan")]) + with pytest.raises(ValueError, match="finite real number"): + fine_tuning._validate_regression_dataset(invalid_regression, split="test") + + +def test_invalid_classification_contract_precedes_model_initialization_and_cleans_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + fine_tuning, + "_load_dataset_immutable", + lambda *args, **kwargs: (_classification_data(valid_labels=[2]), {}), + ) + monkeypatch.setattr( + fine_tuning, + "initialize_model", + lambda *args, **kwargs: pytest.fail("model initialized before data validation"), + ) + output_dir = tmp_path / "classification" + + with pytest.raises(ValueError, match="absent from train"): + fine_tuning.train_classification_model( + dataset_source="offline/classification", + dataset_revision="a" * 40, + output_dir=output_dir, + ) + assert not output_dir.exists() + + +def test_invalid_regression_contract_precedes_model_initialization_and_cleans_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + datasets = iter( + ( + _ColumnDataset(SeqA=["AC"], SeqB=["DE"], labels=[1.0]), + _ColumnDataset(SeqA=["FG"], SeqB=["HI"], labels=[2.0]), + _ColumnDataset(SeqA=["KL"], SeqB=["MN"], labels=[float("inf")]), + ) + ) + monkeypatch.setattr( + fine_tuning, + "_load_dataset_immutable", + lambda *args, **kwargs: (next(datasets), {}), + ) + monkeypatch.setattr( + fine_tuning, + "initialize_model", + lambda *args, **kwargs: pytest.fail("model initialized before data validation"), + ) + output_dir = tmp_path / "regression" + + with pytest.raises(ValueError, match="finite real number"): + fine_tuning.train_regression_model( + train_dataset_source="offline/train", + train_dataset_revision="a" * 40, + validation_dataset_source="offline/validation", + validation_dataset_revision="b" * 40, + test_dataset_source="offline/test", + test_dataset_revision="c" * 40, + output_dir=output_dir, + ) + assert not output_dir.exists() + + +def test_plot_writes_are_exclusive_and_preserve_existing_bytes(tmp_path: Path) -> None: + class FakeFigure: + def savefig(self, handle: Any, **kwargs: Any) -> None: + assert kwargs == {"format": "png", "dpi": 300} + handle.write(b"new-png") + + target = tmp_path / "classification_results.png" + assert fine_tuning._save_figure_exclusive(FakeFigure(), target) == target + assert target.read_bytes() == b"new-png" + target.write_bytes(b"existing-png") + with pytest.raises(FileExistsError, match="Refusing to overwrite"): + fine_tuning._save_figure_exclusive(FakeFigure(), target) + assert target.read_bytes() == b"existing-png" + + +class _OfflineProteinTokenizer: + """Small deterministic tokenizer sufficient for both shipped collators.""" + + pad_token_id = 1 + name_or_path = "offline-cpu-contract" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + @staticmethod + def _encode(sequence: str) -> list[int]: + alphabet = "ACDEFGHIKLMN" + available_ids = (3, 4, *range(6, 16)) + token_ids = dict(zip(alphabet, available_ids, strict=True)) + return [token_ids[residue] for residue in sequence] + + def __call__( + self, + sequences: str | tuple[str, ...] | list[str], + pairs: str | tuple[str, ...] | list[str] | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + self.calls.append(dict(kwargs)) + scalar = isinstance(sequences, str) + sequence_rows = [sequences] if scalar else list(sequences) + if pairs is None: + pair_rows: list[str | None] = [None] * len(sequence_rows) + elif isinstance(pairs, str): + pair_rows = [pairs] + else: + pair_rows = list(pairs) + if len(sequence_rows) != len(pair_rows): + raise ValueError("sequence and pair batches must align") + + rows: list[list[int]] = [] + for sequence, pair in zip(sequence_rows, pair_rows, strict=True): + row = [0, *self._encode(sequence), 2] + if pair is not None: + row.extend((*self._encode(pair), 2)) + max_length = kwargs.get("max_length") + if kwargs.get("truncation") and max_length is not None: + row = row[: int(max_length)] + rows.append(row) + + if kwargs.get("return_tensors") != "pt": + return {"input_ids": rows[0] if scalar else rows} + + width = max(map(len, rows)) + multiple = kwargs.get("pad_to_multiple_of") + if multiple is not None: + width = min( + ((width + int(multiple) - 1) // int(multiple)) * int(multiple), + int(kwargs.get("max_length", width)), + ) + input_ids = torch.tensor( + [row + [self.pad_token_id] * (width - len(row)) for row in rows], + dtype=torch.long, + ) + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(self.pad_token_id).long(), + } + + def save_pretrained(self, save_directory: str | Path) -> tuple[str]: + directory = Path(save_directory) + directory.mkdir(parents=True, exist_ok=True) + path = directory / "tokenizer.json" + path.write_text("{}\n", encoding="utf-8") + return (str(path),) + + +def _tiny_config(attn_backend: str = "eager") -> FastEsmConfig: + return FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + num_labels=2, + pad_token_id=1, + mask_token_id=5, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + position_embedding_type="absolute", + attn_backend=attn_backend, + ) + + +def test_shipped_collators_create_tokenizer_aware_sequence_and_pair_batches() -> None: + tokenizer = _OfflineProteinTokenizer() + sequence_batch = fine_tuning.SequenceCollator( + tokenizer, + regression=False, + max_length=8, + )([("ACD", 0), ("EF", 1)]) + pair_batch = fine_tuning.PairCollator( + tokenizer, + regression=True, + max_length=8, + )([("ACD", "EF", 1.5), ("GH", "IK", 2.5)]) + + assert sequence_batch["input_ids"].shape == (2, 8) + assert sequence_batch["attention_mask"].shape == (2, 8) + assert sequence_batch["labels"].dtype == torch.long + assert pair_batch["input_ids"].shape == (2, 8) + assert pair_batch["attention_mask"].shape == (2, 8) + assert pair_batch["labels"].dtype == torch.float32 + assert tokenizer.calls[0] == { + "padding": "longest", + "return_tensors": "pt", + "truncation": True, + "max_length": 8, + "pad_to_multiple_of": 8, + } + assert tokenizer.calls[1] == { + "padding": "longest", + "return_tensors": "pt", + "truncation": "longest_first", + "max_length": 8, + "pad_to_multiple_of": 8, + } + + +def test_shipped_initializer_drives_one_peft_step_and_atomic_final_reload( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + torch.manual_seed(7) + template = FastEsmForSequenceClassification(_tiny_config()) + base_state = { # parameter name -> checkpoint-shaped tensor + name: tensor.detach().clone() + for name, tensor in template.state_dict().items() + } + tokenizer = _OfflineProteinTokenizer() + loader_calls: list[tuple[str, dict[str, Any]]] = [] + + def load_tiny_model( + model_name: str, + **kwargs: Any, + ) -> FastEsmForSequenceClassification: + loader_calls.append((model_name, dict(kwargs))) + model = FastEsmForSequenceClassification( + _tiny_config(attn_backend=kwargs["attn_implementation"]) + ) + model.load_state_dict(base_state) + model.tokenizer = tokenizer + return model + + monkeypatch.setattr( + fine_tuning.AutoModelForSequenceClassification, + "from_pretrained", + staticmethod(load_tiny_model), + ) + adapter, observed_tokenizer = fine_tuning.initialize_model( + "offline/tiny-esm2", + num_labels=2, + use_lora=True, + lora_config=None, + model_revision="a" * 40, + ) + + assert observed_tokenizer is tokenizer + assert adapter.config.attn_backend == "sdpa" + assert loader_calls == [ + ( + "offline/tiny-esm2", + { + "trust_remote_code": True, + "num_labels": 2, + "attn_implementation": "sdpa", + "revision": "a" * 40, + }, + ) + ] + assert "classifier" in adapter.peft_config["default"].modules_to_save + + batch = fine_tuning.SequenceCollator( # token tensors: (b=2, l) + tokenizer, + regression=False, + max_length=8, + )([("ACD", 0), ("EFG", 1)]) + before = { # parameter name -> parameter-shaped tensor + name: parameter.detach().clone() + for name, parameter in adapter.named_parameters() + } + trainable = {name for name, parameter in adapter.named_parameters() if parameter.requires_grad} + assert trainable + assert all("lora_" in name or "classifier" in name for name in trainable) + + adapter.train() + sdpa_calls: list[dict[str, Any]] = [] + original_sdpa = torch.nn.functional.scaled_dot_product_attention + + def instrumented_sdpa(*args: Any, **kwargs: Any) -> torch.Tensor: + sdpa_calls.append(dict(kwargs)) + return original_sdpa(*args, **kwargs) + + monkeypatch.setattr( + torch.nn.functional, + "scaled_dot_product_attention", + instrumented_sdpa, + ) + optimizer = torch.optim.SGD( + (parameter for parameter in adapter.parameters() if parameter.requires_grad), + lr=0.5, + ) + output = adapter(**batch) # logits: (b=2, c=2); loss: () + assert len(sdpa_calls) == 1 + assert output.loss is not None and torch.isfinite(output.loss) + output.loss.backward() + expected_changed = { + name + for name, parameter in adapter.named_parameters() + if parameter.grad is not None and int(torch.count_nonzero(parameter.grad).item()) > 0 + } + assert expected_changed + optimizer.step() + changed = { + name + for name, parameter in adapter.named_parameters() + if not torch.equal(parameter.detach(), before[name]) + } + assert changed == expected_changed + assert changed <= trainable + assert any("lora_" in name for name in changed) + assert any("classifier" in name for name in changed) + + inference_inputs = { # token tensors: (b, l) + key: value + for key, value in batch.items() + if key != "labels" + } + adapter.eval() + with torch.inference_mode(): + expected = adapter(**inference_inputs).logits # (b, c) + + verification_rows = [("ACD", 0), ("EFG", 1)] + verification_collator = fine_tuning.SequenceCollator( + tokenizer, + regression=False, + max_length=8, + ) + + class _TinyTrainer: + def __init__(self, model: torch.nn.Module) -> None: + self.model = model + self.model_wrapped = model + self.args = SimpleNamespace( + device=torch.device("cpu"), + bf16=False, + fp16=False, + ) + + def predict(self, rows: Any) -> Any: + prediction_batch = verification_collator(rows) + prediction_inputs = { + key: value for key, value in prediction_batch.items() if key != "labels" + } + self.model.eval() + with torch.inference_mode(): + logits = ( # (b, c) + self.model(**prediction_inputs).logits.detach().cpu().numpy() + ) + return SimpleNamespace(predictions=logits) + + def save_model(self, save_directory: str | Path) -> None: + self.model.save_pretrained(save_directory, safe_serialization=True) + + trainer = _TinyTrainer(adapter) + original_model = trainer.model + original_wrapped = trainer.model_wrapped + artifact = fine_tuning._save_reload_verify_final_artifact( + trainer, + tokenizer, + output_dir=str(tmp_path / "training-output"), + model_name="offline/tiny-esm2", + model_revision="a" * 40, + num_labels=2, + use_lora=True, + verification_dataset=verification_rows, + data_collator=verification_collator, + ) + adapter_directory = tmp_path / "training-output" / "final_model" + assert trainer.model is original_model + assert trainer.model_wrapped is original_wrapped + assert (adapter_directory / "adapter_config.json").is_file() + assert (adapter_directory / "adapter_model.safetensors").is_file() + assert (adapter_directory / "artifact_metadata.json").is_file() + assert (adapter_directory / "tokenizer.json").is_file() + assert artifact["path"] == str(adapter_directory.resolve()) + assert len(artifact["tree_sha256"]) == 64 + assert artifact["reload_verified"] is True + assert artifact["held_out_inference"]["rows"] == 2 + assert artifact["held_out_inference"]["max_absolute_error"] <= 1e-6 + assert artifact["verified_parameter_sha256"] == ( + fine_tuning._persisted_parameter_hashes(adapter, use_lora=True) + ) + assert loader_calls[-1] == ( + "offline/tiny-esm2", + { + "trust_remote_code": True, + "num_labels": 2, + "attn_implementation": "sdpa", + "revision": "a" * 40, + }, + ) + with torch.inference_mode(): + observed = adapter(**inference_inputs).logits # (b, c) + torch.testing.assert_close(observed, expected, rtol=0.0, atol=0.0) diff --git a/tests/cpu/test_publication_contracts.py b/tests/cpu/test_publication_contracts.py new file mode 100644 index 0000000..05d949c --- /dev/null +++ b/tests/cpu/test_publication_contracts.py @@ -0,0 +1,393 @@ +"""Mandatory fail-closed artifact and publication security contracts.""" + +import hashlib +import json +import pytest +import torch +from dataclasses import replace +from pathlib import Path +from typing import Any +from huggingface_hub import CommitOperationAdd, CommitOperationDelete +from safetensors.torch import save_file + +from fastplms.registry import ( + CheckpointSource, + FileDigest, + ModelRegistry, + ModelSpec, + get_model_registry, +) +from tests.release import test_artifacts as artifact_contracts +from tests.release import test_publish_files_only as publish_contracts +from tools.artifacts import ArtifactError, hash_file +from tools.artifacts import publish as publish_module +from tools.artifacts.build import ( + _canonical_state_sha256, + _content_manifest, + _provenance, + _runtime_attestation, + validate_artifact, + validate_weight_artifact, +) + + +# Register the release module's immutable source fixture when these contracts are +# collected through this CPU-only allowlist module. +_clean_publication_source = publish_contracts._clean_publication_source + +test_artifact_build_rejects_unknown_runtime_source_extension = ( + artifact_contracts.test_artifact_build_rejects_unknown_runtime_source_extension +) +test_artifact_build_rejects_untracked_runtime_source = ( + artifact_contracts.test_artifact_build_rejects_untracked_runtime_source +) +test_artifact_validation_rejects_unsafe_manifest_paths = ( + artifact_contracts.test_artifact_validation_rejects_unsafe_manifest_paths +) +test_artifact_validation_rejects_self_attested_forged_legal_provenance = ( + artifact_contracts.test_artifact_validation_rejects_self_attested_forged_legal_provenance +) +test_artifact_build_uses_validated_git_blobs_after_worktree_mutation = ( + artifact_contracts.test_artifact_build_uses_validated_git_blobs_after_worktree_mutation +) +test_artifact_build_rejects_concurrent_config_replacement = ( + artifact_contracts.test_artifact_build_rejects_concurrent_config_replacement +) +test_official_tokenizer_copy_rejects_concurrent_replacement = ( + artifact_contracts.test_official_tokenizer_copy_rejects_concurrent_replacement +) +test_weight_snapshot_rejects_concurrent_replacement = ( + artifact_contracts.test_weight_snapshot_rejects_concurrent_replacement +) +test_artifact_rejects_self_attested_forged_canonical_weight = ( + artifact_contracts.test_artifact_rejects_self_attested_forged_canonical_weight +) +test_dplm2_artifact_materializes_non_decoder_cache_config = ( + artifact_contracts.test_dplm2_artifact_materializes_non_decoder_cache_config +) +test_complete_publication_rejects_synthetic_unresolved_checkpoint_license = ( + publish_contracts.test_complete_publication_rejects_synthetic_unresolved_checkpoint_license +) +test_complete_ankh_publish_rejects_missing_probe_binding = ( + publish_contracts.test_complete_ankh_publish_rejects_missing_probe_binding +) +test_complete_plan_rejects_unpinned_competing_remote_weight = ( + publish_contracts.test_complete_plan_rejects_unpinned_competing_remote_weight +) +test_complete_publish_rejects_hand_built_unknown_plan = ( + publish_contracts.test_complete_publish_rejects_hand_built_unknown_plan +) +test_complete_publish_rehashes_every_file_before_atomic_commit = ( + publish_contracts.test_complete_publish_rehashes_every_file_before_atomic_commit +) +test_files_only_plan_rejects_forged_registry_provenance = ( + publish_contracts.test_files_only_plan_rejects_forged_registry_provenance +) +test_files_only_plan_rejects_self_attested_invalid_bundle_data = ( + publish_contracts.test_files_only_plan_rejects_self_attested_invalid_bundle_data +) +test_files_only_plan_rejects_self_attested_stale_release_text = ( + publish_contracts.test_files_only_plan_rejects_self_attested_stale_release_text +) +test_files_only_plan_rejects_self_attested_substituted_bundle_member = ( + publish_contracts.test_files_only_plan_rejects_self_attested_substituted_bundle_member +) +test_files_only_plan_rejects_stale_runtime_revision = ( + publish_contracts.test_files_only_plan_rejects_stale_runtime_revision +) +test_files_only_plan_rejects_unknown_manifest_path = ( + publish_contracts.test_files_only_plan_rejects_unknown_manifest_path +) +test_files_only_plan_rejects_unlisted_and_sensitive_files = ( + publish_contracts.test_files_only_plan_rejects_unlisted_and_sensitive_files +) +test_files_only_plan_supports_every_manifest_model = ( + publish_contracts.test_files_only_plan_supports_every_manifest_model +) +test_files_only_publish_uses_preflighted_bytes_after_local_mutation = ( + publish_contracts.test_files_only_publish_uses_preflighted_bytes_after_local_mutation +) +test_files_only_publish_rejects_release_text_change_after_preflight = ( + publish_contracts.test_files_only_publish_rejects_release_text_change_after_preflight +) +test_files_only_publish_rejects_source_change_after_preflight = ( + publish_contracts.test_files_only_publish_rejects_source_change_after_preflight +) +test_release_snapshot_rejects_untracked_card_symlink_to_tracked_file = ( + publish_contracts.test_release_snapshot_rejects_untracked_card_symlink_to_tracked_file +) +test_required_complete_probe_groups_ankh_views = ( + publish_contracts.test_required_complete_probe_groups_ankh_views +) + + +def _synthetic_complete_ankh_registry( + artifact: Path, + shard_states: tuple[dict[str, torch.Tensor], ...], + tokenizer_payloads: dict[str, bytes], +) -> tuple[ModelRegistry, ModelSpec, dict[str, Any]]: + current_registry = get_model_registry() + current_spec = current_registry["ankh_base"] + shards = tuple( + artifact / f"model-{index:05d}-of-{len(shard_states):05d}.safetensors" + for index in range(1, len(shard_states) + 1) + ) + for path, state in zip(shards, shard_states, strict=True): + save_file(state, path, metadata={"format": "pt"}) + + weight_map = { + key: path.name for path, state in zip(shards, shard_states, strict=True) for key in state + } + total_size = sum( + tensor.numel() * tensor.element_size() + for state in shard_states + for tensor in state.values() + ) + index = artifact / "model.safetensors.index.json" + publish_contracts._write_json( + index, + {"metadata": {"total_size": total_size}, "weight_map": weight_map}, + ) + combined_state = {key: tensor for state in shard_states for key, tensor in state.items()} + canonical_state_sha256 = _canonical_state_sha256(combined_state) + canonical_weights: dict[str, Any] = { + "format": "safetensors", + "index": index.name, + "index_digest": f"sha256:{hash_file(index)}", + "max_shard_bytes": 1024, + "shards": {path.name: f"sha256:{hash_file(path)}" for path in shards}, + "source_schema": "official", + "state_transform": current_spec.family.state_transform, + "state_digest": { + "schema_version": 1, + "algorithm": "sha256", + "sha256": canonical_state_sha256, + }, + "tensor_count": len(weight_map), + "total_size": total_size, + } + official_files = [ + FileDigest(path=index.name, algorithm="sha256", digest=hash_file(index)), + *( + FileDigest(path=path.name, algorithm="sha256", digest=hash_file(path)) + for path in shards + ), + *( + FileDigest( + path=relative_name, + algorithm="sha256", + digest=hashlib.sha256(payload).hexdigest(), + ) + for relative_name, payload in sorted(tokenizer_payloads.items()) + ), + ] + synthetic_spec = replace( + current_spec, + family=replace( + current_spec.family, + requires_complete_weight_publication=True, + ), + official=CheckpointSource( + repo_id=current_spec.official.repo_id, + revision=current_spec.official.revision, + files=tuple(official_files), + ), + canonical_state_sha256=canonical_state_sha256, + ) + models = dict(current_registry) + models[synthetic_spec.id] = synthetic_spec + synthetic_registry = ModelRegistry( + schema_version=current_registry.schema_version, + upstreams=current_registry.upstreams, + families=current_registry.families, + models=models, + runtime_assets=current_registry.runtime_assets, + attention_kernels=current_registry.attention_kernels, + legal_files=current_registry.legal_files, + ) + return synthetic_registry, synthetic_spec, canonical_weights + + +def _complete_ankh_artifact( + root: Path, + registry: ModelRegistry, + spec: ModelSpec, + canonical_weights: dict[str, Any], + prepared_weights: Path, + tokenizer_payloads: dict[str, bytes], +) -> Path: + artifact = publish_contracts._files_only_artifact(root, spec) + stale_shard = artifact / "model-00001-of-00001.safetensors" + stale_shard.unlink() + for relative_name in ( + "model.safetensors.index.json", + *canonical_weights["shards"], + ): + source = prepared_weights / relative_name + (artifact / relative_name).write_bytes(source.read_bytes()) + for relative_name, payload in tokenizer_payloads.items(): + (artifact / relative_name).write_bytes(payload) + + seed_provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + runtime_revision = seed_provenance["runtime_revision"] + source_tree_sha256 = seed_provenance["source_tree_sha256"] + runtime_bundle_sha256 = seed_provenance["runtime_bundle_sha256"] + release_tool_revision = seed_provenance["release_tool_revision"] + release_tool_sha256 = seed_provenance["release_tool_sha256"] + provenance = _provenance( + registry, + spec, + canonical_weights, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + ) + publish_contracts._write_json(artifact / "provenance.json", provenance) + runtime_attestation = _runtime_attestation( + artifact, + spec, + weights_revision=spec.fast.revision, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + ) + publish_contracts._write_json( + artifact / "runtime-attestation.json", + runtime_attestation, + ) + publish_contracts._write_json( + artifact / "artifact-manifest.json", + _content_manifest(artifact), + ) + return artifact + + +def test_complete_ankh_multishard_inventory_is_published_atomically( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + prepared_weights = tmp_path / "prepared-weights" + prepared_weights.mkdir() + shard_states = ( + { + "shared.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), # (vocab, d) + "encoder.block.0.layer.0.SelfAttention.q.weight": torch.ones(4, 4), # (d, d) + }, + { + "decoder.block.0.layer.1.EncDecAttention.q.weight": torch.ones(4, 4), # (d, d) + "lm_head.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), # (vocab, d) + }, + ) + tokenizer_payloads = { + "special_tokens_map.json": b'{"eos_token":"","pad_token":""}\n', + "tokenizer.json": b'{"version":"1.0"}\n', + "tokenizer_config.json": b'{"eos_token":"","pad_token":""}\n', + } + registry, spec, canonical_weights = _synthetic_complete_ankh_registry( + prepared_weights, + shard_states, + tokenizer_payloads, + ) + assert spec.family.requires_complete_weight_publication is True + assert spec.artifact_source == "official" + monkeypatch.setattr(publish_module, "get_model_registry", lambda: registry) + monkeypatch.setattr(publish_contracts, "get_model_registry", lambda: registry) + artifact = _complete_ankh_artifact( + tmp_path, + registry, + spec, + canonical_weights, + prepared_weights, + tokenizer_payloads, + ) + validate_artifact(artifact, spec=spec, registry=registry) + validated_index = validate_weight_artifact(artifact) + replacement_weights = { + "model.safetensors.index.json", + *canonical_weights["shards"], + } + assert set(validated_index["weight_map"].values()) == ( + replacement_weights - {"model.safetensors.index.json"} + ) + + api = publish_contracts.FakeApi(spec) + probe_calls: list[tuple[str, Path, str]] = [] + + def required_autoclass_probe(probe_spec: ModelSpec, probe_artifact: Path) -> tuple[str, ...]: + config = json.loads((probe_artifact / "config.json").read_text(encoding="utf-8")) + auto_map = config["auto_map"] + assert auto_map["AutoModel"].endswith(".FastAnkhModel") + assert auto_map["AutoModelForSeq2SeqLM"].endswith(".FastAnkhForConditionalGeneration") + probe_calls.append( + ( + probe_spec.id, + probe_artifact, + hash_file(probe_artifact / "artifact-manifest.json"), + ) + ) + return ("AutoModel", "AutoModelForSeq2SeqLM") + + monkeypatch.setattr( + publish_module, + "_run_required_complete_autoclass_probe", + required_autoclass_probe, + ) + real_prepare_complete_plan = publish_module.prepare_complete_plan + prepared_plans = [] + + def recording_prepare_complete_plan(*args: Any, **kwargs: Any) -> Any: + prepared_plan = real_prepare_complete_plan(*args, **kwargs) + prepared_plans.append(prepared_plan) + return prepared_plan + + monkeypatch.setattr( + publish_module, + "prepare_complete_plan", + recording_prepare_complete_plan, + ) + plan = publish_module.prepare_complete_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, + ) + assert prepared_plans == [plan] + assert plan.validated_auto_classes == ("AutoModel", "AutoModelForSeq2SeqLM") + assert set(plan.replacement_weight_paths) == replacement_weights + assert replacement_weights.issubset(plan.files) + assert plan.deletes == ("model-00001-of-00001.safetensors",) + + publish_module.publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Atomic ANKH multi-shard publication", + ) + + assert prepared_plans == [plan, plan] + assert len(probe_calls) == 2 + assert probe_calls[0] == probe_calls[1] + assert len(api.create_commit_calls) == 1 + call = api.create_commit_calls[0] + assert call["parent_commit"] == "a" * 40 + operations = call["operations"] + additions = { + operation.path_in_repo + for operation in operations + if isinstance(operation, CommitOperationAdd) + } + deletions = { + operation.path_in_repo + for operation in operations + if isinstance(operation, CommitOperationDelete) + } + assert additions == set(plan.files) + assert replacement_weights.issubset(additions) + assert deletions == {"model-00001-of-00001.safetensors"} + + (artifact / "model-00002-of-00002.safetensors").unlink() + with pytest.raises(ArtifactError, match="index and artifact shard files differ"): + validate_weight_artifact(artifact) diff --git a/tests/cpu/test_remote_autoclass_contracts.py b/tests/cpu/test_remote_autoclass_contracts.py new file mode 100644 index 0000000..c7a9ff6 --- /dev/null +++ b/tests/cpu/test_remote_autoclass_contracts.py @@ -0,0 +1,435 @@ +"""Actual offline Transformers AutoClass dispatch from tiny remote-code artifacts.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import pytest +from pathlib import Path +from typing import Any + +from fastplms.models.ankh.modeling_ankh import FastAnkhConfig +from fastplms.models.boltz.modeling_boltz2 import Boltz2Config +from fastplms.models.dplm.modeling_dplm import DPLMConfig +from fastplms.models.dplm2.modeling_dplm2 import DPLM2Config +from fastplms.models.e1.modeling_e1 import E1Config +from fastplms.models.esm2.modeling_fastesm import FastEsmConfig +from fastplms.models.esm3.modeling_esm3 import FastESM3Config +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ESMplusplusConfig +from fastplms.models.esmfold.modeling_fast_esmfold import FastEsmFoldConfig +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.modeling_esmfold2_common import NUM_RES_TYPES +from fastplms.registry import ModelFamily, get_model_registry +from tools.artifacts.build import ( + ArtifactError, + _artifact_auto_map, + _runtime_source_entries, + _write_bootstrap, + _write_runtime_bundle, + _write_runtime_snapshot, +) +from tools.artifacts.offline_probe import _CPU_CONTRACT_MARKER, _runtime_site_packages +from tools.artifacts.publish import _validate_publishable_non_weight_path + + +_ROOT = Path(__file__).resolve().parents[2] +_PROBE = _ROOT / "tools" / "artifacts" / "offline_probe.py" +_STRUCTURE_FAMILIES = frozenset({"boltz2", "esmfold", "esmfold2"}) +_REMOTE_RESAVE_FAMILIES = frozenset({"ankh", "dplm", "dplm2", "esm2", "esm3", "esm_plusplus"}) + + +def _transformer_values(vocab_size: int) -> dict[str, Any]: + return { + "vocab_size": vocab_size, + "hidden_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "intermediate_size": 16, + "hidden_dropout_prob": 0.0, + "attention_probs_dropout_prob": 0.0, + "max_position_embeddings": 16, + "pad_token_id": 1, + "bos_token_id": 0, + "eos_token_id": 2, + "mask_token_id": min(7, vocab_size - 1), + "position_embedding_type": "rotary", + "attn_backend": "eager", + "num_labels": 3, + } + + +def _tiny_esmfold_config() -> FastEsmFoldConfig: + return FastEsmFoldConfig( + vocab_size=33, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + is_folding_model=True, + attn_backend="eager", + esmfold_config={ + "fp16_esm": False, + "bypass_lm": True, + "lddt_head_hid_dim": 4, + "trunk": { + "num_blocks": 1, + "sequence_state_dim": 8, + "pairwise_state_dim": 4, + "sequence_head_width": 4, + "pairwise_head_width": 2, + "position_bins": 4, + "max_recycles": 1, + "chunk_size": None, + "structure_module": { + "sequence_dim": 8, + "pairwise_dim": 4, + "ipa_dim": 2, + "resnet_dim": 4, + "num_heads_ipa": 2, + "num_qk_points": 1, + "num_v_points": 1, + "dropout_rate": 0.0, + "num_blocks": 1, + "num_transition_layers": 1, + "num_resnet_blocks": 1, + "num_angles": 7, + }, + }, + }, + ) + + +def _tiny_esmfold2_config() -> ESMFold2Config: + atom_token_width = 8 + input_feature_width = atom_token_width // 2 + 2 * NUM_RES_TYPES + 1 + return ESMFold2Config( + type="release", + d_single=8, + d_pair=8, + num_loops=0, + num_diffusion_samples=1, + lm_d_model=8, + lm_num_layers=1, + inputs={ + "d_inputs": input_feature_width, + "atom_encoder": { + "d_atom": 8, + "d_token": atom_token_width, + "n_blocks": 0, + "n_heads": 2, + "swa_window_size": 32, + "expansion_ratio": 2, + "n_spatial_rope_pairs_per_axis": 1, + "n_uid_rope_pairs": 1, + }, + }, + folding_trunk={"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + structure_head={ + "diffusion_module": { + "c_atom": 8, + "c_token": 8, + "c_z": 8, + "c_s_inputs": input_feature_width, + "fourier_dim": 8, + "atom_num_blocks": 0, + "atom_num_heads": 2, + "token_num_blocks": 0, + "token_num_heads": 2, + "transition_multiplier": 2, + }, + "distogram_bins": 8, + "inference_num_steps": 1, + }, + confidence_head={ + "enabled": False, + "folding_trunk": {"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + "num_plddt_bins": 4, + "num_pde_bins": 4, + "num_pae_bins": 4, + "distogram_bins": 8, + }, + msa_encoder={"enabled": False}, + lm_encoder={"enabled": False, "n_layers": 0}, + parcae={"enabled": True, "min_steps": 1, "max_steps": 1, "coda_n_layers": 0}, + ) + + +def _tiny_config(family_id: str) -> Any: + if family_id == "esm2": + values = _transformer_values(16) + values["position_embedding_type"] = "absolute" + return FastEsmConfig(**values) + if family_id == "esm_plusplus": + return ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + dropout=0.0, + pad_token_id=1, + mask_token_id=7, + num_labels=3, + attn_backend="eager", + ) + if family_id == "esm3": + return FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + if family_id == "e1": + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + max_num_sequences=4, + max_num_positions_within_seq=16, + max_num_positions_global=16, + attn_backend="sdpa", + dtype="float32", + num_labels=3, + ) + config.use_cache = False + return config + if family_id == "dplm": + return DPLMConfig(**_transformer_values(16)) + if family_id == "dplm2": + values = _transformer_values(64) + values["attn_backend"] = "sdpa" + return DPLM2Config(**values) + if family_id == "ankh": + return FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + num_decoder_layers=1, + dropout_rate=0.0, + pad_token_id=0, + eos_token_id=1, + decoder_start_token_id=0, + attn_backend="eager", + use_cache=False, + num_labels=3, + ) + if family_id == "boltz2": + return Boltz2Config(core_kwargs={"width": 3}) + if family_id == "esmfold": + return _tiny_esmfold_config() + if family_id == "esmfold2": + return _tiny_esmfold2_config() + raise AssertionError(f"Missing tiny configuration for manifest family {family_id!r}.") + + +def _write_cpu_artifact(root: Path, family: ModelFamily) -> Path: + registry = get_model_registry() + spec = registry[family.representative] + if dict(spec.auto_map) != dict(family.auto_map): + raise AssertionError(f"Representative {spec.id} has a model-specific AutoMap override.") + + artifact = root / family.id / "artifact" + artifact.mkdir(parents=True) + config = _tiny_config(family.id) + config.auto_map = _artifact_auto_map(spec) + config.fastplms_cpu_contract_only = True + config.save_pretrained(artifact) + config_path = artifact / "config.json" + artifact_config = json.loads(config_path.read_text(encoding="utf-8")) + artifact_config["auto_map"] = _artifact_auto_map(spec) + config_path.write_text( + json.dumps(artifact_config, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + (artifact / _CPU_CONTRACT_MARKER).write_text( + json.dumps( + {"release_artifact": False, "schema_version": 1, "scope": "tests/cpu"}, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + + runtime_root = root / family.id / "runtime" / "fastplms" + payloads = { + target.as_posix(): source.read_bytes() + for source, target in _runtime_source_entries(_ROOT, spec) + } + _write_runtime_snapshot(runtime_root, payloads) + runtime_hash = _write_runtime_bundle(artifact / "fastplms_bundle.py", runtime_root) + _write_bootstrap(artifact / "modeling_fastplms.py", spec, runtime_hash) + return artifact + + +def _case_payload(family: ModelFamily) -> list[dict[str, object]]: + return [ + { + "auto_class": auto_class, + "class_path": class_path, + "expected_missing_key_prefixes": [], + "expected_unexpected_key_prefixes": [], + } + for auto_class, class_path in sorted(family.auto_map.items()) + ] + + +def _run_family_probe( + root: Path, + family: ModelFamily, + artifact: Path, +) -> subprocess.CompletedProcess[str]: + family_root = root / family.id + cases_path = family_root / "cases.json" + output_path = family_root / "results.json" + cases_path.write_text( + json.dumps(_case_payload(family), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + command = [ + sys.executable, + "-I", + "-S", + str(_PROBE), + "--artifact", + str(artifact), + "--family", + family.id, + "--bf16-execution", + family.bf16_execution, + "--cases-file", + str(cases_path), + "--implementation", + "artifact", + "--output", + str(output_path), + "--tiny-cpu-contract", + ] + for site_packages in _runtime_site_packages(): + command.extend(("--runtime-site-package", str(site_packages))) + environment = os.environ.copy() + environment.pop("PYTHONHOME", None) + environment.pop("PYTHONPATH", None) + environment.update( + { + "CUDA_VISIBLE_DEVICES": "", + "HF_HOME": str(family_root / "hf-home"), + "HF_HUB_OFFLINE": "1", + "HF_MODULES_CACHE": str(family_root / "modules"), + "PYTHONNOUSERSITE": "1", + "TRANSFORMERS_OFFLINE": "1", + } + ) + return subprocess.run( + command, + cwd=family_root, + env=environment, + capture_output=True, + text=True, + check=False, + timeout=20, + ) + + +_FAMILY_IDS = tuple(sorted(get_model_registry().families)) + + +@pytest.mark.parametrize("family_id", _FAMILY_IDS) +def test_every_family_dispatches_all_advertised_remote_autoclasses_offline( + family_id: str, + tmp_path: Path, +) -> None: + registry = get_model_registry() + family = registry.families[family_id] + artifact = _write_cpu_artifact(tmp_path, family) + completed = _run_family_probe(tmp_path, family, artifact) + output_path = tmp_path / family_id / "results.json" + assert completed.returncode == 0, completed.stdout + completed.stderr + assert output_path.is_file() + output = json.loads(output_path.read_text(encoding="utf-8")) + assert set(output) == set(family.auto_map) + + for auto_class, result in output.items(): + expected = family.auto_map[auto_class] + assert result["class"] == expected.rsplit(".", maxsplit=1)[1] + if auto_class != "AutoConfig" and family_id not in _STRUCTURE_FAMILIES: + assert result["resized_vocab"] >= 9 + assert result["tuple_fields"] >= 1 + if family_id in _STRUCTURE_FAMILIES and auto_class != "AutoConfig": + assert result["structure_forward_delegated"] is True + if family_id in _REMOTE_RESAVE_FAMILIES: + assert output["AutoModel"]["resaved"] is True + marker = json.loads((artifact / _CPU_CONTRACT_MARKER).read_text(encoding="utf-8")) + assert marker == { + "release_artifact": False, + "schema_version": 1, + "scope": "tests/cpu", + } + assert not (artifact / "artifact-manifest.json").exists() + assert not (artifact / "provenance.json").exists() + assert not (artifact / "runtime-attestation.json").exists() + + +def test_structure_auto_dispatch_is_complemented_by_public_forward_contracts() -> None: + from tests.cpu import test_structure_contracts as structure_contracts + + for test_name in ( + "test_boltz_public_forward_honors_output_controls_backward_and_reload", + "test_fast_esmfold_public_forward_honors_output_controls_and_backward", + "test_fast_esmfold_tiny_model_saves_and_reloads_exact_state", + "test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides", + "test_esmfold2_advertised_models_tiny_init_backward_and_save_reload", + ): + assert callable(getattr(structure_contracts, test_name)) + + +def test_grouped_remote_dispatch_covers_exactly_37_family_entries() -> None: + registry = get_model_registry() + assert sum(len(family.auto_map) for family in registry.families.values()) == 37 + + +def test_isolated_probe_blocks_reference_reads_before_remote_code_exec( + tmp_path: Path, +) -> None: + registry = get_model_registry() + family = registry.families["esm2"] + artifact = _write_cpu_artifact(tmp_path, family) + bootstrap = artifact / "modeling_fastplms.py" + forbidden = _ROOT / "vendor" / "upstream" / "forbidden-cpu-probe-read" + bootstrap.write_text( + "from pathlib import Path\n" + f"Path({str(forbidden)!r}).read_bytes()\n" + bootstrap.read_text(encoding="utf-8"), + encoding="utf-8", + ) + + completed = _run_family_probe(tmp_path, family, artifact) + + assert completed.returncode != 0 + assert "may not access submodule/reference path" in completed.stdout + completed.stderr + + +def test_cpu_contract_marker_is_rejected_by_publication_allowlist(tmp_path: Path) -> None: + marker = tmp_path / _CPU_CONTRACT_MARKER + marker.write_text("{}\n", encoding="utf-8") + + with pytest.raises( + ArtifactError, + match=r"sensitive|outside the publication allowlist", + ): + _validate_publishable_non_weight_path( + _CPU_CONTRACT_MARKER, + marker, + declared_assets=frozenset(), + declared_legal_paths=frozenset(), + ) diff --git a/tests/cpu/test_resource_telemetry.py b/tests/cpu/test_resource_telemetry.py new file mode 100644 index 0000000..f5c1869 --- /dev/null +++ b/tests/cpu/test_resource_telemetry.py @@ -0,0 +1,393 @@ +"""Deterministic contracts for concurrent Linux process-tree memory evidence.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +import pytest +from pathlib import Path + +from tests.cpu.resource_telemetry import ( + ConcurrentProcessTreeSampler, + MemoryEvidenceError, + sample_process_tree_memory, + select_concurrent_memory_gate, +) + + +def _write_fake_process( + proc_root: Path, + process_id: int, + *, + children: tuple[int, ...], + rss_kib: int, + pss_kib: int, + state: str = "S", +) -> None: + process_root = proc_root / str(process_id) + task_root = process_root / "task" / str(process_id) + task_root.mkdir(parents=True) + (task_root / "children").write_text( + " ".join(str(child_id) for child_id in children), + encoding="utf-8", + ) + (process_root / "status").write_text( + f"Name:\tpython\nVmRSS:\t{rss_kib} kB\n", + encoding="utf-8", + ) + # Fields after comm begin at field 3 (state); field 22 is starttime. + stat_fields = [state, *(["0"] * 18), str(process_id * 10)] + (process_root / "stat").write_text( + f"{process_id} (python) {' '.join(stat_fields)}\n", + encoding="utf-8", + ) + (process_root / "smaps_rollup").write_text( + f"Pss:\t{pss_kib} kB\n", + encoding="utf-8", + ) + + +def test_overlapping_worker_and_child_roots_are_deduplicated(tmp_path: Path) -> None: + proc_root = tmp_path / "proc" + _write_fake_process( + proc_root, + 100, + children=(200, 300), + rss_kib=100, + pss_kib=50, + ) + _write_fake_process( + proc_root, + 200, + children=(400,), + rss_kib=200, + pss_kib=100, + ) + _write_fake_process( + proc_root, + 300, + children=(), + rss_kib=300, + pss_kib=150, + ) + _write_fake_process( + proc_root, + 400, + children=(), + rss_kib=400, + pss_kib=200, + ) + + snapshot = sample_process_tree_memory( + proc_root=proc_root, + # These roots deliberately overlap: 200 and 400 are descendants of + # 100, and 400 is also a descendant of 200. + root_pids=(100, 200, 400, 200), + ) + + assert snapshot["root_pids"] == [100, 200, 400] + assert snapshot["process_ids"] == [100, 200, 300, 400] + assert snapshot["process_count"] == 4 + assert snapshot["aggregate_rss_bytes"] == 1_000 * 1024 + assert snapshot["aggregate_pss_bytes"] == 500 * 1024 + assert snapshot["pss_complete"] is True + assert len(snapshot["processes"]) == len(set(snapshot["process_ids"])) + + +def test_process_exit_during_pss_read_uses_only_that_process_rss_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tests.cpu.resource_telemetry as telemetry + + proc_root = tmp_path / "proc" + _write_fake_process( + proc_root, + 100, + children=(200,), + rss_kib=100, + pss_kib=50, + ) + _write_fake_process( + proc_root, + 200, + children=(), + rss_kib=200, + pss_kib=100, + ) + original_read_pss = telemetry._read_pss_bytes + + def disappearing_pss(root: Path, process_id: int) -> tuple[int | None, str | None]: + if process_id == 200: + process_root = root / str(process_id) + for path in sorted(process_root.rglob("*"), reverse=True): + if path.is_file(): + path.unlink() + else: + path.rmdir() + process_root.rmdir() + return None, "smaps_rollup disappeared during sampling" + return original_read_pss(root, process_id) + + monkeypatch.setattr(telemetry, "_read_pss_bytes", disappearing_pss) + snapshot = sample_process_tree_memory(root_pids=(100,), proc_root=proc_root) + + assert snapshot["process_ids"] == [100, 200] + assert snapshot["pss_complete"] is False + assert snapshot["aggregate_pss_bytes"] is None + assert snapshot["aggregate_hybrid_bytes"] == 250 * 1024 + child = next(process for process in snapshot["processes"] if process["pid"] == 200) + assert child["accounted_bytes"] == 200 * 1024 + assert child["accounting_metric"] == "resident-set-size-fallback" + assert snapshot["transient_process_event_count"] == 1 + assert "pid 200: exited before PSS sampling completed" in snapshot[ + "transient_process_events" + ] + assert snapshot["pss_fallback_reasons"] == [ + "pid 200: smaps_rollup disappeared during sampling" + ] + + +def test_zombie_without_smaps_is_zero_resident_and_keeps_pss_complete( + tmp_path: Path, +) -> None: + proc_root = tmp_path / "proc" + _write_fake_process( + proc_root, + 100, + children=(200,), + rss_kib=100, + pss_kib=50, + ) + _write_fake_process( + proc_root, + 200, + children=(), + rss_kib=999, + pss_kib=999, + state="Z", + ) + (proc_root / "200" / "smaps_rollup").unlink() + + snapshot = sample_process_tree_memory(root_pids=(100,), proc_root=proc_root) + + zombie = next(process for process in snapshot["processes"] if process["pid"] == 200) + assert zombie["rss_bytes"] == 0 + assert zombie["pss_bytes"] == 0 + assert zombie["accounted_bytes"] == 0 + assert zombie["state"] == "Z" + assert snapshot["pss_complete"] is True + + +def test_persistent_pss_denial_uses_only_denied_process_rss( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tests.cpu.resource_telemetry as telemetry + + proc_root = tmp_path / "proc" + _write_fake_process( + proc_root, + 100, + children=(200,), + rss_kib=100, + pss_kib=50, + ) + _write_fake_process( + proc_root, + 200, + children=(), + rss_kib=200, + pss_kib=100, + ) + original_read_pss = telemetry._read_pss_bytes + + def denied_pss(root: Path, process_id: int) -> tuple[int | None, str | None]: + if process_id == 200: + return None, "smaps_rollup permission denied" + return original_read_pss(root, process_id) + + monkeypatch.setattr(telemetry, "_read_pss_bytes", denied_pss) + snapshot = sample_process_tree_memory(root_pids=(100,), proc_root=proc_root) + gate = select_concurrent_memory_gate( + { + "sample_count": 1, + "peak_concurrent_rss_bytes": snapshot["aggregate_rss_bytes"], + "peak_concurrent_hybrid_bytes": snapshot["aggregate_hybrid_bytes"], + "peak_concurrent_pss_bytes": None, + "pss_complete_for_all_samples": False, + "pss_fallback_reasons": snapshot["pss_fallback_reasons"], + } + ) + + assert snapshot["aggregate_rss_bytes"] == 300 * 1024 + assert snapshot["aggregate_hybrid_bytes"] == 250 * 1024 + assert gate["metric"] == "per-process-pss-rss-hybrid" + assert gate["peak_bytes"] == 250 * 1024 + assert gate["fallback_used"] is True + assert gate["fallback_is_conservative"] is True + + +def test_sampler_captures_a_live_child_process() -> None: + assert sys.platform.startswith("linux") + script = textwrap.dedent( + """ + import os + import sys + + payload = bytearray(8 * 1024 * 1024) + for offset in range(0, len(payload), 4096): + payload[offset] = 1 + print(os.getpid(), flush=True) + sys.stdin.buffer.read(1) + """ + ) + sampler = ConcurrentProcessTreeSampler( + root_pids=(os.getpid(),), + sample_interval_seconds=0.01, + ) + sampler.start() + child = subprocess.Popen( + [sys.executable, "-c", script], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + summary: dict[str, object] + snapshot: dict[str, object] + try: + assert child.stdout is not None + child_id = int(child.stdout.readline().decode("utf-8").strip()) + assert child_id == child.pid + snapshot = sampler.sample_now() + assert child_id in snapshot["process_ids"] + child_record = next( + process for process in snapshot["processes"] if process["pid"] == child_id + ) + assert child_record["rss_bytes"] >= 8 * 1024**2 + finally: + if child.poll() is None and child.stdin is not None: + try: + child.stdin.write(b"x") + child.stdin.flush() + except BrokenPipeError: + pass + try: + child.wait(timeout=3) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=3) + for stream in (child.stdin, child.stdout, child.stderr): + if stream is not None: + stream.close() + summary = sampler.stop() + + assert child.returncode == 0 + assert child.pid in summary["observed_process_ids"] + assert summary["sample_count"] >= 3 + gate = select_concurrent_memory_gate(summary) + assert gate["peak_bytes"] > 0 + if gate["fallback_used"]: + assert gate["fallback_is_conservative"] is True + assert gate["metric"] == "per-process-pss-rss-hybrid" + else: + assert gate["metric"] == "proportional-set-size" + + +def test_pss_is_preferred_and_incomplete_pss_uses_per_process_hybrid() -> None: + preferred = select_concurrent_memory_gate( + { + "sample_count": 4, + "peak_concurrent_rss_bytes": 900, + "peak_concurrent_hybrid_bytes": 600, + "peak_concurrent_pss_bytes": 600, + "pss_complete_for_all_samples": True, + "pss_fallback_reasons": [], + } + ) + assert preferred == { + "metric": "proportional-set-size", + "source": "/proc//smaps_rollup:Pss", + "peak_bytes": 600, + "fallback_used": False, + "fallback_is_conservative": False, + "fallback_reasons": [], + } + + fallback = select_concurrent_memory_gate( + { + "sample_count": 4, + "peak_concurrent_rss_bytes": 900, + "peak_concurrent_hybrid_bytes": 700, + "peak_concurrent_pss_bytes": 600, + "pss_complete_for_all_samples": False, + "pss_fallback_reasons": ["pid 123: smaps_rollup permission denied"], + } + ) + assert fallback["metric"] == "per-process-pss-rss-hybrid" + assert fallback["peak_bytes"] == 700 + assert fallback["fallback_used"] is True + assert fallback["fallback_is_conservative"] is True + assert fallback["fallback_reasons"] == [ + "pid 123: smaps_rollup permission denied" + ] + + with pytest.raises(MemoryEvidenceError, match="no positive RSS peak"): + select_concurrent_memory_gate( + { + "sample_count": 1, + "peak_concurrent_rss_bytes": None, + "peak_concurrent_hybrid_bytes": None, + "pss_complete_for_all_samples": False, + } + ) + + +def test_memory_over_budget_forces_nonzero_pytest_exit(tmp_path: Path) -> None: + test_file = tmp_path / "test_synthetic_gate.py" + test_file.write_text("def test_passes():\n assert True\n", encoding="utf-8") + script = textwrap.dedent( + f""" + import pytest + + from tests.cpu.conftest import _enforce_concurrent_memory_budget + + class SyntheticMemoryGate: + @pytest.hookimpl(trylast=True) + def pytest_sessionfinish(self, session, exitstatus): + del exitstatus + _enforce_concurrent_memory_budget( + session, + {{ + "metric": "per-process-pss-rss-hybrid", + "peak_bytes": 4 * 1024**3 + 1, + }}, + ) + + raise SystemExit( + pytest.main( + [{str(test_file)!r}, "-q", "-p", "no:cacheprovider"], + plugins=[SyntheticMemoryGate()], + ) + ) + """ + ) + workspace = Path(__file__).resolve().parents[2] + environment = dict(os.environ) + environment["PYTHONPATH"] = os.pathsep.join( + (str(workspace), environment.get("PYTHONPATH", "")) + ) + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=workspace, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=8, + ) + + assert completed.returncode == int(pytest.ExitCode.TESTS_FAILED) + assert "concurrent memory budget exceeded" in completed.stdout + completed.stderr diff --git a/tests/cpu/test_sequence_autoclass_contracts.py b/tests/cpu/test_sequence_autoclass_contracts.py new file mode 100644 index 0000000..817774b --- /dev/null +++ b/tests/cpu/test_sequence_autoclass_contracts.py @@ -0,0 +1,671 @@ +"""Tiny end-to-end contracts for sequence-family public and advertised models.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from transformers.modeling_outputs import ModelOutput + +from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmForSequenceClassification, + FastEsmForTokenClassification, + FastEsmModel, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ESMplusplusForSequenceClassification, + ESMplusplusForTokenClassification, + ESMplusplusModel, +) +from tests.integration import test_dplm_generation as dplm_contracts +from tests.unit import test_e1_cache_contract as e1_contracts + + +def _esm2_config() -> FastEsmConfig: + return FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=5, + num_labels=3, + position_embedding_type="absolute", + attn_backend="eager", + ) + + +def _inputs() -> tuple[torch.Tensor, torch.Tensor]: + input_ids = torch.tensor([[0, 3, 4, 2, 1], [0, 6, 2, 1, 1]]) # (b=2, l=5) + return input_ids, input_ids.ne(1) # (b, l), (b, l) + + +def _dplm_config_values(vocab_size: int) -> dict[str, object]: + values = dplm_contracts._common_config(vocab_size) + values.update( + { + "hidden_size": 8, + "num_attention_heads": 2, + "intermediate_size": 16, + "max_position_embeddings": 16, + "attn_backend": "eager", + } + ) + return values + + +def _dplm2_config_values() -> dict[str, object]: + """Return a tiny config using DPLM2's sole manifest backend.""" + + values = _dplm_config_values(64) + values["attn_backend"] = "sdpa" + return values + + +@pytest.mark.parametrize("serialized_backend", (None, "sdpa")) +def test_dplm2_legacy_or_explicit_backend_resolves_to_manifest_sdpa( + serialized_backend: str | None, +) -> None: + values = _dplm2_config_values() + values["attn_backend"] = serialized_backend + model = dplm_contracts.DPLM2Model( + dplm_contracts.DPLM2Config(**values) + ) + + assert model.config.attn_backend == "sdpa" + assert model.attn_backend == "sdpa" + + +def test_dplm2_rejects_eager_instead_of_expanding_its_backend_claim() -> None: + values = _dplm2_config_values() + values["attn_backend"] = "eager" + + with pytest.raises(ValueError, match="does not support 'eager'"): + dplm_contracts.DPLM2Model(dplm_contracts.DPLM2Config(**values)) + + +@pytest.mark.parametrize( + "model_class", + ( + dplm_contracts.DPLM2Model, + dplm_contracts.DPLM2ForMaskedLM, + dplm_contracts.DPLM2ForSequenceClassification, + dplm_contracts.DPLM2ForTokenClassification, + ), +) +def test_dplm2_multimodal_wrappers_require_types_with_precomputed_embeddings( + model_class: type, +) -> None: + model = model_class( + dplm_contracts.DPLM2Config( + **_dplm2_config_values(), + num_labels=3, + ) + ).eval() + input_ids = torch.tensor([[0, 6, 7, 2]]) + inputs_embeds = model.get_input_embeddings()(input_ids) + + with pytest.raises(ValueError, match="type_ids is required"): + model( + inputs_embeds=inputs_embeds, + attention_mask=torch.ones_like(input_ids), + ) + + +def _assert_nested_output_close( + actual: object, + expected: object, + *, + exact: bool = False, +) -> None: + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + if exact: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + else: + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, e1_contracts.DynamicCache): + assert isinstance(actual, e1_contracts.DynamicCache) + assert len(actual.key_cache) == len(expected.key_cache) + assert len(actual.value_cache) == len(expected.value_cache) + for actual_tensor, expected_tensor in zip( + actual.key_cache + actual.value_cache, + expected.key_cache + expected.value_cache, + strict=True, + ): + if exact: + torch.testing.assert_close( + actual_tensor, + expected_tensor, + rtol=0.0, + atol=0.0, + ) + else: + torch.testing.assert_close(actual_tensor, expected_tensor) + return + if isinstance(expected, (tuple, list)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_value, expected_value in zip(actual, expected, strict=True): + _assert_nested_output_close(actual_value, expected_value, exact=exact) + return + assert actual == expected + + +def _assert_exact_state_round_trip( + model: torch.nn.Module, + reloaded: torch.nn.Module, +) -> None: + source_state = model.state_dict() # parameter name -> checkpoint-shaped tensor + restored_state = reloaded.state_dict() # parameter name -> checkpoint-shaped tensor + assert set(restored_state) == set(source_state) + for name, tensor in source_state.items(): + torch.testing.assert_close( + restored_state[name], + tensor, + rtol=0.0, + atol=0.0, + ) + + +def _assert_exact_output_round_trip( + model: torch.nn.Module, + reloaded: torch.nn.Module, + **model_inputs: torch.Tensor, +) -> None: + forward_controls = { + "output_attentions": True, + "output_hidden_states": True, + "output_s_max": True, + } + with torch.inference_mode(): + source_output = model( + **model_inputs, + **forward_controls, + return_dict=True, + ) + restored_output = reloaded( + **model_inputs, + **forward_controls, + return_dict=True, + ) + restored_tuple = reloaded( + **model_inputs, + **forward_controls, + return_dict=False, + ) + + assert type(restored_output) is type(source_output) + assert tuple(restored_output.keys()) == tuple(source_output.keys()) + _assert_nested_output_close( + restored_output.to_tuple(), + source_output.to_tuple(), + exact=True, + ) + _assert_nested_output_close( + restored_tuple, + source_output.to_tuple(), + exact=True, + ) + + +@pytest.mark.parametrize( + "model_class", + ( + FastEsmModel, + FastEsmForMaskedLM, + FastEsmForSequenceClassification, + FastEsmForTokenClassification, + ), +) +def test_esm2_advertised_models_forward_loss_backward_resize_and_reload( + model_class: type, + tmp_path: Path, +) -> None: + model = model_class(_esm2_config()).eval() + input_ids, attention_mask = _inputs() + structured = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + tuple_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + + assert isinstance(structured, ModelOutput) + assert isinstance(tuple_output, tuple) + _assert_nested_output_close(tuple_output, structured.to_tuple()) + primary = ( + structured.last_hidden_state + if model_class is FastEsmModel + else structured.logits + ) + torch.testing.assert_close(tuple_output[0], primary) + assert structured.hidden_states is not None + assert structured.attentions is not None + assert structured.s_max is not None + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model( + input_ids=input_ids, + attention_mask=attention_mask, + unexpected_cpu_contract=True, + ) + + if model_class is FastEsmForMaskedLM: + labels = input_ids.masked_fill(~attention_mask, -100) + elif model_class is FastEsmForSequenceClassification: + labels = torch.tensor([1, 2]) + elif model_class is FastEsmForTokenClassification: + labels = input_ids.remainder(3).masked_fill(~attention_mask, -100) + else: + labels = None + if labels is None: + loss = structured.last_hidden_state.square().mean() + else: + loss_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + loss_tuple = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + _assert_nested_output_close(loss_tuple, loss_output.to_tuple()) + torch.testing.assert_close(loss_tuple[0], loss_output.loss) + torch.testing.assert_close(loss_tuple[1], loss_output.logits) + loss = loss_output.loss + assert loss is not None and torch.isfinite(loss) + loss.backward() + assert any(parameter.grad is not None for parameter in model.parameters()) + + model.resize_token_embeddings(19) + assert model.get_input_embeddings().num_embeddings == 19 + output_embeddings = model.get_output_embeddings() + if output_embeddings is not None: + assert output_embeddings.out_features == 19 + + save_dir = tmp_path / model_class.__name__ + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained(save_dir, local_files_only=True).eval() + assert reloaded.get_input_embeddings().num_embeddings == 19 + _assert_exact_state_round_trip(model, reloaded) + _assert_exact_output_round_trip( + model, + reloaded, + input_ids=input_ids, + attention_mask=attention_mask, + ) + + +@pytest.mark.parametrize( + ("model_class", "kind"), + ( + (ESMplusplusModel, "base"), + (ESMplusplusForMaskedLM, "mlm"), + (ESMplusplusForSequenceClassification, "sequence"), + (ESMplusplusForTokenClassification, "token"), + ), +) +def test_esmc_public_models_forward_loss_backward_resize_and_reload( + model_class: type, + kind: str, + tmp_path: Path, +) -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + dropout=0.0, + pad_token_id=1, + mask_token_id=5, + num_labels=3, + attn_backend="eager", + ) + model = model_class(config).eval() + input_ids, attention_mask = _inputs() + structured = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + tuple_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + + assert isinstance(structured, ModelOutput) + assert isinstance(tuple_output, tuple) + _assert_nested_output_close(tuple_output, structured.to_tuple()) + primary = structured.last_hidden_state if kind == "base" else structured.logits + torch.testing.assert_close(tuple_output[0], primary) + assert structured.hidden_states is not None + assert structured.attentions is not None + assert structured.s_max is not None + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model( + input_ids=input_ids, + attention_mask=attention_mask, + unexpected_cpu_contract=True, + ) + + if kind == "mlm": + labels = input_ids.masked_fill(~attention_mask, -100) + elif kind == "sequence": + labels = torch.tensor([1, 2]) + elif kind == "token": + labels = input_ids.remainder(3).masked_fill(~attention_mask, -100) + else: + labels = None + if labels is not None: + loss_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + loss_tuple = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + _assert_nested_output_close(loss_tuple, loss_output.to_tuple()) + torch.testing.assert_close(loss_tuple[0], loss_output.loss) + torch.testing.assert_close(loss_tuple[1], loss_output.logits) + loss = loss_output.loss + else: + loss = structured.last_hidden_state.square().mean() + assert loss is not None and torch.isfinite(loss) + loss.backward() + assert any(parameter.grad is not None for parameter in model.parameters()) + + model.resize_token_embeddings(19) + assert model.get_input_embeddings().num_embeddings == 19 + output_embeddings = model.get_output_embeddings() + if output_embeddings is not None: + assert output_embeddings.out_features == 19 + + save_dir = tmp_path / model_class.__name__ + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained(save_dir, local_files_only=True).eval() + assert reloaded.get_input_embeddings().num_embeddings == 19 + _assert_exact_state_round_trip(model, reloaded) + _assert_exact_output_round_trip( + model, + reloaded, + input_ids=input_ids, + attention_mask=attention_mask, + ) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size", "kind"), + ( + (dplm_contracts.DPLMModel, dplm_contracts.DPLMConfig, 33, "base"), + (dplm_contracts.DPLMForMaskedLM, dplm_contracts.DPLMConfig, 33, "mlm"), + ( + dplm_contracts.DPLMForSequenceClassification, + dplm_contracts.DPLMConfig, + 33, + "sequence", + ), + ( + dplm_contracts.DPLMForTokenClassification, + dplm_contracts.DPLMConfig, + 33, + "token", + ), + (dplm_contracts.DPLM2Model, dplm_contracts.DPLM2Config, 64, "base"), + (dplm_contracts.DPLM2ForMaskedLM, dplm_contracts.DPLM2Config, 64, "mlm"), + ( + dplm_contracts.DPLM2ForSequenceClassification, + dplm_contracts.DPLM2Config, + 64, + "sequence", + ), + ( + dplm_contracts.DPLM2ForTokenClassification, + dplm_contracts.DPLM2Config, + 64, + "token", + ), + ), +) +def test_dplm_advertised_models_forward_loss_backward_resize_and_reload( + model_class: type, + config_class: type, + vocab_size: int, + kind: str, + tmp_path: Path, +) -> None: + config_values = ( + _dplm2_config_values() + if config_class is dplm_contracts.DPLM2Config + else _dplm_config_values(vocab_size) + ) + config = config_class( + **config_values, + num_labels=3, + return_dict=True, + ) + model = model_class(config).eval() + input_ids = torch.tensor([[0, 6, 7, 2, 1], [0, 8, 2, 1, 1]]) # (b=2, l=5) + attention_mask = input_ids.ne(1) # (b, l) + structured = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + tuple_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + + assert isinstance(structured, ModelOutput) + assert isinstance(tuple_output, tuple) + _assert_nested_output_close(tuple_output, structured.to_tuple()) + primary = structured.last_hidden_state if kind == "base" else structured.logits + torch.testing.assert_close(tuple_output[0], primary) + assert structured.hidden_states is not None + assert structured.attentions is not None + assert structured.s_max is not None + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model( + input_ids=input_ids, + attention_mask=attention_mask, + unexpected_cpu_contract=True, + ) + if kind == "mlm": + labels = input_ids.masked_fill(~attention_mask, -100) + elif kind == "sequence": + labels = torch.tensor([1, 2]) + elif kind == "token": + labels = input_ids.remainder(3).masked_fill(~attention_mask, -100) + else: + labels = None + if labels is None: + loss = structured.last_hidden_state.square().mean() + else: + loss_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + loss_tuple = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + _assert_nested_output_close(loss_tuple, loss_output.to_tuple()) + torch.testing.assert_close(loss_tuple[0], loss_output.loss) + torch.testing.assert_close(loss_tuple[1], loss_output.logits) + loss = loss_output.loss + assert loss is not None and torch.isfinite(loss) + loss.backward() + assert any(parameter.grad is not None for parameter in model.parameters()) + + model.resize_token_embeddings(vocab_size + 3) + assert model.get_input_embeddings().num_embeddings == vocab_size + 3 + output_embeddings = model.get_output_embeddings() + if output_embeddings is not None: + assert output_embeddings.out_features == vocab_size + 3 + save_dir = tmp_path / f"{model_class.__name__}-{vocab_size}" + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained(save_dir, local_files_only=True).eval() + assert reloaded.get_input_embeddings().num_embeddings == vocab_size + 3 + _assert_exact_state_round_trip(model, reloaded) + _assert_exact_output_round_trip( + model, + reloaded, + input_ids=input_ids, + attention_mask=attention_mask, + ) + + +@pytest.mark.parametrize( + ("model_class", "kind"), + ( + (e1_contracts.E1Model, "base"), + (e1_contracts.E1ForMaskedLM, "mlm"), + (e1_contracts.E1ForSequenceClassification, "sequence"), + (e1_contracts.E1ForTokenClassification, "token"), + ), +) +def test_e1_advertised_models_forward_loss_backward_resize_and_reload( + model_class: type, + kind: str, + tmp_path: Path, +) -> None: + config = e1_contracts._tiny_e1_config() + config.num_labels = 3 + model = model_class(config).eval() + batch = e1_contracts._tiny_e1_batch() + structured = model( + **batch, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + tuple_output = model( + **batch, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + + assert isinstance(structured, ModelOutput) + assert isinstance(tuple_output, tuple) + _assert_nested_output_close(tuple_output, structured.to_tuple()) + primary = structured.last_hidden_state if kind == "base" else structured.logits + torch.testing.assert_close(tuple_output[0], primary) + assert structured.hidden_states is not None + assert structured.attentions is not None + assert structured.s_max is not None + with pytest.raises(TypeError, match="unexpected_cpu_contract"): + model(**batch, unexpected_cpu_contract=True) + input_ids = batch["input_ids"] + if kind == "mlm": + labels = input_ids.clone() + elif kind == "sequence": + labels = torch.tensor([1]) + elif kind == "token": + labels = input_ids.remainder(3) + else: + labels = None + if labels is None: + loss = structured.last_hidden_state.square().mean() + else: + loss_output = model( + **batch, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + loss_tuple = model( + **batch, + labels=labels, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=False, + ) + _assert_nested_output_close(loss_tuple, loss_output.to_tuple()) + torch.testing.assert_close(loss_tuple[0], loss_output.loss) + torch.testing.assert_close(loss_tuple[1], loss_output.logits) + loss = loss_output.loss + assert loss is not None and torch.isfinite(loss) + loss.backward() + assert any(parameter.grad is not None for parameter in model.parameters()) + + model.resize_token_embeddings(39) + assert model.get_input_embeddings().num_embeddings == 39 + output_embeddings = model.get_output_embeddings() + if output_embeddings is not None: + assert output_embeddings.out_features == 39 + save_dir = tmp_path / model_class.__name__ + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained(save_dir, local_files_only=True).eval() + assert reloaded.get_input_embeddings().num_embeddings == 39 + _assert_exact_state_round_trip(model, reloaded) + _assert_exact_output_round_trip(model, reloaded, **batch) diff --git a/tests/cpu/test_structure_contracts.py b/tests/cpu/test_structure_contracts.py new file mode 100644 index 0000000..852b13b --- /dev/null +++ b/tests/cpu/test_structure_contracts.py @@ -0,0 +1,432 @@ +"""Portable CPU contracts around public structure and binder helpers.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace + +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model +from fastplms.models.esmfold2.modeling_esmfold2_common import NUM_RES_TYPES +from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, +) +from tests.integration import test_binder_design as binder_contracts +from tests.structure import test_esmfold2_complex_identity as complex_identity +from tests.structure import test_structure_public_helpers as structure_contracts +from tests.structure.test_structure_public_helpers import ( + test_boltz_atom_confidence_is_finite_for_short_uneven_batches_and_multiplicity as _finite, +) +from tests.structure.test_structure_public_helpers import ( + test_boltz_atom_confidence_mapping_preserves_batch_multiplicity_and_atom_order as _order, +) +from tests.unit import test_boltz_checkpoint_io as boltz_checkpoint_contracts +from tests.unit import test_esmfold2_reimplemented_leaves as esmfold2_leaf_contracts +from tests.unit import test_esmfold_api as esmfold_contracts +from tests.unit import test_structure_output_contracts as output_contracts +from tests.unit.test_esmfold2_reimplemented_leaves import ( + test_experimental_top_level_kernel_backend_validates_before_zero_layer_dispatch as _kernel, +) + + +test_boltz_atom_confidence_is_finite_for_short_uneven_batches_and_multiplicity = _finite +test_boltz_atom_confidence_mapping_preserves_batch_multiplicity_and_atom_order = _order +test_experimental_top_level_kernel_backend_validates_before_zero_layer_dispatch = _kernel + +test_boltz_public_helper_is_seeded_and_restores_ambient_rng = ( + structure_contracts.test_boltz_public_helper_is_seeded_and_restores_ambient_rng +) +test_boltz_public_helper_rejects_coerced_seed_types_before_rng_mutation = ( + structure_contracts.test_boltz_public_helper_rejects_coerced_seed_types_before_rng_mutation +) +test_boltz_public_helper_owns_cuda_bf16_autocast_policy = ( + structure_contracts.test_boltz_public_helper_owns_cuda_bf16_autocast_policy +) +test_boltz_public_helper_rejects_non_finite_coordinates = ( + structure_contracts.test_boltz_public_helper_rejects_non_finite_coordinates +) +test_boltz_indexing_matrix_rejects_invalid_public_dimensions = ( + structure_contracts.test_boltz_indexing_matrix_rejects_invalid_public_dimensions +) +test_boltz_piecewise_schedule_validates_and_owns_its_configuration = ( + structure_contracts.test_boltz_piecewise_schedule_validates_and_owns_its_configuration +) +test_boltz_flat_bottom_potential_rejects_invalid_negation_masks = ( + structure_contracts.test_boltz_flat_bottom_potential_rejects_invalid_negation_masks +) +test_boltz_real_features_flow_through_tiny_core_and_structure_loss = ( + structure_contracts.test_boltz_real_features_flow_through_tiny_core_and_structure_loss +) +test_binder_structure_loss_is_finite_and_differentiable = ( + structure_contracts.test_binder_structure_loss_is_finite_and_differentiable +) +test_esmfold_fold_single_uses_linker_masked_mean_plddt = ( + structure_contracts.test_esmfold_fold_single_uses_linker_masked_mean_plddt +) +test_esmfold2_real_features_flow_through_tiny_core_and_tm_loss = ( + structure_contracts.test_esmfold2_real_features_flow_through_tiny_core_and_tm_loss +) +test_pocket_conditioning_is_rejected_instead_of_silently_dropped = ( + esmfold2_leaf_contracts.test_pocket_conditioning_is_rejected_instead_of_silently_dropped +) +test_boltz_save_pretrained_defaults_to_safetensors_and_round_trips = ( + boltz_checkpoint_contracts.test_save_pretrained_defaults_to_safetensors_and_round_trips +) +test_boltz_public_forward_honors_output_controls_backward_and_reload = ( + output_contracts.test_boltz_public_forward_honors_output_controls_backward_and_reload +) +test_esmfold_forward_uses_official_plddt_scale = ( + esmfold_contracts.test_forward_uses_official_plddt_scale +) +test_esmfold_infer_preserves_official_multimer_contract = ( + esmfold_contracts.test_infer_preserves_official_multimer_contract +) +test_fast_esmfold_public_forward_honors_output_controls_and_backward = ( + output_contracts.test_fast_esmfold_public_forward_honors_output_controls_and_backward +) +test_fast_esmfold_output_attentions_uses_masked_per_call_eager_fallback = ( + output_contracts.test_fast_esmfold_output_attentions_uses_masked_per_call_eager_fallback +) +test_fast_esmfold_tiny_model_saves_and_reloads_exact_state = ( + output_contracts.test_fast_esmfold_tiny_model_saves_and_reloads_exact_state +) +test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides = ( + output_contracts.test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides +) +test_molecular_round_trip_preserves_identity_and_repeated_chain_boundaries = ( + complex_identity.test_molecular_round_trip_preserves_identity_and_repeated_chain_boundaries +) +test_backbone_state_dict_does_not_mutate_source_atom_mask = ( + complex_identity.test_backbone_state_dict_does_not_mutate_source_atom_mask +) +test_binder_model_identity_records_selected_kernel_and_mixed_parameter_dtypes = ( + binder_contracts.test_binder_model_identity_records_selected_kernel_and_mixed_parameter_dtypes +) + + +def test_public_binder_workflow_pads_heterogeneous_prepared_atoms_without_truncation() -> None: + """Run real preparation, batching, and forward wiring through an injected tiny core.""" + + from examples import binder_design_fastplms as binder + + class FakeFoldModel: + input_types = SimpleNamespace( + ProteinInput=lambda **values: SimpleNamespace(**values), + StructurePredictionInput=lambda **values: SimpleNamespace(**values), + ) + + def __init__(self) -> None: + self.prepared_sizes: list[int] = [] + self.forward_batches: list[dict[str, torch.Tensor]] = [] + + def prepare_structure_input( + self, + input_data: object, + *, + seed: int | None, + ) -> tuple[dict[str, torch.Tensor], list[str]]: + del input_data, seed + count = (32, 65)[len(self.prepared_sizes)] + marker = float(len(self.prepared_sizes) + 1) + self.prepared_sizes.append(count) + return ( + { + "ref_pos": torch.full((1, count, 3), marker), # (1, n_atom, xyz=3) + "atom_attention_mask": torch.ones( # (1, n_atom) + (1, count), + dtype=torch.bool, + ), + }, + [f"chain-{int(marker)}"], + ) + + def forward(self, **inputs: object) -> dict[str, torch.Tensor]: + ref_pos = inputs["ref_pos"] + atom_attention_mask = inputs["atom_attention_mask"] + res_type_soft = inputs["res_type_soft"] + assert isinstance(ref_pos, torch.Tensor) + assert isinstance(atom_attention_mask, torch.Tensor) + assert isinstance(res_type_soft, torch.Tensor) + self.forward_batches.append( + { + "ref_pos": ref_pos.detach().clone(), + "atom_attention_mask": atom_attention_mask.detach().clone(), + } + ) + batch_size, token_count = res_type_soft.shape[:2] + marker = ref_pos[:, 0, 0] # (b,) + logits = marker[:, None, None, None].expand( # (b, l, l, bins=128) + batch_size, + token_count, + token_count, + 128, + ) + return {"distogram_logits": logits} + + def __call__(self, **inputs: object) -> dict[str, torch.Tensor]: + return self.forward(**inputs) + + model = FakeFoldModel() + design = torch.zeros((2, 2, binder.AA_DIMS), dtype=torch.float32) # (b=2, l=2, aa) + design[0, :, 0] = 1 + design[1, :, 1] = 1 + result = binder.fold_and_get_distogram( + model, + "ACD", + binder.sequence_to_one_hot("ACD", device="cpu"), + design, + seed=7, + ) + + assert model.prepared_sizes == [32, 65] + assert len(model.forward_batches) == 1 + batch = model.forward_batches[0] + assert batch["ref_pos"].shape == (2, 96, 3) + assert batch["atom_attention_mask"].shape == (2, 96) + torch.testing.assert_close(batch["ref_pos"][0, :32], torch.ones((32, 3))) + torch.testing.assert_close(batch["ref_pos"][1, :65], torch.full((65, 3), 2.0)) + assert not batch["ref_pos"][0, 32:].any() + assert not batch["ref_pos"][1, 65:].any() + assert batch["atom_attention_mask"][0, :32].all() + assert batch["atom_attention_mask"][1, :65].all() + assert not batch["atom_attention_mask"][0, 32:].any() + assert not batch["atom_attention_mask"][1, 65:].any() + assert result["inputs"]["ref_pos"].shape == (2, 96, 3) + assert result["chain_info_list"] == [["chain-1"], ["chain-2"]] + assert result["distogram_logits"][:, 0, 0, 0].tolist() == [1.0, 2.0] + + +def test_binder_example_main_wires_explicit_offline_cli_arguments( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from examples import binder_design_fastplms as binder + + observed: dict[str, object] = {} + monkeypatch.setattr( + binder, + "run_local", + lambda arguments: observed.update(vars(arguments)), + ) + output = tmp_path / "binder-output" + + result = binder.main( + [ + "--target-sequence", + "ACD", + "--binder-sequence", + "###", + "--seed", + "9", + "--batch-size", + "2", + "--steps", + "1", + "--output-dir", + str(output), + "--inversion-model", + "Custom/inversion", + "--critic-model", + "Custom/critic", + "--lm-model", + "Custom/lm", + "--model-revision", + f"Custom/inversion={'a' * 40}", + "--model-revision", + f"Custom/critic={'b' * 40}", + "--model-revision", + f"Custom/lm={'c' * 40}", + "--local-files-only", + "--kernel-backend", + "sdpa", + "--compile-model", + "--not-antibody", + ] + ) + + assert result == 0 + assert observed == { + "target_name": None, + "target_sequence": "ACD", + "binder_name": None, + "binder_sequence": "###", + "seed": 9, + "batch_size": 2, + "steps": 1, + "output_dir": str(output), + "inversion_model_names": ["Custom/inversion"], + "critic_model_names": ["Custom/critic"], + "lm_model": "Custom/lm", + "model_revisions": { + "Custom/inversion": "a" * 40, + "Custom/critic": "b" * 40, + "Custom/lm": "c" * 40, + }, + "local_files_only": True, + "kernel_backend": "sdpa", + "compile_model": True, + "is_antibody": False, + } + + +def test_binder_documentation_exposes_pinned_and_offline_loading_contract() -> None: + documentation = (Path(__file__).resolve().parents[2] / "docs" / "binder_design.md").read_text( + encoding="utf-8" + ) + + for required_text in ( + "src/fastplms/models.toml", + "--inversion-model", + "--critic-model", + "--lm-model", + "--model-revision", + "--local-files-only", + "HF_HUB_OFFLINE=1", + "TRANSFORMERS_OFFLINE=1", + "fastplms_weights_revision", + "fastplms_runtime_revision", + ): + assert required_text in documentation + + +def test_esmfold2_public_esmc_loaders_propagate_instance_offline_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastplms.models.esmfold2 import modeling_esmfold2 as release_module + from fastplms.models.esmfold2 import modeling_esmfold2_experimental as experimental_module + + observed: list[tuple[object, str, dict[str, object]]] = [] + + def fake_install(model, source: str, **kwargs: object) -> None: + observed.append((model, source, kwargs)) + + monkeypatch.setattr(release_module, "_install_esmc_backbone", fake_install) + monkeypatch.setattr(experimental_module, "_install_esmc_backbone", fake_install) + release_stub = SimpleNamespace() + experimental_stub = SimpleNamespace() + + ESMFold2Model.load_esmc( + release_stub, + "Synthyra/ESMplusplus_6B", + precision="bf16", + device="cpu", + local_files_only=True, + ) + ESMFold2ExperimentalModel.load_esmc( + experimental_stub, + "Synthyra/ESMplusplus_6B", + precision="bf16", + device="cpu", + local_files_only=True, + ) + + assert observed == [ + ( + release_stub, + "Synthyra/ESMplusplus_6B", + { + "precision": "bf16", + "device": "cpu", + "local_files_only": True, + }, + ), + ( + experimental_stub, + "Synthyra/ESMplusplus_6B", + { + "precision": "bf16", + "device": "cpu", + "local_files_only": True, + }, + ), + ] + + +def _tiny_esmfold2_config(model_type: str) -> ESMFold2Config: + atom_token_width = 8 + input_feature_width = atom_token_width // 2 + 2 * NUM_RES_TYPES + 1 + return ESMFold2Config( + type=model_type, + d_single=8, + d_pair=8, + num_loops=0, + num_diffusion_samples=1, + lm_d_model=8, + lm_num_layers=1, + inputs={ + "d_inputs": input_feature_width, + "atom_encoder": { + "d_atom": 8, + "d_token": atom_token_width, + "n_blocks": 0, + "n_heads": 2, + "swa_window_size": 32, + "expansion_ratio": 2, + "n_spatial_rope_pairs_per_axis": 1, + "n_uid_rope_pairs": 1, + }, + }, + folding_trunk={"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + structure_head={ + "diffusion_module": { + "c_atom": 8, + "c_token": 8, + "c_z": 8, + "c_s_inputs": input_feature_width, + "fourier_dim": 8, + "atom_num_blocks": 0, + "atom_num_heads": 2, + "token_num_blocks": 0, + "token_num_heads": 2, + "transition_multiplier": 2, + }, + "distogram_bins": 8, + "inference_num_steps": 1, + }, + confidence_head={ + "enabled": False, + "folding_trunk": {"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + "num_plddt_bins": 4, + "num_pde_bins": 4, + "num_pae_bins": 4, + "distogram_bins": 8, + }, + msa_encoder={"enabled": False}, + lm_encoder={"enabled": False, "n_layers": 0}, + parcae={"enabled": True, "min_steps": 1, "max_steps": 1, "coda_n_layers": 0}, + ) + + +@pytest.mark.parametrize( + ("model_class", "model_type"), + ( + (ESMFold2Model, "release"), + (ESMFold2ExperimentalModel, "experimental"), + ), +) +def test_esmfold2_advertised_models_tiny_init_backward_and_save_reload( + model_class: type[ESMFold2Model] | type[ESMFold2ExperimentalModel], + model_type: str, + tmp_path: Path, +) -> None: + model = model_class(_tiny_esmfold2_config(model_type)) + pair_state = torch.randn(1, 2, 2, 8, requires_grad=True) # (b=1, l=2, l, d_pair=8) + logits = model.distogram_head( # (b, l, l, distogram_bins=8) + pair_state + pair_state.transpose(1, 2) + ) + logits.square().mean().backward() + + assert torch.isfinite(logits).all() + assert pair_state.grad is not None and torch.isfinite(pair_state.grad).all() + save_dir = tmp_path / model_class.__name__ + model.save_pretrained(save_dir, safe_serialization=True) + reloaded = model_class.from_pretrained( + save_dir, + local_files_only=True, + load_esmc=False, + ) + assert set(reloaded.state_dict()) == set(model.state_dict()) + for name, tensor in model.state_dict().items(): + torch.testing.assert_close(reloaded.state_dict()[name], tensor, rtol=0.0, atol=0.0) diff --git a/tests/cpu/test_ttt_contracts.py b/tests/cpu/test_ttt_contracts.py new file mode 100644 index 0000000..f297cac --- /dev/null +++ b/tests/cpu/test_ttt_contracts.py @@ -0,0 +1,333 @@ +"""Mandatory deterministic test-time-training contracts.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace + +from fastplms.models.ankh.modeling_ankh import FastAnkhForMaskedLMExtension +from fastplms.models.dplm.modeling_dplm import DPLMConfig, DPLMForMaskedLM +from fastplms.models.dplm2.modeling_dplm2 import DPLM2Config, DPLM2ForMaskedLM +from fastplms.models.e1.modeling_e1 import E1ForMaskedLM +from fastplms.models.esm2.modeling_fastesm import FastEsmForMaskedLM +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, +) +from fastplms.models.ttt import LoraInjectedLinear +from tests.integration import test_ttt as contracts +from tests.unit.test_ankh_cpu_contract import _config as _ankh_config +from tests.unit.test_e1_cache_contract import _tiny_e1_batch, _tiny_e1_config + + +test_ttt_first_call_mapping_preserves_explicit_target_override = ( + contracts.test_ttt_first_call_mapping_preserves_explicit_target_override +) +test_ttt_direct_init_mapping_preserves_family_target_class = ( + contracts.test_ttt_direct_init_mapping_preserves_family_target_class +) +test_ttt_config_rejects_invalid_optimizer_and_target_contracts = ( + contracts.test_ttt_config_rejects_invalid_optimizer_and_target_contracts +) +test_ttt_first_call_mapping_preserves_family_target_class = ( + contracts.test_ttt_first_call_mapping_preserves_family_target_class +) +test_seed_and_initial_state_reset_reproduce_losses_and_updates = ( + contracts.test_seed_and_initial_state_reset_reproduce_losses_and_updates +) +test_ttt_adapter_initialization_is_seeded_and_preserves_ambient_rng = ( + contracts.test_ttt_adapter_initialization_is_seeded_and_preserves_ambient_rng +) +test_ttt_generic_replacements_exclude_reserved_vocabulary_ids = ( + contracts.test_ttt_generic_replacements_exclude_reserved_vocabulary_ids +) +test_ttt_lora_injection_is_lazy_and_backbone_scoped = ( + contracts.test_ttt_lora_injection_is_lazy_and_backbone_scoped +) +test_ttt_only_lora_params_change_and_reset_restores_adapter = ( + contracts.test_ttt_only_lora_params_change_and_reset_restores_adapter +) +test_ttt_rejects_all_ignored_inputs_before_adapter_injection = ( + contracts.test_ttt_rejects_all_ignored_inputs_before_adapter_injection +) +test_ttt_rejects_dplm2_structure_tokens_before_adapter_injection = ( + contracts.test_ttt_rejects_dplm2_structure_tokens_before_adapter_injection +) +test_ttt_save_pretrained_round_trip_preserves_adapter_and_reset_state = ( + contracts.test_ttt_save_pretrained_round_trip_preserves_adapter_and_reset_state +) +test_ttt_uneven_batch_samples_only_rows_with_residue_targets = ( + contracts.test_ttt_uneven_batch_samples_only_rows_with_residue_targets +) + + +class _ProteinTokenizer: + pad_token_id = 1 + + def __init__( + self, + mask_token_id: int = 3, + additional_special_ids: tuple[int, ...] = (), + ) -> None: + self.mask_token_id = mask_token_id + self.all_special_ids = sorted( + {0, 1, 2, mask_token_id, *additional_special_ids} + ) + available_ids = [ + token_id + for token_id in range(3, 32) + if token_id not in self.all_special_ids + ][:20] + if len(available_ids) != 20: + raise ValueError("Tiny tokenizer cannot allocate the canonical amino-acid alphabet") + self.vocab = { + amino_acid: token_id + for amino_acid, token_id in zip( + "ACDEFGHIKLMNPQRSTVWY", + available_ids, + strict=True, + ) + } + + def get_vocab(self) -> dict[str, int]: + return dict(self.vocab) + + def convert_tokens_to_ids(self, token: str) -> int: + return self.vocab.get(token, 3) + + +class _DPLM2ProteinTokenizer(_ProteinTokenizer): + aa_mask_token = "" + struct_cls_token = "" + + def __init__(self) -> None: + super().__init__( + mask_token_id=32, + additional_special_ids=(3, 33, 34, 35, 36), + ) + self._token_to_id = { + **self.vocab, + self.aa_mask_token: 32, + self.struct_cls_token: 33, + } + + +def test_dplm2_ttt_replacements_exclude_ambiguous_and_reserved_tokens() -> None: + tokenizer = _DPLM2ProteinTokenizer() + tokenizer._token_to_id.update( + {"X": 24, "B": 25, "U": 26, "Z": 27, "O": 28, "-": 29} + ) + replacements = DPLM2ForMaskedLM._ttt_replacement_tokens( + SimpleNamespace(tokenizer=tokenizer), + torch.tensor([[0, 4, 2]], dtype=torch.long), + ) + + assert replacements.tolist() == [ + tokenizer._token_to_id[residue] + for residue in "ACDEFGHIKLMNPQRSTVWY" + ] + + +_EXPECTED_ADAPTERS = { + "esm2": ( + "esm.encoder.layer.0.attention.self.query", + "esm.encoder.layer.0.attention.self.key", + "esm.encoder.layer.0.attention.self.value", + "esm.encoder.layer.0.attention.output.dense", + ), + "esm_plusplus": ( + "transformer.blocks.0.attn.layernorm_qkv.1", + "transformer.blocks.0.attn.out_proj", + ), + "esm3": ( + "esm3.transformer.blocks.0.attn.layernorm_qkv.1", + "esm3.transformer.blocks.0.attn.out_proj", + ), + "ankh": ( + "encoder.block.0.layer.0.SelfAttention.q", + "encoder.block.0.layer.0.SelfAttention.k", + "encoder.block.0.layer.0.SelfAttention.v", + "encoder.block.0.layer.0.SelfAttention.o", + ), + "dplm": ( + "esm.encoder.layer.0.attention.self.query", + "esm.encoder.layer.0.attention.self.key", + "esm.encoder.layer.0.attention.self.value", + "esm.encoder.layer.0.attention.output.dense", + ), + "dplm2": ( + "esm.encoder.layer.0.attention.self.query", + "esm.encoder.layer.0.attention.self.key", + "esm.encoder.layer.0.attention.self.value", + "esm.encoder.layer.0.attention.output.dense", + ), + "e1": ( + "model.layers.0.norm_attn_norm.self_attn.q_proj", + "model.layers.0.norm_attn_norm.self_attn.k_proj", + "model.layers.0.norm_attn_norm.self_attn.v_proj", + "model.layers.0.norm_attn_norm.self_attn.o_proj", + ), +} + +_EXPECTED_TARGET_CLASS = { + "esm2": "EsmAttention", + "esm_plusplus": "MultiHeadAttention", + "esm3": "MultiHeadAttention", + "ankh": "AnkhSelfAttention", + "dplm": "ModifiedEsmAttention", + "dplm2": "ModifiedEsmAttention", + "e1": "Attention", +} + + +def _family_model_and_inputs( + family: str, +) -> tuple[ + FastEsmForMaskedLM + | ESMplusplusForMaskedLM + | FastESM3Model + | FastAnkhForMaskedLMExtension + | DPLMForMaskedLM + | DPLM2ForMaskedLM + | E1ForMaskedLM, + dict[str, torch.Tensor], +]: + if family == "esm2": + from tests.cpu.test_sequence_autoclass_contracts import _esm2_config + + config = _esm2_config() + config.vocab_size = 32 + tokenizer = _ProteinTokenizer(mask_token_id=config.mask_token_id) + model = FastEsmForMaskedLM(config) + model.tokenizer = tokenizer + residues = [tokenizer.convert_tokens_to_ids(value) for value in "AC"] + return model, {"input_ids": torch.tensor([[0, *residues, 2, 1]])} + if family == "esm_plusplus": + tokenizer = _ProteinTokenizer(mask_token_id=3) + model = ESMplusplusForMaskedLM( + ESMplusplusConfig( + vocab_size=32, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + dropout=0.0, + pad_token_id=1, + mask_token_id=3, + attn_backend="eager", + ) + ) + model.tokenizer = tokenizer + residues = [tokenizer.convert_tokens_to_ids(value) for value in "AC"] + return model, {"input_ids": torch.tensor([[0, *residues, 2, 1]])} + if family == "esm3": + model = FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + ) + return model, {"input_ids": model.encode("AC")["input_ids"]} + if family == "ankh": + tokenizer = _ProteinTokenizer(mask_token_id=3) + model = FastAnkhForMaskedLMExtension( + _ankh_config(vocab_size=32, num_layers=1, num_decoder_layers=1) + ) + model.tokenizer = tokenizer + residues = [tokenizer.convert_tokens_to_ids(value) for value in "AC"] + return model, {"input_ids": torch.tensor([[*residues, 1, 0]])} + if family == "dplm": + from tests.cpu.test_sequence_autoclass_contracts import _dplm_config_values + + config = DPLMConfig(**_dplm_config_values(33)) + tokenizer = _ProteinTokenizer(mask_token_id=config.mask_token_id) + model = DPLMForMaskedLM(config) + model.tokenizer = tokenizer + residues = [tokenizer.convert_tokens_to_ids(value) for value in "AC"] + return model, {"input_ids": torch.tensor([[0, *residues, 2, 1]])} + if family == "dplm2": + from tests.cpu.test_sequence_autoclass_contracts import _dplm2_config_values + + model = DPLM2ForMaskedLM(DPLM2Config(**_dplm2_config_values())) + model.tokenizer = _DPLM2ProteinTokenizer() + return model, {"input_ids": torch.tensor([[0, 4, 5, 2, 1]])} + if family == "e1": + return E1ForMaskedLM(_tiny_e1_config()), _tiny_e1_batch() + raise AssertionError(f"Unhandled TTT family: {family}") + + +@pytest.mark.parametrize("family", tuple(_EXPECTED_ADAPTERS)) +def test_each_sequence_family_first_call_ttt_is_scoped_and_reloadable( + family: str, + tmp_path: Path, +) -> None: + model, model_inputs = _family_model_and_inputs(family) + original_parameters = { # parameter identity -> checkpoint-shaped tensor + id(parameter): parameter.detach().clone() + for parameter in model.parameters() + } + metrics = model.ttt( + **model_inputs, + ttt_config={ + "steps": 1, + "ags": 1, + "batch_size": 1, + "mask_ratio": 1.0, + "bert_leave_prob": 0.0, + "bert_replace_prob": 0.0, + "lora_rank": 2, + "lora_alpha": 1.0, + "seed": 7, + }, + ) + + adapters = tuple( + name + for name, module in model.named_modules() + if isinstance(module, LoraInjectedLinear) + ) + assert model.ttt_config.lora_target_replace_module == _EXPECTED_TARGET_CLASS[family] + assert adapters == _EXPECTED_ADAPTERS[family] + assert len(metrics["losses"]) == 1 + assert torch.isfinite(torch.tensor(metrics["losses"])).all() + for parameter in model.parameters(): + if id(parameter) in original_parameters: + torch.testing.assert_close( + parameter.detach(), + original_parameters[id(parameter)], + rtol=0.0, + atol=0.0, + ) + assert any( + name.endswith("lora_up.weight") + and int(torch.count_nonzero(parameter).item()) > 0 + for name, parameter in model.named_parameters() + ) + + save_directory = tmp_path / family + model.save_pretrained(save_directory, safe_serialization=True) + reloaded = type(model).from_pretrained(save_directory, local_files_only=True) + reloaded_adapters = tuple( + name + for name, module in reloaded.named_modules() + if isinstance(module, LoraInjectedLinear) + ) + assert reloaded_adapters == adapters + source_state = { # LoRA parameter name -> checkpoint-shaped tensor + name: tensor + for name, tensor in model.state_dict().items() + if ".lora_" in name + } + reloaded_state = { # LoRA parameter name -> checkpoint-shaped tensor + name: tensor + for name, tensor in reloaded.state_dict().items() + if ".lora_" in name + } + assert set(reloaded_state) == set(source_state) + for name, tensor in source_state.items(): + torch.testing.assert_close(reloaded_state[name], tensor, rtol=0.0, atol=0.0) diff --git a/tests/goldens/ankh2_large.json b/tests/goldens/ankh2_large.json new file mode 100644 index 0000000..485d91d --- /dev/null +++ b/tests/goldens/ankh2_large.json @@ -0,0 +1,88 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:9286bed4ecbc4f7113024919d16ec9719b0c0748", + "generation_config.json": "git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "pytorch_model.bin": "sha256:2df583f28f111276ee22a7b76007f4297e9a69766d60bccd9c8d7169c06ac606", + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25" + }, + "repo_id": "ElnaggarLab/ankh2-ext2", + "revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"ankh\":\"1.10.0\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"biopython\":\"1.80\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"datasets\":\"2.20.0\",\"dill\":\"0.3.8\",\"exceptiongroup\":\"1.3.1\",\"filelock\":\"3.29.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"networkx\":\"3.4.2\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"2.3.3\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"propcache\":\"0.5.2\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"sentencepiece\":\"0.1.99\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.13.3\",\"tomli\":\"2.4.1\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"transformers\":\"4.25.1\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"2.7.0\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.13.0+cu130" + }, + "fingerprint": "8b5a3c1563644f6c019f497de8a3d707c8f8711499e601e310810f1a599ae5d7" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "ankh2_large" + ], + "input_fingerprint": "8b9b29613bd3a3455c5d5c47eb884cd698a52d8621ae1906571a4a4c97a80ee1", + "model_id": "ankh2_large", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "d654e6701d0c8192dc96c76e4a10b621c193fe085d07db636ecf58a52f62ccbd", + "native/metadata.json": "4cf9ad9fe3064e24334d7966346f65cfcd3935cb2166769b6a87631ea843a7bb" + }, + "sources": [ + { + "id": "ankh", + "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e", + "url": "https://github.com/agemagician/Ankh.git" + } + ], + "tensor_file": { + "path": "ankh2_large.safetensors", + "sha256": "25fe1569f55c635fab8fa49c1d62a889a35a2a738bad921f5764a85b58fd4b5d" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "37b3e21090fc901d75954c7a7d000a6359ab4bbc9dd0bf8354c5634f1bc70a45", + "shape": [ + 3, + 62 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "0c1af56968fd959dbdd7ff32df0e44986766d188b1dc4fbcb0028558e9094b95", + "shape": [ + 3, + 62 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "c4eeec90380d66d51664c6375677e8bf87442db8a859c16ef39b6cd4520bec76", + "shape": [ + 3, + 62, + 1536 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "5cdc874e7bff8e7a0a5752d0a4e9c46e6575f188f6ced645f33930d174a90f78", + "shape": [ + 3, + 62 + ] + } + } +} diff --git a/tests/goldens/ankh2_large.safetensors b/tests/goldens/ankh2_large.safetensors new file mode 100644 index 0000000..e61e470 Binary files /dev/null and b/tests/goldens/ankh2_large.safetensors differ diff --git a/tests/goldens/ankh3_large.json b/tests/goldens/ankh3_large.json new file mode 100644 index 0000000..2d7c06a --- /dev/null +++ b/tests/goldens/ankh3_large.json @@ -0,0 +1,89 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:f5278f77d158cdd8a173df888e3ed365e84a80a3", + "generation_config.json": "git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2", + "pytorch_model.bin": "sha256:26321a345e07a25b21c6c41b651c4db91b420892e52c0dcbc55bd7a8f510f95b", + "special_tokens_map.json": "git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model": "sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json": "git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json": "git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21" + }, + "repo_id": "ElnaggarLab/ankh3-large", + "revision": "2be091622e8a393f0ef21735070084123c874b6e" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"ankh\":\"1.10.0\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"biopython\":\"1.80\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"datasets\":\"2.20.0\",\"dill\":\"0.3.8\",\"exceptiongroup\":\"1.3.1\",\"filelock\":\"3.29.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"networkx\":\"3.4.2\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"2.3.3\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"propcache\":\"0.5.2\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"sentencepiece\":\"0.1.99\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.13.3\",\"tomli\":\"2.4.1\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"transformers\":\"4.25.1\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"2.7.0\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.13.0+cu130" + }, + "fingerprint": "8b5a3c1563644f6c019f497de8a3d707c8f8711499e601e310810f1a599ae5d7" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "ankh3_large" + ], + "input_fingerprint": "b73a37c76776436e87d9bd8fcc17ae05c35abbc80a7ae7f040c988076d64c84f", + "model_id": "ankh3_large", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "aad644cae1b174d673e101106bfb80cd9ef001741ab34778af69d52d32e47d65", + "native/metadata.json": "2c02d3d68eab4c42966f59dc7f04ad2a0009834792791a74bb0115b68fbdf2a4" + }, + "sources": [ + { + "id": "ankh", + "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e", + "url": "https://github.com/agemagician/Ankh.git" + } + ], + "tensor_file": { + "path": "ankh3_large.safetensors", + "sha256": "e5c494ac418e0a2fe7bdad1376676d48960d58ec9e044d19bfffccb8c3288513" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "37b3e21090fc901d75954c7a7d000a6359ab4bbc9dd0bf8354c5634f1bc70a45", + "shape": [ + 3, + 62 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "1cdcea1bb937d17b0c44561161af3fff45b6ec5b267746426571460a1a27d8d0", + "shape": [ + 3, + 62 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "7d86ab8cf5d96de8c426bdea5c3d0272dd69ad927bb745c903be3bd09f6b2eab", + "shape": [ + 3, + 62, + 1536 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "5cdc874e7bff8e7a0a5752d0a4e9c46e6575f188f6ced645f33930d174a90f78", + "shape": [ + 3, + 62 + ] + } + } +} diff --git a/tests/goldens/ankh3_large.safetensors b/tests/goldens/ankh3_large.safetensors new file mode 100644 index 0000000..a73cf61 Binary files /dev/null and b/tests/goldens/ankh3_large.safetensors differ diff --git a/tests/goldens/ankh3_xl.json b/tests/goldens/ankh3_xl.json new file mode 100644 index 0000000..836a41a --- /dev/null +++ b/tests/goldens/ankh3_xl.json @@ -0,0 +1,91 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:f8997040e8913df75fd2eebe71a2a8eb750ed0d0", + "generation_config.json": "git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "pytorch_model-00001-of-00003.bin": "sha256:2c9793cbee16697cd4149debe07d3a27143e280f6e970fa46042aae820fea981", + "pytorch_model-00002-of-00003.bin": "sha256:31c5a860e414513c829ae52affb0970d7cef2c0545df2d6e1338b6806ab7174b", + "pytorch_model-00003-of-00003.bin": "sha256:055a853bdd3623db95a637935aa299427e837cd8ea69fc04708b0262508bec75", + "special_tokens_map.json": "git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model": "sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json": "git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json": "git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21" + }, + "repo_id": "ElnaggarLab/ankh3-xl", + "revision": "e00113df5c95ef71df7ea3f5a73d56bd00e473a4" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"ankh\":\"1.10.0\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"biopython\":\"1.80\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"datasets\":\"2.20.0\",\"dill\":\"0.3.8\",\"exceptiongroup\":\"1.3.1\",\"filelock\":\"3.29.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"networkx\":\"3.4.2\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"2.3.3\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"propcache\":\"0.5.2\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"sentencepiece\":\"0.1.99\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.13.3\",\"tomli\":\"2.4.1\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"transformers\":\"4.25.1\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"2.7.0\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.13.0+cu130" + }, + "fingerprint": "8b5a3c1563644f6c019f497de8a3d707c8f8711499e601e310810f1a599ae5d7" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "ankh3_xl" + ], + "input_fingerprint": "b73a37c76776436e87d9bd8fcc17ae05c35abbc80a7ae7f040c988076d64c84f", + "model_id": "ankh3_xl", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "3307d0435454226f44000413a575bb685e201680bf574ba47f290c8352233b3a", + "native/metadata.json": "089d03fceb1e71caabfacf47bd33e72edb3be6264d17e57f84c5f38b95205cef" + }, + "sources": [ + { + "id": "ankh", + "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e", + "url": "https://github.com/agemagician/Ankh.git" + } + ], + "tensor_file": { + "path": "ankh3_xl.safetensors", + "sha256": "72d34567d0228cb6f1ee701c578ed4039fead4346e3f161a52e0e74df28dc8ae" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "37b3e21090fc901d75954c7a7d000a6359ab4bbc9dd0bf8354c5634f1bc70a45", + "shape": [ + 3, + 62 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "1cdcea1bb937d17b0c44561161af3fff45b6ec5b267746426571460a1a27d8d0", + "shape": [ + 3, + 62 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "53fd63cc76e94c690b358d18cc69a41be2ce8895009687183d79277d998d6a17", + "shape": [ + 3, + 62, + 2560 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "5cdc874e7bff8e7a0a5752d0a4e9c46e6575f188f6ced645f33930d174a90f78", + "shape": [ + 3, + 62 + ] + } + } +} diff --git a/tests/goldens/ankh3_xl.safetensors b/tests/goldens/ankh3_xl.safetensors new file mode 100644 index 0000000..b23b2d1 Binary files /dev/null and b/tests/goldens/ankh3_xl.safetensors differ diff --git a/tests/goldens/ankh_base.json b/tests/goldens/ankh_base.json new file mode 100644 index 0000000..a7ed842 --- /dev/null +++ b/tests/goldens/ankh_base.json @@ -0,0 +1,87 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:abd44a36b5469e9a7cb019e4059b5ac1392d8422", + "pytorch_model.bin": "sha256:9b2a886374f0ff4a893f4e7a989deed76bb2458c8998bd5202ea8e97d92ddcc3", + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0" + }, + "repo_id": "ElnaggarLab/ankh-base", + "revision": "d99cb6b966530dfc2ae96bc69d9255c2a07308b0" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"ankh\":\"1.10.0\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"biopython\":\"1.80\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"datasets\":\"2.20.0\",\"dill\":\"0.3.8\",\"exceptiongroup\":\"1.3.1\",\"filelock\":\"3.29.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"networkx\":\"3.4.2\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"2.3.3\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"propcache\":\"0.5.2\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"sentencepiece\":\"0.1.99\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.13.3\",\"tomli\":\"2.4.1\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"transformers\":\"4.25.1\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"2.7.0\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.13.0+cu130" + }, + "fingerprint": "8b5a3c1563644f6c019f497de8a3d707c8f8711499e601e310810f1a599ae5d7" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "ankh_base" + ], + "input_fingerprint": "8b9b29613bd3a3455c5d5c47eb884cd698a52d8621ae1906571a4a4c97a80ee1", + "model_id": "ankh_base", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "22e13c549d75b3de39021a31672c1db64a3b645e0c0346db64fbc5d7ad71f5ba", + "native/metadata.json": "8635555caa1ed3d8d57c41f35184142af565ca2a2dff51584dc2897a717c5f4b" + }, + "sources": [ + { + "id": "ankh", + "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e", + "url": "https://github.com/agemagician/Ankh.git" + } + ], + "tensor_file": { + "path": "ankh_base.safetensors", + "sha256": "f0e78aa15d11749e0c64ff57f9e88c51cec6538a0adf8951f839df70cc708b65" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "37b3e21090fc901d75954c7a7d000a6359ab4bbc9dd0bf8354c5634f1bc70a45", + "shape": [ + 3, + 62 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "0c1af56968fd959dbdd7ff32df0e44986766d188b1dc4fbcb0028558e9094b95", + "shape": [ + 3, + 62 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "643463daebb89c00b0469b4e7b02d08ec0bc9c4561f2be4c1c5e9441ff82177a", + "shape": [ + 3, + 62, + 768 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "5cdc874e7bff8e7a0a5752d0a4e9c46e6575f188f6ced645f33930d174a90f78", + "shape": [ + 3, + 62 + ] + } + } +} diff --git a/tests/goldens/ankh_base.safetensors b/tests/goldens/ankh_base.safetensors new file mode 100644 index 0000000..f116b3f Binary files /dev/null and b/tests/goldens/ankh_base.safetensors differ diff --git a/tests/goldens/ankh_large.json b/tests/goldens/ankh_large.json new file mode 100644 index 0000000..6c32f84 --- /dev/null +++ b/tests/goldens/ankh_large.json @@ -0,0 +1,87 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:1abf33e52ee3d6be67d780ec57d32ac2b27b5306", + "pytorch_model.bin": "sha256:517b6e8b279dedcb477af240b35c46bd6eb3307723eb281e60d4b2c8a87b889b", + "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json": "git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09" + }, + "repo_id": "ElnaggarLab/ankh-large", + "revision": "74b371dbfa3ee0a05d32ae74df0c2e0b82d6b9a6" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"ankh\":\"1.10.0\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"biopython\":\"1.80\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"datasets\":\"2.20.0\",\"dill\":\"0.3.8\",\"exceptiongroup\":\"1.3.1\",\"filelock\":\"3.29.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"networkx\":\"3.4.2\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"2.3.3\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"propcache\":\"0.5.2\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"sentencepiece\":\"0.1.99\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.13.3\",\"tomli\":\"2.4.1\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"transformers\":\"4.25.1\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"2.7.0\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.13.0+cu130" + }, + "fingerprint": "8b5a3c1563644f6c019f497de8a3d707c8f8711499e601e310810f1a599ae5d7" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "ankh_large" + ], + "input_fingerprint": "8b9b29613bd3a3455c5d5c47eb884cd698a52d8621ae1906571a4a4c97a80ee1", + "model_id": "ankh_large", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "fdba8593cfb54ee4532e8f9fef7d5df110d66d245797ba37bb5c24049d918c2e", + "native/metadata.json": "7df438c5257521541034e7b8e8190c0c4f4cd31d6fb2b00c4f2c502296f68244" + }, + "sources": [ + { + "id": "ankh", + "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e", + "url": "https://github.com/agemagician/Ankh.git" + } + ], + "tensor_file": { + "path": "ankh_large.safetensors", + "sha256": "3fb8d3ac27716d15a9ea92aeef6acf2b977bcc887d9b535000539e523673459b" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "37b3e21090fc901d75954c7a7d000a6359ab4bbc9dd0bf8354c5634f1bc70a45", + "shape": [ + 3, + 62 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "0c1af56968fd959dbdd7ff32df0e44986766d188b1dc4fbcb0028558e9094b95", + "shape": [ + 3, + 62 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "b8f603b15932d06c2aad9baabf6e7811ecf937589078a078ad79745414b21733", + "shape": [ + 3, + 62, + 1536 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "5cdc874e7bff8e7a0a5752d0a4e9c46e6575f188f6ced645f33930d174a90f78", + "shape": [ + 3, + 62 + ] + } + } +} diff --git a/tests/goldens/ankh_large.safetensors b/tests/goldens/ankh_large.safetensors new file mode 100644 index 0000000..1292923 Binary files /dev/null and b/tests/goldens/ankh_large.safetensors differ diff --git a/tests/goldens/dplm2_150m.json b/tests/goldens/dplm2_150m.json new file mode 100644 index 0000000..f05e615 --- /dev/null +++ b/tests/goldens/dplm2_150m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:20f1e55c64fdc4d1d30f7b1df64b6167fa23dc7c", + "pytorch_model.bin": "sha256:be7f5cf9e421f59fcc437e63ce1c7391099a314a4e9a4f10b8688785fa581238", + "special_tokens_map.json": "git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json": "git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt": "git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37" + }, + "repo_id": "airkingbd/dplm2_150m", + "revision": "3451d984d06497f835ed49634bd68c9dfb54d730" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm2_150m" + ], + "input_fingerprint": "a7b239a22465f83adf8de7303bb90b0b85ddccfee476137f54331da231059910", + "model_id": "dplm2_150m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "1c7bf879a98cf293d91d7671ce86f7d2a440515891fa2eef326152995f3a5a03", + "native/metadata.json": "68278e51d81ed79c0ccf05a22848102489a00e04376c10b51b6542ddc3828717" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm2_150m.safetensors", + "sha256": "17fc26600938ba5364b8ecb96750786d33e9f92bcd4ea4df3e12a389340748eb" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "6c19ab143b0587e6766355da1f0527bf7df033f9c5de68afcd60d405a1ce995a", + "shape": [ + 3, + 126 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "11d09bddaf02ce351b86b1c5ce5f05b2e89420605014b337e2c49061d9e57f3d", + "shape": [ + 3, + 126 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "812f85535e9a682122a407c935f8097b08e03ae9bb508f29417a476cdb7047f9", + "shape": [ + 3, + 126, + 640 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "42eacf345357e9ea31bdb6bca3b99f2df6717a4acef6f16400d5fb864c526a66", + "shape": [ + 3, + 126, + 8229 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "ea2dbe6bba30495918cc14a223a43fb6efee780a21cbce1d503cb6200e7ef807", + "shape": [ + 3, + 126 + ] + } + } +} diff --git a/tests/goldens/dplm2_150m.safetensors b/tests/goldens/dplm2_150m.safetensors new file mode 100644 index 0000000..555685e Binary files /dev/null and b/tests/goldens/dplm2_150m.safetensors differ diff --git a/tests/goldens/dplm2_3b.json b/tests/goldens/dplm2_3b.json new file mode 100644 index 0000000..1bfe4d2 --- /dev/null +++ b/tests/goldens/dplm2_3b.json @@ -0,0 +1,108 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:22d51ce44cd6da8d819e0d00566987bb51d74753", + "pytorch_model-00001-of-00004.bin": "sha256:d8c641eae6bf891581ec64d543169891b093e296f5679ac75c695bcf596b4211", + "pytorch_model-00002-of-00004.bin": "sha256:6478ad86ec5fef3d1d26580493af2d8666009d3ff884f3f88548080c8bbf94b5", + "pytorch_model-00003-of-00004.bin": "sha256:dde8f88dac4a6355488c2fb433ee12cd69f1169950566624fba43684d4d99dc6", + "pytorch_model-00004-of-00004.bin": "sha256:17ec0145152bc10e4dd3b4c2edff337979f6b99ee7c7bfd6cf4e6dbd7262d079", + "special_tokens_map.json": "git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json": "git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt": "git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37" + }, + "repo_id": "airkingbd/dplm2_3b", + "revision": "9e77567926f98d1b997ea9131a8eeb035b9bf827" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm2_3b" + ], + "input_fingerprint": "a7b239a22465f83adf8de7303bb90b0b85ddccfee476137f54331da231059910", + "limitations": [ + { + "capability": "generation", + "exception_type": "TypeError", + "public_method": "EsmForDPLM.generate", + "reason": "The checkpoint-selected EsmForDPLM sampler uses tokenizer.cls_token_id as bos_id, but the pinned DPLM2 tokenizer defines no cls_token_id.", + "status": "official_unavailable" + } + ], + "model_id": "dplm2_3b", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "2b0868d8e007ac29aa0674528c5c5de5392151d7c8482d86c4e43946334c24d9", + "native/metadata.json": "feb45f4491402a908dda1f8a1e7c0b35936a65988bb2af6740b93c39f185832d" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm2_3b.safetensors", + "sha256": "838b11824d08f83bcb0c0b3268e579f3a87dbfb965370cfe5c3f8793b96b1964" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "6c19ab143b0587e6766355da1f0527bf7df033f9c5de68afcd60d405a1ce995a", + "shape": [ + 3, + 126 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "11d09bddaf02ce351b86b1c5ce5f05b2e89420605014b337e2c49061d9e57f3d", + "shape": [ + 3, + 126 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "506a9e2ba51142f6653bd599f1025fa2e8f2d320d93a4c5505af464699c4549b", + "shape": [ + 3, + 126, + 2560 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "454a6a8b46a5e687f1338ca83d9e8c43414ffbabbdda905c6977dc168d0b1425", + "shape": [ + 3, + 126, + 8229 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "ea2dbe6bba30495918cc14a223a43fb6efee780a21cbce1d503cb6200e7ef807", + "shape": [ + 3, + 126 + ] + } + } +} diff --git a/tests/goldens/dplm2_3b.safetensors b/tests/goldens/dplm2_3b.safetensors new file mode 100644 index 0000000..6386481 Binary files /dev/null and b/tests/goldens/dplm2_3b.safetensors differ diff --git a/tests/goldens/dplm2_650m.json b/tests/goldens/dplm2_650m.json new file mode 100644 index 0000000..a5da990 --- /dev/null +++ b/tests/goldens/dplm2_650m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:4cce8d9dc212cdace0e20e89169790bcf199c158", + "pytorch_model.bin": "sha256:8d6e08cc05e4858064a714013c74cc88c9caa2cc8b12c34605a3c24bcd877cfb", + "special_tokens_map.json": "git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a", + "tokenizer_config.json": "git-sha1:fc8c21760dcff173955afb106859e5f015d4f757", + "vocab.txt": "git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37" + }, + "repo_id": "airkingbd/dplm2_650m", + "revision": "0bc69b644976c6680ab7e26669854d1979e8876e" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm2_650m" + ], + "input_fingerprint": "a7b239a22465f83adf8de7303bb90b0b85ddccfee476137f54331da231059910", + "model_id": "dplm2_650m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "ef3369eb9d4ae7c143b7041212bfebf9531b5b75eb1eac4170e07126b9a28e8d", + "native/metadata.json": "3fcb1766ae6a2dd7e676ee207dba22c5e73f666b827f63ed2eb44a41686313f6" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm2_650m.safetensors", + "sha256": "c4e0e467c252c3ac813363d2d4b17a5e3bd99e75fad315e76d97689b4655ddac" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "6c19ab143b0587e6766355da1f0527bf7df033f9c5de68afcd60d405a1ce995a", + "shape": [ + 3, + 126 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "11d09bddaf02ce351b86b1c5ce5f05b2e89420605014b337e2c49061d9e57f3d", + "shape": [ + 3, + 126 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "cad36e84d72090e829c3398bc118ec628171d6abca7f2897afef3a9b733599f3", + "shape": [ + 3, + 126, + 1280 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "abe35860b57a618599cd0351e4640e6d2abfe95b8655242bff788b69ed53c2af", + "shape": [ + 3, + 126, + 8229 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "ea2dbe6bba30495918cc14a223a43fb6efee780a21cbce1d503cb6200e7ef807", + "shape": [ + 3, + 126 + ] + } + } +} diff --git a/tests/goldens/dplm2_650m.safetensors b/tests/goldens/dplm2_650m.safetensors new file mode 100644 index 0000000..2c97e6a Binary files /dev/null and b/tests/goldens/dplm2_650m.safetensors differ diff --git a/tests/goldens/dplm_150m.json b/tests/goldens/dplm_150m.json new file mode 100644 index 0000000..3458147 --- /dev/null +++ b/tests/goldens/dplm_150m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:4910cb02f1840e9ac577026f601829604af58c74", + "pytorch_model.bin": "sha256:ea4eaa99536b60ed76f945f71a1a5e604f08447ec3def5104a93ca6001a59961", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "airkingbd/dplm_150m", + "revision": "49b7125a5d28c6418fcc2f3c4fe799352ac1488b" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm_150m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "dplm_150m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "1e759d69828a34f943592cfc3f7f63d1976348d68636a922da6604fba5999a54", + "native/metadata.json": "5937d1229c6af632bbcc66dc6312d5aa887621d1dd9915e68502ba1e2e68eaed" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm_150m.safetensors", + "sha256": "392992235195beed97ab8359b90a2e11e52f4326606f99a471447bed81d146bd" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "cc5e3aa17cdd7e3ec072e8c2411a4ca5649a7a65b26e04ed3ae82ed4eb5226bd", + "shape": [ + 3, + 63, + 640 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "763667a137def055ccf8cd87f1bcbcc1eea89ef3e3cff6bf0d3e1ea2ad70d4f8", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/dplm_150m.safetensors b/tests/goldens/dplm_150m.safetensors new file mode 100644 index 0000000..4df6bfa Binary files /dev/null and b/tests/goldens/dplm_150m.safetensors differ diff --git a/tests/goldens/dplm_3b.json b/tests/goldens/dplm_3b.json new file mode 100644 index 0000000..9044b6b --- /dev/null +++ b/tests/goldens/dplm_3b.json @@ -0,0 +1,99 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:f6206456e8c2f22ebe1d37fce3b5d50fd8073e68", + "pytorch_model-00001-of-00004.bin": "sha256:0bcb86a115fe744ed686756db143f78851304e855e2f83cec58681c6080ced5f", + "pytorch_model-00002-of-00004.bin": "sha256:daf3324f3be949e7dd1c3c84b28da7fec5151b1890cb0904e73427266856a06f", + "pytorch_model-00003-of-00004.bin": "sha256:dbbeb7924a21059854f994931e23590b054aa000b10370a71c052c4aa36e9246", + "pytorch_model-00004-of-00004.bin": "sha256:21c01740d091487db43446489d8a893dea1fcc6f2e1c1991ece13945f7ab4e07", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "airkingbd/dplm_3b", + "revision": "53849d4a7fe944ae0b9cf2bbc0d2cc0054795b51" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm_3b" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "dplm_3b", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "09a98566251cafbdc6daceb9520adcba54430354aa4521e64aaefd04abf46db1", + "native/metadata.json": "953287da8d2b74557f718091b10617e100a1e17292a436f483ad03b6cccef7f0" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm_3b.safetensors", + "sha256": "75b0a0854fc391133920b0feaaeb8f69ab7568a88b3759627aca1556c4338c1e" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "7bd04473fafb3b2b78862690864be2ccfefbed12fdc3cd63162cd652cc292c06", + "shape": [ + 3, + 63, + 2560 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "ece5b588a39984172a16a3b10a23b79163a94fad65efd63c52228e7440dd7132", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/dplm_3b.safetensors b/tests/goldens/dplm_3b.safetensors new file mode 100644 index 0000000..681978f Binary files /dev/null and b/tests/goldens/dplm_3b.safetensors differ diff --git a/tests/goldens/dplm_650m.json b/tests/goldens/dplm_650m.json new file mode 100644 index 0000000..0f670e8 --- /dev/null +++ b/tests/goldens/dplm_650m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:625574d625a4178ca6966e9545fee56026c0b634", + "pytorch_model.bin": "sha256:db4e54343a89e7600f41c3aacbc593db1b0caee82ec28cab25ff2ae090eba39c", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "airkingbd/dplm_650m", + "revision": "7a7e651baa667d094aba05e9dc1cf52a3332110a" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"absl-py\":\"2.5.0\",\"accelerate\":\"1.14.0\",\"aiohappyeyeballs\":\"2.7.1\",\"aiohttp\":\"3.14.1\",\"aiosignal\":\"1.4.0\",\"alembic\":\"1.18.5\",\"annotated-types\":\"0.7.0\",\"antlr4-python3-runtime\":\"4.9.3\",\"async-timeout\":\"5.0.1\",\"attrs\":\"26.1.0\",\"autopage\":\"0.6.0\",\"backports.strenum\":\"1.3.1\",\"biopython\":\"1.79\",\"biotite\":\"1.2.0\",\"biotraj\":\"1.2.2\",\"black\":\"26.5.1\",\"byprot\":\"1.0.0\",\"certifi\":\"2026.6.17\",\"cfgv\":\"3.5.0\",\"cftime\":\"1.6.5\",\"charset-normalizer\":\"2.1.1\",\"click\":\"8.4.2\",\"cliff\":\"4.14.0\",\"cmaes\":\"0.13.0\",\"cmd2\":\"3.5.1\",\"colorlog\":\"6.10.1\",\"contourpy\":\"1.3.2\",\"cycler\":\"0.12.1\",\"datasets\":\"2.20.0\",\"debugpy\":\"1.8.21\",\"deepspeed\":\"0.14.4\",\"dill\":\"0.3.8\",\"distlib\":\"0.4.3\",\"dm-tree\":\"0.1.10\",\"e3nn\":\"0.6.0\",\"einops\":\"0.8.2\",\"exceptiongroup\":\"1.3.1\",\"fair-esm\":\"2.0.0\",\"fastjsonschema\":\"2.21.2\",\"filelock\":\"3.29.0\",\"flake8\":\"7.3.0\",\"fonttools\":\"4.63.0\",\"frozenlist\":\"1.8.0\",\"fsspec\":\"2024.5.0\",\"greenlet\":\"3.5.3\",\"griddataformats\":\"1.0.2\",\"grpcio\":\"1.82.1\",\"hf-xet\":\"1.5.1\",\"hjson\":\"3.1.0\",\"huggingface_hub\":\"0.36.2\",\"hydra-colorlog\":\"1.2.0\",\"hydra-core\":\"1.2.0\",\"hydra-optuna-sweeper\":\"1.2.0\",\"identify\":\"2.6.19\",\"idna\":\"3.4\",\"ihm\":\"2.11\",\"iniconfig\":\"2.3.0\",\"isort\":\"8.0.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"joblib\":\"1.5.3\",\"jsonschema\":\"4.26.0\",\"jsonschema-specifications\":\"2025.9.1\",\"jupyter_core\":\"5.9.1\",\"kiwisolver\":\"1.5.0\",\"lightning\":\"2.2.0\",\"lightning-utilities\":\"0.15.3\",\"lmdb\":\"2.3.0\",\"mako\":\"1.3.12\",\"markdown\":\"3.10.2\",\"markdown-it-py\":\"4.2.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.10.9\",\"mccabe\":\"0.7.0\",\"mda-xdrlib\":\"0.2.0\",\"mdanalysis\":\"2.9.0\",\"mdtraj\":\"1.10.3\",\"mdurl\":\"0.1.2\",\"ml_collections\":\"1.1.0\",\"mmtf-python\":\"1.1.3\",\"modelcif\":\"1.7\",\"mpmath\":\"1.3.0\",\"mrcfile\":\"1.5.4\",\"msgpack\":\"1.2.1\",\"multidict\":\"6.7.1\",\"multiprocess\":\"0.70.16\",\"mypy_extensions\":\"1.1.0\",\"nbformat\":\"5.10.4\",\"nbstripout\":\"0.9.1\",\"netcdf4\":\"1.7.4\",\"networkx\":\"3.4.2\",\"ninja\":\"1.13.0\",\"nodeenv\":\"1.10.0\",\"numpy\":\"1.26.4\",\"nvidia-cublas-cu12\":\"12.1.3.1\",\"nvidia-cuda-cupti-cu12\":\"12.1.105\",\"nvidia-cuda-nvrtc-cu12\":\"12.1.105\",\"nvidia-cuda-runtime-cu12\":\"12.1.105\",\"nvidia-cudnn-cu12\":\"8.9.2.26\",\"nvidia-cufft-cu12\":\"11.0.2.54\",\"nvidia-curand-cu12\":\"10.3.2.106\",\"nvidia-cusolver-cu12\":\"11.4.5.107\",\"nvidia-cusparse-cu12\":\"12.1.0.106\",\"nvidia-ml-py\":\"13.610.43\",\"nvidia-nccl-cu12\":\"2.19.3\",\"nvidia-nvjitlink-cu12\":\"12.9.86\",\"nvidia-nvtx-cu12\":\"12.1.105\",\"omegaconf\":\"2.3.1\",\"opt-einsum-fx\":\"0.1.4\",\"opt_einsum\":\"3.4.0\",\"optuna\":\"2.10.1\",\"packaging\":\"24.2\",\"pandas\":\"2.3.3\",\"parso\":\"0.8.7\",\"pathspec\":\"1.1.1\",\"peft\":\"0.11.1\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"pre_commit\":\"4.6.0\",\"prettytable\":\"3.18.0\",\"propcache\":\"0.5.2\",\"protobuf\":\"7.35.1\",\"psutil\":\"7.2.2\",\"pudb\":\"2025.1.5\",\"py-cpuinfo\":\"9.0.0\",\"pyarrow\":\"25.0.0\",\"pyarrow-hotfix\":\"0.7\",\"pycodestyle\":\"2.14.0\",\"pydantic\":\"2.13.4\",\"pydantic_core\":\"2.46.4\",\"pyflakes\":\"3.4.0\",\"pygments\":\"2.20.0\",\"pyparsing\":\"3.3.2\",\"pyperclip\":\"1.11.0\",\"pyrootutils\":\"1.0.4\",\"pytest\":\"9.1.1\",\"python-dateutil\":\"2.9.0.post0\",\"python-discovery\":\"1.4.4\",\"python-dotenv\":\"1.2.2\",\"pytokens\":\"0.4.1\",\"pytorch-lightning\":\"2.2.0\",\"pytz\":\"2026.2\",\"pyyaml\":\"6.0.3\",\"referencing\":\"0.37.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"rich\":\"15.0.0\",\"rich-argparse\":\"1.8.0\",\"rpds-py\":\"0.30.0\",\"safetensors\":\"0.6.2\",\"scikit-learn\":\"1.7.2\",\"scipy\":\"1.15.3\",\"seaborn\":\"0.13.2\",\"setuptools\":\"59.6.0\",\"sh\":\"2.3.0\",\"six\":\"1.17.0\",\"sqlalchemy\":\"2.0.51\",\"stevedore\":\"5.8.0\",\"sympy\":\"1.14.0\",\"tensorboard\":\"2.21.0\",\"tensorboard-data-server\":\"0.7.2\",\"threadpoolctl\":\"3.6.0\",\"tmtools\":\"0.3.0\",\"tokenizers\":\"0.15.2\",\"tomli\":\"2.4.1\",\"torch\":\"2.2.0+cu121\",\"torch-geometric\":\"2.8.0\",\"torch_scatter\":\"2.1.2+pt22cu121\",\"torchdata\":\"0.7.1\",\"torchmetrics\":\"1.9.0\",\"torchtext\":\"0.17.0+cpu\",\"tqdm\":\"4.66.5\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.39.2\",\"triton\":\"2.2.0\",\"typing-inspection\":\"0.4.2\",\"typing_extensions\":\"4.15.0\",\"tzdata\":\"2026.3\",\"urllib3\":\"1.26.13\",\"urwid\":\"4.0.4\",\"urwid_readline\":\"0.15.1\",\"virtualenv\":\"21.6.1\",\"wcwidth\":\"0.8.2\",\"werkzeug\":\"3.1.8\",\"wrapt\":\"2.2.2\",\"xxhash\":\"3.8.1\",\"yarl\":\"1.24.2\"}", + "python": "3.10.12", + "torch": "2.2.0+cu121" + }, + "fingerprint": "cd54743355027f7ac24d7f456e3807a2037601cdfc9eb57db041a839c3bfc503" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "dplm_650m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "dplm_650m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "ae328ef52a23b5f741389343aeee3f2e9aeb3dee239f88cbb010e51462160e8c", + "native/metadata.json": "e64e455ff06879ae6bed50586c348617757a0257f6b77a39669aab8b1fc39a63" + }, + "sources": [ + { + "id": "dplm", + "revision": "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d", + "url": "https://github.com/bytedance/dplm.git" + } + ], + "tensor_file": { + "path": "dplm_650m.safetensors", + "sha256": "073f0a6abea7e48f28c2d921ff8329a28e22627f01979277cb324908a01b3378" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "d0a0c5607de8dca67c4bd94f1c4774f100e0ba7cb8e261ace5af2fcec66b7d35", + "shape": [ + 3, + 63, + 1280 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "30ba6b9275ff1c61e8aacf76024657c50eb2d7bff1908fe86b64d449eb4aae4e", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/dplm_650m.safetensors b/tests/goldens/dplm_650m.safetensors new file mode 100644 index 0000000..ed6f4e8 Binary files /dev/null and b/tests/goldens/dplm_650m.safetensors differ diff --git a/tests/goldens/e1_150m.json b/tests/goldens/e1_150m.json new file mode 100644 index 0000000..b33c76b --- /dev/null +++ b/tests/goldens/e1_150m.json @@ -0,0 +1,109 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:485e649199b46fe6ee7456bebf7aae9b3d4baeab", + "model.safetensors": "sha256:ba2656339005e6598642836acfdafde480fecc7e145ce0058eb54adf572c3484" + }, + "repo_id": "Profluent-Bio/E1-150m", + "revision": "c4dbfe827e4aa6ed7f95eaef50dc1e084f4d77dc" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.8", + "packages": "{\"asttokens\":\"3.0.2\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"click\":\"8.4.2\",\"comm\":\"0.2.3\",\"debugpy\":\"1.8.21\",\"decorator\":\"5.3.1\",\"e1\":\"1.0.0\",\"einops\":\"0.8.2\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipykernel\":\"7.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jupyter_client\":\"8.9.1\",\"jupyter_core\":\"5.9.1\",\"kernels\":\"0.12.3\",\"markupsafe\":\"3.0.3\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"nest-asyncio2\":\"1.7.2\",\"networkx\":\"3.6.1\",\"numpy\":\"2.5.1\",\"nvidia-cublas-cu12\":\"12.8.4.1\",\"nvidia-cuda-cupti-cu12\":\"12.8.90\",\"nvidia-cuda-nvrtc-cu12\":\"12.8.93\",\"nvidia-cuda-runtime-cu12\":\"12.8.90\",\"nvidia-cudnn-cu12\":\"9.10.2.21\",\"nvidia-cufft-cu12\":\"11.3.3.83\",\"nvidia-cufile-cu12\":\"1.13.1.3\",\"nvidia-curand-cu12\":\"10.3.9.90\",\"nvidia-cusolver-cu12\":\"11.7.3.90\",\"nvidia-cusparse-cu12\":\"12.5.8.93\",\"nvidia-cusparselt-cu12\":\"0.7.1\",\"nvidia-nccl-cu12\":\"2.27.3\",\"nvidia-nvjitlink-cu12\":\"12.8.93\",\"nvidia-nvtx-cu12\":\"12.8.90\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"polars\":\"1.42.1\",\"polars-runtime-32\":\"1.42.1\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"pyzmq\":\"27.1.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"scipy\":\"1.18.0\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.22.1\",\"torch\":\"2.8.0+cu128\",\"tornado\":\"6.5.7\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.56.2\",\"triton\":\"3.4.0\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\"}", + "python": "3.12.3", + "torch": "2.8.0+cu128" + }, + "fingerprint": "190e7e4a5c17d3b86756f71b7788e6bed7a86716fe1a987cc06009dfa54ebee9" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "e1_150m" + ], + "input_fingerprint": "69c301728d054b0ab450b9631d5ff8bac2822e0ff35542c7923b58cae4571da9", + "model_id": "e1_150m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "85cf08b83e1d5026cb6534d4000796440b0cc1f49f53adabdb798c289aa4f48c", + "native/metadata.json": "73f52ea3dc06148a766aa7ce0a9261b38c7c634e810399c0e46c9956bdadff42" + }, + "sources": [ + { + "id": "e1", + "revision": "bfd2620a602248499f3d2583d85a7ecddf0b6e02", + "url": "https://github.com/Profluent-AI/E1.git" + } + ], + "tensor_file": { + "path": "e1_150m.safetensors", + "sha256": "6558bc8f1a7b20629eaaaa6f72601d0c2cdb859a5dc13595549b1773b6e2de41" + }, + "tensors": { + "input__global_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "57283779068d6b70d008b8062820a13e59cd12e3708e3ab109466bbcb48f55b3", + "shape": [ + 3, + 65 + ] + }, + "input__sequence_ids": { + "dtype": "int64", + "sha256": "5381cee8684c258e11ece6fba1ae6897f2897281cae905d59a7ff2d87fedee62", + "shape": [ + 3, + 65 + ] + }, + "input__within_seq_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "52e65a4d037775956a7a65cc4d78ffbbb54d66f06be3695c55ed9f5cebee1efe", + "shape": [ + 3, + 65, + 768 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "37a7e36f520330df7364f9f17c2a40fb1a0aae853312eee2b7e210533f6ccdb9", + "shape": [ + 3, + 65, + 34 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "a622292a616841968bf5ad71c03d53a1f7ac3cbb9b26a1cdaab65c36bb4116a1", + "shape": [ + 3, + 65 + ] + } + } +} diff --git a/tests/goldens/e1_150m.safetensors b/tests/goldens/e1_150m.safetensors new file mode 100644 index 0000000..80d0c9e Binary files /dev/null and b/tests/goldens/e1_150m.safetensors differ diff --git a/tests/goldens/e1_300m.json b/tests/goldens/e1_300m.json new file mode 100644 index 0000000..f451e96 --- /dev/null +++ b/tests/goldens/e1_300m.json @@ -0,0 +1,109 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:918cb09e6e96d4719ed85951f38c693360f9cdb8", + "model.safetensors": "sha256:31e09a2542f45b04e6ce4adafb3b657f21e2d56d12bf68fd2266b1576a80bc9b" + }, + "repo_id": "Profluent-Bio/E1-300m", + "revision": "5a2871c587eadbcc9237bc686ea45e5b4d28dfb3" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.8", + "packages": "{\"asttokens\":\"3.0.2\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"click\":\"8.4.2\",\"comm\":\"0.2.3\",\"debugpy\":\"1.8.21\",\"decorator\":\"5.3.1\",\"e1\":\"1.0.0\",\"einops\":\"0.8.2\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipykernel\":\"7.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jupyter_client\":\"8.9.1\",\"jupyter_core\":\"5.9.1\",\"kernels\":\"0.12.3\",\"markupsafe\":\"3.0.3\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"nest-asyncio2\":\"1.7.2\",\"networkx\":\"3.6.1\",\"numpy\":\"2.5.1\",\"nvidia-cublas-cu12\":\"12.8.4.1\",\"nvidia-cuda-cupti-cu12\":\"12.8.90\",\"nvidia-cuda-nvrtc-cu12\":\"12.8.93\",\"nvidia-cuda-runtime-cu12\":\"12.8.90\",\"nvidia-cudnn-cu12\":\"9.10.2.21\",\"nvidia-cufft-cu12\":\"11.3.3.83\",\"nvidia-cufile-cu12\":\"1.13.1.3\",\"nvidia-curand-cu12\":\"10.3.9.90\",\"nvidia-cusolver-cu12\":\"11.7.3.90\",\"nvidia-cusparse-cu12\":\"12.5.8.93\",\"nvidia-cusparselt-cu12\":\"0.7.1\",\"nvidia-nccl-cu12\":\"2.27.3\",\"nvidia-nvjitlink-cu12\":\"12.8.93\",\"nvidia-nvtx-cu12\":\"12.8.90\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"polars\":\"1.42.1\",\"polars-runtime-32\":\"1.42.1\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"pyzmq\":\"27.1.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"scipy\":\"1.18.0\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.22.1\",\"torch\":\"2.8.0+cu128\",\"tornado\":\"6.5.7\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.56.2\",\"triton\":\"3.4.0\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\"}", + "python": "3.12.3", + "torch": "2.8.0+cu128" + }, + "fingerprint": "190e7e4a5c17d3b86756f71b7788e6bed7a86716fe1a987cc06009dfa54ebee9" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "e1_300m" + ], + "input_fingerprint": "69c301728d054b0ab450b9631d5ff8bac2822e0ff35542c7923b58cae4571da9", + "model_id": "e1_300m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "52338c7641eeb720e87ad35ee48f54f585e38f15e7178163fe5d767e2f49e7e7", + "native/metadata.json": "9404d69bfdcfd1d2b2b63d7f407b93df80115be91bb9b8554ada7c6f1d8d6270" + }, + "sources": [ + { + "id": "e1", + "revision": "bfd2620a602248499f3d2583d85a7ecddf0b6e02", + "url": "https://github.com/Profluent-AI/E1.git" + } + ], + "tensor_file": { + "path": "e1_300m.safetensors", + "sha256": "92778b9ef95a803ddc84b3e3ca764c59e045872a94bcff0eb0cd47647732c188" + }, + "tensors": { + "input__global_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "57283779068d6b70d008b8062820a13e59cd12e3708e3ab109466bbcb48f55b3", + "shape": [ + 3, + 65 + ] + }, + "input__sequence_ids": { + "dtype": "int64", + "sha256": "5381cee8684c258e11ece6fba1ae6897f2897281cae905d59a7ff2d87fedee62", + "shape": [ + 3, + 65 + ] + }, + "input__within_seq_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "7b9127c81f260c8b356ea0cbd5350651ea2f6745ce6ecea0399fc8d437dae058", + "shape": [ + 3, + 65, + 1024 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "ebeb995d44c50e94b864f3a77933bfac4510ab0e8fe7602bdaa309996eaafe1b", + "shape": [ + 3, + 65, + 34 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "a622292a616841968bf5ad71c03d53a1f7ac3cbb9b26a1cdaab65c36bb4116a1", + "shape": [ + 3, + 65 + ] + } + } +} diff --git a/tests/goldens/e1_300m.safetensors b/tests/goldens/e1_300m.safetensors new file mode 100644 index 0000000..42297dc Binary files /dev/null and b/tests/goldens/e1_300m.safetensors differ diff --git a/tests/goldens/e1_600m.json b/tests/goldens/e1_600m.json new file mode 100644 index 0000000..f2d02c3 --- /dev/null +++ b/tests/goldens/e1_600m.json @@ -0,0 +1,109 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:8a0a439ed4201462bc01189c9f8b43523b257b5c", + "model.safetensors": "sha256:cfc108d4b98baaa62932331b40be265eae39dc382595bc3cde4a5ab55db1bf7a" + }, + "repo_id": "Profluent-Bio/E1-600m", + "revision": "52d959fb87a609d15cf223a485127b29ed5c382a" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.8", + "packages": "{\"asttokens\":\"3.0.2\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"click\":\"8.4.2\",\"comm\":\"0.2.3\",\"debugpy\":\"1.8.21\",\"decorator\":\"5.3.1\",\"e1\":\"1.0.0\",\"einops\":\"0.8.2\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipykernel\":\"7.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jupyter_client\":\"8.9.1\",\"jupyter_core\":\"5.9.1\",\"kernels\":\"0.12.3\",\"markupsafe\":\"3.0.3\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"nest-asyncio2\":\"1.7.2\",\"networkx\":\"3.6.1\",\"numpy\":\"2.5.1\",\"nvidia-cublas-cu12\":\"12.8.4.1\",\"nvidia-cuda-cupti-cu12\":\"12.8.90\",\"nvidia-cuda-nvrtc-cu12\":\"12.8.93\",\"nvidia-cuda-runtime-cu12\":\"12.8.90\",\"nvidia-cudnn-cu12\":\"9.10.2.21\",\"nvidia-cufft-cu12\":\"11.3.3.83\",\"nvidia-cufile-cu12\":\"1.13.1.3\",\"nvidia-curand-cu12\":\"10.3.9.90\",\"nvidia-cusolver-cu12\":\"11.7.3.90\",\"nvidia-cusparse-cu12\":\"12.5.8.93\",\"nvidia-cusparselt-cu12\":\"0.7.1\",\"nvidia-nccl-cu12\":\"2.27.3\",\"nvidia-nvjitlink-cu12\":\"12.8.93\",\"nvidia-nvtx-cu12\":\"12.8.90\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pip\":\"26.1.1\",\"platformdirs\":\"4.10.0\",\"pluggy\":\"1.6.0\",\"polars\":\"1.42.1\",\"polars-runtime-32\":\"1.42.1\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"pygments\":\"2.20.0\",\"pytest\":\"8.4.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"pyzmq\":\"27.1.0\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"scipy\":\"1.18.0\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tokenizers\":\"0.22.1\",\"torch\":\"2.8.0+cu128\",\"tornado\":\"6.5.7\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.56.2\",\"triton\":\"3.4.0\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\"}", + "python": "3.12.3", + "torch": "2.8.0+cu128" + }, + "fingerprint": "190e7e4a5c17d3b86756f71b7788e6bed7a86716fe1a987cc06009dfa54ebee9" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "e1_600m" + ], + "input_fingerprint": "69c301728d054b0ab450b9631d5ff8bac2822e0ff35542c7923b58cae4571da9", + "model_id": "e1_600m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "00eb6085764831ba05bbc40a6ab7c3e3613f6ab53ace36e30d5bc25a652bd1fb", + "native/metadata.json": "906ca0df18bf5bf680ddcf3d2e19e208bb6ab6badea6bab6ed7ddd3f18ae3591" + }, + "sources": [ + { + "id": "e1", + "revision": "bfd2620a602248499f3d2583d85a7ecddf0b6e02", + "url": "https://github.com/Profluent-AI/E1.git" + } + ], + "tensor_file": { + "path": "e1_600m.safetensors", + "sha256": "22ed8417a4651ded255099f6d15c63c2c40552e700d2b0470d1adfde3a39c513" + }, + "tensors": { + "input__global_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "57283779068d6b70d008b8062820a13e59cd12e3708e3ab109466bbcb48f55b3", + "shape": [ + 3, + 65 + ] + }, + "input__sequence_ids": { + "dtype": "int64", + "sha256": "5381cee8684c258e11ece6fba1ae6897f2897281cae905d59a7ff2d87fedee62", + "shape": [ + 3, + 65 + ] + }, + "input__within_seq_position_ids": { + "dtype": "int64", + "sha256": "e7cb07fe267b3363e823774271af3eb0b87d2e114dab9f267e87cdde73db0074", + "shape": [ + 3, + 65 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "53df9d156b70af125e64f97637931a1eb7deebe6abd1457533eb3a5cc2a56198", + "shape": [ + 3, + 65, + 1280 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "9cbb71e59df284e45540b39f2553edc16ecb9c51601b75b057bc089c3f1c6dd4", + "shape": [ + 3, + 65, + 34 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "a622292a616841968bf5ad71c03d53a1f7ac3cbb9b26a1cdaab65c36bb4116a1", + "shape": [ + 3, + 65 + ] + } + } +} diff --git a/tests/goldens/e1_600m.safetensors b/tests/goldens/e1_600m.safetensors new file mode 100644 index 0000000..159c53b Binary files /dev/null and b/tests/goldens/e1_600m.safetensors differ diff --git a/tests/goldens/esm2_150m.json b/tests/goldens/esm2_150m.json new file mode 100644 index 0000000..28bfee2 --- /dev/null +++ b/tests/goldens/esm2_150m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:52e04179e6fbad6663a94ea5cc44f09d764c5cd4", + "model.safetensors": "sha256:c3f1da8aea53bddd32c246c86168c23b9fd72341fb9db9a94436f855f5053566", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "facebook/esm2_t30_150M_UR50D", + "revision": "a695f6045e2e32885fa60af20c13cb35398ce30c" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"fair-esm\":\"2.0.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.2\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"pygments\":\"2.20.0\",\"pytest\":\"9.0.2\",\"pyyaml\":\"6.0.3\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"setuptools\":\"78.1.0\",\"sympy\":\"1.14.0\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "bbf86d3a107f16fa44d7d49ae8a9ab208c80a2da6104de22006aaadf018dcf27" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm2_150m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm2_150m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "92d69c4ec9bca3b54e702741d77613fd8910f2829f933f92e55f65f21246a6ce", + "native/metadata.json": "3af2b7a66bb1682f9f0a0a15b55e7e6fe61bb38858b8dc34f6830196429932c1" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + } + ], + "tensor_file": { + "path": "esm2_150m.safetensors", + "sha256": "c03fe9916dba137b452a6bbe944c7dc414db4019a6f0921e87b92d4bb6a8a42f" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "c03427ebaaf7b7f8403fd6f8776b17600dfd54ea2f806a1369339a14bb7aabcc", + "shape": [ + 3, + 63, + 640 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "f2a8325a1950d2b939fdcbd8dfc93cb62ccbdbed5a6c54505fb8d1ca966d2e5c", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm2_150m.safetensors b/tests/goldens/esm2_150m.safetensors new file mode 100644 index 0000000..467a4f5 Binary files /dev/null and b/tests/goldens/esm2_150m.safetensors differ diff --git a/tests/goldens/esm2_35m.json b/tests/goldens/esm2_35m.json new file mode 100644 index 0000000..3effb04 --- /dev/null +++ b/tests/goldens/esm2_35m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:3f64131bb610ed1ce482c4b5421fc358c785278f", + "model.safetensors": "sha256:e35647818e0e064351d4531ed480d225a002567b4b2b93ad3a9246d753150fc0", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "facebook/esm2_t12_35M_UR50D", + "revision": "6fbf070e65b0b7291e7bbcd451118c216cff79d8" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"fair-esm\":\"2.0.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.2\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"pygments\":\"2.20.0\",\"pytest\":\"9.0.2\",\"pyyaml\":\"6.0.3\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"setuptools\":\"78.1.0\",\"sympy\":\"1.14.0\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "bbf86d3a107f16fa44d7d49ae8a9ab208c80a2da6104de22006aaadf018dcf27" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm2_35m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm2_35m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "94a580c6d71a57c92ddb0fc7f5ad9a9cc16d88d5b49a32139f291dedd625e3d3", + "native/metadata.json": "07edcb417ff7878627e5dfb48453d22b09faab318c163ad8d7a340e7b8b73079" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + } + ], + "tensor_file": { + "path": "esm2_35m.safetensors", + "sha256": "c9b8bb616cf884fb7744521a2fcc6eed23586342d11241e6c9ef16454ec31e17" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "51a3a127aaf52cbbefdb062562d90c79ff93108503a36e4b674dad82c1b88cb8", + "shape": [ + 3, + 63, + 480 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "5b647a1a528e7cf17dc5b5057da75408d7e5489602ad619df9bec8c52247e8e1", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm2_35m.safetensors b/tests/goldens/esm2_35m.safetensors new file mode 100644 index 0000000..7a4f69d Binary files /dev/null and b/tests/goldens/esm2_35m.safetensors differ diff --git a/tests/goldens/esm2_3b.json b/tests/goldens/esm2_3b.json new file mode 100644 index 0000000..fbd734f --- /dev/null +++ b/tests/goldens/esm2_3b.json @@ -0,0 +1,97 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:69e7563923f87d2d7439bfb83e5a19b44b46d71b", + "pytorch_model-00001-of-00002.bin": "sha256:0f971f11c449d21422aa982b791619c10351972992c735f4c3cd43fe09790412", + "pytorch_model-00002-of-00002.bin": "sha256:7560b46fc383c691fb74b915b7d4bcef40d3df181447f16ba4b298845e308d0c", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "facebook/esm2_t36_3B_UR50D", + "revision": "476b639933c8baad5ad09a60ac1a87f987b656fc" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"fair-esm\":\"2.0.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.2\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"pygments\":\"2.20.0\",\"pytest\":\"9.0.2\",\"pyyaml\":\"6.0.3\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"setuptools\":\"78.1.0\",\"sympy\":\"1.14.0\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "bbf86d3a107f16fa44d7d49ae8a9ab208c80a2da6104de22006aaadf018dcf27" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm2_3b" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm2_3b", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "d2aa4af21d51eb28436ba1cf9e84d9ed05698c72e8b2f85a82b2f89cf7b03415", + "native/metadata.json": "d1b31384d3f92441490ef16c7daa44b085347616a42264cf7fdd37e8c567c18a" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + } + ], + "tensor_file": { + "path": "esm2_3b.safetensors", + "sha256": "dfd5a8cb05d3e814a080185c4808c8e7ec2277f070f395562fcfbe4376789e4e" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "f33b7320f901536d898f386863969bb5cabee629823e53108a41857acb8086d8", + "shape": [ + 3, + 63, + 2560 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "f181a80c9349476aec4ec4878ee7eaf1cd82b08be7ec1808168510db1ceb3eae", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm2_3b.safetensors b/tests/goldens/esm2_3b.safetensors new file mode 100644 index 0000000..ce46081 Binary files /dev/null and b/tests/goldens/esm2_3b.safetensors differ diff --git a/tests/goldens/esm2_650m.json b/tests/goldens/esm2_650m.json new file mode 100644 index 0000000..82626eb --- /dev/null +++ b/tests/goldens/esm2_650m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:a956a25d277f30bd870d3760b9a116f19ead885e", + "model.safetensors": "sha256:a08adabb949fa67ad3c14b509d04fd60368b35007b0095e3358f81200c4f4db0", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "facebook/esm2_t33_650M_UR50D", + "revision": "08e4846e537177426273712802403f7ba8261b6c" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"fair-esm\":\"2.0.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.2\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"pygments\":\"2.20.0\",\"pytest\":\"9.0.2\",\"pyyaml\":\"6.0.3\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"setuptools\":\"78.1.0\",\"sympy\":\"1.14.0\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "bbf86d3a107f16fa44d7d49ae8a9ab208c80a2da6104de22006aaadf018dcf27" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm2_650m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm2_650m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "65e98c3f3508b0db4b108eb8b4550f7f59250a7b1fb732b75c4d363d022941fe", + "native/metadata.json": "459c6d50138cb6a720654db046416bcb1720ed79591297b9e85f63e4d8796360" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + } + ], + "tensor_file": { + "path": "esm2_650m.safetensors", + "sha256": "c3a66b75add03628e62e238cb63da6a9e4d321f8160e84bdf2a131c096977f86" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "3230510fda222dabeabab727a9d7f207016b90ba30c7e7c0e08ebf11f93cbd69", + "shape": [ + 3, + 63, + 1280 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "acd43879cd264ed6edcbfc2d7333fe9aae6b02b3b246f4bb8f9b8e2301786008", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm2_650m.safetensors b/tests/goldens/esm2_650m.safetensors new file mode 100644 index 0000000..c1766ee Binary files /dev/null and b/tests/goldens/esm2_650m.safetensors differ diff --git a/tests/goldens/esm2_8m.json b/tests/goldens/esm2_8m.json new file mode 100644 index 0000000..7b9f773 --- /dev/null +++ b/tests/goldens/esm2_8m.json @@ -0,0 +1,96 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:c2c6e65a87d9d20d47699ae236d605b80c741dd3", + "model.safetensors": "sha256:24c5fa474c48f3b754b86efe752d5f189d2bcd88190fa2270fc92b2ef3034189", + "special_tokens_map.json": "git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1", + "tokenizer_config.json": "git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e", + "vocab.txt": "git-sha1:6b946952cc35537226f07fd70957ee2f848880d2" + }, + "repo_id": "facebook/esm2_t6_8M_UR50D", + "revision": "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"fair-esm\":\"2.0.1\",\"filelock\":\"3.29.0\",\"fsspec\":\"2026.4.0\",\"hf-xet\":\"1.5.2\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"jinja2\":\"3.1.6\",\"markupsafe\":\"3.0.3\",\"mpmath\":\"1.3.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"pygments\":\"2.20.0\",\"pytest\":\"9.0.2\",\"pyyaml\":\"6.0.3\",\"requests\":\"2.34.2\",\"safetensors\":\"0.6.2\",\"setuptools\":\"78.1.0\",\"sympy\":\"1.14.0\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "bbf86d3a107f16fa44d7d49ae8a9ab208c80a2da6104de22006aaadf018dcf27" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm2_8m" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm2_8m", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "8de8ee1ec82d516fdd61c07738c29bbcaaa11439c8af0bfdc36e0a4741534443", + "native/metadata.json": "08c2dab986ddc0fa14f8ecf85b453da1e0b5842b651081531eee9a70570f9cd5" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + } + ], + "tensor_file": { + "path": "esm2_8m.safetensors", + "sha256": "b40217566c33c71988d28869de353be54a3b3ebfc21fdfd29056e88cf7e99f4c" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "2cf89ef552f208608822f7473e2a0da3b02a9155422dc64c4c3b4b66e19640a1", + "shape": [ + 3, + 63, + 320 + ] + }, + "output__logits": { + "dtype": "float32", + "sha256": "b750b319e312aa253434dac722d6c771e1ccc51c1a9cffcc476282c4da083596", + "shape": [ + 3, + 63, + 33 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm2_8m.safetensors b/tests/goldens/esm2_8m.safetensors new file mode 100644 index 0000000..2a4ab93 Binary files /dev/null and b/tests/goldens/esm2_8m.safetensors differ diff --git a/tests/goldens/esm3_small.json b/tests/goldens/esm3_small.json new file mode 100644 index 0000000..b98cf60 --- /dev/null +++ b/tests/goldens/esm3_small.json @@ -0,0 +1,101 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:0967ef424bce6791893e9a57bb952f80fd536e93", + "data/weights/esm3_function_decoder_v0.pth": "sha256:f76d074efcaccfe21365a4fa96f212dadd66798e1e49d809ab7ffbe025d227c9", + "data/weights/esm3_sm_open_v1.pth": "sha256:5ead5a135c658068db6a4f1b933e72d6110992c4668822e1c0e2dcc53e38acd9", + "data/weights/esm3_structure_decoder_v0.pth": "sha256:3b726258a44274792b40ce7ea307e10c5da09936368a4ffa2970264d909da65b", + "data/weights/esm3_structure_encoder_v0.pth": "sha256:467acbaee703ba3ccde6e75241a912a316952e5ff071355f85c1d33c68704f40" + }, + "repo_id": "biohub/esm3-sm-open-v1", + "revision": "47f0545b2b6daf26a93439a3cd610f4f7f3d5478" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"accelerate\":\"1.13.0\",\"anyio\":\"4.14.2\",\"asttokens\":\"3.0.2\",\"attrs\":\"26.1.0\",\"biopython\":\"1.87\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"boto3\":\"1.43.49\",\"botocore\":\"1.43.49\",\"brotli\":\"1.2.0\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cloudpathlib\":\"0.24.0\",\"comm\":\"0.2.3\",\"contourpy\":\"1.3.3\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"cycler\":\"0.12.1\",\"decorator\":\"5.3.1\",\"dna_features_viewer\":\"3.1.5\",\"einops\":\"0.8.2\",\"esm\":\"3.3.0\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fonttools\":\"4.63.0\",\"fsspec\":\"2026.4.0\",\"h11\":\"0.16.0\",\"hf-xet\":\"1.5.1\",\"httpcore\":\"1.0.9\",\"httpx\":\"0.28.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"ipywidgets\":\"8.1.8\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jmespath\":\"1.1.0\",\"joblib\":\"1.5.3\",\"jupyterlab_widgets\":\"3.0.16\",\"kiwisolver\":\"1.5.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.11.0\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"msgpack-numpy\":\"0.4.8\",\"narwhals\":\"2.24.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"py3dmol\":\"2.5.5\",\"pydssp\":\"0.9.1\",\"pygments\":\"2.20.0\",\"pygtrie\":\"2.5.0\",\"pyparsing\":\"3.3.2\",\"pytest\":\"9.0.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"rdkit\":\"2026.3.4\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"s3transfer\":\"0.19.1\",\"safetensors\":\"0.5.3\",\"scikit-learn\":\"1.9.0\",\"scipy\":\"1.17.1\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tenacity\":\"9.1.4\",\"threadpoolctl\":\"3.6.0\",\"tokenizers\":\"0.22.2\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.57.6\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\",\"widgetsnbextension\":\"4.0.15\",\"zstd\":\"1.5.7.3\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "2ee39b50c310ea21ecbb1e87a69bb9a05076126031d2e6566014ba794dd57c1b" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esm3_small" + ], + "input_fingerprint": "20ef27191e0e1f305f2345f8ed813f3cf685c82856a90fce0fa9e5209ce3696e", + "model_id": "esm3_small", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "a7519c26786ca0de53d80054eb4b8c273c97be291f6453d5b7fb5d2e9d3df573", + "native/metadata.json": "3983df0b109c4a90dd37d359155442d1dd77c75bedf7c92b02ddf7f1212a452e" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + } + ], + "tensor_file": { + "path": "esm3_small.safetensors", + "sha256": "d957922f810c9ab4c557d80d5aaaf6a3aab79a5a45e4638012a634a4134803b1" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "float32", + "sha256": "837b3b53736ca019bb0acf2aacbb8d010957461c958fa3a15b4c8e5953265251", + "shape": [ + 3, + 63, + 1536 + ] + }, + "output__logits": { + "dtype": "bfloat16", + "sha256": "b4e7a076eaeae92de783ad73f8bd2690dd697bbb35c00993cc0fc2488fb67588", + "shape": [ + 3, + 63, + 64 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esm3_small.safetensors b/tests/goldens/esm3_small.safetensors new file mode 100644 index 0000000..04b86df Binary files /dev/null and b/tests/goldens/esm3_small.safetensors differ diff --git a/tests/goldens/esmc_6b.json b/tests/goldens/esmc_6b.json new file mode 100644 index 0000000..5f1c5c0 --- /dev/null +++ b/tests/goldens/esmc_6b.json @@ -0,0 +1,114 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:19f5fb09e4f630fb5b748a497183c22a87ec5102", + "model-00001-of-00006.safetensors": "sha256:bd90149ff223e6ac1a0cac6147a5ae0df20d3a21df4f65356a1f19cd14f4aa8a", + "model-00002-of-00006.safetensors": "sha256:f75e2144d8269fe2eb4b3e0823fb089b94f176d8024153e85b8fb573a42294fa", + "model-00003-of-00006.safetensors": "sha256:f699f01ecc9691d9c6470492765fe54b8b5d2e9f277c139e89427433ffdfe0b2", + "model-00004-of-00006.safetensors": "sha256:46add1b7be098bbfdc3073884851ba3057f1b33ea23a158b650a37007dabd13d", + "model-00005-of-00006.safetensors": "sha256:1e1cb62f060a34e18f54a31a76683ef888b8cec59e73315f5b31d25d45a1f88c", + "model-00006-of-00006.safetensors": "sha256:56c73e13ae96e777ce65eee99364056069ef93b646470f352f83c5f1037b1b18", + "special_tokens_map.json": "git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json": "git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json": "git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61" + }, + "repo_id": "biohub/ESMC-6B", + "revision": "45b0fa5d7fb06faefbd5e3b89bdcef35d564e79a" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"accelerate\":\"1.13.0\",\"anyio\":\"4.14.2\",\"asttokens\":\"3.0.2\",\"attrs\":\"26.1.0\",\"biopython\":\"1.87\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"boto3\":\"1.43.49\",\"botocore\":\"1.43.49\",\"brotli\":\"1.2.0\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cloudpathlib\":\"0.24.0\",\"comm\":\"0.2.3\",\"contourpy\":\"1.3.3\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"cycler\":\"0.12.1\",\"decorator\":\"5.3.1\",\"dna_features_viewer\":\"3.1.5\",\"einops\":\"0.8.2\",\"esm\":\"3.3.0\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fonttools\":\"4.63.0\",\"fsspec\":\"2026.4.0\",\"h11\":\"0.16.0\",\"hf-xet\":\"1.5.1\",\"httpcore\":\"1.0.9\",\"httpx\":\"0.28.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"ipywidgets\":\"8.1.8\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jmespath\":\"1.1.0\",\"joblib\":\"1.5.3\",\"jupyterlab_widgets\":\"3.0.16\",\"kiwisolver\":\"1.5.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.11.0\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"msgpack-numpy\":\"0.4.8\",\"narwhals\":\"2.24.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"py3dmol\":\"2.5.5\",\"pydssp\":\"0.9.1\",\"pygments\":\"2.20.0\",\"pygtrie\":\"2.5.0\",\"pyparsing\":\"3.3.2\",\"pytest\":\"9.0.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"rdkit\":\"2026.3.4\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"s3transfer\":\"0.19.1\",\"safetensors\":\"0.5.3\",\"scikit-learn\":\"1.9.0\",\"scipy\":\"1.17.1\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tenacity\":\"9.1.4\",\"threadpoolctl\":\"3.6.0\",\"tokenizers\":\"0.22.2\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.57.6\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\",\"widgetsnbextension\":\"4.0.15\",\"zstd\":\"1.5.7.3\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "2ee39b50c310ea21ecbb1e87a69bb9a05076126031d2e6566014ba794dd57c1b" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmc_6b" + ], + "input_fingerprint": "182b178e6588beba6a1c681379e3bc8068824eccecc192891173efb89c5068f6", + "model_id": "esmc_6b", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "3c41e3d2b47b03cf573f229df1fb2b15f3d8287ef8ed6eb0336f89f212214d78", + "native/metadata.json": "43b4fc3571d0ae565ddad8423e8e1e90de2daf484bae4192b2f61d867e4e8990" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + } + ], + "tensor_file": { + "path": "esmc_6b.safetensors", + "sha256": "a948945e985c7deaca7be8b7eed09c0a9521a2af3f2b10fc2ec7a7d2a0f99ada" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "input__sequence_id": { + "dtype": "bool", + "sha256": "546670f3d1d77f83906edd5dac50f17e99493f27136e41b2fcae8d6d3d9613ba", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "397f1a0b45fbb7d486e7785264a37d42de19568d16fdb1d575f79c2389763d2b", + "shape": [ + 3, + 63, + 2560 + ] + }, + "output__logits": { + "dtype": "bfloat16", + "sha256": "952de90dae9c1868c2006148bcbe04d4fa2886beec6908bc702a097a96315b1f", + "shape": [ + 3, + 63, + 64 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esmc_6b.safetensors b/tests/goldens/esmc_6b.safetensors new file mode 100644 index 0000000..13f5cbb Binary files /dev/null and b/tests/goldens/esmc_6b.safetensors differ diff --git a/tests/goldens/esmc_large.json b/tests/goldens/esmc_large.json new file mode 100644 index 0000000..b678431 --- /dev/null +++ b/tests/goldens/esmc_large.json @@ -0,0 +1,109 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:71c8241dc28a5fb636248267a0927c0242b264c1", + "model.safetensors": "sha256:e4232c30fd35fe2f57051ec88a703996ac94520580b4b836894207a3d45d9ff8", + "special_tokens_map.json": "git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json": "git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json": "git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61" + }, + "repo_id": "biohub/ESMC-600M", + "revision": "a7e82012c83126b9eedb055fea9fa84b6c02f094" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"accelerate\":\"1.13.0\",\"anyio\":\"4.14.2\",\"asttokens\":\"3.0.2\",\"attrs\":\"26.1.0\",\"biopython\":\"1.87\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"boto3\":\"1.43.49\",\"botocore\":\"1.43.49\",\"brotli\":\"1.2.0\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cloudpathlib\":\"0.24.0\",\"comm\":\"0.2.3\",\"contourpy\":\"1.3.3\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"cycler\":\"0.12.1\",\"decorator\":\"5.3.1\",\"dna_features_viewer\":\"3.1.5\",\"einops\":\"0.8.2\",\"esm\":\"3.3.0\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fonttools\":\"4.63.0\",\"fsspec\":\"2026.4.0\",\"h11\":\"0.16.0\",\"hf-xet\":\"1.5.1\",\"httpcore\":\"1.0.9\",\"httpx\":\"0.28.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"ipywidgets\":\"8.1.8\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jmespath\":\"1.1.0\",\"joblib\":\"1.5.3\",\"jupyterlab_widgets\":\"3.0.16\",\"kiwisolver\":\"1.5.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.11.0\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"msgpack-numpy\":\"0.4.8\",\"narwhals\":\"2.24.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"py3dmol\":\"2.5.5\",\"pydssp\":\"0.9.1\",\"pygments\":\"2.20.0\",\"pygtrie\":\"2.5.0\",\"pyparsing\":\"3.3.2\",\"pytest\":\"9.0.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"rdkit\":\"2026.3.4\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"s3transfer\":\"0.19.1\",\"safetensors\":\"0.5.3\",\"scikit-learn\":\"1.9.0\",\"scipy\":\"1.17.1\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tenacity\":\"9.1.4\",\"threadpoolctl\":\"3.6.0\",\"tokenizers\":\"0.22.2\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.57.6\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\",\"widgetsnbextension\":\"4.0.15\",\"zstd\":\"1.5.7.3\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "2ee39b50c310ea21ecbb1e87a69bb9a05076126031d2e6566014ba794dd57c1b" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmc_large" + ], + "input_fingerprint": "182b178e6588beba6a1c681379e3bc8068824eccecc192891173efb89c5068f6", + "model_id": "esmc_large", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "9591067c1bb9a85177c4e5651ceb2980bf90fe8ec81ed86878889266627ec607", + "native/metadata.json": "354a5461a306568886ba799be1264574b526f1760082932b9a3f57c3f4ca19bb" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + } + ], + "tensor_file": { + "path": "esmc_large.safetensors", + "sha256": "e13302df4cf7e8381552f1043a8fd0f31f3e0d50b2ab6009fb86b7940ae8ff79" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "input__sequence_id": { + "dtype": "bool", + "sha256": "546670f3d1d77f83906edd5dac50f17e99493f27136e41b2fcae8d6d3d9613ba", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "b6b066a58e24b50fbfadb3ed4be47a0ed23c4c6f5e545875a491655a7c2496b0", + "shape": [ + 3, + 63, + 1152 + ] + }, + "output__logits": { + "dtype": "bfloat16", + "sha256": "10204b50f8f6a3062c60bacea653f2af30d0a2644417dd89eca72ed67ab29541", + "shape": [ + 3, + 63, + 64 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esmc_large.safetensors b/tests/goldens/esmc_large.safetensors new file mode 100644 index 0000000..199b2e5 Binary files /dev/null and b/tests/goldens/esmc_large.safetensors differ diff --git a/tests/goldens/esmc_small.json b/tests/goldens/esmc_small.json new file mode 100644 index 0000000..16f0677 --- /dev/null +++ b/tests/goldens/esmc_small.json @@ -0,0 +1,109 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:9a49eacf4e65c39f74381f0f0d240e3b89ef43d7", + "model.safetensors": "sha256:0772d8fe64bb25e14fe6f23b80e3c9a7d215d0da3c6cba5bd356d7c0e0bb22cc", + "special_tokens_map.json": "git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b", + "tokenizer.json": "git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", + "tokenizer_config.json": "git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61" + }, + "repo_id": "biohub/ESMC-300M", + "revision": "a59b831785f907e96e6a246b1d142bfb76df31ee" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "packages": "{\"accelerate\":\"1.13.0\",\"anyio\":\"4.14.2\",\"asttokens\":\"3.0.2\",\"attrs\":\"26.1.0\",\"biopython\":\"1.87\",\"biotite\":\"1.7.1\",\"biotraj\":\"1.2.2\",\"boto3\":\"1.43.49\",\"botocore\":\"1.43.49\",\"brotli\":\"1.2.0\",\"certifi\":\"2026.6.17\",\"charset-normalizer\":\"3.4.9\",\"cloudpathlib\":\"0.24.0\",\"comm\":\"0.2.3\",\"contourpy\":\"1.3.3\",\"cuda-bindings\":\"13.0.3\",\"cuda-pathfinder\":\"1.2.2\",\"cuda-toolkit\":\"13.0.3.0\",\"cycler\":\"0.12.1\",\"decorator\":\"5.3.1\",\"dna_features_viewer\":\"3.1.5\",\"einops\":\"0.8.2\",\"esm\":\"3.3.0\",\"executing\":\"2.2.1\",\"filelock\":\"3.29.0\",\"fonttools\":\"4.63.0\",\"fsspec\":\"2026.4.0\",\"h11\":\"0.16.0\",\"hf-xet\":\"1.5.1\",\"httpcore\":\"1.0.9\",\"httpx\":\"0.28.1\",\"huggingface_hub\":\"0.36.2\",\"idna\":\"3.18\",\"iniconfig\":\"2.3.0\",\"ipython\":\"9.15.0\",\"ipython_pygments_lexers\":\"1.1.1\",\"ipywidgets\":\"8.1.8\",\"jedi\":\"0.20.0\",\"jinja2\":\"3.1.6\",\"jmespath\":\"1.1.0\",\"joblib\":\"1.5.3\",\"jupyterlab_widgets\":\"3.0.16\",\"kiwisolver\":\"1.5.0\",\"markupsafe\":\"3.0.3\",\"matplotlib\":\"3.11.0\",\"matplotlib-inline\":\"0.2.2\",\"mpmath\":\"1.3.0\",\"msgpack\":\"1.2.1\",\"msgpack-numpy\":\"0.4.8\",\"narwhals\":\"2.24.0\",\"networkx\":\"3.6.1\",\"numpy\":\"1.26.4\",\"nvidia-cublas\":\"13.1.1.3\",\"nvidia-cuda-cupti\":\"13.0.85\",\"nvidia-cuda-nvrtc\":\"13.0.88\",\"nvidia-cuda-runtime\":\"13.0.96\",\"nvidia-cudnn-cu13\":\"9.20.0.48\",\"nvidia-cufft\":\"12.0.0.61\",\"nvidia-cufile\":\"1.15.1.6\",\"nvidia-curand\":\"10.4.0.35\",\"nvidia-cusolver\":\"12.0.4.66\",\"nvidia-cusparse\":\"12.6.3.3\",\"nvidia-cusparselt-cu13\":\"0.8.1\",\"nvidia-nccl-cu13\":\"2.29.7\",\"nvidia-nvjitlink\":\"13.2.78\",\"nvidia-nvshmem-cu13\":\"3.4.5\",\"nvidia-nvtx\":\"13.0.85\",\"packaging\":\"26.2\",\"pandas\":\"3.0.3\",\"parso\":\"0.8.7\",\"pexpect\":\"4.9.0\",\"pillow\":\"12.3.0\",\"pip\":\"26.1.1\",\"pluggy\":\"1.6.0\",\"prompt_toolkit\":\"3.0.52\",\"psutil\":\"7.2.2\",\"ptyprocess\":\"0.7.0\",\"pure_eval\":\"0.2.3\",\"py3dmol\":\"2.5.5\",\"pydssp\":\"0.9.1\",\"pygments\":\"2.20.0\",\"pygtrie\":\"2.5.0\",\"pyparsing\":\"3.3.2\",\"pytest\":\"9.0.2\",\"python-dateutil\":\"2.9.0.post0\",\"pyyaml\":\"6.0.3\",\"rdkit\":\"2026.3.4\",\"regex\":\"2026.7.10\",\"requests\":\"2.34.2\",\"s3transfer\":\"0.19.1\",\"safetensors\":\"0.5.3\",\"scikit-learn\":\"1.9.0\",\"scipy\":\"1.17.1\",\"setuptools\":\"78.1.0\",\"six\":\"1.17.0\",\"stack-data\":\"0.6.3\",\"sympy\":\"1.14.0\",\"tenacity\":\"9.1.4\",\"threadpoolctl\":\"3.6.0\",\"tokenizers\":\"0.22.2\",\"torch\":\"2.13.0+cu130\",\"tqdm\":\"4.68.4\",\"traitlets\":\"5.15.1\",\"transformers\":\"4.57.6\",\"triton\":\"3.7.1\",\"typing_extensions\":\"4.15.0\",\"urllib3\":\"2.7.0\",\"uv\":\"0.10.12\",\"wcwidth\":\"0.8.2\",\"widgetsnbextension\":\"4.0.15\",\"zstd\":\"1.5.7.3\"}", + "python": "3.12.3", + "torch": "2.13.0+cu130" + }, + "fingerprint": "2ee39b50c310ea21ecbb1e87a69bb9a05076126031d2e6566014ba794dd57c1b" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmc_small" + ], + "input_fingerprint": "182b178e6588beba6a1c681379e3bc8068824eccecc192891173efb89c5068f6", + "model_id": "esmc_small", + "schema_version": 1, + "source_files": { + "native/bf16.safetensors": "72e12b312e802ff56f56e7a210bbaba8e385a7b96101154240ce8a360e691ef9", + "native/metadata.json": "ff43803a664e0670b7f6be238256cfb84dade36994026ac1c8dcecd53478949f" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + } + ], + "tensor_file": { + "path": "esmc_small.safetensors", + "sha256": "03378d0f0fdd8161178ebb2c1f0da1b9776a726c8e8d3a10c009808a24de5654" + }, + "tensors": { + "input__attention_mask": { + "dtype": "int64", + "sha256": "a4e9ccfdfcb15d75f04569d8635f6121b56d1316e59da2db60416ce57f78551f", + "shape": [ + 3, + 63 + ] + }, + "input__input_ids": { + "dtype": "int64", + "sha256": "7c017cc9226e1262e354dd8494e871929c0ef516214ad8afa58c97df32885bed", + "shape": [ + 3, + 63 + ] + }, + "input__sequence_id": { + "dtype": "bool", + "sha256": "546670f3d1d77f83906edd5dac50f17e99493f27136e41b2fcae8d6d3d9613ba", + "shape": [ + 3, + 63 + ] + }, + "output__last_hidden_state": { + "dtype": "bfloat16", + "sha256": "7a21b7f88b158c6dd3bd2fdd4d4c3582f4ac9d735bc9d6e4414b11be1efdbe1d", + "shape": [ + 3, + 63, + 960 + ] + }, + "output__logits": { + "dtype": "bfloat16", + "sha256": "ac58aa265f493e3fe42a4d7af2f7017c794ccb63d3a6006309e8943d15b3377d", + "shape": [ + 3, + 63, + 64 + ] + }, + "residue_mask": { + "dtype": "bool", + "sha256": "73319c84be8834bd1df5ff869f2137ac3aef45c900f775c33824b61bda50eb24", + "shape": [ + 3, + 63 + ] + } + } +} diff --git a/tests/goldens/esmc_small.safetensors b/tests/goldens/esmc_small.safetensors new file mode 100644 index 0000000..c063471 Binary files /dev/null and b/tests/goldens/esmc_small.safetensors differ diff --git a/tests/goldens/esmfold.json b/tests/goldens/esmfold.json new file mode 100644 index 0000000..8b394a3 --- /dev/null +++ b/tests/goldens/esmfold.json @@ -0,0 +1,183 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:1232d0aee4be551021d8e70e66ed2b062df917bf", + "pytorch_model.bin": "sha256:2ee07356b125d1e3e57503c204111fd7323347fc4735d41d3caac57c2a78e116", + "special_tokens_map.json": "git-sha1:121c8d54f8ea66cdf678f48b3cb37c05b4de5c0d", + "tokenizer_config.json": "git-sha1:aad24fba9f1bad2d74ed79d414ddcd60e6b0f812", + "vocab.txt": "git-sha1:9abfdf5472c0ed970648b683b86ab131256b3e42" + }, + "repo_id": "facebook/esmfold_v1", + "revision": "75a3841ee059df2bf4d56688166c8fb459ddd97a" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "12.1", + "packages": "{\"esm\":\"2.0.1\",\"openfold\":\"unknown\",\"transformers\":null}", + "python": "3.10.12", + "torch": "2.2.2+cu121" + }, + "fingerprint": "7da0cb1455d0769e0c48e6ebbdfc10a1ffaa256d3c66ea53727555a7291746fe" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmfold" + ], + "input_fingerprint": "f04e1b135c58b016c4892874f07ab4f2d0a1d5c7bae8d6b3c793e9a213789124", + "model_id": "esmfold", + "schema_version": 1, + "source_files": { + "native/bundle.safetensors": "c05dca537fb181ba870dc5719f79699aa67d527b0b41df45a0bb1351665f85e5", + "native/metadata.json": "879021f65bc841eec8ca345eed96fdb612ae1a5f00c3de48130f928c8f7b03c8" + }, + "sources": [ + { + "id": "fair-esm", + "revision": "2b369911bb5b4b0dda914521b9475cad1656b2ac", + "url": "https://github.com/facebookresearch/esm.git" + }, + { + "id": "openfold", + "revision": "4b41059694619831a7db195b7e0988fc4ff3a307", + "url": "https://github.com/aqlaboratory/openfold.git" + } + ], + "tensor_file": { + "path": "esmfold.safetensors", + "sha256": "873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb" + }, + "tensors": { + "output__aatype": { + "dtype": "int64", + "sha256": "12e353264f5c6697211bb27df7fa1fd108dc8019355ea1b66baf150efcd89abf", + "shape": [ + 1, + 17 + ] + }, + "output__aligned_confidence_probs": { + "dtype": "float32", + "sha256": "783aefc0eb745dac9102115b12329c46261bd187c0e8cc33ed6cd6699a138a57", + "shape": [ + 1, + 17, + 17, + 64 + ] + }, + "output__atom14_atom_exists": { + "dtype": "float32", + "sha256": "39d00d48fb9e6355df1d745fd05e40079002b33a96d2c02ab35763436cbc78ef", + "shape": [ + 1, + 17, + 14 + ] + }, + "output__atom37_atom_exists": { + "dtype": "float32", + "sha256": "8660de4e7c67484fa52a10ceae7f694aba0f3e22b9019d37477a981e228516a2", + "shape": [ + 1, + 17, + 37 + ] + }, + "output__chain_index": { + "dtype": "int64", + "sha256": "eec3ee5312bee6dad430af43fb71ac95e6ad0f44b4d4ed977a285ca56ef201b0", + "shape": [ + 1, + 17 + ] + }, + "output__distogram_logits": { + "dtype": "bfloat16", + "sha256": "f812607760e27af921f4f31bc5a1cf1ec796752b0d4b7d77a5643457427dc4f9", + "shape": [ + 1, + 17, + 17, + 64 + ] + }, + "output__lm_logits": { + "dtype": "bfloat16", + "sha256": "4bac20c86ded2489f81b4d5216ea4e949b6f5c79a21c664bc339da9addd8569a", + "shape": [ + 1, + 17, + 23 + ] + }, + "output__mean_plddt": { + "dtype": "float32", + "sha256": "7a5d58ad6175772ee49d42728d2b6750079500b69011a2d5b54f25e0c9917582", + "shape": [ + 1 + ] + }, + "output__plddt": { + "dtype": "bfloat16", + "sha256": "fc8052d8207eb5fd4f7911331dc872e93b93fc59d33cdb04616ad4c20794de3c", + "shape": [ + 1, + 17, + 37 + ] + }, + "output__positions": { + "dtype": "float32", + "sha256": "95c5322ee14f387b30184b143615938936411ef8c748ca61a7386a3380a5e215", + "shape": [ + 8, + 1, + 17, + 14, + 3 + ] + }, + "output__predicted_aligned_error": { + "dtype": "float32", + "sha256": "8b407575ddf53428577ee01c87f934532ccbbf4853be3d8a5970e252783b951e", + "shape": [ + 1, + 17, + 17 + ] + }, + "output__ptm": { + "dtype": "float32", + "sha256": "ad4d22c0cede0d19446953b16cd20cf79fc065ca7e784d1c8e0abf481717b4f5", + "shape": [ + 1 + ] + }, + "output__ptm_logits": { + "dtype": "bfloat16", + "sha256": "516eee815e3424be539c82b893a6c23f3a3e13d417448eef0d2102bcabc892e7", + "shape": [ + 1, + 17, + 17, + 64 + ] + }, + "output__residue_index": { + "dtype": "int64", + "sha256": "a4bdd35642ad0ef486b9e1132a861bf10733d26d2dbf2c21e42b7ffc12f0d80f", + "shape": [ + 1, + 17 + ] + } + } +} diff --git a/tests/goldens/esmfold.safetensors b/tests/goldens/esmfold.safetensors new file mode 100644 index 0000000..7209e54 Binary files /dev/null and b/tests/goldens/esmfold.safetensors differ diff --git a/tests/goldens/esmfold2.json b/tests/goldens/esmfold2.json new file mode 100644 index 0000000..93a5d38 --- /dev/null +++ b/tests/goldens/esmfold2.json @@ -0,0 +1,350 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:0300c084b990b2bd600efd9f538aa5de27109fea", + "model.safetensors": "sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e" + }, + "repo_id": "biohub/ESMFold2", + "revision": "1ebf0e3481a5184eb6171d40615c79e384b48796" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "python": "3.12.3", + "torch": "2.13.0+cu130", + "transformer_engine": "null", + "transformers": "4.57.6" + }, + "fingerprint": "aebef9a5edb6b5694d3ec547d636998a94232000ee91639146e9c4acf90d7589" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmfold2" + ], + "input_fingerprint": "9237e3e721124a39a0805c93102f159134ec20709d238f83167e8d670ad2bd6e", + "model_id": "esmfold2", + "schema_version": 1, + "source_files": { + "native/bundle.safetensors": "87b39fb97107afd99e0f27d680b57b7060bccc9da2aabf07dd3d8c8f2ffdb048", + "native/metadata.json": "78201aea6ff0d248022f02944a6d6208684bcc75c443c6e4284b9d2659757599" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + }, + { + "id": "protein-ttt", + "revision": "fde2817cd84b936167cc76ccabf31e5c0fe49962", + "url": "https://github.com/anton-bushuiev/ProteinTTT.git" + } + ], + "tensor_file": { + "path": "esmfold2.safetensors", + "sha256": "e4d6be4344c528e26b13f79a9303549e3de7e582da195c0078db3ce957fad420" + }, + "tensors": { + "feature__asym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__atom_attention_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "feature__atom_to_token": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__deletion_mean": { + "dtype": "float32", + "sha256": "8c51f4825acbf8ca40bd30065492cb000a9489b8245c9447a4e355f1b41b3dbb", + "shape": [ + 1, + 56 + ] + }, + "feature__deletion_value": { + "dtype": "float32", + "sha256": "34688aedb3460cee4d864ac60a112db764823880e1775d653f842a9015c19a92", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__distogram_atom_idx": { + "dtype": "int64", + "sha256": "56d98551f23b897604f00522013b97b939a25c593ac66492ed70bccb4461f663", + "shape": [ + 1, + 56 + ] + }, + "feature__entity_id": { + "dtype": "int64", + "sha256": "568a9a29b407045c55e7793effeaa4b30b336d3e9d02691dda02290107230a1c", + "shape": [ + 1, + 56 + ] + }, + "feature__has_deletion": { + "dtype": "bool", + "sha256": "fde502e113e1d0b810151796abfb4616d3c935a9e22f7d1891e86fdb8c971f6c", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__input_ids": { + "dtype": "int64", + "sha256": "c66b63367006340c1e676e0e2c4440f349286b98786ddc1cceb6f0222b599b95", + "shape": [ + 1, + 56 + ] + }, + "feature__mol_type": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__msa": { + "dtype": "int64", + "sha256": "14efcbac324a4ec831a0af4ec220ef00a57c98eea17525eb16a8ed22e9dde34b", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__msa_attention_mask": { + "dtype": "bool", + "sha256": "fa513f2a8c37868e06a9fbb53ac5f6dd9eb28106680113748e9ff868cfa259d5", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__ref_atom_name_chars": { + "dtype": "int64", + "sha256": "70165370912f9420bfb1a8ebe5b5ba19bf88cd6fbd43395aaaab70189ac3b2ec", + "shape": [ + 1, + 448, + 4 + ] + }, + "feature__ref_charge": { + "dtype": "int8", + "sha256": "cfb9f687b572209e8e604063a450eab9e310e75efea06eea7baceb0a26b2bd5d", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_element": { + "dtype": "int64", + "sha256": "933c7e8d6fe684120a67b7e31f85dfbf8fef55ffa0714975eef4327f43661f7a", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_pos": { + "dtype": "float32", + "sha256": "f62bc0760285da89fba17299242b3065b453dfc43e399b9c91668cd25626b42b", + "shape": [ + 1, + 448, + 3 + ] + }, + "feature__ref_space_uid": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__res_type": { + "dtype": "int64", + "sha256": "2b70c203e4364aab90ede5af0f3f98f9ec96e9f9baca1b37600b9651698d2b94", + "shape": [ + 1, + 56 + ] + }, + "feature__residue_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "feature__sym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__token_attention_mask": { + "dtype": "bool", + "sha256": "8670534efd246dcad8a15296b8b8d3093a42bbfd6180da7331fc591f940cf778", + "shape": [ + 1, + 56 + ] + }, + "feature__token_bonds": { + "dtype": "float32", + "sha256": "c03665bfd44118d16c7e63946c56dab23f28a9f82e7fb9899204558a7396eba7", + "shape": [ + 1, + 56, + 56, + 1 + ] + }, + "feature__token_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "noise__initial_standard_normal": { + "dtype": "float32", + "sha256": "7468cdc59258e5ddf106b6e2cceeef4ffa4d984bc453ea52b452e1a61aeb326e", + "shape": [ + 1, + 448, + 3 + ] + }, + "output__atom_pad_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "output__distogram_logits": { + "dtype": "float32", + "sha256": "77e2c26b6be0f19b7475ea719ecae204a8d23867062faa74c797b0477f904c7b", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__iptm": { + "dtype": "float32", + "sha256": "76e40d7132a867f76cdb3d69b3c6e7716f39c243818fefe1a5924ed437630a99", + "shape": [ + 1 + ] + }, + "output__pae": { + "dtype": "float32", + "sha256": "b2e3e53d2826852f9007ba30624d2ac8de9ee62e4a7f78cdf21c23cf4ebc406e", + "shape": [ + 1, + 56, + 56 + ] + }, + "output__pae_logits": { + "dtype": "float32", + "sha256": "8934a711a31ae98f3f2220a047f43f2cd517280c4feff6d488f0f7064279d8a7", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__pde_logits": { + "dtype": "float32", + "sha256": "7b9c807d2ba3439c4f5d5624bc1db51f15ea0a52db5ded336a8f4cd36a8b3995", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__plddt": { + "dtype": "float32", + "sha256": "73ea5c74841dc1790c9b5308f091906b493829144f44b882138ec64829d6c035", + "shape": [ + 1, + 56 + ] + }, + "output__plddt_logits": { + "dtype": "float32", + "sha256": "cd52df7b2370be2d88cc4ae891abe60768db381b4a816da02d9e321c0c75e10c", + "shape": [ + 1, + 448, + 50 + ] + }, + "output__ptm": { + "dtype": "float32", + "sha256": "7e59dc5485e9632f6f2cdd6bdf51ef268e4908c107fb4d07db24e9e79386e121", + "shape": [ + 1 + ] + }, + "output__sample_atom_coords": { + "dtype": "float32", + "sha256": "3efc43d9b110b93c55a3290b424ba5cf46b4723481a165b682a3c5123a27f2ca", + "shape": [ + 1, + 448, + 3 + ] + } + } +} diff --git a/tests/goldens/esmfold2.safetensors b/tests/goldens/esmfold2.safetensors new file mode 100644 index 0000000..21d8fec Binary files /dev/null and b/tests/goldens/esmfold2.safetensors differ diff --git a/tests/goldens/esmfold2_experimental_cutoff2025.json b/tests/goldens/esmfold2_experimental_cutoff2025.json new file mode 100644 index 0000000..5c6693c --- /dev/null +++ b/tests/goldens/esmfold2_experimental_cutoff2025.json @@ -0,0 +1,340 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:79ed0dc0f867b8f09bfa004d6f77397c2ab9b38d", + "model.safetensors": "sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3" + }, + "repo_id": "biohub/ESMFold2-Experimental-Cutoff2025", + "revision": "56f94f5c1069ecde17512c96928850518340d287" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "python": "3.12.3", + "torch": "2.13.0+cu130", + "transformer_engine": "null", + "transformers": "4.57.6" + }, + "fingerprint": "aebef9a5edb6b5694d3ec547d636998a94232000ee91639146e9c4acf90d7589" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmfold2_experimental_cutoff2025" + ], + "input_fingerprint": "5f428f283d460f924ba72951dc5896266d46c253c9b941a83fd7984422a19601", + "model_id": "esmfold2_experimental_cutoff2025", + "schema_version": 1, + "source_files": { + "native/bundle.safetensors": "364828dfcc1f748c4006dbeff48f148d5433d1d22167f1020027cc11a659bc34", + "native/metadata.json": "7374e66b69ad96fda3018d94f02348a6c665353ad59d74ae8ae7679892e1b674" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + }, + { + "id": "protein-ttt", + "revision": "fde2817cd84b936167cc76ccabf31e5c0fe49962", + "url": "https://github.com/anton-bushuiev/ProteinTTT.git" + } + ], + "tensor_file": { + "path": "esmfold2_experimental_cutoff2025.safetensors", + "sha256": "9347466bbe803b6f5dc82e3356ca6cbbf2c2edd8765f9fd273385bda255019f6" + }, + "tensors": { + "feature__asym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__atom_attention_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "feature__atom_to_token": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__deletion_mean": { + "dtype": "float32", + "sha256": "8c51f4825acbf8ca40bd30065492cb000a9489b8245c9447a4e355f1b41b3dbb", + "shape": [ + 1, + 56 + ] + }, + "feature__deletion_value": { + "dtype": "float32", + "sha256": "34688aedb3460cee4d864ac60a112db764823880e1775d653f842a9015c19a92", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__distogram_atom_idx": { + "dtype": "int64", + "sha256": "56d98551f23b897604f00522013b97b939a25c593ac66492ed70bccb4461f663", + "shape": [ + 1, + 56 + ] + }, + "feature__entity_id": { + "dtype": "int64", + "sha256": "568a9a29b407045c55e7793effeaa4b30b336d3e9d02691dda02290107230a1c", + "shape": [ + 1, + 56 + ] + }, + "feature__has_deletion": { + "dtype": "bool", + "sha256": "fde502e113e1d0b810151796abfb4616d3c935a9e22f7d1891e86fdb8c971f6c", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__input_ids": { + "dtype": "int64", + "sha256": "c66b63367006340c1e676e0e2c4440f349286b98786ddc1cceb6f0222b599b95", + "shape": [ + 1, + 56 + ] + }, + "feature__mol_type": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__msa": { + "dtype": "int64", + "sha256": "14efcbac324a4ec831a0af4ec220ef00a57c98eea17525eb16a8ed22e9dde34b", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__msa_attention_mask": { + "dtype": "bool", + "sha256": "fa513f2a8c37868e06a9fbb53ac5f6dd9eb28106680113748e9ff868cfa259d5", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__ref_atom_name_chars": { + "dtype": "int64", + "sha256": "70165370912f9420bfb1a8ebe5b5ba19bf88cd6fbd43395aaaab70189ac3b2ec", + "shape": [ + 1, + 448, + 4 + ] + }, + "feature__ref_charge": { + "dtype": "int8", + "sha256": "cfb9f687b572209e8e604063a450eab9e310e75efea06eea7baceb0a26b2bd5d", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_element": { + "dtype": "int64", + "sha256": "933c7e8d6fe684120a67b7e31f85dfbf8fef55ffa0714975eef4327f43661f7a", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_pos": { + "dtype": "float32", + "sha256": "f62bc0760285da89fba17299242b3065b453dfc43e399b9c91668cd25626b42b", + "shape": [ + 1, + 448, + 3 + ] + }, + "feature__ref_space_uid": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__res_type": { + "dtype": "int64", + "sha256": "2b70c203e4364aab90ede5af0f3f98f9ec96e9f9baca1b37600b9651698d2b94", + "shape": [ + 1, + 56 + ] + }, + "feature__residue_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "feature__sym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__token_attention_mask": { + "dtype": "bool", + "sha256": "8670534efd246dcad8a15296b8b8d3093a42bbfd6180da7331fc591f940cf778", + "shape": [ + 1, + 56 + ] + }, + "feature__token_bonds": { + "dtype": "float32", + "sha256": "c03665bfd44118d16c7e63946c56dab23f28a9f82e7fb9899204558a7396eba7", + "shape": [ + 1, + 56, + 56, + 1 + ] + }, + "feature__token_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "noise__initial_standard_normal": { + "dtype": "float32", + "sha256": "e3a8639c0782543df8c37ec90776ec1e102ec88ef85c6626e765db4f1a41ae31", + "shape": [ + 1, + 448, + 3 + ] + }, + "output__atom_pad_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "output__distogram_logits": { + "dtype": "bfloat16", + "sha256": "a02399fb080d9abe035d1c5f338b7b5e879e3d6e3b46fcd8923ae375f38058ff", + "shape": [ + 1, + 56, + 56, + 128 + ] + }, + "output__iptm": { + "dtype": "float32", + "sha256": "76e40d7132a867f76cdb3d69b3c6e7716f39c243818fefe1a5924ed437630a99", + "shape": [ + 1 + ] + }, + "output__pae": { + "dtype": "float32", + "sha256": "ec41ef450cbf0b5d2d678f2d54982cda3c83c39e4febb5044336fa79d7c4fe4a", + "shape": [ + 1, + 56, + 56 + ] + }, + "output__pae_logits": { + "dtype": "float32", + "sha256": "21bdafbbe99ce1f66d71c88e2d927dedb0d0a913ba6fc9d32ff49b6e9845d29b", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__plddt": { + "dtype": "float32", + "sha256": "63bf8b3d7014d37d8e29d0341b393ee71a9604abd3ef36c1f05152a0c028b9ee", + "shape": [ + 1, + 56 + ] + }, + "output__plddt_logits": { + "dtype": "float32", + "sha256": "3b74a1d91e5af70c6a717663c98946ff36f9a8acead225d02aec904e00aa37ee", + "shape": [ + 1, + 448, + 50 + ] + }, + "output__ptm": { + "dtype": "float32", + "sha256": "faa0929e2da1b252d05b4da09d8b0392fbdbad2012976f7d338cd6793fe9fb2a", + "shape": [ + 1 + ] + }, + "output__sample_atom_coords": { + "dtype": "float32", + "sha256": "cf7dc3fa75a6a9fa6c9ff336fb82182bb753271db0bae5be67e3c4f83d4ccabe", + "shape": [ + 1, + 448, + 3 + ] + } + } +} diff --git a/tests/goldens/esmfold2_experimental_cutoff2025.safetensors b/tests/goldens/esmfold2_experimental_cutoff2025.safetensors new file mode 100644 index 0000000..31be661 Binary files /dev/null and b/tests/goldens/esmfold2_experimental_cutoff2025.safetensors differ diff --git a/tests/goldens/esmfold2_experimental_fast_cutoff2025.json b/tests/goldens/esmfold2_experimental_fast_cutoff2025.json new file mode 100644 index 0000000..7eee16e --- /dev/null +++ b/tests/goldens/esmfold2_experimental_fast_cutoff2025.json @@ -0,0 +1,340 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:0333d68ddb12ed2f066741dcb801142f466c0a2c", + "model.safetensors": "sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f" + }, + "repo_id": "biohub/ESMFold2-Experimental-Fast-Cutoff2025", + "revision": "74b88548bf19688b8727432db0d698cb2e1d8783" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "python": "3.12.3", + "torch": "2.13.0+cu130", + "transformer_engine": "null", + "transformers": "4.57.6" + }, + "fingerprint": "aebef9a5edb6b5694d3ec547d636998a94232000ee91639146e9c4acf90d7589" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmfold2_experimental_fast_cutoff2025" + ], + "input_fingerprint": "ee1cc56916eb3f9588a1eb4e876bd3e821ee40eb665e3f49fd1cfd1eb271cb39", + "model_id": "esmfold2_experimental_fast_cutoff2025", + "schema_version": 1, + "source_files": { + "native/bundle.safetensors": "2f92c58d9f8d3edd17efe70d35e0a5d21fdcba018bb33687dd0eaf8b611189f6", + "native/metadata.json": "6cea5dd02e98d60d94091380e9dfc4ded513e48148523c73ee1bf0a5101153b2" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + }, + { + "id": "protein-ttt", + "revision": "fde2817cd84b936167cc76ccabf31e5c0fe49962", + "url": "https://github.com/anton-bushuiev/ProteinTTT.git" + } + ], + "tensor_file": { + "path": "esmfold2_experimental_fast_cutoff2025.safetensors", + "sha256": "516e216d05d7e6bee59e77126d3e595e2bb7821929433f00c259c5d5241964bb" + }, + "tensors": { + "feature__asym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__atom_attention_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "feature__atom_to_token": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__deletion_mean": { + "dtype": "float32", + "sha256": "8c51f4825acbf8ca40bd30065492cb000a9489b8245c9447a4e355f1b41b3dbb", + "shape": [ + 1, + 56 + ] + }, + "feature__deletion_value": { + "dtype": "float32", + "sha256": "34688aedb3460cee4d864ac60a112db764823880e1775d653f842a9015c19a92", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__distogram_atom_idx": { + "dtype": "int64", + "sha256": "56d98551f23b897604f00522013b97b939a25c593ac66492ed70bccb4461f663", + "shape": [ + 1, + 56 + ] + }, + "feature__entity_id": { + "dtype": "int64", + "sha256": "568a9a29b407045c55e7793effeaa4b30b336d3e9d02691dda02290107230a1c", + "shape": [ + 1, + 56 + ] + }, + "feature__has_deletion": { + "dtype": "bool", + "sha256": "fde502e113e1d0b810151796abfb4616d3c935a9e22f7d1891e86fdb8c971f6c", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__input_ids": { + "dtype": "int64", + "sha256": "c66b63367006340c1e676e0e2c4440f349286b98786ddc1cceb6f0222b599b95", + "shape": [ + 1, + 56 + ] + }, + "feature__mol_type": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__msa": { + "dtype": "int64", + "sha256": "14efcbac324a4ec831a0af4ec220ef00a57c98eea17525eb16a8ed22e9dde34b", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__msa_attention_mask": { + "dtype": "bool", + "sha256": "fa513f2a8c37868e06a9fbb53ac5f6dd9eb28106680113748e9ff868cfa259d5", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__ref_atom_name_chars": { + "dtype": "int64", + "sha256": "70165370912f9420bfb1a8ebe5b5ba19bf88cd6fbd43395aaaab70189ac3b2ec", + "shape": [ + 1, + 448, + 4 + ] + }, + "feature__ref_charge": { + "dtype": "int8", + "sha256": "cfb9f687b572209e8e604063a450eab9e310e75efea06eea7baceb0a26b2bd5d", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_element": { + "dtype": "int64", + "sha256": "933c7e8d6fe684120a67b7e31f85dfbf8fef55ffa0714975eef4327f43661f7a", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_pos": { + "dtype": "float32", + "sha256": "f62bc0760285da89fba17299242b3065b453dfc43e399b9c91668cd25626b42b", + "shape": [ + 1, + 448, + 3 + ] + }, + "feature__ref_space_uid": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__res_type": { + "dtype": "int64", + "sha256": "2b70c203e4364aab90ede5af0f3f98f9ec96e9f9baca1b37600b9651698d2b94", + "shape": [ + 1, + 56 + ] + }, + "feature__residue_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "feature__sym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__token_attention_mask": { + "dtype": "bool", + "sha256": "8670534efd246dcad8a15296b8b8d3093a42bbfd6180da7331fc591f940cf778", + "shape": [ + 1, + 56 + ] + }, + "feature__token_bonds": { + "dtype": "float32", + "sha256": "c03665bfd44118d16c7e63946c56dab23f28a9f82e7fb9899204558a7396eba7", + "shape": [ + 1, + 56, + 56, + 1 + ] + }, + "feature__token_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "noise__initial_standard_normal": { + "dtype": "float32", + "sha256": "e3a8639c0782543df8c37ec90776ec1e102ec88ef85c6626e765db4f1a41ae31", + "shape": [ + 1, + 448, + 3 + ] + }, + "output__atom_pad_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "output__distogram_logits": { + "dtype": "bfloat16", + "sha256": "a080b4c7781aa784a54958e8b013625548b91689acd87b22e412ae7e6a372c87", + "shape": [ + 1, + 56, + 56, + 128 + ] + }, + "output__iptm": { + "dtype": "float32", + "sha256": "76e40d7132a867f76cdb3d69b3c6e7716f39c243818fefe1a5924ed437630a99", + "shape": [ + 1 + ] + }, + "output__pae": { + "dtype": "float32", + "sha256": "732bd567e0ca0662cb2039f00768fd5876a7025d6d9feb7ceea1a9135d3030c1", + "shape": [ + 1, + 56, + 56 + ] + }, + "output__pae_logits": { + "dtype": "float32", + "sha256": "501f74b2673f4996ae2c7ad23297baed7063b7f78a4f92b0074802f305335b4f", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__plddt": { + "dtype": "float32", + "sha256": "5f422710b3e758d0f0cd5c68afc6d2aef516c07ff86fcae4f42a9b31a42fe936", + "shape": [ + 1, + 56 + ] + }, + "output__plddt_logits": { + "dtype": "float32", + "sha256": "7c7412861338d6d4878cba6cc02574efd635a733cd560baa60a51bf06fc9b5d4", + "shape": [ + 1, + 448, + 50 + ] + }, + "output__ptm": { + "dtype": "float32", + "sha256": "4ad6691be1c4d83df5acac643b3167a016e2d401346705d951a7a7ff7b6146c8", + "shape": [ + 1 + ] + }, + "output__sample_atom_coords": { + "dtype": "float32", + "sha256": "2afd9585a9f90e24ae32b88de4ec058577d7900b63b2883cf5208d9eb4712623", + "shape": [ + 1, + 448, + 3 + ] + } + } +} diff --git a/tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors b/tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors new file mode 100644 index 0000000..78930ac Binary files /dev/null and b/tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors differ diff --git a/tests/goldens/esmfold2_fast.json b/tests/goldens/esmfold2_fast.json new file mode 100644 index 0000000..6e2f6fc --- /dev/null +++ b/tests/goldens/esmfold2_fast.json @@ -0,0 +1,350 @@ +{ + "checkpoint": { + "files": { + "config.json": "git-sha1:c0ca526090fa7f8342ee4666d56e7fe3a4b8cbb2", + "model.safetensors": "sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a" + }, + "repo_id": "biohub/ESMFold2-Fast", + "revision": "b28d8ace5e05e61e5bec1e6820cfd3e221819d12" + }, + "environment": { + "details": { + "cuda_device": "NVIDIA H100 PCIe", + "cuda_runtime": "13.0", + "python": "3.12.3", + "torch": "2.13.0+cu130", + "transformer_engine": "null", + "transformers": "4.57.6" + }, + "fingerprint": "aebef9a5edb6b5694d3ec547d636998a94232000ee91639146e9c4acf90d7589" + }, + "generation_command": [ + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + "esmfold2_fast" + ], + "input_fingerprint": "44275d8ad675fcad7094b1df855a42be7ff2e284f1d50deb5a1e711d9926f670", + "model_id": "esmfold2_fast", + "schema_version": 1, + "source_files": { + "native/bundle.safetensors": "19e8f0b9c06c4ca70a164ac798bf3b22262f254810b4ec2069268ee891862012", + "native/metadata.json": "2454702767117f574b28fa186504ec0b1f9f5b8ce8f120aaa64a8a3ca0484b9c" + }, + "sources": [ + { + "id": "biohub-esm", + "revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "url": "https://github.com/Biohub/esm.git" + }, + { + "id": "biohub-transformers", + "revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "url": "https://github.com/Biohub/transformers.git" + }, + { + "id": "protein-ttt", + "revision": "fde2817cd84b936167cc76ccabf31e5c0fe49962", + "url": "https://github.com/anton-bushuiev/ProteinTTT.git" + } + ], + "tensor_file": { + "path": "esmfold2_fast.safetensors", + "sha256": "6e2e1cd07401538b4d9df994f82abe7a5b38a01e8d1ee26681e1216d44a81990" + }, + "tensors": { + "feature__asym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__atom_attention_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "feature__atom_to_token": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__deletion_mean": { + "dtype": "float32", + "sha256": "8c51f4825acbf8ca40bd30065492cb000a9489b8245c9447a4e355f1b41b3dbb", + "shape": [ + 1, + 56 + ] + }, + "feature__deletion_value": { + "dtype": "float32", + "sha256": "34688aedb3460cee4d864ac60a112db764823880e1775d653f842a9015c19a92", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__distogram_atom_idx": { + "dtype": "int64", + "sha256": "56d98551f23b897604f00522013b97b939a25c593ac66492ed70bccb4461f663", + "shape": [ + 1, + 56 + ] + }, + "feature__entity_id": { + "dtype": "int64", + "sha256": "568a9a29b407045c55e7793effeaa4b30b336d3e9d02691dda02290107230a1c", + "shape": [ + 1, + 56 + ] + }, + "feature__has_deletion": { + "dtype": "bool", + "sha256": "fde502e113e1d0b810151796abfb4616d3c935a9e22f7d1891e86fdb8c971f6c", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__input_ids": { + "dtype": "int64", + "sha256": "c66b63367006340c1e676e0e2c4440f349286b98786ddc1cceb6f0222b599b95", + "shape": [ + 1, + 56 + ] + }, + "feature__mol_type": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__msa": { + "dtype": "int64", + "sha256": "14efcbac324a4ec831a0af4ec220ef00a57c98eea17525eb16a8ed22e9dde34b", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__msa_attention_mask": { + "dtype": "bool", + "sha256": "fa513f2a8c37868e06a9fbb53ac5f6dd9eb28106680113748e9ff868cfa259d5", + "shape": [ + 1, + 1, + 56 + ] + }, + "feature__ref_atom_name_chars": { + "dtype": "int64", + "sha256": "70165370912f9420bfb1a8ebe5b5ba19bf88cd6fbd43395aaaab70189ac3b2ec", + "shape": [ + 1, + 448, + 4 + ] + }, + "feature__ref_charge": { + "dtype": "int8", + "sha256": "cfb9f687b572209e8e604063a450eab9e310e75efea06eea7baceb0a26b2bd5d", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_element": { + "dtype": "int64", + "sha256": "933c7e8d6fe684120a67b7e31f85dfbf8fef55ffa0714975eef4327f43661f7a", + "shape": [ + 1, + 448 + ] + }, + "feature__ref_pos": { + "dtype": "float32", + "sha256": "f62bc0760285da89fba17299242b3065b453dfc43e399b9c91668cd25626b42b", + "shape": [ + 1, + 448, + 3 + ] + }, + "feature__ref_space_uid": { + "dtype": "int64", + "sha256": "a94836d05854aeeae8bd1a84275187ae311164e4a0cc448777b10b66d22a2e65", + "shape": [ + 1, + 448 + ] + }, + "feature__res_type": { + "dtype": "int64", + "sha256": "2b70c203e4364aab90ede5af0f3f98f9ec96e9f9baca1b37600b9651698d2b94", + "shape": [ + 1, + 56 + ] + }, + "feature__residue_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "feature__sym_id": { + "dtype": "int64", + "sha256": "4076ba2ce2248c4a6e16ad1c70d4a4802b013abde7acab3d910da9b5c1846c02", + "shape": [ + 1, + 56 + ] + }, + "feature__token_attention_mask": { + "dtype": "bool", + "sha256": "8670534efd246dcad8a15296b8b8d3093a42bbfd6180da7331fc591f940cf778", + "shape": [ + 1, + 56 + ] + }, + "feature__token_bonds": { + "dtype": "float32", + "sha256": "c03665bfd44118d16c7e63946c56dab23f28a9f82e7fb9899204558a7396eba7", + "shape": [ + 1, + 56, + 56, + 1 + ] + }, + "feature__token_index": { + "dtype": "int64", + "sha256": "f346dfab37ce849e49d4f9ce562075efc8b707f13bc1e88a7f7c6d40a1589d8a", + "shape": [ + 1, + 56 + ] + }, + "noise__initial_standard_normal": { + "dtype": "float32", + "sha256": "7468cdc59258e5ddf106b6e2cceeef4ffa4d984bc453ea52b452e1a61aeb326e", + "shape": [ + 1, + 448, + 3 + ] + }, + "output__atom_pad_mask": { + "dtype": "bool", + "sha256": "311674e8928ea7d4c571757d2f7c9a3ebb17ea71fe840675786ff79ca855d0c5", + "shape": [ + 1, + 448 + ] + }, + "output__distogram_logits": { + "dtype": "float32", + "sha256": "06c7390a85f0f9776bd3f577215a6c9007e2f8b7fc7193e45c0e8569d0f3143b", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__iptm": { + "dtype": "float32", + "sha256": "76e40d7132a867f76cdb3d69b3c6e7716f39c243818fefe1a5924ed437630a99", + "shape": [ + 1 + ] + }, + "output__pae": { + "dtype": "float32", + "sha256": "f3527f8e1aaf68ad187c4ac01749e4facc333da25e9c58cd995fa74d6f831526", + "shape": [ + 1, + 56, + 56 + ] + }, + "output__pae_logits": { + "dtype": "float32", + "sha256": "15d93267ab8e40848dca83291b7d83801ec0e4d2fd390fd54b496f7ec17c2b00", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__pde_logits": { + "dtype": "float32", + "sha256": "6e4a56f94256b1be9353b66edda312fce97b6107db12aa608523680fa31bd006", + "shape": [ + 1, + 56, + 56, + 64 + ] + }, + "output__plddt": { + "dtype": "float32", + "sha256": "e6c1b067c41c9ab6913f9c94757626a1e7d0fbf94c7dacfcc7ac5be7c492ec8c", + "shape": [ + 1, + 56 + ] + }, + "output__plddt_logits": { + "dtype": "float32", + "sha256": "45250ca9436522d7993e209275311c46023e76cecd66e9f0d669a2e2b0f25ecc", + "shape": [ + 1, + 448, + 50 + ] + }, + "output__ptm": { + "dtype": "float32", + "sha256": "917d19eea5d09f0c37791f81687e7cf31921dd6c54fa1cc512fb366078d18811", + "shape": [ + 1 + ] + }, + "output__sample_atom_coords": { + "dtype": "float32", + "sha256": "474aa574d06e58766e4ed468c2bc3804d312052938274ba3d7559c31fa00cfb8", + "shape": [ + 1, + 448, + 3 + ] + } + } +} diff --git a/tests/goldens/esmfold2_fast.safetensors b/tests/goldens/esmfold2_fast.safetensors new file mode 100644 index 0000000..c9d487f Binary files /dev/null and b/tests/goldens/esmfold2_fast.safetensors differ diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..d8c1211 --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1 @@ +"""Local-package integration tests.""" diff --git a/tests/integration/support/__init__.py b/tests/integration/support/__init__.py new file mode 100644 index 0000000..7319d4b --- /dev/null +++ b/tests/integration/support/__init__.py @@ -0,0 +1 @@ +"""Standalone integration utilities that pytest does not collect.""" diff --git a/tests/integration/test_ankh_tokenizer_contract.py b/tests/integration/test_ankh_tokenizer_contract.py new file mode 100644 index 0000000..42e2c16 --- /dev/null +++ b/tests/integration/test_ankh_tokenizer_contract.py @@ -0,0 +1,215 @@ +"""Pinned real-tokenizer contracts for ANKH source, decoder, and TTT paths.""" + +from __future__ import annotations + +import hashlib +import os +import pytest +import torch +from pathlib import Path +from transformers import AutoTokenizer + +from examples.ankh_embeddings import generate_ankh_task +from fastplms.models.ankh.modeling_ankh import ( + FastAnkhConfig, + FastAnkhForConditionalGeneration, + FastAnkhForMaskedLMExtension, + FastAnkhModel, + tokenize_ankh_decoder_prompts, + tokenize_ankh_sequences, +) +from fastplms.registry import FileDigest, get_model_registry + + +pytestmark = [pytest.mark.compliance, pytest.mark.network, pytest.mark.reference] +_SNAPSHOT_ENVIRONMENT = "FASTPLMS_ANKH_TOKENIZER_SNAPSHOT" +_SNAPSHOT_ROOT_ENVIRONMENT = "FASTPLMS_ANKH_TOKENIZER_SNAPSHOT_ROOT" +_ANKH_SPEC_IDS = tuple(spec.id for spec in get_model_registry().by_family("ankh")) +_TOKENIZER_PATHS = ( + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", +) + + +def _file_digest(path: Path, identity: FileDigest) -> str: + payload = path.read_bytes() + if identity.algorithm == "sha256": + return hashlib.sha256(payload).hexdigest() + if identity.algorithm == "git-sha1": + header = f"blob {len(payload)}\0".encode() + return hashlib.sha1(header + payload, usedforsecurity=False).hexdigest() + raise AssertionError(f"Unsupported manifest digest: {identity.algorithm}") + + +@pytest.fixture(scope="module") +def pinned_ankh_tokenizer_snapshot(request: pytest.FixtureRequest) -> tuple[Path, str]: + spec_id = getattr(request, "param", "ankh_base") + spec = get_model_registry()[spec_id] + configured = os.environ.get(_SNAPSHOT_ENVIRONMENT) + configured_root = os.environ.get(_SNAPSHOT_ROOT_ENVIRONMENT) + if configured_root: + snapshot = Path(configured_root).expanduser().resolve() / spec.id + elif configured and spec.id == "ankh_base": + snapshot = Path(configured).expanduser().resolve() + else: + from huggingface_hub import snapshot_download + + snapshot = Path( + snapshot_download( + spec.official.repo_id, + revision=spec.official.revision, + allow_patterns=list(_TOKENIZER_PATHS), + ) + ) + assert snapshot.is_dir(), f"Pinned ANKH tokenizer snapshot is missing: {snapshot}" + for relative_path in _TOKENIZER_PATHS: + identity = spec.official.file_map[relative_path] + path = snapshot / relative_path + assert path.is_file(), f"Pinned ANKH tokenizer asset is missing: {path}" + assert _file_digest(path, identity) == identity.digest, ( + f"Pinned ANKH tokenizer asset does not match the manifest: {relative_path}" + ) + return snapshot, spec.official.revision + + +def _tiny_config(snapshot: Path, revision: str, vocab_size: int) -> FastAnkhConfig: + config = FastAnkhConfig( + vocab_size=vocab_size, + d_model=16, + d_kv=8, + d_ff=32, + num_heads=2, + num_layers=1, + num_decoder_layers=1, + dropout_rate=0.0, + pad_token_id=0, + eos_token_id=1, + decoder_start_token_id=0, + attn_backend="eager", + use_cache=True, + ) + config._name_or_path = str(snapshot) + config._commit_hash = revision + return config + + +@pytest.mark.parametrize( + "pinned_ankh_tokenizer_snapshot", + _ANKH_SPEC_IDS, + indirect=True, +) +def test_every_real_ankh_tokenizer_preserves_raw_residues_and_tight_sentinels( + pinned_ankh_tokenizer_snapshot: tuple[Path, str], +) -> None: + snapshot, _ = pinned_ankh_tokenizer_snapshot + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + + raw = tokenize_ankh_sequences( + tokenizer, + "MSTNPK", + return_tensors="pt", + add_special_tokens=False, + )["input_ids"] + legacy_spaced = tokenize_ankh_sequences( + tokenizer, + "M S T N P K", + return_tensors="pt", + add_special_tokens=False, + )["input_ids"] + tight_prompt = tokenize_ankh_decoder_prompts( + tokenizer, + "M", + return_tensors="pt", + add_special_tokens=False, + )["input_ids"] + spaced_prompt = tokenize_ankh_decoder_prompts( + tokenizer, + "M ", + return_tensors="pt", + add_special_tokens=False, + )["input_ids"] + + assert raw.shape == (1, 6) + assert tokenizer.unk_token_id not in raw + assert torch.equal(raw, legacy_spaced) + assert torch.equal(tight_prompt, spaced_prompt) + assert tokenizer.unk_token_id not in tight_prompt + assert tokenizer.convert_tokens_to_ids("") in tight_prompt + + +def test_real_explicit_tokenizer_matches_model_scoped_tokenizer( + pinned_ankh_tokenizer_snapshot: tuple[Path, str], +) -> None: + snapshot, revision = pinned_ankh_tokenizer_snapshot + explicit = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + model = FastAnkhModel(_tiny_config(snapshot, revision, len(explicit))).eval() + model.__dict__["_fastplms_tokenizer_load_context"] = {"local_files_only": True} + + explicit_ids = tokenize_ankh_sequences( + explicit, + ["MSTNPK", "ACDE"], + return_tensors="pt", + padding=True, + )["input_ids"] + scoped_ids = tokenize_ankh_sequences( + model.tokenizer, + ["MSTNPK", "ACDE"], + return_tensors="pt", + padding=True, + )["input_ids"] + + assert torch.equal(explicit_ids, scoped_ids) + assert explicit.unk_token_id not in explicit_ids + assert model.tokenizer.unk_token_id not in scoped_ids + assert explicit.backend_tokenizer.to_str() == model.tokenizer.backend_tokenizer.to_str() + + +def test_real_sentinel_decoder_extraction_and_generation( + pinned_ankh_tokenizer_snapshot: tuple[Path, str], +) -> None: + snapshot, revision = pinned_ankh_tokenizer_snapshot + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + model = FastAnkhForConditionalGeneration( + _tiny_config(snapshot, revision, len(tokenizer)) + ).eval() + model.tokenizer = tokenizer + + batch = model._embedding_batch( + ["MSTNPK"], + tokenizer=tokenizer, + hidden_state_source="decoder", + decoder_inputs=["M "], + ) + generated = generate_ankh_task( + model, + tokenizer, + "MSTNPK", + "M ", + max_new_tokens=1, + ) + + assert batch.residue_mask.sum().item() == 1 + assert generated.ndim == 2 + assert generated.shape[0] == 1 + assert tokenizer.unk_token_id not in generated[:, :3] + + +def test_real_ankh_ttt_uses_the_shared_raw_sequence_contract( + pinned_ankh_tokenizer_snapshot: tuple[Path, str], +) -> None: + snapshot, revision = pinned_ankh_tokenizer_snapshot + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + model = FastAnkhForMaskedLMExtension(_tiny_config(snapshot, revision, len(tokenizer))).eval() + model.tokenizer = tokenizer + + ttt_ids = model._ttt_tokenize(seq="MSTNPK") + shared_ids = tokenize_ankh_sequences( + tokenizer, + ["MSTNPK"], + return_tensors="pt", + padding=True, + )["input_ids"] + + assert torch.equal(ttt_ids, shared_ids) + assert tokenizer.unk_token_id not in ttt_ids diff --git a/tests/integration/test_backend_consistency.py b/tests/integration/test_backend_consistency.py new file mode 100644 index 0000000..e6d6f44 --- /dev/null +++ b/tests/integration/test_backend_consistency.py @@ -0,0 +1,247 @@ +"""Global BF16 equivalence gates across every advertised attention backend.""" + +from __future__ import annotations + +import contextlib +import importlib +import random +import pytest +import torch +import torch.nn.functional as F +import transformers +from collections.abc import Sequence +from typing import Any + +from fastplms.registry import ModelSpec, get_model_registry +from tests.conftest import CANONICAL_AAS, SEED + + +REGISTRY = get_model_registry() +SEQUENCE_SPECS = tuple( + spec for spec in REGISTRY.values() if spec.family.tokenizer_mode != "structure" +) +NUM_SEQUENCES = 4 +SEQUENCE_LENGTH = 64 +GH200_MEASURED_BACKENDS = ("eager", "sdpa", "flex_attention") + + +def _parameter(spec: ModelSpec) -> Any: + marks: list[Any] = [pytest.mark.gpu] + if spec.size_category in {"large", "xlarge"}: + marks.append(pytest.mark.slow) + if spec.size_category == "xlarge": + marks.append(pytest.mark.large) + return pytest.param(spec, id=spec.id, marks=marks) + + +def _model_class(spec: ModelSpec) -> type[torch.nn.Module]: + """Resolve the repository source class declared by the model manifest. + + Backend consistency is a source contract. Remote-code loading is + covered separately by the artifact suite, where the generated artifact + contains the same source revision as the candidate checkout. + """ + + advertised = set(spec.auto_map) + if spec.family.id == "ankh": + name = "AutoModel" + elif "AutoModelForMaskedLM" in advertised: + name = "AutoModelForMaskedLM" + else: + name = "AutoModel" + assert name in advertised, f"{spec.id} does not advertise {name}" + qualified_name = spec.auto_map[name] + module_name, class_name = qualified_name.rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + assert issubclass(model_class, torch.nn.Module) + return model_class + + +def _sequences() -> list[str]: + generator = random.Random(SEED) + return [ + "M" + "".join(generator.choices(CANONICAL_AAS, k=SEQUENCE_LENGTH - 1)) + for _ in range(NUM_SEQUENCES) + ] + + +def _prepare_inputs( + spec: ModelSpec, + model: torch.nn.Module, + sequences: Sequence[str], + device: torch.device, +) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + if spec.family.tokenizer_mode == "sequence": + batch = model.model.prep_tokens.get_batch_kwargs(sequences, device=device) + inputs = { + "input_ids": batch["input_ids"], + "within_seq_position_ids": batch["within_seq_position_ids"], + "global_position_ids": batch["global_position_ids"], + "sequence_ids": batch["sequence_ids"], + "attention_mask": batch["sequence_ids"].ne(-1).long(), + } + return inputs, batch["sequence_ids"].ge(0) + + tokenizer = getattr(model, "tokenizer", None) + if tokenizer is None: + tokenizer = transformers.AutoTokenizer.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + trust_remote_code=True, + ) + tokenize_kwargs = { + "return_tensors": "pt", + "padding": True, + "truncation": True, + } + sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None) + if callable(sequence_tokenizer): + encoded = sequence_tokenizer( + list(sequences), + tokenizer=tokenizer, + **tokenize_kwargs, + ) + else: + encoded = tokenizer(list(sequences), **tokenize_kwargs) + inputs = {name: value.to(device) for name, value in encoded.items() if torch.is_tensor(value)} + # input_ids: (b, l) + input_ids = inputs["input_ids"] + # residue_mask: (b, l) + residue_mask = inputs["attention_mask"].bool() + for token_id in getattr(tokenizer, "all_special_ids", ()): + residue_mask &= input_ids.ne(token_id) + if spec.family.architecture == "ESMC": + # inputs['sequence_id']: (b, l) + inputs["sequence_id"] = inputs["attention_mask"].bool() + if getattr(model.config, "is_encoder_decoder", False): + inputs["decoder_input_ids"] = input_ids + # inputs['decoder_attention_mask']: (b, l) + inputs["decoder_attention_mask"] = inputs["attention_mask"] + return inputs, residue_mask + + +def _sequence_output(output: object) -> tuple[torch.Tensor, bool]: + for name in ("logits", "sequence_logits"): + value = getattr(output, name, None) + if torch.is_tensor(value): + return value, True + value = getattr(output, "last_hidden_state", None) + if torch.is_tensor(value): + return value, False + raise AssertionError("Advertised sequence model output omitted a residue tensor") + + +def _assert_global_bf16_contract( + candidate: torch.Tensor, + reference: torch.Tensor, + residue_mask: torch.Tensor, + context: str, + *, + has_logits: bool, +) -> None: + # candidate: (...), reference: (...), residue_mask: (b, l) + assert candidate.shape == reference.shape + assert candidate.ndim == 3 + # candidate_f: (...) + candidate_f = candidate.float() + # reference_f: (...) + reference_f = reference.float() + valid_candidate = candidate_f[residue_mask] + valid_reference = reference_f[residue_mask] + difference = valid_candidate - valid_reference + tiny = torch.finfo(torch.float32).tiny + relative_l2 = torch.linalg.vector_norm(difference) / torch.linalg.vector_norm( + valid_reference + ).clamp_min(tiny) + relative_q999 = torch.quantile(difference.abs().reshape(-1), 0.999) / torch.quantile( + valid_reference.abs().reshape(-1), 0.999 + ).clamp_min(tiny) + # residue_cosine_p01: () + residue_cosine_p01 = torch.quantile( + F.cosine_similarity(valid_candidate, valid_reference, dim=-1), + 0.01, + ) + + # M: (...) + M = residue_mask.unsqueeze(-1).float() + candidate_pooled = (candidate_f * M).sum(1) / M.sum(1).clamp_min(1) + reference_pooled = (reference_f * M).sum(1) / M.sum(1).clamp_min(1) + pooled_cosine = F.cosine_similarity(candidate_pooled, reference_pooled, dim=-1) + + assert float(relative_l2) <= 1e-2, f"{context}: relative L2={relative_l2}" + assert float(relative_q999) <= 2.5e-2, f"{context}: relative Q99.9={relative_q999}" + assert float(residue_cosine_p01) >= 0.999, f"{context}: residue cosine p01={residue_cosine_p01}" + assert bool((pooled_cosine >= 0.9995).all()), ( + f"{context}: per-sequence pooled cosine={pooled_cosine.tolist()}" + ) + if has_logits: + # reference_probabilities: (...) + reference_probabilities = reference_f.softmax(-1) + # confidence: (...), reference_top1: (...) + confidence, reference_top1 = reference_probabilities.max(-1) + # confident_mask: (...) + confident_mask = residue_mask & confidence.ge(0.5) + assert bool(confident_mask.any()), f"{context}: no confident biological positions" + # candidate_top1: (...) + candidate_top1 = candidate_f.argmax(-1) + # top1_agreement: () + top1_agreement = ( + (candidate_top1[confident_mask] == reference_top1[confident_mask]).float().mean() + ) + assert float(top1_agreement) >= 0.995, ( + f"{context}: confident top-1 agreement={top1_agreement}" + ) + + +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_SPECS]) +def test_gh200_backends_meet_global_bf16_contract(spec: ModelSpec) -> None: + """Measure only the no-download GH200 eager, SDPA, and Flex matrix.""" + + device = torch.device("cuda") + model_class = _model_class(spec) + use_bf16_autocast = spec.family.bf16_execution == "fp32_parameters_autocast" + load_dtype = torch.float32 if use_bf16_autocast else torch.bfloat16 + model = model_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + dtype=load_dtype, + device_map=device, + ).eval() + inputs, residue_mask = _prepare_inputs(spec, model, _sequences(), device) + + outputs: dict[str, tuple[torch.Tensor, bool]] = {} + measured_backends = tuple( + backend for backend in spec.family.attention if backend in GH200_MEASURED_BACKENDS + ) + assert measured_backends, f"{spec.id}: no GH200 backend is declared" + for backend in measured_backends: + assert hasattr(model, "set_attn_implementation") + model.set_attn_implementation(backend) + resolved = getattr(model.config, "_attn_implementation", None) + if resolved is None: + resolved = getattr(model.config, "attn_implementation", None) + assert resolved == backend, f"{spec.id}: requested {backend}, resolved {resolved}" + numeric_context = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if use_bf16_autocast + else contextlib.nullcontext() + ) + with torch.inference_mode(), numeric_context: + output_tensor, has_logits = _sequence_output(model(**inputs)) + outputs[backend] = output_tensor.detach(), has_logits + + assert "sdpa" in outputs + reference, reference_has_logits = outputs["sdpa"] + for backend, (candidate, has_logits) in outputs.items(): + if backend != "sdpa": + assert has_logits is reference_has_logits + _assert_global_bf16_contract( + candidate, + reference, + residue_mask, + f"{spec.id}:sdpa-vs-{backend}", + has_logits=has_logits, + ) + + del model, outputs + torch.cuda.empty_cache() diff --git a/tests/integration/test_binder_design.py b/tests/integration/test_binder_design.py new file mode 100644 index 0000000..ae33c9b --- /dev/null +++ b/tests/integration/test_binder_design.py @@ -0,0 +1,730 @@ +"""Seeded feature smoke test for the FastPLMs binder-design workflow.""" + +from __future__ import annotations + +import json +import os +import random +import pytest +import torch +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +from examples import binder_design_fastplms as binder + + +APPROVED_CRITIC = "ESMFold2-Experimental-Fast-Cutoff2025" + + +class FakeInputBuilder: + def decode( + self, + output: dict[str, torch.Tensor], + inputs: dict[str, torch.Tensor], + chain_infos: list[Any], + num_diffusion_samples: int, + complex_id: str, + ) -> dict[str, Any]: + del output, inputs, chain_infos, num_diffusion_samples + return {"complex_id": complex_id} + + +@dataclass +class FakeCritic: + input_builder = FakeInputBuilder() + + def result_to_cif(self, complex_result: dict[str, Any]) -> str: + return f"data_{complex_result['complex_id']}\n" + + def result_to_pdb(self, complex_result: dict[str, Any]) -> str: + del complex_result + return "HEADER FASTPLMS TEST\nEND\n" + + +def _fake_fold( + model: Any, + target_seq: str, + target_one_hot: torch.Tensor, + design: torch.Tensor, + num_loops: int = 0, + num_sampling_steps: int = 1, + calculate_confidence: bool = False, + seed: int | None = None, +) -> dict[str, Any]: + # target_one_hot: (...), design: (...) + del model, num_loops, num_sampling_steps, calculate_confidence, seed + b, binder_length, d = design.shape + target_length = target_one_hot.size(1) + # aa_weight: (n,) + aa_weight = torch.linspace(-1.0, 1.0, d, device=design.device) + # binder_signal: (...) + binder_signal = (design * aa_weight).sum(dim=-1) + # token_signal: (...) + token_signal = torch.cat( + (torch.zeros(b, target_length, device=design.device), binder_signal), + dim=1, + ) + pair_signal = token_signal[:, :, None] + token_signal[:, None, :] + # bin_basis: (n,) + bin_basis = torch.linspace(-1.0, 1.0, 128, device=design.device) + # distogram_logits: (..., c) + distogram_logits = pair_signal.unsqueeze(-1) * bin_basis + sequences = [f"{target_seq}|{'A' * binder_length}" for _ in range(b)] + return { + "distogram_logits": distogram_logits, + "inputs": {}, + "chain_info_list": [[] for _ in range(b)], + "output": {"distogram_logits": distogram_logits}, + "seq_list": sequences, + "iptm": torch.ones(b, device=design.device), + "ptm": torch.ones(b, device=design.device), + "plddt": torch.ones(b, 1, device=design.device), + } + + +def _fake_pseudoperplexity( + lm_model: Any, + binder_design: torch.Tensor, + score_mask: torch.Tensor, + batch_size: int = 4, + n_passes: int = 4, + mask_fraction: float = binder.DEFAULT_ESMC_MASK_FRACTION, +) -> torch.Tensor: + # binder_design: (...), score_mask: (...) + del lm_model, score_mask, batch_size, n_passes, mask_fraction + return binder_design.square().mean(dim=(1, 2)) + + +def _run_seeded_workflow() -> tuple[ + list[str], dict[int, dict[str, torch.Tensor]], list[dict[str, Any]] +]: + return binder.design_binder( + inversion_models={APPROVED_CRITIC: FakeCritic()}, + critic_models={APPROVED_CRITIC: FakeCritic()}, + lm_model=object(), + target_name=None, + target_sequence="ACD", + binder_name=None, + binder_sequence="###", + is_antibody=False, + seed=17, + batch_size=1, + steps=1, + device="cpu", + ) + + +def test_public_binder_runner_propagates_loaded_device( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class LoadedModel: + def to(self, **kwargs: Any) -> LoadedModel: + del kwargs + return self + + def eval(self) -> LoadedModel: + return self + + def requires_grad_(self, value: bool) -> LoadedModel: + del value + return self + + monkeypatch.setattr( + binder, + "_load_fold_model", + lambda *args, **kwargs: LoadedModel(), + ) + monkeypatch.setattr( + binder.AutoModelForMaskedLM, + "from_pretrained", + lambda *args, **kwargs: LoadedModel(), + ) + observed: dict[str, Any] = {} + + def fake_design_binder(*args: Any, **kwargs: Any) -> tuple[list, dict, list]: + del args + observed.update(kwargs) + return [], {}, [] + + monkeypatch.setattr(binder, "design_binder", fake_design_binder) + runner = binder.FastPLMsBinderDesign() + runner.load(device="cpu") + runner.design() + + assert runner.device == torch.device("cpu") + assert observed["device"] == torch.device("cpu") + + +def test_fold_loader_pins_revision_and_propagates_local_files_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, Any] = {} + + class LoadedModel: + config = SimpleNamespace(esmc_id="Synthyra/ESMplusplus_6B") + _esmc = None + _esmc_fp8 = False + _esmc_fp8_module_paths: tuple[str, ...] = () + _esmc_source = "Synthyra/ESMplusplus_6B" + _esmc_source_revision = "b" * 40 + _esmc_local_files_only = False + _esmc_precision_policy = "auto" + _esmc_precision_status = SimpleNamespace(resolved="bf16") + + def __init__(self) -> None: + self._esmc_source_files: dict[str, str] = {} + + def load_esmc(self, source: str, **kwargs: Any) -> None: + observed["esmc_load"] = (source, kwargs) + self._esmc = object() + self._esmc_local_files_only = bool(kwargs.get("local_files_only", False)) + + def configure_lm_dropout(self, *args: Any, **kwargs: Any) -> None: + observed["dropout"] = (args, kwargs) + + def to(self, **kwargs: Any) -> LoadedModel: + observed["to"] = kwargs + return self + + def eval(self) -> LoadedModel: + return self + + def requires_grad_(self, value: bool) -> LoadedModel: + observed["requires_grad"] = value + return self + + def fake_from_pretrained(source: str, **kwargs: Any) -> LoadedModel: + observed["source"] = source + observed["load_kwargs"] = kwargs + return LoadedModel() + + monkeypatch.setattr(binder.AutoModel, "from_pretrained", fake_from_pretrained) + monkeypatch.setattr(binder, "_ESMC_CACHE", None) + monkeypatch.setattr(binder, "_ESMC_CACHE_KEY", None) + monkeypatch.setattr(binder, "_ESMC_CACHE_CONTEXT", {}) + revision = "a" * 40 + model = binder._load_fold_model( + "Custom/Fold", + revision=revision, + lm_dropout=0.5, + cache_esmc=True, + device="cpu", + kernel_backend=None, + compile_model=False, + local_files_only=True, + ) + + assert observed["source"] == "Custom/Fold" + assert observed["load_kwargs"]["revision"] == revision + assert observed["load_kwargs"]["local_files_only"] is True + assert observed["load_kwargs"]["trust_remote_code"] is True + assert observed["esmc_load"] == ( + "Synthyra/ESMplusplus_6B", + {"device": "cpu", "local_files_only": True}, + ) + assert model._fastplms_binder_load_identity == { + "repo_id": "Custom/Fold", + "requested_revision": revision, + "local_files_only": True, + } + + +def test_binder_runner_uses_registered_commits_and_owns_offline_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fold_calls: list[tuple[str, dict[str, Any]]] = [] + lm_calls: list[tuple[str, dict[str, Any]]] = [] + + class LoadedModel: + def to(self, **_kwargs: Any) -> LoadedModel: + return self + + def eval(self) -> LoadedModel: + return self + + def requires_grad_(self, _value: bool) -> LoadedModel: + return self + + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.setattr( + binder, + "_load_fold_model", + lambda model_name, **kwargs: fold_calls.append((model_name, kwargs)) or LoadedModel(), + ) + monkeypatch.setattr( + binder.AutoModelForMaskedLM, + "from_pretrained", + lambda source, **kwargs: lm_calls.append((source, kwargs)) or LoadedModel(), + ) + + runner = binder.FastPLMsBinderDesign() + runner.load(device="cpu", local_files_only=True) + + registered = binder._registered_fast_revisions() + assert os.environ["HF_HUB_OFFLINE"] == "1" + assert os.environ["TRANSFORMERS_OFFLINE"] == "1" + assert len(fold_calls) == 3 + for model_name, kwargs in fold_calls: + repo_id = binder._repo_name(model_name) + assert kwargs["revision"] == registered[repo_id] + assert kwargs["local_files_only"] is True + assert lm_calls == [ + ( + "Synthyra/ESMplusplus_6B", + { + "revision": registered["Synthyra/ESMplusplus_6B"], + "local_files_only": True, + "trust_remote_code": True, + "dtype": torch.float32, + }, + ) + ] + assert runner.lm_model._fastplms_binder_load_identity == { + "repo_id": "Synthyra/ESMplusplus_6B", + "requested_revision": registered["Synthyra/ESMplusplus_6B"], + "local_files_only": True, + } + + +def test_custom_binder_model_requires_immutable_revision() -> None: + runner = binder.FastPLMsBinderDesign() + with pytest.raises(ValueError, match="requires an immutable revision"): + runner.load( + device="cpu", + inversion_model_names=("Custom/unpinned",), + ) + + +@pytest.mark.feature +def test_seeded_binder_workflow_is_short_and_reproducible( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(binder, "fold_and_get_distogram", _fake_fold) + monkeypatch.setattr( + binder, + "compute_fastplms_pseudoperplexity_nll", + _fake_pseudoperplexity, + ) + + first_sequences, first_trajectory, first_rows = _run_seeded_workflow() + second_sequences, second_trajectory, second_rows = _run_seeded_workflow() + + assert first_sequences == second_sequences == ["ACD|AAA"] + assert list(first_trajectory) == list(second_trajectory) == [0] + torch.testing.assert_close( + first_trajectory[0]["total_loss"], + second_trajectory[0]["total_loss"], + rtol=0.0, + atol=0.0, + ) + assert len(first_rows) == len(second_rows) == 1 + assert first_rows[0]["critic_name"] == APPROVED_CRITIC + assert first_rows[0]["designed_sequence"] == "ACD|AAA" + assert first_rows[0]["binder_length"] == 3 + assert first_rows[0]["target_length"] == 3 + assert first_rows[0]["iptm"] == second_rows[0]["iptm"] == 1.0 + + +def test_prompt_sampling_preserves_the_callers_random_stream() -> None: + factory = binder.BINDER_PROMPT_FACTORIES["minibinder"] + random.seed(1234) + expected = random.random() + random.seed(1234) + + first = factory.sample(seed=17) + observed = random.random() + + assert first == factory.sample(seed=17) + assert observed == expected + + +def test_atom_batch_padding_uses_the_largest_prepared_table_without_truncation() -> None: + small = { + "ref_pos": torch.arange(32 * 3, dtype=torch.float32).reshape(1, 32, 3), + "atom_attention_mask": torch.ones((1, 32), dtype=torch.bool), + } + large = { + "ref_pos": torch.arange(65 * 3, dtype=torch.float32).reshape(1, 65, 3), + "atom_attention_mask": torch.ones((1, 65), dtype=torch.bool), + } + + padded = binder._pad_prepared_atom_features([(small, ["small"]), (large, ["large"])]) + + assert padded[0][0]["ref_pos"].shape == (1, 96, 3) + assert padded[1][0]["ref_pos"].shape == (1, 96, 3) + torch.testing.assert_close(padded[1][0]["ref_pos"][:, :65], large["ref_pos"]) + assert not padded[1][0]["atom_attention_mask"][:, 65:].any() + with pytest.raises(ValueError, match="Refusing to truncate atom features"): + binder._resize_tensor(large["ref_pos"], dim=1, size=64) + + +def test_binder_output_directory_is_exclusive_and_fail_closed( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + output_dir = tmp_path / "binder-run" + output_dir.mkdir() + (output_dir / "stale-result.txt").write_text("old run\n", encoding="utf-8") + fold_called = False + + def forbidden_fold(*args: Any, **kwargs: Any) -> Any: + nonlocal fold_called + del args, kwargs + fold_called = True + raise AssertionError("stale output reached model execution") + + monkeypatch.setattr(binder, "fold_and_get_distogram", forbidden_fold) + with pytest.raises(FileExistsError, match="never reused"): + binder.design_binder( + inversion_models={APPROVED_CRITIC: FakeCritic()}, + critic_models={APPROVED_CRITIC: FakeCritic()}, + lm_model=object(), + target_name=None, + target_sequence="ACD", + binder_name=None, + binder_sequence="###", + is_antibody=False, + seed=17, + steps=1, + output_dir=output_dir, + device="cpu", + ) + + assert not fold_called + assert (output_dir / "stale-result.txt").read_text(encoding="utf-8") == "old run\n" + + +def test_interrupted_binder_run_cannot_be_reused( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + output_dir = tmp_path / "interrupted-run" + + def failing_fold(*args: Any, **kwargs: Any) -> Any: + del args, kwargs + raise RuntimeError("injected fold failure") + + monkeypatch.setattr(binder, "fold_and_get_distogram", failing_fold) + kwargs = { + "inversion_models": {APPROVED_CRITIC: FakeCritic()}, + "critic_models": {APPROVED_CRITIC: FakeCritic()}, + "lm_model": object(), + "target_name": None, + "target_sequence": "ACD", + "binder_name": None, + "binder_sequence": "###", + "is_antibody": False, + "seed": 17, + "steps": 1, + "output_dir": output_dir, + "device": "cpu", + } + + with pytest.raises(RuntimeError, match="injected fold failure"): + binder.design_binder(**kwargs) + assert output_dir.is_dir() + assert not (output_dir / "run_manifest.json").exists() + + with pytest.raises(FileExistsError, match="never reused"): + binder.design_binder(**kwargs) + + +def test_cli_rejects_stale_output_before_model_loading( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + output_dir = tmp_path / "existing-run" + output_dir.mkdir() + loaded = False + + class ForbiddenRunner: + def load(self, **kwargs: Any) -> None: + nonlocal loaded + del kwargs + loaded = True + + monkeypatch.setattr(binder, "FastPLMsBinderDesign", ForbiddenRunner) + args = binder.parse_args(["--steps", "1", "--output-dir", str(output_dir)]) + with pytest.raises(FileExistsError, match="never reused"): + binder.run_local(args) + + assert not loaded + + +@pytest.mark.feature +def test_selected_sequence_loss_and_logits_share_the_same_optimization_step( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + optimization_designs: list[torch.Tensor] = [] + + def ranked_fake_fold( + model: Any, + target_seq: str, + target_one_hot: torch.Tensor, + design: torch.Tensor, + num_loops: int = 0, + num_sampling_steps: int = 1, + calculate_confidence: bool = False, + seed: int | None = None, + ) -> dict[str, Any]: + # target_one_hot: (...), design: (...) + result = _fake_fold( + model, + target_seq, + target_one_hot, + design, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + calculate_confidence=calculate_confidence, + seed=seed, + ) + if num_sampling_steps != 200: + optimization_step = len(optimization_designs) + optimization_designs.append(design.detach().clone()) + residue = "A" if optimization_step == 0 else "D" + result["seq_list"] = [f"{target_seq}|{residue * design.size(1)}"] + # result['iptm']: (design.size(0),) + result["iptm"] = torch.full( + (design.size(0),), + 0.95 if optimization_step == 0 else 0.20, + device=design.device, + ) + return result + + monkeypatch.setattr(binder, "fold_and_get_distogram", ranked_fake_fold) + monkeypatch.setattr( + binder, + "compute_fastplms_pseudoperplexity_nll", + _fake_pseudoperplexity, + ) + monkeypatch.setattr(binder, "_write_results_table", lambda path, rows: None) + monkeypatch.setattr( + binder, + "_write_official_selection_table", + lambda path, rows, required_hero_critics=None: None, + ) + + output_dir = tmp_path / "binder-run" + sequences, trajectory, rows = binder.design_binder( + inversion_models={APPROVED_CRITIC: FakeCritic()}, + critic_models={APPROVED_CRITIC: FakeCritic()}, + lm_model=object(), + target_name=None, + target_sequence="ACD", + binder_name=None, + binder_sequence="###", + is_antibody=False, + seed=17, + batch_size=1, + steps=2, + output_dir=output_dir, + device="cpu", + ) + + assert sequences == ["ACD|AAA"] + assert rows[0]["designed_sequence"] == "ACD|AAA" + assert rows[0]["selected_step"] == 0 + assert rows[0]["final_loss"] == float(trajectory[0]["total_loss"][0].item()) + + selected_logits = torch.load(rows[0]["logits_path"], weights_only=True) + first_step_temperature = ( + binder.DEFAULT_TEMPERATURE_MIN + (1 - binder.DEFAULT_TEMPERATURE_MIN) * 0.5 + ) + torch.testing.assert_close( + torch.softmax(selected_logits / first_step_temperature, dim=-1), + optimization_designs[0][0], + rtol=0.0, + atol=0.0, + ) + manifest = json.loads((output_dir / "run_manifest.json").read_text(encoding="utf-8")) + assert manifest["schema_version"] == 2 + assert manifest["seed"] == 17 + assert manifest["steps"] == 2 + assert len(manifest["target_sequence_sha256"]) == 64 + assert len(manifest["binder_prompt_sha256"]) == 64 + assert manifest["configuration"]["is_antibody"] is False + assert manifest["configuration"]["loss_weights"] == binder.LOSS_WEIGHTS + assert manifest["command"] + inversion_identity = manifest["models"]["inversion"][0] + assert inversion_identity["requested"] == APPROVED_CRITIC + assert { + "requested_revision", + "hub_revision", + "weights_revision", + "runtime_revision", + "local_files_only", + }.issubset(inversion_identity) + assert { + "hf_hub_offline", + "transformers_offline", + }.issubset(manifest["environment"]) + assert manifest["tokenizer"] is None + + +def test_binder_model_identity_records_selected_kernel_and_mixed_parameter_dtypes() -> None: + class Config: + _name_or_path = "Synthyra/ESMFold2-test" + _commit_hash = "0123456789abcdef" + fastplms_weights_revision = "1" * 40 + fastplms_runtime_revision = "source-tree-sha256:" + "2" * 64 + esmc_attn_backend = "sdpa" + + def to_dict(self) -> dict[str, object]: + return {"model_type": "esmfold2", "d_pair": 8} + + class PrecisionStatus: + def as_dict(self) -> dict[str, object]: + return { + "requested": "auto", + "resolved": "bf16", + "reason": "test contract", + "device": "cuda:0", + "transformer_engine_version": None, + } + + class MixedDtypeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.float_weight = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + self.bfloat_weight = torch.nn.Parameter(torch.ones(3, dtype=torch.bfloat16)) + self.config = Config() + self._kernel_backend = "cuequivariance" + self.esmc_precision_status = PrecisionStatus() + + model = MixedDtypeModel() + binder._record_model_load_identity( + model, + repo_id="Synthyra/ESMFold2-test", + revision="3" * 40, + local_files_only=True, + ) + identity = binder._model_identity("requested-model", model) + + assert identity["kernel_backend"] == "cuequivariance" + assert identity["repo_id"] == "Synthyra/ESMFold2-test" + assert identity["requested_revision"] == "3" * 40 + assert identity["hub_revision"] == "0123456789abcdef" + assert identity["weights_revision"] == "1" * 40 + assert identity["runtime_revision"] == "source-tree-sha256:" + "2" * 64 + assert identity["local_files_only"] is True + assert identity["parameter_dtype"] == "mixed[torch.bfloat16,torch.float32]" + assert identity["parameter_dtypes"] == ["torch.bfloat16", "torch.float32"] + assert identity["parameter_dtype_numel"] == { + "torch.bfloat16": 3, + "torch.float32": 2, + } + assert identity["effective_precision"]["resolved"] == "bf16" + + tokenizer = SimpleNamespace( + get_vocab=lambda: {"": 0, "A": 1}, + init_kwargs={}, + name_or_path=None, + ) + tokenizer_identity = binder._tokenizer_identity(tokenizer, model=model) + assert tokenizer_identity is not None + assert tokenizer_identity["repo_id"] == "Synthyra/ESMFold2-test" + assert tokenizer_identity["requested_revision"] == "3" * 40 + assert tokenizer_identity["weights_revision"] == "1" * 40 + assert tokenizer_identity["runtime_revision"] == "source-tree-sha256:" + "2" * 64 + assert tokenizer_identity["local_files_only"] is True + + +def test_consensus_requires_every_named_hero_critic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.importorskip("pandas") + monkeypatch.setattr( + binder, + "_compute_isoelectric_points", + lambda sequences: [5.0] * len(sequences), + ) + critic_a = "critic-a" + critic_b = "critic-b" + rows = [ + { + "critic_name": critic_a, + "designed_sequence": "ACD|AAA", + "is_antibody": False, + "iptm": 0.95, + "distogram_iptm_proxy": 0.80, + }, + { + "critic_name": critic_b, + "designed_sequence": "ACD|AAA", + "is_antibody": False, + "iptm": 0.96, + "distogram_iptm_proxy": 0.90, + }, + { + "critic_name": critic_a, + "designed_sequence": "ACD|DDD", + "is_antibody": False, + "iptm": 0.99, + "distogram_iptm_proxy": 0.70, + }, + ] + + selected = binder.select_official_designs( + rows, + top_k=2, + required_hero_critics=(critic_a, critic_b), + ).set_index("designed_sequence") + + complete = selected.loc["ACD|AAA"] + assert complete["hero_critic_count"] == 2 + assert complete["required_hero_critic_count"] == 2 + assert complete["iptm_proxy_score"] == pytest.approx(0.85) + assert bool(complete["all_hero_critics_pass"]) + + incomplete = selected.loc["ACD|DDD"] + assert incomplete["hero_critic_count"] == 1 + assert incomplete["required_hero_critic_count"] == 2 + assert not bool(incomplete["all_hero_critics_pass"]) + + +def test_successful_nonempty_binder_run_reaches_selection_summary( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + pytest.importorskip("pandas") + monkeypatch.setattr( + binder, + "_compute_isoelectric_points", + lambda sequences: [5.0] * len(sequences), + ) + rows = [ + { + "critic_name": "critic-a", + "designed_sequence": "ACD|AAA", + "is_antibody": False, + "iptm": 0.95, + "distogram_iptm_proxy": 0.85, + } + ] + + class FakeRunner: + def load(self, **kwargs: Any) -> None: + del kwargs + + def design(self, **kwargs: Any) -> tuple[list[str], None, list[dict[str, Any]]]: + del kwargs + return ["AAA"], None, rows + + messages: list[tuple[str, tuple[Any, ...]]] = [] + monkeypatch.setattr(binder, "FastPLMsBinderDesign", FakeRunner) + monkeypatch.setattr( + binder.logger, + "info", + lambda message, *args: messages.append((message, args)), + ) + + binder.run_local( + binder.parse_args(["--steps", "1", "--output-dir", str(tmp_path / "binder-summary")]) + ) + + summary = [entry for entry in messages if entry[0].startswith("Top official selection")] + assert len(summary) == 1 + assert summary[0][1][2] == pytest.approx(0.85) diff --git a/tests/integration/test_cuda_training_contracts.py b/tests/integration/test_cuda_training_contracts.py new file mode 100644 index 0000000..ade34b2 --- /dev/null +++ b/tests/integration/test_cuda_training_contracts.py @@ -0,0 +1,388 @@ +"""Small real-CUDA contracts that must not be satisfied by CPU-only probes.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path + +from fastplms.attention import _core +from fastplms.models.dplm.modeling_dplm import DPLMConfig, DPLMForMaskedLM +from fastplms.models.dplm2.modeling_dplm2 import DPLM2Config, DPLM2ForMaskedLM +from fastplms.models.esm2.modeling_fastesm import FastEsmConfig, FastEsmModel +from fastplms.models.esm3.modeling_esm3 import ( + FastESM3Config, + FastESM3GenerationConfig, + FastESM3Model, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import TransformerStack +from tests.conftest import strict_fp32_matmul +from tests.integration.test_ttt import ( + DummyPretrainedTTTConfig, + DummyPretrainedTTTModel, +) + + +pytestmark = pytest.mark.gpu + + +def _diffusion_config(vocab_size: int) -> dict[str, object]: + return { + "vocab_size": vocab_size, + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "intermediate_size": 64, + "hidden_dropout_prob": 0.0, + "attention_probs_dropout_prob": 0.0, + "max_position_embeddings": 64, + "pad_token_id": 1, + "bos_token_id": 0, + "eos_token_id": 2, + "mask_token_id": 32, + "position_embedding_type": "rotary", + "attn_backend": "sdpa", + } + + +@pytest.mark.parametrize("family", ("dplm", "dplm2", "esm3")) +def test_seeded_generation_trace_executes_deterministically_on_cuda( + family: str, +) -> None: + assert torch.cuda.is_available(), "Generation CUDA contract requires CUDA." + torch.manual_seed(53) + + if family == "dplm": + model = DPLMForMaskedLM( + DPLMConfig(**_diffusion_config(33)), + dropout=0.0, + ).train().cuda() + # inputs: (1, 5) + inputs = torch.tensor([[0, 6, 32, 8, 2]], device="cuda") + + def run() -> torch.Tensor: + return model.generate(inputs, max_iter=2) + + elif family == "dplm2": + model = DPLM2ForMaskedLM( + DPLM2Config(**_diffusion_config(64)), + dropout=0.0, + ).train().cuda() + # inputs: (1, 8) + inputs = torch.tensor( + [[33, 50, 50, 34, 0, 6, 6, 2]], + device="cuda", + ) + + def run() -> torch.Tensor: + generated = model.generate(inputs, max_iter=2) + return generated["output_tokens"] + + else: + model = FastESM3Model( + FastESM3Config( + hidden_size=64, + num_attention_heads=4, + num_vector_heads=8, + num_hidden_layers=2, + ) + ).train().cuda() + generation_config = FastESM3GenerationConfig( + num_steps=2, + temperature=1.0, + seed=73, + ) + + def run() -> str: + return model.generate("MK__A", generation_config) + + first_child = next(model.children()) + first_child.eval() + expected_training_states = tuple(module.training for module in model.modules()) + + traces: list[torch.Tensor | str] = [] + for _ in range(2): + torch.manual_seed(59) + torch.cuda.manual_seed_all(59) + traces.append(run()) + assert tuple(module.training for module in model.modules()) == expected_training_states + + if isinstance(traces[0], str): + assert traces[0] == traces[1] + else: + assert torch.is_tensor(traces[0]) + assert traces[0].is_cuda + assert torch.equal(traces[0], traces[1]) + + +@pytest.mark.parametrize("family", ("dplm", "dplm2", "esm3")) +def test_generation_restores_mixed_training_state_after_cuda_forward_failure( + family: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert torch.cuda.is_available(), "Generation CUDA contract requires CUDA." + if family == "dplm": + model = DPLMForMaskedLM( + DPLMConfig(**_diffusion_config(33)), + dropout=0.2, + ).train().cuda() + # inputs: (1, 5) + inputs = torch.tensor([[0, 6, 32, 8, 2]], device="cuda") + + def run(): + return model.generate(inputs, max_iter=2) + + elif family == "dplm2": + model = DPLM2ForMaskedLM( + DPLM2Config(**_diffusion_config(64)), + dropout=0.2, + ).train().cuda() + # inputs: (1, 8) + inputs = torch.tensor([[33, 50, 50, 34, 0, 6, 6, 2]], device="cuda") + + def run(): + return model.generate(inputs, max_iter=2) + + else: + model = FastESM3Model( + FastESM3Config( + hidden_size=64, + num_attention_heads=4, + num_vector_heads=8, + num_hidden_layers=1, + ) + ).train().cuda() + + def run(): + return model.generate( + "MK__A", + FastESM3GenerationConfig(num_steps=2, seed=73), + ) + + next(model.children()).eval() + expected_training_states = tuple(module.training for module in model.modules()) + + def fail_forward(*args, **kwargs): + del args, kwargs + raise RuntimeError("synthetic generation forward failure") + + monkeypatch.setattr(model, "forward", fail_forward) + with pytest.raises(RuntimeError, match="synthetic generation forward failure"): + run() + assert tuple(module.training for module in model.modules()) == expected_training_states + + +def test_esmplusplus_flex_sequence_masks_reuse_on_cuda_and_match_sdpa( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert torch.cuda.is_available(), "ESM++ Flex CUDA contract requires CUDA." + torch.manual_seed(79) + sdpa = TransformerStack(32, 2, 1, attn_backend="sdpa").cuda().to(torch.bfloat16) + flex = TransformerStack(32, 2, 1, attn_backend="flex_attention").cuda().to(torch.bfloat16) + flex.load_state_dict(sdpa.state_dict()) + # pattern: (2, 5) + pattern = torch.tensor( + ((True, True, True, False, False), (True, True, False, True, False)), + device="cuda", + ) + created = 0 + original_create_block_mask = _core.create_block_mask + + def count_create_block_mask(*args, **kwargs): + nonlocal created + created += 1 + return original_create_block_mask(*args, **kwargs) + + monkeypatch.setattr(_core, "create_block_mask", count_create_block_mask) + _core.clear_flex_attention_caches() + flex_input = torch.randn(2, 5, 32, device="cuda", dtype=torch.bfloat16).requires_grad_() + sdpa_input = flex_input.detach().clone().requires_grad_() + + flex_output = flex(flex_input, sequence_id=pattern).last_hidden_state + repeated = flex(flex_input.detach(), sequence_id=pattern.clone()).last_hidden_state + assert created == 1 + torch.testing.assert_close(repeated, flex_output.detach(), rtol=0.0, atol=0.0) + flex_output.float().square().mean().backward() + sdpa_output = sdpa(sdpa_input, sequence_id=pattern).last_hidden_state + sdpa_output.float().square().mean().backward() + torch.testing.assert_close(flex_output, sdpa_output, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(flex_input.grad, sdpa_input.grad, rtol=3e-2, atol=3e-2) + assert torch.isfinite(flex_input.grad).all() + + # chain_pattern: (2, 5) + chain_pattern = torch.tensor( + ((0, 0, 1, -1, -1), (0, 1, 1, 2, -1)), + device="cuda", + ) + chain_output = flex(flex_input.detach(), sequence_id=chain_pattern).last_hidden_state + assert created == 2 + assert torch.isfinite(chain_output).all() + assert len(_core._flex_block_masks) == 2 + _core.clear_flex_attention_caches() + assert not _core._flex_block_masks + + +def test_eager_sdpa_and_flex_match_forward_and_backward_on_cuda() -> None: + assert torch.cuda.is_available(), "Attention CUDA contract requires CUDA." + config = FastEsmConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=128, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=64, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + token_dropout=False, + attn_backend="sdpa", + ) + torch.manual_seed(61) + initial = FastEsmModel(config) + state = { + name: tensor.detach().clone() + for name, tensor in initial.state_dict().items() + } + # input_ids: (2, 17) + input_ids = torch.tensor( + ( + (0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2, 1, 1, 1, 1), + (0, 20, 19, 18, 17, 16, 15, 14, 2, 1, 1, 1, 1, 1, 1, 1, 1), + ), + device="cuda", + ) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + outputs: dict[str, torch.Tensor] = {} + parameter_gradients: dict[str, dict[str, torch.Tensor]] = {} + + with strict_fp32_matmul(): + for backend in ("eager", "sdpa", "flex_attention"): + model = FastEsmModel(config).cuda().train() + model.load_state_dict(state) + model.set_attn_implementation(backend) + output = model( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + # loss: () + loss = output[attention_mask].float().square().mean() + assert loss.is_cuda + assert torch.isfinite(loss) + loss.backward() + gradients = { + name: parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + assert gradients + assert all( + torch.isfinite(gradient).all() for gradient in gradients.values() + ) + # outputs[backend]: (...) + outputs[backend] = output.detach() + parameter_gradients[backend] = gradients + + expected_attention_projections = ( + "query", + "key", + "value", + "attention.output.dense", + ) + assert all( + any(projection in name for name in parameter_gradients["eager"]) + for projection in expected_attention_projections + ) + + torch.testing.assert_close(outputs["sdpa"], outputs["eager"], rtol=1e-5, atol=1e-6) + torch.testing.assert_close( + outputs["flex_attention"], + outputs["sdpa"], + rtol=1e-5, + atol=1e-5, + ) + assert parameter_gradients["sdpa"].keys() == parameter_gradients["eager"].keys() + for name, eager_gradient in parameter_gradients["eager"].items(): + torch.testing.assert_close( + parameter_gradients["sdpa"][name], + eager_gradient, + rtol=1e-5, + atol=1e-6, + ) + assert ( + parameter_gradients["flex_attention"].keys() + == parameter_gradients["sdpa"].keys() + ) + for name, sdpa_gradient in parameter_gradients["sdpa"].items(): + torch.testing.assert_close( + parameter_gradients["flex_attention"][name], + sdpa_gradient, + rtol=1e-5, + atol=1e-5, + ) + + +def test_ttt_step_reset_and_save_reload_execute_on_cuda(tmp_path: Path) -> None: + assert torch.cuda.is_available(), "TTT CUDA contract requires CUDA." + model = DummyPretrainedTTTModel(DummyPretrainedTTTConfig()).train().cuda() + model._ttt_ensure_initialized() + initial_state = model._ttt_snapshot_lora_state() + assert initial_state + assert all( + tensor.is_cuda + for module_state in initial_state + for tensor in module_state.values() + ) + ttt_config = { + "steps": 1, + "ags": 1, + "batch_size": 2, + "mask_ratio": 0.5, + "bert_leave_prob": 0.0, + "bert_replace_prob": 0.0, + "seed": 67, + "initial_state_reset": True, + } + + first_metrics = model.ttt(seq="ACDE", ttt_config=ttt_config) + first_adapted = model._ttt_snapshot_lora_state() + assert len(first_metrics["losses"]) == 1 + assert all(torch.isfinite(torch.tensor(first_metrics["losses"]))) + assert any( + not torch.equal(initial[name], adapted[name]) + for initial, adapted in zip(initial_state, first_adapted, strict=True) + for name in initial + ) + + model.ttt_reset() + reset_state = model._ttt_snapshot_lora_state() + for initial, reset in zip(initial_state, reset_state, strict=True): + for name in initial: + torch.testing.assert_close(reset[name], initial[name]) + + second_metrics = model.ttt(seq="ACDE", ttt_config=ttt_config) + second_adapted = model._ttt_snapshot_lora_state() + assert second_metrics["losses"] == pytest.approx(first_metrics["losses"]) + for first, second in zip(first_adapted, second_adapted, strict=True): + for name in first: + torch.testing.assert_close(second[name], first[name]) + + model.save_pretrained(tmp_path, safe_serialization=True) + restored = DummyPretrainedTTTModel.from_pretrained( + tmp_path, + local_files_only=True, + ).cuda() + restored_state = restored._ttt_snapshot_lora_state() + assert restored._ttt_initialized is True + for expected, observed in zip(second_adapted, restored_state, strict=True): + for name in expected: + assert observed[name].is_cuda + torch.testing.assert_close(observed[name], expected[name]) + + restored.ttt_reset() + restored_reset = restored._ttt_snapshot_lora_state() + for initial, observed in zip(initial_state, restored_reset, strict=True): + for name in initial: + torch.testing.assert_close(observed[name], initial[name]) diff --git a/tests/integration/test_dplm_generation.py b/tests/integration/test_dplm_generation.py new file mode 100644 index 0000000..8490233 --- /dev/null +++ b/tests/integration/test_dplm_generation.py @@ -0,0 +1,657 @@ +"""DPLM and DPLM2 diffusion-generation feature contracts.""" + +from __future__ import annotations + +import pytest +import torch +from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput + +from fastplms.models.dplm.modeling_dplm import ( + DPLMConfig, + DPLMForMaskedLM, + DPLMForSequenceClassification, + DPLMForTokenClassification, + DPLMModel, + DPLMSequenceClassifierOutput, + DPLMTokenClassifierOutput, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + FAST_DPLM2_ENCODER, + DPLM2Config, + DPLM2EncoderOutput, + DPLM2ForMaskedLM, + DPLM2ForSequenceClassification, + DPLM2ForTokenClassification, + DPLM2MaskedLMOutput, + DPLM2Model, + DPLM2ModelOutput, + DPLM2SequenceClassifierOutput, + DPLM2TokenClassifierOutput, + ModifiedRotaryEmbedding, +) + + +pytestmark = pytest.mark.feature + + +def _common_config(vocab_size: int) -> dict[str, object]: + return { + "vocab_size": vocab_size, + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "intermediate_size": 64, + "hidden_dropout_prob": 0.0, + "attention_probs_dropout_prob": 0.0, + "max_position_embeddings": 64, + "pad_token_id": 1, + "bos_token_id": 0, + "eos_token_id": 2, + "mask_token_id": 32, + "position_embedding_type": "rotary", + "attn_backend": "sdpa", + } + + +def _assert_nested_close(actual, expected) -> None: + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, (tuple, list)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_value, expected_value in zip(actual, expected, strict=True): + _assert_nested_close(actual_value, expected_value) + return + assert actual == expected + + +def test_dplm_argmax_generation_preserves_fixed_positions() -> None: + torch.manual_seed(13) + model = DPLMForMaskedLM(DPLMConfig(**_common_config(33)), dropout=0.0).eval() + # input_tokens: (1, 6) + input_tokens = torch.tensor([[0, 6, 7, 8, 2, 1]]) + # fixed: (1, 6) + fixed = torch.tensor([[False, True, False, False, False, False]]) + + output_tokens = model.generate( + input_tokens, + max_iter=3, + partial_masks=fixed, + sampling_strategy="argmax", + disable_resample=True, + ) + + assert output_tokens.shape == input_tokens.shape + assert output_tokens[0, 1].item() == 6 + assert torch.equal(output_tokens[0, [0, 4, 5]], input_tokens[0, [0, 4, 5]]) + generated = output_tokens[0, 2:4] + assert not bool(torch.isin(generated, torch.tensor([0, 1, 2, 24, 32])).any()) + + +def test_dplm_vanilla_default_is_zero_temperature() -> None: + model = DPLMForMaskedLM(DPLMConfig(**_common_config(33)), dropout=0.0).eval() + # input_tokens: (1, 5) + input_tokens = torch.tensor([[0, 6, 7, 8, 2]]) + + torch.manual_seed(29) + default_output = model.generate( + input_tokens, + max_iter=2, + sampling_strategy="vanilla", + disable_resample=True, + ) + torch.manual_seed(31) + zero_temperature_output = model.generate( + input_tokens, + max_iter=2, + temperature=0.0, + sampling_strategy="vanilla", + disable_resample=True, + ) + + assert torch.equal(default_output, zero_temperature_output) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size"), + ( + (DPLMForMaskedLM, DPLMConfig, 33), + (DPLM2ForMaskedLM, DPLM2Config, 64), + ), +) +def test_dplm_families_reject_static_bf16_inference( + model_class: type[torch.nn.Module], + config_class: type, + vocab_size: int, +) -> None: + model = ( + model_class( + config_class(**_common_config(vocab_size)), + dropout=0.0, + ) + .to(dtype=torch.bfloat16) + .eval() + ) + # X: (1, 4) + X = torch.tensor([[0, 6, 7, 2]]) + + with pytest.raises(RuntimeError, match="FP32-resident parameters"): + model(input_ids=X, attention_mask=torch.ones_like(X)) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size"), + ( + (DPLMForMaskedLM, DPLMConfig, 33), + (DPLM2ForMaskedLM, DPLM2Config, 64), + ), +) +def test_masked_lm_dropout_round_trips_through_config_and_weights( + model_class: type[torch.nn.Module], + config_class: type, + vocab_size: int, + tmp_path, +) -> None: + config = config_class(**_common_config(vocab_size)) + config.hidden_dropout_prob = 0.37 + model = model_class(config).eval() + model.save_pretrained(tmp_path) + + reloaded = model_class.from_pretrained(tmp_path, local_files_only=True).eval() + + assert model.config.hidden_dropout_prob == pytest.approx(0.37) + assert reloaded.config.hidden_dropout_prob == pytest.approx(0.37) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size"), + ( + (DPLMForMaskedLM, DPLMConfig, 33), + (DPLM2ForMaskedLM, DPLM2Config, 64), + ), +) +def test_masked_lm_resize_updates_input_and_output_projections( + model_class: type[torch.nn.Module], + config_class: type, + vocab_size: int, + tmp_path, +) -> None: + model = model_class(config_class(**_common_config(vocab_size))).eval() + resized_vocab_size = vocab_size + 5 + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 6, 7, 2]]) + with torch.inference_mode(): + original_logits = model(input_ids=input_ids).logits + + model.resize_token_embeddings(resized_vocab_size) + + assert model.get_input_embeddings().num_embeddings == resized_vocab_size + assert model.get_output_embeddings().out_features == resized_vocab_size + assert model.lm_head.bias.shape == (resized_vocab_size,) + assert model.config.vocab_size == resized_vocab_size + assert model.get_output_embeddings().bias is None + + with torch.inference_mode(): + output = model(input_ids=input_ids) + assert output.logits.shape == (*input_ids.shape, resized_vocab_size) + torch.testing.assert_close(output.logits[..., :vocab_size], original_logits) + + save_path = tmp_path / model_class.__name__ + model.save_pretrained(save_path, safe_serialization=True) + reloaded = model_class.from_pretrained(save_path, local_files_only=True).eval() + assert reloaded.get_input_embeddings().num_embeddings == resized_vocab_size + assert reloaded.get_output_embeddings().out_features == resized_vocab_size + assert reloaded.get_output_embeddings().bias is None + assert reloaded.lm_head.bias.shape == (resized_vocab_size,) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size", "output_class", "labels"), + ( + ( + DPLMForSequenceClassification, + DPLMConfig, + 33, + DPLMSequenceClassifierOutput, + [1], + ), + ( + DPLMForTokenClassification, + DPLMConfig, + 33, + DPLMTokenClassifierOutput, + [[1, 1, 1, 1]], + ), + ( + DPLM2ForSequenceClassification, + DPLM2Config, + 64, + DPLM2SequenceClassifierOutput, + [1], + ), + ( + DPLM2ForTokenClassification, + DPLM2Config, + 64, + DPLM2TokenClassifierOutput, + [[1, 1, 1, 1]], + ), + ), +) +def test_dplm_task_heads_honor_config_and_explicit_return_dict( + model_class: type[torch.nn.Module], + config_class: type, + vocab_size: int, + output_class: type, + labels: list[int] | list[list[int]], +) -> None: + config = config_class(**_common_config(vocab_size), num_labels=3, return_dict=False) + model = model_class(config).eval() + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 6, 7, 2]]) + # label_tensor: (...) + label_tensor = torch.tensor(labels) + + with torch.inference_mode(): + unlabeled_tuple = model(input_ids=input_ids) + unlabeled_output = model(input_ids=input_ids, return_dict=True) + labeled_tuple = model( + input_ids=input_ids, + labels=label_tensor, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + ) + labeled_output = model( + input_ids=input_ids, + labels=label_tensor, + output_attentions=True, + output_hidden_states=True, + output_s_max=True, + return_dict=True, + ) + + assert type(unlabeled_output) is output_class + assert tuple(unlabeled_output.keys()) == ("logits",) + assert isinstance(unlabeled_tuple, tuple) + assert len(unlabeled_tuple) == 1 + torch.testing.assert_close(unlabeled_tuple[0], unlabeled_output.logits) + + assert type(labeled_output) is output_class + assert isinstance(labeled_output, (SequenceClassifierOutput, TokenClassifierOutput)) + assert tuple(labeled_output.keys()) == ( + "loss", + "logits", + "hidden_states", + "attentions", + "s_max", + ) + assert isinstance(labeled_tuple, tuple) + _assert_nested_close(labeled_tuple, labeled_output.to_tuple()) + + +@pytest.mark.parametrize( + ("model_class", "output_class", "expected_keys"), + ( + ( + DPLM2Model, + DPLM2ModelOutput, + ("last_hidden_state", "hidden_states", "attentions", "s_max"), + ), + ( + DPLM2ForMaskedLM, + DPLM2MaskedLMOutput, + ( + "loss", + "logits", + "hidden_states", + "attentions", + "s_max", + "last_hidden_state", + ), + ), + ), +) +def test_dplm2_base_and_mlm_preserve_full_structured_tuple_contract( + model_class: type[torch.nn.Module], + output_class: type, + expected_keys: tuple[str, ...], +) -> None: + model = model_class(DPLM2Config(**_common_config(64))).eval() + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 6, 7, 2]]) + call_kwargs = { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "output_attentions": True, + "output_hidden_states": True, + "output_s_max": True, + } + if model_class is DPLM2ForMaskedLM: + call_kwargs["labels"] = input_ids + + with torch.inference_mode(): + structured = model(**call_kwargs, return_dict=True) + tuple_output = model(**call_kwargs, return_dict=False) + + assert type(structured) is output_class + assert tuple(structured.keys()) == expected_keys + assert structured.s_max is not None + assert all(value is not None for value in tuple_output) + _assert_nested_close(tuple_output, structured.to_tuple()) + + +@pytest.mark.parametrize( + "model_class", + ( + DPLMForSequenceClassification, + DPLMForTokenClassification, + DPLM2ForSequenceClassification, + DPLM2ForTokenClassification, + ), +) +@pytest.mark.parametrize( + ("argument", "value"), + ( + ("use_cache", True), + ("past_key_values", ((torch.zeros(1), torch.zeros(1)),)), + ("encoder_hidden_states", torch.zeros(1, 2, 32)), + ("unexpected_option", "typo"), + ), +) +def test_dplm_task_heads_reject_every_unexpected_argument( + model_class: type[torch.nn.Module], + argument: str, + value: object, +) -> None: + config_class = DPLM2Config if model_class.__name__.startswith("DPLM2") else DPLMConfig + vocab_size = 64 if config_class is DPLM2Config else 33 + model = model_class(config_class(**_common_config(vocab_size), num_labels=3)).eval() + + with pytest.raises(TypeError, match=argument): + model(input_ids=torch.tensor([[0, 6, 7, 2]]), **{argument: value}) + + +_DPLM2_INPUT_CONTRACT_MODELS = ( + FAST_DPLM2_ENCODER, + DPLM2Model, + DPLM2ForMaskedLM, + DPLM2ForSequenceClassification, + DPLM2ForTokenClassification, +) + + +@pytest.mark.parametrize("model_class", _DPLM2_INPUT_CONTRACT_MODELS) +def test_dplm2_public_models_require_exactly_one_input_representation( + model_class: type[torch.nn.Module], +) -> None: + model = model_class(DPLM2Config(**_common_config(64))).eval() + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 6, 7, 2]]) + inputs_embeds = model.get_input_embeddings()(input_ids) + + with pytest.raises(ValueError, match="exactly one of input_ids or inputs_embeds"): + model() + with pytest.raises(ValueError, match="exactly one of input_ids or inputs_embeds"): + model(input_ids=input_ids, inputs_embeds=inputs_embeds) + + +@pytest.mark.parametrize("model_class", _DPLM2_INPUT_CONTRACT_MODELS) +@pytest.mark.parametrize( + ("argument", "value"), + ( + ("attention_mask", torch.ones(1, 3, dtype=torch.long)), + ("type_ids", torch.ones(2, 4, dtype=torch.long)), + ), +) +def test_dplm2_public_models_validate_mask_and_type_shapes_before_forward( + model_class: type[torch.nn.Module], + argument: str, + value: torch.Tensor, +) -> None: + # value: (...) + model = model_class(DPLM2Config(**_common_config(64))).eval() + + with pytest.raises(ValueError, match=rf"{argument} must have shape \(1, 4\)"): + model(input_ids=torch.tensor([[0, 6, 7, 2]]), **{argument: value}) + + +@pytest.mark.parametrize( + ("argument", "value"), + ( + ("decoder_input_ids", torch.ones(1, 2, dtype=torch.long)), + ("decoder_attention_mask", torch.ones(1, 2, dtype=torch.long)), + ("decoder_inputs_embeds", torch.ones(1, 2, 32)), + ("encoder_hidden_states", torch.ones(1, 2, 32)), + ("encoder_attention_mask", torch.ones(1, 2, dtype=torch.long)), + ), +) +def test_dplm_masked_lm_rejects_decoder_and_cross_attention_arguments( + argument: str, + value: torch.Tensor, +) -> None: + # value: (...) + model = DPLMForMaskedLM(DPLMConfig(**_common_config(33)), dropout=0.0).eval() + # input_ids: (1, 3) + input_ids = torch.tensor([[0, 6, 2]]) + + with pytest.raises(ValueError, match=argument): + model(input_ids=input_ids, **{argument: value}) + + +@pytest.mark.parametrize( + ("argument", "value"), + ( + ("past_key_values", ((torch.zeros(1), torch.zeros(1)),)), + ("use_cache", True), + ("encoder_hidden_states", torch.ones(1, 2, 32)), + ), +) +def test_dplm_automodel_rejects_decoder_cache_contracts( + argument: str, + value: object, +) -> None: + model = DPLMModel(DPLMConfig(**_common_config(33))).eval() + # input_ids: (1, 3) + input_ids = torch.tensor([[0, 6, 2]]) + + with pytest.raises(ValueError, match=argument): + model(input_ids=input_ids, **{argument: value}) + + +def test_dplm2_rotary_cache_follows_frequency_buffer_dtype() -> None: + rotary = ModifiedRotaryEmbedding(dim=8, aa_type=1, struct_type=0, pad_type=2) + # Q and K are query and key tensors with shape (b, h, l, d). + query = torch.randn(1, 2, 4, 8, dtype=torch.bfloat16) + key = torch.randn_like(query) + # type_ids: (1, 4) + type_ids = torch.ones(1, 4, dtype=torch.long) + + rotary(query, key, type_ids) + assert rotary._cos_cached is not None + assert rotary._cos_cached.dtype == rotary.inv_freq.dtype == torch.float32 + + rotary.align_frequency_buffer(device=query.device, dtype=torch.bfloat16) + assert rotary._cos_cached is None + assert rotary._sin_cached is None + rotary(query, key, type_ids) + assert rotary._cos_cached is not None + assert rotary._cos_cached.dtype == rotary.inv_freq.dtype == torch.bfloat16 + + +def test_dplm2_direct_esm_checkpoint_applies_embeddings_once() -> None: + config = DPLM2Config(**_common_config(64), dplm_type="dplm_esm") + model = DPLM2ForMaskedLM(config, dropout=0.0).eval() + calls = 0 + + def count_embedding_calls( + _module: torch.nn.Module, + _inputs: tuple[object, ...], + _output: object, + ) -> None: + nonlocal calls + calls += 1 + + handle = model.esm.embeddings.register_forward_hook(count_embedding_calls) + try: + model( + input_ids=torch.tensor([[0, 6, 7, 2]]), + attention_mask=torch.ones(1, 4, dtype=torch.long), + ) + finally: + handle.remove() + + assert calls == 1 + + +def test_dplm2_automodel_infers_official_multimodal_types_and_returns_pooling() -> None: + model = object.__new__(DPLM2Model) + torch.nn.Module.__init__(model) + model.config = DPLM2Config(**_common_config(64)) + # input_ids: (1, 8) + input_ids = torch.tensor([[33, 50, 34, 1, 0, 6, 2, 1]]) + # expected_mask: (...) + expected_mask = input_ids.ne(model.config.pad_token_id) + # expected_types: (1, 8) + expected_types = torch.tensor([[0, 0, 0, 2, 1, 1, 1, 2]]) + observed: dict[str, torch.Tensor] = {} + + class CapturingEncoder(torch.nn.Module): + def forward(self, **kwargs: object) -> DPLM2EncoderOutput: + # observed['attention_mask']: (b, l) + observed["attention_mask"] = kwargs["attention_mask"].detach().clone() + # observed['type_ids']: (...) + observed["type_ids"] = kwargs["type_ids"].detach().clone() + return DPLM2EncoderOutput( + last_hidden_state=torch.zeros( + 1, input_ids.shape[1], model.config.hidden_size + ) + ) + + model.esm = CapturingEncoder() + model.pooler = torch.nn.Identity() + output = model(input_ids=input_ids) + tuple_output = model(input_ids=input_ids, return_dict=False) + + assert torch.equal(observed["attention_mask"], expected_mask) + assert torch.equal(observed["type_ids"], expected_types) + assert output.pooler_output is not None + assert tuple_output[1] is not None + + +def test_dplm2_predict_contacts_derives_padding_mask_when_omitted() -> None: + model = DPLM2ForMaskedLM(DPLM2Config(**_common_config(64)), dropout=0.0).eval() + # input_ids: (1, 5) + input_ids = torch.tensor([[0, 6, 7, 2, 1]]) + # expected_mask: (...) + expected_mask = input_ids.ne(model.config.pad_token_id) + observed: dict[str, torch.Tensor] = {} + + def capture_modality_type(input_ids, attention_mask): + # observed['attention_mask']: (b, l) + observed["attention_mask"] = attention_mask.detach().clone() + return torch.ones_like(input_ids) + + model.esm._get_modality_type = capture_modality_type + contacts = model.predict_contacts(input_ids) + + assert torch.equal(observed["attention_mask"], expected_mask) + assert contacts.shape == (1, input_ids.shape[1] - 2, input_ids.shape[1] - 2) + assert torch.isfinite(contacts).all() + + +def test_dplm2_argmax_generation_preserves_modalities_and_fixed_positions() -> None: + torch.manual_seed(17) + model = DPLM2ForMaskedLM(DPLM2Config(**_common_config(64)), dropout=0.0).eval() + # X packs the structure track first and the amino-acid track second, as in + # the official DPLM2 co-generation utility. + # input_tokens: (1, 8) + input_tokens = torch.tensor([[33, 50, 50, 34, 0, 6, 6, 2]]) + # fixed: (1, 8) + fixed = torch.tensor([[False, True, False, False, False, True, False, False]]) + model_inputs: list[torch.Tensor] = [] + + def capture_input( + _module: torch.nn.Module, + _args: tuple[object, ...], + kwargs: dict[str, object], + ) -> None: + # input_tensor: (...) + input_tensor = kwargs["input_ids"] + assert torch.is_tensor(input_tensor) + model_inputs.append(input_tensor.detach().clone()) + + handle = model.register_forward_pre_hook(capture_input, with_kwargs=True) + + try: + output = model.generate( + input_tokens, + max_iter=3, + partial_masks=fixed, + unmasking_strategy="deterministic", + sampling_strategy="argmax", + ) + finally: + handle.remove() + output_tokens = output["output_tokens"] + + assert model_inputs[0][0, 2].item() == model.config.vocab_size - 1 + assert model_inputs[0][0, 6].item() == 32 + assert output_tokens.shape == input_tokens.shape + assert output_tokens[0, 1].item() == 50 + assert output_tokens[0, 5].item() == 6 + assert torch.equal(output_tokens[0, [0, 3, 4, 7]], input_tokens[0, [0, 3, 4, 7]]) + assert int(output_tokens[0, 2]) >= 37 + amino_acid_token = output_tokens[0, 6] + assert int(amino_acid_token) < 33 + assert int(amino_acid_token) not in {0, 1, 2, 3, 24, 25, 26, 27, 28, 32} + + +@pytest.mark.parametrize("family", ("dplm", "dplm2")) +def test_seeded_stochastic_generation_is_repeatable(family: str) -> None: + if family == "dplm": + model = DPLMForMaskedLM(DPLMConfig(**_common_config(33)), dropout=0.0).eval() + # X: (1, 5) + X = torch.tensor([[0, 6, 7, 8, 2]]) + kwargs: dict[str, object] = {"max_iter": 2} + else: + model = DPLM2ForMaskedLM(DPLM2Config(**_common_config(64)), dropout=0.0).eval() + # X: (1, 8) + X = torch.tensor([[33, 50, 50, 34, 0, 6, 6, 2]]) + kwargs = {"max_iter": 2} + + outputs = [] + for _ in range(2): + torch.manual_seed(23) + generated = model.generate(X, **kwargs) + outputs.append(generated["output_tokens"] if isinstance(generated, dict) else generated) + assert torch.equal(outputs[0], outputs[1]) + + +@pytest.mark.parametrize( + ("family", "arguments", "message"), + ( + ("dplm", {"max_iter": 0}, "max_iter"), + ("dplm", {"max_iter": 1, "sampling_strategy": "unknown"}, "sampling strategy"), + ("dplm2", {"max_iter": 1, "unmasking_strategy": "unknown"}, "unmasking strategy"), + ("dplm2", {"max_iter": 1, "sampling_strategy": "unknown"}, "sampling strategy"), + ("dplm2", {"max_iter": 1, "sampling_strategy": "annealing@bad"}, "Annealing"), + ), +) +def test_generation_rejects_invalid_controls( + family: str, + arguments: dict[str, object], + message: str, +) -> None: + if family == "dplm": + model = DPLMForMaskedLM(DPLMConfig(**_common_config(33)), dropout=0.0).eval() + # X: (1, 3) + X = torch.tensor([[0, 32, 2]]) + else: + model = DPLM2ForMaskedLM(DPLM2Config(**_common_config(64)), dropout=0.0).eval() + # X: (1, 6) + X = torch.tensor([[33, 36, 34, 0, 32, 2]]) + with pytest.raises(ValueError, match=message): + model.generate(X, **arguments) diff --git a/tests/integration/test_e1_rag.py b/tests/integration/test_e1_rag.py new file mode 100644 index 0000000..b821541 --- /dev/null +++ b/tests/integration/test_e1_rag.py @@ -0,0 +1,567 @@ +from __future__ import annotations + +import io +import json +import subprocess +import tarfile +import time +import urllib.error +import pytest +import torch +from datetime import UTC, datetime, timedelta +from email.message import Message +from email.utils import format_datetime + +from fastplms.embeddings import EmbeddingResult, load_sqlite_result +from fastplms.models.e1 import retrieval as e1_retrieval +from fastplms.models.e1.modeling_e1 import ( + E1_MSA_SAMPLING_SOURCE_REVISION, + ColabFoldSearcher, + ContextCache, + ContextSpecification, + E1Config, + E1ForMaskedLM, + E1ForSequenceClassification, + E1ForTokenClassification, + HomologueSearcher, + _safe_extract_tar, + get_msa_for_sequence, + get_query_from_a3m, + load_msa_dir, + parse_msa, + sample_context, + sample_multiple_contexts, +) +from fastplms.registry import load_model_registry + + +E1_SAMPLING_GOLDEN_REVISION = "bfd2620a602248499f3d2583d85a7ecddf0b6e02" +E1_SAMPLING_GOLDEN_PROVENANCE = { + "upstream_revision": E1_SAMPLING_GOLDEN_REVISION, + "source_path": "src/E1/msa_sampling.py", + "source_sha256": "9a2acc1932fe494613bbc8de0bea415c075f9e21eaa0caa8eb2b693410471e48", + "generation_command": [ + "python", + "tools/goldens/generate_e1_sampling.py", + "--upstream-root", + "vendor/upstream/e1", + "--output", + "artifacts/goldens/e1_sampling.json", + ], +} +E1_SINGLE_CONTEXT_GOLDEN = { + 0: ("TCDFGHI,ACDEYGH,ACEFGHI", ["mid", "near", "gapped"]), + 3: ("ACDEYGH,TTTTTTTT,TCDFGHI", ["near", "far", "mid"]), + 11: ("ACDEYGH,TCDFGHI,ACEFGHI", ["near", "mid", "gapped"]), +} +E1_MULTIPLE_CONTEXT_GOLDEN = ( + ["TCDFGHI,TTTTTTTT", "ACDEYGH,TCDFGHI,ACEFGHI,ACDEFGHI"], + [["mid", "far"], ["near", "mid", "gapped", "query"]], +) + + +def _write_tiny_a3m(path) -> None: + path.write_text( + ">query\nACDEFG\n>near\nACDEYG\n>far\nTTTTTT\n", + encoding="utf-8", + ) + + +def _write_parity_a3m(path) -> None: + path.write_text( + ">query\nACDEFGHI\n>near\nACDEYGH-\n>gapped\nAC-EFGHI\n>mid\nTCD-FGHI\n>far\nTTTTTTTT\n", + encoding="utf-8", + ) + + +def _tiny_e1_config() -> E1Config: + return E1Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + max_num_sequences=8, + max_num_positions_within_seq=64, + max_num_positions_global=256, + dtype="float32", + ) + + +def _tiny_e1_model(device: torch.device) -> E1ForMaskedLM: + return E1ForMaskedLM(config=_tiny_e1_config()).eval().to(device) + + +def test_a3m_parsing_query_lookup_and_context_sampling(tmp_path) -> None: + a3m_path = tmp_path / "query.a3m" + _write_tiny_a3m(a3m_path) + + records = parse_msa(str(a3m_path)) + assert [record.id for record in records] == ["query", "near", "far"] + assert get_query_from_a3m(str(a3m_path)) == "ACDEFG" + + msa_lookup = load_msa_dir(str(tmp_path)) + assert msa_lookup["ACDEFG"] == str(a3m_path) + assert get_msa_for_sequence("ACDEYG", msa_lookup, min_identity=0.80) == str(a3m_path) + + context, ids = sample_context( + msa_path=str(a3m_path), + max_num_samples=1, + max_token_length=32, + max_query_similarity=0.1, + min_query_similarity=0.0, + seed=0, + device=torch.device("cpu"), + ) + assert context == "TTTTTT" + assert ids == ["far"] + + +def test_context_sampling_matches_pinned_e1_golden(tmp_path) -> None: + registry = load_model_registry() + assert registry.upstreams["e1"].revision == E1_SAMPLING_GOLDEN_REVISION + assert E1_MSA_SAMPLING_SOURCE_REVISION == E1_SAMPLING_GOLDEN_REVISION + assert E1_SAMPLING_GOLDEN_PROVENANCE == { + "upstream_revision": E1_SAMPLING_GOLDEN_REVISION, + "source_path": "src/E1/msa_sampling.py", + "source_sha256": "9a2acc1932fe494613bbc8de0bea415c075f9e21eaa0caa8eb2b693410471e48", + "generation_command": [ + "python", + "tools/goldens/generate_e1_sampling.py", + "--upstream-root", + "vendor/upstream/e1", + "--output", + "artifacts/goldens/e1_sampling.json", + ], + } + + a3m_path = tmp_path / "parity.a3m" + _write_parity_a3m(a3m_path) + + kwargs = { + "msa_path": str(a3m_path), + "max_num_samples": 3, + "max_token_length": 32, + "max_query_similarity": 0.99, + "min_query_similarity": 0.0, + "neighbor_similarity_lower_bound": 0.8, + "device": torch.device("cpu"), + } + for seed in (0, 3, 11): + assert sample_context(seed=seed, **kwargs) == E1_SINGLE_CONTEXT_GOLDEN[seed] + + specs = [ + ContextSpecification( + max_num_samples=3, + max_token_length=16, + max_query_similarity=0.99, + min_query_similarity=0.0, + neighbor_similarity_lower_bound=0.8, + ), + ContextSpecification( + max_num_samples=4, + max_token_length=32, + max_query_similarity=1.0, + min_query_similarity=0.2, + neighbor_similarity_lower_bound=0.8, + ), + ] + contexts, ids = sample_multiple_contexts( + msa_path=str(a3m_path), + context_specifications=specs, + seed=7, + device=torch.device("cpu"), + ) + assert (contexts, ids) == E1_MULTIPLE_CONTEXT_GOLDEN + + +def test_context_cache_round_trip(tmp_path) -> None: + cache = ContextCache(str(tmp_path), specs_hash="abc123", seed=7) + assert cache.load("msa") is None + cache.store("msa", {"ctx": "ACDEFG"}) + assert cache.load("msa") == {"ctx": "ACDEFG"} + cache_files = list(tmp_path.glob("*.json")) + assert len(cache_files) == 1 + payload = json.loads(cache_files[0].read_text(encoding="utf-8")) + assert payload["source_revision"] == E1_SAMPLING_GOLDEN_REVISION + assert payload["contexts"] == {"ctx": "ACDEFG"} + + +def test_context_cache_invalidates_when_msa_content_or_revision_changes(tmp_path) -> None: + a3m_path = tmp_path / "query.a3m" + _write_tiny_a3m(a3m_path) + cache = ContextCache(str(tmp_path / "cache"), specs_hash="abc123", seed=7) + cache.store(str(a3m_path), {"ctx": "ACDEFG"}) + assert cache.load(str(a3m_path)) == {"ctx": "ACDEFG"} + + a3m_path.write_text(">query\nTTTTTT\n", encoding="utf-8") + assert cache.load(str(a3m_path)) is None + cache.store(str(a3m_path), {"ctx": "TTTTTT"}) + assert ( + ContextCache( + str(tmp_path / "cache"), + specs_hash="abc123", + seed=7, + source_revision="different-revision", + ).load(str(a3m_path)) + is None + ) + + +def test_safe_tar_extraction_blocks_traversal(tmp_path) -> None: + tar_path = tmp_path / "bad.tar" + payload = b"bad" + with tarfile.open(tar_path, "w") as tar: + info = tarfile.TarInfo("../escape.a3m") + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + + with tarfile.open(tar_path) as tar, pytest.raises(ValueError): + _safe_extract_tar(tar, str(tmp_path / "out")) + + +@pytest.mark.parametrize("entry_type", [tarfile.SYMTYPE, tarfile.LNKTYPE, tarfile.FIFOTYPE]) +def test_safe_tar_extraction_rejects_links_and_devices(tmp_path, entry_type) -> None: + tar_path = tmp_path / "unsafe.tar" + with tarfile.open(tar_path, "w") as tar: + info = tarfile.TarInfo("unsafe-entry") + info.type = entry_type + if entry_type in {tarfile.SYMTYPE, tarfile.LNKTYPE}: + info.linkname = "../escape.a3m" + tar.addfile(info) + + with tarfile.open(tar_path) as tar, pytest.raises(ValueError): + _safe_extract_tar(tar, str(tmp_path / "out")) + + +def test_public_e1_rag_methods_exist() -> None: + model = _tiny_e1_model(torch.device("cpu")) + methods = [ + model.search_homologues, + model.batch_search_homologues, + model.sample_msa_contexts, + model.score_ppll, + model.embed_with_msa, + model.embed_dataset_with_msa, + ] + for method in methods: + assert callable(method) + assert callable(model.embed_dataset) + + +@pytest.mark.parametrize( + "seq_id", + ["../escape", "nested/query", r"..\escape", "/tmp/escape", r"C:\escape"], +) +def test_homologue_searchers_reject_seq_ids_outside_output_dir(tmp_path, seq_id) -> None: + output_dir = str(tmp_path / "msas") + with pytest.raises(ValueError, match="seq_id"): + ColabFoldSearcher().search("ACDEFG", output_dir, seq_id=seq_id) + with pytest.raises(ValueError, match="seq_id"): + HomologueSearcher(target_db="target_db").search("ACDEFG", output_dir, seq_id=seq_id) + + +def test_e1_sequence_classifier_accepts_explicit_pooling_types() -> None: + model = E1ForSequenceClassification( + _tiny_e1_config(), + pooling_types=["mean"], + ) + assert model.pooler.names == ("mean",) + assert model.classifier[0].in_features == model.config.hidden_size + + +def test_e1_token_classifier_exactly_consumes_official_encoder_width() -> None: + """The extension head consumes the pinned encoder's unmodified H tensor.""" + + config = _tiny_e1_config() + config.num_labels = 3 + model = E1ForTokenClassification(config).eval() + prepared = model.prep_tokens.get_batch_kwargs(["MSTNPKPQ"], device=torch.device("cpu")) + inputs: dict[str, torch.Tensor] = {} + for name in ( + "input_ids", + "within_seq_position_ids", + "global_position_ids", + "sequence_ids", + ): + value = prepared[name] + assert isinstance(value, torch.Tensor) + inputs[name] = value + + with torch.inference_mode(): + encoder_output = model.model(**inputs) + output = model(**inputs) + expected_logits = model.classifier(encoder_output.last_hidden_state) + + assert model.classifier[0].in_features == config.hidden_size + assert output.last_hidden_state.shape[-1] == config.hidden_size + assert torch.equal(output.last_hidden_state, encoder_output.last_hidden_state) + assert torch.equal(output.logits, expected_logits) + + +def test_colabfold_http_errors_are_contextual_and_not_retried( + monkeypatch, +) -> None: + searcher = ColabFoldSearcher(max_retries=1, base_delay=0.0, max_delay=0.0) + headers = Message() + headers["Content-Type"] = "text/plain" + sleeps = [] + + def fail_request(*args, **kwargs): + raise urllib.error.HTTPError( + "https://api.colabfold.com/missing", + 404, + "Not Found", + headers, + io.BytesIO(b"missing"), + ) + + monkeypatch.setattr(e1_retrieval.urllib.request, "urlopen", fail_request) + monkeypatch.setattr(e1_retrieval.time, "sleep", sleeps.append) + with pytest.raises(RuntimeError, match="HTTP 404"): + searcher._request_with_retries("GET", "https://api.colabfold.com/missing") + assert sleeps == [] + + +def test_colabfold_retry_after_is_case_insensitive_date_aware_and_capped() -> None: + searcher = ColabFoldSearcher(max_delay=5.0) + assert searcher._retry_after_delay({"retry-after": "120"}, attempt=0) == 5.0 + retry_at = format_datetime(datetime.now(UTC) + timedelta(minutes=2)) + assert searcher._retry_after_delay({"Retry-After": retry_at}, attempt=0) == 5.0 + + +def test_colabfold_request_respects_expired_deadline(monkeypatch) -> None: + searcher = ColabFoldSearcher(max_retries=1, max_wait_time=1) + + def unexpected_request(*args, **kwargs): + raise AssertionError("expired requests must not reach the network") + + monkeypatch.setattr(e1_retrieval.urllib.request, "urlopen", unexpected_request) + with pytest.raises(TimeoutError, match="deadline"): + searcher._request_with_retries( + "GET", + "https://api.colabfold.com/ticket/expired", + deadline=time.monotonic() - 1.0, + ) + + +def test_mmseqs_searcher_subprocess_path_is_mockable(tmp_path, monkeypatch) -> None: + monkeypatch.chdir(tmp_path) + (tmp_path / "target_db.dbtype").write_bytes(b"test-db") + searcher = HomologueSearcher(target_db="target_db") + calls = [] + + identity = e1_retrieval._DockerImageIdentity( + reference=e1_retrieval.DOCKER_IMAGE, + repository=e1_retrieval.MMSEQS2_IMAGE_REPOSITORY, + version=e1_retrieval.MMSEQS2_VERSION, + manifest_digest=e1_retrieval.MMSEQS2_CPU_MANIFEST_DIGEST, + image_id="sha256:" + "a" * 64, + os="linux", + architecture=e1_retrieval._docker_architecture(), + ) + + def fake_run(cmd, **kwargs): + del kwargs + calls.append(cmd) + if "result2msa" in cmd: + command_index = cmd.index("result2msa") + output = tmp_path / cmd[command_index + 4] + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(">query\nACDEFG\n", encoding="utf-8") + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(searcher, "_ensure_docker_image", lambda: identity) + monkeypatch.setattr(searcher, "_run_docker_command", fake_run) + + a3m_path = searcher.search("ACDEFG", output_dir="msas", seq_id="query") + + assert a3m_path == "msas/query/query.a3m" + assert any("createdb" in call for call in calls) + assert any("search" in call for call in calls) + assert any("result2msa" in call for call in calls) + assert (tmp_path / "msas/query/search-provenance.json").is_file() + + +def test_colabfold_searcher_http_path_is_mockable(tmp_path, monkeypatch) -> None: + searcher = ColabFoldSearcher(inter_request_delay=(0.0, 0.0)) + + def fake_download(ticket_id: str, output_path: str) -> None: + payload = b">query\nACDEFG\n" + with tarfile.open(output_path, "w:gz") as tar: + info = tarfile.TarInfo("uniref.a3m") + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + + monkeypatch.setattr( + searcher, + "_submit", + lambda sequence, deadline=None: {"status": "RUNNING", "id": "ticket"}, + ) + monkeypatch.setattr( + searcher, + "_poll", + lambda ticket_id, deadline=None: {"status": "COMPLETE"}, + ) + monkeypatch.setattr( + searcher, + "_download", + lambda ticket_id, output_path, deadline=None: fake_download(ticket_id, output_path), + ) + + a3m_path = searcher.search("ACDEFG", str(tmp_path), seq_id="query") + + assert a3m_path.endswith("query.a3m") + assert get_query_from_a3m(a3m_path) == "ACDEFG" + + +@pytest.mark.parametrize( + ("searcher", "provider"), + ( + (HomologueSearcher(target_db="target_db"), "mmseqs2"), + (ColabFoldSearcher(inter_request_delay=(0.0, 0.0)), "colabfold"), + ), +) +def test_batch_search_warns_on_partial_failure_without_sequence_leak( + searcher, + provider, + tmp_path, + monkeypatch, + caplog, +) -> None: + sensitive_sequence = "SENSITIVESEQUENCE" + monkeypatch.chdir(tmp_path) + + def fail_search(sequence, output_dir, seq_id=None): + raise RuntimeError(f"provider failure included {sequence}") + + monkeypatch.setattr(searcher, "search", fail_search) + with caplog.at_level("WARNING", logger=e1_retrieval.__name__): + result = searcher.batch_search( + [sensitive_sequence], + "results", + seq_ids=["public-seq-id"], + continue_on_error=True, + ) + + assert result == {} + assert provider in caplog.text + assert "public-seq-id" in caplog.text + assert "RuntimeError" in caplog.text + assert sensitive_sequence not in caplog.text + + +@pytest.mark.parametrize( + "searcher", + ( + HomologueSearcher(target_db="target_db"), + ColabFoldSearcher(inter_request_delay=(0.0, 0.0)), + ), +) +def test_batch_search_preserves_failure_when_continue_is_disabled( + searcher, + tmp_path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + + def fail_search(sequence, output_dir, seq_id=None): + raise RuntimeError("provider failure") + + monkeypatch.setattr(searcher, "search", fail_search) + with pytest.raises(RuntimeError, match="provider failure"): + searcher.batch_search( + ["ACDEFG"], + "results", + seq_ids=["query"], + continue_on_error=False, + ) + + +@pytest.mark.gpu +def test_e1_score_ppll_with_tiny_synthetic_msa(tmp_path) -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _tiny_e1_model(device) + a3m_path = tmp_path / "query.a3m" + _write_tiny_a3m(a3m_path) + + scores = model.score_ppll( + ["ACDEFG"], + a3m_path=str(a3m_path), + max_context_tokens=[64], + similarity_thresholds=[1.0], + min_query_similarity=0.0, + progress=False, + ) + + assert len(scores) == 1 + assert 0.0 <= scores[0] <= 1.0 + + per_context_scores = model.score_ppll( + ["ACDEFG"], + a3m_path=str(a3m_path), + ensemble=False, + max_context_tokens=[64, 128], + similarity_thresholds=[1.0], + min_query_similarity=0.0, + progress=False, + ) + assert len(per_context_scores) == 1 + assert len(per_context_scores[0]) == 2 + for score in per_context_scores[0]: + assert 0.0 <= score <= 1.0 + + +@pytest.mark.gpu +def test_e1_embed_with_msa_shapes(tmp_path) -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _tiny_e1_model(device) + a3m_path = tmp_path / "query.a3m" + _write_tiny_a3m(a3m_path) + + pooled = model.embed_with_msa( + ["ACDEFG"], + a3m_path=str(a3m_path), + pooling_types=["mean"], + progress=False, + ) + matrix = model.embed_with_msa( + ["ACDEFG"], + a3m_path=str(a3m_path), + matrix_embed=True, + progress=False, + ) + + assert pooled.shape == (1, model.config.hidden_size) + assert len(matrix) == 1 + assert matrix[0].shape == (6, model.config.hidden_size) + + +@pytest.mark.gpu +def test_e1_embed_dataset_with_msa_falls_back_without_msa(tmp_path) -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _tiny_e1_model(device) + output = tmp_path / "e1-msa.sqlite" + + embeddings = model.embed_dataset_with_msa( + ["ACDEFG", "ACDEFG"], + msa_lookup={}, + batch_size=1, + max_len=16, + pooling_types=["mean"], + progress=False, + embed_dtype=torch.float32, + output=output, + format="sqlite", + ) + + assert isinstance(embeddings, EmbeddingResult) + assert [(record.id, record.sequence) for record in embeddings] == [ + ("0", "ACDEFG"), + ("1", "ACDEFG"), + ] + assert all(record.load_tensor().shape == (model.config.hidden_size,) for record in embeddings) + assert embeddings.metadata["descriptor_index"] == "sqlite-records" + assert embeddings.metadata["family_adapter"]["kind"] == "e1-msa-v1" + reopened = load_sqlite_result(output) + assert [record.sequence for record in reopened] == ["ACDEFG", "ACDEFG"] diff --git a/tests/integration/test_esm3.py b/tests/integration/test_esm3.py new file mode 100644 index 0000000..6a7d657 --- /dev/null +++ b/tests/integration/test_esm3.py @@ -0,0 +1,791 @@ +import base64 +import hashlib +import io +import json +import os +import runpy +import stat +import subprocess +import sys +import textwrap +import pytest +import torch +from pathlib import Path +from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo +from transformers import AutoModel + +import fastplms.models.esm3.modeling_esm3 as esm3_module +from fastplms.models.esm3.modeling_esm3 import ( + _MAX_SAVED_RUNTIME_FILE_BYTES, + _SAVED_RUNTIME_FILES, + SEQUENCE_BOS_TOKEN, + SEQUENCE_EOS_TOKEN, + SEQUENCE_MASK_TOKEN, + FastESM3Config, + FastESM3GenerationConfig, + FastESM3Model, + _build_saved_runtime_archive, + _render_saved_runtime_bundle, + _saved_runtime_tree_hash, + _validate_saved_runtime_relative_path, +) +from tests.conftest import strict_fp32_matmul + + +def _small_config() -> FastESM3Config: + return FastESM3Config( + hidden_size=64, + num_attention_heads=4, + num_vector_heads=8, + num_hidden_layers=2, + ) + + +def _small_model() -> FastESM3Model: + return FastESM3Model(_small_config()).eval() + + +def _write_synthetic_runtime(package_root: Path) -> None: + for relative in _SAVED_RUNTIME_FILES: + path = package_root.joinpath(*relative.split("/")) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(f"validated runtime source: {relative}\n".encode()) + + +def _saved_bundle_namespace(model_path: Path) -> dict[str, object]: + return runpy.run_path(str(model_path / "fastplms_bundle.py")) + + +def _rewrite_saved_runtime_archive( + model_path: Path, + *, + first_name: str | None = None, + first_mode: int | None = None, + corrupt_first_hash: bool = False, +) -> None: + namespace = _saved_bundle_namespace(model_path) + old_archive_hash = namespace["RUNTIME_HASH"] + old_tree_hash = namespace["RUNTIME_TREE_HASH"] + manifest = json.loads(json.dumps(namespace["RUNTIME_MANIFEST"])) + archive = base64.b85decode("".join(namespace["RUNTIME_DATA"])) + + if corrupt_first_hash: + first_relative = _SAVED_RUNTIME_FILES[0] + manifest["files"][first_relative]["sha256"] = "0" * 64 + buffer = io.BytesIO() + with ( + ZipFile(io.BytesIO(archive)) as source, + ZipFile( + buffer, + mode="w", + compression=ZIP_DEFLATED, + compresslevel=9, + ) as destination, + ): + for index, member in enumerate(source.infolist()): + name = first_name if index == 0 and first_name is not None else member.filename + info = ZipInfo(name, date_time=member.date_time) + info.create_system = member.create_system + info.compress_type = member.compress_type + info.external_attr = ( + first_mode << 16 if index == 0 and first_mode is not None else member.external_attr + ) + destination.writestr( + info, + source.read(member), + compress_type=ZIP_DEFLATED, + compresslevel=9, + ) + poisoned_archive = buffer.getvalue() + new_tree_hash = _saved_runtime_tree_hash(manifest) + new_archive_hash, bundle = _render_saved_runtime_bundle( + poisoned_archive, + manifest, + new_tree_hash, + ) + (model_path / "fastplms_bundle.py").write_bytes(bundle) + bridge_path = model_path / "modeling_fastplms.py" + bridge = bridge_path.read_text(encoding="utf-8") + assert isinstance(old_archive_hash, str) + assert isinstance(old_tree_hash, str) + bridge = bridge.replace(old_archive_hash, new_archive_hash).replace( + old_tree_hash, + new_tree_hash, + ) + bridge_path.write_text(bridge, encoding="utf-8", newline="\n") + + +def _run_isolated_bridge_probe( + model_path: Path, + tmp_path: Path, + body: str = "", +) -> subprocess.CompletedProcess[str]: + probe = tmp_path / "esm3_runtime_bridge_probe.py" + source = textwrap.dedent( + """\ + import importlib.abc + import importlib.util + import sys + import types + from pathlib import Path + + + class BlockInstalledFastPLMs(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastplms": + raise ModuleNotFoundError("external FastPLMs source is blocked") + return None + + + artifact = Path(sys.argv[1]) + sys.modules.pop("fastplms", None) + for name in tuple(sys.modules): + if name.startswith("fastplms."): + sys.modules.pop(name) + sys.meta_path.insert(0, BlockInstalledFastPLMs()) + + def load_bridge(package_name): + package = types.ModuleType(package_name) + package.__package__ = package_name + package.__path__ = [str(artifact)] + sys.modules[package_name] = package + module_name = f"{package_name}.modeling_fastplms" + spec = importlib.util.spec_from_file_location( + module_name, + artifact / "modeling_fastplms.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load generated ESM3 bridge") + bridge = importlib.util.module_from_spec(spec) + sys.modules[module_name] = bridge + spec.loader.exec_module(bridge) + return bridge + """ + ) + if body.strip(): + source += "\n" + textwrap.dedent(body).strip() + "\n" + probe.write_text( + source, + encoding="utf-8", + newline="\n", + ) + environment = os.environ.copy() + environment.pop("PYTHONPATH", None) + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + return subprocess.run( + [sys.executable, "-I", str(probe), str(model_path)], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=120, + check=False, + ) + + +def test_esm3_sequence_only_forward() -> None: + model = _small_model() + batch = model.tokenize_sequences(["MKTAYIAKQ", "GGGG"], device=model.device) + + with torch.inference_mode(): + output = model(**batch) + + assert output.logits is not None + assert output.function_logits is not None + assert output.residue_logits is not None + assert output.logits.shape[:2] == batch["input_ids"].shape + assert output.logits.shape[-1] == 64 + assert output.structure_logits.shape[-1] == 4096 + assert output.function_logits.shape[-2:] == (8, 260) + assert output.residue_logits.shape[-1] == 1478 + assert not torch.isnan(output.logits).any() + + +def test_esm3_uses_hugging_face_initialization_and_only_retains_requested_states() -> None: + model = _small_model() + model.attn_backend = "eager" + batch = model.tokenize_sequences(["MKTAYIAKQ"], device=model.device) + + embedding_std = model.esm3.encoder.sequence_embed.weight.detach().std().item() + assert embedding_std == pytest.approx(model.config.initializer_range, rel=0.2) + + with torch.inference_mode(): + default_output = model(**batch) + full_output = model( + **batch, + output_hidden_states=True, + output_attentions=True, + ) + tuple_output = model( + **batch, + output_hidden_states=True, + output_attentions=True, + return_dict=False, + ) + + assert default_output.hidden_states is None + assert full_output.hidden_states is not None + assert full_output.attentions is not None + assert len(full_output.hidden_states) == model.config.num_hidden_layers + assert len(full_output.attentions) == model.config.num_hidden_layers + assert tuple(full_output.keys()) == ( + "last_hidden_state", + "hidden_states", + "attentions", + "logits", + "sequence_logits", + "structure_logits", + "secondary_structure_logits", + "sasa_logits", + "function_logits", + "residue_logits", + "embeddings", + ) + assert isinstance(tuple_output, tuple) + assert torch.equal(tuple_output[0], full_output.last_hidden_state) + assert isinstance(tuple_output[1], tuple) + assert isinstance(tuple_output[2], tuple) + for actual, expected in zip(tuple_output, full_output.to_tuple(), strict=True): + if isinstance(expected, tuple): + assert isinstance(actual, tuple) + for actual_tensor, expected_tensor in zip(actual, expected, strict=True): + torch.testing.assert_close(actual_tensor, expected_tensor) + continue + torch.testing.assert_close(actual, expected) + + # labels: (b, l) + labels = batch["input_ids"].clone() + with torch.inference_mode(): + labeled_output = model( + **batch, + labels=labels, + output_hidden_states=True, + output_attentions=True, + ) + labeled_tuple = model( + **batch, + labels=labels, + output_hidden_states=True, + output_attentions=True, + return_dict=False, + ) + assert tuple(labeled_output.keys())[:3] == ( + "loss", + "last_hidden_state", + "hidden_states", + ) + assert labeled_output.loss is not None + assert torch.isfinite(labeled_output.loss) + torch.testing.assert_close(labeled_tuple[0], labeled_output.loss) + torch.testing.assert_close(labeled_tuple[1], labeled_output.last_hidden_state) + assert len(labeled_tuple) == len(labeled_output.to_tuple()) + for actual, expected in zip( + labeled_tuple, + labeled_output.to_tuple(), + strict=True, + ): + if isinstance(expected, tuple): + assert isinstance(actual, tuple) + for actual_tensor, expected_tensor in zip(actual, expected, strict=True): + torch.testing.assert_close(actual_tensor, expected_tensor) + continue + torch.testing.assert_close(actual, expected) + + +def test_esm3_config_drives_hidden_state_output() -> None: + config = _small_config() + config.output_hidden_states = True + model = FastESM3Model(config).eval() + batch = model.tokenize_sequences(["MKTAYIAKQ"], device=model.device) + + with torch.inference_mode(): + output = model(**batch) + + assert output.hidden_states is not None + assert len(output.hidden_states) == config.num_hidden_layers + + +def test_esm3_resize_updates_sequence_input_and_output_embeddings() -> None: + model = _small_model() + original_vocab_size = model.config.vocab_size + resized_vocab_size = original_vocab_size + 7 + + model.resize_token_embeddings(resized_vocab_size) + + assert model.get_input_embeddings().num_embeddings == resized_vocab_size + assert model.get_output_embeddings().out_features == resized_vocab_size + assert model.config.vocab_size == resized_vocab_size + with torch.inference_mode(): + output = model( + input_ids=torch.tensor([[SEQUENCE_BOS_TOKEN, original_vocab_size, SEQUENCE_EOS_TOKEN]]) + ) + assert output.sequence_logits.shape[-1] == resized_vocab_size + + +def test_esm3_accepts_function_tokens_argument() -> None: + model = _small_model() + batch = model.tokenize_sequences(["MKTAYIAKQ"], device=model.device) + # function_tokens: (...) + function_tokens = batch["input_ids"].new_zeros((*batch["input_ids"].shape, 8)) + + with torch.inference_mode(): + output = model(**batch, function_tokens=function_tokens) + + assert output.logits is not None + assert output.logits.shape[:2] == batch["input_ids"].shape + + +def test_esm3_rejects_attention_mask_row_without_a_valid_key() -> None: + model = _small_model() + # input_ids: (2, 3) + input_ids = torch.tensor( + [ + [SEQUENCE_BOS_TOKEN, 4, SEQUENCE_EOS_TOKEN], + [SEQUENCE_BOS_TOKEN, 5, SEQUENCE_EOS_TOKEN], + ] + ) + # attention_mask: (2, 3) + attention_mask = torch.tensor([[1, 1, 1], [0, 0, 0]]) + + with pytest.raises( + ValueError, + match="attention_mask must keep at least one valid key per batch row", + ): + model(input_ids=input_ids, attention_mask=attention_mask) + + +def test_esm3_saved_runtime_archive_is_fixed_bounded_and_deterministic( + tmp_path: Path, +) -> None: + package_root = tmp_path / "fastplms" + _write_synthetic_runtime(package_root) + (package_root / ".secrets.env").write_text("TOKEN=excluded\n", encoding="utf-8") + injected = package_root / "models" / "esm3" / "untracked_injected.py" + injected.write_text("raise RuntimeError('must not be bundled')\n", encoding="utf-8") + external = tmp_path / "external.py" + external.write_text("raise RuntimeError('must not be followed')\n", encoding="utf-8") + (package_root / "unknown_symlink.py").symlink_to(external) + + first_archive, first_manifest, first_tree_hash = _build_saved_runtime_archive(package_root) + second_archive, second_manifest, second_tree_hash = _build_saved_runtime_archive(package_root) + + assert first_archive == second_archive + assert first_manifest == second_manifest + assert first_tree_hash == second_tree_hash + assert first_manifest["schema_version"] == 1 + assert set(first_manifest["files"]) == set(_SAVED_RUNTIME_FILES) + assert first_manifest["total_size"] == sum( + record["size"] for record in first_manifest["files"].values() + ) + assert ".secrets.env" not in first_manifest["files"] + assert "models/esm3/untracked_injected.py" not in first_manifest["files"] + assert "unknown_symlink.py" not in first_manifest["files"] + with ZipFile(io.BytesIO(first_archive)) as archive: + assert set(archive.namelist()) == { + f"fastplms/{relative}" for relative in _SAVED_RUNTIME_FILES + } + for relative, record in first_manifest["files"].items(): + payload = archive.read(f"fastplms/{relative}") + assert len(payload) == record["size"] + assert hashlib.sha256(payload).hexdigest() == record["sha256"] + + +@pytest.mark.parametrize( + "value", + ( + "../escape.py", + "/absolute.py", + "C:/escape.py", + "models\\escape.py", + "models//escape.py", + ), +) +def test_esm3_saved_runtime_rejects_noncanonical_allowlist_paths(value: str) -> None: + with pytest.raises(RuntimeError, match="runtime path is unsafe"): + _validate_saved_runtime_relative_path(value) + + +def test_esm3_saved_runtime_rejects_missing_allowlisted_file(tmp_path: Path) -> None: + package_root = tmp_path / "fastplms" + _write_synthetic_runtime(package_root) + package_root.joinpath(*_SAVED_RUNTIME_FILES[-1].split("/")).unlink() + + with pytest.raises(RuntimeError, match="runtime file is missing"): + _build_saved_runtime_archive(package_root) + + +def test_esm3_saved_runtime_rejects_allowlisted_symlink(tmp_path: Path) -> None: + package_root = tmp_path / "fastplms" + _write_synthetic_runtime(package_root) + target = tmp_path / "outside.py" + target.write_text("outside = True\n", encoding="utf-8") + allowlisted = package_root / "__init__.py" + allowlisted.unlink() + allowlisted.symlink_to(target) + + with pytest.raises(RuntimeError, match="must not contain a symlink"): + _build_saved_runtime_archive(package_root) + + +def test_esm3_saved_runtime_rejects_oversize_allowlisted_file(tmp_path: Path) -> None: + package_root = tmp_path / "fastplms" + _write_synthetic_runtime(package_root) + (package_root / "__init__.py").write_bytes(b"x" * (_MAX_SAVED_RUNTIME_FILE_BYTES + 1)) + + with pytest.raises(RuntimeError, match="exceeds its size limit"): + _build_saved_runtime_archive(package_root) + + +def test_esm3_saved_runtime_rejects_oversize_total( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + package_root = tmp_path / "fastplms" + _write_synthetic_runtime(package_root) + monkeypatch.setattr(esm3_module, "_MAX_SAVED_RUNTIME_TOTAL_BYTES", 1) + + with pytest.raises(RuntimeError, match="total expanded size limit"): + _build_saved_runtime_archive(package_root) + + +def test_esm3_loads_with_automodel(tmp_path: Path) -> None: + model = _small_model() + model.save_pretrained(tmp_path) + assert (tmp_path / "modeling_esm3.py").is_file() + assert (tmp_path / "modeling_fastplms.py").is_file() + assert (tmp_path / "fastplms_bundle.py").is_file() + assert not (tmp_path / "fastplms").exists() + assert not list(tmp_path.glob("_fastplms_runtime_*")) + bundle = _saved_bundle_namespace(tmp_path) + assert set(bundle["RUNTIME_MANIFEST"]["files"]) == set(_SAVED_RUNTIME_FILES) + bridge = (tmp_path / "modeling_fastplms.py").read_text(encoding="utf-8") + assert "extractall" not in bridge + assert "rglob" not in bridge + assert "TemporaryDirectory" in bridge + config = json.loads((tmp_path / "config.json").read_text(encoding="utf-8")) + assert config["auto_map"] == { + "AutoConfig": "modeling_fastplms.FastESM3Config", + "AutoModel": "modeling_fastplms.FastESM3Model", + } + + loaded = AutoModel.from_pretrained(tmp_path, trust_remote_code=True).eval() + batch = loaded.tokenize_sequences(["MKTAYIAKQ"], device=loaded.device) + + with torch.inference_mode(): + output = loaded(**batch) + + assert output.logits is not None + assert output.logits.shape[:2] == batch["input_ids"].shape + + +def test_esm3_repeated_save_removes_stale_runtime_outputs(tmp_path: Path) -> None: + model = _small_model() + model.save_pretrained(tmp_path) + expected_bundle = (tmp_path / "fastplms_bundle.py").read_bytes() + expected_bridge = (tmp_path / "modeling_fastplms.py").read_bytes() + + stale_tree = tmp_path / "fastplms" / "models" / "esm3" + stale_tree.mkdir(parents=True) + (stale_tree / "poison.py").write_text("raise RuntimeError('stale')\n", encoding="utf-8") + stale_cache = tmp_path / "_fastplms_runtime_deadbeef" / "fastplms" + stale_cache.mkdir(parents=True) + (stale_cache / "__init__.py").write_text( + "raise RuntimeError('stale cache')\n", + encoding="utf-8", + ) + (tmp_path / "fastplms_bundle.py").write_text("stale bundle\n", encoding="utf-8") + (tmp_path / "modeling_fastplms.py").write_text("stale bridge\n", encoding="utf-8") + + model.save_pretrained(tmp_path) + + assert not (tmp_path / "fastplms").exists() + assert not (tmp_path / "_fastplms_runtime_deadbeef").exists() + assert (tmp_path / "fastplms_bundle.py").read_bytes() == expected_bundle + assert (tmp_path / "modeling_fastplms.py").read_bytes() == expected_bridge + + +def test_esm3_saved_bridge_reuses_same_runtime_in_process(tmp_path: Path) -> None: + model_path = tmp_path / "saved" + _small_model().save_pretrained(model_path) + result = _run_isolated_bridge_probe( + model_path, + tmp_path, + """ + first = load_bridge("artifact_first") + runtime = sys.modules["fastplms"] + runtime_root = Path(runtime.__file__).absolute().parent.parent + second = load_bridge("artifact_second") + assert sys.modules["fastplms"] is runtime + assert len(first._RUNTIME_TEMPORARIES) == 1 + assert second._RUNTIME_TEMPORARIES == [] + assert runtime.__fastplms_saved_runtime_tree_hash__ == first.RUNTIME_TREE_HASH + first._cleanup_runtime_temporaries() + assert not runtime_root.exists() + """, + ) + + assert result.returncode == 0, result.stdout + result.stderr + + +def test_esm3_saved_bridge_rejects_preimported_runtime_mismatch(tmp_path: Path) -> None: + model_path = tmp_path / "saved" + _small_model().save_pretrained(model_path) + result = _run_isolated_bridge_probe( + model_path, + tmp_path, + """ + fake_root = artifact / "mismatched_fastplms" + fake_root.mkdir() + fake_init = fake_root / "__init__.py" + fake_init.write_text("__version__ = 'different'\\n", encoding="utf-8") + fake = types.ModuleType("fastplms") + fake.__file__ = str(fake_init) + fake.__path__ = [str(fake_root)] + sys.modules["fastplms"] = fake + load_bridge("artifact_mismatch") + """, + ) + + assert result.returncode != 0 + assert "Loaded FastPLMs version/runtime mismatch" in result.stderr + + +@pytest.mark.parametrize( + ("poison", "message"), + ( + ({"first_name": "fastplms/../escape.py"}, "unsafe path"), + ({"first_name": "fastplms/unknown.py"}, "inventory is unexpected"), + ({"first_mode": stat.S_IFLNK | 0o777}, "member is not canonical"), + ({"corrupt_first_hash": True}, "member hash mismatch"), + ), +) +def test_esm3_saved_bridge_rejects_poisoned_archive( + tmp_path: Path, + poison: dict[str, object], + message: str, +) -> None: + model_path = tmp_path / "saved" + _small_model().save_pretrained(model_path) + _rewrite_saved_runtime_archive(model_path, **poison) + + result = _run_isolated_bridge_probe( + model_path, + tmp_path, + 'load_bridge("artifact_poisoned")', + ) + + assert result.returncode != 0 + assert message in result.stderr + + +def test_esm3_seeded_generation_is_repeatable_and_preserves_context() -> None: + model = _small_model() + config = FastESM3GenerationConfig(num_steps=2, temperature=1.0, seed=73) + + first = model.generate("MK__A", config) + second = model.generate("MK__A", config) + + assert isinstance(first, str) + assert first == second + assert len(first) == 5 + assert first[:2] == "MK" + assert first[-1] == "A" + assert "_" not in first + + +@pytest.mark.parametrize("num_steps", (0, -1)) +def test_esm3_generation_rejects_nonpositive_num_steps(num_steps: int) -> None: + model = _small_model() + + with pytest.raises(ValueError, match="num_steps must be positive"): + model.generate("M_K", FastESM3GenerationConfig(num_steps=num_steps)) + + +@pytest.mark.parametrize("num_steps", (True, False, 1.5, "2")) +def test_esm3_generation_rejects_noninteger_num_steps(num_steps: object) -> None: + model = _small_model() + + with pytest.raises(TypeError, match="num_steps must be an integer or None"): + model.generate("M_K", FastESM3GenerationConfig(num_steps=num_steps)) + + +def test_esm3_generation_none_num_steps_uses_mask_count() -> None: + model = _small_model() + observed_steps = 0 + original_forward = model.forward + + def count_forward_calls(*args, **kwargs): + nonlocal observed_steps + observed_steps += 1 + return original_forward(*args, **kwargs) + + model.forward = count_forward_calls + generated = model.generate("M__K", FastESM3GenerationConfig(num_steps=None, seed=19)) + + assert isinstance(generated, str) + assert "_" not in generated + assert observed_steps == 2 + + +def test_esm3_generation_preserves_every_supported_conditioning_track() -> None: + model = _small_model() + # input_ids: (1, 3) + input_ids = torch.tensor([[SEQUENCE_BOS_TOKEN, SEQUENCE_MASK_TOKEN, SEQUENCE_EOS_TOKEN]]) + shape = input_ids.shape + conditioning = { + "attention_mask": torch.ones(shape, dtype=torch.long), + "structure_tokens": torch.zeros(shape, dtype=torch.long), + "ss8_tokens": torch.zeros(shape, dtype=torch.long), + "sasa_tokens": torch.zeros(shape, dtype=torch.long), + "function_tokens": torch.zeros((*shape, 8), dtype=torch.long), + "residue_annotation_tokens": torch.zeros((*shape, 16), dtype=torch.long), + "average_plddt": torch.ones(shape), + "per_res_plddt": torch.zeros(shape), + "structure_coords": torch.full((*shape, 3, 3), float("nan")), + "chain_id": torch.zeros(shape, dtype=torch.long), + "sequence_id": torch.ones(shape, dtype=torch.bool), + } + observed: list[dict[str, torch.Tensor]] = [] + original_forward = model.forward + + def capture_forward(*args, **kwargs): + observed.append({name: value for name, value in kwargs.items() if torch.is_tensor(value)}) + return original_forward(*args, **kwargs) + + model.forward = capture_forward + generated = model.generate( + {"sequence_tokens": input_ids, **conditioning}, + FastESM3GenerationConfig(num_steps=1, seed=17), + ) + + assert torch.is_tensor(generated) + assert not generated.eq(SEQUENCE_MASK_TOKEN).any() + assert len(observed) == 1 + assert set(conditioning).issubset(observed[0]) + for name, expected in conditioning.items(): + torch.testing.assert_close(observed[0][name], expected, equal_nan=True) + + +def test_esm3_generation_rejects_unknown_or_ambiguous_inputs() -> None: + model = _small_model() + # input_ids: (1, 3) + input_ids = torch.tensor([[SEQUENCE_BOS_TOKEN, SEQUENCE_MASK_TOKEN, SEQUENCE_EOS_TOKEN]]) + + with pytest.raises(TypeError, match="Unsupported ESM3 generation inputs: labels"): + model.generate({"input_ids": input_ids, "labels": input_ids}) + with pytest.raises(ValueError, match="only one of input_ids or sequence_tokens"): + model.generate({"input_ids": input_ids, "sequence_tokens": input_ids}) + + +def test_esm3_saved_model_loads_without_installed_fastplms(tmp_path: Path) -> None: + model_path = tmp_path / "saved" + _small_model().save_pretrained(model_path) + runtime_hash = _saved_bundle_namespace(model_path)["RUNTIME_HASH"] + assert isinstance(runtime_hash, str) + poisoned_cache = model_path / f"_fastplms_runtime_{runtime_hash[:16]}" / "fastplms" + poisoned_cache.mkdir(parents=True) + (poisoned_cache / "__init__.py").write_text( + "raise RuntimeError('persistent model-local cache was trusted')\n", + encoding="utf-8", + ) + runtime_record = tmp_path / "runtime-path.txt" + script = textwrap.dedent( + """ + import importlib.abc + import sys + from pathlib import Path + + class BlockInstalledFastPLMs(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastplms": + raise ModuleNotFoundError("external FastPLMs source is blocked") + return None + + sys.modules.pop("fastplms", None) + for name in tuple(sys.modules): + if name.startswith("fastplms."): + sys.modules.pop(name) + sys.meta_path.insert(0, BlockInstalledFastPLMs()) + + import torch + from transformers import AutoModel + + model = AutoModel.from_pretrained( + sys.argv[1], + trust_remote_code=True, + local_files_only=True, + ).eval() + assert type(model).__module__ == "fastplms.models.esm3.modeling_esm3" + package_file = Path(sys.modules["fastplms"].__file__).resolve() + package_root = package_file.parent + runtime_root = package_root.parent + model_root = Path(sys.argv[1]).resolve() + assert runtime_root.name.startswith("fastplms-esm3-runtime-") + assert runtime_root != model_root + assert model_root not in runtime_root.parents + assert not any(path.name == "__pycache__" for path in package_root.rglob("*")) + Path(sys.argv[2]).write_text(str(runtime_root), encoding="utf-8") + batch = model.tokenize_sequences(["MKTAYIAKQ"], device=model.device) + with torch.inference_mode(): + output = model(**batch) + assert output.logits is not None + assert output.logits.shape[:2] == batch["input_ids"].shape + """ + ) + environment = os.environ.copy() + environment.pop("PYTHONPATH", None) + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + result = subprocess.run( + [sys.executable, "-I", "-c", script, str(model_path), str(runtime_record)], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=120, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + extracted_runtime = Path(runtime_record.read_text(encoding="utf-8")) + assert not extracted_runtime.exists() + assert ( + (poisoned_cache / "__init__.py") + .read_text(encoding="utf-8") + .startswith("raise RuntimeError") + ) + + +def test_esm3_embed_dataset(tmp_path: Path) -> None: + model = _small_model() + save_path = tmp_path / "embeddings" + + result = model.embed_dataset( + inputs=["MKTAYIAKQ", "GGGG"], + batch_size=2, + max_length=16, + pooling=("mean", "cls"), + output=save_path, + ) + + embeddings = result.as_dict(key="sequence") + assert set(embeddings) == {"MKTAYIAKQ", "GGGG"} + assert embeddings["MKTAYIAKQ"].shape == (128,) + assert (save_path / "index.json").is_file() + + +@pytest.mark.gpu +def test_esm3_flex_matches_sdpa() -> None: + model = _small_model().to(torch.device("cuda")) + batch = model.tokenize_sequences(["MKTAYIAKQ", "GGGG"], device=model.device) + + with torch.inference_mode(), strict_fp32_matmul(): + model.set_attn_implementation("sdpa") + sdpa_output = model(**batch).last_hidden_state + model.set_attn_implementation("flex_attention") + flex_output = model(**batch).last_hidden_state + + max_abs = (sdpa_output - flex_output).float().abs().max().item() + mse = ((sdpa_output - flex_output).float() ** 2).mean().item() + assert max_abs < 1e-4 + assert mse < 1e-8 diff --git a/tests/integration/test_flash_attention_backends.py b/tests/integration/test_flash_attention_backends.py new file mode 100644 index 0000000..1e9c9c6 --- /dev/null +++ b/tests/integration/test_flash_attention_backends.py @@ -0,0 +1,477 @@ +"""Execution contracts for precompiled Hugging Face FlashAttention kernels.""" + +from __future__ import annotations + +import importlib.util +import pytest +import torch +from collections.abc import Callable +from pathlib import Path +from typing import Any +from torch.nn import functional as F + +from fastplms.attention import _core +from fastplms.models.dplm.modeling_dplm import DPLMConfig, DPLMModel +from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmModel, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ESMplusplusModel, +) +from fastplms.registry import get_model_registry +from tools.debug.probe_flash_attention_forward import _model_results, _shared_results +from tools.debug.probe_flash_checkpoint_forward import _run_checkpoint + + +_FLASH_BACKENDS = ("flash_attention_2", "flash_attention_3") +_ESMC_FLASH_BACKENDS = _FLASH_BACKENDS + + +def _tiny_model_spec(family_id: str) -> tuple[type[torch.nn.Module], Any]: + if family_id == "esm2": + return FastEsmModel, FastEsmConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=128, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=64, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + token_dropout=False, + attn_backend="sdpa", + ) + if family_id == "esm_plusplus": + return ESMplusplusModel, ESMplusplusConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + dropout=0.0, + pad_token_id=1, + attn_backend="sdpa", + ) + raise AssertionError(f"Unexpected FlashAttention family: {family_id}") + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + # actual: (...), expected: (...) + relative_l2 = torch.linalg.vector_norm( + actual.float() - expected.float() + ) / torch.linalg.vector_norm(expected.float()).clamp_min(1e-12) + # minimum_cosine: () + minimum_cosine = F.cosine_similarity(actual.float(), expected.float(), dim=-1).min() + assert relative_l2.item() <= 1e-2 + assert minimum_cosine.item() >= 0.999 + + +def test_manifest_flash_inventory_is_exact() -> None: + registry = get_model_registry() + advertised = { + family.id: tuple(name for name in family.attention if name in _FLASH_BACKENDS) + for family in registry.families.values() + if set(family.attention).intersection(_FLASH_BACKENDS) + } + assert advertised == { + "esm2": _FLASH_BACKENDS, + "esm_plusplus": _ESMC_FLASH_BACKENDS, + "dplm": ("flash_attention_3",), + } + assert importlib.util.find_spec("flash_attn") is None + + +@pytest.mark.gpu +@pytest.mark.parametrize("family_id", ("esm2", "esm_plusplus")) +def test_explicit_flash_from_pretrained_uses_only_pinned_kernels( + family_id: str, + tmp_path: Path, +) -> None: + assert torch.cuda.is_available(), "FlashAttention integration requires CUDA." + model_class, config = _tiny_model_spec(family_id) + model_path = tmp_path / family_id + torch.manual_seed(29) + model_class(config).save_pretrained(model_path) + + # input_ids: (2, 17) + input_ids = torch.tensor( + ( + (0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2, 1, 1, 1, 1), + (0, 20, 19, 18, 17, 16, 15, 14, 2, 1, 1, 1, 1, 1, 1, 1, 1), + ), + device="cuda", + ) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + # reference: (...) + reference = ( + model_class.from_pretrained( + model_path, + attn_implementation="sdpa", + dtype=torch.bfloat16, + ) + .eval() + .to("cuda") + ) + with torch.inference_mode(): + expected = reference( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + + advertised = tuple( + name + for name in get_model_registry().families[family_id].attention + if name in _FLASH_BACKENDS + ) + expected_backends = _FLASH_BACKENDS if family_id == "esm2" else _ESMC_FLASH_BACKENDS + assert advertised == expected_backends + for backend in advertised: + _core._FLASH_KERNELS.pop(backend, None) + model = ( + model_class.from_pretrained( + model_path, + attn_implementation=backend, + dtype=torch.bfloat16, + ) + .eval() + .to("cuda") + ) + assert model.config._attn_implementation == backend + assert backend not in _core._FLASH_KERNELS + with torch.inference_mode(): + actual = model( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + assert backend in _core._FLASH_KERNELS + assert torch.isfinite(actual).all() + _assert_close(actual[attention_mask], expected[attention_mask]) + + +@pytest.mark.gpu +def test_precompiled_flash_kernels_match_sdpa_for_dense_and_mixed_padding() -> None: + assert torch.cuda.is_available(), "FlashAttention integration requires CUDA." + results = _shared_results() + for backend in _FLASH_BACKENDS: + assert results[backend]["dense"]["relative_l2"] <= 1e-2 + assert results[backend]["dense"]["minimum_cosine"] >= 0.999 + assert results[backend]["mixed_padding"]["relative_l2"] <= 1e-2 + assert results[backend]["mixed_padding"]["minimum_cosine"] >= 0.999 + assert results[backend]["padding_is_zero"] is True + + +@pytest.mark.gpu +@pytest.mark.parametrize("mixed_padding", (False, True), ids=("dense", "mixed-padding")) +def test_precompiled_flash_attention_2_dense_and_varlen_backward( + mixed_padding: bool, +) -> None: + """The pinned FA2 autograd wrappers propagate gradients through Q, K, and V.""" + + assert torch.cuda.is_available(), "FlashAttention integration requires CUDA." + torch.manual_seed(31) + # query: (2, 17, 4, 16) + query = torch.randn( + (2, 17, 4, 16), + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + key = torch.randn_like(query, requires_grad=True) + value = torch.randn_like(query, requires_grad=True) + attention_mask = None + if mixed_padding: + # attention_mask: (2, 17) + attention_mask = torch.tensor( + ( + (1,) * 13 + (0,) * 4, + (1,) * 9 + (0,) * 8, + ), + device="cuda", + dtype=torch.bool, + ) + + output = _core.kernels_flash_attention_func( + query, + key, + value, + attention_mask_2d=attention_mask, + implementation="flash_attention_2", + ) + selected = output if attention_mask is None else output[attention_mask] + selected.float().square().mean().backward() + + for tensor in (query, key, value): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + + +@pytest.mark.gpu +def test_flash_attention_2_mixed_padding_lora_step_and_reload( + tmp_path: Path, +) -> None: + """A real pinned FA2 kernel preserves a complete PEFT training graph.""" + + peft = pytest.importorskip("peft") + assert torch.cuda.is_available(), "FlashAttention PEFT integration requires CUDA." + + def make_base() -> FastEsmForMaskedLM: + config = FastEsmConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=128, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=64, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + token_dropout=False, + attn_backend="flash_attention_2", + ) + return FastEsmForMaskedLM(config) + + torch.manual_seed(41) + base = make_base() + base_state = { + name: tensor.detach().clone() + for name, tensor in base.state_dict().items() + } + model = peft.get_peft_model( + base, + peft.LoraConfig( + task_type=peft.TaskType.TOKEN_CLS, + r=2, + lora_alpha=4, + lora_dropout=0.0, + target_modules=["query", "value"], + ), + ).to("cuda").train() + model.base_model.model.set_attn_implementation("flash_attention_2") + + # input_ids: (2, 17) + input_ids = torch.tensor( + ( + (0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2, 1, 1, 1, 1), + (0, 20, 19, 18, 17, 16, 15, 14, 2, 1, 1, 1, 1, 1, 1, 1, 1), + ), + device="cuda", + ) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + # labels: (b, l) + labels = input_ids.masked_fill(~attention_mask, -100) + before = { + name: parameter.detach().clone() + for name, parameter in model.named_parameters() + } + optimizer = torch.optim.AdamW( + (parameter for parameter in model.parameters() if parameter.requires_grad), + lr=1e-2, + ) + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + loss = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + ).loss + assert loss.is_cuda + assert torch.isfinite(loss) + loss.backward() + + trainable_gradients = { + name: parameter.grad + for name, parameter in model.named_parameters() + if parameter.requires_grad + } + assert trainable_gradients + assert all(gradient is not None for gradient in trainable_gradients.values()) + assert all( + torch.isfinite(gradient).all() + for gradient in trainable_gradients.values() + if gradient is not None + ) + optimizer.step() + + changed = { + name + for name, parameter in model.named_parameters() + if not torch.equal(parameter.detach(), before[name]) + } + trainable = { + name + for name, parameter in model.named_parameters() + if parameter.requires_grad + } + assert changed + assert changed <= trainable + assert any("lora_" in name for name in changed) + assert all( + torch.equal(parameter.detach(), before[name]) + for name, parameter in model.named_parameters() + if not parameter.requires_grad + ) + + model.eval() + with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + expected = model( + input_ids=input_ids, + attention_mask=attention_mask, + ).logits + model.save_pretrained(tmp_path) + + restored_base = make_base() + restored_base.load_state_dict(base_state) + restored = peft.PeftModel.from_pretrained(restored_base, tmp_path).to("cuda").eval() + restored.base_model.model.set_attn_implementation("flash_attention_2") + with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + observed = restored( + input_ids=input_ids, + attention_mask=attention_mask, + ).logits + torch.testing.assert_close(observed, expected, rtol=0.0, atol=0.0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("backend", _FLASH_BACKENDS) +def test_precompiled_flash_accepts_fp32_storage_only_under_cuda_bf16_autocast( + backend: str, +) -> None: + assert torch.cuda.is_available(), "FlashAttention integration requires CUDA." + # X: (2, 17, 4, 16) + X = torch.randn((2, 17, 4, 16), device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match=r"bfloat16.*received float32"): + _core.kernels_flash_attention_func(X, X, X, implementation=backend) + + with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + Y = _core.kernels_flash_attention_func(X, X, X, implementation=backend) + + assert Y.dtype == torch.bfloat16 + assert torch.isfinite(Y).all() + + +@pytest.mark.gpu +def test_advertised_small_models_run_finite_flash_forwards() -> None: + assert torch.cuda.is_available(), "FlashAttention integration requires CUDA." + results = _model_results() + assert set(results) == {"esm2", "esm_plusplus"} + for family_id, family_results in results.items(): + expected_backends = _FLASH_BACKENDS if family_id == "esm2" else _ESMC_FLASH_BACKENDS + assert set(family_results) == set(expected_backends) + for backend in expected_backends: + metrics = family_results[backend] + assert metrics["finite"] is True + assert metrics["vs_sdpa"]["relative_l2"] <= 1e-2 + assert metrics["vs_sdpa"]["minimum_cosine"] >= 0.999 + + +@pytest.mark.gpu +@pytest.mark.parametrize("mixed_padding", (False, True), ids=("dense", "mixed-padding")) +def test_dplm_tiny_flash_attention_3_forward_parity_and_backward( + mixed_padding: bool, +) -> None: + """DPLM's advertised FA3 path executes under its official BF16 policy.""" + + assert torch.cuda.is_available(), "DPLM FlashAttention integration requires CUDA." + config = DPLMConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=128, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=64, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + token_dropout=False, + attn_backend="sdpa", + ) + torch.manual_seed(37) + model = DPLMModel(config).to(device="cuda", dtype=torch.float32) + # input_ids: (2, 17) + input_ids = torch.tensor( + ( + (0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2, 1, 1, 1, 1), + (0, 20, 19, 18, 17, 16, 15, 14, 2, 1, 1, 1, 1, 1, 1, 1, 1), + ), + device="cuda", + ) + attention_mask = input_ids.ne(1) if mixed_padding else torch.ones_like(input_ids).bool() + + model.eval() + outputs: dict[str, torch.Tensor] = {} + with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for backend in ("sdpa", "flash_attention_3"): + model.set_attn_implementation(backend) + outputs[backend] = model( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + assert torch.isfinite(outputs["flash_attention_3"]).all() + _assert_close( + outputs["flash_attention_3"][attention_mask], + outputs["sdpa"][attention_mask], + ) + + # Use a fresh training instance so rotary factors created under + # `inference_mode` cannot leak into the autograd contract. + training_model = DPLMModel(config).to(device="cuda", dtype=torch.float32).train() + training_model.set_attn_implementation("flash_attention_3") + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = training_model( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + # loss: () + loss = output[attention_mask].float().square().mean() + loss.backward() + gradients = [ + parameter.grad + for parameter in training_model.parameters() + if parameter.grad is not None + ] + assert gradients + assert all(torch.isfinite(gradient).all() for gradient in gradients) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.parametrize( + ("model_id", "model_class", "backends"), + ( + ("esm2_8m", FastEsmForMaskedLM, _FLASH_BACKENDS), + ("esmc_small", ESMplusplusForMaskedLM, _ESMC_FLASH_BACKENDS), + ("dplm_150m", DPLMModel, ("flash_attention_3",)), + ), +) +def test_manifest_checkpoint_flash_mixed_padding_parity( + model_id: str, + model_class: type[torch.nn.Module], + backends: tuple[str, ...], + record_property: Callable[[str, object], None], +) -> None: + assert torch.cuda.is_available(), "Checkpoint FlashAttention parity requires CUDA." + spec = get_model_registry()[model_id] + result = _run_checkpoint(model_id, model_class) + assert result["checkpoint"] == spec.fast.repo_id + assert result["revision"] == spec.fast.revision + assert set(result["backends"]) == set(backends) + for backend in backends: + metrics = result["backends"][backend] + record_property(f"{backend}_relative_l2", metrics["relative_l2"]) + record_property(f"{backend}_minimum_cosine", metrics["minimum_cosine"]) + assert metrics["finite"] is True + assert metrics["relative_l2"] <= 1e-2 + assert metrics["minimum_cosine"] >= 0.999 diff --git a/tests/integration/test_label_device_contracts.py b/tests/integration/test_label_device_contracts.py new file mode 100644 index 0000000..42fd9ba --- /dev/null +++ b/tests/integration/test_label_device_contracts.py @@ -0,0 +1,110 @@ +"""Accelerator contracts for direct loss calls with host-resident labels.""" + +from __future__ import annotations + +import pytest +import torch + +from fastplms.models.e1.modeling_e1 import E1Config, E1ForTokenClassification +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ESMplusplusForTokenClassification, +) + + +def _esmc_config() -> ESMplusplusConfig: + return ESMplusplusConfig( + vocab_size=40, + hidden_size=8, + num_attention_heads=2, + num_hidden_layers=1, + num_labels=3, + dropout=0.0, + attn_backend="sdpa", + ) + + +def _e1_config() -> E1Config: + return E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + max_num_sequences=4, + max_num_positions_within_seq=32, + max_num_positions_global=64, + num_labels=3, + attn_backend="sdpa", + dtype="float32", + use_cache=False, + ) + + +def _esm3_config() -> FastESM3Config: + return FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="sdpa", + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("family", ("esmc_mlm", "esmc_token", "e1_token", "esm3")) +def test_advertised_heads_accept_host_labels_with_cuda_logits(family: str) -> None: + assert torch.cuda.is_available(), "label device contracts require CUDA" + device = torch.device("cuda") + + if family == "esmc_mlm": + model = ESMplusplusForMaskedLM(_esmc_config()).to(device).train() + # input_ids: (2, 4) + input_ids = torch.tensor(((0, 4, 5, 2), (0, 6, 2, 1)), device=device) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + # labels: (b, l) + labels = input_ids.cpu().masked_fill(~attention_mask.cpu(), -100) + output = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + elif family == "esmc_token": + model = ESMplusplusForTokenClassification(_esmc_config()).to(device).train() + # input_ids: (2, 4) + input_ids = torch.tensor(((0, 4, 5, 2), (0, 6, 2, 1)), device=device) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + # labels: (b, l) + labels = input_ids.cpu().remainder(3).masked_fill(~attention_mask.cpu(), -100) + output = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + elif family == "e1_token": + model = E1ForTokenClassification(_e1_config()).to(device).train() + # input_ids: (1, 4) + input_ids = torch.tensor(((1, 5, 6, 2),), device=device) + # positions: (..., 3) + positions = torch.arange(input_ids.shape[1], device=device).unsqueeze(0) + labels = input_ids.cpu().remainder(3) + output = model( + input_ids=input_ids, + within_seq_position_ids=positions, + global_position_ids=positions, + sequence_ids=torch.zeros_like(input_ids), + labels=labels, + ) + else: + model = FastESM3Model(_esm3_config()).to(device).train() + # input_ids: (2, 4) + input_ids = torch.tensor(((0, 4, 5, 2), (0, 6, 2, 1)), device=device) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + # labels: (b, l) + labels = input_ids.cpu().masked_fill(~attention_mask.cpu(), -100) + output = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + + assert output.loss is not None + assert output.loss.device.type == "cuda" + assert torch.isfinite(output.loss) + output.loss.backward() + gradients = [parameter.grad for parameter in model.parameters() if parameter.grad is not None] + assert gradients + assert all(bool(torch.isfinite(gradient).all()) for gradient in gradients) diff --git a/tests/integration/test_meta_loader_rotary.py b/tests/integration/test_meta_loader_rotary.py new file mode 100644 index 0000000..d46a3c9 --- /dev/null +++ b/tests/integration/test_meta_loader_rotary.py @@ -0,0 +1,83 @@ +"""Fresh Transformers v5 loader regressions for non-persistent rotary buffers.""" + +from __future__ import annotations + +import torch +from transformers import PreTrainedModel + +from fastplms.models.e1.modeling_e1 import E1Config, E1ForMaskedLM +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model + + +def test_e1_fresh_pretrained_load_rebuilds_rotary_cache(tmp_path) -> None: + """A fresh local v5 load produces finite E1 outputs on its first forward.""" + + config = E1Config( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_num_sequences=4, + max_num_positions_within_seq=16, + max_num_positions_global=16, + attn_backend="sdpa", + ) + source = E1ForMaskedLM(config).eval() + source.save_pretrained(tmp_path, safe_serialization=True) + reloaded = E1ForMaskedLM.from_pretrained( + tmp_path, + local_files_only=True, + attn_implementation="sdpa", + ).eval() + rotary = reloaded.model.layers[0].norm_attn_norm.self_attn.rotary_emb + assert rotary.inv_freq.numel() == 0 + + # input_ids: (1, 5) + input_ids = torch.tensor([[1, 4, 5, 6, 2]]) + # position_ids: (b, l) + position_ids = torch.arange(5).unsqueeze(0) + output = reloaded( + input_ids=input_ids, + within_seq_position_ids=position_ids, + global_position_ids=position_ids, + sequence_ids=torch.zeros_like(input_ids), + ) + + assert torch.isfinite(output.last_hidden_state).all() + assert torch.isfinite(output.logits).all() + assert rotary.inv_freq.numel() == 4 + + +def test_esm3_fresh_pretrained_load_rebuilds_rotary_frequency(tmp_path) -> None: + """A fresh local v5 load cannot reuse a materialized garbage frequency.""" + + config = FastESM3Config( + hidden_size=16, + num_attention_heads=2, + num_vector_heads=4, + num_hidden_layers=1, + attn_backend="sdpa", + ) + source = FastESM3Model(config).eval() + PreTrainedModel.save_pretrained(source, tmp_path, safe_serialization=True) + reloaded = FastESM3Model.from_pretrained( + tmp_path, + local_files_only=True, + attn_implementation="sdpa", + ).eval() + rotary = reloaded.esm3.transformer.blocks[0].attn.rotary + assert rotary._cos_cached is None + assert rotary._sin_cached is None + + # input_ids: (1, 5) + input_ids = torch.tensor([[0, 5, 6, 7, 2]]) + output = reloaded( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + ) + + expected_inv_freq = rotary._compute_inv_freq(torch.device("cpu")) + assert torch.equal(rotary.inv_freq, expected_inv_freq) + assert torch.isfinite(output.last_hidden_state).all() + assert torch.isfinite(output.logits).all() diff --git a/tests/integration/test_official_goldens.py b/tests/integration/test_official_goldens.py new file mode 100644 index 0000000..1af3752 --- /dev/null +++ b/tests/integration/test_official_goldens.py @@ -0,0 +1,134 @@ +"""Fast candidate regression against manifest-declared official goldens.""" + +from __future__ import annotations + +import contextlib +import gc +import importlib +import json +import pytest +import torch +from pathlib import Path +from safetensors.torch import load_file + +from fastplms.registry import ModelSpec, get_model_registry +from tests.parity.test_model_parity import ( + _assert_logits_contract, + _assert_tensor_contract, + _last_hidden, + _numeric_contract, +) +from tests.structure.support.hardware import assert_recorded_hopper_device_matches +from tools.goldens import validate_golden_bundle + + +ROOT = Path(__file__).resolve().parents[2] +REGISTRY = get_model_registry() +SEQUENCE_GOLDENS = tuple( + spec + for spec in REGISTRY.values() + if spec.official_golden is not None and spec.family.tokenizer_mode != "structure" +) + + +def _parameter(spec: ModelSpec) -> object: + marks = [pytest.mark.large] if spec.size_category == "xlarge" else [] + return pytest.param(spec, id=spec.id, marks=marks) + + +def _model_class(spec: ModelSpec) -> type[torch.nn.Module]: + """Resolve the current package implementation declared by the manifest.""" + + if spec.family.id == "ankh": + auto_class = "AutoModel" + elif "AutoModelForMaskedLM" in spec.auto_map: + auto_class = "AutoModelForMaskedLM" + else: + auto_class = "AutoModel" + qualified_name = spec.auto_map[auto_class] + module_name, class_name = qualified_name.rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + assert issubclass(model_class, torch.nn.Module) + return model_class + + +@pytest.mark.gpu +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_GOLDENS]) +def test_declared_sequence_golden_matches_candidate(spec: ModelSpec) -> None: + """Run one compact BF16 regression without importing an official package.""" + + declaration = spec.official_golden + assert declaration is not None + metadata_path = ROOT / declaration.metadata.path + tensors_path = ROOT / declaration.tensors.path + validate_golden_bundle( + spec, + REGISTRY, + metadata_path=metadata_path, + tensors_path=tensors_path, + declaration=declaration, + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + recorded_environment = metadata["environment"]["details"] + assert isinstance(recorded_environment, dict) + properties = torch.cuda.get_device_properties(0) + assert_recorded_hopper_device_matches( + { + "cuda_device": properties.name, + "cuda_device_capability": list(torch.cuda.get_device_capability(0)), + "cuda_total_memory": int(properties.total_memory), + }, + recorded_environment, + ) + tensors = load_file(tensors_path, device="cpu") + device = torch.device("cuda") + use_bf16_autocast = spec.family.bf16_execution == "fp32_parameters_autocast" + load_dtype = torch.float32 if use_bf16_autocast else torch.bfloat16 + model = ( + _model_class(spec) + .from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + dtype=load_dtype, + device_map=device, + ) + .eval() + ) + inputs = { + name.removeprefix("input__"): T.to(device) + for name, T in tensors.items() + if name.startswith("input__") + } + # residue_mask: (b, l) + residue_mask = tensors["residue_mask"].to(device).bool() + numeric_context = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if use_bf16_autocast + else contextlib.nullcontext() + ) + with torch.inference_mode(), numeric_context: + output = model(**inputs, output_hidden_states=True) + contract = _numeric_contract(spec, torch.bfloat16, None) + _assert_tensor_contract( + _last_hidden(output), + tensors["output__last_hidden_state"].to(device), + residue_mask, + contract, + f"{spec.id}:bf16:golden:last_hidden_state", + ) + official_logits = tensors.get("output__logits") + candidate_logits = getattr(output, "logits", None) + assert (candidate_logits is None) == (official_logits is None), ( + f"{spec.id}: golden and candidate output-head contracts differ" + ) + if official_logits is not None: + _assert_logits_contract( + candidate_logits, + official_logits.to(device), + residue_mask, + contract, + f"{spec.id}:bf16:golden:logits", + ) + del model, output, tensors + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/integration/test_sequence_output_contracts.py b/tests/integration/test_sequence_output_contracts.py new file mode 100644 index 0000000..29a64ac --- /dev/null +++ b/tests/integration/test_sequence_output_contracts.py @@ -0,0 +1,213 @@ +"""Hugging Face output contracts for the ESM2 and ESMC families.""" + +from __future__ import annotations + +import pytest +import torch +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + +from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmForSequenceClassification, + FastEsmForTokenClassification, + FastEsmModel, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ESMplusplusForSequenceClassification, + ESMplusplusForTokenClassification, + ESMplusplusModel, +) + + +def _esm2_config() -> FastEsmConfig: + return FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=5, + num_labels=3, + position_embedding_type="absolute", + attn_backend="eager", + return_dict=False, + output_hidden_states=True, + ) + + +def _esmc_config() -> ESMplusplusConfig: + return ESMplusplusConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + dropout=0.0, + pad_token_id=1, + mask_token_id=5, + num_labels=3, + attn_backend="eager", + return_dict=False, + output_hidden_states=True, + ) + + +def _assert_nested_close(actual, expected) -> None: + if torch.is_tensor(expected): + assert torch.is_tensor(actual) + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, (tuple, list)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_value, expected_value in zip(actual, expected, strict=True): + _assert_nested_close(actual_value, expected_value) + return + assert actual == expected + + +@pytest.mark.parametrize( + ("model_class", "config_factory", "kind", "output_class"), + ( + (FastEsmModel, _esm2_config, "base", ModelOutput), + (FastEsmForMaskedLM, _esm2_config, "mlm", MaskedLMOutput), + ( + FastEsmForSequenceClassification, + _esm2_config, + "sequence", + SequenceClassifierOutput, + ), + ( + FastEsmForTokenClassification, + _esm2_config, + "token", + TokenClassifierOutput, + ), + (ESMplusplusModel, _esmc_config, "base", ModelOutput), + (ESMplusplusForMaskedLM, _esmc_config, "mlm", MaskedLMOutput), + ( + ESMplusplusForSequenceClassification, + _esmc_config, + "sequence", + SequenceClassifierOutput, + ), + ( + ESMplusplusForTokenClassification, + _esmc_config, + "token", + TokenClassifierOutput, + ), + ), +) +def test_sequence_models_honor_config_and_explicit_output_controls( + model_class: type[torch.nn.Module], + config_factory, + kind: str, + output_class: type[ModelOutput], +) -> None: + model = model_class(config_factory()).eval() + # input_ids: (2, 5) + input_ids = torch.tensor([[0, 3, 4, 2, 1], [0, 6, 2, 1, 1]]) + # attention_mask: (b, l) + attention_mask = input_ids.ne(1) + labels = { + "base": None, + "mlm": input_ids.masked_fill(~attention_mask, -100), + "sequence": torch.tensor([1, 2]), + "token": input_ids.remainder(3).masked_fill(~attention_mask, -100), + }[kind] + full_kwargs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "output_attentions": True, + "output_hidden_states": True, + "output_s_max": True, + } + if labels is not None: + full_kwargs["labels"] = labels + + with torch.inference_mode(): + default_output = model(input_ids=input_ids, attention_mask=attention_mask) + default_structured = model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + ) + tuple_output = model(**full_kwargs, return_dict=False) + structured = model(**full_kwargs, return_dict=True) + no_states = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=False, + return_dict=True, + ) + + assert isinstance(default_output, tuple) + _assert_nested_close(default_output, default_structured.to_tuple()) + assert isinstance(structured, output_class) + _assert_nested_close(tuple_output, structured.to_tuple()) + assert structured.hidden_states is not None + assert structured.attentions is not None + assert structured.s_max is not None + assert no_states.hidden_states is None + if kind != "base": + assert tuple(structured.keys())[:4] == ( + "loss", + "logits", + "hidden_states", + "attentions", + ) + assert tuple(structured.keys())[4] == "s_max" + + +@pytest.mark.parametrize( + ("model_class", "config_factory"), + ( + (FastEsmModel, _esm2_config), + (FastEsmForMaskedLM, _esm2_config), + (FastEsmForSequenceClassification, _esm2_config), + (FastEsmForTokenClassification, _esm2_config), + (ESMplusplusModel, _esmc_config), + (ESMplusplusForMaskedLM, _esmc_config), + (ESMplusplusForSequenceClassification, _esmc_config), + (ESMplusplusForTokenClassification, _esmc_config), + ), +) +def test_sequence_models_reject_unexpected_forward_arguments( + model_class: type[torch.nn.Module], + config_factory, +) -> None: + model = model_class(config_factory()).eval() + + with pytest.raises(TypeError, match="unexpected_contract"): + model( + input_ids=torch.tensor([[0, 3, 2]]), + unexpected_contract=True, + ) + + +def test_esm2_resize_preserves_existing_logits_and_bias_contract() -> None: + model = FastEsmForMaskedLM(_esm2_config()).eval() + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 3, 4, 2]]) + with torch.inference_mode(): + original_logits = model(input_ids=input_ids, return_dict=True).logits + + model.resize_token_embeddings(19) + + with torch.inference_mode(): + resized_logits = model(input_ids=input_ids, return_dict=True).logits + assert model.get_output_embeddings().bias is None + assert model.lm_head.bias.shape == (19,) + torch.testing.assert_close(resized_logits[..., :16], original_logits) diff --git a/tests/integration/test_source_archive_provenance.py b/tests/integration/test_source_archive_provenance.py new file mode 100644 index 0000000..5c54ffb --- /dev/null +++ b/tests/integration/test_source_archive_provenance.py @@ -0,0 +1,364 @@ +"""End-to-end contracts for Git-free remote source archives.""" + +from __future__ import annotations + +import json +import subprocess +import tarfile +import pytest +from collections.abc import Iterable +from pathlib import Path, PurePosixPath +from types import SimpleNamespace + +import tools.artifacts.build as artifact_build +from tools.artifacts.build import ( + ArtifactError, + _validate_vendor_revisions, + _validated_runtime_snapshot, +) +from tools.remote.run import create_source_archive +from tools.source_provenance import ( + ARCHIVE_PROVENANCE_NAME, + SourceProvenanceError, + validate_archived_root, +) + + +def _git(repository: Path, *arguments: str) -> str: + completed = subprocess.run( + [ + "git", + "-c", + "user.name=FastPLMs Tests", + "-c", + "user.email=fastplms-tests@example.invalid", + "-c", + "protocol.file.allow=always", + *arguments, + ], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + return completed.stdout.strip() + + +def _revision_contract(revision: str) -> tuple[SimpleNamespace, SimpleNamespace]: + source = SimpleNamespace(path="vendor/upstream/toy", revision=revision) + registry = SimpleNamespace(upstreams={"toy": source}) + spec = SimpleNamespace(family=SimpleNamespace(upstreams=("toy",))) + return registry, spec + + +def _create_repository(tmp_path: Path) -> tuple[Path, str]: + upstream = tmp_path / "upstream" + upstream.mkdir() + _git(upstream, "init", "--initial-branch=main") + (upstream / "LICENSE").write_text("Synthetic license\n", encoding="utf-8") + (upstream / "weights.py").write_text("scale = 2\n", encoding="utf-8") + _git(upstream, "add", "LICENSE", "weights.py") + _git(upstream, "commit", "-m", "Create pinned upstream") + revision = _git(upstream, "rev-parse", "HEAD") + + repository = tmp_path / "repository" + repository.mkdir() + _git(repository, "init", "--initial-branch=main") + (repository / "README.md").write_text("Synthetic FastPLMs tree\n", encoding="utf-8") + runtime = repository / "src" / "fastplms" / "toy_runtime" + runtime.mkdir(parents=True) + (runtime / "core.py").write_text("VALUE = 1\n", encoding="utf-8") + (repository / ".secrets.env").write_text("TOKEN=not-archived\n", encoding="utf-8") + _git(repository, "add", "README.md", "src/fastplms/toy_runtime/core.py") + _git( + repository, + "submodule", + "add", + str(upstream), + "vendor/upstream/toy", + ) + _git(repository, "commit", "-m", "Pin synthetic upstream") + return repository, revision + + +def test_source_archive_preserves_revision_proof_without_git_metadata( + tmp_path: Path, +) -> None: + repository, revision = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + member_names = [member.name for member in archive.getmembers()] + archive.extractall(extracted, filter="data") + + assert ARCHIVE_PROVENANCE_NAME in member_names + assert ".secrets.env" not in member_names + assert not any(".git" in Path(name).parts for name in member_names) + assert not any(path.name == ".git" for path in extracted.rglob(".git")) + + provenance = json.loads((extracted / ARCHIVE_PROVENANCE_NAME).read_text(encoding="utf-8")) + root_record = provenance["root"] + assert root_record["head_revision"] == _git(repository, "rev-parse", "HEAD") + assert ".secrets.env" not in root_record["files"] + runtime_record = root_record["files"]["src/fastplms/toy_runtime/core.py"] + assert runtime_record["mode"] == "100644" + assert runtime_record["size"] == len("VALUE = 1\n") + assert len(runtime_record["sha256"]) == 64 + validate_archived_root(extracted) + archived_runtime = extracted / "src" / "fastplms" / "toy_runtime" / "core.py" + archived_runtime.write_text("VALUE = 2\n", encoding="utf-8") + with pytest.raises(SourceProvenanceError, match="tracked bytes or modes differ"): + validate_archived_root(extracted) + archived_runtime.write_text("VALUE = 1\n", encoding="utf-8") + record = provenance["submodules"]["vendor/upstream/toy"] + assert record["gitlink_revision"] == revision + assert record["head_revision"] == revision + assert record["tracked_files"] == ["LICENSE", "weights.py"] + + registry, spec = _revision_contract(revision) + _validate_vendor_revisions(extracted, registry, spec) + + (extracted / "vendor" / "upstream" / "toy" / "weights.py").write_text( + "scale = 3\n", + encoding="utf-8", + ) + with pytest.raises(ArtifactError, match="tracked-tree digest differs"): + _validate_vendor_revisions(extracted, registry, spec) + + marker = extracted / ARCHIVE_PROVENANCE_NAME + marker.unlink() + marker.mkdir() + with pytest.raises(ArtifactError, match="not a regular file"): + _validate_vendor_revisions(extracted, registry, spec) + + +def _runtime_contract() -> tuple[SimpleNamespace, SimpleNamespace]: + registry = SimpleNamespace() + family = SimpleNamespace(runtime_paths=("toy_runtime",), attention=()) + return registry, SimpleNamespace(family=family) + + +def test_git_free_runtime_snapshot_is_content_addressed_and_rejects_unknown_files( + tmp_path: Path, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + registry, spec = _runtime_contract() + + runtime_revision, payloads, source_tree_sha256 = _validated_runtime_snapshot( + extracted, + registry, + spec, + ) + + assert runtime_revision == f"source-tree-sha256:{source_tree_sha256}" + assert payloads == {"toy_runtime/core.py": b"VALUE = 1\n"} + assert len(source_tree_sha256) == 64 + + extra = extracted / "src" / "fastplms" / "toy_runtime" / "extra.py" + extra.write_text("EXTRA = True\n", encoding="utf-8") + with pytest.raises(ArtifactError, match="inventory differs"): + _validated_runtime_snapshot(extracted, registry, spec) + + +def test_git_free_runtime_revision_does_not_trust_diagnostic_head( + tmp_path: Path, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + marker = extracted / ARCHIVE_PROVENANCE_NAME + provenance = json.loads(marker.read_text(encoding="utf-8")) + provenance["root"]["head_revision"] = "f" * 40 + marker.write_text(json.dumps(provenance), encoding="utf-8") + + diagnostic_head, _inventory = validate_archived_root(extracted) + registry, spec = _runtime_contract() + runtime_revision, _payloads, source_tree_sha256 = _validated_runtime_snapshot( + extracted, + registry, + spec, + ) + + assert diagnostic_head == "f" * 40 + assert runtime_revision == f"source-tree-sha256:{source_tree_sha256}" + assert runtime_revision != diagnostic_head + + +def test_git_free_runtime_snapshot_rejects_mutation_between_validation_and_snapshot( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + registry, spec = _runtime_contract() + original_snapshot = artifact_build._snapshot_runtime_sources + + def tampered_snapshot( + source_root: Path, + entries: Iterable[tuple[Path, PurePosixPath]], + revision: str | None, + ) -> dict[str, bytes]: + payloads = original_snapshot(source_root, entries, revision) + payloads["toy_runtime/core.py"] = b"VALUE = 2\n" + return payloads + + monkeypatch.setattr(artifact_build, "_snapshot_runtime_sources", tampered_snapshot) + + with pytest.raises(ArtifactError, match="mutated during snapshot"): + _validated_runtime_snapshot(extracted, registry, spec) + + +@pytest.mark.parametrize( + "mutation", + ("schema", "revision", "digest", "size", "path"), +) +def test_archived_root_rejects_malformed_or_forged_metadata( + tmp_path: Path, + mutation: str, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + marker = extracted / ARCHIVE_PROVENANCE_NAME + provenance = json.loads(marker.read_text(encoding="utf-8")) + root_record = provenance["root"] + runtime_name = "src/fastplms/toy_runtime/core.py" + if mutation == "schema": + provenance["unknown"] = True + elif mutation == "revision": + root_record["head_revision"] = "not-a-commit" + elif mutation == "digest": + root_record["files"][runtime_name]["sha256"] = "0" * 64 + elif mutation == "size": + root_record["files"][runtime_name]["size"] += 1 + else: + root_record["files"]["../escape.py"] = root_record["files"].pop(runtime_name) + marker.write_text(json.dumps(provenance), encoding="utf-8") + + with pytest.raises(SourceProvenanceError): + validate_archived_root(extracted) + + +@pytest.mark.parametrize( + "unsafe_name", + ("C:\\escape.py", "src\\fastplms\\escape.py", "payload:stream.py", "."), +) +def test_archived_root_rejects_nonportable_paths( + tmp_path: Path, + unsafe_name: str, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + marker = extracted / ARCHIVE_PROVENANCE_NAME + provenance = json.loads(marker.read_text(encoding="utf-8")) + files = provenance["root"]["files"] + files[unsafe_name] = files.pop("src/fastplms/toy_runtime/core.py") + marker.write_text(json.dumps(provenance), encoding="utf-8") + + with pytest.raises(SourceProvenanceError, match="Non-portable tracked path"): + validate_archived_root(extracted) + + +def test_archived_root_rejects_parent_symlink_traversal(tmp_path: Path) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + runtime = extracted / "src" / "fastplms" / "toy_runtime" + preserved = extracted / "src" / "fastplms" / "toy_runtime-preserved" + runtime.rename(preserved) + runtime.symlink_to(preserved, target_is_directory=True) + + with pytest.raises(SourceProvenanceError, match="traverses a symlink"): + validate_archived_root(extracted) + + +def test_git_free_runtime_snapshot_rejects_symlinks_and_unknown_extensions( + tmp_path: Path, +) -> None: + repository, _ = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + extracted = tmp_path / "extracted" + extracted.mkdir() + with tarfile.open(archive_path, "r:gz") as archive: + archive.extractall(extracted, filter="data") + registry, spec = _runtime_contract() + runtime = extracted / "src" / "fastplms" / "toy_runtime" + core = runtime / "core.py" + core.unlink() + core.symlink_to(extracted / "README.md") + with pytest.raises(ArtifactError, match="Symlinks are not allowed"): + _validated_runtime_snapshot(extracted, registry, spec) + + core.unlink() + core.write_text("VALUE = 1\n", encoding="utf-8") + (runtime / "payload.bin").write_bytes(b"unknown") + with pytest.raises(ArtifactError, match="unapproved extension"): + _validated_runtime_snapshot(extracted, registry, spec) + + +def test_archive_provenance_is_not_used_when_git_metadata_exists(tmp_path: Path) -> None: + repository, revision = _create_repository(tmp_path) + archive_path = tmp_path / "source.tar.gz" + create_source_archive(repository, archive_path) + with tarfile.open(archive_path, "r:gz") as archive: + marker = archive.extractfile(ARCHIVE_PROVENANCE_NAME) + assert marker is not None + (repository / ARCHIVE_PROVENANCE_NAME).write_bytes(marker.read()) + + (repository / "vendor" / "upstream" / "toy" / ".git").unlink() + wrong_revision = "0" * 40 if revision != "0" * 40 else "1" * 40 + registry, spec = _revision_contract(wrong_revision) + with pytest.raises(ArtifactError, match="not initialized"): + _validate_vendor_revisions(repository, registry, spec) + + +def test_source_archive_rejects_modified_tracked_submodule_files(tmp_path: Path) -> None: + repository, _ = _create_repository(tmp_path) + (repository / "vendor" / "upstream" / "toy" / "weights.py").write_text( + "scale = 3\n", + encoding="utf-8", + ) + + with pytest.raises(RuntimeError, match="modified tracked files"): + create_source_archive(repository, tmp_path / "source.tar.gz") + + +def test_source_archive_rejects_a_tracked_sensitive_path(tmp_path: Path) -> None: + repository, _ = _create_repository(tmp_path) + _git(repository, "add", ".secrets.env") + _git(repository, "commit", "-m", "Track a forbidden credential-shaped file") + + with pytest.raises(RuntimeError, match="tracks forbidden source path"): + create_source_archive(repository, tmp_path / "source.tar.gz") diff --git a/tests/integration/test_ttt.py b/tests/integration/test_ttt.py new file mode 100644 index 0000000..db2ac05 --- /dev/null +++ b/tests/integration/test_ttt.py @@ -0,0 +1,651 @@ +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar +from transformers import PretrainedConfig, PreTrainedModel + +from fastplms.models.ankh.modeling_ankh import FastAnkhForMaskedLMExtension +from fastplms.models.dplm.modeling_dplm import DPLMForMaskedLM +from fastplms.models.dplm2.modeling_dplm2 import DPLM2ForMaskedLM +from fastplms.models.e1.modeling_e1 import E1ForMaskedLM +from fastplms.models.esm2.modeling_fastesm import FastEsmForMaskedLM +from fastplms.models.esm3.modeling_esm3 import FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM +from fastplms.models.esmfold.modeling_fast_esmfold import FastEsmForProteinFolding +from fastplms.models.ttt import ( + FastPLMTestTimeTrainingMixin, + LoraInjectedLinear, + TTTConfig, +) +from tests.conftest import MODEL_REGISTRY, STRUCTURE_MODEL_REGISTRY + + +TEST_SEQUENCE = "MSTNPKPQRKTKRNT" +LOCAL_MODEL_CLASSES = { + "esm2": FastEsmForMaskedLM, + "esmc": ESMplusplusForMaskedLM, + "esm3": FastESM3Model, + "e1": E1ForMaskedLM, + "dplm": DPLMForMaskedLM, + "dplm2": DPLM2ForMaskedLM, + "ankh": FastAnkhForMaskedLMExtension, +} + + +@pytest.mark.parametrize("method", ["ttt", "ttt_reset", "fold_protein_ttt"]) +def test_esmfold_ttt_entry_points_reject_without_an_untrained_head(method: str) -> None: + model = object.__new__(FastEsmForProteinFolding) + kwargs = {"seq": "ACDE"} if method == "ttt" else {} + if method == "fold_protein_ttt": + kwargs = {"sequence": "ACDE"} + + with pytest.raises(RuntimeError, match="does not contain a trained masked-language-model head"): + getattr(model, method)(**kwargs) + + +def test_esmfold_fold_protein_ttt_flag_rejects_before_folding() -> None: + model = object.__new__(FastEsmForProteinFolding) + + with pytest.raises(RuntimeError, match="does not contain a trained masked-language-model head"): + model.fold_protein("ACDE", ttt=True) + + +class DummyConfig: + vocab_size = 8 + + +class DummyTokenizer: + pad_token_id = 0 + cls_token_id = 1 + eos_token_id = 2 + mask_token_id = 3 + all_special_ids: ClassVar[list[int]] = [0, 1, 2, 3] + + def __init__(self) -> None: + self.vocab = { + "A": 4, + "C": 5, + "D": 6, + "E": 7, + } + + def __call__( + self, + seq: str | list[str], + return_tensors: str = "pt", + padding: bool = True, + ) -> dict[str, torch.Tensor]: + del return_tensors, padding + sequences = [seq] if isinstance(seq, str) else seq + encoded = [] + for sequence in sequences: + encoded.append( + [self.cls_token_id] + [self.vocab[aa] for aa in sequence] + [self.eos_token_id] + ) + max_len = max(len(ids) for ids in encoded) + # input_ids: (len(encoded), max_len) + input_ids = torch.full((len(encoded), max_len), self.pad_token_id) + for row, ids in enumerate(encoded): + # input_ids[row, :len(ids)]: (...) + input_ids[row, : len(ids)] = torch.tensor(ids) + return {"input_ids": input_ids.long()} + + +class DummyTTTModel(FastPLMTestTimeTrainingMixin, nn.Module): + def __init__(self) -> None: + nn.Module.__init__(self) + self.config = DummyConfig() + self.tokenizer = DummyTokenizer() + self.embed = nn.Embedding(self.config.vocab_size, 8) + self.backbone = nn.Sequential( + nn.Linear(8, 8), + nn.GELU(), + nn.Dropout(p=0.5), + nn.Linear(8, 8), + ) + self.lm_head = nn.Linear(8, self.config.vocab_size) + self.init_ttt( + { + "steps": 1, + "ags": 1, + "batch_size": 1, + "mask_ratio": 1.0, + "bert_leave_prob": 0.0, + "bert_replace_prob": 0.0, + "lora_rank": 2, + "lora_alpha": 1.0, + } + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.backbone] + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ): + # input_ids: (b, l) + del attention_mask + hidden = self.backbone(self.embed(input_ids)) + return SimpleNamespace(logits=self.lm_head(hidden)) + + +class FamilyAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.query = nn.Linear(8, 8) + self.value = nn.Linear(8, 8) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # hidden_states: (..., d) + return self.query(hidden_states) + self.value(hidden_states) + + +class DummyFamilyTargetTTTModel(FastPLMTestTimeTrainingMixin, nn.Module): + def __init__(self) -> None: + nn.Module.__init__(self) + self.config = DummyConfig() + self.tokenizer = DummyTokenizer() + self.embed = nn.Embedding(self.config.vocab_size, 8) + self.backbone = nn.ModuleDict( + { + "attention": FamilyAttention(), + "feed_forward": nn.Linear(8, 8), + } + ) + self.lm_head = nn.Linear(8, self.config.vocab_size) + self.init_ttt({"lora_target_replace_module": "FamilyAttention"}) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.backbone] + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ): + # input_ids: (b, l) + del attention_mask + hidden = self.embed(input_ids) + hidden = self.backbone["attention"](hidden) + hidden = self.backbone["feed_forward"](hidden) + return SimpleNamespace(logits=self.lm_head(hidden)) + + +class DummyPretrainedTTTConfig(PretrainedConfig): + model_type = "dummy_pretrained_ttt" + + def __init__(self, vocab_size: int = 8, **kwargs) -> None: + super().__init__(**kwargs) + self.vocab_size = vocab_size + + +class DummyPretrainedTTTModel(FastPLMTestTimeTrainingMixin, PreTrainedModel): + config_class = DummyPretrainedTTTConfig + + def __init__(self, config: DummyPretrainedTTTConfig) -> None: + PreTrainedModel.__init__(self, config) + self.tokenizer = DummyTokenizer() + self.embed = nn.Embedding(config.vocab_size, 8) + self.backbone = nn.Sequential(nn.Linear(8, 8), nn.GELU(), nn.Linear(8, 8)) + self.lm_head = nn.Linear(8, config.vocab_size) + self.post_init() + self.init_ttt( + { + "seed": 17, + "lora_rank": 2, + "lora_alpha": 1.0, + "lora_target_modules": ("0", "2"), + } + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.backbone] + + def forward(self, input_ids: torch.Tensor, **kwargs): + # input_ids: (b, l) + del kwargs + hidden = self.backbone(self.embed(input_ids)) + return SimpleNamespace(logits=self.lm_head(hidden)) + + +def test_ttt_masking_masks_only_residue_tokens() -> None: + model = DummyTTTModel() + tokenized = model._ttt_tokenize(seq="ACDE") + generator = torch.Generator() + generator.manual_seed(0) + + batch, labels = model._ttt_sample_batch(tokenized, generator) + + assert isinstance(batch, torch.Tensor) + assert labels[0, 0].item() == -100 + assert labels[0, -1].item() == -100 + assert torch.all(batch[labels != -100] == model.tokenizer.mask_token_id) + + +def test_ttt_lora_injection_is_lazy_and_backbone_scoped() -> None: + model = DummyTTTModel() + + assert all("lora_" not in name for name in model.state_dict()) + + model._ttt_ensure_initialized() + + assert any(isinstance(module, LoraInjectedLinear) for module in model.backbone.modules()) + assert not any(isinstance(module, LoraInjectedLinear) for module in model.lm_head.modules()) + + +def test_ttt_lora_alpha_is_the_proteinttt_direct_multiplier() -> None: + adapter = LoraInjectedLinear( + nn.Linear(3, 2, bias=False), + rank=2, + alpha=6.0, + ) + # inputs: (1, 3) + inputs = torch.tensor([[1.0, 2.0, 3.0]]) + with torch.no_grad(): + adapter.linear.weight.zero_() + adapter.lora_down.weight.copy_( + torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + ] + ) + ) + adapter.lora_up.weight.copy_( + torch.tensor( + [ + [1.0, 1.0], + [2.0, -1.0], + ] + ) + ) + + unscaled_delta = adapter.lora_up(adapter.lora_down(inputs)) + torch.testing.assert_close(adapter(inputs), unscaled_delta * 6.0) + assert not torch.equal(adapter(inputs), unscaled_delta * (6.0 / adapter.rank)) + + +def test_ttt_first_call_mapping_preserves_family_target_class() -> None: + model = DummyFamilyTargetTTTModel() + base_weights = { + "query": model.backbone["attention"].query.weight.detach().clone(), + "value": model.backbone["attention"].value.weight.detach().clone(), + "feed_forward": model.backbone["feed_forward"].weight.detach().clone(), + "lm_head": model.lm_head.weight.detach().clone(), + } + + metrics = model.ttt( + seq="ACDE", + ttt_config={ + "steps": 1, + "ags": 1, + "batch_size": 1, + "mask_ratio": 1.0, + "bert_leave_prob": 0.0, + "bert_replace_prob": 0.0, + "lora_rank": 2, + "lora_alpha": 1.0, + "seed": 23, + }, + ) + + assert model.ttt_config.lora_target_replace_module == "FamilyAttention" + adapter_names = [ + name for name, module in model.named_modules() if isinstance(module, LoraInjectedLinear) + ] + assert adapter_names == [ + "backbone.attention.query", + "backbone.attention.value", + ] + assert len(metrics["losses"]) == 1 + torch.testing.assert_close( + model.backbone["attention"].query.linear.weight, + base_weights["query"], + ) + torch.testing.assert_close( + model.backbone["attention"].value.linear.weight, + base_weights["value"], + ) + torch.testing.assert_close( + model.backbone["feed_forward"].weight, + base_weights["feed_forward"], + ) + torch.testing.assert_close(model.lm_head.weight, base_weights["lm_head"]) + assert any( + not torch.equal(module.lora_up.weight, module._ttt_initial_lora_up) + for module in model._ttt_lora_modules() + ) + + +def test_ttt_direct_init_mapping_preserves_family_target_class() -> None: + model = DummyFamilyTargetTTTModel() + + model.init_ttt({"steps": 2, "seed": 29}) + model._ttt_ensure_initialized() + + assert model.ttt_config.steps == 2 + assert model.ttt_config.seed == 29 + assert model.ttt_config.lora_target_replace_module == "FamilyAttention" + assert [ + name for name, module in model.named_modules() if isinstance(module, LoraInjectedLinear) + ] == [ + "backbone.attention.query", + "backbone.attention.value", + ] + + +def test_ttt_first_call_mapping_preserves_explicit_target_override() -> None: + model = DummyFamilyTargetTTTModel() + + model.ttt( + seq="ACDE", + ttt_config={ + "steps": 1, + "ags": 1, + "batch_size": 1, + "mask_ratio": 1.0, + "bert_leave_prob": 0.0, + "bert_replace_prob": 0.0, + "lora_rank": 2, + "lora_alpha": 1.0, + "lora_target_replace_module": None, + "lora_target_modules": ("feed_forward",), + }, + ) + + assert model.ttt_config.lora_target_replace_module is None + assert model.ttt_config.lora_target_modules == ("feed_forward",) + assert [ + name for name, module in model.named_modules() if isinstance(module, LoraInjectedLinear) + ] == ["backbone.feed_forward"] + + +@pytest.mark.parametrize( + ("values", "exception"), + ( + ({"lr": float("nan")}, ValueError), + ({"lora_alpha": float("inf")}, ValueError), + ({"momentum": "0.9"}, TypeError), + ({"momentum": -0.1}, ValueError), + ({"weight_decay": float("-inf")}, ValueError), + ({"weight_decay": -0.1}, ValueError), + ({"seed": True}, TypeError), + ({"seed": 1.5}, TypeError), + ({"initial_state_reset": 1}, TypeError), + ({"automatic_best_state_reset": None}, TypeError), + ({"eval_each_step": 0}, TypeError), + ({"gradient_clip": "false"}, TypeError), + ({"lora_target_replace_module": ""}, ValueError), + ({"lora_target_replace_module": 7}, TypeError), + ({"lora_target_modules": []}, TypeError), + ({"lora_target_modules": ()}, ValueError), + ({"lora_target_modules": ("query", 7)}, TypeError), + ({"lora_target_modules": ("query", "")}, ValueError), + ({"lora_target_modules": ("query", "query")}, ValueError), + ), +) +def test_ttt_config_rejects_invalid_optimizer_and_target_contracts( + values: dict[str, object], + exception: type[Exception], +) -> None: + with pytest.raises(exception): + TTTConfig(**values) + + +def test_ttt_adapter_initialization_is_seeded_and_preserves_ambient_rng() -> None: + first = DummyTTTModel() + torch.manual_seed(11) + # expected_next: (4,) + expected_next = torch.rand(4) + torch.manual_seed(11) + first._ttt_ensure_initialized() + # actual_next: (4,) + actual_next = torch.rand(4) + + torch.manual_seed(987654) + second = DummyTTTModel() + second._ttt_ensure_initialized() + + torch.testing.assert_close(actual_next, expected_next) + first_state = first._ttt_snapshot_lora_state() + second_state = second._ttt_snapshot_lora_state() + for first_module, second_module in zip(first_state, second_state, strict=True): + for name in first_module: + assert torch.equal(first_module[name], second_module[name]), name + + +def test_ttt_generic_replacements_exclude_reserved_vocabulary_ids() -> None: + model = DummyTTTModel() + canonical = {aa: idx + 4 for idx, aa in enumerate("ACDEFGHIKLMNPQRSTVWY")} + tokenizer = DummyTokenizer() + tokenizer.vocab = { + **canonical, + "": 24, + "": 25, + "": 26, + } + model.tokenizer = tokenizer + model.config = SimpleNamespace(vocab_size=27, model_type="esm3") + # input_ids: (1, 6) + input_ids = torch.tensor([[1, canonical["A"], 24, 25, 26, 2]]) + + replacements = model._ttt_replacement_tokens(input_ids) + trainable = model._ttt_non_special_mask(input_ids) + + assert replacements.tolist() == list(canonical.values()) + assert trainable.tolist() == [[False, True, False, False, False, False]] + + +def test_ttt_uneven_batch_samples_only_rows_with_residue_targets() -> None: + model = DummyTTTModel() + model._ttt_cfg.batch_size = 4 + # tokenized: (2, 3) + tokenized = torch.tensor( + [ + [model.tokenizer.cls_token_id, model.tokenizer.vocab["A"], 2], + [model.tokenizer.cls_token_id, model.tokenizer.eos_token_id, 0], + ] + ) + generator = torch.Generator().manual_seed(3) + + _, labels = model._ttt_sample_batch(tokenized, generator) + + assert labels.ne(-100).any(dim=1).all() + assert torch.equal(labels[labels.ne(-100)], torch.full((4,), model.tokenizer.vocab["A"])) + + +def test_ttt_rejects_all_ignored_inputs_before_adapter_injection() -> None: + model = DummyTTTModel() + # input_ids: (1, 3) + input_ids = torch.tensor( + [[model.tokenizer.cls_token_id, model.tokenizer.eos_token_id, model.tokenizer.pad_token_id]] + ) + + with pytest.raises(ValueError, match="no trainable biological residue"): + model.ttt(input_ids=input_ids) + + assert model._ttt_initialized is False + + +def test_ttt_rejects_dplm2_structure_tokens_before_adapter_injection() -> None: + model = DummyTTTModel() + tokenizer = DummyTokenizer() + tokenizer.struct_cls_token = "" + tokenizer._token_to_id = {"": 33} + model.tokenizer = tokenizer + model.config = SimpleNamespace(vocab_size=64, model_type="dplm2", struct_type=0) + # input_ids: (1, 5) + input_ids = torch.tensor([[1, tokenizer.vocab["A"], 33, 40, 2]]) + + with pytest.raises(ValueError, match="amino-acid-only"): + model.ttt(input_ids=input_ids) + + assert model._ttt_initialized is False + + +def test_ttt_only_lora_params_change_and_reset_restores_adapter() -> None: + model = DummyTTTModel() + model._ttt_ensure_initialized() + initial = {name: parameter.detach().clone() for name, parameter in model.named_parameters()} + + metrics = model.ttt(seq="ACDE") + + changed = [ + name + for name, parameter in model.named_parameters() + if not torch.equal(parameter.detach(), initial[name]) + ] + assert len(metrics["losses"]) == 1 + assert len(changed) > 0 + assert all("lora_" in name for name in changed) + + model.ttt_reset() + for name, parameter in model.named_parameters(): + torch.testing.assert_close(parameter.detach(), initial[name]) + + +def test_seed_and_initial_state_reset_reproduce_losses_and_updates() -> None: + """Repeat one seeded adaptation from the same initial adapter state.""" + + model = DummyTTTModel() + config = { + "steps": 2, + "ags": 2, + "batch_size": 2, + "mask_ratio": 0.5, + "bert_leave_prob": 0.1, + "bert_replace_prob": 0.2, + "seed": 7, + "initial_state_reset": True, + } + + first = model.ttt(seq="ACDE", ttt_config=config) + first_state = model._ttt_snapshot_lora_state() + second = model.ttt(seq="ACDE", ttt_config=config) + second_state = model._ttt_snapshot_lora_state() + + assert first == second + for first_module, second_module in zip(first_state, second_state, strict=True): + assert first_module.keys() == second_module.keys() + for name in first_module: + assert torch.equal(first_module[name], second_module[name]), name + + +def test_ttt_save_pretrained_round_trip_preserves_adapter_and_reset_state( + tmp_path: Path, +) -> None: + model = DummyPretrainedTTTModel(DummyPretrainedTTTConfig()).eval() + model._ttt_ensure_initialized() + with torch.no_grad(): + for index, module in enumerate(model._ttt_lora_modules(), start=1): + module.lora_up.weight.fill_(index / 10) + adapted_state = model._ttt_snapshot_lora_state() + + model.save_pretrained(tmp_path, safe_serialization=True) + reloaded = DummyPretrainedTTTModel.from_pretrained(tmp_path, local_files_only=True).eval() + + assert reloaded._ttt_initialized is True + assert reloaded.ttt_config.seed == 17 + assert reloaded.ttt_config.lora_target_modules == ("0", "2") + reloaded_state = reloaded._ttt_snapshot_lora_state() + for expected_module, actual_module in zip(adapted_state, reloaded_state, strict=True): + for name in expected_module: + torch.testing.assert_close(actual_module[name], expected_module[name]) + + expected_initial = [ + ( + module._ttt_initial_lora_down.detach().clone(), + module._ttt_initial_lora_up.detach().clone(), + ) + for module in reloaded._ttt_lora_modules() + ] + reloaded.ttt_reset() + for module, (expected_down, expected_up) in zip( + reloaded._ttt_lora_modules(), expected_initial, strict=True + ): + torch.testing.assert_close(module.lora_down.weight, expected_down) + torch.testing.assert_close(module.lora_up.weight, expected_up) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.network +@pytest.mark.checkpoint +@pytest.mark.parametrize("model_key", list(MODEL_REGISTRY)) +def test_sequence_model_ttt_smoke(model_key: str) -> None: + config = MODEL_REGISTRY[model_key] + model_cls = LOCAL_MODEL_CLASSES[model_key] + model = ( + model_cls.from_pretrained( + config["fast_path"], + dtype=torch.float32, + ) + .eval() + .cuda() + ) + metrics = model.ttt( + seq=TEST_SEQUENCE, + ttt_config={ + "steps": 1, + "ags": 1, + "batch_size": 1, + "crop_size": 64, + "lora_rank": 2, + "lora_alpha": 1.0, + }, + ) + + assert len(metrics["losses"]) == 1 + assert callable(model.ttt_reset) + model.ttt_reset() + del model + torch.cuda.empty_cache() + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +def test_esmfold2_ttt_smoke() -> None: + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model + + config = STRUCTURE_MODEL_REGISTRY["esmfold2_fast"] + model = ( + ESMFold2Model.from_pretrained( + config["fast_path"], + load_esmc=True, + dtype=torch.float32, + ) + .eval() + .cuda() + ) + + result = model.fold_protein( + TEST_SEQUENCE, + num_loops=1, + num_sampling_steps=1, + num_diffusion_samples=1, + seed=0, + ttt=True, + ttt_config={ + "steps": 1, + "ags": 1, + "batch_size": 1, + "crop_size": 64, + "lora_rank": 2, + "lora_alpha": 1.0, + }, + ) + + assert result.ttt_metrics is not None + assert len(result.ttt_metrics["losses"]) == 1 + assert len(result.ttt_metrics["step_plddts"]) == 2 + assert result.ttt_metrics["best_step"] in {0, 1} + + del model, result + torch.cuda.empty_cache() diff --git a/tests/parity/__init__.py b/tests/parity/__init__.py new file mode 100644 index 0000000..ea56895 --- /dev/null +++ b/tests/parity/__init__.py @@ -0,0 +1 @@ +"""Numerical and behavioral parity tests.""" diff --git a/tests/parity/fixtures/esmc_biological_holdout.json b/tests/parity/fixtures/esmc_biological_holdout.json new file mode 100644 index 0000000..da1926e --- /dev/null +++ b/tests/parity/fixtures/esmc_biological_holdout.json @@ -0,0 +1,40 @@ +{ + "cases": [ + { + "case_id": "1crn-a", + "sequence": "TTCCPSIVARSNFNVCRLPGTPEAICATYTGCIIIPGATCPGDYAN", + "sequence_sha256": "6f284e8edba221b51718d8f6864cb976e35f7708dd5bccd0c144b6ebe307be6d", + "source": "RCSB 1CRN:A", + "source_sha256": "23787562c427d7c1abe5420e86d5f1d0a6c7007dec1e8ce85645a6d69c32e8ba" + }, + { + "case_id": "1pga-a", + "sequence": "MTYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE", + "sequence_sha256": "7e859d82171047700fd3e9632f7a47eab4a39baedc8c3316d2fc62d3ce2260bb", + "source": "RCSB 1PGA:A", + "source_sha256": "e8d9e171a9ca2b7e66b10568bf6d0122da8450e57ad632a7b4c13c6a71f7edf2" + }, + { + "case_id": "5pti-a", + "sequence": "RPDFCLEPPYTGPCKARIIRYFYNAKAGLCQTFVYGGCRAKRNNFKSAEDCMRTCGGA", + "sequence_sha256": "70b8ee8cfe717fd0ba7d4d5a9b53a0c7587ddc9b2623765f9866fe1773f4febb", + "source": "RCSB 5PTI:A", + "source_sha256": "d39881b81af60445d9cb7ca59c6b4858a1efd989d5c404e99730c475b26fb55b" + }, + { + "case_id": "1ubq-a", + "sequence": "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG", + "sequence_sha256": "233b4b0b8c4616095bc3249f9375fc345d32af7deaef07117d0383c51d6f19aa", + "source": "RCSB 1UBQ:A", + "source_sha256": "056f98710cb2b36f633c45e41902a02eb446e82871da21ff2dd44f74a56ca0f6" + }, + { + "case_id": "1ema-a", + "sequence": "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKRHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK", + "sequence_sha256": "01926ebded8ef06ae535f928f56150e3694474de61ad9e8a5288c7bdb0e13855", + "source": "RCSB 1EMA:A", + "source_sha256": "b7d583424b6f960a7421fd4c88cdb6bbb10025667cd863d8cd14f70bc775aceb" + } + ], + "schema_version": 1 +} diff --git a/tests/parity/support/__init__.py b/tests/parity/support/__init__.py new file mode 100644 index 0000000..2b34ea9 --- /dev/null +++ b/tests/parity/support/__init__.py @@ -0,0 +1 @@ +"""Parity fixtures and reference adapters.""" diff --git a/tests/parity/support/esmc_calibration.py b/tests/parity/support/esmc_calibration.py new file mode 100644 index 0000000..13d5a36 --- /dev/null +++ b/tests/parity/support/esmc_calibration.py @@ -0,0 +1,171 @@ +"""Compact, immutable biological sequences for ESMC backend calibration.""" + +from __future__ import annotations + +import hashlib +import json +import random +from collections.abc import Mapping +from pathlib import Path +from typing import Any + + +FIXTURE_PATH = Path(__file__).parents[1] / "fixtures" / "esmc_biological_holdout.json" +CASE_IDS = ("1crn-a", "1pga-a", "5pti-a", "1ubq-a", "1ema-a") +ESMC_BOUNDARY_LENGTHS = (13, 15, 16, 17, 29, 31, 32, 33, 61, 127, 128, 129) +ESMC_CALIBRATION_SEED = 42 +ESMC_PANEL_DEFINITION_SCHEMA_VERSION = 1 +CANONICAL_AA_ALPHABET = "ACDEFGHIKLMNPQRSTVWY" +CANONICAL_AAS = frozenset(CANONICAL_AA_ALPHABET) +PANEL_KINDS = ("generated_kernel_boundary", "real_biological_holdout") + + +def load_esmc_biological_holdout( + path: Path = FIXTURE_PATH, +) -> tuple[dict[str, str], ...]: + """Load the five pinned sequence/source pairs and fail closed on drift.""" + + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("schema_version") != 1: + raise ValueError("Unsupported ESMC biological-holdout schema") + cases = payload.get("cases") + if not isinstance(cases, list) or tuple(case.get("case_id") for case in cases) != CASE_IDS: + raise ValueError("ESMC biological-holdout inventory differs from the release contract") + for case in cases: + _validate_case(case) + return tuple(dict(case) for case in cases) + + +def _validate_case(case: Mapping[str, Any]) -> None: + case_id = str(case.get("case_id")) + expected_fields = { + "case_id", + "sequence", + "sequence_sha256", + "source", + "source_sha256", + } + if set(case) != expected_fields: + raise ValueError(f"{case_id}: biological-holdout fields differ from the contract") + sequence = case.get("sequence") + if ( + not isinstance(sequence, str) + or not sequence + or not sequence.isupper() + or not set(sequence).issubset(CANONICAL_AAS) + ): + raise ValueError(f"{case_id}: sequence must contain canonical uppercase amino acids") + sequence_sha256 = hashlib.sha256(sequence.encode("ascii")).hexdigest() + if case.get("sequence_sha256") != sequence_sha256: + raise ValueError(f"{case_id}: sequence digest mismatch") + source = case.get("source") + if not isinstance(source, str) or not source.startswith("RCSB "): + raise ValueError(f"{case_id}: source must identify an RCSB chain") + source_sha256 = case.get("source_sha256") + if not isinstance(source_sha256, str) or len(source_sha256) != 64: + raise ValueError(f"{case_id}: source digest is not pinned") + try: + bytes.fromhex(source_sha256) + except ValueError as error: + raise ValueError(f"{case_id}: source digest is not hexadecimal") from error + + +def _sequence_sha256(sequence: str) -> str: + return hashlib.sha256(sequence.encode("ascii")).hexdigest() + + +def generated_esmc_boundary_cases() -> tuple[dict[str, object], ...]: + """Return the exact seed-locked kernel-boundary sequence panel.""" + + generator = random.Random(ESMC_CALIBRATION_SEED) + cases: list[dict[str, object]] = [] + for length in ESMC_BOUNDARY_LENGTHS: + sequence = "M" + "".join(generator.choices(CANONICAL_AA_ALPHABET, k=length - 1)) + cases.append( + { + "case_id": f"generated-boundary-{length}", + "sequence": sequence, + "sequence_length": length, + "sequence_sha256": _sequence_sha256(sequence), + } + ) + return tuple(cases) + + +def biological_esmc_holdout_cases() -> tuple[dict[str, object], ...]: + """Return the exact source- and sequence-pinned biological panel.""" + + return tuple( + { + **case, + "sequence_length": len(case["sequence"]), + } + for case in load_esmc_biological_holdout() + ) + + +def esmc_calibration_batches() -> tuple[dict[str, object], ...]: + """Build the two native-reference batches without embedding their digest.""" + + return ( + { + "kind": "generated_kernel_boundary", + "seed": ESMC_CALIBRATION_SEED, + "cases": [dict(case) for case in generated_esmc_boundary_cases()], + }, + { + "kind": "real_biological_holdout", + "seed": ESMC_CALIBRATION_SEED, + "cases": [dict(case) for case in biological_esmc_holdout_cases()], + }, + ) + + +def validate_esmc_calibration_batch(batch: Mapping[str, Any]) -> dict[str, object]: + """Validate one immutable panel and return its canonical identity and digest.""" + + if set(batch) != {"kind", "seed", "cases"}: + raise ValueError("ESMC calibration batch fields differ from the release contract") + kind = batch.get("kind") + if kind not in PANEL_KINDS: + raise ValueError(f"Unsupported ESMC calibration panel: {kind!r}") + if batch.get("seed") != ESMC_CALIBRATION_SEED: + raise ValueError("ESMC calibration seed differs from the release contract") + expected = next(item for item in esmc_calibration_batches() if item["kind"] == kind) + if batch != expected: + raise ValueError(f"ESMC calibration panel {kind!r} differs from the release contract") + expected_cases = expected["cases"] + if not isinstance(expected_cases, list): + raise ValueError(f"ESMC calibration panel {kind!r} cases are not an ordered list") + + definition = { + "schema_version": ESMC_PANEL_DEFINITION_SCHEMA_VERSION, + "kind": kind, + "seed": ESMC_CALIBRATION_SEED, + "cases": [dict(case) for case in expected_cases], + } + encoded = json.dumps( + definition, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return { + **definition, + "definition_sha256": hashlib.sha256(encoded).hexdigest(), + } + + +__all__ = [ + "CANONICAL_AA_ALPHABET", + "CASE_IDS", + "ESMC_BOUNDARY_LENGTHS", + "ESMC_CALIBRATION_SEED", + "ESMC_PANEL_DEFINITION_SCHEMA_VERSION", + "FIXTURE_PATH", + "PANEL_KINDS", + "biological_esmc_holdout_cases", + "esmc_calibration_batches", + "generated_esmc_boundary_cases", + "load_esmc_biological_holdout", + "validate_esmc_calibration_batch", +] diff --git a/tests/parity/support/native_reference.py b/tests/parity/support/native_reference.py new file mode 100644 index 0000000..ef70be2 --- /dev/null +++ b/tests/parity/support/native_reference.py @@ -0,0 +1,930 @@ +"""Execute pinned official models inside their native reference containers. + +This module intentionally has no FastPLMs import. It invokes only an official +adapter, applies the independent compliance state transform, and writes a +normalized result that the candidate container can consume later. +""" + +from __future__ import annotations + +import argparse +import contextlib +import gc +import hashlib +import importlib +import importlib.metadata +import json +import os +import platform +import shutil +import subprocess +import tempfile +import torch +import torch.nn as nn +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any +from safetensors.torch import save_file + +from tests.parity.support.reference_adapters import ( + OfficialGenerationUnavailable, + snapshot_path, +) +from tests.parity.support.semantic_config import transformed_semantic_config +from tests.parity.support.state_transforms import ( + transform_parameter_names, + transform_preserves_aliases, + transform_state, +) +from tools.remote.biohub_reference_environment import ( + validate_biohub_reference_environment_evidence, +) +from tools.remote.reference_source_attestation import ( + validate_reference_sources_evidence, +) + + +SCHEMA_VERSION = 1 +_ADAPTER_PREFIX = "tests.parity.support.reference_adapters." +_BIOHUB_REFERENCE_FAMILIES = frozenset({"esm_plusplus", "esm3", "esmfold2"}) +_BIOHUB_REFERENCE_SOURCE_NAMES = ("biohub-esm", "biohub-transformers") +_SPECIAL_TOKEN_FIELDS = ( + "pad_token_id", + "bos_token_id", + "cls_token_id", + "eos_token_id", + "mask_token_id", + "unk_token_id", +) +_TOKENIZER_SETTINGS = ( + {"padding": True}, + {"padding": "max_length", "truncation": True, "max_length": 12}, + {"padding": True, "truncation": True, "max_length": 5}, +) + + +def _tensor_digest(tensor: torch.Tensor) -> dict[str, Any]: + # tensor: (...) + # value: (...) + value = tensor.detach().cpu().contiguous() + raw = value.view(torch.uint8).numpy().tobytes() + return { + "shape": list(value.shape), + "dtype": str(value.dtype), + "sha256": hashlib.sha256(raw).hexdigest(), + } + + +def _cuda_driver_version() -> str: + """Read the exact host driver exposed to the reference container.""" + + try: + completed = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (FileNotFoundError, subprocess.SubprocessError) as error: + raise RuntimeError("Native compliance requires the exact NVIDIA driver version.") from error + versions = { + line.strip() for line in completed.stdout.splitlines() if line.strip() + } + if len(versions) != 1: + raise RuntimeError( + "Native compliance requires one unambiguous NVIDIA driver version." + ) + return versions.pop() + + +def _environment_metadata() -> dict[str, object]: + """Describe the isolated native environment without host-specific paths.""" + + distributions: dict[str, str] = {} + for distribution in importlib.metadata.distributions(): + name = distribution.metadata.get("Name") + if isinstance(name, str) and name: + distributions[name.lower()] = distribution.version + cuda_properties = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None + uname = platform.uname() + return { + "cuda_device": cuda_properties.name if cuda_properties is not None else "unavailable", + "cuda_device_capability": ( + list(torch.cuda.get_device_capability(0)) if cuda_properties is not None else None + ), + "cuda_total_memory": ( + int(cuda_properties.total_memory) if cuda_properties is not None else None + ), + "cuda_runtime": str(torch.version.cuda or "unavailable"), + "cuda_driver": _cuda_driver_version(), + "packages": json.dumps(distributions, separators=(",", ":"), sort_keys=True), + "platform_machine": platform.machine(), + "python": platform.python_version(), + "torch": torch.__version__, + "uname": { + "system": uname.system, + "release": uname.release, + "version": uname.version, + "machine": uname.machine, + }, + } + + +def _adapter_reference_sources( + adapter: Any, + request: Mapping[str, Any], +) -> dict[str, dict[str, object]] | None: + """Read and validate an optional named official-source provenance hook.""" + + hook = getattr(adapter, "reference_sources", None) + required = request.get("family") in _BIOHUB_REFERENCE_FAMILIES + if hook is None: + if required: + raise RuntimeError( + f"{request.get('model_id')}: Biohub adapter omits source attestations." + ) + return None + if not callable(hook): + raise RuntimeError("Official adapter reference-sources hook is not callable.") + return validate_reference_sources_evidence( + hook(), + required_sources=_BIOHUB_REFERENCE_SOURCE_NAMES, + ) + + +def _adapter_reference_environment( + adapter: Any, + request: Mapping[str, Any], +) -> dict[str, object] | None: + """Read and validate the locked runtime/image evidence for Biohub adapters.""" + + hook = getattr(adapter, "reference_environment", None) + required = request.get("family") in _BIOHUB_REFERENCE_FAMILIES + if hook is None: + if required: + raise RuntimeError( + f"{request.get('model_id')}: Biohub adapter omits reference environment." + ) + return None + if not callable(hook): + raise RuntimeError("Official adapter reference-environment hook is not callable.") + try: + lock_root = Path(os.environ["FASTPLMS_BIOHUB_LOCK_ROOT"]) + contract = Path(os.environ["FASTPLMS_BIOHUB_LOCK_CONTRACT"]) + except KeyError as error: + raise RuntimeError("Biohub lock validation environment is incomplete.") from error + return validate_biohub_reference_environment_evidence( + hook(), + repository_root=lock_root, + contract_path=contract, + ) + + +def _tokenizer_asset_contract(request: Mapping[str, Any]) -> dict[str, Any]: + files = request.get("tokenizer_files", []) + if request["tokenizer_mode"] != "tokenizer": + return {} + if not files: + # Some official implementations, including ESM3, define their exact + # tokenizer vocabulary in pinned upstream source rather than checkpoint + # assets. The source revision and behavior contract below cover that + # case; checkpoint-backed tokenizers still hash every declared file. + return {} + snapshot = snapshot_path(request["reference_repo_id"], request["reference_revision"]) + result: dict[str, Any] = {} + for relative_name in files: + relative = Path(relative_name) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"Unsafe tokenizer asset path: {relative_name!r}") + path = snapshot.joinpath(*relative.parts) + if not path.is_file(): + raise FileNotFoundError(f"Official tokenizer asset is missing: {path}") + content = path.read_bytes() + result[relative.as_posix()] = { + "size": len(content), + "sha256": hashlib.sha256(content).hexdigest(), + } + return result + + +def _state_contract(model: nn.Module, transform_name: str) -> dict[str, Any]: + state = transform_state(transform_name, model.state_dict()) + tensors: dict[str, Any] = {} + for name, value in sorted(state.items()): + if not torch.is_tensor(value): + raise TypeError(f"Official state entry {name!r} is not a tensor") + tensors[name] = _tensor_digest(value) + + by_parameter: dict[int, set[str]] = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + mapped = transform_parameter_names(transform_name, name) + by_parameter.setdefault(id(parameter), set()).update(mapped) + aliases = ( + sorted(sorted(names) for names in by_parameter.values() if len(names) > 1) + if transform_preserves_aliases(transform_name) + else [] + ) + return {"tensors": tensors, "aliases": aliases} + + +def _normalize_tokenizer_error(message: str) -> str: + """Remove a dependency-list difference between Transformers v4 and v5.""" + + return message.replace( + "python, numpy, pytorch or tensorflow object.", + "python, numpy or pytorch object.", + ).replace("python, numpy, or pytorch object.", "python, numpy or pytorch object.") + + +def _token_result(tokenizer: object, sequences: Sequence[str], options: Mapping[str, Any]) -> Any: + try: + encoded = tokenizer(sequences, return_tensors="pt", **options) + except Exception as error: + return [ + "error", + type(error).__module__, + type(error).__qualname__, + _normalize_tokenizer_error(str(error)), + ] + normalized = { + key: value.tolist() if torch.is_tensor(value) else value for key, value in encoded.items() + } + return ["ok", normalized] + + +def _tokenizer_contract( + tokenizer: object, + edge_sequences: Sequence[str], + tokenizer_mode: str, +) -> dict[str, Any] | None: + if tokenizer_mode != "tokenizer": + return None + return { + "vocab": tokenizer.get_vocab(), + "special_ids": {name: getattr(tokenizer, name, None) for name in _SPECIAL_TOKEN_FIELDS}, + "behavior": [ + {"options": options, "result": _token_result(tokenizer, edge_sequences, options)} + for options in _TOKENIZER_SETTINGS + ], + } + + +def _to_device(values: Mapping[str, Any], device: torch.device) -> dict[str, torch.Tensor]: + return {name: value.to(device) for name, value in values.items() if torch.is_tensor(value)} + + +def _prepare_dplm2_inputs( + sequences: Sequence[str], + tokenizer: object, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Construct aligned structure and amino-acid tracks for DPLM2.""" + + vocabulary = tokenizer.get_vocab() + required = { + "", + "", + "", + "", + "", + } + required.update({residue for sequence in sequences for residue in sequence}) + missing = sorted(required.difference(vocabulary)) + if missing: + raise RuntimeError(f"Official DPLM2 tokenizer omits input tokens: {missing}") + structure_token_id = 50 + if structure_token_id >= len(vocabulary): + raise RuntimeError("Official DPLM2 vocabulary omits structure token 50") + track_length = max(map(len, sequences)) + 2 + pad_id = vocabulary[""] + # input_ids: (len(sequences), 2 * track_length) + input_ids = torch.full( + (len(sequences), 2 * track_length), + pad_id, + dtype=torch.long, + device=device, + ) + residue_mask = torch.zeros_like(input_ids, dtype=torch.bool) + for row_index, sequence in enumerate(sequences): + residue_count = len(sequence) + structure = [ + vocabulary[""], + *([structure_token_id] * residue_count), + vocabulary[""], + ] + amino_acids = [ + vocabulary[""], + *(vocabulary[residue] for residue in sequence), + vocabulary[""], + ] + # input_ids[row_index, :len(structure)]: (...) + input_ids[row_index, : len(structure)] = torch.tensor(structure, device=device) + aa_start = track_length + # input_ids[row_index, aa_start:aa_start + len(amino_acids)]: (...) + input_ids[row_index, aa_start : aa_start + len(amino_acids)] = torch.tensor( + amino_acids, + device=device, + ) + residue_mask[row_index, 1 : 1 + residue_count] = True + residue_mask[ + row_index, + aa_start + 1 : aa_start + 1 + residue_count, + ] = True + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(pad_id).long(), + }, residue_mask + + +def _prepare_inputs( + request: Mapping[str, Any], + tokenizer: object, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + sequences = request["sequences"] + if request["tokenizer_mode"] == "sequence": + prepared = dict(tokenizer.get_batch_kwargs(sequences, device=device)) + required = ( + "input_ids", + "within_seq_position_ids", + "global_position_ids", + "sequence_ids", + ) + missing = [name for name in required if not torch.is_tensor(prepared.get(name))] + if missing: + raise RuntimeError(f"Official sequence adapter omits tensors: {missing}") + # E1's preparer also returns labels and human-readable context records. + # They are data-loader outputs, not arguments to the public inference + # computation, and therefore do not belong in a tensor golden. + inputs = {name: prepared[name] for name in required} + # residue_mask: (b, l) + residue_mask = inputs["sequence_ids"].ge(0) + return inputs, residue_mask + if request["family"] == "dplm2": + return _prepare_dplm2_inputs(sequences, tokenizer, device) + + encoded = _to_device( + tokenizer(sequences, return_tensors="pt", padding=True), + device, + ) + # input_ids: (b, l) + input_ids = encoded["input_ids"] + # residue_mask: (b, l) + residue_mask = encoded["attention_mask"].bool() + for token_id in getattr(tokenizer, "all_special_ids", ()): + residue_mask &= input_ids.ne(token_id) + inputs = { + name: value for name, value in encoded.items() if name in {"input_ids", "attention_mask"} + } + if request["architecture"] == "ESMC": + # inputs['sequence_id']: (b, l) + inputs["sequence_id"] = encoded["attention_mask"].bool() + return inputs, residue_mask + + +def _output_tensors(output: object) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + raw_hidden_states = getattr(output, "hidden_states", None) + if torch.is_tensor(raw_hidden_states): + hidden_states = tuple(raw_hidden_states) + else: + hidden_states = tuple(raw_hidden_states or ()) + if not hidden_states: + raise RuntimeError("Official inference omitted hidden states") + for index, value in enumerate(hidden_states): + # tensors[f'output__hidden_{index:04d}']: (..., d) + tensors[f"output__hidden_{index:04d}"] = value.detach().cpu().contiguous().clone() + last_hidden = getattr(output, "last_hidden_state", None) + if last_hidden is None: + # last_hidden: (..., d) + last_hidden = hidden_states[-1] + # tensors['output__last_hidden_state']: (..., d) + tensors["output__last_hidden_state"] = last_hidden.detach().cpu().contiguous().clone() + logits = getattr(output, "logits", None) + if logits is not None: + # tensors['output__logits']: (..., c) + tensors["output__logits"] = logits.detach().cpu().contiguous().clone() + return tensors + + +@contextlib.contextmanager +def _strict_fp32_matmul(): + """Disable TF32 locally while producing FP32 compliance outputs.""" + + try: + old_fp32_precision = torch.backends.fp32_precision + old_matmul_precision = torch.backends.cuda.matmul.fp32_precision + old_cudnn_precision = torch.backends.cudnn.fp32_precision + except AttributeError: + old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + old_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32 + torch.backends.cudnn.allow_tf32 = old_cudnn_tf32 + return + torch.backends.fp32_precision = "ieee" + torch.backends.cuda.matmul.fp32_precision = "ieee" + torch.backends.cudnn.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.fp32_precision = old_fp32_precision + torch.backends.cuda.matmul.fp32_precision = old_matmul_precision + torch.backends.cudnn.fp32_precision = old_cudnn_precision + + +def _inference_tensors( + model: nn.Module, + tokenizer: object, + request: Mapping[str, Any], + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + use_native_autocast = ( + request["family"] in {"dplm", "dplm2", "esm2", "esm3"} and dtype == torch.bfloat16 + ) + if use_native_autocast: + # These pinned implementations use AMP for native mixed precision. + # A static BF16 cast breaks their intentional FP32 softmax and rotary + # operations before the following matrix reduction. + model = model.to(device=device, dtype=torch.float32).eval() + else: + model = model.to(device=device, dtype=dtype).eval() + torch.manual_seed(int(request["seed"])) + inputs, residue_mask = _prepare_inputs(request, tokenizer, device) + if dtype == torch.float32: + numeric_context = _strict_fp32_matmul() + elif use_native_autocast: + numeric_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16) + else: + numeric_context = contextlib.nullcontext() + with torch.inference_mode(), numeric_context: + output = model(**inputs, output_hidden_states=True) + tensors = { + f"input__{name}": value.detach().cpu().contiguous().clone() + for name, value in inputs.items() + } + # tensors['residue_mask']: (b, l) + tensors["residue_mask"] = residue_mask.detach().cpu().contiguous().clone() + tensors.update(_output_tensors(output)) + del output + return tensors + + +def _ankh_generation_contract( + adapter: Any, + request: Mapping[str, Any], + device: torch.device, +) -> dict[str, Any]: + """Run official ANKH generation from an explicit task prompt. + + ANKH is a T5 checkpoint. Its decoder input is task-specific, so this + contract deliberately supplies a short decoder prompt instead of shifting + or otherwise reusing the encoder source tokens. + """ + + load_seq2seq = getattr(adapter, "load_official_seq2seq", None) + if not callable(load_seq2seq): + raise RuntimeError("The ANKH reference adapter omits load_official_seq2seq().") + generation_model, generation_tokenizer = load_seq2seq( + reference_repo_id=request["reference_repo_id"], + reference_revision=request["reference_revision"], + device=device, + dtype=torch.float32, + ) + source_text = "M S T N P K" + decoder_prompt_text = "A C" + try: + encoded = _to_device( + generation_tokenizer(source_text, return_tensors="pt"), + device, + ) + prompt = _to_device( + generation_tokenizer( + decoder_prompt_text, + return_tensors="pt", + add_special_tokens=False, + ), + device, + ) + prompt_ids = prompt.get("input_ids") + if not torch.is_tensor(prompt_ids) or prompt_ids.ndim != 2: + raise RuntimeError("Official ANKH tokenizer returned invalid decoder prompt IDs.") + decoder_start_token_id = getattr( + generation_model.config, + "decoder_start_token_id", + None, + ) + if not isinstance(decoder_start_token_id, int): + raise RuntimeError("Official ANKH config omits decoder_start_token_id.") + # decoder_input_ids: (...) + decoder_input_ids = torch.cat( + ( + prompt_ids.new_full((prompt_ids.shape[0], 1), decoder_start_token_id), + prompt_ids, + ), + dim=1, + ) + decoder_attention_mask = torch.ones_like(decoder_input_ids) + kwargs = { + "do_sample": False, + "max_new_tokens": 4, + "num_beams": 1, + "use_cache": True, + } + torch.manual_seed(int(request["seed"])) + torch.cuda.manual_seed_all(int(request["seed"])) + with torch.inference_mode(), _strict_fp32_matmul(): + generated = generation_model.generate( + input_ids=encoded["input_ids"], + attention_mask=encoded["attention_mask"], + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + **kwargs, + ) + if not torch.is_tensor(generated): + raise RuntimeError("Official ANKH generation did not return output tokens.") + decoder_fingerprint = _tensor_digest(decoder_input_ids)["sha256"] + return { + "interface": "T5ForConditionalGeneration.generate", + "source_text": source_text, + "input_ids": encoded["input_ids"].detach().cpu().tolist(), + "attention_mask": encoded["attention_mask"].detach().cpu().tolist(), + "decoder_prompt_text": decoder_prompt_text, + "decoder_prompt_contract": "explicit-task-prompt", + "decoder_input_ids": decoder_input_ids.detach().cpu().tolist(), + "decoder_attention_mask": decoder_attention_mask.detach().cpu().tolist(), + "decoder_input_fingerprint": decoder_fingerprint, + "kwargs": kwargs, + "output_tokens": generated.detach().cpu().tolist(), + "seed": int(request["seed"]), + } + finally: + del generation_model + gc.collect() + torch.cuda.empty_cache() + + +def _generation_contract( + model: nn.Module | None, + tokenizer: object, + request: Mapping[str, Any], + device: torch.device, + *, + adapter: Any | None = None, +) -> dict[str, Any] | None: + """Run a deterministic official generation call required by the manifest.""" + + family = request["family"] + generation_policy = request.get("generation_policy", "required") + if generation_policy == "not_applicable": + return None + if generation_policy not in {"required", "official_unavailable"}: + raise RuntimeError( + f"{request['model_id']}: unknown generation policy {generation_policy!r}." + ) + if family == "ankh": + if adapter is None: + raise RuntimeError("ANKH generation requires the pinned official adapter.") + return _ankh_generation_contract(adapter, request, device) + if family not in {"dplm", "dplm2"}: + raise RuntimeError( + f"{request['model_id']}: generation is {generation_policy} but the " + f"{family!r} adapter has no generation contract." + ) + if model is None: + raise RuntimeError(f"{request['model_id']}: official generation model is missing.") + max_iter = 4 + if family == "dplm": + encoded = tokenizer("ACDEFG", return_tensors="pt") + # input_tokens: (...) + input_tokens = encoded["input_ids"].to(device) + kwargs: dict[str, Any] = { + "max_iter": max_iter, + "sampling_strategy": "argmax", + "disable_resample": True, + } + else: + vocabulary = tokenizer.get_vocab() + required = ( + "", + "", + "", + "", + "A", + ) + missing = sorted(name for name in required if name not in vocabulary) + if missing: + raise RuntimeError(f"Official DPLM2 tokenizer omits generation tokens: {missing}") + structure = [vocabulary[""], 50, 50, 50, 50, vocabulary[""]] + amino_acids = [vocabulary[""], *([vocabulary["A"]] * 4), vocabulary[""]] + # input_tokens: (1, 12) + input_tokens = torch.tensor([structure + amino_acids], device=device) + kwargs = { + "max_iter": max_iter, + "sampling_strategy": "argmax", + "unmasking_strategy": "deterministic", + } + + model = model.to(device=device, dtype=torch.float32).eval() + torch.manual_seed(int(request["seed"])) + torch.cuda.manual_seed_all(int(request["seed"])) + with torch.inference_mode(), _strict_fp32_matmul(): + generated = model.generate(input_tokens=input_tokens, **kwargs) + if isinstance(generated, Mapping): + generated = generated.get("output_tokens") + if not torch.is_tensor(generated): + raise RuntimeError("Official DPLM generation did not return output tokens") + return { + "input_tokens": input_tokens.detach().cpu().tolist(), + "kwargs": kwargs, + "output_tokens": generated.detach().cpu().tolist(), + "seed": int(request["seed"]), + } + + +def _record_generation_contract( + metadata: dict[str, Any], + model: nn.Module | None, + tokenizer: object, + request: Mapping[str, Any], + device: torch.device, + *, + adapter: Any, +) -> None: + """Apply the manifest generation policy and fail closed on missing evidence.""" + + policy = request.get("generation_policy", "required") + try: + generation = _generation_contract( + model, + tokenizer, + request, + device, + adapter=adapter, + ) + except OfficialGenerationUnavailable as error: + metadata["generation_limitation"] = _validated_generation_limitation( + request, + error, + ) + return + + if policy == "official_unavailable": + raise RuntimeError( + f"{request['model_id']}: official sampler executed despite an " + "official_unavailable request." + ) + if policy == "required": + if not isinstance(generation, dict): + raise RuntimeError( + f"{request['model_id']}: required official generation evidence is missing." + ) + metadata["generation"] = generation + return + if policy == "not_applicable": + if generation is not None: + raise RuntimeError( + f"{request['model_id']}: not_applicable generation produced evidence." + ) + return + raise RuntimeError(f"{request['model_id']}: unknown generation policy {policy!r}.") + + +def _validated_generation_limitation( + request: Mapping[str, Any], + error: OfficialGenerationUnavailable, +) -> dict[str, str]: + """Accept only the exact limitation declared by a native request.""" + + limitation = error.as_record() + expected = request.get("official_generation_limitation") + if request.get("generation_policy", "required") != "official_unavailable": + raise RuntimeError( + f"{request['model_id']}: official generation is required; " + "a public sampler failure cannot become a native result." + ) from error + if expected != limitation: + raise RuntimeError( + f"{request['model_id']}: official generation limitation differs " + "from the manifest-derived request." + ) from error + return limitation + + +def run_request(request_path: Path, output_root: Path) -> Path: + """Execute one official request and atomically publish its normalized result.""" + + request = json.loads(request_path.read_text(encoding="utf-8")) + if request.get("schema_version") != SCHEMA_VERSION: + raise ValueError(f"Unsupported native-reference schema in {request_path}") + adapter_name = request.get("adapter") + if not isinstance(adapter_name, str) or not adapter_name.startswith(_ADAPTER_PREFIX): + raise ValueError(f"Invalid official adapter in {request_path}") + if not torch.cuda.is_available(): + raise RuntimeError("Native BF16 compliance requires CUDA") + + adapter = importlib.import_module(adapter_name) + load_kwargs: dict[str, Any] = {} + if request.get("oracle_assets"): + load_kwargs["oracle_assets"] = request["oracle_assets"] + model, tokenizer = adapter.load_official_model( + reference_repo_id=request["reference_repo_id"], + reference_revision=request["reference_revision"], + device=torch.device("cpu"), + dtype=None, + **load_kwargs, + ) + reference_sources = _adapter_reference_sources(adapter, request) + reference_environment = _adapter_reference_environment(adapter, request) + core = getattr(model, "model", model) + metadata: dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "model_id": request["model_id"], + "family": request["family"], + "reference_repo_id": request["reference_repo_id"], + "reference_revision": request["reference_revision"], + "reference_files": request["reference_files"], + "state_transform": request["state_transform"], + "environment": _environment_metadata(), + "semantic_config": transformed_semantic_config(core, request["state_transform"]), + "state": _state_contract(core, request["state_transform"]), + "tokenizer": _tokenizer_contract( + tokenizer, + request["edge_sequences"], + request["tokenizer_mode"], + ), + "tokenizer_assets": _tokenizer_asset_contract(request), + } + if reference_sources is not None: + metadata["reference_sources"] = reference_sources + if reference_environment is not None: + metadata["reference_environment"] = reference_environment + + device = torch.device("cuda") + # ANKH's native encoder wrapper intentionally has no decoder. Defer its + # generation contract until encoder inference is complete, then release the + # encoder before loading the complete official T5 checkpoint. + defer_generation = ( + request["family"] == "ankh" and request.get("generation_policy", "required") == "required" + ) + if not defer_generation: + _record_generation_contract( + metadata, + model, + tokenizer, + request, + device, + adapter=adapter, + ) + precision_tensors: dict[str, dict[str, torch.Tensor]] = {} + if request["deep_reference"]: + precision_tensors["fp32"] = _inference_tensors( + model, tokenizer, request, device, torch.float32 + ) + precision_tensors["bf16"] = _inference_tensors( + model, tokenizer, request, device, torch.bfloat16 + ) + metadata["precision_tensor_keys"] = { + precision: sorted(tensors) for precision, tensors in precision_tensors.items() + } + calibration_tensors: dict[str, dict[str, torch.Tensor]] = {} + calibration_batches = request.get("calibration_batches", []) + if calibration_batches: + if request["family"] != "esm_plusplus": + raise ValueError("Calibration batches are reserved for ESM++/ESMC requests") + for batch in calibration_batches: + kind = batch.get("kind") + cases = batch.get("cases") + if not isinstance(kind, str) or not isinstance(cases, list) or not cases: + raise ValueError(f"{request['model_id']}: invalid ESMC calibration batch") + sequences = [case["sequence"] for case in cases] + calibration_request = {**request, "sequences": sequences} + calibration_tensors[kind] = _inference_tensors( + model, + tokenizer, + calibration_request, + device, + torch.bfloat16, + ) + metadata["calibration_batches"] = calibration_batches + metadata["calibration_tensor_keys"] = { + kind: sorted(tensors) for kind, tensors in calibration_tensors.items() + } + + if defer_generation: + del model, core + model = None + core = None + gc.collect() + torch.cuda.empty_cache() + _record_generation_contract( + metadata, + None, + tokenizer, + request, + device, + adapter=adapter, + ) + + output_root.mkdir(parents=True, exist_ok=True) + destination = output_root / request["model_id"] + if destination.exists(): + raise FileExistsError(f"Native reference result already exists: {destination}") + temporary_path = Path(tempfile.mkdtemp(dir=output_root, prefix=".native-")) + try: + for precision, tensors in precision_tensors.items(): + save_file(tensors, temporary_path / f"{precision}.safetensors") + calibration_root = temporary_path / "calibration" + for kind, tensors in calibration_tensors.items(): + calibration_root.mkdir(parents=True, exist_ok=True) + save_file(tensors, calibration_root / f"{kind}.safetensors") + (temporary_path / "metadata.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary_path.replace(destination) + except BaseException: + shutil.rmtree(temporary_path, ignore_errors=True) + raise + if model is not None: + del model + if core is not None: + del core + del precision_tensors, calibration_tensors + gc.collect() + torch.cuda.empty_cache() + return destination + + +def _select_requests( + request_dir: Path, + model_ids: Sequence[str] | None, + *, + deep_only: bool = False, +) -> tuple[Path, ...]: + """Select explicit request files without importing the FastPLMs manifest.""" + + available = {path.stem: path for path in sorted(request_dir.glob("*.json"))} + if not available: + raise FileNotFoundError(f"No native reference requests in {request_dir}") + if model_ids is not None: + if len(set(model_ids)) != len(model_ids): + raise ValueError("Native reference model selections must be unique") + unknown = sorted(set(model_ids).difference(available)) + if unknown: + raise FileNotFoundError( + f"Native reference requests are missing selected models: {unknown}" + ) + candidates = (available[model_id] for model_id in model_ids) + else: + candidates = iter(available.values()) + + selected: list[Path] = [] + for path in candidates: + request = json.loads(path.read_text(encoding="utf-8")) + if request.get("model_id") != path.stem: + raise ValueError(f"Native reference request filename and model ID differ: {path}") + if deep_only and request.get("deep_reference") is not True: + continue + selected.append(path) + return tuple(selected) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--request-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--model", + action="append", + dest="model_ids", + help="Run only this request ID; repeat to select multiple checkpoints.", + ) + parser.add_argument( + "--deep-only", + action="store_true", + help="Run only manifest-declared deep architecture representatives.", + ) + arguments = parser.parse_args(argv) + requests = _select_requests( + arguments.request_dir, + arguments.model_ids, + deep_only=arguments.deep_only, + ) + for request in requests: + print(run_request(request, arguments.output_dir)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/parity/support/reference_adapters/__init__.py b/tests/parity/support/reference_adapters/__init__.py new file mode 100644 index 0000000..9f6d9e2 --- /dev/null +++ b/tests/parity/support/reference_adapters/__init__.py @@ -0,0 +1,165 @@ +"""Helpers for loading official implementations and immutable Hub snapshots. + +Adapters in this package are deliberately independent of :mod:`fastplms`. +Their only job is to invoke an upstream public API at the source and checkpoint +revisions declared by the model manifest, then normalize output containers for +the compliance harness. +""" + +from __future__ import annotations + +import os +import sys +import torch +import torch.nn as nn +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from types import ModuleType +from typing import Any + + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_ESM_SUBMODULE = _REPOSITORY_ROOT / "vendor" / "upstream" / "biohub-esm" + + +class OfficialGenerationUnavailable(RuntimeError): + """Normalized evidence that an official public sampler cannot execute.""" + + def __init__( + self, + *, + public_method: str, + exception_type: str, + reason: str, + ) -> None: + super().__init__(reason) + self.public_method = public_method + self.exception_type = exception_type + self.reason = reason + + def as_record(self) -> dict[str, str]: + """Return portable metadata without a traceback or host path.""" + + return { + "status": "official_unavailable", + "public_method": self.public_method, + "exception_type": self.exception_type, + "reason": self.reason, + } + + +def install_byprot_sequence_namespace(source_root: Path) -> None: + """Load only ByProt packages required by its official sequence models. + + ByProt's top-level package recursively imports every data, structure, and + task module. That discovery path requires optional compiled OpenFold + kernels even when only the public DPLM sequence classes are requested. + This namespace preserves the upstream package paths and exact registry + decorator while preventing unrelated modules from being imported. Model + classes and their public loaders still execute directly from pinned source. + """ + + package_root = source_root / "byprot" + if not package_root.is_dir(): + raise FileNotFoundError(f"Pinned ByProt source is missing: {package_root}") + existing = sys.modules.get("byprot") + marker = str(package_root.resolve()) + if existing is not None: + if getattr(existing, "_fastplms_source_root", None) != marker: + raise RuntimeError("A different ByProt package is already imported") + return + + package_paths = { + "byprot": package_root, + "byprot.models": package_root / "models", + "byprot.models.dplm": package_root / "models" / "dplm", + "byprot.models.dplm2": package_root / "models" / "dplm2", + "byprot.datamodules": package_root / "datamodules", + "byprot.datamodules.dataset": package_root / "datamodules" / "dataset", + } + for name, path in package_paths.items(): + module = ModuleType(name) + module.__file__ = str(path / "__init__.py") + module.__package__ = name + module.__dict__["__path__"] = [str(path)] + sys.modules[name] = module + parent_name, _, child_name = name.rpartition(".") + if parent_name: + setattr(sys.modules[parent_name], child_name, module) + + registry: dict[str, type[Any]] = {} + + def register_model(name: str) -> Callable[[type[Any]], type[Any]]: + def decorator(model_class: type[Any]) -> type[Any]: + registry[name] = model_class + return model_class + + return decorator + + models = sys.modules["byprot.models"] + models.__dict__["MODEL_REGISTRY"] = registry + models.__dict__["register_model"] = register_model + sys.modules["byprot"].__dict__["_fastplms_source_root"] = marker + + +def use_esm_submodule() -> None: + """Load ``esm`` from the pinned Biohub submodule instead of site-packages. + + The Biohub esm package uses the same top-level `esm` import as fair-esm. + """ + path = str(_ESM_SUBMODULE) + if not _ESM_SUBMODULE.is_dir(): + raise FileNotFoundError( + "Biohub ESM submodule is missing; run git submodule update --init --recursive" + ) + if path not in sys.path: + sys.path.insert(0, path) + + +def snapshot_path(repo_id: str, revision: str) -> Path: + """Resolve one immutable Hub snapshot or fail before loading an oracle.""" + + if not revision: + raise ValueError("A non-empty immutable reference revision is required") + from huggingface_hub import snapshot_download + + return Path(snapshot_download(repo_id=repo_id, revision=revision)).resolve() + + +def move_model( + model: nn.Module, + device: torch.device, + dtype: torch.dtype | None, +) -> nn.Module: + """Move an oracle without masking its stored dtype when dtype is omitted.""" + + if dtype is None: + return model.to(device=device) + return model.to(device=device, dtype=dtype) + + +@contextmanager +def pinned_biohub_snapshot(repo_id: str, revision: str) -> Iterator[Path]: + """Expose a pinned snapshot through Biohub's supported infra-provider path. + + Biohub's official builders resolve weights relative to the process working + directory when ``INFRA_PROVIDER`` is set. The context uses that public + deployment mode so the official implementation reads only the requested + immutable Hub revision. It does not modify an upstream class or forward + implementation. + """ + + snapshot = snapshot_path(repo_id, revision) + previous_directory = Path.cwd() + previous_provider = os.environ.get("INFRA_PROVIDER") + os.environ["INFRA_PROVIDER"] = "1" + os.chdir(snapshot) + try: + yield snapshot + finally: + os.chdir(previous_directory) + if previous_provider is None: + os.environ.pop("INFRA_PROVIDER", None) + else: + os.environ["INFRA_PROVIDER"] = previous_provider diff --git a/tests/parity/support/reference_adapters/ankh.py b/tests/parity/support/reference_adapters/ankh.py new file mode 100644 index 0000000..21c6d83 --- /dev/null +++ b/tests/parity/support/reference_adapters/ankh.py @@ -0,0 +1,102 @@ +"""Load ANKH through the pinned official package.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from pathlib import Path +from typing import Any + +from tests.parity.support.reference_adapters import move_model, snapshot_path + + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_ANKH_SUBMODULE = _REPOSITORY_ROOT / "vendor" / "upstream" / "ankh" / "src" + + +class _OfficialAnkhForwardWrapper(nn.Module): + """Normalize encoder output names; ANKH encoders intentionally have no LM head.""" + + def __init__(self, model: nn.Module, tokenizer: Any) -> None: + super().__init__() + self.model = model + self.tokenizer = tokenizer + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + **_kwargs: Any, + ) -> Any: + # input_ids: (b, l), attention_mask: (b, l) + return self.model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load an official ANKH encoder without synthesizing an output head.""" + + import sys + + if not _ANKH_SUBMODULE.is_dir(): + raise FileNotFoundError( + "ANKH submodule is missing; run git submodule update --init --recursive" + ) + source = str(_ANKH_SUBMODULE) + if source not in sys.path: + sys.path.insert(0, source) + + from ankh.models.ankh_transformers import get_specified_model + from transformers import AutoTokenizer + + snapshot = snapshot_path(reference_repo_id, reference_revision) + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + model = get_specified_model( + path=str(snapshot), + generation=False, + output_attentions=False, + framework="pt", + ) + wrapped = _OfficialAnkhForwardWrapper(model, tokenizer) + return move_model(wrapped, device, dtype).eval(), tokenizer + + +def load_official_seq2seq( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load ANKH's official sequence-to-sequence head at a pinned revision.""" + + import sys + + if not _ANKH_SUBMODULE.is_dir(): + raise FileNotFoundError( + "ANKH submodule is missing; run git submodule update --init --recursive" + ) + source = str(_ANKH_SUBMODULE) + if source not in sys.path: + sys.path.insert(0, source) + + from ankh.models.ankh_transformers import get_specified_model + from transformers import AutoTokenizer + + snapshot = snapshot_path(reference_repo_id, reference_revision) + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + model = get_specified_model( + path=str(snapshot), + generation=True, + output_attentions=False, + framework="pt", + ) + return move_model(model, device, dtype).eval(), tokenizer diff --git a/tests/parity/support/reference_adapters/biohub_loader_reproducer.py b/tests/parity/support/reference_adapters/biohub_loader_reproducer.py new file mode 100644 index 0000000..a8b29dc --- /dev/null +++ b/tests/parity/support/reference_adapters/biohub_loader_reproducer.py @@ -0,0 +1,142 @@ +"""Reproduce the pinned Biohub ESMC public-loader meta-tensor failure.""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import torch +from collections.abc import Sequence +from pathlib import Path + +from tests.parity.support.reference_adapters import ( + pinned_biohub_snapshot, + use_esm_submodule, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_sources, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo-id", required=True) + parser.add_argument("--revision", required=True) + parser.add_argument("--model-name", default="esmc_300m") + parser.add_argument( + "--construction", + choices=("public", "normal"), + default="public", + help="Use the failing public builder or the independent normal constructor probe.", + ) + return parser + + +def _load_with_normal_construction( + repo_id: str, + revision: str, + model_name: str, +) -> torch.nn.Module: + """Construct ESMC normally, then apply the pinned official loader exactly.""" + + from esm.models.esmc import ESMC + from esm.tokenization import get_esmc_model_tokenizers + from huggingface_hub import load_torch_model + from safetensors import safe_open + + configurations = { + "esmc_300m": (960, 15, 30), + "esmc_600m": (1152, 18, 36), + "esmc_6b": (2560, 40, 80), + } + try: + d_model, n_heads, n_layers = configurations[model_name] + except KeyError as error: + raise ValueError(f"Unsupported ESMC model name: {model_name!r}") from error + + with pinned_biohub_snapshot(repo_id, revision) as snapshot: + model = ESMC( + d_model=d_model, + n_heads=n_heads, + n_layers=n_layers, + tokenizer=get_esmc_model_tokenizers(), + use_flash_attn=False, + ).eval() + load_torch_model(model, snapshot) + checkpoint_path = Path(snapshot) / "model.safetensors" + state = model.state_dict() + with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint: + checkpoint_keys = set(checkpoint.keys()) + state_keys = set(state) + if checkpoint_keys != state_keys: + raise RuntimeError( + "Normal construction changed the official state-key set: " + f"missing={sorted(checkpoint_keys - state_keys)}, " + f"unexpected={sorted(state_keys - checkpoint_keys)}." + ) + for name in sorted(checkpoint_keys): + if not torch.equal(state[name], checkpoint.get_tensor(name)): + raise RuntimeError(f"Normal construction changed official tensor {name!r}.") + print( + f"Normal construction preserved all {len(state)} official state tensors exactly.", + flush=True, + ) + return model + + +def main(argv: Sequence[str] | None = None) -> int: + """Invoke only the pinned public loader and fail if parameters remain meta.""" + + arguments = _parser().parse_args(argv) + sources = reference_sources() + use_esm_submodule() + from esm.models.esmc import ESMC + + environment = { + "accelerate": importlib.metadata.version("accelerate"), + "huggingface_hub": importlib.metadata.version("huggingface-hub"), + "safetensors": importlib.metadata.version("safetensors"), + "torch": torch.__version__, + "reference_sources": sources, + } + print(json.dumps(environment, sort_keys=True), flush=True) + if arguments.construction == "normal": + model = _load_with_normal_construction( + arguments.repo_id, + arguments.revision, + arguments.model_name, + ) + else: + with pinned_biohub_snapshot(arguments.repo_id, arguments.revision): + model = ESMC.from_pretrained( + arguments.model_name, + device=torch.device("cpu"), + use_flash_attn=False, + ) + meta_parameters = [name for name, parameter in model.named_parameters() if parameter.is_meta] + if meta_parameters: + raise RuntimeError( + "Pinned Biohub public loader left meta parameters: " + ", ".join(meta_parameters[:10]) + ) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device=device) + if device.type == "cuda": + model = model.to(dtype=torch.bfloat16) + token_ids = model.tokenizer.encode("MSTNPKPQ", add_special_tokens=True) + # sequence_tokens: (1,) + sequence_tokens = torch.tensor([token_ids], device=device) + with torch.inference_mode(): + output = model(sequence_tokens=sequence_tokens) + for name in ("sequence_logits", "embeddings", "hidden_states"): + value = getattr(output, name) + if value is None or not torch.isfinite(value).all(): + raise RuntimeError(f"Official ESMC forward returned invalid {name}.") + print( + f"Biohub ESMC {arguments.construction} loader and unmodified forward passed.", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/parity/support/reference_adapters/biohub_source.py b/tests/parity/support/reference_adapters/biohub_source.py new file mode 100644 index 0000000..132deb3 --- /dev/null +++ b/tests/parity/support/reference_adapters/biohub_source.py @@ -0,0 +1,124 @@ +"""Shared runtime provenance gate for Biohub-backed official adapters.""" + +from __future__ import annotations + +import os +from pathlib import Path + +from tools.remote.biohub_reference_environment import ( + capture_biohub_reference_environment, +) +from tools.remote.reference_source_attestation import ( + validate_reference_sources_evidence, + verify_reference_source, +) + + +BIOHUB_ESM_REVISION = "82ee35553d39169d678f784c8d3f8712ffd7d2c4" +BIOHUB_ESM_TREE_SHA256 = ( + "c5489f1fc58de200978803de2c38e1a78f769cb183a2ee90be833f0f4a0212e8" +) +BIOHUB_TRANSFORMERS_REVISION = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf" +BIOHUB_TRANSFORMERS_TREE_SHA256 = ( + "28b910cc18b821870db2fb6d1c50376c2d14287ae18485080699e03fa4ba4f43" +) +BIOHUB_REFERENCE_SOURCE_NAMES = ("biohub-esm", "biohub-transformers") + + +def _reference_source( + *, + name: str, + environment_prefix: str, + expected_revision: str, + expected_tree_sha256: str, +) -> dict[str, object]: + """Verify one exact Biohub source tree and its runtime import origin.""" + + configured_revision = os.environ.get(f"{environment_prefix}_REVISION") + if configured_revision != expected_revision: + raise RuntimeError( + f"Biohub reference container does not declare the pinned {name} revision: " + f"expected {expected_revision}, received {configured_revision!r}." + ) + required_environment = ( + f"{environment_prefix}_SOURCE", + f"{environment_prefix}_ATTESTATION", + f"{environment_prefix}_CONTRACT", + ) + try: + source_root, attestation, contract = ( + Path(os.environ[name]) for name in required_environment + ) + except KeyError as error: + raise RuntimeError( + "Biohub reference source attestation environment is incomplete." + ) from error + evidence = verify_reference_source( + source_root, + attestation, + contract, + expected_revision=expected_revision, + ) + if evidence["tree_sha256"] != expected_tree_sha256: + raise RuntimeError(f"{name} source digest differs from the adapter pin.") + return evidence + + +def reference_sources() -> dict[str, dict[str, object]]: + """Verify both source trees that together define the Biohub oracle.""" + + evidence = { + "biohub-esm": _reference_source( + name="Biohub ESM", + environment_prefix="FASTPLMS_BIOHUB_ESM", + expected_revision=BIOHUB_ESM_REVISION, + expected_tree_sha256=BIOHUB_ESM_TREE_SHA256, + ), + "biohub-transformers": _reference_source( + name="Biohub Transformers", + environment_prefix="FASTPLMS_BIOHUB_TRANSFORMERS", + expected_revision=BIOHUB_TRANSFORMERS_REVISION, + expected_tree_sha256=BIOHUB_TRANSFORMERS_TREE_SHA256, + ), + } + return validate_reference_sources_evidence( + evidence, + required_sources=BIOHUB_REFERENCE_SOURCE_NAMES, + ) + + +def reference_environment() -> dict[str, object]: + """Verify the exact GH200 lock, installed inventory, and image identities.""" + + required_environment = ( + "FASTPLMS_BIOHUB_LOCK_ROOT", + "FASTPLMS_BIOHUB_LOCK_CONTRACT", + "FASTPLMS_REFERENCE_CONTAINER_IDENTITIES", + "FASTPLMS_REFERENCE_CONTAINER_TARGET", + ) + try: + lock_root, contract, container_identities = ( + Path(os.environ[name]) for name in required_environment[:3] + ) + reference_target = os.environ[required_environment[3]] + except KeyError as error: + raise RuntimeError( + "Biohub reference environment identity is incomplete." + ) from error + return capture_biohub_reference_environment( + lock_root, + contract, + container_identities, + reference_target=reference_target, + ) + + +__all__ = [ + "BIOHUB_ESM_REVISION", + "BIOHUB_ESM_TREE_SHA256", + "BIOHUB_REFERENCE_SOURCE_NAMES", + "BIOHUB_TRANSFORMERS_REVISION", + "BIOHUB_TRANSFORMERS_TREE_SHA256", + "reference_environment", + "reference_sources", +] diff --git a/tests/parity/support/reference_adapters/boltz.py b/tests/parity/support/reference_adapters/boltz.py new file mode 100644 index 0000000..a27641a --- /dev/null +++ b/tests/parity/support/reference_adapters/boltz.py @@ -0,0 +1,35 @@ +"""Load Boltz2 through the pinned official Lightning checkpoint API.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from pathlib import Path + +from tests.parity.support.reference_adapters import move_model + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, None]: + """Load the official Boltz2 class from an immutable Hub revision.""" + + from boltz.model.models.boltz2 import Boltz2 + from huggingface_hub import hf_hub_download + + checkpoint = Path( + hf_hub_download( + repo_id=reference_repo_id, + filename="boltz2_conf.ckpt", + revision=reference_revision, + ) + ) + model = Boltz2.load_from_checkpoint( + checkpoint, + strict=True, + map_location="cpu", + ) + return move_model(model, device, dtype).eval(), None diff --git a/tests/parity/support/reference_adapters/dplm.py b/tests/parity/support/reference_adapters/dplm.py new file mode 100644 index 0000000..be6bc5d --- /dev/null +++ b/tests/parity/support/reference_adapters/dplm.py @@ -0,0 +1,115 @@ +"""Official DPLM parity adapter backed by the pinned Bytedance source tree. + +The adapter invokes ``DiffusionProteinLanguageModel.forward`` and only +normalizes the returned container. Forward hooks observe intermediate tensors; +they do not replace or reconstruct any upstream computation. +""" + +from __future__ import annotations + +import sys +import torch +import torch.nn as nn +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from tests.parity.support.reference_adapters import ( + install_byprot_sequence_namespace, + move_model, + snapshot_path, +) + + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_DPLM_SOURCE = _REPOSITORY_ROOT / "vendor" / "upstream" / "dplm" / "src" + + +def _install_source_path() -> None: + if not _DPLM_SOURCE.is_dir(): + raise FileNotFoundError( + "DPLM submodule is missing; run git submodule update --init --recursive" + ) + source = str(_DPLM_SOURCE) + if source not in sys.path: + sys.path.insert(0, source) + install_byprot_sequence_namespace(_DPLM_SOURCE) + + +class _OfficialDPLMForwardWrapper(nn.Module): + """Expose Hugging Face-style output names around the official public API.""" + + def __init__(self, oracle: nn.Module) -> None: + super().__init__() + self.oracle = oracle + # The checkpoint conversion targets the official network state exactly. + self.model = oracle.net + self.tokenizer = oracle.tokenizer + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **_kwargs: Any, + ) -> SimpleNamespace: + # input_ids: (b, l) + del attention_mask + captured: list[torch.Tensor] = [] + + def capture(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> None: + value = output[0] if isinstance(output, tuple) else output + if torch.is_tensor(value): + captured.append(value) + + handles = [self.model.esm.embeddings.register_forward_hook(capture)] + handles.extend( + layer.register_forward_hook(capture) + for layer in self.model.esm.encoder.layer[:-1] + ) + try: + # This is the upstream model's public inference entry point. + logits, last_hidden_state = self.oracle( + input_ids=input_ids, + return_last_hidden_state=True, + ) + finally: + for handle in handles: + handle.remove() + + hidden_states = tuple(captured) + if not hidden_states or hidden_states[-1] is not last_hidden_state: + hidden_states = (*hidden_states, last_hidden_state) + return SimpleNamespace( + logits=logits, + last_hidden_state=last_hidden_state, + hidden_states=hidden_states, + ) + + def generate(self, input_tokens: torch.Tensor, **kwargs: Any) -> torch.Tensor: + """Invoke the pinned implementation's public diffusion sampler.""" + + # input_tokens: (...) + return self.oracle.generate(input_tokens=input_tokens, **kwargs) + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load DPLM through its pinned official ``from_pretrained`` method.""" + + _install_source_path() + # Register the exact official network class without triggering ByProt's + # package-wide discovery of unrelated structure models. + from byprot.models.dplm.dplm import DiffusionProteinLanguageModel + from byprot.models.dplm.modules import dplm_modeling_esm as _dplm_modeling_esm + + del _dplm_modeling_esm + + snapshot = snapshot_path(reference_repo_id, reference_revision) + oracle = DiffusionProteinLanguageModel.from_pretrained(str(snapshot)) + oracle = move_model(oracle, device, dtype).eval() + wrapped = move_model(_OfficialDPLMForwardWrapper(oracle), device, dtype).eval() + return wrapped, wrapped.tokenizer diff --git a/tests/parity/support/reference_adapters/dplm2.py b/tests/parity/support/reference_adapters/dplm2.py new file mode 100644 index 0000000..48836ec --- /dev/null +++ b/tests/parity/support/reference_adapters/dplm2.py @@ -0,0 +1,317 @@ +"""Official DPLM2 parity adapter backed by the pinned Bytedance source tree. + +No FastPLMs loader, token remapping, or reconstructed forward pass is used. +The official tokenizer and multimodal model receive the inputs directly. +""" + +from __future__ import annotations + +import inspect +import sys +import torch +import torch.nn as nn +from collections.abc import Mapping, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Protocol, cast + +from tests.parity.support.reference_adapters import ( + OfficialGenerationUnavailable, + install_byprot_sequence_namespace, + move_model, + snapshot_path, +) + + +DPLM2_3B_GENERATION_LIMITATION = { + "status": "official_unavailable", + "public_method": "EsmForDPLM.generate", + "exception_type": "TypeError", + "reason": ( + "The checkpoint-selected EsmForDPLM sampler uses tokenizer.cls_token_id " + "as bos_id, but the pinned DPLM2 tokenizer defines no cls_token_id." + ), +} + +# Exact evidence from the pinned 150M official checkpoint. These trained head +# tensors are part of the release state contract and must never be replaced by +# random initialization when a mirror or local artifact is loaded. +DPLM2_150M_OFFICIAL_HEAD_CONTRACT = { + "esm.contact_head.regression.bias": { + "dtype": "torch.float32", + "sha256": "bbce9db798883b08550850be32ea0a60cde4e06adb02e0a0ac686469a419311e", + "shape": [1], + }, + "esm.contact_head.regression.weight": { + "dtype": "torch.float32", + "sha256": "8037b6e221939baa4fdf62b3a89d9fd4a3b2430494daa85c58bb760e0514a6fc", + "shape": [1, 600], + }, + "esm.embeddings.word_embeddings.weight": { + "dtype": "torch.float32", + "sha256": "58662f66967b04570801ca4bd4c49c4bc610df0feadbe92128fe416a2fa23325", + "shape": [8229, 640], + }, + "lm_head.bias": { + "dtype": "torch.float32", + "sha256": "4ee4c69d1b4d6beea9c28a70d6440e5faff338b2e1cf927867b4e653cfa2f0f6", + "shape": [8229], + }, + "lm_head.decoder.weight": { + "dtype": "torch.float32", + "sha256": "25f3f82396eca43ba601bb7062458750a9b08177e7edc4ae24ae4d424ce2aea2", + "shape": [8229, 640], + }, + "lm_head.dense.bias": { + "dtype": "torch.float32", + "sha256": "e880655edb59ea7774c250b81870e9cefa85738258185c23d3b5846004a68daf", + "shape": [640], + }, + "lm_head.dense.weight": { + "dtype": "torch.float32", + "sha256": "f0dfa1b0a2d85e4cdc2e13727d21da29e94f09fb2267fdc3d3cfc5f2d0fd3450", + "shape": [640, 640], + }, + "lm_head.layer_norm.bias": { + "dtype": "torch.float32", + "sha256": "7c82d73e9be1b191ec39ffdfd5aa10b8da854d575daf1c334cdee261eb6406d3", + "shape": [640], + }, + "lm_head.layer_norm.weight": { + "dtype": "torch.float32", + "sha256": "edf9e1e32bec22be5eb81e2819e34324742eacf6f8cd5f410b1297705ab019b9", + "shape": [640], + }, +} + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_DPLM_SOURCE = _REPOSITORY_ROOT / "vendor" / "upstream" / "dplm" / "src" + + +class _DPLM2Encoder(Protocol): + """Static shape of the encoder fields used by the parity adapter.""" + + layer: Sequence[nn.Module] + + +class _DPLM2Esm(Protocol): + """Static shape of the ESM fields used by the parity adapter.""" + + embeddings: nn.Module + encoder: _DPLM2Encoder + + +class _DPLM2Generative(Protocol): + """Static shape of the public sampler used by the parity adapter.""" + + def generate(self, **kwargs: Any) -> Any: ... + + +class _DPLM2ModelWithEsm(Protocol): + """Static shape of the checkpoint-selected model wrapper.""" + + esm: _DPLM2Esm + + +def _install_source_path() -> None: + if not _DPLM_SOURCE.is_dir(): + raise FileNotFoundError( + "DPLM submodule is missing; run git submodule update --init --recursive" + ) + source = str(_DPLM_SOURCE) + if source not in sys.path: + sys.path.insert(0, source) + install_byprot_sequence_namespace(_DPLM_SOURCE) + + +def _field(output: Any, name: str) -> Any: + if isinstance(output, Mapping): + return output.get(name) + return getattr(output, name, None) + + +def _accepts_type_ids(module: nn.Module) -> bool: + """Return whether the pinned network's public forward accepts ``type_ids``.""" + + parameters = inspect.signature(module.forward).parameters.values() + return any( + parameter.name == "type_ids" + or parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters + ) + + +def _call_checkpoint_forward( + oracle: nn.Module, + network: nn.Module, + input_ids: torch.Tensor, + kwargs: Mapping[str, Any], +) -> Any: + """Call the supported public forward selected by the official checkpoint. + + The pinned DPLM2-3B checkpoint selects the official ``dplm_esm`` network. + Its public forward does not accept ``type_ids``, while the official + multimodal wrapper passes that keyword unconditionally. In that one case, + the adapter invokes the checkpoint-selected public network forward with its + supported signature. No hook is registered and no upstream object or class + is modified. + """ + + # input_ids: (b, l) + target = oracle if _accepts_type_ids(network) else network + return target(input_ids=input_ids, **kwargs) + + +def _call_checkpoint_generate( + oracle: nn.Module, + network: nn.Module, + input_tokens: torch.Tensor, + kwargs: Mapping[str, Any], +) -> Any: + """Call the checkpoint-selected implementation's public sampler. + + The DPLM2-3B checkpoint selects ``dplm_esm``. Its public sampler accepts a + batch mapping, while the multimodal wrapper's sampler re-enters the broken + ``type_ids`` forward path. Use the selected network's sampler directly for + that checkpoint, passing only keywords declared by its public signature. + Other DPLM2 checkpoints retain the multimodal wrapper's public sampler. + """ + + # input_tokens: (...) + oracle_generate = cast(_DPLM2Generative, oracle).generate + generate = getattr(network, "generate", None) + if _accepts_type_ids(network) or not callable(generate): + return oracle_generate(input_tokens=input_tokens, **kwargs) + + parameters = inspect.signature(generate).parameters + supported_kwargs = { + name: value for name, value in kwargs.items() if name in parameters + } + try: + generated = generate( + batch={"input_ids": input_tokens}, + **supported_kwargs, + ) + except TypeError as error: + is_pinned_public_failure = ( + type(network).__name__ == "EsmForDPLM" + and getattr(network, "bos_id", object()) is None + and "NoneType" in str(error) + and "ne()" in str(error) + ) + if not is_pinned_public_failure: + raise + raise OfficialGenerationUnavailable( + public_method=DPLM2_3B_GENERATION_LIMITATION["public_method"], + exception_type=type(error).__name__, + reason=DPLM2_3B_GENERATION_LIMITATION["reason"], + ) from error + if isinstance(generated, tuple): + return generated[0] + return generated + + +class _OfficialDPLM2ForwardWrapper(nn.Module): + """Normalize names while retaining the official public forward computation.""" + + def __init__(self, oracle: nn.Module) -> None: + super().__init__() + self.oracle = oracle + self.model = cast(nn.Module, oracle.net) + self.tokenizer = oracle.tokenizer + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Any, + ) -> SimpleNamespace: + # input_ids: (b, l) + del attention_mask + captured: list[torch.Tensor] = [] + + def capture(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> None: + value = output[0] if isinstance(output, tuple) else output + if torch.is_tensor(value): + captured.append(value) + + esm = cast(_DPLM2ModelWithEsm, self.model).esm + handles = [esm.embeddings.register_forward_hook(capture)] + handles.extend( + layer.register_forward_hook(capture) + for layer in esm.encoder.layer[:-1] + ) + try: + output = _call_checkpoint_forward( + self.oracle, + self.model, + input_ids, + kwargs, + ) + finally: + for handle in handles: + handle.remove() + + logits = _field(output, "logits") + last_hidden_state = _field(output, "last_hidden_state") + if logits is None or last_hidden_state is None: + raise RuntimeError("Official DPLM2 output omitted logits or last_hidden_state") + # The official multimodal wrapper calls the embedding block once, then + # its inner ESM model calls the same block again. The second result is + # the encoder's semantic input hidden state. + # The multimodal wrapper computes embeddings before calling its inner + # network, which computes them a second time. The pinned 3B checkpoint + # selects the inner network directly, so that path has no duplicate. + hidden_states = tuple(captured[1:] if _accepts_type_ids(self.model) else captured) + if not hidden_states or hidden_states[-1] is not last_hidden_state: + hidden_states = (*hidden_states, last_hidden_state) + return SimpleNamespace( + logits=logits, + last_hidden_state=last_hidden_state, + hidden_states=hidden_states, + ) + + def generate(self, input_tokens: torch.Tensor, **kwargs: Any) -> Any: + """Invoke the checkpoint-selected implementation's public sampler.""" + + # input_tokens: (...) + return _call_checkpoint_generate( + self.oracle, + self.model, + input_tokens, + kwargs, + ) + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load DPLM2 through its pinned official ``from_pretrained`` method.""" + + _install_source_path() + # The 3B DPLM2 checkpoint declares the upstream ``dplm_esm`` network + # architecture, so its official registry module must be imported before + # the public loader resolves the class name. + from byprot.datamodules.dataset.tokenized_protein import DPLM2Tokenizer + from byprot.models.dplm.modules import dplm_modeling_esm as _dplm_modeling_esm + from byprot.models.dplm2.dplm2 import MultimodalDiffusionProteinLanguageModel + from transformers import AutoTokenizer, EsmConfig + + del _dplm_modeling_esm + # This checkpoint stores the official class name in tokenizer_config.json. + # Registering that unmodified upstream class restores the public + # AutoTokenizer lookup expected by the pinned DPLM loader. + AutoTokenizer.register( # type: ignore[no-untyped-call] + EsmConfig, + slow_tokenizer_class=DPLM2Tokenizer, + exist_ok=True, + ) + + snapshot = snapshot_path(reference_repo_id, reference_revision) + oracle = MultimodalDiffusionProteinLanguageModel.from_pretrained(str(snapshot)) + oracle = move_model(oracle, device, dtype).eval() + wrapped = move_model(_OfficialDPLM2ForwardWrapper(oracle), device, dtype).eval() + return wrapped, wrapped.tokenizer diff --git a/testing/official/e1.py b/tests/parity/support/reference_adapters/e1.py similarity index 60% rename from testing/official/e1.py rename to tests/parity/support/reference_adapters/e1.py index 93157f2..d8c006a 100644 --- a/testing/official/e1.py +++ b/tests/parity/support/reference_adapters/e1.py @@ -1,7 +1,9 @@ """Load official E1 model from the e1 package for comparison.""" + import torch import torch.nn as nn -from typing import Tuple + +from tests.parity.support.reference_adapters import move_model class _OfficialE1ForwardWrapper(nn.Module): @@ -15,9 +17,12 @@ def forward( within_seq_position_ids: torch.LongTensor, global_position_ids: torch.LongTensor, sequence_ids: torch.LongTensor, - attention_mask: torch.LongTensor, + attention_mask: torch.LongTensor | None = None, **kwargs, ): + # input_ids: (b, l); within_seq_position_ids: (b, l) + # global_position_ids: (b, l); sequence_ids: (b, l) + del attention_mask, kwargs batch = { "input_ids": input_ids, "within_seq_position_ids": within_seq_position_ids, @@ -30,9 +35,10 @@ def forward( def load_official_model( reference_repo_id: str, + reference_revision: str, device: torch.device, - dtype: torch.dtype = torch.float32, -) -> Tuple[nn.Module, object]: + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: """Load the official E1 model from the e1 submodule. Args: @@ -43,20 +49,19 @@ def load_official_model( Returns (official_model, batch_preparer) where batch_preparer is an E1BatchPreparer. The official model is E1ForMaskedLM with standard HF forward interface. """ - from E1.modeling import E1ForMaskedLM from E1.batch_preparer import E1BatchPreparer + from E1.modeling import E1ForMaskedLM - model = E1ForMaskedLM.from_pretrained( - reference_repo_id, - tie_word_embeddings=False, - device_map=device, - dtype=dtype, - ).eval() + load_kwargs = { + "revision": reference_revision, + "tie_word_embeddings": False, + } + if dtype is not None: + load_kwargs["dtype"] = dtype + # Load through the official public API on CPU, then transfer once below. + # A Transformers ``device_map`` requires Accelerate even for one device and + # adds no semantic value for these reference checkpoints. + model = E1ForMaskedLM.from_pretrained(reference_repo_id, **load_kwargs).eval() batch_preparer = E1BatchPreparer() - wrapped = _OfficialE1ForwardWrapper(model).eval() + wrapped = move_model(_OfficialE1ForwardWrapper(model), device, dtype).eval() return wrapped, batch_preparer - - -if __name__ == "__main__": - model, batch_preparer = load_official_model("Profluent-Bio/E1-150m", torch.device("cpu")) - print(model) \ No newline at end of file diff --git a/tests/parity/support/reference_adapters/esm2.py b/tests/parity/support/reference_adapters/esm2.py new file mode 100644 index 0000000..5d45d49 --- /dev/null +++ b/tests/parity/support/reference_adapters/esm2.py @@ -0,0 +1,190 @@ +"""Load ESM2 through the pinned Meta ESM implementation.""" + +from __future__ import annotations + +import hashlib +import os +import sys +import tempfile +import urllib.request +import torch +import torch.nn as nn +from collections.abc import Mapping, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from tests.parity.support.reference_adapters import move_model, snapshot_path + + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_FAIR_ESM_SUBMODULE = _REPOSITORY_ROOT / "vendor" / "upstream" / "fair-esm" + + +def _asset_field(asset: object, name: str) -> Any: + if isinstance(asset, Mapping): + return asset[name] + return getattr(asset, name) + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _verified_asset(asset: object) -> Path: + relative = Path(str(_asset_field(asset, "path"))) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"Unsafe fair-esm oracle asset path: {relative}") + torch_home = Path(os.environ.get("TORCH_HOME", "~/.cache/torch")).expanduser() + cache_root = Path(os.environ.get("FASTPLMS_ORACLE_CACHE", str(torch_home / "fair-esm"))) + destination = cache_root / relative + expected_size = int(_asset_field(asset, "size")) + expected_sha256 = str(_asset_field(asset, "sha256")) + + def valid() -> bool: + return ( + destination.is_file() + and destination.stat().st_size == expected_size + and _file_sha256(destination) == expected_sha256 + ) + + if valid(): + return destination + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile(dir=destination.parent, delete=False) as temporary: + temporary_path = Path(temporary.name) + try: + urllib.request.urlretrieve(str(_asset_field(asset, "url")), temporary_path) + if temporary_path.stat().st_size != expected_size: + raise RuntimeError(f"Size mismatch for fair-esm oracle asset {relative}") + if _file_sha256(temporary_path) != expected_sha256: + raise RuntimeError(f"SHA-256 mismatch for fair-esm oracle asset {relative}") + temporary_path.replace(destination) + finally: + temporary_path.unlink(missing_ok=True) + return destination + + +class _AlphabetTokenizer: + """Expose the official Alphabet through the subset used by parity tests.""" + + def __init__(self, alphabet: Any) -> None: + self.alphabet = alphabet + self.pad_token_id = alphabet.padding_idx + self.cls_token_id = alphabet.cls_idx + self.bos_token_id = alphabet.cls_idx + self.eos_token_id = alphabet.eos_idx + self.mask_token_id = alphabet.mask_idx + self.unk_token_id = alphabet.unk_idx + self.all_special_ids = [ + self.pad_token_id, + self.cls_token_id, + self.eos_token_id, + self.mask_token_id, + self.unk_token_id, + ] + + def get_vocab(self) -> dict[str, int]: + return dict(self.alphabet.tok_to_idx) + + def __call__( + self, + sequences: str | Sequence[str], + *, + return_tensors: str = "pt", + padding: bool | str = True, + truncation: bool = False, + max_length: int | None = None, + **_kwargs: Any, + ) -> dict[str, torch.Tensor]: + if return_tensors != "pt": + raise ValueError("The official ESM2 parity tokenizer returns PyTorch tensors") + values = [sequences] if isinstance(sequences, str) else list(sequences) + if truncation and max_length is not None: + residue_limit = max(1, max_length - 2) + values = [sequence[:residue_limit] for sequence in values] + converter = self.alphabet.get_batch_converter() + _, _, input_ids = converter( + [(str(index), sequence) for index, sequence in enumerate(values)] + ) + if padding == "max_length" and max_length is not None and input_ids.shape[1] < max_length: + # pad: (input_ids.shape[0], max_length - input_ids.shape[1]) + pad = torch.full( + (input_ids.shape[0], max_length - input_ids.shape[1]), + self.pad_token_id, + dtype=input_ids.dtype, + ) + # input_ids: (...) + input_ids = torch.cat((input_ids, pad), dim=1) + # attention_mask: (b, l) + attention_mask = input_ids.ne(self.pad_token_id).long() + return {"input_ids": input_ids, "attention_mask": attention_mask} + + +class _OfficialESM2ForwardWrapper(nn.Module): + def __init__(self, model: nn.Module, alphabet: Any) -> None: + super().__init__() + self.model = model + self.tokenizer = _AlphabetTokenizer(alphabet) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **_kwargs: Any, + ) -> Any: + # input_ids: (b, l) + del attention_mask + layers = list(range(self.model.num_layers + 1)) + output = self.model(input_ids, repr_layers=layers, return_contacts=False) + hidden_states = tuple(output["representations"][index] for index in layers) + return SimpleNamespace( + logits=output["logits"], + last_hidden_state=hidden_states[-1], + hidden_states=hidden_states, + ) + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, + oracle_assets: Sequence[object] = (), +) -> tuple[nn.Module, _AlphabetTokenizer]: + """Load Meta's exact fair-esm code and hash-pinned native checkpoint.""" + + if not _FAIR_ESM_SUBMODULE.is_dir(): + raise FileNotFoundError( + "Meta ESM submodule is missing; run git submodule update --init --recursive" + ) + source = str(_FAIR_ESM_SUBMODULE) + if source not in sys.path: + sys.path.insert(0, source) + import esm + + # Resolve the declared immutable Hub snapshot as a provenance gate even + # though fair-esm's native oracle consumes Meta's hash-pinned `.pt` files. + snapshot_path(reference_repo_id, reference_revision) + by_role = {str(_asset_field(asset, "role")): asset for asset in oracle_assets} + if set(by_role) != {"weights", "contact_regression"}: + raise RuntimeError( + "ESM2 live parity requires hash-pinned weights and contact-regression assets" + ) + weights_path = _verified_asset(by_role["weights"]) + regression_path = _verified_asset(by_role["contact_regression"]) + model_name = reference_repo_id.rsplit("/", 1)[-1] + model_data = torch.load(weights_path, map_location="cpu", weights_only=False) + regression_data = torch.load(regression_path, map_location="cpu", weights_only=False) + model, alphabet = esm.pretrained.load_model_and_alphabet_core( + model_name, + model_data, + regression_data, + ) + wrapped = _OfficialESM2ForwardWrapper(model, alphabet) + wrapped = move_model(wrapped, device, dtype).eval() + return wrapped, wrapped.tokenizer diff --git a/tests/parity/support/reference_adapters/esm3.py b/tests/parity/support/reference_adapters/esm3.py new file mode 100644 index 0000000..99de916 --- /dev/null +++ b/tests/parity/support/reference_adapters/esm3.py @@ -0,0 +1,140 @@ +"""Load ESM3 through the pinned Biohub implementation's public forward API.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from typing import Any + +from tests.parity.support.reference_adapters import ( + move_model, + pinned_biohub_snapshot, + use_esm_submodule, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_environment as _reference_environment, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_sources, +) + + +reference_environment = _reference_environment + +use_esm_submodule() + + +class _ESM3ComplianceOutput: + """Normalize official output names without changing official computation.""" + + def __init__(self, output: Any, hidden_states: tuple[torch.Tensor, ...]) -> None: + self.logits = output.sequence_logits + self.last_hidden_state = output.embeddings + self.hidden_states = hidden_states + self.sequence_logits = output.sequence_logits + self.structure_logits = output.structure_logits + self.function_logits = output.function_logits + self.residue_logits = output.residue_logits + + +class _ESM3StateDictRoot(nn.Module): + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.esm3 = model + + +class _OfficialESM3ForwardWrapper(nn.Module): + """Adapt Hugging Face-style names to the official ESM3 keyword names.""" + + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.model = _ESM3StateDictRoot(model) + self.tokenizer = model.tokenizers.sequence + + @property + def esm3(self) -> nn.Module: + return self.model.esm3 + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_tokens: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs: Any, + ) -> _ESM3ComplianceOutput: + if sequence_tokens is None: + sequence_tokens = input_ids + if sequence_id is None and attention_mask is not None: + # sequence_id: (b, l) + sequence_id = attention_mask.to(dtype=torch.bool) + + captured: dict[str, tuple[torch.Tensor, ...]] = {} + + def capture_transformer_output( + _module: nn.Module, + _inputs: tuple[Any, ...], + output: tuple[Any, ...], + ) -> None: + captured["hidden_states"] = tuple(output[2]) + + handle = self.esm3.transformer.register_forward_hook(capture_transformer_output) + try: + output = self.esm3( + sequence_tokens=sequence_tokens, + sequence_id=sequence_id, + **{ + key: value + for key, value in kwargs.items() + if key + in { + "structure_tokens", + "ss8_tokens", + "sasa_tokens", + "function_tokens", + "residue_annotation_tokens", + "average_plddt", + "per_res_plddt", + "structure_coords", + "chain_id", + "output_attentions", + } + }, + ) + finally: + handle.remove() + + hidden_states = captured.get("hidden_states") + if output_hidden_states and hidden_states is None: + raise RuntimeError("Official ESM3 transformer did not expose hidden states") + return _ESM3ComplianceOutput(output, hidden_states or ()) + + +def _normalize_reference_repo_id(reference_repo_id: str) -> str: + aliases = { + "biohub/esm3-sm-open-v1": "esm3-sm-open-v1", + "EvolutionaryScale/esm3-sm-open-v1": "esm3-sm-open-v1", + } + return aliases.get(reference_repo_id, reference_repo_id) + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load official weights, then expose only name-normalized outputs.""" + + reference_sources() + from esm.pretrained import load_local_model + from esm.utils.constants.models import normalize_model_name + + model_name = normalize_model_name(_normalize_reference_repo_id(reference_repo_id)) + # Load on CPU to avoid an implicit BF16 conversion before FP32 parity runs. + with pinned_biohub_snapshot(reference_repo_id, reference_revision): + model = load_local_model(model_name, device=torch.device("cpu")) + model = move_model(model, device, dtype).eval() + wrapped = move_model(_OfficialESM3ForwardWrapper(model), device, dtype).eval() + return wrapped, wrapped.tokenizer diff --git a/tests/parity/support/reference_adapters/esm_plusplus.py b/tests/parity/support/reference_adapters/esm_plusplus.py new file mode 100644 index 0000000..f9d8f2e --- /dev/null +++ b/tests/parity/support/reference_adapters/esm_plusplus.py @@ -0,0 +1,46 @@ +"""Load ESMC through the pinned Biohub Transformers public API.""" + +import torch +import torch.nn as nn + +from tests.parity.support.reference_adapters import ( + move_model, + snapshot_path, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_environment as _reference_environment, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_sources, +) + + +reference_environment = _reference_environment + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, object]: + """Load the official ESMC model from the pinned Biohub Transformers fork. + + Args: + reference_repo_id: e.g. "biohub/ESMC-300M" + device: target device + dtype: target dtype (should be float32 for comparison) + + Returns (wrapped_model, tokenizer). + """ + reference_sources() + from transformers import AutoTokenizer + from transformers.models.esmc.modeling_esmc import ESMCForMaskedLM + + snapshot = snapshot_path(reference_repo_id, reference_revision) + load_kwargs: dict[str, object] = {"local_files_only": True} + if dtype is not None: + load_kwargs["dtype"] = dtype + official_model = ESMCForMaskedLM.from_pretrained(snapshot, **load_kwargs) + tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) + return move_model(official_model, device, dtype).eval(), tokenizer diff --git a/tests/parity/support/reference_adapters/esmfold.py b/tests/parity/support/reference_adapters/esmfold.py new file mode 100644 index 0000000..fb466b3 --- /dev/null +++ b/tests/parity/support/reference_adapters/esmfold.py @@ -0,0 +1,56 @@ +"""Load ESMFold through Meta's public API and a hash-pinned native asset.""" + +from __future__ import annotations + +import sys +import torch +import torch.nn as nn +from collections.abc import Sequence +from pathlib import Path + +from tests.parity.support.reference_adapters import move_model, snapshot_path +from tests.parity.support.reference_adapters.esm2 import _asset_field, _verified_asset + + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_FAIR_ESM_SUBMODULE = _REPOSITORY_ROOT / "vendor" / "upstream" / "fair-esm" + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, + oracle_assets: Sequence[object] = (), +) -> tuple[nn.Module, None]: + """Load Meta ESMFold v1 without allowing its mutable download path.""" + + if not _FAIR_ESM_SUBMODULE.is_dir(): + raise FileNotFoundError( + "Meta ESM submodule is missing; run git submodule update --init --recursive" + ) + source = str(_FAIR_ESM_SUBMODULE) + if source not in sys.path: + sys.path.insert(0, source) + + # The Hub revision records the immutable converted packaging checkpoint. + # The live fair-esm oracle consumes Meta's independently hash-pinned `.pt`. + snapshot_path(reference_repo_id, reference_revision) + by_role = {str(_asset_field(asset, "role")): asset for asset in oracle_assets} + if set(by_role) != {"weights"}: + raise RuntimeError("ESMFold live parity requires its hash-pinned native weights") + weights_path = _verified_asset(by_role["weights"]) + + # `esm.pretrained.esmfold_v1()` is the public official constructor. Seed + # its standard Torch Hub cache with the already verified file so the + # constructor cannot resolve the mutable URL over the network. + checkpoint = Path(torch.hub.get_dir()) / "checkpoints" / "esmfold_3B_v1.pt" + checkpoint.parent.mkdir(parents=True, exist_ok=True) + if checkpoint.exists() or checkpoint.is_symlink(): + checkpoint.unlink() + checkpoint.symlink_to(weights_path) + + import esm + + model = esm.pretrained.esmfold_v1() + return move_model(model, device, dtype).eval(), None diff --git a/tests/parity/support/reference_adapters/esmfold2.py b/tests/parity/support/reference_adapters/esmfold2.py new file mode 100644 index 0000000..f1977a1 --- /dev/null +++ b/tests/parity/support/reference_adapters/esmfold2.py @@ -0,0 +1,91 @@ +"""Load ESMFold2 through the pinned Biohub Transformers implementation.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from tests.parity.support.reference_adapters import move_model, snapshot_path +from tests.parity.support.reference_adapters.biohub_source import ( + reference_environment as _reference_environment, +) +from tests.parity.support.reference_adapters.biohub_source import ( + reference_sources, +) + + +reference_environment = _reference_environment + + +class _OfficialESMFold2Wrapper(nn.Module): + """Expose projection-only output while delegating folding to Biohub.""" + + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.model = model + + def forward(self, *args, **kwargs): + return self.model(*args, **kwargs) + + def project_esmc_hidden_states( + self, + hidden_states: torch.Tensor, + residue_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return Biohub's learned sequence summary before pair expansion.""" + + # hidden_states: (..., d) + shim = self.model.language_model + captured: list[torch.Tensor] = [] + + def capture_pair_input(_module: nn.Module, args: tuple[torch.Tensor, ...]) -> None: + if len(args) != 1: + raise RuntimeError("Biohub base_z_mlp received an unexpected input signature.") + captured.append(args[0]) + + # Observe the official shim's public forward path at the boundary before + # pair expansion. This avoids reproducing Biohub's learned projection. + handle = shim.base_z_mlp.register_forward_pre_hook(capture_pair_input) + try: + shim(hidden_states, lm_dropout=0.0) + finally: + handle.remove() + if len(captured) != 1: + raise RuntimeError("Biohub LanguageModelShim did not expose one sequence summary.") + projected = captured[0] + if residue_mask is not None: + projected = projected * residue_mask.to( + device=projected.device, + dtype=projected.dtype, + ).unsqueeze(-1) + return projected + + +def load_official_model( + reference_repo_id: str, + reference_revision: str, + device: torch.device, + dtype: torch.dtype | None = None, +) -> tuple[nn.Module, None]: + """Load one of the four supported ESMFold2 snapshots exactly.""" + + reference_sources() + from transformers.models.esmfold2.configuration_esmfold2 import ESMFold2Config + from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model + from transformers.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, + ) + + snapshot = snapshot_path(reference_repo_id, reference_revision) + config = ESMFold2Config.from_pretrained(snapshot, local_files_only=True) + model_class = ESMFold2ExperimentalModel if config.type == "experimental" else ESMFold2Model + load_kwargs = { + "config": config, + "local_files_only": True, + "load_esmc": False, + } + if dtype is not None: + load_kwargs["torch_dtype"] = dtype + model = model_class.from_pretrained(snapshot, **load_kwargs) + wrapped = _OfficialESMFold2Wrapper(model) + return move_model(wrapped, device, dtype).eval(), None diff --git a/tests/parity/support/semantic_config.py b/tests/parity/support/semantic_config.py new file mode 100644 index 0000000..56133a6 --- /dev/null +++ b/tests/parity/support/semantic_config.py @@ -0,0 +1,120 @@ +"""Shared semantic configuration extraction for native and candidate parity. + +Keep this module independent of FastPLMs so the same extractor runs unchanged +inside isolated official-reference containers and candidate containers. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from typing import Any + + +SEMANTIC_PATHS: dict[str, tuple[str, ...]] = { + "vocab_size": ( + "config.vocab_size", + "vocab_size", + "alphabet_size", + "embed.num_embeddings", + "embeddings.word_embeddings.num_embeddings", + "encoder.sequence_embed.num_embeddings", + ), + "d_model": ( + "config.hidden_size", + "config.d_model", + "hidden_size", + "d_model", + "embed_dim", + "embed.embedding_dim", + "embeddings.word_embeddings.embedding_dim", + "encoder.sequence_embed.embedding_dim", + ), + "n_layers": ( + "config.num_hidden_layers", + "config.num_layers", + "config.n_layers", + "num_layers", + "transformer.blocks", + "layers", + "encoder.layer", + "encoder.block", + ), + "n_heads": ( + "config.num_attention_heads", + "config.num_heads", + "config.n_heads", + "attention_heads", + "transformer.blocks.0.attn.n_heads", + "layers.0.self_attn.num_heads", + "encoder.layer.0.attention.self.num_attention_heads", + ), + "d_ff": ("config.intermediate_size", "config.d_ff"), + "layer_norm_epsilon": ("config.layer_norm_eps", "config.layer_norm_epsilon"), + "max_positions": ("config.max_position_embeddings",), + "relative_buckets": ("config.relative_attention_num_buckets",), + "relative_max_distance": ("config.relative_attention_max_distance",), + "pad_token_id": ("config.pad_token_id", "padding_idx"), + "bos_token_id": ("config.bos_token_id", "cls_idx"), + "eos_token_id": ("config.eos_token_id", "eos_idx"), + "mask_token_id": ("config.mask_token_id", "mask_idx"), + "token_dropout": ("config.token_dropout", "token_dropout"), + "initializer_range": ("config.initializer_range",), + "classifier_dropout": ("config.classifier_dropout",), + "tie_word_embeddings": ("config.tie_word_embeddings",), +} + + +def _attribute(root: object, path: str) -> Any: + current = root + for part in path.split("."): + if part.isdigit() and hasattr(current, "__len__") and hasattr(current, "__getitem__"): + index = int(part) + if index >= len(current): + return None + current = current[index] + elif hasattr(current, part): + current = getattr(current, part) + else: + return None + return current + + +def semantic_config(model: nn.Module) -> dict[str, Any]: + """Extract the common inference semantics from an official or mirror model.""" + + roots: list[object] = [model] + if hasattr(model, "esm3"): + roots.insert(0, model.esm3) + result: dict[str, Any] = {} + for semantic_name, paths in SEMANTIC_PATHS.items(): + for root in roots: + for path in paths: + value = _attribute(root, path) + if value is None: + continue + if isinstance(value, (nn.ModuleList, list, tuple)): + value = len(value) + if torch.is_tensor(value) and value.numel() == 1: + value = value.item() + if isinstance(value, (str, int, float, bool)): + result[semantic_name] = value + break + if semantic_name in result: + break + missing = sorted({"vocab_size", "d_model", "n_layers", "n_heads"}.difference(result)) + if missing: + raise RuntimeError(f"Could not extract required semantic configuration fields: {missing}") + return result + + +def transformed_semantic_config(model: nn.Module, transform_name: str) -> dict[str, Any]: + """Extract semantics after applying a declared checkpoint conversion.""" + + result = semantic_config(model) + if transform_name == "dplm_to_fastplms_v1": + result["tie_word_embeddings"] = False + return result + + +__all__ = ["SEMANTIC_PATHS", "semantic_config", "transformed_semantic_config"] diff --git a/tests/parity/support/state_transforms.py b/tests/parity/support/state_transforms.py new file mode 100644 index 0000000..12f64aa --- /dev/null +++ b/tests/parity/support/state_transforms.py @@ -0,0 +1,193 @@ +"""Independent deterministic transforms for official checkpoint state. + +The transform name comes exclusively from ``models.toml``. A missing transform +is a compliance failure, not an invitation to compare only intersecting keys. +""" + +from __future__ import annotations + +import re +import torch +from collections.abc import Callable, Mapping + + +State = Mapping[str, torch.Tensor] +Transform = Callable[[State], dict[str, torch.Tensor]] + + +def _identity(state: State) -> dict[str, torch.Tensor]: + return dict(state) + + +def _cast_floating(state: State, dtype: torch.dtype) -> dict[str, torch.Tensor]: + return { + key: value.to(dtype=dtype) if value.is_floating_point() else value + for key, value in state.items() + } + + +def _drop_unused_rotary_position_table(state: State) -> dict[str, torch.Tensor]: + return { + key: value + for key, value in state.items() + if key != "esm.embeddings.position_embeddings.weight" + } + + +def _esm2_fair_to_fastplms(state: State) -> dict[str, torch.Tensor]: + """Map the pinned Meta ESM2 module names to the Hugging Face ESM schema.""" + + mapped: dict[str, torch.Tensor] = {} + projection_names = {"q_proj": "query", "k_proj": "key", "v_proj": "value"} + for key, value in state.items(): + target: str | None = None + if key == "embed_tokens.weight": + target = "esm.embeddings.word_embeddings.weight" + elif key.startswith("layers."): + match = re.fullmatch(r"layers\.(\d+)\.(.+)", key) + if match is None: + raise AssertionError(f"Unrecognized official ESM2 layer key: {key}") + layer, suffix = match.groups() + prefix = f"esm.encoder.layer.{layer}." + if suffix == "self_attn.rot_emb.inv_freq": + target = f"{prefix}attention.self.rotary_embeddings.inv_freq" + for source_name, target_name in projection_names.items(): + if suffix.startswith(f"self_attn.{source_name}."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.self.{target_name}.{parameter}" + break + if suffix.startswith("self_attn.out_proj."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.output.dense.{parameter}" + elif suffix.startswith("self_attn_layer_norm."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.LayerNorm.{parameter}" + elif suffix.startswith("fc1."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}intermediate.dense.{parameter}" + elif suffix.startswith("fc2."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}output.dense.{parameter}" + elif suffix.startswith("final_layer_norm."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}LayerNorm.{parameter}" + elif key.startswith("emb_layer_norm_after."): + target = f"esm.encoder.{key}" + elif key.startswith("contact_head."): + target = f"esm.{key}" + elif key == "lm_head.weight": + target = "lm_head.decoder.weight" + elif key == "lm_head.bias": + mapped["lm_head.bias"] = value + continue + elif key.startswith("lm_head."): + target = key + + if target is None: + raise AssertionError(f"Unrecognized official ESM2 state key: {key}") + mapped[target] = value + return mapped + + +def _esmc_to_fastplms(state: State) -> dict[str, torch.Tensor]: + """Apply the declared ESMC checkpoint-key normalization exactly once.""" + + replacements = ( + (".attn.layernorm_qkv.layer_norm_bias", ".attn.layernorm_qkv.0.bias"), + (".attn.layernorm_qkv.layer_norm_weight", ".attn.layernorm_qkv.0.weight"), + (".attn.layernorm_qkv.weight", ".attn.layernorm_qkv.1.weight"), + (".ffn.layer_norm_bias", ".ffn.0.bias"), + (".ffn.layer_norm_weight", ".ffn.0.weight"), + (".ffn.fc1_weight", ".ffn.1.weight"), + (".ffn.fc2_weight", ".ffn.3.weight"), + ) + mapped: dict[str, torch.Tensor] = {} + for raw_key, value in state.items(): + if raw_key.endswith("._extra_state"): + continue + key = raw_key.removeprefix("esmc.") + if key.startswith("lm_head."): + key = f"sequence_head.{key.removeprefix('lm_head.')}" + for source, target in replacements: + key = key.replace(source, target) + mapped[key] = value + return mapped + + +_ESMFOLD_DERIVED_BUFFERS = frozenset( + { + "positional_encoding._float_tensor", + "trunk.structure_module.atom_mask", + "trunk.structure_module.default_frames", + "trunk.structure_module.group_idx", + "trunk.structure_module.lit_positions", + } +) + + +def _esmfold_meta_to_fastplms(state: State) -> dict[str, torch.Tensor]: + """Map Meta ESMFold state and omit heads unused by structure inference.""" + + if any(key.startswith("esm.encoder.") for key in state): + return { + key: value + for key, value in state.items() + if key not in _ESMFOLD_DERIVED_BUFFERS + and not key.startswith(("mlm_head.", "esm.contact_head.")) + } + + folding: dict[str, torch.Tensor] = {} + native_esm: dict[str, torch.Tensor] = {} + for key, value in state.items(): + if key in _ESMFOLD_DERIVED_BUFFERS: + continue + if key.startswith("esm."): + inner = key.removeprefix("esm.") + if inner.startswith(("lm_head.", "contact_head.")): + continue + native_esm[inner] = value + else: + folding[key] = value + mapped_esm = _esm2_fair_to_fastplms(native_esm) + overlap = set(folding).intersection(mapped_esm) + if overlap: + raise AssertionError(f"ESMFold state-key collision: {sorted(overlap)[:20]}") + return {**folding, **mapped_esm} + + +TRANSFORMS: dict[str, Transform] = { + "identity": _identity, + "esm2_hf_to_fastplms_v1": _esm2_fair_to_fastplms, + "esmc_to_fastplms_v1": _esmc_to_fastplms, + "esm3_to_fastplms_v1": lambda state: _cast_floating(state, torch.float32), + "e1_to_fastplms_v1": lambda state: _cast_floating(state, torch.bfloat16), + "dplm_to_fastplms_v1": _drop_unused_rotary_position_table, + "dplm2_to_fastplms_v1": _drop_unused_rotary_position_table, + "ankh_t5_to_fastplms_v1": _identity, + "esmfold_meta_to_fastplms_v1": _esmfold_meta_to_fastplms, +} + + +def transform_state(name: str, state: State) -> dict[str, torch.Tensor]: + """Return transformed state or fail if the manifest names no implementation.""" + + try: + transform = TRANSFORMS[name] + except KeyError as error: + raise AssertionError(f"No compliance state transform is registered for {name!r}") from error + return transform(state) + + +def transform_parameter_names(name: str, parameter_name: str) -> tuple[str, ...]: + """Map one parameter name, including intentionally duplicated alias names.""" + + # marker: (0,) + marker = torch.empty(0) + transformed = transform_state(name, {parameter_name: marker}) + return tuple(transformed) + + +def transform_preserves_aliases(name: str) -> bool: + """Return whether the declared conversion retains source parameter aliases.""" + + return name not in {"dplm_to_fastplms_v1", "esm2_hf_to_fastplms_v1"} diff --git a/tests/parity/test_ankh_seq2seq_parity.py b/tests/parity/test_ankh_seq2seq_parity.py new file mode 100644 index 0000000..981c2bc --- /dev/null +++ b/tests/parity/test_ankh_seq2seq_parity.py @@ -0,0 +1,138 @@ +"""Official ANKH sequence-to-sequence head and alias compliance.""" + +from __future__ import annotations + +import gc +import pytest +import torch +import torch.nn.functional as F +from pathlib import Path +from typing import Any +from transformers import AutoModelForSeq2SeqLM + +from fastplms.models.ankh.modeling_ankh import tokenize_ankh_sequences +from fastplms.registry import ModelSpec, get_model_registry +from tests.parity.support.reference_adapters.ankh import load_official_seq2seq + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.slow] +ANKH_SPECS = get_model_registry().by_family("ankh") + + +def _parameter(spec: ModelSpec) -> Any: + marks: list[Any] = [] + if spec.size_category == "xlarge": + marks.append(pytest.mark.large) + return pytest.param(spec, id=spec.id, marks=marks) + + +def _alias_groups(model: torch.nn.Module) -> set[frozenset[str]]: + groups: dict[int, set[str]] = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + groups.setdefault(id(parameter), set()).add(name) + return {frozenset(names) for names in groups.values() if len(names) > 1} + + +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in ANKH_SPECS]) +def test_ankh_official_seq2seq_state_aliases_and_seeded_inference( + spec: ModelSpec, + tmp_path: Path, +) -> None: + device = torch.device("cuda") + fast = AutoModelForSeq2SeqLM.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + trust_remote_code=True, + dtype=torch.bfloat16, + device_map=device, + attn_implementation="eager", + ).eval() + official, tokenizer = load_official_seq2seq( + reference_repo_id=spec.official.repo_id, + reference_revision=spec.official.revision, + device=device, + dtype=torch.bfloat16, + ) + + fast_state = fast.state_dict() + official_state = official.state_dict() + assert set(fast_state) == set(official_state) + for name in sorted(fast_state): + assert fast_state[name].shape == official_state[name].shape + assert fast_state[name].dtype == official_state[name].dtype + assert torch.equal(fast_state[name], official_state[name]), ( + f"{spec.id}:{name}: sequence-to-sequence weight differs" + ) + assert _alias_groups(fast) == _alias_groups(official), ( + f"{spec.id}: sequence-to-sequence tied-weight contract differs" + ) + + encoded = tokenize_ankh_sequences( + tokenizer, + ["MSTNPK", "ACDE"], + return_tensors="pt", + padding=True, + ) + inputs = {name: value.to(device) for name, value in encoded.items() if torch.is_tensor(value)} + # decoder_input_ids: (b, l) + decoder_input_ids = inputs["input_ids"] + with torch.inference_mode(): + # fast_logits: (..., c) + fast_logits = fast( + **inputs, + decoder_input_ids=decoder_input_ids, + return_dict=True, + ).logits.float() + # official_logits: (..., c) + official_logits = official( + **inputs, + decoder_input_ids=decoder_input_ids, + return_dict=True, + ).logits.float() + relative_l2 = torch.linalg.vector_norm(fast_logits - official_logits) / ( + torch.linalg.vector_norm(official_logits).clamp_min(torch.finfo(torch.float32).tiny) + ) + cosine = F.cosine_similarity(fast_logits, official_logits, dim=-1) + assert float(relative_l2) <= 1e-2 + assert float(torch.quantile(cosine, 0.01)) >= 0.999 + + generation_kwargs = { + **inputs, + "do_sample": True, + "top_k": 5, + "max_new_tokens": 4, + } + torch.manual_seed(42) + fast_tokens = fast.generate(**generation_kwargs) + torch.manual_seed(42) + official_tokens = official.generate(**generation_kwargs) + assert torch.equal(fast_tokens, official_tokens), ( + f"{spec.id}: seeded sequence-to-sequence generation differs" + ) + + save_path = tmp_path / "seq2seq" + fast.save_pretrained(save_path, safe_serialization=True) + tokenizer.save_pretrained(save_path) + del fast + torch.cuda.empty_cache() + reloaded = AutoModelForSeq2SeqLM.from_pretrained( + save_path, + trust_remote_code=True, + local_files_only=True, + dtype=torch.bfloat16, + device_map=device, + attn_implementation="eager", + ).eval() + reloaded_state = reloaded.state_dict() + assert set(reloaded_state) == set(official_state) + for name, official_tensor in official_state.items(): + assert torch.equal(reloaded_state[name], official_tensor) + torch.manual_seed(42) + reloaded_tokens = reloaded.generate(**generation_kwargs) + assert torch.equal(reloaded_tokens, official_tokens), ( + f"{spec.id}: save/reload changed seeded sequence-to-sequence generation" + ) + + del reloaded, official + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/parity/test_boltz_source_refactor.py b/tests/parity/test_boltz_source_refactor.py new file mode 100644 index 0000000..683810c --- /dev/null +++ b/tests/parity/test_boltz_source_refactor.py @@ -0,0 +1,1359 @@ +"""Focused parity checks for independently maintained Boltz runtime helpers.""" + +from __future__ import annotations + +import importlib +import importlib.util +import sys +import numpy as np +import pytest +import torch +from difflib import SequenceMatcher +from pathlib import Path +from types import ModuleType + +from fastplms.models.boltz import vb_const +from fastplms.models.boltz import vb_layers_attention as pair_attention +from fastplms.models.boltz import vb_layers_attentionv2 as pair_attention_v2 +from fastplms.models.boltz import vb_layers_confidence_utils as confidence +from fastplms.models.boltz import vb_layers_dropout as dropout +from fastplms.models.boltz import vb_layers_outer_product_mean as outer_product +from fastplms.models.boltz import vb_layers_pair_averaging as pair_averaging +from fastplms.models.boltz import vb_layers_pairformer as pairformer +from fastplms.models.boltz import vb_layers_transition as transition +from fastplms.models.boltz import vb_layers_triangular_mult as triangular_mult +from fastplms.models.boltz import vb_loss_diffusionv2 as diffusion_loss +from fastplms.models.boltz import vb_modules_diffusion_conditioning as conditioning +from fastplms.models.boltz import vb_modules_encodersv2 as encoders +from fastplms.models.boltz import vb_modules_transformersv2 as diffusion_transformers +from fastplms.models.boltz import vb_modules_utils as module_utils +from fastplms.models.boltz import vb_potentials_potentials as potentials +from fastplms.models.boltz import vb_potentials_schedules as schedules +from fastplms.models.boltz import vb_tri_attn_attention as triangle_attention +from fastplms.models.boltz import vb_tri_attn_primitives as primitives +from fastplms.models.boltz import vb_tri_attn_utils as attention_utils +from fastplms.models.boltz.minimal_featurizer import build_boltz2_features + + +pytestmark = pytest.mark.structure + +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +UPSTREAM_ROOT = REPOSITORY_ROOT / "vendor" / "upstream" / "boltz" / "src" + +REWRITTEN_SOURCE_PAIRS = ( + ("vb_const.py", "boltz/data/const.py"), + ("vb_layers_attention.py", "boltz/model/layers/attention.py"), + ("vb_layers_attentionv2.py", "boltz/model/layers/attentionv2.py"), + ("vb_layers_initialize.py", "boltz/model/layers/initialize.py"), + ("vb_layers_outer_product_mean.py", "boltz/model/layers/outer_product_mean.py"), + ("vb_layers_pair_averaging.py", "boltz/model/layers/pair_averaging.py"), + ("vb_layers_pairformer.py", "boltz/model/layers/pairformer.py"), + ("vb_layers_transition.py", "boltz/model/layers/transition.py"), + ("vb_layers_triangular_mult.py", "boltz/model/layers/triangular_mult.py"), + ("vb_layers_dropout.py", "boltz/model/layers/dropout.py"), + ("vb_potentials_schedules.py", "boltz/model/potentials/schedules.py"), + ( + "vb_tri_attn_attention.py", + "boltz/model/layers/triangular_attention/attention.py", + ), + ( + "vb_tri_attn_utils.py", + "boltz/model/layers/triangular_attention/utils.py", + ), + ( + "vb_tri_attn_primitives.py", + "boltz/model/layers/triangular_attention/primitives.py", + ), + ("vb_layers_confidence_utils.py", "boltz/model/layers/confidence_utils.py"), + ("vb_loss_diffusionv2.py", "boltz/model/loss/diffusionv2.py"), + ( + "vb_modules_diffusion_conditioning.py", + "boltz/model/modules/diffusion_conditioning.py", + ), + ("vb_modules_transformersv2.py", "boltz/model/modules/transformersv2.py"), + ("vb_modules_trunkv2.py", "boltz/model/modules/trunkv2.py"), + ("vb_modules_utils.py", "boltz/model/modules/utils.py"), + ("vb_modules_encodersv2.py", "boltz/model/modules/encodersv2.py"), + ("vb_potentials_potentials.py", "boltz/model/potentials/potentials.py"), +) + + +def _install_import_only_dependency_stubs() -> None: + """Provide initialization-only reference shims outside the core extras.""" + + sys.modules.setdefault("einx", ModuleType("einx")) + if "scipy.stats" in sys.modules: + return + + class _InitializationOnlyTruncatedNormal: + @staticmethod + def std(*args: object, **kwargs: object) -> float: + return 1.0 + + @staticmethod + def rvs(*args: object, **kwargs: object) -> np.ndarray: + del args + return np.zeros(kwargs["size"], dtype=np.float32) + + scipy = ModuleType("scipy") + scipy.__path__ = [] # type: ignore[attr-defined] + stats = ModuleType("scipy.stats") + stats.truncnorm = _InitializationOnlyTruncatedNormal # type: ignore[attr-defined] + scipy.stats = stats # type: ignore[attr-defined] + sys.modules["scipy"] = scipy + sys.modules["scipy.stats"] = stats + + +def _load_standalone(relative_path: str, name: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(name, UPSTREAM_ROOT / relative_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize(("local_name", "upstream_name"), REWRITTEN_SOURCE_PAIRS) +def test_boltz_runtime_is_not_an_upstream_source_relocation( + local_name: str, + upstream_name: str, +) -> None: + local_path = REPOSITORY_ROOT / "src" / "fastplms" / "models" / "boltz" / local_name + upstream_path = UPSTREAM_ROOT / upstream_name + local_lines = local_path.read_text(encoding="utf-8").splitlines() + upstream_lines = upstream_path.read_text(encoding="utf-8").splitlines() + similarity = SequenceMatcher(None, local_lines, upstream_lines).ratio() + assert similarity < 0.75, f"{local_name} has line similarity {similarity:.3f}" + + +@pytest.fixture(scope="module") +def upstream_const() -> ModuleType: + return _load_standalone("boltz/data/const.py", "fastplms_test_upstream_const") + + +@pytest.fixture(scope="module") +def upstream_dropout() -> ModuleType: + return _load_standalone( + "boltz/model/layers/dropout.py", + "fastplms_test_upstream_dropout", + ) + + +@pytest.fixture(scope="module") +def upstream_schedules() -> ModuleType: + return _load_standalone( + "boltz/model/potentials/schedules.py", + "fastplms_test_upstream_schedules", + ) + + +@pytest.fixture(scope="module") +def upstream_attention_utils() -> ModuleType: + return _load_standalone( + "boltz/model/layers/triangular_attention/utils.py", + "fastplms_test_upstream_attention_utils", + ) + + +@pytest.fixture(scope="module") +def upstream_diffusion_loss() -> ModuleType: + _install_import_only_dependency_stubs() + return _load_standalone( + "boltz/model/loss/diffusionv2.py", + "fastplms_test_upstream_diffusion_loss", + ) + + +@pytest.fixture(scope="module") +def upstream_package() -> ModuleType: + sys.path.insert(0, str(UPSTREAM_ROOT)) + try: + yield importlib.import_module("boltz") + finally: + sys.path.remove(str(UPSTREAM_ROOT)) + + +def test_runtime_constants_match_upstream(upstream_const: ModuleType) -> None: + retained_names = ( + "chain_types", + "chain_type_ids", + "canonical_tokens", + "tokens", + "token_ids", + "num_tokens", + "prot_letter_to_token", + "ref_atoms", + "protein_backbone_atom_names", + "nucleic_backbone_atom_names", + "protein_backbone_atom_index", + "nucleic_backbone_atom_index", + "res_to_center_atom", + "res_to_disto_atom", + "num_elements", + "bond_types", + "contact_conditioning_info", + "chunk_size_threshold", + "method_types_ids", + "num_method_types", + "vdw_radii", + ) + for name in retained_names: + assert getattr(vb_const, name) == getattr(upstream_const, name), name + + +@pytest.mark.parametrize("columnwise", [False, True]) +@pytest.mark.parametrize("training", [False, True]) +def test_dropout_mask_matches_upstream( + upstream_dropout: ModuleType, + columnwise: bool, + training: bool, +) -> None: + # pair: (2, 5, 7, 3) + pair = torch.empty(2, 5, 7, 3) + torch.manual_seed(917) + expected = upstream_dropout.get_dropout_mask(0.2, pair, training, columnwise) + torch.manual_seed(917) + actual = dropout.get_dropout_mask(0.2, pair, training, columnwise) + assert torch.equal(actual, expected) + + +def test_parameter_schedules_match_upstream(upstream_schedules: ModuleType) -> None: + reference_exp = upstream_schedules.ExponentialInterpolation(0.1, 4.0, 2.5) + local_exp = schedules.ExponentialInterpolation(0.1, 4.0, 2.5) + reference_step = upstream_schedules.PiecewiseStepFunction((0.2, 0.7), (1, 2, 5)) + local_step = schedules.PiecewiseStepFunction((0.2, 0.7), (1, 2, 5)) + for time in (0.0, 0.2, 0.21, 0.7, 0.9, 1.0): + assert local_exp.compute(time) == reference_exp.compute(time) + assert local_step.compute(time) == reference_step.compute(time) + + +@pytest.mark.parametrize("low_mem", [False, True]) +def test_chunk_layer_matches_upstream( + upstream_attention_utils: ModuleType, + low_mem: bool, +) -> None: + inputs = { + "left": torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4), + "right": torch.tensor([[[2.0, 1.0, 0.0, -1.0]]]), + } + + def layer(left: torch.Tensor, right: torch.Tensor) -> dict[str, torch.Tensor]: + # left: (...), right: (...) + return {"sum": left + right, "product": left * right} + + expected = upstream_attention_utils.chunk_layer( + layer, + inputs, + chunk_size=4, + no_batch_dims=2, + low_mem=low_mem, + ) + actual = attention_utils.chunk_layer( + layer, + inputs, + chunk_size=4, + no_batch_dims=2, + low_mem=low_mem, + ) + assert torch.equal(actual["sum"], expected["sum"]) + assert torch.equal(actual["product"], expected["product"]) + + +@pytest.mark.parametrize( + ("start", "stop"), + [(0, 1), (1, 7), (3, 19), (0, 24), (17, 24)], +) +def test_low_memory_flat_slice_matches_upstream( + upstream_attention_utils: ModuleType, + start: int, + stop: int, +) -> None: + # tensor: (2, 3, 4, 5) + tensor = torch.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5) + expected = upstream_attention_utils._chunk_slice(tensor, start, stop, 3) + actual = attention_utils._chunk_slice(tensor, start, stop, 3) + assert torch.equal(actual, expected) + + +def test_rigid_alignment_matches_upstream(upstream_diffusion_loss: ModuleType) -> None: + generator = torch.Generator().manual_seed(761) + # true_coords: (2, 8, 3) + true_coords = torch.randn(2, 8, 3, generator=generator) + # pred_coords: (2, 8, 3) + pred_coords = torch.randn(2, 8, 3, generator=generator) + # weights: (2, 8) + weights = torch.rand(2, 8, generator=generator) + # mask: (2, 8) + mask = torch.ones(2, 8) + expected = upstream_diffusion_loss.weighted_rigid_align( + true_coords, + pred_coords, + weights, + mask, + ) + actual = diffusion_loss.weighted_rigid_align( + true_coords, + pred_coords, + weights, + mask, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_smooth_lddt_matches_upstream(upstream_diffusion_loss: ModuleType) -> None: + generator = torch.Generator().manual_seed(177) + # pred_coords: (2, 10, 3) + pred_coords = torch.randn(2, 10, 3, generator=generator) + # true_coords: (2, 10, 3) + true_coords = torch.randn(2, 10, 3, generator=generator) + # is_nucleotide: (2, 10) + is_nucleotide = torch.tensor([[0.0] * 5 + [1.0] * 5] * 2) + # coords_mask: (2, 10) + coords_mask = torch.ones(2, 10) + expected = upstream_diffusion_loss.smooth_lddt_loss( + pred_coords, + true_coords, + is_nucleotide, + coords_mask, + ) + actual = diffusion_loss.smooth_lddt_loss( + pred_coords, + true_coords, + is_nucleotide, + coords_mask, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_confidence_scalar_helpers_match_upstream(upstream_package: ModuleType) -> None: + del upstream_package + reference = importlib.import_module("boltz.model.layers.confidence_utils") + # logits: (2, 4, 50) + logits = torch.randn(2, 4, 50, generator=torch.Generator().manual_seed(93)) + expected = reference.compute_aggregated_metric(logits) + actual = confidence.compute_aggregated_metric(logits) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + # distances: (n,) + distances = torch.linspace(0, 32, 17) + # residues: (2, 1) + residues = torch.tensor([[25.0], [300.0]]) + expected_tm = reference.tm_function(distances, residues) + actual_tm = confidence.tm_function(distances, residues) + torch.testing.assert_close(actual_tm, expected_tm, rtol=0, atol=0) + + +def test_attention_state_and_forward_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.layers.triangular_attention.primitives") + torch.manual_seed(12) + reference = reference_module.Attention(8, 8, 8, 4, 2, gating=True) + torch.manual_seed(12) + local = primitives.Attention(8, 8, 8, 4, 2, gating=True) + assert local.state_dict().keys() == reference.state_dict().keys() + for name, tensor in reference.state_dict().items(): + assert torch.equal(local.state_dict()[name], tensor), name + + generator = torch.Generator().manual_seed(33) + # query: (2, 5, 8) + query = torch.randn(2, 5, 8, generator=generator) + # key_value: (2, 7, 8) + key_value = torch.randn(2, 7, 8, generator=generator) + # triangle_bias: (2, 2, 5, 7) + triangle_bias = torch.randn(2, 2, 5, 7, generator=generator) + # mask_bias: (2, 2, 5, 7) + mask_bias = torch.zeros(2, 2, 5, 7) + # mask: (2, 5, 7) + mask = torch.ones(2, 5, 7) + expected = reference(query, key_value, triangle_bias, mask_bias, mask) + actual = local(query, key_value, triangle_bias, mask_bias, mask) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_transition_state_and_dense_or_chunked_forward_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module("boltz.model.layers.transition").Transition + local = transition.Transition(dim=8, hidden=10, out_dim=6).eval() + reference = reference_class(dim=8, hidden=10, out_dim=6).eval() + reference.load_state_dict(local.state_dict(), strict=True) + assert reference.state_dict().keys() == local.state_dict().keys() + + # X is a deterministic transition input with shape (b, l, d). + # input_tensor: (2, 5, 8) + input_tensor = torch.randn(2, 5, 8, generator=torch.Generator().manual_seed(271)) + for chunk_size in (None, 4, 16): + expected = reference(input_tensor, chunk_size=chunk_size) + actual = local(input_tensor, chunk_size=chunk_size) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def _fill_parameters(module: torch.nn.Module, seed: int) -> None: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in module.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator)) + + +def test_pair_biased_attention_matches_upstream_and_cache_contract( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module("boltz.model.layers.attention").AttentionPairBias + local = pair_attention.AttentionPairBias(8, 4, 2).eval() + _fill_parameters(local, seed=281) + reference = reference_class(8, 4, 2).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + generator = torch.Generator().manual_seed(283) + # sequence_states: (2, 5, 8) + sequence_states = torch.randn(2, 5, 8, generator=generator) + # pair_states: (2, 5, 5, 4) + pair_states = torch.randn(2, 5, 5, 4, generator=generator) + # mask: (2, 5) + mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], dtype=torch.float32) + expected_cache: dict[str, torch.Tensor] = {} + actual_cache: dict[str, torch.Tensor] = {} + expected = reference(sequence_states, pair_states, mask, model_cache=expected_cache) + actual = local(sequence_states, pair_states, mask, model_cache=actual_cache) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual_cache["z"], expected_cache["z"], rtol=0, atol=0) + + # replacement_pairs: (2, 5, 5, 4) + replacement_pairs = torch.randn(2, 5, 5, 4, generator=generator) + expected_cached = reference( + sequence_states, + replacement_pairs, + mask, + model_cache=expected_cache, + ) + actual_cached = local( + sequence_states, + replacement_pairs, + mask, + model_cache=actual_cache, + ) + torch.testing.assert_close(actual_cached, expected_cached, rtol=0, atol=0) + + +@pytest.mark.parametrize("compute_pair_bias", [False, True]) +def test_pair_biased_cross_attention_matches_upstream( + upstream_package: ModuleType, + compute_pair_bias: bool, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module("boltz.model.layers.attentionv2").AttentionPairBias + kwargs = { + "c_s": 8, + "c_z": 4, + "num_heads": 2, + "compute_pair_bias": compute_pair_bias, + } + local = pair_attention_v2.AttentionPairBias(**kwargs).eval() + _fill_parameters(local, seed=293) + reference = reference_class(**kwargs).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + generator = torch.Generator().manual_seed(307) + # query_states: (2, 3, 8) + query_states = torch.randn(2, 3, 8, generator=generator) + # key_states: (2, 5, 8) + key_states = torch.randn(2, 5, 8, generator=generator) + pair_width = 4 if compute_pair_bias else 2 + # pair_states: (2, 3, 5, pair_width) + pair_states = torch.randn(2, 3, 5, pair_width, generator=generator) + # mask: (2, 5) + mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], dtype=torch.float32) + expected = reference(query_states, pair_states, mask, key_states) + actual = local(query_states, pair_states, mask, key_states) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_outer_product_mean_dense_and_chunked_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module( + "boltz.model.layers.outer_product_mean" + ).OuterProductMean + local = outer_product.OuterProductMean(c_in=5, c_hidden=3, c_out=7).eval() + _fill_parameters(local, seed=311) + reference = reference_class(c_in=5, c_hidden=3, c_out=7).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + generator = torch.Generator().manual_seed(313) + # msa_states: (2, 3, 4, 5) + msa_states = torch.randn(2, 3, 4, 5, generator=generator) + # mask: (2, 3, 4) + mask = torch.tensor( + [ + [[1, 1, 1, 1], [1, 1, 1, 0], [1, 1, 0, 0]], + [[1, 1, 1, 1], [1, 0, 1, 0], [1, 1, 1, 0]], + ], + dtype=torch.float32, + ) + for chunk_size in (None, 2, 8): + expected = reference(msa_states, mask, chunk_size=chunk_size) + actual = local(msa_states, mask, chunk_size=chunk_size) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("chunk_heads", [False, True]) +def test_pair_weighted_averaging_matches_upstream( + upstream_package: ModuleType, + chunk_heads: bool, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module( + "boltz.model.layers.pair_averaging" + ).PairWeightedAveraging + kwargs = {"c_m": 6, "c_z": 5, "c_h": 2, "num_heads": 3} + local = pair_averaging.PairWeightedAveraging(**kwargs).eval() + _fill_parameters(local, seed=317) + reference = reference_class(**kwargs).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + generator = torch.Generator().manual_seed(331) + # msa_states: (2, 3, 4, 6) + msa_states = torch.randn(2, 3, 4, 6, generator=generator) + # pair_states: (2, 4, 4, 5) + pair_states = torch.randn(2, 4, 4, 5, generator=generator) + # mask: (2, 4, 4) + mask = torch.tensor( + [ + [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 0], [1, 1, 0, 0]], + [[1, 1, 1, 1], [1, 1, 1, 0], [1, 1, 1, 0], [1, 0, 0, 0]], + ], + dtype=torch.float32, + ) + expected = reference(msa_states, pair_states, mask, chunk_heads=chunk_heads) + actual = local(msa_states, pair_states, mask, chunk_heads=chunk_heads) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("class_name", "local_class"), + [ + ( + "TriangleMultiplicationOutgoing", + triangular_mult.TriangleMultiplicationOutgoing, + ), + ( + "TriangleMultiplicationIncoming", + triangular_mult.TriangleMultiplicationIncoming, + ), + ], +) +def test_triangular_multiplication_matches_upstream( + upstream_package: ModuleType, + class_name: str, + local_class: type[torch.nn.Module], +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = getattr( + importlib.import_module("boltz.model.layers.triangular_mult"), + class_name, + ) + local = local_class(dim=4).eval() + _fill_parameters(local, seed=337) + reference = reference_class(dim=4).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + # pair_states: (2, 3, 3, 4) + pair_states = torch.randn(2, 3, 3, 4, generator=torch.Generator().manual_seed(347)) + # mask: (2, 3, 3) + mask = torch.tensor( + [ + [[1, 1, 1], [1, 1, 1], [1, 1, 0]], + [[1, 1, 1], [1, 1, 0], [1, 0, 0]], + ], + dtype=torch.float32, + ) + expected = reference(pair_states, mask) + actual = local(pair_states, mask) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("class_name", "local_class"), + [ + ( + "TriangleAttentionStartingNode", + triangle_attention.TriangleAttentionStartingNode, + ), + ( + "TriangleAttentionEndingNode", + triangle_attention.TriangleAttentionEndingNode, + ), + ], +) +def test_triangle_attention_dense_and_chunked_match_upstream( + upstream_package: ModuleType, + class_name: str, + local_class: type[torch.nn.Module], +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = getattr( + importlib.import_module("boltz.model.layers.triangular_attention.attention"), + class_name, + ) + local = local_class(c_in=8, c_hidden=4, no_heads=2).eval() + _fill_parameters(local, seed=349) + reference = reference_class(c_in=8, c_hidden=4, no_heads=2).eval() + reference.load_state_dict(local.state_dict(), strict=True) + + # pair_states: (2, 3, 3, 8) + pair_states = torch.randn(2, 3, 3, 8, generator=torch.Generator().manual_seed(353)) + # mask: (2, 3, 3) + mask = torch.tensor( + [ + [[1, 1, 1], [1, 1, 1], [1, 1, 0]], + [[1, 1, 1], [1, 1, 0], [1, 0, 0]], + ], + dtype=torch.float32, + ) + for chunk_size in (None, 2): + expected = reference(pair_states, mask, chunk_size=chunk_size) + actual = local(pair_states, mask, chunk_size=chunk_size) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_conditioned_diffusion_transformer_matches_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.modules.transformersv2") + kwargs = { + "depth": 2, + "heads": 2, + "dim": 8, + "dim_single_cond": 6, + "post_layer_norm": True, + } + local = diffusion_transformers.DiffusionTransformer(**kwargs).eval() + _fill_parameters(local, seed=359) + reference = reference_module.DiffusionTransformer(**kwargs).eval() + reference.load_state_dict(local.state_dict(), strict=True) + assert reference.state_dict().keys() == local.state_dict().keys() + + generator = torch.Generator().manual_seed(367) + # activations: (2, 4, 8) + activations = torch.randn(2, 4, 8, generator=generator) + # conditioning_states: (2, 4, 6) + conditioning_states = torch.randn(2, 4, 6, generator=generator) + # pair_bias: (2, 4, 4, 4) + pair_bias = torch.randn(2, 4, 4, 4, generator=generator) + # mask: (2, 4) + mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=torch.float32) + expected = reference(activations, conditioning_states, pair_bias, mask) + actual = local(activations, conditioning_states, pair_bias, mask) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_rotation_and_coordinate_augmentation_utilities_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + reference = importlib.import_module("boltz.model.modules.utils") + for function_name in ("random_quaternions", "random_rotations"): + torch.manual_seed(373) + expected = getattr(reference, function_name)(5, dtype=torch.float64) + torch.manual_seed(373) + actual = getattr(module_utils, function_name)(5, dtype=torch.float64) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + generator = torch.Generator().manual_seed(379) + # coordinates: (2, 5, 3) + coordinates = torch.randn(2, 5, 3, generator=generator) + # second: (2, 5, 3) + second = torch.randn(2, 5, 3, generator=generator) + # mask: (2, 5) + mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], dtype=torch.float32) + torch.manual_seed(383) + expected = reference.center_random_augmentation( + coordinates, + mask, + s_trans=0.4, + return_second_coords=True, + second_coords=second, + ) + torch.manual_seed(383) + actual = module_utils.center_random_augmentation( + coordinates, + mask, + s_trans=0.4, + return_second_coords=True, + second_coords=second, + ) + assert isinstance(actual, tuple) and isinstance(expected, tuple) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1], expected[1], rtol=0, atol=0) + + +def test_exponential_moving_average_matches_upstream(upstream_package: ModuleType) -> None: + del upstream_package + reference_class = importlib.import_module("boltz.model.modules.utils").ExponentialMovingAverage + local_parameters = [ + torch.nn.Parameter(torch.tensor([1.0, 2.0])), + torch.nn.Parameter(torch.tensor([[3.0], [4.0]])), + ] + reference_parameters = [ + torch.nn.Parameter(parameter.detach().clone()) for parameter in local_parameters + ] + local = module_utils.ExponentialMovingAverage(local_parameters, decay=0.95) + reference = reference_class(reference_parameters, decay=0.95) + with torch.no_grad(): + for local_parameter, reference_parameter in zip( + local_parameters, + reference_parameters, + strict=True, + ): + local_parameter.add_(0.75) + reference_parameter.add_(0.75) + local.update(local_parameters) + reference.update(reference_parameters) + assert local.num_updates == reference.num_updates + for actual, expected in zip( + local.shadow_params, + reference.shadow_params, + strict=True, + ): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_joint_and_pair_only_pairformer_layers_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.layers.pairformer") + joint_kwargs = { + "token_s": 8, + "token_z": 4, + "num_heads": 2, + "dropout": 0.0, + "pairwise_head_width": 2, + "pairwise_num_heads": 2, + "post_layer_norm": True, + "v2": True, + } + local_joint = pairformer.PairformerLayer(**joint_kwargs).eval() + _fill_parameters(local_joint, seed=389) + reference_joint = reference_module.PairformerLayer(**joint_kwargs).eval() + reference_joint.load_state_dict(local_joint.state_dict(), strict=True) + assert reference_joint.state_dict().keys() == local_joint.state_dict().keys() + + generator = torch.Generator().manual_seed(397) + # sequence_states: (2, 3, 8) + sequence_states = torch.randn(2, 3, 8, generator=generator) + # pair_states: (2, 3, 3, 4) + pair_states = torch.randn(2, 3, 3, 4, generator=generator) + # mask: (2, 3) + mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.float32) + # pair_mask: (...) + pair_mask = mask[:, :, None] * mask[:, None, :] + expected_s, expected_z = reference_joint( + sequence_states, + pair_states, + mask, + pair_mask, + chunk_size_tri_attn=2, + ) + actual_s, actual_z = local_joint( + sequence_states, + pair_states, + mask, + pair_mask, + chunk_size_tri_attn=2, + ) + torch.testing.assert_close(actual_s, expected_s, rtol=0, atol=0) + torch.testing.assert_close(actual_z, expected_z, rtol=0, atol=0) + + pair_kwargs = { + "token_z": 4, + "dropout": 0.0, + "pairwise_head_width": 2, + "pairwise_num_heads": 2, + } + local_pair = pairformer.PairformerNoSeqLayer(**pair_kwargs).eval() + _fill_parameters(local_pair, seed=401) + reference_pair = reference_module.PairformerNoSeqLayer(**pair_kwargs).eval() + reference_pair.load_state_dict(local_pair.state_dict(), strict=True) + expected_z = reference_pair(pair_states, pair_mask, chunk_size_tri_attn=2) + actual_z = local_pair(pair_states, pair_mask, chunk_size_tri_attn=2) + torch.testing.assert_close(actual_z, expected_z, rtol=0, atol=0) + + +def test_diffusion_conditioning_state_schema_matches_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_class = importlib.import_module( + "boltz.model.modules.diffusion_conditioning" + ).DiffusionConditioning + kwargs = { + "token_s": 8, + "token_z": 6, + "atom_s": 4, + "atom_z": 5, + "atoms_per_window_queries": 2, + "atoms_per_window_keys": 4, + "atom_encoder_depth": 2, + "atom_encoder_heads": 2, + "token_transformer_depth": 3, + "token_transformer_heads": 2, + "atom_decoder_depth": 2, + "atom_decoder_heads": 2, + "atom_feature_dim": 10, + "conditioning_transition_layers": 1, + } + torch.manual_seed(41) + reference = reference_class(**kwargs) + torch.manual_seed(41) + local = conditioning.DiffusionConditioning(**kwargs) + assert local.state_dict().keys() == reference.state_dict().keys() + for name, tensor in reference.state_dict().items(): + assert local.state_dict()[name].shape == tensor.shape, name + + +def _to_device( + tree: dict[str, torch.Tensor], + device: torch.device, +) -> dict[str, torch.Tensor]: + return {name: tensor.to(device) for name, tensor in tree.items()} + + +def _load_local_state(reference: torch.nn.Module, local: torch.nn.Module) -> None: + assert reference.state_dict().keys() == local.state_dict().keys() + reference.load_state_dict(local.state_dict(), strict=True) + + +@pytest.mark.gpu +def test_encoder_components_match_upstream_on_gpu( + upstream_package: ModuleType, +) -> None: + del upstream_package + assert torch.cuda.is_available(), "CUDA is required for Boltz encoder parity" + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.modules.encodersv2") + device = torch.device("cuda") + + torch.manual_seed(101) + local_fourier = encoders.FourierEmbedding(8).to(device) + reference_fourier = reference_module.FourierEmbedding(8).to(device) + _load_local_state(reference_fourier, local_fourier) + # times: (2,) + times = torch.tensor([0.1, 0.7], device=device) + torch.testing.assert_close( + local_fourier(times), + reference_fourier(times), + rtol=0, + atol=0, + ) + + relative_kwargs = { + "token_z": 7, + "r_max": 4, + "s_max": 2, + "fix_sym_check": True, + "cyclic_pos_enc": True, + } + local_relative = encoders.RelativePositionEncoder(**relative_kwargs).to(device) + reference_relative = reference_module.RelativePositionEncoder(**relative_kwargs).to(device) + _load_local_state(reference_relative, local_relative) + relative_features = { + "asym_id": torch.tensor([[0, 0, 1, 1]], device=device), + "residue_index": torch.tensor([[0, 1, 0, 1]], device=device), + "entity_id": torch.tensor([[0, 0, 1, 1]], device=device), + "token_index": torch.tensor([[0, 1, 2, 3]], device=device), + "sym_id": torch.tensor([[0, 1, 0, 1]], device=device), + "cyclic_period": torch.tensor([[0.0, 2.0, 0.0, 0.0]], device=device), + } + torch.testing.assert_close( + local_relative(relative_features), + reference_relative(relative_features), + rtol=0, + atol=0, + ) + + local_single = encoders.SingleConditioning( + sigma_data=16.0, + token_s=4, + dim_fourier=6, + num_transitions=2, + ).to(device) + reference_single = reference_module.SingleConditioning( + sigma_data=16.0, + token_s=4, + dim_fourier=6, + num_transitions=2, + ).to(device) + _load_local_state(reference_single, local_single) + generator = torch.Generator(device=device).manual_seed(401) + # s_trunk: (2, 5, 4) + s_trunk = torch.randn(2, 5, 4, device=device, generator=generator) + # s_inputs: (2, 5, 4) + s_inputs = torch.randn(2, 5, 4, device=device, generator=generator) + local_single_output = local_single(times, s_trunk, s_inputs) + reference_single_output = reference_single(times, s_trunk, s_inputs) + torch.testing.assert_close( + local_single_output[0], + reference_single_output[0], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + local_single_output[1], + reference_single_output[1], + rtol=0, + atol=0, + ) + + local_pair = encoders.PairwiseConditioning(6, 3, num_transitions=2).to(device) + reference_pair = reference_module.PairwiseConditioning( + 6, + 3, + num_transitions=2, + ).to(device) + _load_local_state(reference_pair, local_pair) + # pair_trunk: (2, 5, 5, 6) + pair_trunk = torch.randn(2, 5, 5, 6, device=device, generator=generator) + # relative_pair: (2, 5, 5, 3) + relative_pair = torch.randn(2, 5, 5, 3, device=device, generator=generator) + torch.testing.assert_close( + local_pair(pair_trunk, relative_pair), + reference_pair(pair_trunk, relative_pair), + rtol=0, + atol=0, + ) + + local_index = encoders.get_indexing_matrix(3, 4, 8, device) + reference_index = reference_module.get_indexing_matrix(3, 4, 8, device) + assert torch.equal(local_index, reference_index) + # single: (2, 12, 5) + single = torch.randn(2, 12, 5, device=device, generator=generator) + torch.testing.assert_close( + encoders.single_to_keys(single, local_index, 4, 8), + reference_module.single_to_keys(single, reference_index, 4, 8), + rtol=0, + atol=0, + ) + + +@pytest.mark.gpu +def test_atom_encoder_matches_upstream_on_gpu(upstream_package: ModuleType) -> None: + del upstream_package + assert torch.cuda.is_available(), "CUDA is required for Boltz atom parity" + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.modules.encodersv2") + device = torch.device("cuda") + features, _ = build_boltz2_features("ACDE") + features = _to_device(features, device) + kwargs = { + "atom_s": 8, + "atom_z": 6, + "token_s": 8, + "token_z": 6, + "atoms_per_window_queries": 32, + "atoms_per_window_keys": 128, + "atom_feature_dim": 388, + "structure_prediction": True, + } + local = encoders.AtomEncoder(**kwargs).to(device) + reference = reference_module.AtomEncoder(**kwargs).to(device) + _load_local_state(reference, local) + generator = torch.Generator(device=device).manual_seed(812) + num_tokens = features["token_index"].shape[1] + # s_trunk: (1, num_tokens, 8) + s_trunk = torch.randn(1, num_tokens, 8, device=device, generator=generator) + # z_trunk: (1, num_tokens, num_tokens, 6) + z_trunk = torch.randn( + 1, + num_tokens, + num_tokens, + 6, + device=device, + generator=generator, + ) + local_output = local(features, s_trunk=s_trunk, z=z_trunk) + reference_output = reference(features, s_trunk=s_trunk, z=z_trunk) + for local_tensor, reference_tensor in zip(local_output[:3], reference_output[:3], strict=True): + torch.testing.assert_close(local_tensor, reference_tensor, rtol=0, atol=0) + # probe: (1, features['ref_pos'].shape[1], 3) + probe = torch.randn( + 1, + features["ref_pos"].shape[1], + 3, + device=device, + generator=generator, + ) + torch.testing.assert_close( + local_output[3](probe), + reference_output[3](probe), + rtol=0, + atol=0, + ) + + +@pytest.mark.gpu +def test_atom_attention_encoder_and_decoder_match_upstream_on_gpu( + upstream_package: ModuleType, +) -> None: + del upstream_package + assert torch.cuda.is_available(), "CUDA is required for Boltz atom parity" + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.modules.encodersv2") + device = torch.device("cuda") + features, _ = build_boltz2_features("ACDE") + features = _to_device(features, device) + generator = torch.Generator(device=device).manual_seed(912) + num_atoms = features["ref_pos"].shape[1] + num_tokens = features["token_index"].shape[1] + + atom_kwargs = { + "atom_s": 8, + "atom_z": 2, + "token_s": 8, + "token_z": 6, + "atoms_per_window_queries": 32, + "atoms_per_window_keys": 128, + "atom_feature_dim": 388, + "structure_prediction": True, + } + atom_encoder = encoders.AtomEncoder(**atom_kwargs).to(device) + # s_trunk: (1, num_tokens, 8) + s_trunk = torch.randn(1, num_tokens, 8, device=device, generator=generator) + # z_trunk: (1, num_tokens, num_tokens, 6) + z_trunk = torch.randn( + 1, + num_tokens, + num_tokens, + 6, + device=device, + generator=generator, + ) + q, c, _, to_keys = atom_encoder(features, s_trunk=s_trunk, z=z_trunk) + # atom_bias: (1, num_atoms // 32, 32, 128, 2) + atom_bias = torch.randn( + 1, + num_atoms // 32, + 32, + 128, + 2, + device=device, + generator=generator, + ) + + encoder_kwargs = { + "atom_s": 8, + "token_s": 8, + "atoms_per_window_queries": 32, + "atoms_per_window_keys": 128, + "atom_encoder_depth": 1, + "atom_encoder_heads": 2, + "structure_prediction": True, + } + local_encoder = encoders.AtomAttentionEncoder(**encoder_kwargs).to(device) + reference_encoder = reference_module.AtomAttentionEncoder(**encoder_kwargs).to(device) + _load_local_state(reference_encoder, local_encoder) + # coordinates: (1, num_atoms, 3) + coordinates = torch.randn(1, num_atoms, 3, device=device, generator=generator) + local_encoded = local_encoder( + features, + q, + c, + atom_bias, + to_keys, + r=coordinates, + ) + reference_encoded = reference_encoder( + features, + q, + c, + atom_bias, + to_keys, + r=coordinates, + ) + for local_tensor, reference_tensor in zip( + local_encoded[:3], + reference_encoded[:3], + strict=True, + ): + torch.testing.assert_close(local_tensor, reference_tensor, rtol=0, atol=0) + + decoder_kwargs = { + "atom_s": 8, + "token_s": 8, + "attn_window_queries": 32, + "attn_window_keys": 128, + "atom_decoder_depth": 1, + "atom_decoder_heads": 2, + } + local_decoder = encoders.AtomAttentionDecoder(**decoder_kwargs).to(device) + reference_decoder = reference_module.AtomAttentionDecoder(**decoder_kwargs).to(device) + _load_local_state(reference_decoder, local_decoder) + # token_update: (1, num_tokens, 16) + token_update = torch.randn( + 1, + num_tokens, + 16, + device=device, + generator=generator, + ) + local_decoded = local_decoder( + token_update, + local_encoded[1], + local_encoded[2], + atom_bias, + features, + to_keys, + ) + reference_decoded = reference_decoder( + token_update, + reference_encoded[1], + reference_encoded[2], + atom_bias, + features, + to_keys, + ) + torch.testing.assert_close(local_decoded, reference_decoded, rtol=0, atol=0) + + +def _assert_tree_equal(local: object, reference: object) -> None: + if isinstance(reference, torch.Tensor): + assert isinstance(local, torch.Tensor) + assert local.dtype == reference.dtype + assert local.device == reference.device + torch.testing.assert_close(local, reference, rtol=0, atol=0, equal_nan=True) + return + if isinstance(reference, (tuple, list)): + assert isinstance(local, type(reference)) + assert len(local) == len(reference) + for local_item, reference_item in zip(local, reference, strict=True): + _assert_tree_equal(local_item, reference_item) + return + assert local == reference + + +@pytest.mark.gpu +def test_potential_geometry_and_gradients_match_upstream_on_gpu( + upstream_package: ModuleType, +) -> None: + del upstream_package + assert torch.cuda.is_available(), "CUDA is required for Boltz potential parity" + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.potentials.potentials") + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(1211) + # coordinates: (2, 9, 3) + coordinates = torch.randn(2, 9, 3, device=device, generator=generator) + + # pair_index: (2, 3) + pair_index = torch.tensor([[0, 2, 4], [1, 5, 8]], device=device) + for local, reference, is_distance in ( + ( + potentials.ConnectionsPotential(), + reference_module.ConnectionsPotential(), + True, + ), + ( + potentials.ChiralAtomPotential(), + reference_module.ChiralAtomPotential(), + False, + ), + ( + potentials.PlanarBondPotential(), + reference_module.PlanarBondPotential(), + False, + ), + ): + index = ( + pair_index + if is_distance + else torch.tensor( + [[0, 1], [1, 2], [3, 4], [5, 6]], + device=device, + ) + ) + _assert_tree_equal( + local.compute_variable(coordinates, index, compute_gradient=False), + reference.compute_variable(coordinates, index, compute_gradient=False), + ) + _assert_tree_equal( + local.compute_variable(coordinates, index, compute_gradient=True), + reference.compute_variable(coordinates, index, compute_gradient=True), + ) + + # values: (1, 3) + values = torch.tensor([[-2.0, 0.5, 4.0]], device=device) + # k: (3,) + k = torch.tensor([1.0, 2.0, 3.0], device=device) + # lower: (3,) + lower = torch.tensor([-1.0, 0.0, 1.0], device=device) + # upper: (3,) + upper = torch.tensor([1.0, 2.0, 3.0], device=device) + _assert_tree_equal( + potentials.FlatBottomPotential.compute_function( + object(), + values, + k, + lower, + upper, + compute_derivative=True, + ), + reference_module.FlatBottomPotential.compute_function( + object(), + values, + k, + lower, + upper, + compute_derivative=True, + ), + ) + + connection_features = {"connected_atom_index": pair_index[None]} + parameters = {"buffer": 2.5} + local_connection = potentials.ConnectionsPotential() + reference_connection = reference_module.ConnectionsPotential() + _assert_tree_equal( + local_connection.compute(coordinates, connection_features, parameters), + reference_connection.compute(coordinates, connection_features, parameters), + ) + _assert_tree_equal( + local_connection.compute_gradient( + coordinates, + connection_features, + parameters, + ), + reference_connection.compute_gradient( + coordinates, + connection_features, + parameters, + ), + ) + + +@pytest.mark.gpu +def test_potential_argument_builders_match_upstream_on_gpu( + upstream_package: ModuleType, +) -> None: + del upstream_package + assert torch.cuda.is_available(), "CUDA is required for Boltz potential parity" + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.potentials.potentials") + device = torch.device("cuda") + atom_to_token = torch.tensor( + [ + [1, 0, 0, 0], + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + [0, 0, 0, 1], + ], + dtype=torch.float32, + device=device, + )[None] + # ref_element: (1, 8, 128) + ref_element = torch.zeros(1, 8, 128, device=device) + ref_element[..., 6] = 1 + features = { + "atom_to_token": atom_to_token, + "asym_id": torch.tensor([[0, 0, 1, 1]], device=device), + "atom_pad_mask": torch.ones(1, 8, device=device), + "ref_element": ref_element, + "connected_chain_index": torch.tensor([[[0], [1]]], device=device), + "symmetric_chain_index": torch.tensor([[[0], [1]]], device=device), + "rdkit_bounds_index": torch.tensor( + [[[0, 2, 4], [1, 3, 5]]], + device=device, + ), + "rdkit_lower_bounds": torch.tensor([[1.0, 1.5, 2.0]], device=device), + "rdkit_upper_bounds": torch.tensor([[2.0, 2.5, 3.0]], device=device), + "rdkit_bounds_bond_mask": torch.tensor( + [[True, False, True]], + device=device, + ), + "rdkit_bounds_angle_mask": torch.tensor( + [[False, True, True]], + device=device, + ), + "stereo_bond_index": torch.tensor( + [[[[0, 1], [1, 2], [2, 3], [3, 4]]]], + device=device, + ).squeeze(1), + "stereo_bond_orientations": torch.tensor([[True, False]], device=device), + "chiral_atom_index": torch.tensor( + [[[[0, 1], [1, 2], [2, 3], [3, 4]]]], + device=device, + ).squeeze(1), + "chiral_atom_orientations": torch.tensor([[False, True]], device=device), + "planar_bond_index": torch.tensor( + [[[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6]]], + device=device, + ), + "contact_pair_index": torch.tensor( + [[[0, 2, 4], [1, 3, 5]]], + device=device, + ), + "contact_union_index": torch.tensor([[0, 0, 1]], device=device), + "contact_negation_mask": torch.tensor( + [[False, True, False]], + device=device, + ), + "contact_thresholds": torch.tensor([[3.0, 4.0, 5.0]], device=device), + } + cases = ( + ( + "PoseBustersPotential", + {"bond_buffer": 0.1, "angle_buffer": 0.2, "clash_buffer": 0.15}, + ), + ("VDWOverlapPotential", {"buffer": 0.225}), + ("SymmetricChainCOMPotential", {"buffer": 1.5}), + ("StereoBondPotential", {"buffer": 0.3}), + ("ChiralAtomPotential", {"buffer": 0.3}), + ("PlanarBondPotential", {"buffer": 0.2}), + ("ContactPotentital", {"union_lambda": 2.0}), + ) + for class_name, parameters in cases: + local = getattr(potentials, class_name)() + reference = getattr(reference_module, class_name)() + _assert_tree_equal( + local.compute_args(features, parameters), + reference.compute_args(features, parameters), + ) + + +def test_potential_stack_and_schedules_match_upstream( + upstream_package: ModuleType, +) -> None: + del upstream_package + _install_import_only_dependency_stubs() + reference_module = importlib.import_module("boltz.model.potentials.potentials") + steering = { + "fk_steering": True, + "physical_guidance_update": True, + "contact_guidance_update": True, + } + local_stack = potentials.get_potentials(steering, boltz2=True) + reference_stack = reference_module.get_potentials(steering, boltz2=True) + assert [type(item).__name__ for item in local_stack] == [ + type(item).__name__ for item in reference_stack + ] + for local, reference in zip(local_stack, reference_stack, strict=True): + for time in (0.0, 0.5, 1.0): + assert local.compute_parameters(time) == reference.compute_parameters(time) diff --git a/tests/parity/test_e1_source_independence_parity.py b/tests/parity/test_e1_source_independence_parity.py new file mode 100644 index 0000000..ceff28a --- /dev/null +++ b/tests/parity/test_e1_source_independence_parity.py @@ -0,0 +1,372 @@ +"""Exact parity tests for independently maintained E1 runtime contracts.""" + +from __future__ import annotations + +import importlib +import sys +import pytest +import torch +import torch.nn as nn +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from types import ModuleType +from typing import Any, cast + +from fastplms.models.e1.modeling_e1 import ( + FAST_E1_ENCODER, + AttentionArgs, + DataPrepConfig, + DecoderLayer, + E1BatchPreparer, + E1Config, + E1PreTrainedModel, + RMSNorm, + _get_unpad_data, + build_block_causal_mask_4d, + build_within_seq_mask_4d, + create_within_seq_block_mask, + direct_block_mask, + get_overlapping_blocks, + get_tokenizer, +) + + +ROOT = Path(__file__).resolve().parents[2] +UPSTREAM_SRC = ROOT / "vendor/upstream/e1/src" + + +@contextmanager +def _official_e1_modules() -> Iterator[tuple[ModuleType, ModuleType, ModuleType]]: + assert (UPSTREAM_SRC / "E1").is_dir(), "pinned E1 submodule is missing" + previous_path = list(sys.path) + previous_modules = { + name: module + for name, module in sys.modules.items() + if name == "E1" or name.startswith("E1.") + } + for name in previous_modules: + sys.modules.pop(name, None) + sys.path.insert(0, str(UPSTREAM_SRC)) + try: + yield ( + importlib.import_module("E1.batch_preparer"), + importlib.import_module("E1.model.varlen_flex_attention"), + importlib.import_module("E1.model.flash_attention_utils"), + ) + finally: + sys.path[:] = previous_path + for name in tuple(sys.modules): + if name == "E1" or name.startswith("E1."): + sys.modules.pop(name, None) + sys.modules.update(previous_modules) + + +@pytest.fixture(scope="module") +def official_e1() -> Iterator[tuple[ModuleType, ModuleType, ModuleType]]: + with _official_e1_modules() as modules: + yield modules + + +def _assert_encoding_equal(candidate: dict[str, Any], official: dict[str, Any]) -> None: + assert candidate.keys() == official.keys() + for key in candidate: + if isinstance(candidate[key], torch.Tensor): + assert torch.equal(candidate[key], official[key]), key + else: + assert candidate[key] == official[key], key + + +@pytest.mark.parametrize("remove_x_tokens", (False, True)) +@pytest.mark.parametrize("preserve_context_labels", (False, True)) +def test_e1_sequence_preparation_is_exact( + official_e1: tuple[ModuleType, ModuleType, ModuleType], + remove_x_tokens: bool, + preserve_context_labels: bool, +) -> None: + official_batch, _, _ = official_e1 + tokenizer = get_tokenizer() + local = E1BatchPreparer( + data_prep_config=DataPrepConfig(remove_X_tokens=remove_x_tokens), + tokenizer=tokenizer, + preserve_context_labels=preserve_context_labels, + ) + official = official_batch.E1BatchPreparer( + data_prep_config=official_batch.DataPrepConfig(remove_X_tokens=remove_x_tokens), + tokenizer=tokenizer, + preserve_context_labels=preserve_context_labels, + device=torch.device("cpu"), + ) + + sequences = ("ACD?X", "ACDX,EF?G", "XACD,XXE") + for sequence in sequences: + _assert_encoding_equal( + local.prepare_multiseq(sequence), + official.prepare_multiseq(sequence), + ) + _assert_encoding_equal( + local.get_batch_kwargs(list(sequences)), + official.get_batch_kwargs(list(sequences)), + ) + + local_single = local.prepare_singleseq("AX?D") + official_single = official.prepare_singleseq("AX?D") + _assert_encoding_equal(local_single, official_single) + assert local_single["input_ids"].data_ptr() == local_single["labels"].data_ptr() + assert official_single["input_ids"].data_ptr() == official_single["labels"].data_ptr() + + +@pytest.mark.parametrize( + ("sequence", "max_sequences", "max_positions"), + ( + ("AcD", 4, 16), + ("ABCDE", 4, 4), + ("A,C,D", 2, 16), + ), +) +def test_e1_sequence_preparation_errors_are_exact( + official_e1: tuple[ModuleType, ModuleType, ModuleType], + sequence: str, + max_sequences: int, + max_positions: int, +) -> None: + official_batch, _, _ = official_e1 + tokenizer = get_tokenizer() + local = E1BatchPreparer( + data_prep_config=DataPrepConfig( + max_num_sequences=max_sequences, + max_num_positions_within_seq=max_positions, + ), + tokenizer=tokenizer, + ) + official = official_batch.E1BatchPreparer( + data_prep_config=official_batch.DataPrepConfig( + max_num_sequences=max_sequences, + max_num_positions_within_seq=max_positions, + ), + tokenizer=tokenizer, + device=torch.device("cpu"), + ) + + with pytest.raises(ValueError) as local_error: + local.prepare_multiseq(sequence) + with pytest.raises(ValueError) as official_error: + official.prepare_multiseq(sequence) + assert str(local_error.value) == str(official_error.value) + + +@pytest.mark.parametrize( + "sequence_ids", + ( + torch.tensor([[0, 0, 1, 1, -1], [0, 0, 0, -1, -1]]), + torch.tensor([[3, -1, 3, 4, -1], [7, 8, 8, 8, 9]]), + torch.tensor([[0, 0, 0, 0], [0, 1, 2, 3]]), + ), +) +def test_e1_unpadding_metadata_is_exact( + official_e1: tuple[ModuleType, ModuleType, ModuleType], + sequence_ids: torch.Tensor, +) -> None: + # sequence_ids: (b, l) + _, _, official_flash = official_e1 + candidate = _get_unpad_data(sequence_ids) + official = official_flash._get_unpad_data(sequence_ids) + assert torch.equal(candidate[0], official[0]) + assert torch.equal(candidate[1], official[1]) + assert candidate[2] == official[2] + + +@pytest.mark.parametrize( + ("q_lengths", "k_lengths"), + ( + (torch.tensor([64, 128, 65]), torch.tensor([127, 2, 128])), + (torch.tensor([1, 255, 1]), torch.tensor([129, 128])), + (torch.tensor([128, 128]), torch.tensor([256])), + ), +) +def test_e1_block_classification_and_mask_are_exact( + official_e1: tuple[ModuleType, ModuleType, ModuleType], + q_lengths: torch.Tensor, + k_lengths: torch.Tensor, +) -> None: + # q_lengths: (...), k_lengths: (...) + _, official_flex, _ = official_e1 + candidate_full, candidate_partial = get_overlapping_blocks(q_lengths, k_lengths) + official_full, official_partial = official_flex.get_overlapping_blocks(q_lengths, k_lengths) + assert torch.equal(candidate_full, official_full) + assert torch.equal(candidate_partial, official_partial) + + candidate_mask = direct_block_mask(q_lengths, k_lengths) + official_mask = official_flex.direct_block_mask(q_lengths, k_lengths) + assert candidate_mask.shape == official_mask.shape + assert torch.equal(candidate_mask.to_dense(), official_mask.to_dense()) + for attribute in ( + "kv_num_blocks", + "kv_indices", + "full_kv_num_blocks", + "full_kv_indices", + "q_num_blocks", + "q_indices", + "full_q_num_blocks", + "full_q_indices", + ): + candidate_value = getattr(candidate_mask, attribute) + official_value = getattr(official_mask, attribute) + if candidate_value is None or official_value is None: + assert candidate_value is official_value + else: + assert torch.equal(candidate_value, official_value), attribute + + q_index = torch.arange(int(q_lengths.sum().item()))[:, None] + k_index = torch.arange(int(k_lengths.sum().item()))[None, :] + # zero: (...) + zero = torch.tensor(0) + assert torch.equal( + candidate_mask.mask_mod(zero, zero, q_index, k_index), + official_mask.mask_mod(zero, zero, q_index, k_index), + ) + + +def _tiny_config(attn_backend: str = "sdpa") -> E1Config: + return E1Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_num_sequences=8, + max_num_positions_within_seq=64, + max_num_positions_global=256, + dtype="float32", + attn_backend=attn_backend, + ) + + +def test_e1_parameter_initialization_is_exact() -> None: + model = FAST_E1_ENCODER(_tiny_config()).eval() # type: ignore[no-untyped-call] + + actual_linear = nn.Linear(9, 7) + expected_linear = nn.Linear(9, 7) + torch.manual_seed(41) + expected_linear.weight.data.normal_(mean=0.0, std=model.config.initializer_range) + expected_linear.bias.data.zero_() + torch.manual_seed(41) + model._init_weights(actual_linear) + assert torch.equal(actual_linear.weight, expected_linear.weight) + assert torch.equal(actual_linear.bias, expected_linear.bias) + + actual_embedding = nn.Embedding(11, 5, padding_idx=3) + expected_embedding = nn.Embedding(11, 5, padding_idx=3) + torch.manual_seed(43) + expected_embedding.weight.data.normal_(mean=0.0, std=model.config.initializer_range) + expected_embedding.weight.data[3].zero_() + torch.manual_seed(43) + model._init_weights(actual_embedding) + assert torch.equal(actual_embedding.weight, expected_embedding.weight) + + norm = RMSNorm(13) + norm.weight.data.fill_(7.0) + model._init_weights(norm) + assert torch.equal(norm.weight, torch.ones_like(norm.weight)) + + +def test_e1_keeps_flash_flags_disabled() -> None: + assert E1PreTrainedModel._supports_flash_attn_2 is False + assert E1PreTrainedModel._supports_flash_attn_3 is False + assert E1PreTrainedModel._fastplms_attention_implementations == ( + "eager", + "sdpa", + "flex_attention", + ) + + +def _manual_encoder_forward( + model: FAST_E1_ENCODER, + batch: dict[str, torch.Tensor], +) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: + sequence_ids = batch["sequence_ids"] + hidden_states = model.embed_tokens(batch["input_ids"]) + # hidden_states: (..., d) + hidden_states = hidden_states + model.embed_seq_id(sequence_ids.clamp_min(0)) + first_layer = cast(DecoderLayer, model.layers[0]) + # hidden_states: (..., d) + hidden_states = hidden_states.to(first_layer.norm_attn_norm.self_attn.q_proj.weight.dtype) + + use_flex = model._attn_backend.value == "flex_attention" + use_dense = model._attn_backend.value in {"eager", "sdpa"} + attention_args = AttentionArgs( + block_causal_block_mask=None, + within_seq_block_mask=(create_within_seq_block_mask(sequence_ids) if use_flex else None), + within_seq_mask_4d=(build_within_seq_mask_4d(sequence_ids) if use_dense else None), + block_causal_mask_4d=(build_block_causal_mask_4d(sequence_ids) if use_dense else None), + ) + hidden_history = [] + for raw_layer in model.layers: + layer = cast(DecoderLayer, raw_layer) + hidden_history.append(hidden_states) + hidden_states, _, _, _ = layer( + hidden_states, + within_seq_position_ids=batch["within_seq_position_ids"], + global_position_ids=batch["global_position_ids"], + sequence_ids=sequence_ids, + attention_args=attention_args, + ) + hidden_states = model.norm(hidden_states) + hidden_history.append(hidden_states) + return hidden_states, tuple(hidden_history) + + +@pytest.mark.gpu +@pytest.mark.parametrize("attn_backend", ("eager", "sdpa", "flex_attention")) +def test_e1_refactored_forward_is_exact_on_h100(attn_backend: str) -> None: + assert torch.cuda.is_available(), "E1 BF16 forward parity requires CUDA" + torch.manual_seed(47) + model = ( + FAST_E1_ENCODER(_tiny_config(attn_backend)) + .eval() + .to( # type: ignore[no-untyped-call] + "cuda", + dtype=torch.bfloat16, + ) + ) + prepared = E1BatchPreparer(tokenizer=get_tokenizer()).get_batch_kwargs( + ["ACDEFG", "ACD,EFGH"], + device=torch.device("cuda"), + ) + batch: dict[str, torch.Tensor] = {} + for key in ( + "input_ids", + "within_seq_position_ids", + "global_position_ids", + "sequence_ids", + ): + value = prepared[key] + assert isinstance(value, torch.Tensor) + batch[key] = value + before = {name: tensor.detach().clone() for name, tensor in model.state_dict().items()} + + with torch.no_grad(): + expected_last, expected_history = _manual_encoder_forward(model, batch) + output = model(**batch, output_hidden_states=True) + combined_embeddings = model.embed_tokens(batch["input_ids"]) + # combined_embeddings: (...) + combined_embeddings = combined_embeddings + model.embed_seq_id( + batch["sequence_ids"].clamp_min(0) + ) + soft_output = model( + inputs_embeds=combined_embeddings, + within_seq_position_ids=batch["within_seq_position_ids"], + global_position_ids=batch["global_position_ids"], + sequence_ids=batch["sequence_ids"], + output_hidden_states=True, + ) + + assert torch.equal(output.last_hidden_state, expected_last) + assert output.hidden_states is not None + assert len(output.hidden_states) == len(expected_history) + for actual, expected in zip(output.hidden_states, expected_history, strict=True): + assert torch.equal(actual, expected) + assert torch.equal(soft_output.last_hidden_state, output.last_hidden_state) + assert model.state_dict().keys() == before.keys() + for name, tensor in model.state_dict().items(): + assert torch.equal(tensor, before[name]), name diff --git a/tests/parity/test_esmfold2_common_parity.py b/tests/parity/test_esmfold2_common_parity.py new file mode 100644 index 0000000..a02b6e7 --- /dev/null +++ b/tests/parity/test_esmfold2_common_parity.py @@ -0,0 +1,359 @@ +"""Exact differentials for the independently organized ESMFold2 core blocks.""" + +from __future__ import annotations + +import importlib.util +import inspect +import sys +import types +import pytest +import torch +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any +from torch import Tensor, nn + +from fastplms.models.esmfold2 import modeling_esmfold2_common as local + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +OFFICIAL_ROOT = ROOT / "vendor/upstream/biohub-transformers/src/transformers/models/esmfold2" +_MISSING = object() + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source(name: str, path: Path, aliases: dict[str, types.ModuleType]) -> types.ModuleType: + assert path.is_file(), f"pinned source is missing: {path}" + specification = importlib.util.spec_from_file_location(name, path) + assert specification is not None and specification.loader is not None + module = importlib.util.module_from_spec(specification) + with _temporary_modules({**aliases, name: module}): + specification.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def official() -> types.ModuleType: + import transformers.configuration_utils as configuration_utils + + root_name = "_fastplms_pinned_esmfold2_common" + aliases = { + root_name: _package(root_name), + f"{root_name}.models": _package(f"{root_name}.models"), + f"{root_name}.models.esmfold2": _package(f"{root_name}.models.esmfold2"), + f"{root_name}.configuration_utils": configuration_utils, + } + config_name = f"{root_name}.models.esmfold2.configuration_esmfold2" + config = _load_source(config_name, OFFICIAL_ROOT / "configuration_esmfold2.py", aliases) + aliases[config_name] = config + return _load_source( + f"{root_name}.models.esmfold2.modeling_esmfold2_common", + OFFICIAL_ROOT / "modeling_esmfold2_common.py", + aliases, + ) + + +def _assert_tensor_exact(actual: Tensor, expected: Tensor) -> None: + # actual: (...), expected: (...) + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + assert actual.device == expected.device + assert torch.equal(actual, expected) + + +def _alias_groups(module: nn.Module) -> set[tuple[str, ...]]: + names_by_id: dict[int, list[str]] = {} + for name, parameter in module.named_parameters(remove_duplicate=False): + names_by_id.setdefault(id(parameter), []).append(name) + return {tuple(names) for names in names_by_id.values() if len(names) > 1} + + +_STATE_SPECS: tuple[tuple[str, tuple[Any, ...], dict[str, Any]], ...] = ( + ("TransitionLayer", (8, 2), {}), + ("AdaptiveLayerNorm", (8, 6), {}), + ("FourierEmbedding", (8,), {}), + ("SwiGLUMLP", (8, 2), {}), + ("SWA3DRoPEAttention", (32, 4), {"half_window": 2}), + ("AttentionPairBias", (8, 6, 2), {"d_cond": 8}), + ("ConditionedTransitionBlock", (8,), {"d_cond": 8}), + ("DiffusionTransformer", (8, 6, 2, 2), {"d_cond": 8}), + ("RowAttentionPooling", (6, 8), {}), + ("ResIdxAsymIdSymIdEntityIdEncoding", (2, 1, 8), {}), + ("SingleToPair", (8, 4, 6), {}), + ("LanguageModelShim", (8, 12, 2), {}), + ("TriangleMultiplicativeBlock", (8, 4, "outgoing"), {}), + ("TriangleMultiplicativeUpdate", (8, True), {}), + ("Transition", (8, 2), {}), + ("PairUpdateBlock", (8, 2), {}), + ("FoldingTrunk", (2, 8, 2), {}), + ("OuterProductMean", (6, 4, 8), {}), + ("MSAPairWeightedAveraging", (6, 8, 2, 4), {}), +) + + +@pytest.mark.parametrize(("name", "args", "kwargs"), _STATE_SPECS, ids=lambda x: x) +def test_state_schema_and_aliases_match_pinned_biohub( + official: types.ModuleType, + name: str, + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> None: + torch.manual_seed(1907) + expected = getattr(official, name)(*args, **kwargs) + torch.manual_seed(1907) + actual = getattr(local, name)(*args, **kwargs) + actual_state = actual.state_dict() + expected_state = expected.state_dict() + assert tuple(actual_state) == tuple(expected_state) + for key in expected_state: + _assert_tensor_exact(actual_state[key], expected_state[key]) + assert _alias_groups(actual) == _alias_groups(expected) + + +def _paired_modules( + official: types.ModuleType, + name: str, + *args: Any, + **kwargs: Any, +) -> tuple[nn.Module, nn.Module]: + # expected: (...) + expected = getattr(official, name)(*args, **kwargs).eval().cuda() + # actual: (...) + actual = getattr(local, name)(*args, **kwargs).eval().cuda() + actual.load_state_dict(expected.state_dict(), strict=True) + return actual, expected + + +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) +def test_projection_and_pair_blocks_match_exactly( + official: types.ModuleType, dtype: torch.dtype +) -> None: + assert torch.cuda.is_available(), "the pinned common-module differential needs CUDA" + torch.manual_seed(8675309) + with ( + torch.no_grad(), + torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=dtype == torch.bfloat16), + ): + actual_lm, expected_lm = _paired_modules(official, "LanguageModelShim", 8, 12, 2) + actual_lm.to(dtype=dtype) + expected_lm.to(dtype=dtype) + # hidden_states: (2, 5, 3, 12) + hidden_states = torch.randn(2, 5, 3, 12, device="cuda", dtype=dtype) + _assert_tensor_exact(actual_lm(hidden_states), expected_lm(hidden_states)) + sequence_projection = actual_lm.project_sequence(hidden_states) + # expected_projection: (...) + expected_projection = ( + expected_lm.base_z_combine.softmax(0) @ expected_lm.base_z_linear(hidden_states) + ).squeeze(-2) + _assert_tensor_exact(sequence_projection, expected_projection) + + actual_triangle, expected_triangle = _paired_modules( + official, "TriangleMultiplicativeBlock", 8, 4, "outgoing" + ) + actual_triangle.to(dtype=dtype) + expected_triangle.to(dtype=dtype) + actual_triangle.set_chunk_size(None) + expected_triangle.set_chunk_size(None) + # pair: (1, 4, 4, 8) + pair = torch.randn(1, 4, 4, 8, device="cuda", dtype=dtype) + # pair_mask: (1, 4, 4) + pair_mask = torch.tensor( + [[[1, 1, 1, 0], [1, 1, 1, 0], [1, 1, 1, 0], [0, 0, 0, 0]]], + device="cuda", + dtype=torch.bool, + ) + _assert_tensor_exact(actual_triangle(pair, pair_mask), expected_triangle(pair, pair_mask)) + + actual_trunk, expected_trunk = _paired_modules(official, "FoldingTrunk", 2, 8, 2) + actual_trunk.to(dtype=dtype) + expected_trunk.to(dtype=dtype) + actual_trunk.set_chunk_size(None) + expected_trunk.set_chunk_size(None) + _assert_tensor_exact(actual_trunk(pair, pair_mask), expected_trunk(pair, pair_mask)) + + +def test_attention_and_diffusion_blocks_match_exactly( + official: types.ModuleType, +) -> None: + torch.manual_seed(4401) + with torch.no_grad(): + actual, expected = _paired_modules(official, "DiffusionTransformer", 8, 6, 2, 2, d_cond=8) + # token: (2, 5, 8) + token = torch.randn(2, 5, 8, device="cuda") + # condition: (2, 5, 8) + condition = torch.randn(2, 5, 8, device="cuda") + # pair: (2, 5, 5, 6) + pair = torch.randn(2, 5, 5, 6, device="cuda") + # mask: (2, 5) + mask = torch.tensor( + [[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], + device="cuda", + dtype=torch.bool, + ) + actual_out, actual_intermediates = actual( + token, condition, pair, attention_mask=mask, return_intermediates=True + ) + expected_out, expected_intermediates = expected( + token, condition, pair, attention_mask=mask, return_intermediates=True + ) + _assert_tensor_exact(actual_out, expected_out) + assert len(actual_intermediates) == len(expected_intermediates) + for actual_value, expected_value in zip( + actual_intermediates, expected_intermediates, strict=True + ): + _assert_tensor_exact(actual_value, expected_value) + + +def test_input_pair_and_msa_blocks_match_exactly( + official: types.ModuleType, +) -> None: + torch.manual_seed(923) + with torch.no_grad(): + actual_relative, expected_relative = _paired_modules( + official, "ResIdxAsymIdSymIdEntityIdEncoding", 2, 1, 8 + ) + # residue_index: (1, 4) + residue_index = torch.tensor([[0, 1, 2, 0]], device="cuda") + # asym_id: (1, 4) + asym_id = torch.tensor([[0, 0, 0, 1]], device="cuda") + # sym_id: (1, 4) + sym_id = torch.tensor([[0, 0, 0, 1]], device="cuda") + # entity_id: (1, 4) + entity_id = torch.tensor([[0, 0, 0, 0]], device="cuda") + # token_index: (1, 4) + token_index = torch.tensor([[0, 1, 2, 0]], device="cuda") + inputs = (residue_index, asym_id, sym_id, entity_id, token_index) + _assert_tensor_exact(actual_relative(*inputs), expected_relative(*inputs)) + + actual_opm, expected_opm = _paired_modules(official, "OuterProductMean", 6, 4, 8) + actual_opm.set_chunk_size(None) + expected_opm.set_chunk_size(None) + # msa: (2, 4, 3, 6) + msa = torch.randn(2, 4, 3, 6, device="cuda") + # msa_mask: (2, 4, 3) + msa_mask = torch.tensor( + [ + [[1, 1, 1], [1, 1, 0], [1, 1, 1], [1, 0, 0]], + [[1, 1, 1], [1, 1, 1], [1, 0, 0], [1, 1, 0]], + ], + device="cuda", + dtype=torch.bool, + ) + _assert_tensor_exact(actual_opm(msa, msa_mask), expected_opm(msa, msa_mask)) + + actual_avg, expected_avg = _paired_modules(official, "MSAPairWeightedAveraging", 6, 8, 2, 4) + # pair: (2, 4, 4, 8) + pair = torch.randn(2, 4, 4, 8, device="cuda") + # pair_mask: (2, 4, 4) + pair_mask = torch.ones(2, 4, 4, device="cuda", dtype=torch.bool) + _assert_tensor_exact(actual_avg(msa, pair, pair_mask), expected_avg(msa, pair, pair_mask)) + + +def test_tensor_utilities_match_pinned_biohub(official: types.ModuleType) -> None: + # token: (2, 4, 6) + token = torch.arange(48, device="cuda", dtype=torch.float32).reshape(2, 4, 6) + # atom_to_token: (2, 5) + atom_to_token = torch.tensor([[0, 0, 2, 3, 3], [0, 1, 1, 2, 3]], device="cuda") + for name, args in ( + ("gather_token_to_atom", (token, atom_to_token)), + ( + "scatter_atom_to_token", + ( + torch.randn(2, 5, 6, device="cuda"), + atom_to_token, + 4, + torch.tensor( + [[1, 1, 1, 1, 0], [1, 1, 0, 1, 1]], + device="cuda", + dtype=torch.bool, + ), + ), + ), + ( + "gather_rep_atom_coords", + ( + torch.randn(2, 5, 3, device="cuda"), + torch.tensor([[0, 2, 4], [1, 3, 4]], device="cuda"), + ), + ), + ): + _assert_tensor_exact(getattr(local, name)(*args), getattr(official, name)(*args)) + + +def test_atom_rope_attention_and_rigid_alignment_match_exactly( + official: types.ModuleType, +) -> None: + torch.manual_seed(314159) + # ref_pos: (2, 7, 3) + ref_pos = torch.randn(2, 7, 3, device="cuda") + # space_uid: (2, 7) + space_uid = torch.tensor([[0, 0, 1, 1, 2, 2, 2], [3, 3, 3, 4, 4, 5, 5]], device="cuda") + actual_cos, actual_sin = local.build_3d_rope( + ref_pos, space_uid, head_dim=16, n_spatial_per_axis=2, n_uid_pairs=1 + ) + expected_cos, expected_sin = official.build_3d_rope( + ref_pos, space_uid, head_dim=16, n_spatial_per_axis=2, n_uid_pairs=1 + ) + _assert_tensor_exact(actual_cos, expected_cos) + _assert_tensor_exact(actual_sin, expected_sin) + + with torch.no_grad(): + actual_attention, expected_attention = _paired_modules( + official, "SWA3DRoPEAttention", 64, 4, half_window=2 + ) + # atom: (2, 7, 64) + atom = torch.randn(2, 7, 64, device="cuda") + _assert_tensor_exact( + actual_attention(atom, (actual_cos, actual_sin)), + expected_attention(atom, (expected_cos, expected_sin)), + ) + + # mobile: (2, 7, 3) + mobile = torch.randn(2, 7, 3, device="cuda") + # target: (2, 7, 3) + target = torch.randn(2, 7, 3, device="cuda") + # weights: (2, 7) + weights = torch.rand(2, 7, device="cuda") + # mask: (2, 7) + mask = torch.tensor( + [[1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1]], + device="cuda", + dtype=torch.bool, + ) + _assert_tensor_exact( + local.DiffusionStructureHead._weighted_rigid_align(mobile, target, weights, mask), + official.DiffusionStructureHead._weighted_rigid_align(mobile, target, weights, mask), + ) + + +def test_shared_public_class_surface_is_preserved(official: types.ModuleType) -> None: + def classes(module: types.ModuleType) -> set[str]: + return { + name + for name, value in vars(module).items() + if inspect.isclass(value) and value.__module__ == module.__name__ + } + + assert classes(local) == classes(official) diff --git a/tests/parity/test_esmfold2_protein_data_parity.py b/tests/parity/test_esmfold2_protein_data_parity.py new file mode 100644 index 0000000..6f27b4f --- /dev/null +++ b/tests/parity/test_esmfold2_protein_data_parity.py @@ -0,0 +1,381 @@ +"""Pinned Biohub parity for ESMFold2 protein-chain and complex data APIs.""" + +from __future__ import annotations + +import importlib.util +import inspect +import io +import sys +import types +import numpy as np +import pytest +import torch +from collections.abc import Iterator +from contextlib import contextmanager +from functools import cache +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from fastplms.models.esmfold2 import esmfold2_affine3d as local_affine +from fastplms.models.esmfold2 import esmfold2_aligner as local_aligner +from fastplms.models.esmfold2 import esmfold2_atom_indexer as local_atom_indexer +from fastplms.models.esmfold2 import esmfold2_metrics as local_metrics +from fastplms.models.esmfold2 import esmfold2_misc as local_misc +from fastplms.models.esmfold2 import esmfold2_mmcif_parsing as local_mmcif +from fastplms.models.esmfold2 import ( + esmfold2_normalize_coordinates as local_normalize, +) +from fastplms.models.esmfold2 import esmfold2_protein_chain as local_chain +from fastplms.models.esmfold2 import esmfold2_protein_complex as local_complex +from fastplms.models.esmfold2 import esmfold2_protein_structure as local_structure +from fastplms.models.esmfold2 import esmfold2_residue_constants as local_residues +from fastplms.models.esmfold2 import esmfold2_utils_types as local_types + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +BIOHUB_ESM = ROOT / "vendor/upstream/biohub-esm/esm" +_MISSING = object() +SOURCE_PAIRS = { + "esmfold2_protein_chain.py": BIOHUB_ESM / "utils/structure/protein_chain.py", + "esmfold2_protein_complex.py": BIOHUB_ESM / "utils/structure/protein_complex.py", +} + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source( + module_name: str, + path: Path, + aliases: dict[str, types.ModuleType], +) -> types.ModuleType: + assert path.is_file(), f"pinned source is missing: {path}" + specification = importlib.util.spec_from_file_location(module_name, path) + assert specification is not None and specification.loader is not None + module = importlib.util.module_from_spec(specification) + with _temporary_modules({**aliases, module_name: module}): + specification.loader.exec_module(module) + return module + + +def _base_aliases() -> dict[str, types.ModuleType]: + return { + "esm": _package("esm"), + "esm.utils": _package("esm.utils"), + "esm.utils.structure": _package("esm.utils.structure"), + "esm.utils.residue_constants": local_residues, + "esm.utils.misc": local_misc, + "esm.utils.structure.affine3d": local_affine, + "esm.utils.structure.aligner": local_aligner, + "esm.utils.structure.atom_indexer": local_atom_indexer, + "esm.utils.structure.metrics": local_metrics, + "esm.utils.structure.mmcif_parsing": local_mmcif, + "esm.utils.structure.normalize_coordinates": local_normalize, + "esm.utils.structure.protein_structure": local_structure, + "esm.utils.types": local_types, + } + + +@cache +def _official_chain() -> types.ModuleType: + return _load_source( + "_fastplms_pinned_biohub_protein_chain", + SOURCE_PAIRS["esmfold2_protein_chain.py"], + _base_aliases(), + ) + + +@cache +def _official_complex() -> types.ModuleType: + official_chain = _official_chain() + return _load_source( + "_fastplms_pinned_biohub_protein_complex", + SOURCE_PAIRS["esmfold2_protein_complex.py"], + { + **_base_aliases(), + "esm.utils.structure.protein_chain": official_chain, + }, + ) + + +def _assert_equal(actual: Any, expected: Any) -> None: + if isinstance(actual, np.ndarray): + assert isinstance(expected, np.ndarray) + np.testing.assert_array_equal(actual, expected) + elif isinstance(actual, torch.Tensor): + assert isinstance(expected, torch.Tensor) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + elif isinstance(actual, dict): + assert actual.keys() == expected.keys() + for key in actual: + _assert_equal(actual[key], expected[key]) + elif isinstance(actual, (list, tuple)): + assert type(actual) is type(expected) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_equal(actual_item, expected_item) + else: + assert actual == expected + + +def _assert_chain_equal(actual: Any, expected: Any) -> None: + for field_name in ( + "id", + "sequence", + "chain_id", + "entity_id", + "residue_index", + "insertion_code", + "atom37_positions", + "atom37_mask", + "confidence", + "atom37_confidence", + ): + _assert_equal(getattr(actual, field_name), getattr(expected, field_name)) + + +def _assert_complex_equal(actual: Any, expected: Any) -> None: + for field_name in ( + "id", + "sequence", + "entity_id", + "chain_id", + "sym_id", + "residue_index", + "insertion_code", + "atom37_positions", + "atom37_mask", + "confidence", + "atom37_confidence", + ): + _assert_equal(getattr(actual, field_name), getattr(expected, field_name)) + _assert_equal(actual.metadata.entity_lookup, expected.metadata.entity_lookup) + _assert_equal(actual.metadata.chain_lookup, expected.metadata.chain_lookup) + _assert_equal( + actual.metadata.assembly_composition, + expected.metadata.assembly_composition, + ) + + +def _atom37_coordinates(offset: float = 0.0) -> np.ndarray: + # X: (4, 37, 3) + X = np.full((4, 37, 3), np.nan, dtype=np.float32) + for residue_index in range(4): + # x: (...) + x = offset + 3.8 * residue_index + atoms = { + "N": (x, 0.1, 0.0), + "CA": (x + 1.4, 0.2, 0.1), + "C": (x + 2.4, 1.1, 0.2), + "O": (x + 2.1, 2.2, 0.3), + "CB": (x + 1.5, -0.8, 1.0), + } + for atom_name, coordinate in atoms.items(): + X[residue_index, local_residues.atom_order[atom_name]] = coordinate + return X + + +def _make_chain(module: types.ModuleType, *, chain_id: str, entity_id: int, offset: float): + return module.ProteinChain.from_atom37( + _atom37_coordinates(offset), + id="fixture", + sequence="AGST", + chain_id=chain_id, + entity_id=entity_id, + residue_index=np.asarray([4, 5, 5, 9], dtype=np.int64), + insertion_code=np.asarray(["", "", "A", ""]), + confidence=np.asarray([0.91, 0.83, 0.72, 0.65], dtype=np.float32), + ) + + +def test_protein_data_source_inventory_is_pinned() -> None: + runtime = ROOT / "src/fastplms/models/esmfold2" + for runtime_name, upstream_path in SOURCE_PAIRS.items(): + assert (runtime / runtime_name).is_file() + assert upstream_path.is_file() + gitmodules = (ROOT / ".gitmodules").read_text(encoding="utf-8") + assert "vendor/upstream/biohub-esm" in gitmodules + + +def test_public_class_surfaces_match_pinned_biohub() -> None: + official_chain = _official_chain() + official_complex = _official_complex() + for actual, expected in ( + (local_chain.ProteinChain, official_chain.ProteinChain), + (local_complex.ProteinComplex, official_complex.ProteinComplex), + ): + actual_names = { + name + for name, value in inspect.getmembers(actual) + if not name.startswith("__") and (callable(value) or inspect.isdatadescriptor(value)) + } + expected_names = { + name + for name, value in inspect.getmembers(expected) + if not name.startswith("__") and (callable(value) or inspect.isdatadescriptor(value)) + } + assert actual_names == expected_names + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + ("1", [("1",)]), + ("1-3", [("1",), ("2",), ("3",)]), + ("(1-2)(4,6)", [("4", "1"), ("4", "2"), ("6", "1"), ("6", "2")]), + ], +) +def test_assembly_operation_expansion_matches_pinned_biohub( + expression: str, expected: list[tuple[str, ...]] +) -> None: + official = _official_complex() + assert local_complex._parse_operation_expression(expression) == expected + assert official._parse_operation_expression(expression) == expected + + +def test_assembly_transform_application_matches_pinned_biohub() -> None: + official = _official_complex() + local_input = _make_chain(local_chain, chain_id="A", entity_id=1, offset=0.0) + official_input = _make_chain(_official_chain(), chain_id="A", entity_id=1, offset=0.0) + transforms = { + "1": SimpleNamespace(rotation=np.eye(3), target_translation=np.asarray([1.0, -2.0, 0.5])), + "2": SimpleNamespace( + rotation=np.diag([-1.0, 1.0, -1.0]), + target_translation=np.asarray([0.25, 0.5, 1.0]), + ), + } + operations = [("2", "1"), ("1",)] + actual = local_complex._apply_transformations_fast([local_input], transforms, operations) + expected = official._apply_transformations_fast([official_input], transforms, operations) + assert len(actual) == len(expected) == 2 + for actual_chain, expected_chain in zip(actual, expected, strict=True): + _assert_chain_equal(actual_chain, expected_chain) + + +def test_chain_construction_slicing_and_atom_views_match_pinned_biohub() -> None: + official = _official_chain() + actual = _make_chain(local_chain, chain_id="Q", entity_id=7, offset=0.0) + expected = _make_chain(official, chain_id="Q", entity_id=7, offset=0.0) + _assert_chain_equal(actual, expected) + selection = np.asarray([True, False, True, True]) + _assert_chain_equal(actual[selection], expected[selection]) + _assert_equal(actual.atoms[["N", "CA", "C"]], expected.atoms[["N", "CA", "C"]]) + _assert_equal(actual.atom_mask["CB"], expected.atom_mask["CB"]) + _assert_equal(actual.residue_index_no_insertions, expected.residue_index_no_insertions) + _assert_equal(actual.cbeta_contacts(), expected.cbeta_contacts()) + + +def test_chain_geometry_and_encoder_inputs_match_pinned_biohub() -> None: + official = _official_chain() + actual = _make_chain(local_chain, chain_id="Q", entity_id=7, offset=0.0) + expected = _make_chain(official, chain_id="Q", entity_id=7, offset=0.0) + _assert_chain_equal(actual.infer_cbeta(), expected.infer_cbeta()) + _assert_chain_equal(actual.infer_oxygen(), expected.infer_oxygen()) + _assert_chain_equal(actual.normalize_coordinates(), expected.normalize_coordinates()) + for actual_tensor, expected_tensor in zip( + actual.to_structure_encoder_inputs(), + expected.to_structure_encoder_inputs(), + strict=True, + ): + _assert_equal(actual_tensor, expected_tensor) + assert actual.rmsd(actual, only_compute_backbone_rmsd=True) == pytest.approx( + expected.rmsd(expected, only_compute_backbone_rmsd=True) + ) + + +def test_chain_compact_storage_matches_pinned_biohub() -> None: + official = _official_chain() + actual = _make_chain(local_chain, chain_id="Q", entity_id=7, offset=0.0) + expected = _make_chain(official, chain_id="Q", entity_id=7, offset=0.0) + _assert_equal(actual.state_dict(), expected.state_dict()) + assert actual.to_blob() == expected.to_blob() + _assert_chain_equal( + local_chain.ProteinChain.from_blob(actual.to_blob()), + official.ProteinChain.from_blob(expected.to_blob()), + ) + + +def test_chain_structure_interchange_matches_pinned_biohub() -> None: + official = _official_chain() + actual = _make_chain(local_chain, chain_id="Q", entity_id=7, offset=0.0) + expected = _make_chain(official, chain_id="Q", entity_id=7, offset=0.0) + assert actual.to_pdb_string() == expected.to_pdb_string() + assert actual.to_mmcif_string() == expected.to_mmcif_string() + _assert_equal(actual.atom_array.b_factor, expected.atom_array.b_factor) + _assert_equal(actual.atom_array.occupancy, expected.atom_array.occupancy) + _assert_chain_equal( + local_chain.ProteinChain.from_pdb(io.StringIO(actual.to_pdb_string()), is_predicted=True), + official.ProteinChain.from_pdb(io.StringIO(expected.to_pdb_string()), is_predicted=True), + ) + + +def _make_complexes() -> tuple[Any, Any]: + official_chain = _official_chain() + official_complex = _official_complex() + local_chains = [ + _make_chain(local_chain, chain_id="A", entity_id=1, offset=0.0), + _make_chain(local_chain, chain_id="B", entity_id=2, offset=2.0), + ] + official_chains = [ + _make_chain(official_chain, chain_id="A", entity_id=1, offset=0.0), + _make_chain(official_chain, chain_id="B", entity_id=2, offset=2.0), + ] + return ( + local_complex.ProteinComplex.from_chains(local_chains), + official_complex.ProteinComplex.from_chains(official_chains), + ) + + +def test_complex_construction_views_and_topology_match_pinned_biohub() -> None: + actual, expected = _make_complexes() + _assert_complex_equal(actual, expected) + assert actual.chain_boundaries == expected.chain_boundaries + _assert_equal(actual.chain_lengths, expected.chain_lengths) + _assert_equal(actual.chain_adjacency(), expected.chain_adjacency()) + _assert_equal(actual.chain_adjacency_by_index(0), expected.chain_adjacency_by_index(0)) + _assert_chain_equal(actual.get_chain_by_index(1), expected.get_chain_by_index(1)) + mask = np.asarray([True, True, False, False, False, True, True, True, True]) + _assert_complex_equal(actual[mask], expected[mask]) + + +def test_complex_geometry_and_compact_storage_match_pinned_biohub() -> None: + actual, expected = _make_complexes() + _assert_complex_equal(actual.infer_cbeta(), expected.infer_cbeta()) + _assert_complex_equal(actual.infer_oxygen(), expected.infer_oxygen()) + _assert_equal(actual.state_dict(), expected.state_dict()) + assert actual.to_blob() == expected.to_blob() + _assert_complex_equal( + local_complex.ProteinComplex.from_blob(actual.to_blob()), + _official_complex().ProteinComplex.from_blob(expected.to_blob()), + ) + assert actual.to_mmcif_string() == expected.to_mmcif_string() + + +def test_chain_to_complex_adapter_matches_pinned_biohub() -> None: + official_chain = _official_chain() + official_complex = _official_complex() + actual_chain = _make_chain(local_chain, chain_id="A", entity_id=1, offset=0.0) + expected_chain = _make_chain(official_chain, chain_id="A", entity_id=1, offset=0.0) + actual = local_complex.protein_chain_to_protein_complex(actual_chain) + expected = official_complex.protein_chain_to_protein_complex(expected_chain) + _assert_complex_equal(actual, expected) diff --git a/tests/parity/test_esmfold2_reimplemented_source_parity.py b/tests/parity/test_esmfold2_reimplemented_source_parity.py new file mode 100644 index 0000000..f8db555 --- /dev/null +++ b/tests/parity/test_esmfold2_reimplemented_source_parity.py @@ -0,0 +1,405 @@ +"""Parity checks for independently implemented ESMFold2 source utilities.""" + +from __future__ import annotations + +import importlib.util +import random +import sys +import types +import numpy as np +import pytest +import torch +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from fastplms.models.esmfold2 import configuration_esmfold2 as local_configuration +from fastplms.models.esmfold2 import esmfold2_affine3d as local_affine +from fastplms.models.esmfold2 import esmfold2_conformers as local_conformers +from fastplms.models.esmfold2 import esmfold2_constants as local_constants +from fastplms.models.esmfold2 import esmfold2_metrics as local_metrics +from fastplms.models.esmfold2 import esmfold2_misc as local_misc +from fastplms.models.esmfold2 import esmfold2_molecular_complex as local_complex +from fastplms.models.esmfold2 import esmfold2_output as local_output +from fastplms.models.esmfold2 import esmfold2_paired_msa as local_paired_msa +from fastplms.models.esmfold2 import esmfold2_prepare_input as local_prepare_input +from fastplms.models.esmfold2 import esmfold2_processor as local_processor +from fastplms.models.esmfold2 import esmfold2_protein_structure as local_structure +from fastplms.models.esmfold2 import esmfold2_residue_constants as local_residue_constants +from fastplms.models.esmfold2 import esmfold2_types as local_types +from fastplms.models.esmfold2.esmfold2_msa import MSA +from fastplms.models.esmfold2.esmfold2_parsing import FastaEntry + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +BIOHUB_ESM = ROOT / "vendor/upstream/biohub-esm/esm" +BIOHUB_TRANSFORMERS = ROOT / "vendor/upstream/biohub-transformers/src/transformers/models/esmfold2" +_MISSING = object() + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source( + module_name: str, + path: Path, + aliases: dict[str, types.ModuleType], +) -> types.ModuleType: + assert path.is_file(), f"pinned parity source is missing: {path}" + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + with _temporary_modules({**aliases, module_name: module}): + spec.loader.exec_module(module) + return module + + +def _biohub_packages() -> dict[str, types.ModuleType]: + return { + "esm": _package("esm"), + "esm.models": _package("esm.models"), + "esm.models.esmfold2": _package("esm.models.esmfold2"), + "esm.utils": _package("esm.utils"), + "esm.utils.msa": _package("esm.utils.msa"), + "esm.utils.structure": _package("esm.utils.structure"), + } + + +def _msa_compatibility_module() -> types.ModuleType: + module = types.ModuleType("esm.utils.msa.msa") + module.MSA = MSA + module.is_a3m_insertion = lambda character: character == "." or character.islower() + return module + + +def _load_official_configuration() -> types.ModuleType: + import transformers.configuration_utils as configuration_utils + + root_name = "_fastplms_pinned_transformers" + aliases = { + root_name: _package(root_name), + f"{root_name}.models": _package(f"{root_name}.models"), + f"{root_name}.models.esmfold2": _package(f"{root_name}.models.esmfold2"), + f"{root_name}.configuration_utils": configuration_utils, + } + return _load_source( + f"{root_name}.models.esmfold2.configuration_esmfold2", + BIOHUB_TRANSFORMERS / "configuration_esmfold2.py", + aliases, + ) + + +def _load_official_structure() -> types.ModuleType: + aliases = { + **_biohub_packages(), + "esm.utils.residue_constants": local_residue_constants, + "esm.utils.misc": local_misc, + "esm.utils.structure.affine3d": local_affine, + } + return _load_source( + "_fastplms_pinned_biohub_protein_structure", + BIOHUB_ESM / "utils/structure/protein_structure.py", + aliases, + ) + + +def _load_official_metrics(official_structure: types.ModuleType) -> types.ModuleType: + aliases = { + **_biohub_packages(), + "esm.utils.residue_constants": local_residue_constants, + "esm.utils.misc": local_misc, + "esm.utils.structure.protein_structure": official_structure, + } + return _load_source( + "_fastplms_pinned_biohub_metrics", + BIOHUB_ESM / "utils/structure/metrics.py", + aliases, + ) + + +def _load_official_paired_msa() -> types.ModuleType: + aliases = { + **_biohub_packages(), + "esm.models.esmfold2.constants": local_constants, + "esm.utils.msa.msa": _msa_compatibility_module(), + } + return _load_source( + "_fastplms_pinned_biohub_paired_msa", + BIOHUB_ESM / "models/esmfold2/paired_msa.py", + aliases, + ) + + +def _load_official_processor() -> types.ModuleType: + aliases = { + **_biohub_packages(), + "esm.models.esmfold2.conformers": local_conformers, + "esm.models.esmfold2.output": local_output, + "esm.models.esmfold2.prepare_input": local_prepare_input, + "esm.models.esmfold2.types": local_types, + "esm.utils.structure.molecular_complex": local_complex, + } + return _load_source( + "_fastplms_pinned_biohub_processor", + BIOHUB_ESM / "models/esmfold2/processor.py", + aliases, + ) + + +def _assert_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + # actual: (...), expected: (...) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +def test_configuration_matches_pinned_biohub_schema() -> None: + official = _load_official_configuration() + kwargs = { + "type": "release", + "d_single": 320, + "d_pair": 192, + "esmc_id": "Synthyra/ESMplusplus_6B", + "inputs": {"d_inputs": 777, "atom_encoder": {"n_blocks": 5}}, + "folding_trunk": {"n_layers": 8, "n_heads": 4}, + "structure_head": {"diffusion_module": {"token_num_blocks": 3}}, + "confidence_head": {"folding_trunk": {"n_layers": 2}}, + "msa_encoder": {"enabled": True, "d_msa": 64}, + "parcae": {"max_steps": None}, + "lm_encoder": {"per_loop_lm_dropout": False}, + "msa_encoder_overwrite": False, + } + actual = local_configuration.ESMFold2Config(**kwargs) + expected = official.ESMFold2Config(**kwargs) + + scalar_fields = ( + "type", + "d_single", + "d_pair", + "n_relative_residx_bins", + "n_relative_chain_bins", + "num_loops", + "num_diffusion_samples", + "disable_msa_features", + "lm_dropout", + "force_lm_dropout_during_inference", + "lm_d_model", + "lm_num_layers", + "esmc_id", + "msa_encoder_overwrite", + ) + assert {name: getattr(actual, name) for name in scalar_fields} == { + name: getattr(expected, name) for name in scalar_fields + } + nested_fields = ( + "inputs", + "folding_trunk", + "structure_head", + "confidence_head", + "msa_encoder", + "parcae", + "lm_encoder", + ) + assert {name: asdict(getattr(actual, name)) for name in nested_fields} == { + name: asdict(getattr(expected, name)) for name in nested_fields + } + + +def test_protein_geometry_matches_pinned_biohub_on_h100() -> None: + assert torch.cuda.is_available(), "the ESMFold2 compliance suite requires CUDA" + official = _load_official_structure() + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260714) + # mobile: (3, 12, 3) + mobile = torch.randn((3, 12, 3), generator=generator, device=device) + target = mobile + 0.05 * torch.randn( + mobile.shape, + generator=generator, + device=device, + ) + # mask: (3, 12) + mask = torch.tensor( + [[1] * 12, [1] * 9 + [0] * 3, [1, 0] * 6], + dtype=torch.bool, + device=device, + ) + + actual_alignment = local_structure.compute_alignment_tensors(mobile, target, mask) + expected_alignment = official.compute_alignment_tensors(mobile, target, mask) + for actual, expected in zip(actual_alignment, expected_alignment, strict=True): + _assert_equal(actual, expected) + + actual_affine, actual_rmsd = local_structure.compute_affine_and_rmsd(mobile, target, mask) + expected_affine, expected_rmsd = official.compute_affine_and_rmsd(mobile, target, mask) + _assert_equal(actual_affine.tensor, expected_affine.tensor) + _assert_equal(actual_rmsd, expected_rmsd) + _assert_equal( + local_structure.compute_gdt_ts_no_alignment(mobile, target, mask), + official.compute_gdt_ts_no_alignment(mobile, target, mask), + ) + + +def test_structure_metrics_match_pinned_biohub_on_h100() -> None: + assert torch.cuda.is_available(), "the ESMFold2 compliance suite requires CUDA" + official = _load_official_metrics(_load_official_structure()) + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(712) + # predicted: (2, 10, 3) + predicted = torch.randn((2, 10, 3), generator=generator, device=device) + target = predicted + 0.2 * torch.randn( + predicted.shape, + generator=generator, + device=device, + ) + # atom_mask: (2, 10) + atom_mask = torch.tensor( + [[1] * 10, [1] * 7 + [0] * 3], + dtype=torch.float32, + device=device, + ) + # sequence_id: (2, 10) + sequence_id = torch.tensor( + [[0] * 5 + [1] * 5, [0] * 4 + [1] * 6], + device=device, + ) + for per_residue in (False, True): + _assert_equal( + local_metrics.compute_lddt( + predicted, + target, + atom_mask, + per_residue=per_residue, + sequence_id=sequence_id, + ), + official.compute_lddt( + predicted, + target, + atom_mask, + per_residue=per_residue, + sequence_id=sequence_id, + ), + ) + + # predictions: (2, 14, 14) + predictions = torch.rand((2, 14, 14), generator=generator, device=device) + # targets: (...) + targets = torch.randint(0, 2, (2, 14, 14), generator=generator, device=device).float() + targets[1, 11:] = -1 + # lengths: (2,) + lengths = torch.tensor([14, 11], device=device) + actual_contacts = local_metrics.contact_precision( + predictions, targets, lengths, minsep=3, maxsep=10 + ) + expected_contacts = official.contact_precision( + predictions, targets, lengths, minsep=3, maxsep=10 + ) + for name, expected in expected_contacts.items(): + _assert_equal(actual_contacts[name], expected) + + +def test_paired_msa_matches_pinned_biohub() -> None: + official = _load_official_paired_msa() + msa_a = MSA( + [ + FastaEntry("query", "ACD-E"), + FastaEntry("a key=10", "AqCD-E"), + FastaEntry("b key=20", "ACD-E"), + FastaEntry("unpaired", "A-CDX"), + ] + ) + msa_b = MSA( + [ + FastaEntry("query", "FGHI"), + FastaEntry("a key=10", "FGHI"), + FastaEntry("b key=20", "FgGHI"), + FastaEntry("extra key=10", "F-HI"), + ] + ) + arguments = { + "chain_msas": {1: msa_a, 2: msa_b, 3: None}, + "chain_query_res_types": { + 1: np.asarray([0, 1, 2, 3, 4]), + 2: np.asarray([5, 6, 7, 8]), + 3: np.asarray([9, 10]), + }, + "token_asym_ids": np.asarray([1, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3]), + "token_res_ids": np.asarray([0, 1, 2, 3, 4, 0, 1, 2, 3, 0, 1]), + "max_pairs": 8, + "max_total": 12, + "max_seqs": 10, + } + actual = local_paired_msa.construct_paired_msa(**arguments) + official_arguments = { + **arguments, + "chain_msas": { + chain_id: ( + None + if msa is None + else types.SimpleNamespace( + entries=msa.entries, + depth=msa.depth, + deletions=None, + ) + ) + for chain_id, msa in arguments["chain_msas"].items() + }, + } + expected = official.construct_paired_msa(**official_arguments) + for actual_array, expected_array in zip(actual, expected, strict=True): + np.testing.assert_array_equal(actual_array, expected_array) + + +def _random_trace(module: types.ModuleType) -> tuple[Any, Any]: + random.seed(42) + np.random.seed(42) + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + with module._seed_context(19): + inside = (random.random(), float(np.random.random()), torch.rand(3, device="cuda")) + after = (random.random(), float(np.random.random()), torch.rand(3, device="cuda")) + return inside, after + + +def test_processor_cleaning_and_rng_match_pinned_biohub_on_h100() -> None: + assert torch.cuda.is_available(), "the ESMFold2 compliance suite requires CUDA" + official = _load_official_processor() + source_msa = MSA.from_sequences(["AAA-AAA-BBB", "AAA-AAA-BBB"]) + input_value = local_types.StructurePredictionInput( + sequences=[ + local_types.ProteinInput( + id=["entity"], + sequence="AAA|AAA|BBB", + modifications=[local_types.Modification(position=1, ccd="MSE")], + msa=source_msa, + ) + ] + ) + assert local_processor.clean_esmfold2_input(input_value) == official.clean_esmfold2_input( + input_value + ) + + actual_trace = _random_trace(local_processor) + expected_trace = _random_trace(official) + assert actual_trace[0][:2] == expected_trace[0][:2] + assert actual_trace[1][:2] == expected_trace[1][:2] + _assert_equal(actual_trace[0][2], expected_trace[0][2]) + _assert_equal(actual_trace[1][2], expected_trace[1][2]) diff --git a/tests/parity/test_esmfold2_residue_config_parity.py b/tests/parity/test_esmfold2_residue_config_parity.py new file mode 100644 index 0000000..466fd34 --- /dev/null +++ b/tests/parity/test_esmfold2_residue_config_parity.py @@ -0,0 +1,457 @@ +"""Exact pinned parity for ESMFold2 residue constants and configuration.""" + +from __future__ import annotations + +import importlib.util +import re +import struct +import sys +import tokenize +import types +import numpy as np +import pytest +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import asdict +from difflib import SequenceMatcher +from io import StringIO +from pathlib import Path +from typing import Any + +from fastplms.models.esmfold2 import configuration_esmfold2 as local_config +from fastplms.models.esmfold2 import esmfold2_residue_constants as local_residues + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +LOCAL_ROOT = ROOT / "src/fastplms/models/esmfold2" +OFFICIAL_RESIDUES = ROOT / "vendor/upstream/biohub-esm/esm/utils/residue_constants.py" +OFFICIAL_CONFIG = ( + ROOT + / "vendor/upstream/biohub-transformers/src/transformers/models/esmfold2" + / "configuration_esmfold2.py" +) +_MISSING = object() + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source( + module_name: str, + path: Path, + aliases: dict[str, types.ModuleType] | None = None, +) -> types.ModuleType: + assert path.is_file(), f"pinned source is missing: {path}" + specification = importlib.util.spec_from_file_location(module_name, path) + assert specification is not None and specification.loader is not None + module = importlib.util.module_from_spec(specification) + with _temporary_modules({**(aliases or {}), module_name: module}): + specification.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def official_residues() -> types.ModuleType: + return _load_source("_fastplms_pinned_biohub_residue_constants", OFFICIAL_RESIDUES) + + +@pytest.fixture(scope="module") +def official_config() -> types.ModuleType: + import transformers.configuration_utils as configuration_utils + + root_name = "_fastplms_pinned_biohub_transformers" + aliases = { + root_name: _package(root_name), + f"{root_name}.models": _package(f"{root_name}.models"), + f"{root_name}.models.esmfold2": _package(f"{root_name}.models.esmfold2"), + f"{root_name}.configuration_utils": configuration_utils, + } + return _load_source( + f"{root_name}.models.esmfold2.configuration_esmfold2", + OFFICIAL_CONFIG, + aliases, + ) + + +def _assert_exact(actual: Any, expected: Any) -> None: + if isinstance(expected, np.ndarray): + assert isinstance(actual, np.ndarray) + assert actual.dtype == expected.dtype + assert actual.shape == expected.shape + assert actual.tobytes() == expected.tobytes() + return + if isinstance(expected, float): + assert isinstance(actual, float) + assert struct.pack("!d", actual) == struct.pack("!d", expected) + return + if isinstance(expected, dict): + assert isinstance(actual, dict) + assert tuple(actual) == tuple(expected) + for key in expected: + _assert_exact(actual[key], expected[key]) + return + if isinstance(expected, (list, tuple)): + assert type(actual) is type(expected) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_exact(actual_item, expected_item) + return + assert type(actual) is type(expected) + assert actual == expected + + +def _public_data(module: types.ModuleType) -> dict[str, Any]: + return { + name: value + for name, value in vars(module).items() + if not name.startswith("_") + and name != "annotations" + and not isinstance(value, types.ModuleType) + and not callable(value) + } + + +def test_all_public_residue_tables_match_exactly( + official_residues: types.ModuleType, +) -> None: + expected = _public_data(official_residues) + actual = _public_data(local_residues) + assert actual.keys() == expected.keys(), ( + f"local-only={sorted(actual.keys() - expected.keys())}, " + f"official-only={sorted(expected.keys() - actual.keys())}" + ) + for name in expected: + _assert_exact(actual[name], expected[name]) + assert local_residues.Bond._fields == official_residues.Bond._fields + assert local_residues.BondAngle._fields == official_residues.BondAngle._fields + + +@pytest.mark.parametrize("atom_index", range(4)) +def test_chi_selectors_match_exactly(official_residues: types.ModuleType, atom_index: int) -> None: + _assert_exact( + local_residues.chi_angle_atom(atom_index), + official_residues.chi_angle_atom(atom_index), + ) + + +@pytest.mark.parametrize( + ("sequence", "mapping", "map_unknown_to_x"), + [ + ("", {"A": 0, "X": 1}, False), + ("AXA", {"A": 0, "X": 1}, False), + ("AZ", {"A": 0, "X": 1}, True), + ("ACDEFGHIKLMNPQRSTVWY", local_residues.restype_order_with_x, False), + ], +) +def test_one_hot_encoding_matches_exactly( + official_residues: types.ModuleType, + sequence: str, + mapping: dict[str, int], + map_unknown_to_x: bool, +) -> None: + actual = local_residues.sequence_to_onehot(sequence, mapping, map_unknown_to_x) + expected = official_residues.sequence_to_onehot(sequence, mapping, map_unknown_to_x) + _assert_exact(actual, expected) + + +@pytest.mark.parametrize( + ("sequence", "mapping", "map_unknown_to_x"), + [ + ("a", {"A": 0, "X": 1}, True), + ("?", {"A": 0, "X": 1}, True), + ("B", {"A": 0, "X": 2}, True), + ("B", {"A": 0}, False), + ], +) +def test_one_hot_failures_match_pinned_behavior( + official_residues: types.ModuleType, + sequence: str, + mapping: dict[str, int], + map_unknown_to_x: bool, +) -> None: + with pytest.raises(Exception) as local_error: + local_residues.sequence_to_onehot(sequence, mapping, map_unknown_to_x) + with pytest.raises(Exception) as official_error: + official_residues.sequence_to_onehot(sequence, mapping, map_unknown_to_x) + assert type(local_error.value) is type(official_error.value) + assert str(local_error.value) == str(official_error.value) + + +def test_rigid_transform_and_mapping_builders_match_exactly( + official_residues: types.ModuleType, +) -> None: + ex = np.asarray((1.25, -0.5, 0.75), dtype=np.float64) + ey = np.asarray((-0.25, 1.5, 0.125), dtype=np.float64) + translation = np.asarray((7.0, -3.0, 2.0), dtype=np.float64) + _assert_exact( + local_residues._make_rigid_transformation_4x4(ex, ey, translation), + official_residues._make_rigid_transformation_4x4(ex, ey, translation), + ) + for function_name in ( + "_make_standard_atom_mask", + "_make_restype_atom14_to_atom37", + "_make_restype_atom37_to_atom14", + ): + _assert_exact( + getattr(local_residues, function_name)(), + getattr(official_residues, function_name)(), + ) + indices = np.asarray((0, 4, 7, 20, 2), dtype=np.int64) + assert local_residues.aatype_to_str_sequence(indices) == ( + official_residues.aatype_to_str_sequence(indices) + ) + + +def test_stereo_chemical_derivation_matches_exactly( + monkeypatch: pytest.MonkeyPatch, + official_residues: types.ModuleType, +) -> None: + table = """bond residue length stddev +N-CA ALA 1.458 0.020 +CA-C ALA 1.525 0.021 +- + +angle residue degrees stddev +N-CA-C ALA 111.2 2.1 +- +""" + + class _Table: + @staticmethod + def read_text() -> str: + return table + + monkeypatch.setattr(local_residues, "_STEREO_CHEMICAL_PROPS_PATH", _Table()) + monkeypatch.setattr(official_residues, "Path", lambda _path: _Table()) + local_residues.load_stereo_chemical_props.cache_clear() + official_residues.load_stereo_chemical_props.cache_clear() + actual = local_residues.load_stereo_chemical_props() + expected = official_residues.load_stereo_chemical_props() + for actual_group, expected_group in zip(actual, expected, strict=True): + assert tuple(actual_group) == tuple(expected_group) + for residue_name in expected_group: + assert [tuple(record) for record in actual_group[residue_name]] == [ + tuple(record) for record in expected_group[residue_name] + ] + for actual_record, expected_record in zip( + actual_group[residue_name], expected_group[residue_name], strict=True + ): + for actual_value, expected_value in zip( + actual_record, expected_record, strict=True + ): + _assert_exact(actual_value, expected_value) + local_residues.load_stereo_chemical_props.cache_clear() + official_residues.load_stereo_chemical_props.cache_clear() + + +@pytest.mark.parametrize( + ("overlap_tolerance", "bond_length_tolerance_factor"), + [(1.5, 15.0), (0.75, 4.0)], +) +def test_atom14_distance_bounds_match_exactly( + monkeypatch: pytest.MonkeyPatch, + official_residues: types.ModuleType, + overlap_tolerance: float, + bond_length_tolerance_factor: float, +) -> None: + def tables(module: types.ModuleType) -> tuple[dict, dict, dict]: + bonds = {name: [] for name in module.resnames} + virtual = {name: [] for name in module.resnames} + angles = {name: [] for name in module.resnames} + bonds["ALA"] = [module.Bond("N", "CA", 1.458, 0.02)] + return bonds, virtual, angles + + monkeypatch.setattr( + local_residues, "load_stereo_chemical_props", lambda: tables(local_residues) + ) + monkeypatch.setattr( + official_residues, + "load_stereo_chemical_props", + lambda: tables(official_residues), + ) + actual = local_residues.make_atom14_dists_bounds( + overlap_tolerance, bond_length_tolerance_factor + ) + expected = official_residues.make_atom14_dists_bounds( + overlap_tolerance, bond_length_tolerance_factor + ) + assert actual.keys() == expected.keys() + for name in expected: + _assert_exact(actual[name], expected[name]) + + +_NESTED_NAMES = ( + "AtomAttentionConfig", + "DiffusionModuleConfig", + "FoldingTrunkConfig", + "InputsEmbedderConfig", + "DiffusionStructureHeadConfig", + "ConfidenceHeadConfig", + "MSAEncoderConfig", + "LMEncoderConfig", + "ParcaeConfig", +) + + +def test_nested_configuration_defaults_match_exactly( + official_config: types.ModuleType, +) -> None: + for class_name in _NESTED_NAMES: + actual = getattr(local_config, class_name)() + expected = getattr(official_config, class_name)() + assert asdict(actual) == asdict(expected) + + +def test_full_configuration_schema_matches_pinned_biohub( + official_config: types.ModuleType, +) -> None: + kwargs = { + "type": "experimental", + "d_single": 320, + "d_pair": 192, + "n_relative_residx_bins": 21, + "n_relative_chain_bins": 3, + "num_loops": 7, + "num_diffusion_samples": 3, + "disable_msa_features": True, + "lm_dropout": 0.125, + "force_lm_dropout_during_inference": True, + "lm_d_model": 1280, + "lm_num_layers": 40, + "esmc_id": "Synthyra/ESMplusplus_6B", + "inputs": {"d_inputs": 777, "atom_encoder": {"n_blocks": 5}}, + "folding_trunk": {"n_layers": 8, "n_heads": 4}, + "structure_head": {"diffusion_module": {"token_num_blocks": 3}}, + "confidence_head": {"folding_trunk": {"n_layers": 2}}, + "msa_encoder": {"enabled": True, "d_msa": 64}, + "parcae": {"max_steps": None}, + "lm_encoder": {"per_loop_lm_dropout": False}, + "msa_encoder_overwrite": False, + } + actual = local_config.ESMFold2Config(**kwargs) + expected = official_config.ESMFold2Config(**kwargs) + shared_scalars = ( + "type", + "d_single", + "d_pair", + "n_relative_residx_bins", + "n_relative_chain_bins", + "num_loops", + "num_diffusion_samples", + "disable_msa_features", + "lm_dropout", + "force_lm_dropout_during_inference", + "lm_d_model", + "lm_num_layers", + "esmc_id", + "msa_encoder_overwrite", + ) + assert {name: getattr(actual, name) for name in shared_scalars} == { + name: getattr(expected, name) for name in shared_scalars + } + for name in ( + "inputs", + "folding_trunk", + "structure_head", + "confidence_head", + "msa_encoder", + "parcae", + "lm_encoder", + ): + assert asdict(getattr(actual, name)) == asdict(getattr(expected, name)) + + +def test_fastplms_configuration_extensions_are_strict(tmp_path: Path) -> None: + config = local_config.ESMFold2Config( + type="release", + esmc_id="biohub/ESMC-6B", + attn_implementation={"": "flex"}, + esmc_precision="fp8", + ) + assert config.esmc_id == "Synthyra/ESMplusplus_6B" + assert config.esmc_attn_backend == "flex_attention" + assert config.esmc_precision == "fp8" + config.save_pretrained(tmp_path) + restored = local_config.ESMFold2Config.from_pretrained(tmp_path) + assert restored.to_dict() == config.to_dict() + with pytest.raises(ValueError, match="Unsupported ESMFold2 attention"): + local_config.ESMFold2Config(type="release", attn_implementation="unknown") + with pytest.raises(ValueError, match="esmc_precision"): + local_config.ESMFold2Config(type="release", esmc_precision="int8") + + +def test_modules_have_standalone_artifact_import_closure() -> None: + residues = _load_source( + "_fastplms_artifact_residue_constants", + LOCAL_ROOT / "esmfold2_residue_constants.py", + ) + configuration = _load_source( + "_fastplms_artifact_configuration_esmfold2", + LOCAL_ROOT / "configuration_esmfold2.py", + ) + assert residues.STANDARD_ATOM_MASK.shape == (21, 37) + config = configuration.ESMFold2Config(type="release") + assert config.model_type == "esmfold2" + assert config.esmc_id == "Synthyra/ESMplusplus_6B" + + +def _meaningful_lines(text: str) -> list[str]: + return [ + " ".join(line.strip().split()) + for line in text.splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + +@pytest.mark.parametrize( + ("local_name", "official_path"), + [ + ("esmfold2_residue_constants.py", OFFICIAL_RESIDUES), + ("configuration_esmfold2.py", OFFICIAL_CONFIG), + ], +) +def test_runtime_source_is_independently_organized(local_name: str, official_path: Path) -> None: + local_path = LOCAL_ROOT / local_name + similarity = SequenceMatcher( + None, + _meaningful_lines(local_path.read_text(encoding="utf-8")), + _meaningful_lines(official_path.read_text(encoding="utf-8")), + autojunk=False, + ).ratio() + assert similarity < 0.75, f"{local_name} has line similarity {similarity:.3f}" + + +def test_comments_and_docstrings_follow_shape_notation() -> None: + square_shape = re.compile(r"\[\s*[A-Z](?:\s*,\s*[A-Z])+(?:\s*,[^]]*)?\]") + upper_dimensions = re.compile(r"\(\s*[BLDNH](?:\s*,\s*[BLDNH])+(?:\s*,[^)]*)?\)") + for path in ( + LOCAL_ROOT / "esmfold2_residue_constants.py", + LOCAL_ROOT / "configuration_esmfold2.py", + ): + source = path.read_text(encoding="utf-8") + prose = "\n".join( + token.string + for token in tokenize.generate_tokens(StringIO(source).readline) + if token.type in {tokenize.COMMENT, tokenize.STRING} + ) + assert square_shape.search(prose) is None + assert upper_dimensions.search(prose) is None diff --git a/tests/parity/test_esmfold2_source_slice3_parity.py b/tests/parity/test_esmfold2_source_slice3_parity.py new file mode 100644 index 0000000..78eee7d --- /dev/null +++ b/tests/parity/test_esmfold2_source_slice3_parity.py @@ -0,0 +1,667 @@ +"""Differential contracts for independently organized ESMFold2 utilities.""" + +from __future__ import annotations + +import importlib.util +import io +import sys +import types +import biotite.structure as bs +import numpy as np +import pytest +import torch +import zstandard +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +from fastplms.models.esmfold2 import esmfold2_affine3d as local_affine +from fastplms.models.esmfold2 import esmfold2_misc as local_misc +from fastplms.models.esmfold2 import esmfold2_mmcif_parsing as local_mmcif +from fastplms.models.esmfold2 import esmfold2_msa as local_msa +from fastplms.models.esmfold2 import esmfold2_msa_filter_sequences as local_filter +from fastplms.models.esmfold2 import esmfold2_parsing as local_parsing +from fastplms.models.esmfold2 import esmfold2_residue_constants as local_residues +from fastplms.models.esmfold2 import esmfold2_sequential_dataclass as local_sequential +from fastplms.models.esmfold2 import esmfold2_system as local_system +from fastplms.models.esmfold2 import esmfold2_utils_types as local_types +from fastplms.models.esmfold2.esmfold2_constants_esm3 import CHAIN_BREAK_STR + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +BIOHUB_ESM = ROOT / "vendor/upstream/biohub-esm/esm" +_MISSING = object() + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source( + module_name: str, + path: Path, + aliases: dict[str, types.ModuleType], +) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + with _temporary_modules({**aliases, module_name: module}): + spec.loader.exec_module(module) + return module + + +def _base_aliases() -> dict[str, types.ModuleType]: + return { + "esm": _package("esm"), + "esm.utils": _package("esm.utils"), + "esm.utils.constants": _package("esm.utils.constants"), + "esm.utils.msa": _package("esm.utils.msa"), + "esm.utils.structure": _package("esm.utils.structure"), + } + + +def _official_misc() -> types.ModuleType: + constants = types.ModuleType("esm.utils.constants.esm3") + constants.CHAIN_BREAK_STR = CHAIN_BREAK_STR + zstd_adapter = types.ModuleType("zstd") + decompress = zstandard.ZstdDecompressor().decompress + zstd_adapter.decompress = decompress # type: ignore[attr-defined] + zstd_adapter.ZSTD_uncompress = decompress # type: ignore[attr-defined] + return _load_source( + "_fastplms_pinned_biohub_misc", + BIOHUB_ESM / "utils/misc.py", + { + **_base_aliases(), + "esm.utils.constants.esm3": constants, + "esm.utils.types": local_types, + "zstd": zstd_adapter, + }, + ) + + +def _official_affine() -> types.ModuleType: + return _load_source( + "_fastplms_pinned_biohub_affine3d", + BIOHUB_ESM / "utils/structure/affine3d.py", + {**_base_aliases(), "esm.utils.misc": local_misc}, + ) + + +def _official_msa() -> types.ModuleType: + return _load_source( + "_fastplms_pinned_biohub_msa", + BIOHUB_ESM / "utils/msa/msa.py", + { + **_base_aliases(), + "esm.utils.misc": local_misc, + "esm.utils.msa.filter_sequences": local_filter, + "esm.utils.parsing": local_parsing, + "esm.utils.sequential_dataclass": local_sequential, + "esm.utils.system": local_system, + }, + ) + + +def _official_mmcif() -> types.ModuleType: + return _load_source( + "_fastplms_pinned_biohub_mmcif", + BIOHUB_ESM / "utils/structure/mmcif_parsing.py", + {**_base_aliases(), "esm.utils.residue_constants": local_residues}, + ) + + +def _assert_tensor_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + # actual: (...), expected: (...) + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +def _assert_array_equal(actual: np.ndarray, expected: np.ndarray) -> None: + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_misc_tensor_contracts_match_pinned_biohub(device: str) -> None: + assert torch.cuda.is_available(), "the utility parity suite requires CUDA" + official = _official_misc() + torch.manual_seed(730) + # values: (2, 4) + values = torch.randn((2, 4), device=device) + _assert_tensor_equal( + local_misc.rbf(values, -2.0, 3.0, 9), + official.rbf(values, -2.0, 3.0, 9), + ) + + # data: (2, 3, 5, 2) + data = torch.arange(2 * 3 * 5 * 2, device=device).reshape(2, 3, 5, 2) + # indices: (2, 3, 2) + indices = torch.tensor( + [[[0, 3], [2, 1], [4, 0]], [[4, 1], [0, 2], [3, 3]]], + device=device, + ) + _assert_tensor_equal( + local_misc.batched_gather(data, indices, dim=2, no_batch_dims=2), + official.batched_gather(data, indices, dim=2, no_batch_dims=2), + ) + + # coords: (2, 3, 3) + coords = torch.tensor( + [ + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [float("nan"), 0.0, 0.0]], + [[0.0, 0.0, 0.0], [3.0, 0.0, 0.0], [4.0, 0.0, 0.0]], + ], + device=device, + ) + # coord_mask: (2, 3) + coord_mask = torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.bool, device=device) + # padding_mask: (2, 3) + padding_mask = torch.tensor([[0, 0, 1], [0, 0, 0]], dtype=torch.bool, device=device) + # sequence_id: (2, 3) + sequence_id = torch.tensor([[0, 0, 0], [0, 1, 1]], device=device) + actual_edges = local_misc.knn_graph( + coords, + coord_mask, + padding_mask, + sequence_id, + no_knn=3, + ) + expected_edges = official.knn_graph( + coords, + coord_mask, + padding_mask, + sequence_id, + no_knn=3, + ) + for actual, expected in zip(actual_edges, expected_edges, strict=True): + _assert_tensor_equal(actual, expected) + + rows = [torch.arange(3, device=device), torch.arange(5, device=device)] + _assert_tensor_equal( + local_misc.stack_variable_length_tensors(rows, -1), + official.stack_variable_length_tensors(rows, -1), + ) + # packed: (5, 4, 2) + packed = torch.arange(5 * 4 * 2, device=device).reshape(5, 4, 2) + # ids: (2, 4) + ids = torch.tensor([[0, 0, 1, 1], [0, 1, 1, 2]], device=device) + actual_bin = local_misc.binpack(packed, ids, -7) + expected_bin = official.binpack(packed, ids, -7) + _assert_tensor_equal(actual_bin, expected_bin) + _assert_tensor_equal( + local_misc.unbinpack(actual_bin, ids, -7), + official.unbinpack(expected_bin, ids, -7), + ) + + +def test_misc_python_contracts_match_pinned_biohub() -> None: + official = _official_misc() + indices = np.asarray([True, False, True, False]) + for value in ("ABCD", [1, 2, 3, 4], (1, 2, 3, 4)): + assert local_misc.slice_python_object_as_numpy(value, indices) == ( + official.slice_python_object_as_numpy(value, indices) + ) + assert local_misc.merge_ranges( + [range(8, 10), range(1, 3), range(4, 8)], merge_gap_max=1 + ) == official.merge_ranges([range(8, 10), range(1, 3), range(4, 8)], merge_gap_max=1) + annotations = [ + local_types.FunctionAnnotation(label="helix", start=2, end=5), + local_types.FunctionAnnotation(label="helix", start=7, end=8), + local_types.FunctionAnnotation(label="site", start=4, end=4), + ] + assert local_misc.merge_annotations(annotations, 1) == official.merge_annotations( + annotations, 1 + ) + sequence = list(f"AC{CHAIN_BREAK_STR}DE") + _assert_array_equal( + local_misc.get_chainbreak_boundaries_from_sequence(sequence), + official.get_chainbreak_boundaries_from_sequence(sequence), + ) + for value in (None, [1, float("inf"), -float("inf")]): + assert local_misc.replace_inf(value) == official.replace_inf(value) + assert local_misc.join_lists([[1, 2], [3], [4]], [0]) == official.join_lists( + [[1, 2], [3], [4]], [0] + ) + _assert_array_equal( + local_misc.concat_objects([np.asarray([1, 2]), np.asarray([3])], separator=0), + official.concat_objects([np.asarray([1, 2]), np.asarray([3])], separator=0), + ) + + +def test_misc_conversion_serialization_and_concat_match_pinned_biohub() -> None: + official = _official_misc() + tensor_list = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])] + _assert_tensor_equal( + local_misc.maybe_tensor(tensor_list), + official.maybe_tensor(tensor_list), + ) + nested = [[1.0, None], [float("inf"), -2.0]] + _assert_tensor_equal( + local_misc.maybe_tensor(nested, convert_none_to_nan=True), + official.maybe_tensor(nested, convert_none_to_nan=True), + ) + # values: (2, 2) + values = torch.tensor([[1.0, float("nan")], [3.0, 4.0]]) + assert local_misc.maybe_list(values, convert_nan_to_none=True) == ( + official.maybe_list(values, convert_nan_to_none=True) + ) + assert local_misc.concat_objects(["AB", "CD"], "|") == ( + official.concat_objects(["AB", "CD"], "|") + ) + assert local_misc.concat_objects([[1, 2], [3]], 0) == official.concat_objects([[1, 2], [3]], 0) + _assert_tensor_equal( + local_misc.concat_objects([torch.tensor([1, 2]), torch.tensor([3])], separator=0), + official.concat_objects([torch.tensor([1, 2]), torch.tensor([3])], separator=0), + ) + assert list(local_misc.iterate_with_intermediate([1, 2, 3], 0)) == list( + official.iterate_with_intermediate([1, 2, 3], 0) + ) + + buffer = io.BytesIO() + torch.save({"X": torch.arange(6).reshape(2, 3)}, buffer) + compressed = zstandard.ZstdCompressor().compress(buffer.getvalue()) + actual = local_misc.deserialize_tensors(compressed) + expected = official.deserialize_tensors(compressed) + assert actual.keys() == expected.keys() + _assert_tensor_equal(actual["X"], expected["X"]) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_affine_rotation_contracts_match_pinned_biohub(device: str) -> None: + assert torch.cuda.is_available(), "the affine parity suite requires CUDA" + official = _official_affine() + generator = torch.Generator(device=device).manual_seed(2026) + # quaternions: (2, 5, 4) + quaternions = torch.randn((2, 5, 4), generator=generator, device=device) + # points: (2, 5, 3) + points = torch.randn((2, 5, 3), generator=generator, device=device) + + actual_quat = local_affine.RotationQuat(quaternions, normalized=True) + expected_quat = official.RotationQuat(quaternions, normalized=True) + _assert_tensor_equal(actual_quat.tensor, expected_quat.tensor) + _assert_tensor_equal(actual_quat.as_matrix().tensor, expected_quat.as_matrix().tensor) + _assert_tensor_equal(actual_quat.apply(points), expected_quat.apply(points)) + _assert_tensor_equal(actual_quat.invert().tensor, expected_quat.invert().tensor) + _assert_tensor_equal( + actual_quat.compose(actual_quat).tensor, + expected_quat.compose(expected_quat).tensor, + ) + + actual_matrix = actual_quat.as_matrix() + expected_matrix = expected_quat.as_matrix() + _assert_tensor_equal(actual_matrix.as_quat().tensor, expected_matrix.as_quat().tensor) + _assert_tensor_equal(actual_matrix.apply(points), expected_matrix.apply(points)) + _assert_tensor_equal(actual_matrix.invert().tensor, expected_matrix.invert().tensor) + _assert_tensor_equal( + actual_matrix.compose(actual_matrix).tensor, + expected_matrix.compose(expected_matrix).tensor, + ) + + # batched_points: (2, 7, 3) + batched_points = torch.randn((2, 7, 3), generator=generator, device=device) + actual_single = local_affine.RotationMatrix.identity((2, 1), device=device) + expected_single = official.RotationMatrix.identity((2, 1), device=device) + _assert_tensor_equal( + actual_single.apply(batched_points), + expected_single.apply(batched_points), + ) + + +def test_affine_frames_and_coordinate_fallback_match_pinned_biohub() -> None: + assert torch.cuda.is_available(), "the affine parity suite requires CUDA" + official = _official_affine() + generator = torch.Generator(device="cuda").manual_seed(912) + # translation: (2, 4, 3) + translation = torch.randn((2, 4, 3), generator=generator, device="cuda") + # quaternion: (2, 4, 4) + quaternion = torch.randn((2, 4, 4), generator=generator, device="cuda") + # encoded: (...) + encoded = torch.cat((quaternion, translation), dim=-1) + actual = local_affine.Affine3D.from_tensor(encoded) + expected = official.Affine3D.from_tensor(encoded) + _assert_tensor_equal(actual.tensor, expected.tensor) + _assert_tensor_equal(actual.invert().tensor, expected.invert().tensor) + _assert_tensor_equal( + actual.compose(actual).tensor, + expected.compose(expected).tensor, + ) + # mask: (2, 4) + mask = torch.tensor( + [[1, 0, 1, 0], [0, 1, 1, 0]], + dtype=torch.bool, + device="cuda", + ) + _assert_tensor_equal(actual.mask(mask).tensor, expected.mask(mask).tensor) + _assert_tensor_equal( + actual.mask(mask, with_zero=True).tensor, + expected.mask(mask, with_zero=True).tensor, + ) + + # coords: (2, 6, 3, 3) + coords = torch.randn((2, 6, 3, 3), generator=generator, device="cuda") + coords[0, 2] = torch.nan + coords[1, 4] = 2e6 + actual_frame, actual_mask = local_affine.build_affine3d_from_coordinates(coords) + expected_frame, expected_mask = official.build_affine3d_from_coordinates(coords) + _assert_tensor_equal(actual_frame.tensor, expected_frame.tensor) + _assert_tensor_equal(actual_mask, expected_mask) + + +def test_affine_encodings_and_collection_operations_match_pinned_biohub() -> None: + assert torch.cuda.is_available(), "the affine parity suite requires CUDA" + official = _official_affine() + generator = torch.Generator(device="cuda").manual_seed(661) + # translations: (2, 3) + translations = torch.randn((2, 3), generator=generator, device="cuda") + # compact_quat: (2, 3) + compact_quat = torch.randn((2, 3), generator=generator, device="cuda") + # full_quat: (2, 4) + full_quat = torch.randn((2, 4), generator=generator, device="cuda") + # matrix: (...) + matrix = torch.eye(3, device="cuda").expand(2, -1, -1) + # matrix4: (...) + matrix4 = torch.eye(4, device="cuda").expand(2, -1, -1).clone() + matrix4[..., :3, 3] = translations + encodings = ( + matrix4, + torch.cat((compact_quat, translations), dim=-1), + torch.cat((full_quat, translations), dim=-1), + torch.cat((matrix.flatten(-2), translations), dim=-1), + ) + for encoded in encodings: + actual = local_affine.Affine3D.from_tensor(encoded) + expected = official.Affine3D.from_tensor(encoded) + _assert_tensor_equal(actual.tensor, expected.tensor) + + # x_axis: (2, 3) + x_axis = torch.randn((2, 3), generator=generator, device="cuda") + # origin: (2, 3) + origin = torch.randn((2, 3), generator=generator, device="cuda") + # plane: (2, 3) + plane = torch.randn((2, 3), generator=generator, device="cuda") + actual = local_affine.Affine3D.from_graham_schmidt(x_axis, origin, plane) + expected = official.Affine3D.from_graham_schmidt(x_axis, origin, plane) + _assert_tensor_equal(actual.tensor, expected.tensor) + # points: (2, 3) + points = torch.randn((2, 3), generator=generator, device="cuda") + _assert_tensor_equal(actual.apply(points), expected.apply(points)) + _assert_tensor_equal(actual.scale(2.5).tensor, expected.scale(2.5).tensor) + _assert_tensor_equal( + actual.compose_rotation(actual.rot).tensor, + expected.compose_rotation(expected.rot).tensor, + ) + _assert_tensor_equal( + local_affine.Affine3D.cat([actual, actual], dim=0).tensor, + official.Affine3D.cat([expected, expected], dim=0).tensor, + ) + _assert_tensor_equal( + actual.tensor_apply(lambda component: component + 1).tensor, + expected.tensor_apply(lambda component: component + 1).tensor, + ) + + torch.manual_seed(444) + actual_random = local_affine.Affine3D.random((2, 3), device="cuda") + torch.manual_seed(444) + expected_random = official.Affine3D.random((2, 3), device="cuda") + _assert_tensor_equal(actual_random.tensor, expected_random.tensor) + + +def _normalize_entries(msa: Any) -> list[tuple[str, str]]: + return [(entry.header, entry.sequence) for entry in msa.entries] + + +def _assert_msa_equal(actual: Any, expected: Any) -> None: + assert _normalize_entries(actual) == _normalize_entries(expected) + if actual.deletions is None or expected.deletions is None: + assert actual.deletions is expected.deletions + else: + _assert_array_equal(actual.deletions, expected.deletions) + + +def test_msa_a3m_state_and_selection_match_pinned_biohub() -> None: + official = _official_msa() + a3m = ">query\nACD-EF\n>hit one\nAqCD-EF\n>hit two\nACdD-EF\n" + actual = local_msa.MSA.from_a3m(io.StringIO(a3m)) + expected = official.MSA.from_a3m(io.StringIO(a3m)) + _assert_msa_equal(actual, expected) + _assert_array_equal( + local_msa.a3m_deletion_counts("AqrC.D"), + official.a3m_deletion_counts("AqrC.D"), + ) + assert [local_msa.is_a3m_insertion(value) for value in "a.A-"] == [ + official.is_a3m_insertion(value) for value in "a.A-" + ] + assert actual.to_bytes() == expected.to_bytes() + _assert_msa_equal( + local_msa.MSA.from_bytes(actual.to_bytes()), + official.MSA.from_bytes(expected.to_bytes()), + ) + assert actual.to_sequence_bytes() == expected.to_sequence_bytes() + _assert_msa_equal(actual.select_sequences([0, 2]), expected.select_sequences([0, 2])) + _assert_msa_equal(actual.select_positions([0, 2, 4]), expected.select_positions([0, 2, 4])) + _assert_msa_equal(actual[[1, 3, 5]], expected[[1, 3, 5]]) + _assert_msa_equal(actual.pad_to_depth(5), expected.pad_to_depth(5)) + assert actual.state_dict(json_serializable=True) == expected.state_dict(json_serializable=True) + _assert_msa_equal( + local_msa.MSA.from_state_dict(actual.state_dict()), + official.MSA.from_state_dict(expected.state_dict()), + ) + + +def test_msa_composition_and_fast_representation_match_pinned_biohub() -> None: + official = _official_msa() + left_sequences = ["ACD", "A-D", "AC-"] + right_sequences = ["EF", "E-", "-F"] + # left_deletions: (3, 3) + left_deletions = np.arange(9, dtype=np.float32).reshape(3, 3) + # right_deletions: (3, 2) + right_deletions = np.arange(6, dtype=np.float32).reshape(3, 2) + actual_left = local_msa.MSA.from_state_dict( + {"sequences": left_sequences, "deletions": left_deletions} + ) + expected_left = official.MSA.from_state_dict( + {"sequences": left_sequences, "deletions": left_deletions} + ) + actual_right = local_msa.MSA.from_state_dict( + {"sequences": right_sequences, "deletions": right_deletions} + ) + expected_right = official.MSA.from_state_dict( + {"sequences": right_sequences, "deletions": right_deletions} + ) + _assert_msa_equal( + local_msa.MSA.concat([actual_left, actual_right], join_token=""), + official.MSA.concat([expected_left, expected_right], join_token=""), + ) + _assert_msa_equal( + local_msa.MSA.stack([actual_left, actual_left]), + official.MSA.stack([expected_left, expected_left]), + ) + _assert_array_equal(actual_left.seqid, expected_left.seqid) + assert repr(actual_left) == repr(expected_left) + assert local_msa.remove_insertions_from_sequence("AqCdeD") == ( + official.remove_insertions_from_sequence("AqCdeD") + ) + np.random.seed(71) + actual_random = actual_left.select_random_sequences(2) + np.random.seed(71) + expected_random = expected_left.select_random_sequences(2) + _assert_msa_equal(actual_random, expected_random) + + actual_fast = actual_left.to_fast_msa() + expected_fast = expected_left.to_fast_msa() + _assert_array_equal(actual_fast.array, expected_fast.array) + assert actual_fast.headers == expected_fast.headers + _assert_array_equal( + actual_fast.pad_to_depth(5).array, + expected_fast.pad_to_depth(5).array, + ) + actual_fast_concat = local_msa.FastMSA.concat([actual_fast, actual_fast]) + expected_fast_concat = official.FastMSA.concat([expected_fast, expected_fast]) + _assert_array_equal(actual_fast_concat.array, expected_fast_concat.array) + assert actual_fast_concat.headers == expected_fast_concat.headers + assert _normalize_entries(actual_fast.to_msa()) == _normalize_entries(expected_fast.to_msa()) + actual_fast_bytes = local_msa.FastMSA.from_bytes(actual_left.to_bytes()) + expected_fast_bytes = official.FastMSA.from_bytes(expected_left.to_bytes()) + _assert_array_equal(actual_fast_bytes.array, expected_fast_bytes.array) + assert actual_fast_bytes.headers == expected_fast_bytes.headers + actual_sequence_only = local_msa.FastMSA.from_sequence_bytes(actual_left.to_sequence_bytes()) + expected_sequence_only = official.FastMSA.from_sequence_bytes(expected_left.to_sequence_bytes()) + _assert_array_equal(actual_sequence_only.array, expected_sequence_only.array) + + short_actual = local_msa.MSA.from_sequences(["AC", "A-"]) + short_expected = official.MSA.from_sequences(["AC", "A-"]) + _assert_msa_equal( + local_msa.MSA.concat([actual_left, short_actual], allow_depth_mismatch=True), + official.MSA.concat([expected_left, short_expected], allow_depth_mismatch=True), + ) + + +class _FakeColumn: + def __init__(self, values, mask=None): + self.values = np.asarray(values) + self.mask = mask + + def as_array(self, dtype): + return self.values.astype(dtype) + + def as_item(self): + return self.values.item() + + +def _category(**columns): + return {name: _FakeColumn(values) for name, values in columns.items()} + + +def _structure() -> bs.AtomArray: + atoms = bs.AtomArray(7) + # coord: (7, 3) + atoms.coord = np.arange(21, dtype=np.float32).reshape(7, 3) + atoms.chain_id = np.asarray(["A", "A", "A", "B", "B", "L", "L"]) + atoms.res_id = np.asarray([10, 10, 11, 5, 6, 1, 1]) + atoms.res_name = np.asarray(["ALA", "ALA", "CYS", "GLY", "SER", "ATP", "ATP"]) + atoms.atom_name = np.asarray(["N", "CA", "N", "N", "N", "P", "O1"]) + atoms.element = np.asarray(["N", "C", "N", "N", "N", "P", "O"]) + atoms.hetero = np.asarray([False, False, False, False, False, True, True]) + return atoms + + +def _fake_block() -> dict[str, Any]: + return { + "pdbx_database_status": _category(recvd_initial_deposition_date="2025-06-04"), + "refine": _category(ls_d_res_high="1.75"), + "exptl": _category(method="electron microscopy"), + "entity": _category( + id=["1", "2", "3"], + type=["polymer", "polymer", "non-polymer"], + ), + "entity_poly": _category( + entity_id=["1", "2"], + pdbx_seq_one_letter_code_can=["AC D", "GS"], + pdbx_strand_id=["A", "B"], + ), + "struct_asym": _category(id=["A", "B", "L"], entity_id=["1", "2", "3"]), + "pdbx_poly_seq_scheme": _category( + asym_id=["A", "A", "A", "B", "B"], + seq_id=["1", "2", "3", "1", "2"], + auth_seq_num=["10", "10", "?", "5", "6"], + pdb_ins_code=[".", "A", "?", ".", "."], + hetflag=["N", "Y", "N", "N", "N"], + pdb_strand_id=["A", "A", "A", "B", "B"], + ), + "pdbx_entity_nonpoly": _category(entity_id=["3"], comp_id=["ATP"]), + "atom_site": _category( + label_asym_id=["A", "A", "A", "B", "B", "L", "L"], + label_entity_id=["1", "1", "1", "2", "2", "3", "3"], + label_comp_id=["ALA", "ALA", "CYS", "GLY", "SER", "ATP", "ATP"], + ), + } + + +def _normalized_mapping(wrapper: Any): + return { + chain: { + index: (residue.residue_number, residue.insertion_code, residue.hetflag) + for index, residue in mapping.items() + } + for chain, mapping in wrapper.seqres_to_structure.items() + } + + +def test_mmcif_metadata_sequence_and_ligand_contracts_match_pinned_biohub() -> None: + official = _official_mmcif() + block = _fake_block() + raw = types.SimpleNamespace(block=block) + actual = local_mmcif.MmcifWrapper("case") + expected = official.MmcifWrapper("case") + actual.raw = raw + expected.raw = raw + actual.structure = _structure() + expected.structure = _structure() + for wrapper in (actual, expected): + wrapper._parse_header() + wrapper._parse_entities() + wrapper._parse_sequences() + assert ( + actual.header.release_date, + actual.header.resolution, + actual.header.structure_method, + ) == ( + expected.header.release_date, + expected.header.resolution, + expected.header.structure_method, + ) + assert actual.entities == expected.entities + assert actual.chain_to_seqres == expected.chain_to_seqres + assert _normalized_mapping(actual) == _normalized_mapping(expected) + actual_ligands = actual._parse_nonpoly_from_mmcif() + expected_ligands = expected._parse_nonpoly_from_mmcif() + assert actual_ligands.keys() == expected_ligands.keys() + for key in actual_ligands: + _assert_array_equal(actual_ligands[key].coord, expected_ligands[key].coord) + actual_fallback = actual._parse_nonpoly_fallback() + expected_fallback = expected._parse_nonpoly_fallback() + assert actual_fallback.keys() == expected_fallback.keys() + for key in actual_fallback: + _assert_array_equal(actual_fallback[key].coord, expected_fallback[key].coord) + + +def test_mmcif_rounding_matches_pinned_biohub() -> None: + official = _official_mmcif() + columns = { + "Cartn_x": _FakeColumn([1.23456, -2.34567]), + "Cartn_y": _FakeColumn([3.45678, 4.56789]), + "Cartn_z": _FakeColumn([5.67891, 6.78912]), + "B_iso_or_equiv": _FakeColumn([92.345, 10.005]), + "label_atom_id": _FakeColumn(["CA", "N"]), + } + actual_file = types.SimpleNamespace(block={"atom_site": dict(columns)}) + expected_file = types.SimpleNamespace(block={"atom_site": dict(columns)}) + local_mmcif.round_mmcif_columns(actual_file) + official.round_mmcif_columns(expected_file) + for name in ("Cartn_x", "Cartn_y", "Cartn_z", "B_iso_or_equiv"): + _assert_array_equal( + actual_file.block["atom_site"][name].as_array(str), + expected_file.block["atom_site"][name].as_array(str), + ) + actual_empty = types.SimpleNamespace(block={}) + expected_empty = types.SimpleNamespace(block={}) + local_mmcif.round_mmcif_columns(actual_empty) + official.round_mmcif_columns(expected_empty) + assert actual_empty.block == expected_empty.block diff --git a/tests/parity/test_esmfold2_source_slice4_parity.py b/tests/parity/test_esmfold2_source_slice4_parity.py new file mode 100644 index 0000000..edf2225 --- /dev/null +++ b/tests/parity/test_esmfold2_source_slice4_parity.py @@ -0,0 +1,420 @@ +"""Pinned Biohub parity for ESMFold2 input preparation and flat complexes.""" + +from __future__ import annotations + +import importlib.util +import sys +import types +import numpy as np +import pytest +import torch +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +from fastplms.models.esmfold2 import esmfold2_conformers as local_conformers +from fastplms.models.esmfold2 import esmfold2_constants as local_constants +from fastplms.models.esmfold2 import esmfold2_metrics as local_metrics +from fastplms.models.esmfold2 import esmfold2_mmcif_parsing as local_mmcif +from fastplms.models.esmfold2 import esmfold2_molecular_complex as local_complex +from fastplms.models.esmfold2 import esmfold2_paired_msa as local_paired_msa +from fastplms.models.esmfold2 import esmfold2_prepare_input as local_prepare +from fastplms.models.esmfold2 import esmfold2_protein_complex as local_protein_complex +from fastplms.models.esmfold2 import esmfold2_residue_constants as local_residues +from fastplms.models.esmfold2 import esmfold2_types as local_types + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.structure] + +ROOT = Path(__file__).resolve().parents[2] +BIOHUB_ESM = ROOT / "vendor/upstream/biohub-esm/esm" +_MISSING = object() + + +def _package(name: str) -> types.ModuleType: + package = types.ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + return package + + +@contextmanager +def _temporary_modules(modules: dict[str, types.ModuleType]) -> Iterator[None]: + previous = {name: sys.modules.get(name, _MISSING) for name in modules} + sys.modules.update(modules) + try: + yield + finally: + for name, module in previous.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module # type: ignore[assignment] + + +def _load_source( + module_name: str, + path: Path, + aliases: dict[str, types.ModuleType], +) -> types.ModuleType: + assert path.is_file(), f"pinned source is missing: {path}" + specification = importlib.util.spec_from_file_location(module_name, path) + assert specification is not None and specification.loader is not None + module = importlib.util.module_from_spec(specification) + with _temporary_modules({**aliases, module_name: module}): + specification.loader.exec_module(module) + return module + + +def _biohub_packages() -> dict[str, types.ModuleType]: + return { + "esm": _package("esm"), + "esm.models": _package("esm.models"), + "esm.models.esmfold2": _package("esm.models.esmfold2"), + "esm.utils": _package("esm.utils"), + "esm.utils.structure": _package("esm.utils.structure"), + } + + +def _prepare_aliases() -> dict[str, types.ModuleType]: + return { + **_biohub_packages(), + "esm.models.esmfold2.conformers": local_conformers, + "esm.models.esmfold2.constants": local_constants, + "esm.models.esmfold2.paired_msa": local_paired_msa, + "esm.models.esmfold2.types": local_types, + } + + +def _official_prepare() -> types.ModuleType: + return _load_source( + "_fastplms_pinned_biohub_prepare_input", + BIOHUB_ESM / "models/esmfold2/prepare_input.py", + _prepare_aliases(), + ) + + +def _official_complex() -> types.ModuleType: + aliases = _complex_aliases() + return _load_source( + "_fastplms_pinned_biohub_molecular_complex", + BIOHUB_ESM / "utils/structure/molecular_complex.py", + aliases, + ) + + +def _complex_aliases() -> dict[str, types.ModuleType]: + return { + **_biohub_packages(), + "esm.utils.residue_constants": local_residues, + "esm.utils.structure.metrics": local_metrics, + "esm.utils.structure.mmcif_parsing": local_mmcif, + "esm.utils.structure.protein_complex": local_protein_complex, + } + + +def _assert_value_equal(actual: Any, expected: Any) -> None: + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + elif isinstance(actual, np.ndarray): + np.testing.assert_array_equal(actual, expected) + else: + assert actual == expected + + +def _assert_records_equal(actual: Any, expected: Any) -> None: + assert vars(actual).keys() == vars(expected).keys() + for field_name, actual_value in vars(actual).items(): + _assert_value_equal(actual_value, vars(expected)[field_name]) + + +def _assert_prepared_equal( + actual: tuple[list[Any], list[Any], list[Any]], + expected: tuple[list[Any], list[Any], list[Any]], +) -> None: + actual_chains, actual_tokens, actual_atoms = actual + expected_chains, expected_tokens, expected_atoms = expected + assert len(actual_chains) == len(expected_chains) + assert len(actual_tokens) == len(expected_tokens) + assert len(actual_atoms) == len(expected_atoms) + for actual_token, expected_token in zip(actual_tokens, expected_tokens, strict=True): + _assert_records_equal(actual_token, expected_token) + for actual_atom, expected_atom in zip(actual_atoms, expected_atoms, strict=True): + _assert_records_equal(actual_atom, expected_atom) + for actual_chain, expected_chain in zip(actual_chains, expected_chains, strict=True): + for field_name in ( + "chain_id", + "asym_id", + "entity_id", + "sym_id", + "mol_type", + "ligand_bonds", + ): + _assert_value_equal( + getattr(actual_chain, field_name), getattr(expected_chain, field_name) + ) + assert [vars(token) for token in actual_chain.tokens] == [ + vars(token) for token in expected_chain.tokens + ] + + +def _install_fake_ccd(monkeypatch: pytest.MonkeyPatch, official: types.ModuleType) -> None: + atom_records = { + "MSE": [ + ("N", "N", 0), + ("CA", "C", 0), + ("C", "C", 0), + ("O", "O", 0), + ("SE", "Se", 0), + ], + "PSU": [("P", "P", 0), ("C1'", "C", 0), ("N1", "N", 0)], + "LIG": [("C1", "C", 0), ("N1", "N", 1), ("O1", "O", -1)], + } + bonds = { + "MSE": [("N", "CA"), ("CA", "C"), ("C", "O"), ("CA", "SE")], + "PSU": [("P", "C1'"), ("C1'", "N1")], + "LIG": [("C1", "N1"), ("N1", "O1")], + } + + def idealized(residue_type: int, atom_name: str) -> np.ndarray: + base = residue_type + sum(map(ord, atom_name)) / 1000 + return np.asarray([base, base + 1, base + 2], dtype=np.float32) + + def ligand_position(residue_name: str, atom_name: str) -> np.ndarray: + base = (sum(map(ord, residue_name + atom_name)) % 97) / 10 + return np.asarray([base, base + 0.5, base + 1], dtype=np.float32) + + replacements = { + "get_idealized_atom_pos": idealized, + "get_ligand_idealized_atom_pos": ligand_position, + "get_ligand_ccd_atoms_with_charges": atom_records.get, + "get_ligand_ccd_bonds": bonds.get, + "get_ccd_leaving_atoms": lambda name: {"O1"} if name == "LIG" else set(), + } + for name, replacement in replacements.items(): + monkeypatch.setattr(local_prepare, name, replacement) + monkeypatch.setattr(official, name, replacement) + + +def _mixed_input() -> local_types.StructurePredictionInput: + msa = local_types.MSA.from_sequences(["AMG", "A-G"]) + return local_types.StructurePredictionInput( + sequences=[ + local_types.ProteinInput( + id=["A", "B"], + sequence="AMG", + modifications=[local_types.Modification(position=1, ccd="MSE")], + msa=msa, + ), + local_types.DNAInput(id="C", sequence="ATN"), + local_types.RNAInput( + id="D", + sequence="AUN", + modifications=[local_types.Modification(position=1, ccd="PSU")], + ), + local_types.LigandInput(id="E", ccd=["LIG"]), + ], + covalent_bonds=[ + local_types.CovalentBond( + chain_id1="A", + res_idx1=1, + atom_idx1=0, + chain_id2="E", + res_idx2=0, + atom_idx2=0, + ) + ], + ) + + +def test_mixed_input_pipeline_matches_pinned_biohub( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert torch.cuda.is_available(), "the ESMFold2 parity suite requires CUDA" + official = _official_prepare() + _install_fake_ccd(monkeypatch, official) + input_value = _mixed_input() + actual_parts = local_prepare.build_chains_from_input(input_value, seed=71) + expected_parts = official.build_chains_from_input(input_value, seed=71) + _assert_prepared_equal(actual_parts, expected_parts) + with _temporary_modules(_prepare_aliases()): + actual = local_prepare.build_feature_tensors(*actual_parts, input_value) + expected = official.build_feature_tensors(*expected_parts, input_value) + assert actual.keys() == expected.keys() + for feature_name in actual: + _assert_value_equal(actual[feature_name], expected[feature_name]) + + +def test_distogram_conditioning_matches_pinned_biohub( + monkeypatch: pytest.MonkeyPatch, +) -> None: + official = _official_prepare() + _install_fake_ccd(monkeypatch, official) + msa = local_types.MSA.from_sequences(["ACD"]) + input_value = local_types.StructurePredictionInput( + sequences=[local_types.ProteinInput(id="A", sequence="ACD", msa=msa)], + distogram_conditioning=[ + local_types.DistogramConditioning( + chain_id="A", + distogram=np.asarray( + [[0.0, 4.0, 12.0], [4.0, 0.0, 22.0], [12.0, 22.0, 0.0]], + dtype=np.float32, + ), + ) + ], + ) + actual_parts = local_prepare.build_chains_from_input(input_value) + expected_parts = official.build_chains_from_input(input_value) + actual = local_prepare.compute_distogram_conditioning( + input_value, actual_parts[0], actual_parts[1], torch.zeros(3, 3) + ) + expected = official.compute_distogram_conditioning( + input_value, expected_parts[0], expected_parts[1], torch.zeros(3, 3) + ) + for actual_tensor, expected_tensor in zip(actual, expected, strict=True): + _assert_value_equal(actual_tensor, expected_tensor) + + +def test_smiles_topology_matches_pinned_biohub() -> None: + official = _official_prepare() + arguments = { + "smiles": "CC(=O)N", + "entity_id": 2, + "asym_id": 3, + "sym_id": 0, + "token_offset": 7, + "atom_offset": 19, + "space_uid_offset": 5, + "seed": 919, + } + actual_tokens, actual_atoms, actual_bonds = local_prepare.tokenize_ligand_smiles(**arguments) + expected_tokens, expected_atoms, expected_bonds = official.tokenize_ligand_smiles(**arguments) + for actual, expected in zip(actual_tokens, expected_tokens, strict=True): + _assert_records_equal(actual, expected) + for actual, expected in zip(actual_atoms, expected_atoms, strict=True): + _assert_records_equal(actual, expected) + assert actual_bonds == expected_bonds + + +def _protein_fixture() -> local_protein_complex.ProteinComplex: + sequence = "AC|GG" + n_positions = len(sequence) + # positions: (n_positions, 37, 3) + positions = np.full((n_positions, 37, 3), np.nan, dtype=np.float32) + # mask: (n_positions, 37) + mask = np.zeros((n_positions, 37), dtype=bool) + atom_names = ("N", "CA", "C", "O", "CB", "SG") + for sequence_index in (0, 1, 3, 4): + for atom_offset, atom_name in enumerate(atom_names): + atom_index = local_residues.atom_order[atom_name] + positions[sequence_index, atom_index] = np.asarray( + [sequence_index, atom_offset, sequence_index + atom_offset / 10], + dtype=np.float32, + ) + mask[sequence_index, atom_index] = True + return local_protein_complex.ProteinComplex( + id="fixture", + sequence=sequence, + entity_id=np.asarray([0, 0, -1, 1, 1], dtype=np.int64), + chain_id=np.asarray([0, 0, -1, 1, 1], dtype=np.int64), + sym_id=np.zeros(n_positions, dtype=np.int64), + residue_index=np.asarray([1, 2, 0, 1, 2], dtype=np.int64), + insertion_code=np.asarray([""] * n_positions, dtype=object), + atom37_positions=positions, + atom37_mask=mask, + confidence=np.asarray([0.8, 0.7, 0.0, 0.9, 0.6], dtype=np.float32), + metadata=local_protein_complex.ProteinComplexMetadata( + entity_lookup={0: 0, 1: 1}, + chain_lookup={0: "A", 1: "ligand_1"}, + assembly_composition={"1": ["A", "ligand_1"]}, + ), + ) + + +def _assert_complex_equal(actual: Any, expected: Any) -> None: + assert actual.id == expected.id + assert actual.sequence == expected.sequence + for field_name in ( + "atom_positions", + "atom_elements", + "token_to_atoms", + "chain_id", + "plddt", + "atom_names", + "atom_hetero", + ): + _assert_value_equal(getattr(actual, field_name), getattr(expected, field_name)) + assert vars(actual.metadata) == vars(expected.metadata) + + +def _assert_protein_equal(actual: Any, expected: Any) -> None: + assert actual.id == expected.id + assert actual.sequence == expected.sequence + for field_name in ( + "entity_id", + "chain_id", + "sym_id", + "residue_index", + "insertion_code", + "atom37_positions", + "atom37_mask", + "confidence", + ): + _assert_value_equal(getattr(actual, field_name), getattr(expected, field_name)) + assert actual.metadata.entity_lookup == expected.metadata.entity_lookup + assert actual.metadata.chain_lookup == expected.metadata.chain_lookup + assert actual.metadata.assembly_composition == expected.metadata.assembly_composition + + +def test_molecular_complex_conversion_extends_pinned_biohub_with_identity_storage() -> None: + official = _official_complex() + protein = _protein_fixture() + actual = local_complex.MolecularComplex.from_protein_complex(protein) + with _temporary_modules(_complex_aliases()): + expected = official.MolecularComplex.from_protein_complex(protein) + _assert_complex_equal(actual, expected) + _assert_records_equal(actual[1], expected[1]) + with _temporary_modules(_complex_aliases()): + expected_protein = expected.to_protein_complex() + _assert_protein_equal(actual.to_protein_complex(), expected_protein) + + actual_mmcif = actual.to_mmcif() + expected_mmcif = expected.to_mmcif() + assert actual_mmcif == expected_mmcif + _assert_complex_equal( + local_complex.MolecularComplex.from_mmcif(expected_mmcif, id="roundtrip"), + official.MolecularComplex.from_mmcif(expected_mmcif, id="roundtrip"), + ) + actual_state = actual.state_dict() + identity_state = {key: actual_state.pop(key) for key in ("entity_id", "sym_id")} + assert actual_state == expected.state_dict() + assert identity_state == { + "entity_id": [0, 0, 1, 1], + "sym_id": [0, 0, 0, 0], + } + restored = local_complex.MolecularComplex.from_blob(actual.to_blob()) + _assert_complex_equal(restored, actual) + _assert_value_equal(restored.entity_id, actual.entity_id) + _assert_value_equal(restored.sym_id, actual.sym_id) + _assert_complex_equal( + local_complex.MolecularComplex.from_blob(expected.to_blob()), + official.MolecularComplex.from_blob(expected.to_blob()), + ) + + +def test_molecular_complex_metrics_match_pinned_biohub() -> None: + official = _official_complex() + protein = _protein_fixture() + actual = local_complex.MolecularComplex.from_protein_complex(protein) + with _temporary_modules(_complex_aliases()): + expected = official.MolecularComplex.from_protein_complex(protein) + shifted_positions = actual.atom_positions.copy() + shifted_positions[:, 0] += np.linspace(0.0, 0.2, len(shifted_positions)) + actual_target = local_complex.MolecularComplex( + **{**vars(actual), "atom_positions": shifted_positions} + ) + expected_target = official.MolecularComplex( + **{**vars(expected), "atom_positions": shifted_positions} + ) + assert actual.rmsd(actual_target) == expected.rmsd(expected_target) + assert actual.lddt_ca(actual_target) == expected.lddt_ca(expected_target) diff --git a/tests/parity/test_model_parity.py b/tests/parity/test_model_parity.py new file mode 100644 index 0000000..bfc94ae --- /dev/null +++ b/tests/parity/test_model_parity.py @@ -0,0 +1,1045 @@ +"""Strict manifest-driven equivalence against pinned official implementations. + +Configuration, tokenizer behavior, state, aliases, and inference are release +gates. Numeric contracts are shared except for explicit, evidence-backed +model/backend calibrations; this suite never silently falls back. +""" + +from __future__ import annotations + +import contextlib +import gc +import importlib +import os +import random +import warnings +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from transformers import AutoModel, AutoModelForMaskedLM + +from fastplms.registry import ModelSpec, get_model_registry +from tests.conftest import CANONICAL_AAS, SEED, strict_fp32_matmul +from tests.parity.support.semantic_config import ( + semantic_config, + transformed_semantic_config, +) +from tests.parity.support.state_transforms import ( + TRANSFORMS, + transform_parameter_names, + transform_preserves_aliases, + transform_state, +) + + +_semantic_config = semantic_config + +pytestmark = pytest.mark.compliance + +REGISTRY = get_model_registry() +SEQUENCE_SPECS = tuple( + spec for spec in REGISTRY.values() if spec.family.tokenizer_mode != "structure" +) +DEEP_SPECS = tuple(spec for spec in SEQUENCE_SPECS if spec.is_deep_reference) +MIXED_LENGTHS = (61, 29, 13) +EDGE_SEQUENCES = ( + "ACDEFGHIKLMNPQRSTVWY", + "AXBJOUZ", + "acdefghik", + "A C\nD\tE", + "", +) + + +@dataclass(frozen=True, slots=True) +class NumericContract: + """Fixed engineering target and hard release boundary.""" + + relative_l2_target: float + relative_l2_hard: float + relative_q999_target: float + relative_q999_hard: float + residue_cosine_target: float + residue_cosine_hard: float + pooled_cosine_target: float + pooled_cosine_hard: float + top1_target: float + top1_hard: float + jsd_target: float + jsd_hard: float + + +FP32_CONTRACT = NumericContract( + relative_l2_target=2e-6, + relative_l2_hard=2e-5, + relative_q999_target=1e-5, + relative_q999_hard=1e-4, + residue_cosine_target=0.999999, + residue_cosine_hard=0.9999, + pooled_cosine_target=0.999999, + pooled_cosine_hard=0.9999, + top1_target=0.9999, + top1_hard=0.999, + jsd_target=1e-8, + jsd_hard=1e-6, +) +BF16_CONTRACT = NumericContract( + relative_l2_target=1e-2, + relative_l2_hard=3e-2, + relative_q999_target=2.5e-2, + relative_q999_hard=5e-2, + residue_cosine_target=0.999, + residue_cosine_hard=0.995, + pooled_cosine_target=0.9995, + pooled_cosine_hard=0.995, + top1_target=0.995, + top1_hard=0.99, + jsd_target=1e-4, + jsd_hard=1e-3, +) +ESMC_ALTERNATE_BF16_CONTRACT = NumericContract( + relative_l2_target=2.9e-2, + relative_l2_hard=BF16_CONTRACT.relative_l2_hard, + relative_q999_target=4.9e-2, + relative_q999_hard=BF16_CONTRACT.relative_q999_hard, + residue_cosine_target=0.997, + residue_cosine_hard=BF16_CONTRACT.residue_cosine_hard, + pooled_cosine_target=BF16_CONTRACT.pooled_cosine_target, + pooled_cosine_hard=BF16_CONTRACT.pooled_cosine_hard, + top1_target=BF16_CONTRACT.top1_target, + top1_hard=BF16_CONTRACT.top1_hard, + jsd_target=4e-4, + jsd_hard=BF16_CONTRACT.jsd_hard, +) +# Flex and FA3 are supported ESMC implementations whose backend-specific BF16 +# arithmetic is reported diagnostically. These deliberately broad limits catch +# corrupt dispatch, broken masking, non-finite outputs, or catastrophic +# biological disagreement without turning known backend drift into an xfail. +ESMC_CATASTROPHIC_BF16_CONTRACT = NumericContract( + relative_l2_target=0.25, + relative_l2_hard=0.25, + relative_q999_target=0.50, + relative_q999_hard=0.50, + residue_cosine_target=0.90, + residue_cosine_hard=0.90, + pooled_cosine_target=0.95, + pooled_cosine_hard=0.95, + top1_target=0.80, + top1_hard=0.80, + jsd_target=0.05, + jsd_hard=0.05, +) +ESM2_OPTIMIZED_BF16_CONTRACT = NumericContract( + relative_l2_target=2e-2, + relative_l2_hard=BF16_CONTRACT.relative_l2_hard, + relative_q999_target=BF16_CONTRACT.relative_q999_target, + relative_q999_hard=BF16_CONTRACT.relative_q999_hard, + residue_cosine_target=BF16_CONTRACT.residue_cosine_target, + residue_cosine_hard=BF16_CONTRACT.residue_cosine_hard, + pooled_cosine_target=BF16_CONTRACT.pooled_cosine_target, + pooled_cosine_hard=BF16_CONTRACT.pooled_cosine_hard, + top1_target=BF16_CONTRACT.top1_target, + top1_hard=BF16_CONTRACT.top1_hard, + jsd_target=BF16_CONTRACT.jsd_target, + jsd_hard=BF16_CONTRACT.jsd_hard, +) +ESM2_3B_SDPA_BF16_CONTRACT = NumericContract( + # Calibrated on the pinned 3B checkpoint: exact weights and logits retain + # perfect confident-token agreement while deep BF16 SDPA layers accumulate + # more rounding drift than the smaller ESM2 variants. + relative_l2_target=6e-2, + relative_l2_hard=7e-2, + relative_q999_target=1.5e-1, + relative_q999_hard=1.8e-1, + residue_cosine_target=0.994, + residue_cosine_hard=0.992, + pooled_cosine_target=0.998, + pooled_cosine_hard=0.997, + top1_target=BF16_CONTRACT.top1_target, + top1_hard=BF16_CONTRACT.top1_hard, + jsd_target=BF16_CONTRACT.jsd_target, + jsd_hard=BF16_CONTRACT.jsd_hard, +) + + +@dataclass(frozen=True, slots=True) +class TensorMetrics: + """Normalized metrics over biological residues only.""" + + relative_l2: float + relative_q999: float + residue_cosine_p01: float + pooled_cosine_min: float + + +@dataclass(frozen=True, slots=True) +class TensorMetricRecord: + """Metrics and identity for one layer or output tensor.""" + + context: str + metrics: TensorMetrics + + +@dataclass(frozen=True, slots=True) +class LogitsMetrics: + """Distribution-level metrics for a masked-language-model head.""" + + confident_top1_agreement: float + mean_jsd: float + + +def _numeric_contract( + spec: ModelSpec, + dtype: torch.dtype, + backend: str | None, +) -> NumericContract: + """Resolve the fixed contract without weakening the global BF16 policy.""" + + if dtype == torch.float32: + return FP32_CONTRACT + if dtype != torch.bfloat16: + raise ValueError(f"Unsupported parity dtype: {dtype}") + if spec.id == "esm2_3b" and backend in (None, "sdpa"): + return ESM2_3B_SDPA_BF16_CONTRACT + if spec.family.id == "esm2" and backend != "eager": + return ESM2_OPTIMIZED_BF16_CONTRACT + if spec.family.architecture == "ESMC" and backend not in (None, "sdpa"): + return ESMC_ALTERNATE_BF16_CONTRACT + return BF16_CONTRACT + + +def _parameter(spec: ModelSpec) -> Any: + marks: list[Any] = [pytest.mark.slow] + if spec.size_category == "xlarge": + marks.append(pytest.mark.large) + return pytest.param(spec, id=spec.id, marks=marks) + + +def _deep_parameter(spec: ModelSpec, dtype: torch.dtype, backend: str) -> Any: + marks: list[Any] = [pytest.mark.gpu, pytest.mark.slow] + if spec.size_category == "xlarge": + marks.append(pytest.mark.large) + dtype_name = "fp32" if dtype == torch.float32 else "bf16" + return pytest.param(spec, dtype, backend, id=f"{spec.id}-{dtype_name}-{backend}", marks=marks) + + +def _sequence_batch(lengths: Sequence[int] = MIXED_LENGTHS) -> list[str]: + rng = random.Random(SEED) + return ["M" + "".join(rng.choices(CANONICAL_AAS, k=length - 1)) for length in lengths] + + +def _load_fast( + spec: ModelSpec, + device: torch.device, + dtype: torch.dtype | None, +) -> nn.Module: + # ANKH parity is the official encoder contract. Its masked-LM class is a + # separately named FastPLMs extension and is not presented as upstream-equivalent. + auto_class_name = ( + "AutoModel" + if spec.family.id == "ankh" or "AutoModelForMaskedLM" not in spec.auto_map + else "AutoModelForMaskedLM" + ) + auto_class = AutoModel if auto_class_name == "AutoModel" else AutoModelForMaskedLM + artifact_root = os.environ.get("FASTPLMS_CANDIDATE_ARTIFACTS") + if artifact_root: + repository_name = spec.fast.repo_id.split("/", maxsplit=1)[-1] + model_source = Path(artifact_root) / repository_name + if not model_source.is_dir(): + raise FileNotFoundError( + f"Candidate compliance artifact is missing for {spec.id}: {model_source}" + ) + # Native compliance already imports the current package to read the + # typed manifest. Loading the generated remote-code bridge in this same + # interpreter would deliberately fail its runtime-isolation guard. + # Resolve the manifest-declared package classes directly while keeping + # the artifact as the sole config, tokenizer-asset, and weight source. + config_path = spec.auto_map["AutoConfig"] + model_path = spec.auto_map[auto_class_name] + config_module, config_name = config_path.rsplit(".", maxsplit=1) + model_module, model_name = model_path.rsplit(".", maxsplit=1) + config_class = getattr(importlib.import_module(config_module), config_name) + model_class = getattr(importlib.import_module(model_module), model_name) + config = config_class.from_pretrained(model_source, local_files_only=True) + load_kwargs = { + "config": config, + "local_files_only": True, + "device_map": device, + } + else: + model_source = spec.fast.repo_id + load_kwargs = { + "revision": spec.fast.revision, + "trust_remote_code": True, + "device_map": device, + } + if dtype is not None: + load_kwargs["dtype"] = dtype + loader = model_class if artifact_root else auto_class + model = loader.from_pretrained(model_source, **load_kwargs) + return model.eval() + + +def _load_reference( + spec: ModelSpec, + device: torch.device, + dtype: torch.dtype | None, +) -> tuple[nn.Module, object]: + adapter = importlib.import_module(spec.family.reference_adapter) + kwargs: dict[str, Any] = {} + if spec.oracle_assets: + kwargs["oracle_assets"] = spec.oracle_assets + return adapter.load_official_model( + reference_repo_id=spec.official.repo_id, + reference_revision=spec.official.revision, + device=device, + dtype=dtype, + **kwargs, + ) + + +def _reference_core(reference: nn.Module) -> nn.Module: + core = getattr(reference, "model", reference) + return core + + +def _assert_semantic_config_equal(spec: ModelSpec, fast: nn.Module, reference: nn.Module) -> None: + fast_config = _semantic_config(fast) + reference_config = transformed_semantic_config( + _reference_core(reference), spec.family.state_transform + ) + missing = sorted(set(reference_config).difference(fast_config)) + assert not missing, ( + f"{spec.id}: FastPLMs configuration omits official semantic fields {missing}" + ) + compared_fast = {name: fast_config[name] for name in reference_config} + assert compared_fast == reference_config, ( + f"{spec.id}: semantic configuration differs: " + f"fast={compared_fast}, official={reference_config}" + ) + + +def _assert_state_equal(spec: ModelSpec, fast: nn.Module, reference: nn.Module) -> None: + official_state = transform_state( + spec.family.state_transform, + _reference_core(reference).state_dict(), + ) + fast_state = fast.state_dict() + assert set(fast_state) == set(official_state), ( + f"{spec.id}: state-key set differs; " + f"only_fast={sorted(set(fast_state) - set(official_state))[:20]}, " + f"only_official={sorted(set(official_state) - set(fast_state))[:20]}" + ) + for name in sorted(fast_state): + # candidate: (...) + candidate = fast_state[name].detach().cpu() + # official: (...) + official = official_state[name].detach().cpu() + assert candidate.shape == official.shape, ( + f"{spec.id}:{name}: shape {tuple(candidate.shape)} != {tuple(official.shape)}" + ) + assert candidate.dtype == official.dtype, ( + f"{spec.id}:{name}: dtype {candidate.dtype} != {official.dtype}" + ) + assert torch.equal(candidate, official), f"{spec.id}:{name}: tensor values are not exact" + + +def _alias_groups(model: nn.Module) -> set[frozenset[str]]: + by_parameter: dict[int, set[str]] = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + by_parameter.setdefault(id(parameter), set()).add(name) + return {frozenset(names) for names in by_parameter.values() if len(names) > 1} + + +def _transformed_alias_groups(spec: ModelSpec, model: nn.Module) -> set[frozenset[str]]: + if not transform_preserves_aliases(spec.family.state_transform): + return set() + by_parameter: dict[int, set[str]] = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + mapped = transform_parameter_names(spec.family.state_transform, name) + by_parameter.setdefault(id(parameter), set()).update(mapped) + return {frozenset(names) for names in by_parameter.values() if len(names) > 1} + + +def _assert_aliases_equal(spec: ModelSpec, fast: nn.Module, reference: nn.Module) -> None: + candidate = _alias_groups(fast) + official = _transformed_alias_groups(spec, _reference_core(reference)) + assert candidate == official, ( + f"{spec.id}: tied-parameter aliases differ; " + f"only_fast={sorted(map(sorted, candidate - official))}, " + f"only_official={sorted(map(sorted, official - candidate))}" + ) + + +def _normalize_tokenizer_error(message: str) -> str: + """Remove a dependency-list difference between Transformers v4 and v5.""" + + return message.replace( + "python, numpy, pytorch or tensorflow object.", + "python, numpy or pytorch object.", + ).replace("python, numpy, or pytorch object.", "python, numpy or pytorch object.") + + +def _token_result(tokenizer: object, sequences: Sequence[str], **kwargs: Any) -> Any: + try: + encoded = tokenizer(sequences, return_tensors="pt", **kwargs) + except Exception as error: # Exact error behavior is part of the token contract. + return ( + "error", + type(error).__module__, + type(error).__qualname__, + _normalize_tokenizer_error(str(error)), + ) + normalized: dict[str, Any] = {} + for key, value in encoded.items(): + normalized[key] = value.tolist() if torch.is_tensor(value) else value + return ("ok", normalized) + + +def _assert_tokenizer_equal( + spec: ModelSpec, + fast_tokenizer: object, + official_tokenizer: object, +) -> None: + assert fast_tokenizer.get_vocab() == official_tokenizer.get_vocab(), ( + f"{spec.id}: tokenizer vocabulary or token IDs differ" + ) + for name in ( + "pad_token_id", + "bos_token_id", + "cls_token_id", + "eos_token_id", + "mask_token_id", + "unk_token_id", + ): + assert getattr(fast_tokenizer, name, None) == getattr(official_tokenizer, name, None), ( + f"{spec.id}: tokenizer {name} differs" + ) + + settings = ( + {"padding": True}, + {"padding": "max_length", "truncation": True, "max_length": 12}, + {"padding": True, "truncation": True, "max_length": 5}, + ) + for options in settings: + fast_result = _token_result(fast_tokenizer, EDGE_SEQUENCES, **options) + official_result = _token_result(official_tokenizer, EDGE_SEQUENCES, **options) + assert fast_result == official_result, ( + f"{spec.id}: tokenizer behavior differs for options={options}: " + f"fast={fast_result}, official={official_result}" + ) + + +def _to_device(values: Mapping[str, Any], device: torch.device) -> dict[str, torch.Tensor]: + return {name: value.to(device) for name, value in values.items() if torch.is_tensor(value)} + + +def _prepare_inputs( + spec: ModelSpec, + fast: nn.Module, + reference_tokenizer: object, + sequences: Sequence[str], + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor], torch.Tensor]: + if spec.family.tokenizer_mode == "sequence": + fast_batch = fast.model.prep_tokens.get_batch_kwargs(sequences, device=device) + official_batch = reference_tokenizer.get_batch_kwargs(sequences, device=device) + assert set(fast_batch) == set(official_batch), f"{spec.id}: sequence-adapter keys differ" + for name in fast_batch: + assert torch.equal(fast_batch[name], official_batch[name]), ( + f"{spec.id}: sequence-adapter tensor {name!r} differs" + ) + # residue_mask: (b, l) + residue_mask = fast_batch["sequence_ids"].ge(0) + fast_inputs = dict(fast_batch) + official_inputs = dict(official_batch) + # fast_inputs['attention_mask']: (b, l) + fast_inputs["attention_mask"] = residue_mask.long() + # official_inputs['attention_mask']: (b, l) + official_inputs["attention_mask"] = residue_mask.long() + return fast_inputs, official_inputs, residue_mask + + fast_tokenizer = fast.tokenizer + fast_encoded = _to_device( + fast_tokenizer(sequences, return_tensors="pt", padding=True), + device, + ) + official_encoded = _to_device( + reference_tokenizer(sequences, return_tensors="pt", padding=True), + device, + ) + assert set(fast_encoded) == set(official_encoded), f"{spec.id}: tokenized keys differ" + for name in fast_encoded: + assert torch.equal(fast_encoded[name], official_encoded[name]), ( + f"{spec.id}: tokenized tensor {name!r} differs" + ) + + # input_ids: (b, l) + input_ids = fast_encoded["input_ids"] + # residue_mask: (b, l) + residue_mask = fast_encoded["attention_mask"].bool() + for token_id in getattr(fast_tokenizer, "all_special_ids", ()): + residue_mask &= input_ids.ne(token_id) + + allowed = {"input_ids", "attention_mask"} + fast_inputs = {name: value for name, value in fast_encoded.items() if name in allowed} + official_inputs = {name: value for name, value in official_encoded.items() if name in allowed} + if spec.family.architecture == "ESMC": + # sequence_id: (b, l) + sequence_id = fast_encoded["attention_mask"].bool() + fast_inputs["sequence_id"] = sequence_id + official_inputs["sequence_id"] = sequence_id + return fast_inputs, official_inputs, residue_mask + + +def _hidden_state_tuple(output: object) -> tuple[torch.Tensor, ...]: + """Normalize tuple-valued and layer-stacked hidden-state outputs.""" + + raw = getattr(output, "hidden_states", None) + if torch.is_tensor(raw): + return tuple(raw.unbind(dim=0)) + return tuple(raw or ()) + + +def _last_hidden(output: object) -> torch.Tensor: + value = getattr(output, "last_hidden_state", None) + if value is not None: + return value + hidden_states = _hidden_state_tuple(output) + assert hidden_states, "Model output omitted last_hidden_state and hidden_states" + return hidden_states[-1] + + +def tensor_metrics( + candidate: torch.Tensor, + official: torch.Tensor, + residue_mask: torch.Tensor, +) -> TensorMetrics: + """Compute normalized errors and cosine metrics on valid residues.""" + + # candidate: (...), official: (...), residue_mask: (b, l) + assert candidate.shape == official.shape + assert candidate.ndim == 3 + valid_candidate = candidate.float()[residue_mask] + valid_official = official.float()[residue_mask] + assert valid_candidate.numel() > 0, "Parity batch contains no biological residues" + difference = valid_candidate - valid_official + # denominator: (...) + denominator = torch.linalg.vector_norm(valid_official).clamp_min( + torch.finfo(torch.float32).tiny + ) + relative_l2 = torch.linalg.vector_norm(difference) / denominator + # reference_q999: (...) + reference_q999 = torch.quantile(valid_official.abs().reshape(-1), 0.999) + # difference_q999: () + difference_q999 = torch.quantile(difference.abs().reshape(-1), 0.999) + relative_q999 = difference_q999 / reference_q999.clamp_min(torch.finfo(torch.float32).tiny) + residue_cosines = F.cosine_similarity(valid_candidate, valid_official, dim=-1) + # residue_cosine_p01: () + residue_cosine_p01 = torch.quantile(residue_cosines, 0.01) + + # mask: (...) + mask = residue_mask.unsqueeze(-1) + # denominator: (...) + denominator = mask.sum(1).clamp_min(1) + # candidate_values: (...) + candidate_values = candidate.float() + # official_values: (...) + official_values = official.float() + candidate_pooled = torch.where(mask, candidate_values, 0.0).sum(1) / denominator + official_pooled = torch.where(mask, official_values, 0.0).sum(1) / denominator + # pooled_cosine_min: () + pooled_cosine_min = F.cosine_similarity(candidate_pooled, official_pooled, dim=-1).min() + return TensorMetrics( + relative_l2=float(relative_l2), + relative_q999=float(relative_q999), + residue_cosine_p01=float(residue_cosine_p01), + pooled_cosine_min=float(pooled_cosine_min), + ) + + +def _assert_upper(name: str, value: float, target: float, hard: float, context: str) -> None: + assert value <= hard, f"{context}: {name}={value:.6g} exceeds hard limit {hard:.6g}" + assert value <= target, f"{context}: {name}={value:.6g} misses target {target:.6g}" + + +def _assert_lower(name: str, value: float, target: float, hard: float, context: str) -> None: + assert value >= hard, f"{context}: {name}={value:.6g} violates hard limit {hard:.6g}" + assert value >= target, f"{context}: {name}={value:.6g} misses target {target:.6g}" + + +def _assert_tensor_contract( + candidate: torch.Tensor, + official: torch.Tensor, + residue_mask: torch.Tensor, + contract: NumericContract, + context: str, +) -> None: + # candidate: (...), official: (...), residue_mask: (b, l) + metrics = tensor_metrics(candidate, official, residue_mask) + _assert_tensor_metrics(metrics, contract, context) + + +def _assert_tensor_metrics( + metrics: TensorMetrics, + contract: NumericContract, + context: str, +) -> None: + _assert_upper( + "relative_l2", + metrics.relative_l2, + contract.relative_l2_target, + contract.relative_l2_hard, + context, + ) + _assert_upper( + "relative_q999", + metrics.relative_q999, + contract.relative_q999_target, + contract.relative_q999_hard, + context, + ) + _assert_lower( + "residue_cosine_p01", + metrics.residue_cosine_p01, + contract.residue_cosine_target, + contract.residue_cosine_hard, + context, + ) + _assert_lower( + "pooled_cosine_min", + metrics.pooled_cosine_min, + contract.pooled_cosine_target, + contract.pooled_cosine_hard, + context, + ) + + +def _assert_tensor_metric_records( + records: Sequence[TensorMetricRecord], + contract: NumericContract, +) -> None: + """Assert aggregate extrema after every output tensor has been measured.""" + + assert records, "No output tensor metrics were collected" + upper_metrics = ( + ( + "relative_l2", + "relative_l2", + contract.relative_l2_target, + contract.relative_l2_hard, + ), + ( + "relative_q999", + "relative_q999", + contract.relative_q999_target, + contract.relative_q999_hard, + ), + ) + lower_metrics = ( + ( + "residue_cosine_p01", + "residue_cosine_p01", + contract.residue_cosine_target, + contract.residue_cosine_hard, + ), + ( + "pooled_cosine_min", + "pooled_cosine_min", + contract.pooled_cosine_target, + contract.pooled_cosine_hard, + ), + ) + for name, attribute, target, hard in upper_metrics: + worst = max(records, key=lambda record: getattr(record.metrics, attribute)) + _assert_upper(name, getattr(worst.metrics, attribute), target, hard, worst.context) + for name, attribute, target, hard in lower_metrics: + worst = min(records, key=lambda record: getattr(record.metrics, attribute)) + _assert_lower(name, getattr(worst.metrics, attribute), target, hard, worst.context) + + +def _logits_metrics( + candidate: torch.Tensor, + official: torch.Tensor, + residue_mask: torch.Tensor, + context: str, +) -> LogitsMetrics: + """Collect logits semantics before any numeric threshold is asserted.""" + + # candidate: (...), official: (...), residue_mask: (b, l) + # official_probabilities: (...) + official_probabilities = official.float().softmax(-1) + # candidate_probabilities: (...) + candidate_probabilities = candidate.float().softmax(-1) + # confidence: (...), official_top1: (...) + confidence, official_top1 = official_probabilities.max(-1) + # confident_mask: (...) + confident_mask = residue_mask & confidence.ge(0.5) + assert bool(confident_mask.any()), ( + f"{context}: no positions meet the fixed confidence threshold" + ) + # candidate_top1: (...) + candidate_top1 = candidate_probabilities.argmax(-1) + # top1_agreement: () + top1_agreement = ( + (candidate_top1[confident_mask] == official_top1[confident_mask]).float().mean() + ) + + midpoint = 0.5 * (official_probabilities + candidate_probabilities) + official_log = official_probabilities.clamp_min(1e-12).log() + candidate_log = candidate_probabilities.clamp_min(1e-12).log() + midpoint_log = midpoint.clamp_min(1e-12).log() + jsd = 0.5 * ( + (official_probabilities * (official_log - midpoint_log)).sum(-1) + + (candidate_probabilities * (candidate_log - midpoint_log)).sum(-1) + ) + return LogitsMetrics( + confident_top1_agreement=float(top1_agreement), + mean_jsd=float(jsd[residue_mask].mean()), + ) + + +def _assert_logits_contract( + candidate: torch.Tensor, + official: torch.Tensor, + residue_mask: torch.Tensor, + contract: NumericContract, + context: str, +) -> None: + # candidate: (...), official: (...), residue_mask: (b, l) + _assert_tensor_contract(candidate, official, residue_mask, contract, context) + metrics = _logits_metrics(candidate, official, residue_mask, context) + _assert_lower( + "confident_top1_agreement", + metrics.confident_top1_agreement, + contract.top1_target, + contract.top1_hard, + context, + ) + _assert_upper( + "mean_jsd", + metrics.mean_jsd, + contract.jsd_target, + contract.jsd_hard, + context, + ) + + +def _collect_output_metrics( + spec: ModelSpec, + fast_output: object, + official_output: object, + residue_mask: torch.Tensor, + context: str, +) -> tuple[list[TensorMetricRecord], LogitsMetrics | None]: + """Validate output structure/finite values and collect every parity metric.""" + + # residue_mask: (b, l) + fast_hidden = _hidden_state_tuple(fast_output) + official_hidden = _hidden_state_tuple(official_output) + assert len(fast_hidden) == len(official_hidden), ( + f"{context}: hidden-state count {len(fast_hidden)} != {len(official_hidden)}" + ) + assert fast_hidden, f"{context}: hidden states were not returned" + metric_records: list[TensorMetricRecord] = [] + for layer, (candidate, official) in enumerate(zip(fast_hidden, official_hidden, strict=True)): + assert torch.isfinite(candidate).all(), f"{context}:layer={layer}: non-finite candidate" + assert torch.isfinite(official).all(), f"{context}:layer={layer}: non-finite reference" + metric_records.append( + TensorMetricRecord( + context=f"{context}:layer={layer}", + metrics=tensor_metrics(candidate, official, residue_mask), + ) + ) + fast_last = _last_hidden(fast_output) + official_last = _last_hidden(official_output) + assert torch.isfinite(fast_last).all(), f"{context}:last_hidden_state: non-finite candidate" + assert torch.isfinite(official_last).all(), f"{context}:last_hidden_state: non-finite reference" + metric_records.append( + TensorMetricRecord( + context=f"{context}:last_hidden_state", + metrics=tensor_metrics(fast_last, official_last, residue_mask), + ) + ) + + fast_logits = getattr(fast_output, "logits", None) + official_logits = getattr(official_output, "logits", None) + assert (fast_logits is None) == (official_logits is None), ( + f"{spec.id}: official and FastPLMs output-head contracts differ" + ) + logits_context = f"{context}:logits" + logits_metrics = None + if fast_logits is not None: + assert official_logits is not None + assert torch.isfinite(fast_logits).all(), f"{logits_context}: non-finite candidate" + assert torch.isfinite(official_logits).all(), f"{logits_context}: non-finite reference" + metric_records.append( + TensorMetricRecord( + context=logits_context, + metrics=tensor_metrics(fast_logits, official_logits, residue_mask), + ) + ) + logits_metrics = _logits_metrics( + fast_logits, + official_logits, + residue_mask, + logits_context, + ) + + return metric_records, logits_metrics + + +def _assert_outputs( + spec: ModelSpec, + fast_output: object, + official_output: object, + residue_mask: torch.Tensor, + contract: NumericContract, + context: str, +) -> None: + # residue_mask: (b, l) + metric_records, logits_metrics = _collect_output_metrics( + spec, + fast_output, + official_output, + residue_mask, + context, + ) + + _assert_tensor_metric_records(metric_records, contract) + if logits_metrics is not None: + logits_context = f"{context}:logits" + _assert_lower( + "confident_top1_agreement", + logits_metrics.confident_top1_agreement, + contract.top1_target, + contract.top1_hard, + logits_context, + ) + _assert_upper( + "mean_jsd", + logits_metrics.mean_jsd, + contract.jsd_target, + contract.jsd_hard, + logits_context, + ) + + +def _assert_esmc_alternate_backend_outputs( + spec: ModelSpec, + fast_output: object, + official_output: object, + residue_mask: torch.Tensor, + context: str, +) -> None: + """Warn on published-band drift while retaining catastrophic hard gates.""" + + # residue_mask: (b, l) + records, logits = _collect_output_metrics( + spec, + fast_output, + official_output, + residue_mask, + context, + ) + _assert_tensor_metric_records(records, ESMC_CATASTROPHIC_BF16_CONTRACT) + if logits is not None: + assert logits.confident_top1_agreement >= ESMC_CATASTROPHIC_BF16_CONTRACT.top1_hard + assert logits.mean_jsd <= ESMC_CATASTROPHIC_BF16_CONTRACT.jsd_hard + try: + _assert_tensor_metric_records(records, ESMC_ALTERNATE_BF16_CONTRACT) + if logits is not None: + _assert_lower( + "confident_top1_agreement", + logits.confident_top1_agreement, + ESMC_ALTERNATE_BF16_CONTRACT.top1_target, + ESMC_ALTERNATE_BF16_CONTRACT.top1_hard, + f"{context}:logits", + ) + _assert_upper( + "mean_jsd", + logits.mean_jsd, + ESMC_ALTERNATE_BF16_CONTRACT.jsd_target, + ESMC_ALTERNATE_BF16_CONTRACT.jsd_hard, + f"{context}:logits", + ) + except AssertionError as error: + warnings.warn( + f"{context}: supported ESMC backend is outside its published diagnostic band " + f"but passed catastrophic biological gates: {error}", + UserWarning, + stacklevel=2, + ) + + +def _assert_esmc_sdpa_exact( + fast_output: object, + official_output: object, + context: str, +) -> None: + """Require exact ESMC SDPA equality, including special and padding tokens.""" + + fast_hidden = _hidden_state_tuple(fast_output) + official_hidden = _hidden_state_tuple(official_output) + assert len(fast_hidden) == len(official_hidden) + for layer, (candidate, official) in enumerate( + zip(fast_hidden, official_hidden, strict=True) + ): + assert torch.equal(candidate, official), ( + f"{context}:layer={layer}: full hidden states are not exact" + ) + + assert torch.equal(_last_hidden(fast_output), _last_hidden(official_output)), ( + f"{context}: full last_hidden_state is not exact" + ) + fast_logits = getattr(fast_output, "logits", None) + official_logits = getattr(official_output, "logits", None) + if fast_logits is not None or official_logits is not None: + assert fast_logits is not None and official_logits is not None + assert torch.equal(fast_logits, official_logits), ( + f"{context}: full logits are not exact" + ) + + +def _run_inference_contract( + spec: ModelSpec, + dtype: torch.dtype, + backend: str | None, +) -> None: + device = torch.device("cuda") + torch.manual_seed(SEED) + # The manifest distinguishes static BF16 parameters from FP32-resident + # parameters evaluated under CUDA BF16 autocast. Parity, serving probes, + # and benchmarks all consume this same typed execution contract. + use_bf16_autocast = ( + dtype == torch.bfloat16 and spec.family.bf16_execution == "fp32_parameters_autocast" + ) + load_dtype = torch.float32 if use_bf16_autocast else dtype + fast = _load_fast(spec, device, load_dtype) + reference, reference_tokenizer = _load_reference(spec, device, load_dtype) + if backend is not None: + assert hasattr(fast, "set_attn_implementation"), ( + f"{spec.id}: advertised backend API set_attn_implementation is missing" + ) + fast.set_attn_implementation(backend) + resolved = getattr(fast.config, "_attn_implementation", None) + if resolved is None: + resolved = getattr(fast.config, "attn_implementation", None) + assert resolved == backend, ( + f"{spec.id}: requested {backend!r}, resolved {resolved!r}; silent fallback is forbidden" + ) + + fast_inputs, official_inputs, residue_mask = _prepare_inputs( + spec, + fast, + reference_tokenizer, + _sequence_batch(), + device, + ) + if dtype == torch.float32: + numeric_context = strict_fp32_matmul() + elif use_bf16_autocast: + numeric_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16) + else: + numeric_context = contextlib.nullcontext() + with torch.inference_mode(), numeric_context: + fast_output = fast(**fast_inputs, output_hidden_states=True) + official_output = reference(**official_inputs, output_hidden_states=True) + contract = _numeric_contract(spec, dtype, backend) + dtype_name = "fp32" if dtype == torch.float32 else "bf16" + context = f"{spec.id}:{dtype_name}:{backend or 'default'}" + if ( + spec.family.architecture == "ESMC" + and dtype == torch.bfloat16 + and backend in {"flex_attention", "flash_attention_3"} + ): + _assert_esmc_alternate_backend_outputs( + spec, + fast_output, + official_output, + residue_mask, + context, + ) + else: + _assert_outputs( + spec, + fast_output, + official_output, + residue_mask, + contract, + context, + ) + if spec.family.architecture == "ESMC" and backend in (None, "sdpa"): + _assert_esmc_sdpa_exact( + fast_output, + official_output, + context, + ) + del fast, reference, fast_output, official_output + gc.collect() + torch.cuda.empty_cache() + + +def test_manifest_state_transforms_are_registered() -> None: + declared = { + spec.family.state_transform + for spec in REGISTRY.values() + if spec.family.tokenizer_mode != "structure" + } + assert declared.issubset(TRANSFORMS), ( + f"Missing deterministic state transforms: {sorted(declared.difference(TRANSFORMS))}" + ) + + +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_SPECS]) +def test_exact_checkpoint_contract(spec: ModelSpec) -> None: + """Every checkpoint has exact semantic config, state, and aliases.""" + + device = torch.device("cpu") + fast = _load_fast(spec, device, None) + reference, reference_tokenizer = _load_reference(spec, device, None) + _assert_semantic_config_equal(spec, fast, reference) + _assert_state_equal(spec, fast, reference) + _assert_aliases_equal(spec, fast, reference) + if spec.family.tokenizer_mode == "tokenizer": + _assert_tokenizer_equal(spec, fast.tokenizer, reference_tokenizer) + del fast, reference + gc.collect() + + +@pytest.mark.gpu +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_SPECS]) +def test_every_checkpoint_live_bf16_inference(spec: ModelSpec) -> None: + """Every checkpoint passes one live mixed-length BF16 official comparison.""" + + _run_inference_contract(spec, torch.bfloat16, backend=None) + + +DEEP_CASES = [ + _deep_parameter(spec, dtype, backend) + for spec in DEEP_SPECS + for backend in spec.family.attention + for dtype in ( + { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + }[dtype_name] + for dtype_name in REGISTRY.supported_attention_dtypes(spec.family.id, backend) + ) +] + + +@pytest.mark.parametrize(("spec", "dtype", "backend"), DEEP_CASES) +def test_representative_deep_backend_parity( + spec: ModelSpec, + dtype: torch.dtype, + backend: str, +) -> None: + """Representatives pass all-layer parity for every advertised backend.""" + + _run_inference_contract(spec, dtype, backend) diff --git a/tests/parity/test_native_results.py b/tests/parity/test_native_results.py new file mode 100644 index 0000000..4aec661 --- /dev/null +++ b/tests/parity/test_native_results.py @@ -0,0 +1,1912 @@ +"""Consume native-container results without importing an official package.""" + +from __future__ import annotations + +import contextlib +import gc +import hashlib +import importlib.metadata +import json +import math +import os +import platform +import re +import subprocess +import tempfile +import warnings +import pytest +import torch +import transformers +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from safetensors.torch import load_file + +from fastplms.registry import ModelSpec, get_model_registry +from tests.conftest import strict_fp32_matmul +from tests.parity.support.esmc_calibration import ( + ESMC_BOUNDARY_LENGTHS, + esmc_calibration_batches, + load_esmc_biological_holdout, + validate_esmc_calibration_batch, +) +from tests.parity.support.native_reference import _tensor_digest, _token_result +from tests.parity.support.reference_adapters.biohub_source import ( + BIOHUB_ESM_REVISION, + BIOHUB_ESM_TREE_SHA256, + BIOHUB_REFERENCE_SOURCE_NAMES, + BIOHUB_TRANSFORMERS_REVISION, + BIOHUB_TRANSFORMERS_TREE_SHA256, +) +from tests.parity.support.reference_adapters.dplm2 import ( + DPLM2_3B_GENERATION_LIMITATION, + DPLM2_150M_OFFICIAL_HEAD_CONTRACT, +) +from tests.parity.test_model_parity import ( + BF16_CONTRACT, + EDGE_SEQUENCES, + ESMC_ALTERNATE_BF16_CONTRACT, + ESMC_CATASTROPHIC_BF16_CONTRACT, + LogitsMetrics, + TensorMetricRecord, + _alias_groups, + _assert_esmc_alternate_backend_outputs, + _assert_esmc_sdpa_exact, + _assert_outputs, + _assert_tensor_metric_records, + _collect_output_metrics, + _hidden_state_tuple, + _last_hidden, + _load_fast, + _numeric_contract, + _semantic_config, +) +from tools.remote.biohub_reference_environment import ( + BiohubReferenceEnvironmentError, + validate_biohub_reference_environment_evidence, +) +from tools.remote.reference_source_attestation import validate_reference_sources_evidence + + +pytestmark = [pytest.mark.compliance, pytest.mark.gpu, pytest.mark.slow] +REGISTRY = get_model_registry() +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +BIOHUB_REFERENCE_LOCK = REPOSITORY_ROOT / "docker/constraints/biohub-reference-lock.json" +SEQUENCE_SPECS = tuple( + spec for spec in REGISTRY.values() if spec.family.tokenizer_mode != "structure" +) + + +def _parameter(spec: ModelSpec) -> Any: + marks = [pytest.mark.large] if spec.size_category == "xlarge" else [] + return pytest.param(spec, id=spec.id, marks=marks) + + +def _validated_biohub_reference_sources( + value: object, +) -> dict[str, dict[str, object]]: + sources = validate_reference_sources_evidence( + value, + required_sources=BIOHUB_REFERENCE_SOURCE_NAMES, + ) + expected = { + "biohub-esm": { + "source_revision": BIOHUB_ESM_REVISION, + "tree_sha256": BIOHUB_ESM_TREE_SHA256, + "import_name": "esm", + "import_root": "esm", + "import_file": "esm/__init__.py", + "package_version": "3.3.0", + }, + "biohub-transformers": { + "source_revision": BIOHUB_TRANSFORMERS_REVISION, + "tree_sha256": BIOHUB_TRANSFORMERS_TREE_SHA256, + "import_name": "transformers", + "import_root": "src/transformers", + "import_file": "src/transformers/__init__.py", + "package_version": "4.57.6", + }, + } + for source_name, source_expected in expected.items(): + source = sources[source_name] + for field, expected_value in source_expected.items(): + if source[field] != expected_value: + raise ValueError( + f"{source_name} reference source evidence {field} differs from " + f"{expected_value!r}: {source[field]!r}" + ) + return sources + + +def _validated_biohub_reference_environment(value: object) -> dict[str, object]: + try: + return validate_biohub_reference_environment_evidence( + value, + repository_root=REPOSITORY_ROOT, + contract_path=BIOHUB_REFERENCE_LOCK, + ) + except BiohubReferenceEnvironmentError as error: + raise ValueError(f"Biohub reference environment is invalid: {error}") from error + + +def _result(spec: ModelSpec) -> tuple[dict[str, Any], Path]: + root = os.environ.get("FASTPLMS_REFERENCE_RESULTS") + if not root: + raise RuntimeError("FASTPLMS_REFERENCE_RESULTS is required for native compliance") + directory = Path(root) / spec.id + metadata_path = directory / "metadata.json" + if not metadata_path.is_file(): + raise FileNotFoundError(f"Native reference result is missing for {spec.id}: {directory}") + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if not isinstance(metadata, dict): + raise ValueError(f"Native reference metadata for {spec.id} is not an object") + expected_identity = { + "reference_repo_id": spec.official.repo_id, + "reference_revision": spec.official.revision, + "state_transform": spec.family.state_transform, + } + for name, expected in expected_identity.items(): + if metadata.get(name) != expected: + raise ValueError( + f"Native reference {name} for {spec.id} differs from {expected!r}: " + f"{metadata.get(name)!r}" + ) + if spec.family.id in {"esm_plusplus", "esm3", "esmfold2"}: + _validated_biohub_reference_sources(metadata.get("reference_sources")) + _validated_biohub_reference_environment(metadata.get("reference_environment")) + return metadata, directory + + +def _load_package_generation_model( + spec: ModelSpec, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.nn.Module: + """Load pinned mirror weights into the current package implementation.""" + + if spec.family.id == "dplm": + from fastplms.models.dplm.modeling_dplm import DPLMConfig, DPLMForMaskedLM + + config_class = DPLMConfig + model_class = DPLMForMaskedLM + elif spec.family.id == "dplm2": + from fastplms.models.dplm2.modeling_dplm2 import DPLM2Config, DPLM2ForMaskedLM + + config_class = DPLM2Config + model_class = DPLM2ForMaskedLM + elif spec.family.id == "ankh": + from fastplms.models.ankh.modeling_ankh import ( + FastAnkhConfig, + FastAnkhForConditionalGeneration, + ) + + config_class = FastAnkhConfig + model_class = FastAnkhForConditionalGeneration + else: + raise ValueError(f"Generation loading is unsupported for {spec.family.id!r}") + config = config_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + ) + model_kwargs: dict[str, Any] = {} + if spec.family.id == "ankh": + model_kwargs["attn_implementation"] = "eager" + model = model_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + config=config, + dtype=dtype, + **model_kwargs, + ) + return model.to(device).eval() + + +def _normalized_token_result(tokenizer: object, options: dict[str, Any]) -> Any: + result = _token_result(tokenizer, EDGE_SEQUENCES, options) + return json.loads(json.dumps(result)) + + +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_SPECS]) +def test_native_exact_checkpoint_contract(spec: ModelSpec) -> None: + """Candidate config, weight bytes, aliases, and tokenizer match native output.""" + + metadata, _ = _result(spec) + fast = _load_fast(spec, torch.device("cpu"), None) + + official_config = metadata["semantic_config"] + candidate_config = _semantic_config(fast) + assert {name: candidate_config[name] for name in official_config} == official_config + + candidate_state = { + name: _tensor_digest(tensor) for name, tensor in sorted(fast.state_dict().items()) + } + assert candidate_state == metadata["state"]["tensors"], ( + f"{spec.id}: candidate state differs from the native official state" + ) + candidate_aliases = sorted(sorted(group) for group in _alias_groups(fast)) + assert candidate_aliases == metadata["state"]["aliases"] + + if spec.family.tokenizer_mode == "tokenizer": + contract = metadata["tokenizer"] + tokenizer = fast.tokenizer + assert tokenizer.get_vocab() == contract["vocab"] + assert { + name: getattr(tokenizer, name, None) for name in contract["special_ids"] + } == contract["special_ids"] + for case in contract["behavior"]: + assert _normalized_token_result(tokenizer, case["options"]) == case["result"] + artifact_root = os.environ.get("FASTPLMS_CANDIDATE_ARTIFACTS") + if not artifact_root: + raise RuntimeError("FASTPLMS_CANDIDATE_ARTIFACTS is required for tokenizer assets") + repository_name = spec.fast.repo_id.split("/", maxsplit=1)[-1] + artifact = Path(artifact_root) / repository_name + for relative_name, expected in metadata["tokenizer_assets"].items(): + path = artifact.joinpath(*Path(relative_name).parts) + assert path.is_file(), f"{spec.id}: missing tokenizer asset {relative_name}" + content = path.read_bytes() + assert len(content) == expected["size"] + assert hashlib.sha256(content).hexdigest() == expected["sha256"], ( + f"{spec.id}: tokenizer asset bytes differ for {relative_name}" + ) + del fast + gc.collect() + + +def test_native_dplm2_150m_exact_head_contract() -> None: + """Pinned trained heads remain exact, complete, and independent.""" + + metadata, _ = _result(REGISTRY["dplm2_150m"]) + state = metadata["state"]["tensors"] + observed = {name: state[name] for name in DPLM2_150M_OFFICIAL_HEAD_CONTRACT} + assert observed == DPLM2_150M_OFFICIAL_HEAD_CONTRACT + assert metadata["state"]["aliases"] == [] + + +def _official_output(tensors: dict[str, torch.Tensor], device: torch.device) -> object: + hidden_names = sorted(name for name in tensors if name.startswith("output__hidden_")) + hidden_states = tuple(tensors[name].to(device) for name in hidden_names) + values: dict[str, Any] = { + "hidden_states": hidden_states, + "last_hidden_state": tensors["output__last_hidden_state"].to(device), + } + if "output__logits" in tensors: + # values['logits']: (..., c) + values["logits"] = tensors["output__logits"].to(device) + return SimpleNamespace(**values) + + +ESMC_DIAGNOSTIC_SCHEMA_VERSION = 3 +ESMC_MEASURED_BACKENDS = ("eager", "sdpa", "flex_attention") +ESMC_UNAVAILABLE_BACKENDS = ("flash_attention_2", "flash_attention_3") +_SHA256_HEX_LENGTH = 64 +_CANDIDATE_IDENTITY_FIELDS = ( + "fastplms_model_id", + "fastplms_checkpoint_repo_id", + "fastplms_checkpoint_revision", + "fastplms_weights_revision", + "fastplms_runtime_revision", + "fastplms_source_tree_sha256", + "fastplms_runtime_bundle_sha256", +) + + +def _require_sha256(value: object, context: str) -> str: + if not isinstance(value, str) or len(value) != _SHA256_HEX_LENGTH or value != value.lower(): + raise ValueError(f"{context} must be a 64-character SHA-256 digest") + try: + bytes.fromhex(value) + except ValueError as error: + raise ValueError(f"{context} must be hexadecimal") from error + return value + + +def _require_runtime_revision( + value: object, + source_tree_sha256: str, + context: str, +) -> str: + """Accept the two immutable identities emitted by the artifact builder.""" + + if isinstance(value, str) and re.fullmatch(r"[0-9a-f]{40}", value) is not None: + return value + content_addressed = f"source-tree-sha256:{source_tree_sha256}" + if value == content_addressed: + return content_addressed + raise ValueError(f"{context} must be a clean Git revision or the exact source-tree digest") + + +def _esmc_unavailability_identity( + backend: str, + reference_environment: Mapping[str, object], +) -> dict[str, str]: + runtime = reference_environment.get("runtime") + if not isinstance(runtime, Mapping): + raise ValueError("ESMC unavailability requires the locked runtime identity") + operating_system = runtime.get("operating_system") + architecture = runtime.get("architecture") + gpu = runtime.get("gpu") + if ( + not isinstance(operating_system, str) + or not operating_system.strip() + or not isinstance(architecture, str) + or not architecture.strip() + or not isinstance(gpu, Mapping) + ): + raise ValueError("ESMC unavailability runtime platform identity is malformed") + gpu_name = gpu.get("name") + capability = gpu.get("capability") + if ( + not isinstance(gpu_name, str) + or not gpu_name.strip() + or not isinstance(capability, list) + or len(capability) != 2 + or any(isinstance(value, bool) or not isinstance(value, int) for value in capability) + ): + raise ValueError("ESMC unavailability accelerator identity is malformed") + platform_identity = f"{operating_system.lower()}/{architecture.lower()}" + accelerator_identity = f"{gpu_name}/SM{capability[0]}{capability[1]}" + if backend == "flash_attention_2": + historical_evidence = "separate_historical_focused_evidence_only" + reason = ( + f"The locked {platform_identity} {accelerator_identity} release environment " + "has no validated FlashAttention 2 " + "kernel. Prior focused execution evidence is historical and is not part of " + "the current ESMC release distribution." + ) + elif backend == "flash_attention_3": + historical_evidence = "none" + reason = ( + "The manifest-pinned FlashAttention 3 kernel has no validated artifact for " + f"the locked {platform_identity} {accelerator_identity} release environment." + ) + else: + raise ValueError(f"ESMC backend {backend!r} is not a structured unavailable backend") + return { + "code": "locked_platform_kernel_unavailable", + "platform": platform_identity, + "accelerator": accelerator_identity, + "dispatch_contract": "fail_closed_without_dispatch", + "historical_evidence": historical_evidence, + "reason": reason, + } + + +def _optional_package_version(distribution: str) -> str | None: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def _cuda_driver_version() -> str: + try: + completed = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (FileNotFoundError, subprocess.SubprocessError) as error: + raise RuntimeError("ESMC diagnostics require the exact NVIDIA driver version") from error + versions = tuple(line.strip() for line in completed.stdout.splitlines() if line.strip()) + if not versions: + raise RuntimeError("nvidia-smi did not report an NVIDIA driver version") + return versions[0] + + +def _candidate_environment_identity() -> dict[str, object]: + if not torch.cuda.is_available(): + raise RuntimeError("ESMC diagnostic evidence requires a CUDA device") + device = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(device) + return { + "python": platform.python_version(), + "torch": torch.__version__, + "transformers": transformers.__version__, + "cuda_runtime": str(torch.version.cuda or "unavailable"), + "cuda_driver": _cuda_driver_version(), + "gpu": { + "name": properties.name, + "capability": list(torch.cuda.get_device_capability(device)), + "total_memory_bytes": int(properties.total_memory), + }, + "packages": { + distribution: _optional_package_version(distribution) + for distribution in ( + "fastplms", + "huggingface-hub", + "kernels", + "tokenizers", + "transformer-engine", + "transformer-engine-torch", + ) + }, + } + + +def _candidate_identity(spec: ModelSpec, model: torch.nn.Module) -> dict[str, object]: + config = getattr(model, "config", None) + if config is None: + raise ValueError(f"{spec.id}: candidate model has no configuration identity") + observed = {name: getattr(config, name, None) for name in _CANDIDATE_IDENTITY_FIELDS} + expected = { + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "fastplms_checkpoint_revision": spec.artifact_checkpoint.revision, + "fastplms_weights_revision": spec.artifact_checkpoint.revision, + } + for name, expected_value in expected.items(): + value = observed[name] + if value != expected_value: + raise ValueError( + f"{spec.id}: candidate {name}={value!r} differs from {expected_value!r}" + ) + source_tree_sha256 = _require_sha256( + observed["fastplms_source_tree_sha256"], + f"{spec.id} candidate fastplms_source_tree_sha256", + ) + runtime_bundle_sha256 = _require_sha256( + observed["fastplms_runtime_bundle_sha256"], + f"{spec.id} candidate fastplms_runtime_bundle_sha256", + ) + runtime_revision = _require_runtime_revision( + observed["fastplms_runtime_revision"], + source_tree_sha256, + f"{spec.id} candidate runtime revision", + ) + resolved_commit = getattr(config, "_commit_hash", None) + if resolved_commit is not None and ( + not isinstance(resolved_commit, str) or not resolved_commit.strip() + ): + raise ValueError(f"{spec.id}: candidate resolved commit is invalid") + return { + "repo_id": spec.fast.repo_id, + "manifest_revision": spec.fast.revision, + "resolved_commit": resolved_commit, + "checkpoint_repo_id": observed["fastplms_checkpoint_repo_id"], + "checkpoint_revision": observed["fastplms_checkpoint_revision"], + "weights_revision": observed["fastplms_weights_revision"], + "runtime_revision": runtime_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + } + + +def _reference_identity( + spec: ModelSpec, + metadata: Mapping[str, object], +) -> dict[str, object]: + expected = { + "reference_repo_id": spec.official.repo_id, + "reference_revision": spec.official.revision, + "state_transform": spec.family.state_transform, + } + for name, expected_value in expected.items(): + if metadata.get(name) != expected_value: + raise ValueError( + f"{spec.id}: native {name}={metadata.get(name)!r} differs from {expected_value!r}" + ) + environment = metadata.get("environment") + if not isinstance(environment, Mapping): + raise ValueError(f"{spec.id}: native result omits its environment identity") + required_environment = { + "cuda_device", + "cuda_device_capability", + "cuda_total_memory", + "cuda_runtime", + "packages", + "python", + "torch", + } + if not required_environment.issubset(environment): + missing = sorted(required_environment.difference(environment)) + raise ValueError(f"{spec.id}: native environment omits {missing!r}") + reference_sources = _validated_biohub_reference_sources(metadata.get("reference_sources")) + reference_environment = _validated_biohub_reference_environment( + metadata.get("reference_environment") + ) + return { + "repo_id": spec.official.repo_id, + "revision": spec.official.revision, + "state_transform": spec.family.state_transform, + "environment": dict(environment), + "reference_environment": reference_environment, + "reference_sources": reference_sources, + } + + +def _kernel_identity(backend: str) -> dict[str, object]: + kernel = REGISTRY.attention_kernels.get(backend) + if kernel is None: + return { + "implementation": backend, + "provider": "torch", + "torch_version": torch.__version__, + } + return { + "implementation": backend, + "provider": "huggingface_kernels", + "repository": kernel.repository, + "revision": kernel.revision, + "version": kernel.version, + "expected_variant": kernel.expected_variant, + "supported_dtypes": list(kernel.dtypes), + "kernels_package_version": _optional_package_version("kernels"), + } + + +def _metric_payload( + record: TensorMetricRecord, + *, + output: str, + layer_index: int | None, +) -> dict[str, object]: + metrics = record.metrics + return { + "context": record.context, + "output": output, + "layer_index": layer_index, + "relative_l2": metrics.relative_l2, + "relative_q999": metrics.relative_q999, + "residue_cosine_p01": metrics.residue_cosine_p01, + "pooled_cosine_min": metrics.pooled_cosine_min, + } + + +def _structured_metric_payloads( + output: object, + records: list[TensorMetricRecord], +) -> list[dict[str, object]]: + hidden_count = len(_hidden_state_tuple(output)) + has_logits = getattr(output, "logits", None) is not None + expected_count = hidden_count + 1 + int(has_logits) + if len(records) != expected_count: + raise ValueError( + f"ESMC metric record count {len(records)} differs from expected {expected_count}" + ) + result = [ + _metric_payload(record, output="hidden_state", layer_index=layer) + for layer, record in enumerate(records[:hidden_count]) + ] + result.append( + _metric_payload( + records[hidden_count], + output="last_hidden_state", + layer_index=None, + ) + ) + if has_logits: + result.append(_metric_payload(records[-1], output="logits", layer_index=None)) + return result + + +def _logits_metric_payload(metrics: LogitsMetrics | None) -> dict[str, float] | None: + if metrics is None: + return None + return { + "confident_top1_agreement": metrics.confident_top1_agreement, + "mean_jsd": metrics.mean_jsd, + } + + +def _case_identity(case: Mapping[str, object]) -> dict[str, object]: + return { + "case_id": case["case_id"], + "sequence_length": case["sequence_length"], + "sequence_sha256": case["sequence_sha256"], + "source": case.get("source"), + "source_sha256": case.get("source_sha256"), + } + + +def _public_panel_identity(panel: Mapping[str, object]) -> dict[str, object]: + cases = panel.get("cases") + if not isinstance(cases, list): + raise ValueError("ESMC panel identity omits its ordered cases") + return { + "schema_version": panel["schema_version"], + "kind": panel["kind"], + "seed": panel["seed"], + "definition_sha256": panel["definition_sha256"], + "cases": [_case_identity(case) for case in cases], + } + + +def _slice_output(output: object, index: int) -> SimpleNamespace: + values: dict[str, object] = { + "hidden_states": tuple(value[index : index + 1] for value in _hidden_state_tuple(output)), + "last_hidden_state": _last_hidden(output)[index : index + 1], + } + logits = getattr(output, "logits", None) + if logits is not None: + # values['logits']: (..., c) + values["logits"] = logits[index : index + 1] + return SimpleNamespace(**values) + + +def _case_metric_distributions( + spec: ModelSpec, + candidate: object, + official: object, + residue_mask: torch.Tensor, + panel: Mapping[str, object], + context: str, +) -> tuple[list[dict[str, object]], list[str]]: + # residue_mask: (b, l) + cases = panel.get("cases") + if not isinstance(cases, list) or residue_mask.ndim != 2: + raise ValueError("ESMC panel cases and residue mask must be batch aligned") + if len(cases) != residue_mask.shape[0]: + raise ValueError( + f"ESMC panel has {len(cases)} cases for batch size {residue_mask.shape[0]}" + ) + result: list[dict[str, object]] = [] + violations: list[str] = [] + for index, case in enumerate(cases): + if not isinstance(case, Mapping): + raise ValueError(f"ESMC panel case {index} is not an object") + case_id = str(case["case_id"]) + expected_length = int(case["sequence_length"]) + observed_length = int(residue_mask[index].sum().item()) + if observed_length != expected_length: + raise ValueError( + f"{case_id}: residue mask length {observed_length} != {expected_length}" + ) + candidate_case = _slice_output(candidate, index) + official_case = _slice_output(official, index) + # case_mask: (...) + case_mask = residue_mask[index : index + 1] + records, logits = _collect_output_metrics( + spec, + candidate_case, + official_case, + case_mask, + f"{context}:case={case_id}", + ) + _assert_esmc_catastrophic_metrics( + records, + logits, + f"{context}:case={case_id}", + ) + violations.extend(_esmc_published_band_violations(records, logits)) + result.append( + { + **_case_identity(case), + "tensor_metrics": _structured_metric_payloads(candidate_case, records), + "logits_metrics": _logits_metric_payload(logits), + } + ) + return result, violations + + +def _esmc_published_band_violations( + records: list[TensorMetricRecord], + logits: LogitsMetrics | None, +) -> list[str]: + contract = ESMC_ALTERNATE_BF16_CONTRACT + violations: list[str] = [] + for record in records: + metrics = record.metrics + if metrics.relative_l2 > contract.relative_l2_hard: + violations.append( + f"{record.context}:relative_l2={metrics.relative_l2:.6g}>" + f"{contract.relative_l2_hard:.6g}" + ) + if metrics.relative_q999 > contract.relative_q999_hard: + violations.append( + f"{record.context}:relative_q999={metrics.relative_q999:.6g}>" + f"{contract.relative_q999_hard:.6g}" + ) + if metrics.residue_cosine_p01 < contract.residue_cosine_hard: + violations.append( + f"{record.context}:residue_cosine_p01={metrics.residue_cosine_p01:.6g}<" + f"{contract.residue_cosine_hard:.6g}" + ) + if metrics.pooled_cosine_min < contract.pooled_cosine_hard: + violations.append( + f"{record.context}:pooled_cosine_min={metrics.pooled_cosine_min:.6g}<" + f"{contract.pooled_cosine_hard:.6g}" + ) + if logits is not None: + if logits.confident_top1_agreement < contract.top1_hard: + violations.append( + "logits:confident_top1_agreement=" + f"{logits.confident_top1_agreement:.6g}<{contract.top1_hard:.6g}" + ) + if logits.mean_jsd > contract.jsd_hard: + violations.append(f"logits:mean_jsd={logits.mean_jsd:.6g}>{contract.jsd_hard:.6g}") + return violations + + +def _assert_esmc_catastrophic_metrics( + records: list[TensorMetricRecord], + logits: LogitsMetrics | None, + context: str, +) -> None: + _assert_tensor_metric_records(records, ESMC_CATASTROPHIC_BF16_CONTRACT) + if logits is None: + return + assert logits.confident_top1_agreement >= ESMC_CATASTROPHIC_BF16_CONTRACT.top1_hard, ( + f"{context}: catastrophic top-1 disagreement" + ) + assert logits.mean_jsd <= ESMC_CATASTROPHIC_BF16_CONTRACT.jsd_hard, ( + f"{context}: catastrophic Jensen-Shannon divergence" + ) + + +def _release_gate_identity(backend: str) -> dict[str, str]: + if backend == "sdpa": + mode = "exact" + elif backend == "eager": + mode = "strict_numeric" + elif backend == "flex_attention": + mode = "diagnostic_with_catastrophe_gate" + else: + raise ValueError(f"Unsupported ESMC diagnostic backend: {backend!r}") + return {"mode": mode, "status": "passed"} + + +def _validate_metric_payload(payload: Mapping[str, object]) -> None: + expected = { + "context", + "output", + "layer_index", + "relative_l2", + "relative_q999", + "residue_cosine_p01", + "pooled_cosine_min", + } + if set(payload) != expected: + raise ValueError("ESMC tensor-metric fields differ from schema v3") + context = payload.get("context") + if not isinstance(context, str) or not context.strip(): + raise ValueError("ESMC tensor metric context is invalid") + output = payload.get("output") + layer_index = payload.get("layer_index") + if output == "hidden_state": + if isinstance(layer_index, bool) or not isinstance(layer_index, int) or layer_index < 0: + raise ValueError("ESMC hidden-state metrics require a nonnegative layer index") + elif output in {"last_hidden_state", "logits"}: + if layer_index is not None: + raise ValueError(f"ESMC {output} metrics must not carry a layer index") + else: + raise ValueError(f"Unsupported ESMC metric output: {output!r}") + for name in ( + "relative_l2", + "relative_q999", + "residue_cosine_p01", + "pooled_cosine_min", + ): + value = payload.get(name) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise ValueError(f"ESMC metric {name} must be a finite number") + upper_bounds = { + "relative_l2": ESMC_CATASTROPHIC_BF16_CONTRACT.relative_l2_hard, + "relative_q999": ESMC_CATASTROPHIC_BF16_CONTRACT.relative_q999_hard, + } + lower_bounds = { + "residue_cosine_p01": ESMC_CATASTROPHIC_BF16_CONTRACT.residue_cosine_hard, + "pooled_cosine_min": ESMC_CATASTROPHIC_BF16_CONTRACT.pooled_cosine_hard, + } + for name, limit in upper_bounds.items(): + if float(payload[name]) < 0 or float(payload[name]) > limit: + raise ValueError(f"ESMC metric {name} fails the catastrophe gate") + for name, limit in lower_bounds.items(): + if float(payload[name]) < limit or float(payload[name]) > 1.000001: + raise ValueError(f"ESMC metric {name} fails the catastrophe gate") + + +def _validate_logits_metric_payload(payload: object) -> None: + if payload is None: + return + if not isinstance(payload, Mapping) or set(payload) != { + "confident_top1_agreement", + "mean_jsd", + }: + raise ValueError("ESMC logits-metric fields differ from schema v3") + for name in ("confident_top1_agreement", "mean_jsd"): + value = payload.get(name) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise ValueError(f"ESMC logits metric {name} must be a finite number") + agreement = float(payload["confident_top1_agreement"]) + mean_jsd = float(payload["mean_jsd"]) + if not ESMC_CATASTROPHIC_BF16_CONTRACT.top1_hard <= agreement <= 1.000001: + raise ValueError("ESMC top-1 agreement fails the catastrophe gate") + if not -1e-7 <= mean_jsd <= ESMC_CATASTROPHIC_BF16_CONTRACT.jsd_hard: + raise ValueError("ESMC Jensen-Shannon divergence fails the catastrophe gate") + + +def _validate_metric_distribution( + metrics: object, + *, + context: str, +) -> None: + if not isinstance(metrics, list) or not metrics: + raise ValueError(f"{context} tensor metrics are missing") + hidden_layers: list[int] = [] + output_counts = { + "last_hidden_state": 0, + "logits": 0, + } + for metric in metrics: + if not isinstance(metric, Mapping): + raise ValueError(f"{context} tensor metric is not an object") + _validate_metric_payload(metric) + output = metric["output"] + if output == "hidden_state": + layer_index = metric["layer_index"] + if not isinstance(layer_index, int): + raise ValueError(f"{context} hidden-state layer index is invalid") + hidden_layers.append(layer_index) + else: + output_counts[str(output)] += 1 + if not hidden_layers: + raise ValueError(f"{context} contains no hidden-state layer metrics") + if hidden_layers != list(range(len(hidden_layers))): + raise ValueError(f"{context} hidden-state layers are incomplete or unordered") + if output_counts["last_hidden_state"] != 1: + raise ValueError(f"{context} must contain one last-hidden-state metric") + if output_counts["logits"] not in {0, 1}: + raise ValueError(f"{context} contains duplicate logits metrics") + + +def _validate_candidate_environment(environment: object) -> None: + if not isinstance(environment, Mapping) or set(environment) != { + "python", + "torch", + "transformers", + "cuda_runtime", + "cuda_driver", + "gpu", + "packages", + }: + raise ValueError("ESMC candidate environment differs from schema v3") + for name in ("python", "torch", "transformers", "cuda_runtime", "cuda_driver"): + value = environment.get(name) + if not isinstance(value, str) or not value.strip() or value == "unavailable": + raise ValueError(f"ESMC candidate environment has invalid {name}") + + packages = environment.get("packages") + expected_packages = { + "fastplms", + "huggingface-hub", + "kernels", + "tokenizers", + "transformer-engine", + "transformer-engine-torch", + } + if not isinstance(packages, Mapping) or set(packages) != expected_packages: + raise ValueError("ESMC candidate package versions differ from schema v3") + for name, value in packages.items(): + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ValueError(f"ESMC candidate package version is invalid for {name}") + + gpu = environment.get("gpu") + if not isinstance(gpu, Mapping) or set(gpu) != { + "name", + "capability", + "total_memory_bytes", + }: + raise ValueError("ESMC GPU identity differs from schema v3") + capability = gpu.get("capability") + if ( + not isinstance(gpu.get("name"), str) + or not str(gpu["name"]).strip() + or not isinstance(capability, list) + or len(capability) != 2 + or any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in capability + ) + or isinstance(gpu.get("total_memory_bytes"), bool) + or not isinstance(gpu.get("total_memory_bytes"), int) + or int(gpu["total_memory_bytes"]) <= 0 + ): + raise ValueError("ESMC GPU identity is malformed") + if environment != _candidate_environment_identity(): + raise ValueError("ESMC candidate environment differs from the active runtime") + + +def _validate_reference_environment(environment: object) -> None: + if not isinstance(environment, Mapping): + raise ValueError("ESMC reference environment is missing") + required = { + "cuda_device", + "cuda_device_capability", + "cuda_total_memory", + "cuda_runtime", + "packages", + "python", + "torch", + } + if not required.issubset(environment): + raise ValueError("ESMC reference environment fields are incomplete") + for name in ("cuda_device", "cuda_runtime", "python", "torch"): + value = environment.get(name) + if not isinstance(value, str) or not value.strip() or value == "unavailable": + raise ValueError(f"ESMC reference environment has invalid {name}") + capability = environment.get("cuda_device_capability") + if ( + not isinstance(capability, list) + or len(capability) != 2 + or any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in capability + ) + ): + raise ValueError("ESMC reference CUDA capability is invalid") + memory = environment.get("cuda_total_memory") + if isinstance(memory, bool) or not isinstance(memory, int) or memory <= 0: + raise ValueError("ESMC reference CUDA memory is invalid") + packages = environment.get("packages") + if not isinstance(packages, str) or not packages: + raise ValueError("ESMC reference package inventory is invalid") + try: + package_inventory = json.loads(packages) + except json.JSONDecodeError as error: + raise ValueError("ESMC reference package inventory is not JSON") from error + if not isinstance(package_inventory, Mapping): + raise ValueError("ESMC reference package inventory is not an object") + + +def _validate_esmc_environment_binding( + candidate_environment: Mapping[str, object], + reference_environment: Mapping[str, object], + locked_reference_environment: Mapping[str, object], +) -> None: + candidate_gpu = candidate_environment.get("gpu") + locked_runtime = locked_reference_environment.get("runtime") + locked_gpu = locked_runtime.get("gpu") if isinstance(locked_runtime, Mapping) else None + if not isinstance(candidate_gpu, Mapping) or not isinstance(locked_gpu, Mapping): + raise ValueError("ESMC candidate/reference GPU binding is malformed") + candidate_identity = { + "python": candidate_environment.get("python"), + "torch": candidate_environment.get("torch"), + "cuda_runtime": candidate_environment.get("cuda_runtime"), + "cuda_driver": candidate_environment.get("cuda_driver"), + "gpu": dict(candidate_gpu), + } + dynamic_reference_identity = { + "python": reference_environment.get("python"), + "torch": reference_environment.get("torch"), + "cuda_runtime": reference_environment.get("cuda_runtime"), + "cuda_driver": candidate_environment.get("cuda_driver"), + "gpu": { + "name": reference_environment.get("cuda_device"), + "capability": reference_environment.get("cuda_device_capability"), + "total_memory_bytes": reference_environment.get("cuda_total_memory"), + }, + } + locked_identity = { + "python": locked_runtime.get("python_version") + if isinstance(locked_runtime, Mapping) + else None, + "torch": locked_runtime.get("torch") if isinstance(locked_runtime, Mapping) else None, + "cuda_runtime": locked_runtime.get("cuda_runtime") + if isinstance(locked_runtime, Mapping) + else None, + "cuda_driver": locked_runtime.get("cuda_driver") + if isinstance(locked_runtime, Mapping) + else None, + "gpu": dict(locked_gpu), + } + if candidate_identity != dynamic_reference_identity: + raise ValueError("ESMC candidate and native reference environments differ") + if candidate_identity != locked_identity: + raise ValueError("ESMC candidate environment differs from the locked reference runtime") + + +def _validate_kernel_identity(kernel: object, backend: str) -> None: + if not isinstance(kernel, Mapping): + raise ValueError("ESMC kernel identity is malformed") + expected = _kernel_identity(backend) + if kernel != expected: + raise ValueError("ESMC kernel identity differs from the manifest or runtime") + + +def _expected_panel_identity(kind: object) -> dict[str, object]: + if not isinstance(kind, str): + raise ValueError("ESMC panel kind is invalid") + try: + batch = next( + candidate for candidate in esmc_calibration_batches() if candidate["kind"] == kind + ) + except StopIteration as error: + raise ValueError(f"Unsupported ESMC calibration panel: {kind!r}") from error + return _public_panel_identity(validate_esmc_calibration_batch(batch)) + + +def _report_sha256(payload: Mapping[str, object]) -> str: + digest_payload = dict(payload) + digest_payload.pop("report_sha256", None) + encoded = json.dumps( + digest_payload, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _validate_esmc_diagnostic_report( + payload: Mapping[str, object], + spec: ModelSpec, + *, + expected_candidate: Mapping[str, object] | None = None, +) -> None: + expected = { + "schema_version", + "model_id", + "candidate", + "reference", + "record_status", + "unavailability", + "configured_backend", + "effective_backend", + "dtype", + "panel", + "environment", + "kernel", + "panel_tensor_metrics", + "panel_logits_metrics", + "cases", + "published_band_violations", + "catastrophic_gate", + "release_gate", + "report_sha256", + } + if set(payload) != expected or payload.get("schema_version") != ESMC_DIAGNOSTIC_SCHEMA_VERSION: + raise ValueError("ESMC diagnostic fields differ from schema v3") + if payload.get("model_id") != spec.id or payload.get("dtype") != "bfloat16": + raise ValueError("ESMC diagnostic model or dtype identity differs from the request") + + backend = payload.get("configured_backend") + record_status = payload.get("record_status") + if backend not in (*ESMC_MEASURED_BACKENDS, *ESMC_UNAVAILABLE_BACKENDS): + raise ValueError("ESMC diagnostic backend identity is invalid") + if record_status not in {"measured", "unavailable"}: + raise ValueError("ESMC diagnostic record status is invalid") + + candidate = payload.get("candidate") + if not isinstance(candidate, Mapping) or set(candidate) != { + "repo_id", + "manifest_revision", + "resolved_commit", + "checkpoint_repo_id", + "checkpoint_revision", + "weights_revision", + "runtime_revision", + "source_tree_sha256", + "runtime_bundle_sha256", + }: + raise ValueError("ESMC candidate identity differs from schema v3") + if ( + candidate.get("repo_id") != spec.fast.repo_id + or candidate.get("manifest_revision") != spec.fast.revision + ): + raise ValueError("ESMC candidate repository identity differs from the manifest") + expected_checkpoint = { + "checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "checkpoint_revision": spec.artifact_checkpoint.revision, + "weights_revision": spec.artifact_checkpoint.revision, + } + for name, expected_value in expected_checkpoint.items(): + value = candidate.get(name) + if value != expected_value: + raise ValueError(f"ESMC candidate {name} differs from the manifest") + resolved_commit = candidate.get("resolved_commit") + if resolved_commit is not None and resolved_commit != spec.fast.revision: + raise ValueError("ESMC candidate resolved commit differs from the manifest") + source_tree_sha256 = candidate.get("source_tree_sha256") + runtime_bundle_sha256 = candidate.get("runtime_bundle_sha256") + runtime_revision = candidate.get("runtime_revision") + source_digest = _require_sha256( + source_tree_sha256, + "ESMC candidate source tree", + ) + _require_sha256(runtime_bundle_sha256, "ESMC candidate runtime bundle") + _require_runtime_revision( + runtime_revision, + source_digest, + "ESMC candidate runtime revision", + ) + if expected_candidate is not None and candidate != expected_candidate: + raise ValueError("ESMC candidate identity differs from the validated artifact identity") + + reference = payload.get("reference") + if not isinstance(reference, Mapping) or set(reference) != { + "repo_id", + "revision", + "state_transform", + "environment", + "reference_environment", + "reference_sources", + }: + raise ValueError("ESMC reference identity differs from schema v3") + if ( + reference.get("repo_id") != spec.official.repo_id + or reference.get("revision") != spec.official.revision + or reference.get("state_transform") != spec.family.state_transform + ): + raise ValueError("ESMC reference identity differs from the manifest") + dynamic_reference_environment = reference.get("environment") + _validate_reference_environment(dynamic_reference_environment) + locked_reference_environment = _validated_biohub_reference_environment( + reference.get("reference_environment") + ) + _validated_biohub_reference_sources(reference.get("reference_sources")) + + candidate_environment = payload.get("environment") + _validate_candidate_environment(candidate_environment) + if not isinstance(candidate_environment, Mapping) or not isinstance( + dynamic_reference_environment, Mapping + ): + raise ValueError("ESMC candidate/reference environment binding is malformed") + _validate_esmc_environment_binding( + candidate_environment, + dynamic_reference_environment, + locked_reference_environment, + ) + + _validate_kernel_identity(payload.get("kernel"), str(backend)) + + panel = payload.get("panel") + if not isinstance(panel, Mapping) or set(panel) != { + "schema_version", + "kind", + "seed", + "definition_sha256", + "cases", + }: + raise ValueError("ESMC panel identity differs from schema v3") + if panel != _expected_panel_identity(panel.get("kind")): + raise ValueError("ESMC panel identity differs from the immutable definition") + panel_cases = panel.get("cases") + cases = payload.get("cases") + if ( + not isinstance(panel_cases, list) + or not isinstance(cases, list) + or len(panel_cases) != len(cases) + ): + raise ValueError("ESMC panel and metric cases are not aligned") + + identity_fields = { + "case_id", + "sequence_length", + "sequence_sha256", + "source", + "source_sha256", + } + violations = payload.get("published_band_violations") + if not isinstance(violations, list) or any( + not isinstance(value, str) or not value.strip() for value in violations + ): + raise ValueError("ESMC published-band violations must be a string list") + if record_status == "unavailable": + if backend not in ESMC_UNAVAILABLE_BACKENDS: + raise ValueError("Only locked Flash backends may use unavailable records") + if payload.get("effective_backend") is not None: + raise ValueError("Unavailable ESMC records must not claim effective dispatch") + if payload.get("unavailability") != _esmc_unavailability_identity( + str(backend), locked_reference_environment + ): + raise ValueError("ESMC structured unavailability identity is invalid") + if payload.get("catastrophic_gate") != "not_run": + raise ValueError("Unavailable ESMC records must mark the catastrophe gate not run") + if payload.get("release_gate") != {"mode": "availability", "status": "unavailable"}: + raise ValueError("Unavailable ESMC release-gate identity is invalid") + if ( + payload.get("panel_tensor_metrics") is not None + or payload.get("panel_logits_metrics") is not None + or violations + ): + raise ValueError("Unavailable ESMC records must not contain numerical measurements") + if cases != panel_cases: + raise ValueError("Unavailable ESMC cases must be immutable panel identities only") + else: + if backend not in ESMC_MEASURED_BACKENDS: + raise ValueError( + "Current GH200 release measurements are limited to eager, SDPA, and Flex" + ) + if payload.get("effective_backend") != backend: + raise ValueError("ESMC diagnostic backend identity indicates fallback") + if payload.get("unavailability") is not None: + raise ValueError("Measured ESMC records must not carry unavailability metadata") + if payload.get("catastrophic_gate") != "passed": + raise ValueError("Measured ESMC records require a passed catastrophe gate") + if payload.get("release_gate") != _release_gate_identity(str(backend)): + raise ValueError("ESMC diagnostic release-gate identity is invalid") + + metrics = payload.get("panel_tensor_metrics") + _validate_metric_distribution(metrics, context="ESMC panel") + panel_logits_metrics = payload.get("panel_logits_metrics") + _validate_logits_metric_payload(panel_logits_metrics) + if not isinstance(metrics, list): + raise ValueError("ESMC panel tensor metrics must be an ordered list") + panel_layout = [ + (metric["output"], metric["layer_index"]) + for metric in metrics + if isinstance(metric, Mapping) + ] + for panel_case, case in zip(panel_cases, cases, strict=True): + if not isinstance(panel_case, Mapping) or set(panel_case) != identity_fields: + raise ValueError("ESMC panel case identity differs from schema v3") + if not isinstance(case, Mapping) or set(case) != identity_fields | { + "tensor_metrics", + "logits_metrics", + }: + raise ValueError("ESMC case distribution differs from schema v3") + if any(case.get(name) != panel_case.get(name) for name in identity_fields): + raise ValueError("ESMC case distribution is misaligned with the panel") + _require_sha256(case.get("sequence_sha256"), "ESMC case sequence") + source_sha256 = case.get("source_sha256") + if source_sha256 is not None: + _require_sha256(source_sha256, "ESMC case source") + case_metrics = case.get("tensor_metrics") + _validate_metric_distribution( + case_metrics, + context=f"ESMC case {case.get('case_id')}", + ) + if not isinstance(case_metrics, list): + raise ValueError("ESMC case tensor metrics must be an ordered list") + case_layout = [ + (metric["output"], metric["layer_index"]) + for metric in case_metrics + if isinstance(metric, Mapping) + ] + if case_layout != panel_layout: + raise ValueError("ESMC case metric layout differs from the panel") + case_logits_metrics = case.get("logits_metrics") + _validate_logits_metric_payload(case_logits_metrics) + if (case_logits_metrics is None) != (panel_logits_metrics is None): + raise ValueError("ESMC case logits metrics differ from the panel") + report_digest = _require_sha256(payload.get("report_sha256"), "ESMC report") + if report_digest != _report_sha256(payload): + raise ValueError("ESMC report digest does not match its payload") + + +def _build_esmc_diagnostic_report( + spec: ModelSpec, + candidate: object, + official: object, + residue_mask: torch.Tensor, + *, + backend: str, + effective_backend: str, + context: str, + calibration_batch: Mapping[str, object], + model: torch.nn.Module, + reference_metadata: Mapping[str, object], +) -> dict[str, object]: + # residue_mask: (b, l) + panel = validate_esmc_calibration_batch(calibration_batch) + records, logits = _collect_output_metrics( + spec, + candidate, + official, + residue_mask, + context, + ) + _assert_esmc_catastrophic_metrics(records, logits, context) + violations = _esmc_published_band_violations(records, logits) + case_metrics, case_violations = _case_metric_distributions( + spec, + candidate, + official, + residue_mask, + panel, + context, + ) + violations.extend(case_violations) + candidate_identity = _candidate_identity(spec, model) + reference_identity = _reference_identity(spec, reference_metadata) + candidate_environment = _candidate_environment_identity() + payload: dict[str, object] = { + "schema_version": ESMC_DIAGNOSTIC_SCHEMA_VERSION, + "model_id": spec.id, + "candidate": candidate_identity, + "reference": reference_identity, + "record_status": "measured", + "unavailability": None, + "configured_backend": backend, + "effective_backend": effective_backend, + "dtype": "bfloat16", + "panel": _public_panel_identity(panel), + "environment": candidate_environment, + "kernel": _kernel_identity(backend), + "panel_tensor_metrics": _structured_metric_payloads(candidate, records), + "panel_logits_metrics": _logits_metric_payload(logits), + "cases": case_metrics, + "published_band_violations": violations, + "catastrophic_gate": "passed", + "release_gate": _release_gate_identity(backend), + } + payload["report_sha256"] = _report_sha256(payload) + _validate_esmc_diagnostic_report( + payload, + spec, + expected_candidate=candidate_identity, + ) + return payload + + +def _build_esmc_unavailable_report( + spec: ModelSpec, + *, + backend: str, + calibration_batch: Mapping[str, object], + model: torch.nn.Module, + reference_metadata: Mapping[str, object], +) -> dict[str, object]: + if backend not in ESMC_UNAVAILABLE_BACKENDS: + raise ValueError(f"ESMC backend {backend!r} is not unavailable on the locked target") + panel = _public_panel_identity(validate_esmc_calibration_batch(calibration_batch)) + panel_cases = panel["cases"] + if not isinstance(panel_cases, list): + raise ValueError("ESMC unavailable report panel cases are not an ordered list") + candidate_identity = _candidate_identity(spec, model) + reference_identity = _reference_identity(spec, reference_metadata) + locked_reference_environment = reference_identity["reference_environment"] + if not isinstance(locked_reference_environment, Mapping): + raise ValueError("ESMC unavailable report omits its locked reference environment") + candidate_environment = _candidate_environment_identity() + payload: dict[str, object] = { + "schema_version": ESMC_DIAGNOSTIC_SCHEMA_VERSION, + "model_id": spec.id, + "candidate": candidate_identity, + "reference": reference_identity, + "record_status": "unavailable", + "unavailability": _esmc_unavailability_identity(backend, locked_reference_environment), + "configured_backend": backend, + "effective_backend": None, + "dtype": "bfloat16", + "panel": panel, + "environment": candidate_environment, + "kernel": _kernel_identity(backend), + "panel_tensor_metrics": None, + "panel_logits_metrics": None, + "cases": [dict(case) for case in panel_cases], + "published_band_violations": [], + "catastrophic_gate": "not_run", + "release_gate": {"mode": "availability", "status": "unavailable"}, + } + payload["report_sha256"] = _report_sha256(payload) + _validate_esmc_diagnostic_report( + payload, + spec, + expected_candidate=candidate_identity, + ) + return payload + + +def _write_esmc_diagnostic_report( + spec: ModelSpec, + payload: Mapping[str, object], +) -> Path: + _validate_esmc_diagnostic_report(payload, spec) + report_root = Path(os.environ.get("FASTPLMS_DIAGNOSTIC_REPORTS", "artifacts/diagnostics/esmc")) + report_root.mkdir(parents=True, exist_ok=True) + panel = payload["panel"] + assert isinstance(panel, Mapping) + report_name = f"{spec.id}-{payload['configured_backend']}-{panel['kind']}.json" + report_path = report_root / report_name + encoded = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if report_path.exists(): + if report_path.read_text(encoding="utf-8") == encoded: + return report_path + raise RuntimeError(f"Refusing to replace different ESMC evidence: {report_path}") + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=report_root, + prefix=f".{report_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary.write(encoded) + temporary.flush() + os.fsync(temporary.fileno()) + temporary_path = Path(temporary.name) + os.replace(temporary_path, report_path) + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + return report_path + + +def _record_esmc_diagnostic( + spec: ModelSpec, + candidate: object, + official: object, + residue_mask: torch.Tensor, + *, + backend: str, + effective_backend: str, + context: str, + calibration_batch: Mapping[str, object], + model: torch.nn.Module, + reference_metadata: Mapping[str, object], + warn_on_published_band: bool, +) -> dict[str, object]: + # residue_mask: (b, l) + payload = _build_esmc_diagnostic_report( + spec, + candidate, + official, + residue_mask, + backend=backend, + effective_backend=effective_backend, + context=context, + calibration_batch=calibration_batch, + model=model, + reference_metadata=reference_metadata, + ) + report_path = _write_esmc_diagnostic_report(spec, payload) + violations = payload["published_band_violations"] + assert isinstance(violations, list) + if warn_on_published_band and violations: + warnings.warn( + f"{spec.id} configured backend={backend}, effective backend={effective_backend}: " + f"{len(violations)} diagnostic metric(s) are outside the published ESMC " + f"backend bands; catastrophic biological gates passed. Report: {report_path}", + UserWarning, + stacklevel=2, + ) + return payload + + +def _assert_and_record_esmc_diagnostic( + spec: ModelSpec, + candidate: object, + official: object, + residue_mask: torch.Tensor, + *, + backend: str, + effective_backend: str, + context: str, + calibration_batch: Mapping[str, object], + model: torch.nn.Module, + reference_metadata: Mapping[str, object], +) -> dict[str, object]: + # residue_mask: (b, l) + return _record_esmc_diagnostic( + spec, + candidate, + official, + residue_mask, + backend=backend, + effective_backend=effective_backend, + context=context, + calibration_batch=calibration_batch, + model=model, + reference_metadata=reference_metadata, + warn_on_published_band=True, + ) + + +def _run_native_inference( + spec: ModelSpec, + result_dir: Path, + *, + precision: str, + backend: str | None, + package_source: bool = False, + tensor_path: Path | None = None, + context_suffix: str = "native", + calibration_batch: Mapping[str, object] | None = None, + reference_metadata: Mapping[str, object] | None = None, +) -> None: + device = torch.device("cuda") + dtype = torch.float32 if precision == "fp32" else torch.bfloat16 + tensors = load_file(tensor_path or result_dir / f"{precision}.safetensors", device="cpu") + use_bf16_autocast = ( + dtype == torch.bfloat16 and spec.family.bf16_execution == "fp32_parameters_autocast" + ) + model_dtype = torch.float32 if use_bf16_autocast else dtype + if package_source: + fast = _load_package_generation_model(spec, device, model_dtype) + else: + fast = _load_fast(spec, device, model_dtype) + effective_backend: str | None = backend + if backend is not None: + fast.set_attn_implementation(backend) + effective_backend = getattr(fast.config, "_attn_implementation", None) + if effective_backend is None: + effective_backend = getattr(fast.config, "attn_implementation", None) + assert effective_backend == backend, ( + f"{spec.id}: requested {backend!r}, resolved {effective_backend!r}" + ) + + inputs = { + name.removeprefix("input__"): value.to(device) + for name, value in tensors.items() + if name.startswith("input__") + } + # residue_mask: (b, l) + residue_mask = tensors["residue_mask"].to(device).bool() + if dtype == torch.float32: + numeric_context = strict_fp32_matmul() + elif use_bf16_autocast: + numeric_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16) + else: + numeric_context = contextlib.nullcontext() + with torch.inference_mode(), numeric_context: + candidate = fast(**inputs, output_hidden_states=True) + official = _official_output(tensors, device) + contract = _numeric_contract(spec, dtype, backend) + context = f"{spec.id}:{precision}:{backend or 'default'}:{context_suffix}" + if ( + spec.family.architecture == "ESMC" + and dtype == torch.bfloat16 + and backend == "flex_attention" + ): + assert backend is not None + if calibration_batch is None: + _assert_esmc_alternate_backend_outputs( + spec, + candidate, + official, + residue_mask, + context, + ) + else: + if reference_metadata is None: + raise ValueError("ESMC calibration requires native reference metadata") + _assert_and_record_esmc_diagnostic( + spec, + candidate, + official, + residue_mask, + backend=backend, + effective_backend=str(effective_backend), + context=context, + calibration_batch=calibration_batch, + model=fast, + reference_metadata=reference_metadata, + ) + else: + _assert_outputs( + spec, + candidate, + official, + residue_mask, + contract, + context, + ) + if spec.family.architecture == "ESMC" and backend in (None, "sdpa"): + _assert_esmc_sdpa_exact( + candidate, + official, + context, + ) + if ( + spec.family.architecture == "ESMC" + and dtype == torch.bfloat16 + and backend in {"eager", "sdpa"} + and calibration_batch is not None + ): + if reference_metadata is None: + raise ValueError("ESMC calibration requires native reference metadata") + _record_esmc_diagnostic( + spec, + candidate, + official, + residue_mask, + backend=backend, + effective_backend=str(effective_backend), + context=context, + calibration_batch=calibration_batch, + model=fast, + reference_metadata=reference_metadata, + warn_on_published_band=False, + ) + if backend == "sdpa": + for unavailable_backend in ESMC_UNAVAILABLE_BACKENDS: + _write_esmc_diagnostic_report( + spec, + _build_esmc_unavailable_report( + spec, + backend=unavailable_backend, + calibration_batch=calibration_batch, + model=fast, + reference_metadata=reference_metadata, + ), + ) + del fast, candidate, official, tensors + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in SEQUENCE_SPECS]) +def test_native_every_checkpoint_bf16_inference(spec: ModelSpec) -> None: + """Every checkpoint matches one native mixed-length BF16 result.""" + + _, result_dir = _result(spec) + _run_native_inference(spec, result_dir, precision="bf16", backend=None) + + +NATIVE_REPRESENTATIVE_CASES = [ + pytest.param( + spec, + precision, + backend, + id=f"{spec.id}-{precision}-{backend}", + marks=[pytest.mark.large] if spec.size_category == "xlarge" else [], + ) + for spec in SEQUENCE_SPECS + if spec.is_deep_reference + for backend in spec.family.attention + if not (spec.family.id == "esm_plusplus" and backend in ESMC_UNAVAILABLE_BACKENDS) + for precision, dtype_name in (("fp32", "float32"), ("bf16", "bfloat16")) + if dtype_name in REGISTRY.supported_attention_dtypes(spec.family.id, backend) +] + + +@pytest.mark.parametrize(("spec", "precision", "backend"), NATIVE_REPRESENTATIVE_CASES) +def test_native_representatives_all_backends( + spec: ModelSpec, + precision: str, + backend: str, +) -> None: + """Each representative matches native FP32 and BF16 for advertised backends.""" + + _, result_dir = _result(spec) + _run_native_inference(spec, result_dir, precision=precision, backend=backend) + + +def _esmc_calibration_marks(spec: ModelSpec, backend: str, kind: str) -> list[Any]: + del backend, kind + return [pytest.mark.large] if spec.size_category == "xlarge" else [] + + +@pytest.mark.parametrize( + ("spec", "backend", "kind"), + [ + pytest.param( + spec, + backend, + kind, + id=f"{spec.id}-{backend}-{kind}", + marks=_esmc_calibration_marks(spec, backend, kind), + ) + for spec in SEQUENCE_SPECS + if spec.family.id == "esm_plusplus" + for backend in ESMC_MEASURED_BACKENDS + if backend in spec.family.attention + and "bfloat16" in REGISTRY.supported_attention_dtypes(spec.family.id, backend) + for kind in ("generated_kernel_boundary", "real_biological_holdout") + ], +) +def test_esmc_bf16_calibration_and_biological_holdout( + spec: ModelSpec, + backend: str, + kind: str, +) -> None: + """Calibrate BF16 parity on pinned shape and biological panels.""" + + metadata, result_dir = _result(spec) + batches = metadata.get("calibration_batches") + assert isinstance(batches, list) + batch = next(item for item in batches if item["kind"] == kind) + assert isinstance(batch, Mapping) + validate_esmc_calibration_batch(batch) + assert batch["seed"] == 42 + expected_biological = {case["case_id"]: case for case in load_esmc_biological_holdout()} + if kind == "generated_kernel_boundary": + observed_lengths = tuple(case["sequence_length"] for case in batch["cases"]) + assert observed_lengths == ESMC_BOUNDARY_LENGTHS + else: + assert tuple(case["case_id"] for case in batch["cases"]) == tuple(expected_biological) + for case in batch["cases"]: + sequence = case["sequence"] + assert len(sequence) == case["sequence_length"] + assert hashlib.sha256(sequence.encode("ascii")).hexdigest() == case["sequence_sha256"] + if kind == "real_biological_holdout": + expected = expected_biological[case["case_id"]] + assert { + name: case[name] + for name in ( + "case_id", + "sequence", + "sequence_sha256", + "source", + "source_sha256", + ) + } == expected + _run_native_inference( + spec, + result_dir, + precision="bf16", + backend=backend, + tensor_path=result_dir / "calibration" / f"{kind}.safetensors", + context_suffix=kind, + calibration_batch=batch, + reference_metadata=metadata, + ) + + +@pytest.mark.parametrize( + "spec", + [_parameter(spec) for spec in SEQUENCE_SPECS if spec.id in {"dplm_150m", "dplm2_150m"}], +) +def test_native_dplm_package_source_fp32(spec: ModelSpec) -> None: + """Current repository source matches native DPLM-family FP32 inference.""" + + _, result_dir = _result(spec) + _run_native_inference( + spec, + result_dir, + precision="fp32", + backend=None, + package_source=True, + ) + + +@pytest.mark.parametrize( + ("spec", "backend"), + [ + pytest.param(spec, backend, id=f"{spec.id}-{backend}") + for spec in SEQUENCE_SPECS + if spec.id in {"dplm_150m", "dplm2_150m"} + for backend in spec.family.attention + if "bfloat16" in REGISTRY.supported_attention_dtypes(spec.family.id, backend) + ], +) +def test_native_dplm_package_source_bf16(spec: ModelSpec, backend: str) -> None: + """Current repository source matches native DPLM BF16 on every supported backend.""" + + _, result_dir = _result(spec) + _run_native_inference( + spec, + result_dir, + precision="bf16", + backend=backend, + package_source=True, + ) + + +@pytest.mark.parametrize( + "spec", + [ + pytest.param(REGISTRY["dplm_150m"], id="dplm_150m"), + pytest.param(REGISTRY["dplm2_150m"], id="dplm2_150m"), + ], +) +def test_native_dplm_sdpa_uses_fp32_storage_and_meets_every_hidden_target( + spec: ModelSpec, +) -> None: + """Each manifest-declared DPLM AMP path passes every hidden-state target.""" + + assert spec.family.bf16_execution == "fp32_parameters_autocast" + _, result_dir = _result(spec) + device = torch.device("cuda") + tensors = load_file(result_dir / "bf16.safetensors", device="cpu") + model = _load_package_generation_model(spec, device, torch.float32) + assert {parameter.dtype for parameter in model.parameters()} == {torch.float32} + model.set_attn_implementation("sdpa") + inputs = { + name.removeprefix("input__"): value.to(device) + for name, value in tensors.items() + if name.startswith("input__") + } + # residue_mask: (b, l) + residue_mask = tensors["residue_mask"].to(device).bool() + + with ( + torch.inference_mode(), + torch.autocast(device_type="cuda", dtype=torch.bfloat16), + ): + candidate = model(**inputs, output_hidden_states=True) + official = _official_output(tensors, device) + assert len(candidate.hidden_states) == len(official.hidden_states) == 31 + _assert_outputs( + spec, + candidate, + official, + residue_mask, + BF16_CONTRACT, + f"{spec.id}:bf16-autocast:sdpa:fp32-storage", + ) + del model, candidate, official, tensors + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize( + "spec", + [ + _parameter(spec) + for spec in SEQUENCE_SPECS + if spec.family.id in {"dplm", "dplm2"} and spec.id != "dplm2_3b" + ], +) +def test_native_dplm_generation(spec: ModelSpec) -> None: + """DPLM-family output tokens match the isolated official public sampler.""" + + metadata, _ = _result(spec) + contract = metadata.get("generation") + assert isinstance(contract, dict), f"{spec.id}: native result omits generation" + device = torch.device("cuda") + fast = _load_package_generation_model(spec, device) + # input_tokens: (...) + input_tokens = torch.tensor(contract["input_tokens"], device=device) + torch.manual_seed(int(contract["seed"])) + torch.cuda.manual_seed_all(int(contract["seed"])) + with torch.inference_mode(), strict_fp32_matmul(): + generated = fast.generate(input_tokens=input_tokens, **contract["kwargs"]) + if isinstance(generated, dict): + generated = generated["output_tokens"] + # expected: (...) + expected = torch.tensor(contract["output_tokens"], device=device) + assert torch.equal(generated, expected), f"{spec.id}: generated tokens differ" + del fast, generated + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize( + "spec", + [_parameter(spec) for spec in SEQUENCE_SPECS if spec.family.id == "ankh"], +) +def test_native_ankh_explicit_decoder_prompt_generation(spec: ModelSpec) -> None: + """ANKH tokens match native T5 generation from the recorded decoder prompt.""" + + metadata, _ = _result(spec) + contract = metadata.get("generation") + assert isinstance(contract, dict), f"{spec.id}: native result omits generation" + assert contract["interface"] == "T5ForConditionalGeneration.generate" + assert contract["decoder_prompt_contract"] == "explicit-task-prompt" + device = torch.device("cuda") + fast = _load_package_generation_model(spec, device) + # input_ids: (...) + input_ids = torch.tensor(contract["input_ids"], device=device) + # attention_mask: (...) + attention_mask = torch.tensor(contract["attention_mask"], device=device) + # decoder_input_ids: (...) + decoder_input_ids = torch.tensor(contract["decoder_input_ids"], device=device) + # decoder_attention_mask: (...) + decoder_attention_mask = torch.tensor( + contract["decoder_attention_mask"], + device=device, + ) + assert _tensor_digest(decoder_input_ids)["sha256"] == (contract["decoder_input_fingerprint"]) + torch.manual_seed(int(contract["seed"])) + torch.cuda.manual_seed_all(int(contract["seed"])) + with torch.inference_mode(), strict_fp32_matmul(): + generated = fast.generate( + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + **contract["kwargs"], + ) + # expected: (...) + expected = torch.tensor(contract["output_tokens"], device=device) + assert torch.equal(generated, expected), f"{spec.id}: generated tokens differ" + del fast, generated + gc.collect() + torch.cuda.empty_cache() + + +def test_native_dplm2_3b_official_generation_limitation() -> None: + """The pinned public 3B sampler is unavailable, not parity-passing.""" + + spec = REGISTRY["dplm2_3b"] + metadata, _ = _result(spec) + assert "generation" not in metadata + assert metadata.get("generation_limitation") == DPLM2_3B_GENERATION_LIMITATION diff --git a/tests/release/__init__.py b/tests/release/__init__.py new file mode 100644 index 0000000..d8c908c --- /dev/null +++ b/tests/release/__init__.py @@ -0,0 +1 @@ +"""Tests for built remote-code artifacts and published releases.""" diff --git a/tests/release/test_artifacts.py b/tests/release/test_artifacts.py new file mode 100644 index 0000000..2f8ce9f --- /dev/null +++ b/tests/release/test_artifacts.py @@ -0,0 +1,2575 @@ +from __future__ import annotations + +import configparser +import hashlib +import json +import os +import subprocess +import sys +import textwrap +import pytest +import torch +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from safetensors.torch import load_file, save_file + +from fastplms.registry import ( + CheckpointSource, + FileDigest, + ModelFamily, + ModelRegistry, + ModelSpec, + UpstreamSource, + get_model_registry, + load_model_registry, +) +from tools.artifacts import ( + ArtifactError, + build_artifact, + canonicalize_checkpoint_weights, + hash_file, + validate_artifact, + validate_repository_legal_inventory, + validate_weight_artifact, + verify_checkpoint, +) +from tools.artifacts.build import ( + _RELEASE_TOOL_SCOPE_PATHS, + _canonical_state_sha256, + _checkpoint_identity_hash, + _conversion_equality_attestation, + _copy_attention_kernel_lock, + _copy_official_tokenizer_assets, + _git_runtime_revision, + _is_weight_file, + _materialize_model_card, + _provenance, + _render_artifact_requirements, + _tokenizer_checkpoint, + _validate_vendor_revisions, + _validated_release_tool_snapshot, + render_model_card, +) + + +ROOT = Path(__file__).resolve().parents[2] + + +def _canonical_text_sha256(path: Path) -> str: + raw = path.read_bytes() + canonical = raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n") + return hashlib.sha256(canonical).hexdigest() + + +def _initialize_release_tool_repository(root: Path) -> None: + for relative_name in _RELEASE_TOOL_SCOPE_PATHS: + path = root.joinpath(*relative_name.split("/")) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"# release tool {relative_name}\n", encoding="utf-8") + git = ["git", "-c", f"safe.directory={root.as_posix()}"] + subprocess.run( + [*git, "init", "--initial-branch=main"], + cwd=root, + check=True, + capture_output=True, + ) + subprocess.run([*git, "config", "user.email", "tests@example.invalid"], cwd=root, check=True) + subprocess.run([*git, "config", "user.name", "FastPLMs Tests"], cwd=root, check=True) + subprocess.run([*git, "add", "."], cwd=root, check=True) + subprocess.run( + [*git, "commit", "-m", "immutable release tools"], + cwd=root, + check=True, + capture_output=True, + ) + + +@pytest.mark.parametrize( + "relative_name", + ( + "tools/artifacts/build.py", + "tools/artifacts/publish.py", + "tools/artifacts/offline_probe.py", + "tools/conversion/state_transforms.py", + ), +) +def test_release_tool_snapshot_rejects_dirty_critical_tool( + tmp_path: Path, + relative_name: str, +) -> None: + _initialize_release_tool_repository(tmp_path) + path = tmp_path.joinpath(*relative_name.split("/")) + path.write_text(path.read_text(encoding="utf-8") + "# dirty\n", encoding="utf-8") + + with pytest.raises(ArtifactError, match="release tools must be tracked and clean"): + _validated_release_tool_snapshot(tmp_path) + + +def test_release_tool_snapshot_rejects_untracked_scope_growth(tmp_path: Path) -> None: + _initialize_release_tool_repository(tmp_path) + (tmp_path / "tools" / "artifacts" / "new_validator.py").write_text( + "# untracked validation bypass\n", + encoding="utf-8", + ) + + with pytest.raises(ArtifactError, match="release tools must be tracked and clean"): + _validated_release_tool_snapshot(tmp_path) + + +def test_materialized_model_card_keeps_runtime_identity_out_of_user_facing_text() -> None: + template = render_model_card(get_model_registry()["esm2_8m"]) + runtime_revision = "a" * 40 + source_tree_sha256 = "b" * 64 + runtime_bundle_sha256 = "c" * 64 + + card = _materialize_model_card( + template, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ) + + assert "" not in card + assert "FastPLMs.git@" not in card + assert runtime_revision not in card + assert source_tree_sha256 not in card + assert runtime_bundle_sha256 not in card + assert "- Runtime revision: recorded separately" in card + assert "- Source-tree and runtime-bundle SHA-256: recorded" in card + + +def test_shared_sources_are_in_runtime_artifacts() -> None: + """Keep remote-code artifacts closed over their package-source imports.""" + + registry = get_model_registry() + required = { + "esm2": {"models/_esm_rotary.py"}, + "dplm": {"models/_diffusion_generation.py", "models/_esm_rotary.py"}, + "dplm2": {"models/_diffusion_generation.py", "models/_esm_rotary.py"}, + "esmfold": {"models/_esm_rotary.py"}, + "esmfold2": {"models/esm_plusplus"}, + } + package_root = ROOT / "src" / "fastplms" + for family_id, paths in required.items(): + family = registry.families[family_id] + assert paths.issubset(family.runtime_paths) + for relative_path in paths: + assert (package_root / relative_path).exists() + + +@pytest.mark.parametrize( + ("model_id", "required", "excluded"), + ( + ("esm2_8m", ("torch>=2.13,<2.14", "kernels>=0.15,<0.16"), ("biotite",)), + ("ankh_base", ("torch>=2.13,<2.14",), ("kernels", "biotite")), + ("boltz2", ("torch>=2.13,<2.14", "biotite>=1.4,<2"), ("kernels",)), + ), +) +def test_artifact_requirements_match_advertised_runtime( + model_id: str, + required: tuple[str, ...], + excluded: tuple[str, ...], +) -> None: + payloads = { + relative_name: (ROOT / relative_name).read_bytes() + for relative_name in ( + "requirements/core.in", + "requirements/features/flash.in", + "requirements/features/structure.in", + ) + } + rendered = _render_artifact_requirements(get_model_registry()[model_id], payloads) + + for requirement in required: + assert requirement in rendered + for requirement in excluded: + assert requirement not in rendered + assert "fastplms" not in "\n".join( + line for line in rendered.splitlines() if not line.startswith("#") + ).lower() + + +def test_esmfold2_runtime_asset_provenance_records_trust_and_offline_boundary() -> None: + registry = get_model_registry() + provenance = _provenance( + registry, + registry["esmfold2"], + {}, + runtime_revision="a" * 40, + source_tree_sha256="b" * 64, + runtime_bundle_sha256="c" * 64, + release_tool_revision="d" * 40, + release_tool_sha256="e" * 64, + ) + + assert provenance["runtime_assets"] == [ + { + "id": "esmfold2_ccd", + "repository": "biohub/ESMFold2", + "revision": "1ebf0e3481a5184eb6171d40615c79e384b48796", + "path": "ccd.pkl", + "sha256": "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5", + "size": 417306584, + "license": "MIT", + "consumer_family": "esmfold2", + "trust_kind": "hash_pinned_pickle", + "offline_behavior": "requires_cached_verified_file", + "cache_identity": hashlib.sha256( + b"biohub/ESMFold2@1ebf0e3481a5184eb6171d40615c79e384b48796:" + b"ccd.pkl:9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5:" + b"417306584" + ).hexdigest(), + } + ] + assert len(provenance["runtime_assets"][0]["cache_identity"]) == 64 + + +def test_e1_runtime_artifact_closes_over_split_modules() -> None: + """Keep every E1 responsibility module inside the bundled runtime tree.""" + + family = get_model_registry().families["e1"] + assert "models/e1" in family.runtime_paths + source_root = ROOT / "src" / "fastplms" / "models" / "e1" + for source_name in ( + "attention.py", + "cache.py", + "modeling_e1.py", + "preparation.py", + "retrieval.py", + ): + assert (source_root / source_name).is_file() + + +@pytest.mark.parametrize("model_id", ("esm2_8m", "esmc_small", "dplm_150m")) +def test_flash_artifact_build_embeds_the_kernel_lock( + model_id: str, + tmp_path: Path, +) -> None: + """Exercise the kernel-lock build step without requiring checkpoint weights.""" + + registry = get_model_registry() + _copy_attention_kernel_lock( + ROOT, + tmp_path, + registry, + registry[model_id], + ) + assert (tmp_path / "kernels.lock").read_bytes() == (ROOT / "kernels.lock").read_bytes() + + +def _synthetic_registry(source_root: Path, checkpoint: Path) -> tuple[ModelRegistry, ModelSpec]: + requirements = source_root / "requirements" + (requirements / "features").mkdir(parents=True) + (requirements / "core.in").write_text( + "torch>=2.13,<2.14\ntransformers>=5.13,<5.14\n", + encoding="utf-8", + ) + (requirements / "features" / "flash.in").write_text( + "kernels>=0.15,<0.16\n", + encoding="utf-8", + ) + (requirements / "features" / "structure.in").write_text( + "biotite>=1.4,<2\n", + encoding="utf-8", + ) + package = source_root / "src" / "fastplms" + (package / "models" / "toy").mkdir(parents=True) + (package / "__init__.py").write_text("__version__ = '1.0.0'\n", encoding="utf-8") + (package / "models" / "toy" / "modeling_toy.py").write_text( + "class ToyConfig: pass\nclass ToyModel: pass\n", encoding="utf-8" + ) + upstream_root = source_root / "vendor" / "upstream" / "toy" + upstream_root.mkdir(parents=True) + canonical_license = upstream_root / "LICENSE" + canonical_license.write_text("Synthetic test license\n", encoding="utf-8") + inventory_root = source_root / "LICENSES" / "toy" + inventory_root.mkdir(parents=True) + distribution_license = inventory_root / "LICENSE" + distribution_license.write_text("Synthetic test license\n", encoding="utf-8") + project_license = source_root / "LICENSE" + project_license.write_text("FastPLMs test license\n", encoding="utf-8") + third_party_notices = source_root / "THIRD_PARTY_NOTICES.md" + third_party_notices.write_text("Synthetic test notice\n", encoding="utf-8") + + config = checkpoint / "config.json" + weight = checkpoint / "model.safetensors" + config.write_text('{"model_type": "toy"}\n', encoding="utf-8") + save_file( + { + "linear.bias": torch.arange(4, dtype=torch.float32), + "linear.weight": torch.arange(16, dtype=torch.float32).reshape(4, 4), + }, + weight, + metadata={"format": "pt"}, + ) + fast = CheckpointSource( + repo_id="Synthyra/ToyModel", + revision="1" * 40, + files=( + FileDigest("config.json", "git-sha1", hash_file(config, "git-sha1")), + FileDigest("model.safetensors", "sha256", hash_file(weight)), + ), + ) + official = CheckpointSource( + repo_id="upstream/ToyModel", + revision="2" * 40, + files=(FileDigest("model.safetensors", "sha256", "3" * 64),), + ) + upstream = UpstreamSource( + id="toy", + path="vendor/upstream/toy", + url="https://github.com/example/toy.git", + revision="4" * 40, + license_expression="MIT", + license_files=("LICENSE",), + license_digests=( + FileDigest("LICENSE", "sha256", _canonical_text_sha256(canonical_license)), + ), + distribution_files=( + FileDigest("LICENSE", "sha256", _canonical_text_sha256(distribution_license)), + ), + ) + family = ModelFamily( + id="toy", + architecture="Toy", + upstreams=("toy",), + tokenizer_mode="sequence", + public_input="Synthetic token IDs", + extra="core", + reference_container="reference-toy", + reference_adapter="tests.parity.support.reference_adapters.toy", + attention=("eager",), + dtypes=("float32",), + bf16_execution="static_parameters", + precisions=("default",), + vram_tier="sequence", + checkpoint_license="MIT", + hub_license="mit", + state_transform="identity", + representative="toy", + documentation="docs/toy.md", + test_tiers=("artifact",), + runtime_paths=("__init__.py", "models/toy"), + auto_map_items=( + ("AutoConfig", "fastplms.models.toy.modeling_toy.ToyConfig"), + ("AutoModel", "fastplms.models.toy.modeling_toy.ToyModel"), + ), + weights_publication_allowed=True, + conversion_provenance=( + "Input: synthetic official state. Transformation: identity. " + "Output: synthetic FastPLMs state. Validation: exact hash equality. " + "Limitation: synthetic test only." + ), + ) + spec = ModelSpec( + id="toy", + family=family, + fast=fast, + official=official, + size_category="small", + ) + registry = ModelRegistry( + schema_version=1, + upstreams={"toy": upstream}, + families={"toy": family}, + models={"toy": spec}, + legal_files=( + FileDigest("LICENSE", "sha256", _canonical_text_sha256(project_license)), + FileDigest( + "THIRD_PARTY_NOTICES.md", + "sha256", + _canonical_text_sha256(third_party_notices), + ), + ), + ) + return registry, spec + + +def _inject_checkpoint_race( + monkeypatch: pytest.MonkeyPatch, + target: Path, + *, + replacement: bytes, + in_place: bool = False, +) -> None: + """Mutate a pinned source after selection but before the builder copies it.""" + + from tools.artifacts import build as build_module + + original_copy = build_module._copy_file + fired = False + + def racing_copy(source: Path, destination: Path) -> None: + nonlocal fired + if not fired and source.resolve() == target.resolve(): + fired = True + if in_place: + target.write_bytes(replacement) + else: + staged = target.with_name(target.name + ".concurrent-replacement") + staged.write_bytes(replacement) + staged.replace(target) + original_copy(source, destination) + + monkeypatch.setattr(build_module, "_copy_file", racing_copy) + + +def test_artifact_build_rejects_concurrent_config_replacement( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + _inject_checkpoint_race( + monkeypatch, + checkpoint / "config.json", + replacement=b'{"model_type":"forged"}\n', + in_place=True, + ) + + with pytest.raises(ArtifactError, match="Preserved checkpoint bytes differ"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + +def test_official_tokenizer_copy_rejects_concurrent_replacement( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + snapshot = tmp_path / "snapshot" + destination = tmp_path / "artifact" + snapshot.mkdir() + tokenizer = snapshot / "tokenizer_config.json" + tokenizer.write_bytes(b'{"tokenizer_class":"Pinned"}\n') + source = CheckpointSource( + repo_id="upstream/tokenizer", + revision="a" * 40, + files=( + FileDigest( + tokenizer.name, + "git-sha1", + hash_file(tokenizer, "git-sha1"), + ), + ), + ) + _inject_checkpoint_race( + monkeypatch, + tokenizer, + replacement=b'{"tokenizer_class":"Forged"}\n', + ) + + with pytest.raises(ArtifactError, match="Preserved checkpoint bytes differ"): + _copy_official_tokenizer_assets(snapshot, destination, source) + + +def test_weight_snapshot_rejects_concurrent_replacement( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + # weights: (...) + weights = snapshot / "model.safetensors" + save_file({"weight": torch.arange(8, dtype=torch.float32)}, weights) + source = CheckpointSource( + repo_id="upstream/weights", + revision="b" * 40, + files=(FileDigest(weights.name, "sha256", hash_file(weights)),), + ) + _inject_checkpoint_race( + monkeypatch, + weights, + replacement=b"not-the-pinned-safetensors-bytes", + ) + + with pytest.raises(ArtifactError, match="Preserved checkpoint bytes differ"): + canonicalize_checkpoint_weights(snapshot, source, tmp_path / "artifact") + + +def _build_synthetic_for_replacement( + spec: ModelSpec, + registry: ModelRegistry, + checkpoint: Path, + output_root: Path, + source_root: Path, + *, + replace_existing: bool = False, +) -> Path: + return build_artifact( + spec, + registry, + checkpoint, + output_root, + source_root, + replace=replace_existing, + _allow_untracked_runtime_for_tests=True, + ) + + +def _change_synthetic_runtime(source_root: Path, marker: str) -> None: + runtime = source_root / "src" / "fastplms" / "models" / "toy" / "modeling_toy.py" + runtime.write_text( + runtime.read_text(encoding="utf-8") + f"\n# {marker}\n", + encoding="utf-8", + newline="\n", + ) + + +def test_artifact_replace_restores_prior_version_after_backup_rename_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + original_manifest = (artifact / "artifact-manifest.json").read_bytes() + _change_synthetic_runtime(source_root, "replacement generation") + + from tools.artifacts import build as build_module + + original_rename = build_module._atomic_artifact_rename + backup = output_root / ".ToyModel.backup" + fired = False + + def fail_after_backup_rename(source: Path, destination: Path, root: Path) -> None: + nonlocal fired + original_rename(source, destination, root) + if destination == backup and not fired: + fired = True + raise OSError("injected failure after old-to-backup rename") + + monkeypatch.setattr(build_module, "_atomic_artifact_rename", fail_after_backup_rename) + with pytest.raises(ArtifactError, match="prior artifact was restored"): + _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + + assert fired + assert (artifact / "artifact-manifest.json").read_bytes() == original_manifest + validate_artifact(artifact, spec=spec, registry=registry) + assert not (output_root / ".ToyModel.tmp").exists() + assert not backup.exists() + + +def test_artifact_replace_restores_prior_version_when_new_install_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + original_manifest = (artifact / "artifact-manifest.json").read_bytes() + _change_synthetic_runtime(source_root, "replacement generation") + + from tools.artifacts import build as build_module + + original_rename = build_module._atomic_artifact_rename + temporary = output_root / ".ToyModel.tmp" + fired = False + + def fail_new_install(source: Path, destination: Path, root: Path) -> None: + nonlocal fired + if source == temporary and destination == artifact and not fired: + fired = True + raise OSError("injected failure during temporary-to-destination rename") + original_rename(source, destination, root) + + monkeypatch.setattr(build_module, "_atomic_artifact_rename", fail_new_install) + with pytest.raises(ArtifactError, match="prior validated artifact was restored"): + _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + + assert fired + assert (artifact / "artifact-manifest.json").read_bytes() == original_manifest + validate_artifact(artifact, spec=spec, registry=registry) + assert not temporary.exists() + assert not (output_root / ".ToyModel.backup").exists() + + +def test_artifact_replace_recovers_backup_on_next_invocation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + original_manifest = (artifact / "artifact-manifest.json").read_bytes() + backup = output_root / ".ToyModel.backup" + temporary = output_root / ".ToyModel.tmp" + artifact.rename(backup) + temporary.mkdir() + (temporary / "partial-write").write_text("incomplete", encoding="utf-8") + + from tools.artifacts import build as build_module + + def stop_after_recovery(*args: object, **kwargs: object) -> None: + raise ArtifactError("injected build stop after transaction recovery") + + monkeypatch.setattr(build_module, "_copy_checkpoint_assets", stop_after_recovery) + with pytest.raises(ArtifactError, match="injected build stop after transaction recovery"): + _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + + assert artifact.is_dir() + assert (artifact / "artifact-manifest.json").read_bytes() == original_manifest + validate_artifact(artifact, spec=spec, registry=registry) + assert not temporary.exists() + assert not backup.exists() + + +def test_artifact_replace_success_cleans_transaction_slots(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + original_manifest = (artifact / "artifact-manifest.json").read_bytes() + _change_synthetic_runtime(source_root, "successful replacement") + + replaced = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + + assert replaced == artifact + assert (artifact / "artifact-manifest.json").read_bytes() != original_manifest + assert "successful replacement" in ( + artifact / "fastplms" / "models" / "toy" / "modeling_toy.py" + ).read_text(encoding="utf-8") + validate_artifact(artifact, spec=spec, registry=registry) + assert not (output_root / ".ToyModel.tmp").exists() + assert not (output_root / ".ToyModel.backup").exists() + + +def test_repeated_artifact_replace_does_not_retain_stale_files(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + + for generation in range(2): + stale = artifact / f"stale-generation-{generation}.bin" + stale.write_bytes(b"must not survive replacement") + _change_synthetic_runtime(source_root, f"replacement {generation}") + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + assert not stale.exists() + assert not any(artifact.glob("stale-generation-*.bin")) + validate_artifact(artifact, spec=spec, registry=registry) + + assert not (output_root / ".ToyModel.tmp").exists() + assert not (output_root / ".ToyModel.backup").exists() + + +@pytest.mark.parametrize("slot_suffix", (".tmp", ".backup")) +def test_artifact_replace_rejects_symlink_transaction_slots( + tmp_path: Path, + slot_suffix: str, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + output_root = tmp_path / "artifacts" + artifact = _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + outside = tmp_path / "outside" + outside.mkdir() + sentinel = outside / "sentinel" + sentinel.write_text("preserve", encoding="utf-8") + os.symlink(outside, output_root / f".ToyModel{slot_suffix}", target_is_directory=True) + + with pytest.raises(ArtifactError, match="must not be a symlink"): + _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + replace_existing=True, + ) + + assert sentinel.read_text(encoding="utf-8") == "preserve" + assert artifact.is_dir() + + +def test_artifact_build_rejects_symlink_output_root(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + outside = tmp_path / "outside-output" + outside.mkdir() + sentinel = outside / "sentinel" + sentinel.write_text("preserve", encoding="utf-8") + output_root = tmp_path / "artifact-link" + os.symlink(outside, output_root, target_is_directory=True) + + with pytest.raises(ArtifactError, match="output root must not be a symlink"): + _build_synthetic_for_replacement( + spec, + registry, + checkpoint, + output_root, + source_root, + ) + assert sentinel.read_text(encoding="utf-8") == "preserve" + assert not (outside / "ToyModel").exists() + + +def test_artifact_build_rejects_repository_name_path_escape(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + escaped_fast = replace(spec.fast, repo_id="Synthyra/../../escaped-artifact") + escaped_spec = replace(spec, fast=escaped_fast) + escaped_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={spec.id: escaped_spec}, + runtime_assets=registry.runtime_assets, + attention_kernels=registry.attention_kernels, + legal_files=registry.legal_files, + ) + output_root = tmp_path / "artifacts" + + with pytest.raises(ArtifactError, match="Invalid artifact repository name"): + _build_synthetic_for_replacement( + escaped_spec, + escaped_registry, + checkpoint, + output_root, + source_root, + ) + assert not (tmp_path / "escaped-artifact").exists() + + +def test_artifact_build_is_deterministic_and_self_verifying(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + + first = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "first", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + second = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "second", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + validate_artifact(first) + validate_artifact(second) + + first_manifest = json.loads((first / "artifact-manifest.json").read_text(encoding="utf-8")) + second_manifest = json.loads((second / "artifact-manifest.json").read_text(encoding="utf-8")) + assert first_manifest == second_manifest + config = json.loads((first / "config.json").read_text(encoding="utf-8")) + assert config["auto_map"] == { + "AutoConfig": "modeling_fastplms.ToyConfig", + "AutoModel": "modeling_fastplms.ToyModel", + } + assert config["fastplms_model_id"] == spec.id + assert config["fastplms_checkpoint_repo_id"] == spec.artifact_checkpoint.repo_id + assert config["fastplms_checkpoint_revision"] == spec.artifact_checkpoint.revision + assert config["fastplms_weights_revision"] == spec.artifact_checkpoint.revision + assert config["fastplms_runtime_revision"].startswith("source-tree-sha256:") + assert len(config["fastplms_source_tree_sha256"]) == 64 + assert len(config["fastplms_runtime_bundle_sha256"]) == 64 + assert config["fastplms_checkpoint_hash"] == _checkpoint_identity_hash( + spec.artifact_checkpoint + ) + assert (first / "fastplms" / "models" / "toy" / "modeling_toy.py").is_file() + assert (first / "fastplms_bundle.py").is_file() + assert (first / "requirements.txt").read_text(encoding="utf-8") == ( + "# Direct runtime dependencies for Synthyra/ToyModel.\n" + "# FastPLMs source is embedded in this model repository.\n" + "torch>=2.13,<2.14\n" + "transformers>=5.13,<5.14\n" + ) + bridge = (first / "modeling_fastplms.py").read_text(encoding="utf-8") + assert "from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH" in bridge + assert "from .fastplms." not in bridge + assert not (first / "vendor").exists() + assert (first / "LICENSES" / "toy" / "LICENSE").is_file() + assert (first / "THIRD_PARTY_NOTICES.md").is_file() + assert (first / "model.safetensors.index.json").is_file() + assert not (first / "model.safetensors").exists() + assert len(list(first.glob("model-*.safetensors"))) == 1 + provenance = json.loads((first / "provenance.json").read_text(encoding="utf-8")) + assert provenance["schema_version"] == 4 + assert provenance["generator"] == { + "name": "tools.artifacts.build", + "version": 4, + } + assert provenance["weights_license_status"] == "resolved" + assert provenance["redistributable"] is True + assert provenance["weights_revision"] == spec.artifact_checkpoint.revision + assert provenance["runtime_revision"] == config["fastplms_runtime_revision"] + assert provenance["source_tree_sha256"] == config["fastplms_source_tree_sha256"] + assert provenance["runtime_bundle_sha256"] == config["fastplms_runtime_bundle_sha256"] + assert provenance["release_tool_revision"] == config[ + "fastplms_release_tool_revision" + ] + assert provenance["release_tool_sha256"] == config["fastplms_release_tool_sha256"] + assert "" not in (first / "README.md").read_text(encoding="utf-8") + assert provenance["attestations"]["complete_artifact"]["scope"] == "weights+runtime" + runtime_attestation = json.loads( + (first / "runtime-attestation.json").read_text(encoding="utf-8") + ) + assert runtime_attestation["scope"] == "runtime-only" + assert runtime_attestation["weights_license_status"] == "resolved" + assert runtime_attestation["redistributable"] is True + assert runtime_attestation["weights"] == { + "repo_id": spec.fast.repo_id, + "revision": spec.fast.revision, + } + assert "provenance.json" not in runtime_attestation["files"] + assert "requirements.txt" in runtime_attestation["files"] + assert not any(_is_weight_file(path) for path in runtime_attestation["files"]) + assert provenance["bf16_execution"] == "static_parameters" + assert provenance["canonical_weights"]["source_schema"] == "canonical" + assert provenance["canonical_weights"]["state_transform"] == "identity" + assert provenance["hub_license_metadata"] == {"license": "mit"} + assert "`static_parameters`" in (first / "README.md").read_text(encoding="utf-8") + + second_config_path = second / "config.json" + second_config = json.loads(second_config_path.read_text(encoding="utf-8")) + second_config["fastplms_model_id"] = "wrong-model" + second_config_path.write_text( + json.dumps(second_config, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + second_manifest["config.json"] = f"sha256:{hash_file(second_config_path)}" + (second / "artifact-manifest.json").write_text( + json.dumps(second_manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + with pytest.raises(ArtifactError, match="packaging identity differs"): + validate_artifact(second) + + readme_path = first / "README.md" + manifest_path = first / "artifact-manifest.json" + original_readme = readme_path.read_bytes() + original_manifest = manifest_path.read_bytes() + tampered_readme = original_readme.decode("utf-8").replace( + 'license: "mit"', + 'license: "apache-2.0"', + 1, + ) + assert tampered_readme.encode("utf-8") != original_readme + readme_path.write_text(tampered_readme, encoding="utf-8", newline="\n") + tampered_manifest = json.loads(original_manifest) + tampered_manifest["README.md"] = f"sha256:{hash_file(readme_path)}" + manifest_path.write_text( + json.dumps(tampered_manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + with pytest.raises(ArtifactError, match="differs from provenance"): + validate_artifact(first) + readme_path.write_bytes(original_readme) + manifest_path.write_bytes(original_manifest) + + next(first.glob("model-*.safetensors")).write_bytes(b"tampered") + with pytest.raises(ArtifactError, match="digest mismatch"): + validate_artifact(first) + + +def test_complete_artifact_rejects_self_attested_card_runtime_placeholder( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + readme = artifact / "README.md" + readme.write_text( + readme.read_text(encoding="utf-8").replace( + "requirements.txt", + "requirements.txt@", + 1, + ), + encoding="utf-8", + newline="\n", + ) + runtime_attestation_path = artifact / "runtime-attestation.json" + runtime_attestation = json.loads(runtime_attestation_path.read_text(encoding="utf-8")) + runtime_attestation["files"]["README.md"] = f"sha256:{hash_file(readme)}" + runtime_attestation_path.write_text( + json.dumps(runtime_attestation, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["README.md"] = f"sha256:{hash_file(readme)}" + manifest["runtime-attestation.json"] = ( + f"sha256:{hash_file(runtime_attestation_path)}" + ) + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + + with pytest.raises(ArtifactError, match="unresolved runtime-revision placeholder"): + validate_artifact(artifact, spec=spec, registry=registry) + + +def test_generated_bridge_uses_private_verified_runtime(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + config = json.loads((artifact / "config.json").read_text(encoding="utf-8")) + runtime_hash = config["fastplms_runtime_bundle_sha256"] + poisoned_cache = tmp_path / f"_fastplms_runtime_{runtime_hash}" / "fastplms" + poisoned_cache.mkdir(parents=True) + (poisoned_cache / "__init__.py").write_text( + "raise RuntimeError('shared cache was trusted')\n", + encoding="utf-8", + ) + probe = tmp_path / "private_runtime_probe.py" + probe.write_text( + textwrap.dedent( + """\ + import importlib.util + import sys + import types + from pathlib import Path + + + artifact = Path(sys.argv[1]) + poisoned_cache = Path(sys.argv[2]).resolve() + package = types.ModuleType("artifact_private") + package.__package__ = "artifact_private" + package.__path__ = [str(artifact)] + sys.modules["artifact_private"] = package + module_name = "artifact_private.modeling_fastplms" + spec = importlib.util.spec_from_file_location( + module_name, + artifact / "modeling_fastplms.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load generated bridge") + bridge = importlib.util.module_from_spec(spec) + sys.modules[module_name] = bridge + spec.loader.exec_module(bridge) + + assert len(bridge._RUNTIME_TEMPORARIES) == 1 + private_root = Path(bridge._RUNTIME_TEMPORARIES[0].name).resolve() + package_root = private_root / "fastplms" + runtime = sys.modules["fastplms"] + assert runtime.__fastplms_artifact_runtime_temporaries__ == tuple( + bridge._RUNTIME_TEMPORARIES + ) + assert private_root.name.startswith("fastplms-artifact-runtime-") + assert poisoned_cache != package_root + assert poisoned_cache not in package_root.parents + assert not any(path.name == "__pycache__" for path in package_root.rglob("*")) + + before = bridge._runtime_file_hashes(package_root) + source = next(path for path in package_root.rglob("*.py") if path.is_file()) + relative = source.relative_to(package_root).as_posix() + original = source.read_bytes() + source.write_bytes(original + b"\\n# in-place mutation\\n") + after = bridge._runtime_file_hashes(package_root) + assert after[relative] != before[relative] + source.write_bytes(original) + + bytecode = package_root / "injected.pyc" + bytecode.write_bytes(b"not trusted bytecode") + try: + bridge._runtime_file_hashes(package_root) + except RuntimeError as error: + assert "contains bytecode" in str(error) + else: + raise AssertionError("Private runtime bytecode was accepted") + bytecode.unlink() + + link_probe = package_root / "pretend_symlink.py" + link_probe.write_text("# symlink stand-in\\n", encoding="utf-8") + original_is_symlink = Path.is_symlink + Path.is_symlink = lambda path: path == link_probe or original_is_symlink(path) + try: + try: + bridge._runtime_file_hashes(package_root) + except RuntimeError as error: + assert "contains a symlink" in str(error) + else: + raise AssertionError("Private runtime symlink was accepted") + finally: + Path.is_symlink = original_is_symlink + """ + ), + encoding="utf-8", + newline="\n", + ) + environment = dict(os.environ) + environment.update( + { + "TMPDIR": str(tmp_path), + "TMP": str(tmp_path), + "TEMP": str(tmp_path), + } + ) + completed = subprocess.run( + [sys.executable, "-I", "-S", str(probe), str(artifact), str(poisoned_cache)], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.parametrize( + "relative_name", + ( + "../outside.txt", + "/absolute.txt", + "C:/absolute.txt", + "nested//non-normalized.txt", + r"nested\windows-path.txt", + "NUL", + "nested/con.py", + "nested/bad:name.py", + "nested/trailing.", + ), +) +def test_artifact_validation_rejects_unsafe_manifest_paths( + tmp_path: Path, + relative_name: str, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest[relative_name] = "sha256:" + "0" * 64 + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + + with pytest.raises(ArtifactError, match="invalid artifact manifest path"): + validate_artifact(artifact) + + +def test_different_runtime_bundles_fail_without_replacing_loaded_runtime( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + + first = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "first", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + runtime_source = source_root / "src" / "fastplms" / "models" / "toy" / "modeling_toy.py" + runtime_source.write_text( + runtime_source.read_text(encoding="utf-8") + "\n# Distinct runtime identity.\n", + encoding="utf-8", + newline="\n", + ) + second = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "second", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + probe = tmp_path / "mixed_runtime_probe.py" + probe.write_text( + textwrap.dedent( + """\ + import importlib.util + import sys + import types + from pathlib import Path + + + def load_bridge(root, package_name): + package = types.ModuleType(package_name) + package.__package__ = package_name + package.__path__ = [str(root)] + sys.modules[package_name] = package + module_name = f"{package_name}.modeling_fastplms" + spec = importlib.util.spec_from_file_location( + module_name, + root / "modeling_fastplms.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load {root}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + + first = load_bridge(Path(sys.argv[1]), "artifact_first") + runtime = sys.modules["fastplms"] + runtime_hash = runtime.__fastplms_artifact_runtime_hash__ + try: + load_bridge(Path(sys.argv[2]), "artifact_second") + except RuntimeError as error: + if "incompatible runtime sources" not in str(error): + raise + else: + raise AssertionError("A different runtime bundle loaded silently") + assert sys.modules["fastplms"] is runtime + assert runtime.__fastplms_artifact_runtime_hash__ == runtime_hash + assert first.ToyConfig().__class__ is first.ToyConfig + """ + ), + encoding="utf-8", + newline="\n", + ) + completed = subprocess.run( + [sys.executable, "-I", "-S", str(probe), str(first), str(second)], + cwd=tmp_path, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +def test_complementary_fastplms_artifacts_load_in_one_process(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, first_spec = _synthetic_registry(source_root, checkpoint) + + second_module = source_root / "src" / "fastplms" / "models" / "toy_second" + second_module.mkdir(parents=True) + (second_module / "modeling_toy_second.py").write_text( + "class ToySecondConfig: pass\nclass ToySecondModel: pass\n", + encoding="utf-8", + ) + second_family = replace( + first_spec.family, + id="toy_second", + runtime_paths=("__init__.py", "models/toy_second"), + auto_map_items=( + ( + "AutoConfig", + "fastplms.models.toy_second.modeling_toy_second.ToySecondConfig", + ), + ( + "AutoModel", + "fastplms.models.toy_second.modeling_toy_second.ToySecondModel", + ), + ), + ) + second_spec = replace( + first_spec, + id="toy_second", + family=second_family, + fast=replace(first_spec.fast, repo_id="Synthyra/ToyModelSecond"), + ) + registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={ + first_spec.family.id: first_spec.family, + second_family.id: second_family, + }, + models={ + first_spec.id: first_spec, + second_spec.id: second_spec, + }, + runtime_assets=registry.runtime_assets, + attention_kernels=registry.attention_kernels, + legal_files=registry.legal_files, + ) + first = build_artifact( + first_spec, + registry, + checkpoint, + tmp_path / "first", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + second = build_artifact( + second_spec, + registry, + checkpoint, + tmp_path / "second", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + probe = tmp_path / "compatible_runtime_probe.py" + probe.write_text( + textwrap.dedent( + """\ + import importlib.util + import sys + import types + from pathlib import Path + + + def load_bridge(root, package_name): + package = types.ModuleType(package_name) + package.__package__ = package_name + package.__path__ = [str(root)] + sys.modules[package_name] = package + module_name = f"{package_name}.modeling_fastplms" + spec = importlib.util.spec_from_file_location( + module_name, + root / "modeling_fastplms.py", + ) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + + first = load_bridge(Path(sys.argv[1]), "artifact_first") + runtime = sys.modules["fastplms"] + second = load_bridge(Path(sys.argv[2]), "artifact_second") + assert sys.modules["fastplms"] is runtime + assert len(runtime.__fastplms_artifact_runtime_hashes__) == 2 + assert first.ToyConfig().__class__ is first.ToyConfig + assert second.ToySecondConfig().__class__ is second.ToySecondConfig + """ + ), + encoding="utf-8", + newline="\n", + ) + completed = subprocess.run( + [sys.executable, "-I", "-S", str(probe), str(first), str(second)], + cwd=tmp_path, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +def test_preloaded_non_artifact_fastplms_runtime_is_rejected(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + source_package = source_root / "src" / "fastplms" + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + probe = tmp_path / "external_runtime_probe.py" + probe.write_text( + textwrap.dedent( + """\ + import importlib.util + import sys + import types + from pathlib import Path + + source_package = Path(sys.argv[1]) + artifact = Path(sys.argv[2]) + external_spec = importlib.util.spec_from_file_location( + "fastplms", + source_package / "__init__.py", + submodule_search_locations=[str(source_package)], + ) + external = importlib.util.module_from_spec(external_spec) + sys.modules["fastplms"] = external + external_spec.loader.exec_module(external) + + artifact_package = types.ModuleType("artifact") + artifact_package.__package__ = "artifact" + artifact_package.__path__ = [str(artifact)] + sys.modules["artifact"] = artifact_package + bridge_spec = importlib.util.spec_from_file_location( + "artifact.modeling_fastplms", + artifact / "modeling_fastplms.py", + ) + bridge = importlib.util.module_from_spec(bridge_spec) + sys.modules["artifact.modeling_fastplms"] = bridge + try: + bridge_spec.loader.exec_module(bridge) + except RuntimeError as error: + if "non-artifact fastplms module" not in str(error): + raise + else: + raise AssertionError("External FastPLMs source loaded into an artifact runtime") + """ + ), + encoding="utf-8", + newline="\n", + ) + completed = subprocess.run( + [ + sys.executable, + "-I", + "-S", + str(probe), + str(source_root / "src" / "fastplms"), + str(artifact), + ], + cwd=tmp_path, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + + +def test_legal_texts_use_canonical_lf_across_checkouts(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + legal_paths = ( + source_root / "LICENSE", + source_root / "THIRD_PARTY_NOTICES.md", + source_root / "vendor" / "upstream" / "toy" / "LICENSE", + source_root / "LICENSES" / "toy" / "LICENSE", + ) + for path in legal_paths: + raw = path.read_bytes().replace(b"\r\n", b"\n").replace(b"\r", b"\n") + path.write_bytes(raw.replace(b"\n", b"\r\n")) + + validate_repository_legal_inventory(source_root, registry) + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + distributed = ( + artifact / "LICENSES" / "toy" / "LICENSE", + artifact / "LICENSES" / "FastPLMs-Apache-2.0.txt", + artifact / "THIRD_PARTY_NOTICES.md", + ) + for path in distributed: + assert b"\r" not in path.read_bytes() + + canonical = legal_paths[2] + canonical.write_text("Changed license content\n", encoding="utf-8") + with pytest.raises(ArtifactError, match="canonical LF normalization"): + validate_repository_legal_inventory(source_root, registry) + + +def test_manifest_distributes_required_modified_file_notices() -> None: + registry = load_model_registry() + distribution = { + source_id: {item.path for item in source.distribution_files} + for source_id, source in registry.upstreams.items() + } + assert {"Apache-2.0.txt", "BSD-3-Clause.txt", "MODIFICATIONS.md"}.issubset(distribution["e1"]) + assert {"LICENSE", "PROVENANCE.md"}.issubset(distribution["dplm"]) + assert {"LICENSE", "MODIFICATIONS.md", "PROVENANCE.md"}.issubset(distribution["openfold"]) + + +def test_artifact_rejects_stale_checked_in_hub_license_metadata(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + card_path = source_root / "model_cards" / "toy.md" + card_path.parent.mkdir() + card_path.write_text( + '---\nlibrary_name: transformers\nlicense: "apache-2.0"\n---\n\n# Toy\n', + encoding="utf-8", + ) + + with pytest.raises(ArtifactError, match=r"license metadata differs from models\.toml"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + +def test_artifact_copies_official_tokenizer_bytes_exactly(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + official_snapshot = tmp_path / "official" + checkpoint.mkdir() + official_snapshot.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + + candidate_tokenizer = checkpoint / "tokenizer.json" + official_tokenizer = official_snapshot / "tokenizer.json" + candidate_tokenizer_config = checkpoint / "tokenizer_config.json" + official_tokenizer_config = official_snapshot / "tokenizer_config.json" + candidate_tokenizer.write_bytes(b'{"source":"candidate"}\n') + official_tokenizer.write_bytes(b'{"source":"official"}\n') + candidate_tokenizer_config.write_bytes(b'{"tokenizer_class":"CandidateTokenizer"}\n') + official_tokenizer_config.write_bytes(b'{"tokenizer_class":"BuiltInTokenizer"}\n') + fast = replace( + spec.fast, + files=( + *spec.fast.files, + FileDigest("tokenizer.json", "sha256", hash_file(candidate_tokenizer)), + FileDigest( + "tokenizer_config.json", + "sha256", + hash_file(candidate_tokenizer_config), + ), + ), + ) + official = replace( + spec.official, + files=( + *spec.official.files, + FileDigest("tokenizer.json", "sha256", hash_file(official_tokenizer)), + FileDigest( + "tokenizer_config.json", + "sha256", + hash_file(official_tokenizer_config), + ), + ), + ) + family = replace(spec.family, tokenizer_mode="tokenizer") + tokenizer_spec = replace(spec, family=family, fast=fast, official=official) + tokenizer_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={family.id: family}, + models={tokenizer_spec.id: tokenizer_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + tokenizer_spec, + tokenizer_registry, + checkpoint, + tmp_path / "artifact", + source_root, + tokenizer_dir=official_snapshot, + _allow_untracked_runtime_for_tests=True, + ) + assert (artifact / "tokenizer.json").read_bytes() == official_tokenizer.read_bytes() + assert ( + artifact / "tokenizer_config.json" + ).read_bytes() == official_tokenizer_config.read_bytes() + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + assert provenance["tokenizer_checkpoint"]["repo_id"] == official.repo_id + assert provenance["tokenizer_checkpoint"]["revision"] == official.revision + assert provenance["tokenizer_checkpoint"]["files"] == { + "tokenizer.json": official.files[-2].encoded, + "tokenizer_config.json": official.files[-1].encoded, + } + + +def test_artifact_rewrites_custom_tokenizer_auto_map_to_local_bridge(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + official_snapshot = tmp_path / "official" + checkpoint.mkdir() + official_snapshot.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + + runtime_path = source_root / "src" / "fastplms" / "models" / "toy" / "modeling_toy.py" + runtime_path.write_text( + runtime_path.read_text(encoding="utf-8") + "\nclass ToyTokenizer: pass\n", + encoding="utf-8", + newline="\n", + ) + candidate_tokenizer = checkpoint / "tokenizer.json" + official_tokenizer = official_snapshot / "tokenizer.json" + official_tokenizer_config = official_snapshot / "tokenizer_config.json" + candidate_tokenizer.write_text('{"source":"candidate"}\n', encoding="utf-8") + official_tokenizer.write_text('{"source":"official"}\n', encoding="utf-8") + official_tokenizer_config.write_text( + json.dumps( + { + "auto_map": { + "AutoProcessor": "upstream_processing.UpstreamProcessor", + "AutoTokenizer": [ + "upstream_tokenization.UpstreamTokenizer", + None, + ], + }, + "preserved": True, + "tokenizer_class": "UpstreamTokenizer", + } + ) + + "\n", + encoding="utf-8", + ) + fast = replace( + spec.fast, + files=( + *spec.fast.files, + FileDigest("tokenizer.json", "sha256", hash_file(candidate_tokenizer)), + ), + ) + official = replace( + spec.official, + files=( + *spec.official.files, + FileDigest("tokenizer.json", "sha256", hash_file(official_tokenizer)), + FileDigest( + "tokenizer_config.json", + "sha256", + hash_file(official_tokenizer_config), + ), + ), + ) + family = replace( + spec.family, + tokenizer_mode="tokenizer", + tokenizer_class="fastplms.models.toy.modeling_toy.ToyTokenizer", + ) + tokenizer_spec = replace(spec, family=family, fast=fast, official=official) + tokenizer_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={family.id: family}, + models={tokenizer_spec.id: tokenizer_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + tokenizer_spec, + tokenizer_registry, + checkpoint, + tmp_path / "artifact", + source_root, + tokenizer_dir=official_snapshot, + _allow_untracked_runtime_for_tests=True, + ) + + tokenizer_config = json.loads( + (artifact / "tokenizer_config.json").read_text(encoding="utf-8") + ) + assert tokenizer_config["auto_map"] == { + "AutoProcessor": "upstream_processing.UpstreamProcessor", + "AutoTokenizer": ["modeling_fastplms.ToyTokenizer", None], + } + assert tokenizer_config["preserved"] is True + assert "ToyTokenizer =" in (artifact / "modeling_fastplms.py").read_text(encoding="utf-8") + validate_artifact(artifact) + + +def test_esm3_uses_the_manifest_pinned_official_esmc_tokenizer() -> None: + registry = load_model_registry() + spec = registry["esm3_small"] + tokenizer_checkpoint = _tokenizer_checkpoint(registry, spec) + + assert spec.tokenizer_source_id == "esmc_small" + assert tokenizer_checkpoint == registry["esmc_small"].official + assert tokenizer_checkpoint.repo_id == "biohub/ESMC-300M" + assert {Path(item.path).name for item in tokenizer_checkpoint.files}.issuperset( + {"special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"} + ) + + +def test_fast_checkpoint_artifact_requires_explicit_official_tokenizer_snapshot( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + candidate_tokenizer = checkpoint / "tokenizer.json" + candidate_tokenizer.write_bytes(b'{"source":"candidate"}\n') + fast = replace( + spec.fast, + files=( + *spec.fast.files, + FileDigest("tokenizer.json", "sha256", hash_file(candidate_tokenizer)), + ), + ) + official = replace( + spec.official, + files=( + *spec.official.files, + FileDigest("tokenizer.json", "sha256", "0" * 64), + ), + ) + family = replace(spec.family, tokenizer_mode="tokenizer") + tokenizer_spec = replace(spec, family=family, fast=fast, official=official) + tokenizer_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={family.id: family}, + models={tokenizer_spec.id: tokenizer_spec}, + legal_files=registry.legal_files, + ) + + with pytest.raises(ArtifactError, match="requires the pinned official tokenizer snapshot"): + build_artifact( + tokenizer_spec, + tokenizer_registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_checkpoint_verification_reports_hash_mismatch(tmp_path: Path) -> None: + weight = tmp_path / "model.safetensors" + weight.write_bytes(b"content") + source = CheckpointSource( + repo_id="Synthyra/ToyModel", + revision="1" * 40, + files=(FileDigest("model.safetensors", "sha256", "0" * 64),), + ) + with pytest.raises(ArtifactError, match="Checkpoint verification failed"): + verify_checkpoint(tmp_path, source) + + +def test_artifact_build_rejects_unresolved_provenance(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + unresolved_fast = replace(spec.fast, unresolved_files=("tokenizer.json",)) + unresolved_spec = replace(spec, fast=unresolved_fast) + unresolved_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={spec.id: unresolved_spec}, + legal_files=registry.legal_files, + ) + + with pytest.raises(ArtifactError, match="Release provenance is unresolved"): + build_artifact( + unresolved_spec, + unresolved_registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_rejects_missing_legal_inventory(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + (source_root / "LICENSES" / "toy" / "LICENSE").unlink() + + with pytest.raises(ArtifactError, match="Missing required toy distribution legal file"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_rejects_missing_conversion_record(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + family = replace(spec.family, conversion_provenance="") + invalid_spec = replace(spec, family=family) + invalid_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={family.id: family}, + models={invalid_spec.id: invalid_spec}, + legal_files=registry.legal_files, + ) + + with pytest.raises(ArtifactError, match="missing conversion provenance"): + build_artifact( + invalid_spec, + invalid_registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_stamps_unresolved_checkpoint_as_nonredistributable( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + family = replace(spec.family, weights_publication_allowed=False) + invalid_spec = replace(spec, family=family) + invalid_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={family.id: family}, + models={invalid_spec.id: invalid_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + invalid_spec, + invalid_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + attestation = json.loads( + (artifact / "runtime-attestation.json").read_text(encoding="utf-8") + ) + assert provenance["weights_license_status"] == "unresolved" + assert provenance["redistributable"] is False + assert attestation["weights_license_status"] == "unresolved" + assert attestation["redistributable"] is False + + +@pytest.mark.parametrize("relative_name", ("credentials.pem", "secrets.py")) +def test_artifact_build_rejects_unknown_runtime_source_extension( + tmp_path: Path, + relative_name: str, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + secret = source_root / "src" / "fastplms" / "models" / "toy" / relative_name + secret.write_text("private material", encoding="utf-8") + + with pytest.raises(ArtifactError, match="sensitive path"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_rejects_untracked_runtime_source(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + subprocess.run(["git", "init", "--initial-branch=main"], cwd=source_root, check=True) + subprocess.run( + ["git", "config", "user.email", "tests@example.invalid"], + cwd=source_root, + check=True, + ) + subprocess.run(["git", "config", "user.name", "FastPLMs Tests"], cwd=source_root, check=True) + subprocess.run(["git", "add", "src/fastplms"], cwd=source_root, check=True) + subprocess.run(["git", "commit", "-m", "runtime fixture"], cwd=source_root, check=True) + untracked = source_root / "src" / "fastplms" / "models" / "toy" / "injected.py" + untracked.write_text("TOKEN = 'unsafe'\n", encoding="utf-8") + + with pytest.raises(ArtifactError, match="tracked and clean"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_rejects_runtime_source_without_git_provenance( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + injected = source_root / "src" / "fastplms" / "models" / "toy" / "injected.py" + injected.write_text("APPROVED_EXTENSION_BUT_UNTRACKED = True\n", encoding="utf-8") + + with pytest.raises(ArtifactError, match="verifiable Git worktree"): + build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + ) + + +def test_artifact_build_retains_validated_names_when_worktree_file_is_deleted( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + subprocess.run(["git", "init", "--initial-branch=main"], cwd=source_root, check=True) + subprocess.run( + ["git", "config", "user.email", "tests@example.invalid"], + cwd=source_root, + check=True, + ) + subprocess.run( + ["git", "config", "user.name", "FastPLMs Tests"], + cwd=source_root, + check=True, + ) + subprocess.run(["git", "add", "src/fastplms"], cwd=source_root, check=True) + subprocess.run(["git", "commit", "-m", "runtime fixture"], cwd=source_root, check=True) + runtime_source = ( + source_root / "src" / "fastplms" / "models" / "toy" / "modeling_toy.py" + ) + committed = runtime_source.read_bytes() + + def delete_after_validation(*args: object, **kwargs: object) -> str | None: + revision = _git_runtime_revision(*args, **kwargs) # type: ignore[arg-type] + runtime_source.unlink() + return revision + + monkeypatch.setattr( + "tools.artifacts.build._git_runtime_revision", + delete_after_validation, + ) + + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + assert (artifact / "fastplms" / "models" / "toy" / "modeling_toy.py").read_bytes() == ( + committed + ) + + +def test_artifact_build_uses_validated_git_blobs_after_worktree_mutation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + subprocess.run(["git", "init", "--initial-branch=main"], cwd=source_root, check=True) + subprocess.run( + ["git", "config", "user.email", "tests@example.invalid"], + cwd=source_root, + check=True, + ) + subprocess.run(["git", "config", "user.name", "FastPLMs Tests"], cwd=source_root, check=True) + subprocess.run(["git", "add", "src/fastplms"], cwd=source_root, check=True) + subprocess.run(["git", "commit", "-m", "runtime fixture"], cwd=source_root, check=True) + runtime_source = ( + source_root / "src" / "fastplms" / "models" / "toy" / "modeling_toy.py" + ) + committed = runtime_source.read_bytes() + + def mutate_during_checkpoint_build(*args: object, **kwargs: object) -> object: + runtime_source.write_bytes(b"MUTATED_DURING_BUILD = True\n") + try: + return canonicalize_checkpoint_weights(*args, **kwargs) + finally: + runtime_source.write_bytes(committed) + + monkeypatch.setattr( + "tools.artifacts.build.canonicalize_checkpoint_weights", + mutate_during_checkpoint_build, + ) + + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + assert (artifact / "fastplms" / "models" / "toy" / "modeling_toy.py").read_bytes() == ( + committed + ) + + +@pytest.mark.parametrize( + ("field", "forged"), + ( + ("checkpoint_license", "Forged-License"), + ("legal_files", {}), + ("upstreams", []), + ), +) +def test_artifact_validation_rejects_self_attested_forged_legal_provenance( + tmp_path: Path, + field: str, + forged: object, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + artifact = build_artifact( + spec, + registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + provenance_path = artifact / "provenance.json" + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + provenance[field] = forged + provenance_path.write_text( + json.dumps(provenance, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["provenance.json"] = f"sha256:{hash_file(provenance_path)}" + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + + with pytest.raises(ArtifactError, match="differs from the current registry"): + validate_artifact(artifact, spec=spec, registry=registry) + + +def test_dplm2_artifact_materializes_non_decoder_cache_config(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + source_config_path = checkpoint / "config.json" + source_config_path.write_text( + json.dumps( + { + "model_type": "toy", + "is_decoder": True, + "add_cross_attention": True, + "use_cache": True, + } + ) + + "\n", + encoding="utf-8", + newline="\n", + ) + source_files = tuple( + FileDigest(item.path, item.algorithm, hash_file(source_config_path, item.algorithm)) + if item.path == "config.json" + else item + for item in spec.fast.files + ) + source_config_digest = next(item for item in source_files if item.path == "config.json") + official = replace( + spec.official, + files=source_files, + repo_id="upstream/DPLM2Official", + revision="5" * 40, + ) + dplm2_family = replace(spec.family, id="dplm2", architecture="DPLM2") + selected_spec = replace( + spec, + family=dplm2_family, + official=official, + artifact_source="official", + canonical_state_sha256=_canonical_state_sha256( + load_file(checkpoint / "model.safetensors") + ), + ) + selected_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families={dplm2_family.id: dplm2_family}, + models={selected_spec.id: selected_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + selected_spec, + selected_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + raw_config = json.loads((artifact / "config.json").read_text(encoding="utf-8")) + assert raw_config["is_decoder"] is False + assert raw_config["add_cross_attention"] is False + assert raw_config["use_cache"] is False + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + assert provenance["artifact_checkpoint"]["files"]["config.json"] == ( + source_config_digest.encoded + ) + assert provenance["official_checkpoint"]["files"]["config.json"] == ( + source_config_digest.encoded + ) + assert hash_file(artifact / "config.json", source_config_digest.algorithm) != ( + source_config_digest.digest + ) + + +def test_non_dplm2_artifact_preserves_source_cache_fields(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + source_config_path = checkpoint / "config.json" + source_values = { + "model_type": "toy", + "is_decoder": True, + "add_cross_attention": True, + "use_cache": True, + } + source_config_path.write_text( + json.dumps(source_values) + "\n", + encoding="utf-8", + newline="\n", + ) + fast = replace( + spec.fast, + files=tuple( + FileDigest(item.path, item.algorithm, hash_file(source_config_path, item.algorithm)) + if item.path == "config.json" + else item + for item in spec.fast.files + ), + ) + selected_spec = replace(spec, fast=fast) + selected_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={selected_spec.id: selected_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + selected_spec, + selected_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + raw_config = json.loads((artifact / "config.json").read_text(encoding="utf-8")) + assert {key: raw_config[key] for key in source_values} == source_values + + +def test_artifact_uses_manifest_selected_official_checkpoint(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + official = replace( + spec.official, + files=spec.fast.files, + repo_id="upstream/ToyOfficial", + revision="5" * 40, + ) + selected_spec = replace( + spec, + official=official, + artifact_source="official", + canonical_state_sha256=_canonical_state_sha256( + load_file(checkpoint / "model.safetensors") + ), + ) + selected_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={selected_spec.id: selected_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + selected_spec, + selected_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + assert provenance["artifact_source"] == "official" + assert provenance["artifact_checkpoint"]["repo_id"] == "upstream/ToyOfficial" + assert provenance["artifact_checkpoint"]["revision"] == "5" * 40 + assert provenance["canonical_weights"]["source_schema"] == "official" + assert provenance["canonical_weights"]["state_digest"] == ( + provenance["conversion_equality_attestation"]["canonical_state"] + ) + + +def test_artifact_rejects_self_attested_forged_canonical_weight( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + official = replace( + spec.official, + files=spec.fast.files, + repo_id="upstream/ToyOfficial", + revision="5" * 40, + ) + selected_spec = replace( + spec, + official=official, + artifact_source="official", + canonical_state_sha256=_canonical_state_sha256( + load_file(checkpoint / "model.safetensors") + ), + ) + selected_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={selected_spec.id: selected_spec}, + legal_files=registry.legal_files, + ) + artifact = build_artifact( + selected_spec, + selected_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + + shard = next(artifact.glob("model-*.safetensors")) + forged_state = load_file(shard) + # forged_state['linear.bias']: (...) + forged_state["linear.bias"] = forged_state["linear.bias"].clone() + forged_state["linear.bias"][0] += 1 + save_file(forged_state, shard, metadata={"format": "pt"}) + forged_state_sha256 = _canonical_state_sha256(forged_state) + + provenance_path = artifact / "provenance.json" + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + provenance["canonical_weights"]["shards"][shard.name] = ( + f"sha256:{hash_file(shard)}" + ) + provenance["canonical_weights"]["state_digest"]["sha256"] = forged_state_sha256 + forged_spec = replace(selected_spec, canonical_state_sha256=forged_state_sha256) + provenance["conversion_equality_attestation"] = _conversion_equality_attestation( + forged_spec + ) + provenance_path.write_text( + json.dumps(provenance, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest[shard.name] = f"sha256:{hash_file(shard)}" + manifest["provenance.json"] = f"sha256:{hash_file(provenance_path)}" + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + with pytest.raises( + ArtifactError, + match="conversion equality attestation differs from the current registry", + ): + validate_artifact( + artifact, + spec=selected_spec, + registry=selected_registry, + ) + + +def test_hash_pinned_bin_is_canonicalized_with_safe_loading(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + (checkpoint / "model.safetensors").unlink() + bin_path = checkpoint / "pytorch_model.bin" + torch.save( + { + "linear.bias": torch.arange(4, dtype=torch.float32), + "linear.weight": torch.arange(16, dtype=torch.float32).reshape(4, 4), + }, + bin_path, + ) + config_digest = spec.fast.file_map["config.json"] + bin_source = replace( + spec.fast, + files=( + config_digest, + FileDigest("pytorch_model.bin", "sha256", hash_file(bin_path)), + ), + ) + bin_spec = replace(spec, fast=bin_source) + bin_registry = ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={bin_spec.id: bin_spec}, + legal_files=registry.legal_files, + ) + + artifact = build_artifact( + bin_spec, + bin_registry, + checkpoint, + tmp_path / "artifact", + source_root, + _allow_untracked_runtime_for_tests=True, + ) + validate_weight_artifact(artifact) + assert not (artifact / "pytorch_model.bin").exists() + assert list(artifact.glob("model-*.safetensors")) + + +def test_canonical_weight_sharding_and_index_validation(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + output = tmp_path / "output" + checkpoint.mkdir() + weight = checkpoint / "model.safetensors" + save_file( + {f"layer.{index}.weight": torch.arange(40, dtype=torch.float32) for index in range(4)}, + weight, + metadata={"format": "pt"}, + ) + source = CheckpointSource( + repo_id="Synthyra/ShardedToy", + revision="6" * 40, + files=(FileDigest("model.safetensors", "sha256", hash_file(weight)),), + ) + + record = canonicalize_checkpoint_weights( + checkpoint, + source, + output, + max_shard_bytes=512, + ) + index = validate_weight_artifact(output, max_shard_bytes=512) + assert len(record["shards"]) == 2 + assert len(set(index["weight_map"].values())) == 2 + assert all(path.stat().st_size <= 512 for path in output.glob("*.safetensors")) + + index["weight_map"].pop("layer.0.weight") + (output / "model.safetensors.index.json").write_text(json.dumps(index), encoding="utf-8") + with pytest.raises(ArtifactError, match="keys differ from the weight index"): + validate_weight_artifact(output, max_shard_bytes=512) + + +def test_canonicalization_applies_declared_esm2_transform_before_sharding( + tmp_path: Path, +) -> None: + checkpoint = tmp_path / "checkpoint" + output = tmp_path / "output" + checkpoint.mkdir() + weight = checkpoint / "model.safetensors" + source_state = { + "embed_tokens.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "layers.0.self_attn.q_proj.weight": torch.tensor([[1.0, 2.0]]), + "lm_head.weight": torch.tensor([[3.0, 4.0]]), + "lm_head.bias": torch.tensor([5.0, 6.0]), + } + save_file(source_state, weight, metadata={"format": "pt"}) + source = CheckpointSource( + repo_id="facebook/esm2-synthetic", + revision="7" * 40, + files=(FileDigest("model.safetensors", "sha256", hash_file(weight)),), + ) + + record = canonicalize_checkpoint_weights( + checkpoint, + source, + output, + state_transform="esm2_hf_to_fastplms_v1", + ) + converted: dict[str, torch.Tensor] = {} + for shard in sorted(output.glob("model-*.safetensors")): + converted.update(load_file(shard, device="cpu")) + + assert set(converted) == { + "esm.embeddings.word_embeddings.weight", + "esm.encoder.layer.0.attention.self.query.weight", + "lm_head.bias", + "lm_head.decoder.weight", + } + assert torch.equal(converted["lm_head.bias"], source_state["lm_head.bias"]) + assert record["state_transform"] == "esm2_hf_to_fastplms_v1" + + +def test_canonical_esmfold_artifact_drops_only_declared_unused_state( + tmp_path: Path, +) -> None: + checkpoint = tmp_path / "checkpoint" + output = tmp_path / "output" + checkpoint.mkdir() + weight = checkpoint / "model.safetensors" + source_state = { + "esm.encoder.layer.0.weight": torch.tensor([1.0]), + "esm.contact_head.regression.bias": torch.tensor([2.0]), + "mlm_head.bias": torch.tensor([3.0]), + "positional_encoding._float_tensor": torch.tensor([4.0]), + } + save_file(source_state, weight, metadata={"format": "pt"}) + source = CheckpointSource( + repo_id="Synthyra/FastESMFold-synthetic", + revision="8" * 40, + files=(FileDigest("model.safetensors", "sha256", hash_file(weight)),), + ) + + canonicalize_checkpoint_weights( + checkpoint, + source, + output, + state_transform="esmfold_meta_to_fastplms_v1", + source_is_canonical=True, + ) + converted: dict[str, torch.Tensor] = {} + for shard in sorted(output.glob("model-*.safetensors")): + converted.update(load_file(shard, device="cpu")) + + assert set(converted) == {"esm.encoder.layer.0.weight"} + + +def test_official_submodule_worktrees_match_manifest_revisions() -> None: + registry = load_model_registry() + parser = configparser.ConfigParser(interpolation=None) + assert parser.read(ROOT / ".gitmodules", encoding="utf-8") + declared = { + parser.get(section, "path"): parser.get(section, "url") for section in parser.sections() + } + expected = {source.path: source.url for source in registry.upstreams.values()} + assert declared == expected + + # The portable remote runner deliberately strips every .git entry from its + # source archive. Archive validation can still require the manifest-selected + # source directories and exact .gitmodules declarations; a full checkout + # additionally verifies the Git-link objects and worktree revisions below. + if not (ROOT / ".git").exists(): + for source in registry.upstreams.values(): + checkout = ROOT / source.path + assert checkout.is_dir() + assert not (checkout / ".git").exists() + return + + for source in registry.upstreams.values(): + checkout = ROOT / source.path + gitlink = subprocess.run( + [ + "git", + "-c", + f"safe.directory={ROOT.as_posix()}", + "-C", + str(ROOT), + "ls-files", + "--stage", + "--", + source.path, + ], + check=False, + capture_output=True, + text=True, + ) + assert gitlink.returncode == 0, gitlink.stderr + mode, revision, stage_and_path = gitlink.stdout.strip().split(maxsplit=2) + assert mode == "160000" + assert revision == source.revision + assert stage_and_path == f"0\t{source.path}" + result = subprocess.run( + [ + "git", + "-c", + f"safe.directory={checkout.as_posix()}", + "-C", + str(checkout), + "rev-parse", + "HEAD", + ], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == source.revision + + +def test_artifact_build_rejects_dirty_official_source(tmp_path: Path) -> None: + """A matching HEAD is insufficient when tracked oracle bytes were modified.""" + + source_root = tmp_path / "source" + checkout = source_root / "vendor" / "upstream" / "toy" + checkout.mkdir(parents=True) + subprocess.run(["git", "init", "--initial-branch=main"], cwd=source_root, check=True) + subprocess.run(["git", "init", "--initial-branch=main"], cwd=checkout, check=True) + tracked = checkout / "oracle.py" + tracked.write_text("scale = 1\n", encoding="utf-8") + subprocess.run(["git", "add", "oracle.py"], cwd=checkout, check=True) + subprocess.run( + [ + "git", + "-c", + "user.name=FastPLMs Tests", + "-c", + "user.email=fastplms-tests@example.invalid", + "commit", + "-m", + "Pin oracle", + ], + cwd=checkout, + check=True, + ) + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=checkout, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + source = SimpleNamespace(path="vendor/upstream/toy", revision=revision) + registry = SimpleNamespace(upstreams={"toy": source}) + spec = SimpleNamespace(family=SimpleNamespace(upstreams=("toy",))) + + _validate_vendor_revisions(source_root, registry, spec) + tracked.write_text("scale = 2\n", encoding="utf-8") + with pytest.raises(ArtifactError, match="must have a clean worktree"): + _validate_vendor_revisions(source_root, registry, spec) + + +def test_repository_legal_inventory_matches_manifest_digests() -> None: + validate_repository_legal_inventory(ROOT, load_model_registry()) diff --git a/tests/release/test_binder_example_contracts.py b/tests/release/test_binder_example_contracts.py new file mode 100644 index 0000000..2a672e3 --- /dev/null +++ b/tests/release/test_binder_example_contracts.py @@ -0,0 +1,109 @@ +"""Release contracts for the binder example's validation and environment.""" + +from __future__ import annotations + +import ast +import os +import subprocess +import sys +import textwrap +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +EXAMPLE = ROOT / "examples" / "binder_design_fastplms.py" +GUIDE = ROOT / "docs" / "binder_design.md" + + +def test_binder_example_has_no_optimized_away_or_private_validation() -> None: + source = EXAMPLE.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(EXAMPLE)) + runtime_asserts = [node for node in ast.walk(tree) if isinstance(node, ast.Assert)] + + assert not runtime_asserts, "binder validation must survive python -O" + assert "abnumber.common" not in source + assert "_anarci_align" not in source + assert "Chain.multiple_domains" in source + assert "use_anarcii=True" in source + + +def test_binder_example_uses_the_binder_dependency_profile() -> None: + source = EXAMPLE.read_text(encoding="utf-8") + guide = GUIDE.read_text(encoding="utf-8") + profile = (ROOT / "requirements" / "profiles" / "binder.in").read_text( + encoding="utf-8" + ) + + assert "# /// script" not in source + assert "requires-python" not in source + assert profile.splitlines() == [ + "-r ../core.in", + "-r ../features/structure.in", + "-r ../features/binder.in", + ] + assert "Python 3.11-3.14" in guide + assert "no standalone PEP 723 dependency block" in guide + for fragment in ( + "requirements/profiles/binder.in", + "requirements/constraints/validation.txt", + "PYTHONPATH=src python examples/binder_design_fastplms.py", + ): + assert fragment in guide + + +def test_binder_validation_survives_python_optimized_mode() -> None: + program = textwrap.dedent( + """ + import torch + from examples import binder_design_fastplms as binder + + checks = ( + lambda: binder.build_initial_soft_sequence_logits("A?", batch_size=1), + lambda: binder.compute_distogram_iptm_proxy( + torch.zeros(3, 3, 128), + target_length=2, + binder_sequence="AA", + is_antibody=False, + ), + lambda: binder._binder_sequence_from_designed_sequence("missing-separator"), + ) + for check in checks: + try: + check() + except ValueError: + continue + raise SystemExit("validation disappeared under python -O") + """ + ) + environment = os.environ.copy() + environment.update( + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + OPENBLAS_NUM_THREADS="1", + PYTHONPATH=os.pathsep.join((str(ROOT / "src"), str(ROOT))), + ) + result = subprocess.run( + [sys.executable, "-O", "-c", program], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, result.stderr or result.stdout + + +def test_binder_output_contract_is_documented_fail_closed() -> None: + source = EXAMPLE.read_text(encoding="utf-8") + guide = GUIDE.read_text(encoding="utf-8") + + assert "_require_fresh_output_directory(args.output_dir)" in source + assert source.index("_require_fresh_output_directory(args.output_dir)") < source.index( + "runner.load(" + ) + assert source.index("_write_official_selection_table(") < source.rindex("_write_run_manifest(") + assert "must not already exist, including as an empty directory" in guide + assert "written atomically and last" in guide + assert "treat the\ndirectory as an incomplete run" in guide diff --git a/tests/release/test_conversion_tools.py b/tests/release/test_conversion_tools.py new file mode 100644 index 0000000..4ced423 --- /dev/null +++ b/tests/release/test_conversion_tools.py @@ -0,0 +1,339 @@ +from __future__ import annotations + +import json +import re +import pytest +import torch +from pathlib import Path + +from fastplms.registry import load_model_registry +from tools.conversion import ( + StateTransformError, + apply_state_transform, + available_state_transforms, +) +from tools.conversion.extract_esmfold2_geometry import extract_geometry +from tools.conversion.state_validation import assert_state_dict_equal + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_every_manifest_state_transform_has_a_pure_implementation() -> None: + registry = load_model_registry() + declared = {family.state_transform for family in registry.families.values()} + assert declared == set(available_state_transforms()) + + +def test_conversion_tools_contain_no_hub_mutation_or_authentication_code() -> None: + forbidden = re.compile( + r"push_to_hub|create_repo|delete_repo|upload_(?:file|folder)|" + r"\blogin\s*\(|\bHfApi\b|update_HF|snapshot_download|hf_hub_download" + ) + files = sorted((ROOT / "tools" / "conversion").glob("*.py")) + assert {path.name for path in files} == { + "__init__.py", + "extract_esmfold2_geometry.py", + "state_transforms.py", + "state_validation.py", + } + for path in files: + assert forbidden.search(path.read_text(encoding="utf-8")) is None, path + + +def test_esmfold2_geometry_extractor_is_literal_only_and_reproducible( + tmp_path: Path, +) -> None: + source = ( + ROOT + / "vendor" + / "upstream" + / "biohub-transformers" + / "src" + / "transformers" + / "models" + / "esmfold2" + / "protein_utils.py" + ) + expected_path = ( + ROOT / "src" / "fastplms" / "models" / "esmfold2" / "protein_reference_geometry.json" + ) + generated = ( + json.dumps( + extract_geometry(source), + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ) + assert generated == expected_path.read_text(encoding="utf-8") + + marker = tmp_path / "executed" + inert_source = tmp_path / "inert.py" + inert_source.write_text( + f"open({str(marker)!r}, 'w').write('executed')\n" + "PROTEIN_REF_POS: dict = {'UNK': {'CA': (0.0, 0.0, 0.0)}}\n", + encoding="utf-8", + ) + assert extract_geometry(inert_source)["residues"] == {"UNK": {"CA": (0.0, 0.0, 0.0)}} + assert not marker.exists() + + dynamic_source = tmp_path / "dynamic.py" + dynamic_source.write_text( + "PROTEIN_REF_POS: dict = make_geometry()\n", + encoding="utf-8", + ) + with pytest.raises(ValueError, match="must be a Python literal"): + extract_geometry(dynamic_source) + + +def test_identity_key_transform_is_value_exact_and_non_aliasing() -> None: + transform_id = "identity" + source = { + "decoder.block.0.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "encoder.embed_tokens.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + "lm_head.weight": torch.arange(4, dtype=torch.float32).reshape(2, 2), + } + transformed = apply_state_transform( + transform_id, + source, + expected_keys=source, + ) + + assert_state_dict_equal(source, transformed, context=transform_id) + assert list(transformed) == sorted(source) + for key in source: + assert transformed[key].data_ptr() != source[key].data_ptr() + + +def _tiny_complete_ankh_state() -> dict[str, torch.Tensor]: + return { + "shared.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), + "encoder.embed_tokens.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), + "encoder.block.0.layer.0.SelfAttention.q.weight": torch.ones(4, 4), + "decoder.embed_tokens.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), + "decoder.block.0.layer.0.SelfAttention.q.weight": torch.ones(4, 4), + "decoder.block.0.layer.1.EncDecAttention.q.weight": torch.ones(4, 4), + "lm_head.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), + } + + +def test_ankh_transform_requires_and_preserves_complete_t5_state() -> None: + source = _tiny_complete_ankh_state() + transformed = apply_state_transform( + "ankh_t5_to_fastplms_v1", + source, + expected_keys=source, + ) + + assert_state_dict_equal(source, transformed, context="ankh_t5_to_fastplms_v1") + assert list(transformed) == sorted(source) + assert all(transformed[key].data_ptr() != source[key].data_ptr() for key in source) + + +def test_ankh_transform_rejects_encoder_only_publication_state() -> None: + with pytest.raises(StateTransformError, match="complete official T5 state"): + apply_state_transform( + "ankh_t5_to_fastplms_v1", + { + "shared.weight": torch.ones(2, 4), + "encoder.embed_tokens.weight": torch.ones(2, 4), + "encoder.block.0.layer.0.SelfAttention.q.weight": torch.ones(4, 4), + }, + ) + + +def test_precision_and_rotary_table_transforms_match_published_artifacts() -> None: + source = { + "encoder.weight": torch.tensor([1.25], dtype=torch.float32), + "esm.embeddings.position_embeddings.weight": torch.ones(2, 2), + } + expected = {"encoder.weight"} + + e1 = apply_state_transform( + "e1_to_fastplms_v1", + {"encoder.weight": source["encoder.weight"]}, + expected_keys=expected, + ) + assert e1["encoder.weight"].dtype is torch.bfloat16 + + for transform_id in ("dplm_to_fastplms_v1", "dplm2_to_fastplms_v1"): + transformed = apply_state_transform( + transform_id, + source, + expected_keys=expected, + ) + assert set(transformed) == expected + + +def test_esm2_transform_matches_parity_mapping_and_is_idempotent() -> None: + from tests.parity.support.state_transforms import ( + transform_preserves_aliases, + transform_state, + ) + + source = { + "embed_tokens.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "layers.0.self_attn.q_proj.weight": torch.tensor([[1.0, 2.0]]), + "layers.0.self_attn.rot_emb.inv_freq": torch.tensor([0.5, 0.25]), + "layers.0.self_attn.out_proj.bias": torch.tensor([3.0]), + "layers.0.self_attn_layer_norm.weight": torch.tensor([4.0]), + "layers.0.fc1.weight": torch.tensor([[5.0]]), + "layers.0.fc2.bias": torch.tensor([6.0]), + "layers.0.final_layer_norm.weight": torch.tensor([7.0]), + "emb_layer_norm_after.weight": torch.tensor([8.0]), + "contact_head.regression.weight": torch.tensor([[9.0]]), + "lm_head.dense.weight": torch.tensor([[10.0]]), + "lm_head.weight": torch.tensor([[11.0]]), + "lm_head.bias": torch.tensor([12.0]), + } + expected = transform_state("esm2_hf_to_fastplms_v1", source) + transformed = apply_state_transform( + "esm2_hf_to_fastplms_v1", + source, + expected_keys=expected, + ) + + assert_state_dict_equal(expected, transformed, context="ESM2 conversion") + assert torch.equal( + transformed["esm.encoder.layer.0.attention.self.rotary_embeddings.inv_freq"], + source["layers.0.self_attn.rot_emb.inv_freq"], + ) + assert transformed["lm_head.bias"].data_ptr() != source["lm_head.bias"].data_ptr() + assert not transform_preserves_aliases("esm2_hf_to_fastplms_v1") + + canonical = apply_state_transform( + "esm2_hf_to_fastplms_v1", + transformed, + expected_keys=transformed, + ) + assert_state_dict_equal(transformed, canonical, context="canonical ESM2 conversion") + + +def test_esmc_transform_maps_official_keys_exactly() -> None: + source = { + "esmc.transformer.blocks.0.attn.layernorm_qkv.layer_norm_weight": torch.tensor([1.0, 2.0]), + "esmc.transformer.blocks.0.ffn.fc1_weight": torch.tensor([[3.0, 4.0]]), + "esmc.transformer.blocks.0.rotary._extra_state": torch.tensor(1), + "lm_head.weight": torch.tensor([[5.0, 6.0]]), + } + expected = { + "sequence_head.weight": source["lm_head.weight"], + "transformer.blocks.0.attn.layernorm_qkv.0.weight": source[ + "esmc.transformer.blocks.0.attn.layernorm_qkv.layer_norm_weight" + ], + "transformer.blocks.0.ffn.1.weight": source["esmc.transformer.blocks.0.ffn.fc1_weight"], + } + + transformed = apply_state_transform( + "esmc_to_fastplms_v1", + source, + expected_keys=expected, + ) + assert_state_dict_equal(expected, transformed, context="ESMC conversion") + + +def test_esmfold_transform_maps_native_state_and_removes_untrained_heads() -> None: + from tests.parity.support.state_transforms import transform_state + + source = { + "esm.embed_tokens.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "esm.layers.0.self_attn.q_proj.weight": torch.tensor([[1.0, 2.0]]), + "esm.contact_head.regression.weight": torch.tensor([[2.0]]), + "esm.contact_head.regression.bias": torch.tensor([2.5]), + "esm.lm_head.weight": torch.tensor([[3.0, 4.0]]), + "trunk.blocks.0.weight": torch.tensor([5.0]), + "trunk.structure_module.atom_mask": torch.tensor([True]), + "positional_encoding._float_tensor": torch.tensor([6.0]), + "lm_head.weight": torch.tensor([[7.0, 8.0]]), + } + expected = { + "esm.embeddings.word_embeddings.weight": source["esm.embed_tokens.weight"], + "esm.encoder.layer.0.attention.self.query.weight": source[ + "esm.layers.0.self_attn.q_proj.weight" + ], + "trunk.blocks.0.weight": source["trunk.blocks.0.weight"], + "lm_head.weight": source["lm_head.weight"], + } + transformed = apply_state_transform( + "esmfold_meta_to_fastplms_v1", + source, + expected_keys=expected, + ) + assert_state_dict_equal(expected, transformed, context="native ESMFold conversion") + assert_state_dict_equal( + expected, + transform_state("esmfold_meta_to_fastplms_v1", source), + context="independent native ESMFold parity transform", + ) + + canonical_with_obsolete_head = { + **transformed, + "mlm_head.weight": torch.tensor([[9.0]]), + "esm.contact_head.regression.weight": torch.tensor([[10.0]]), + "esm.contact_head.regression.bias": torch.tensor([11.0]), + "trunk.structure_module.atom_mask": torch.tensor([True]), + } + canonical = apply_state_transform( + "esmfold_meta_to_fastplms_v1", + canonical_with_obsolete_head, + expected_keys=expected, + ) + assert_state_dict_equal(expected, canonical, context="canonical ESMFold conversion") + assert_state_dict_equal( + expected, + transform_state("esmfold_meta_to_fastplms_v1", canonical_with_obsolete_head), + context="independent canonical ESMFold parity transform", + ) + assert not any(name.startswith("esm.contact_head.") for name in transformed) + assert not any(name.startswith("esm.contact_head.") for name in canonical) + + +def test_esm3_transform_adds_only_the_fastplms_wrapper_prefix() -> None: + source = {"encoder.sequence_embed.weight": torch.arange(6).reshape(2, 3)} + expected = {"esm3.encoder.sequence_embed.weight": source["encoder.sequence_embed.weight"]} + transformed = apply_state_transform( + "esm3_to_fastplms_v1", + source, + expected_keys=expected, + ) + assert_state_dict_equal(expected, transformed, context="ESM3 conversion") + + +def test_boltz2_transform_selects_only_the_declared_inference_core() -> None: + source = { + "model.module.input_embedder.weight": torch.tensor([1.0]), + "model.template_module.weight": torch.tensor([2.0]), + "ema.input_embedder.weight": torch.tensor([3.0]), + } + expected = {"core.input_embedder.weight": source["model.module.input_embedder.weight"]} + transformed = apply_state_transform( + "boltz2_inference_core_v1", + source, + expected_keys=expected, + ) + assert_state_dict_equal(expected, transformed, context="Boltz2 conversion") + + with pytest.raises(StateTransformError, match="undeclared non-inference"): + apply_state_transform( + "boltz2_inference_core_v1", + {**source, "model.training_only.weight": torch.tensor([4.0])}, + expected_keys=expected, + ) + + +@pytest.mark.parametrize( + "candidate", + [ + {"weight": torch.tensor([1.0]), "extra": torch.tensor([1.0])}, + {}, + {"weight": torch.tensor([1], dtype=torch.int64)}, + {"weight": torch.tensor([2.0])}, + ], +) +def test_exact_state_validation_rejects_schema_or_value_drift( + candidate: dict[str, torch.Tensor], +) -> None: + reference = {"weight": torch.tensor([1.0])} + with pytest.raises(AssertionError, match="state_dict parity failed"): + assert_state_dict_equal(reference, candidate, context="exact conversion") diff --git a/tests/release/test_dependency_contracts.py b/tests/release/test_dependency_contracts.py new file mode 100644 index 0000000..926f48d --- /dev/null +++ b/tests/release/test_dependency_contracts.py @@ -0,0 +1,304 @@ +"""Release contracts for intentionally scoped optional dependencies.""" + +from __future__ import annotations + +import pytest +from pathlib import Path + +from tools.remote.runtime_import_closure import ( + RuntimeImportClosureError, + inspect_runtime_import_closure, +) + + +ROOT = Path(__file__).resolve().parents[2] +REQUIREMENTS = ROOT / "requirements" + + +def _requirements(relative_path: str) -> list[str]: + path = REQUIREMENTS / relative_path + return [ + line.strip() + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + +def _package_name(requirement: str) -> str: + name = requirement.partition(";")[0] + for operator in ("==", ">="): + name = name.partition(operator)[0] + return name.strip() + + +def test_core_dependencies_are_direct_and_bounded() -> None: + assert _requirements("core.in") == [ + "torch>=2.13,<2.14", + "transformers>=5.13,<5.14", + "huggingface-hub>=0.34,<2", + "tokenizers>=0.22,<0.23", + "safetensors>=0.5,<1", + "numpy>=1.26,<3", + "einops>=0.8,<1", + "tqdm>=4.67,<5", + ] + + +def test_cpu_validation_profile_is_explicit_and_cuda_free() -> None: + assert _requirements("features/cpu.in") == ["torch==2.13.0"] + assert _requirements("constraints/validation.txt") == [ + "torch==2.13.0", + "transformers==5.13.0", + ] + assert _requirements("profiles/cpu-validation.in") == [ + "-r ../core.in", + "-r ../features/cpu.in", + "-r ../features/dev.in", + "-r ../features/structure.in", + "-r ../features/train.in", + ] + for profile in (REQUIREMENTS / "profiles").glob("*.in"): + declarations = profile.read_text(encoding="utf-8") + if "features/cpu.in" not in declarations: + continue + assert "features/cueq.in" not in declarations + assert "features/fp8.in" not in declarations + instructions = (REQUIREMENTS / "README.md").read_text(encoding="utf-8") + assert "--torch-backend cpu" in instructions + assert "requirements/constraints/validation.txt" in instructions + + +def test_structure_dependencies_are_runtime_owned_or_documented_integrations() -> None: + assert _requirements("features/structure.in") == [ + "accelerate>=1.10,<2", # Transformers device_map in the 6B quick start. + "biopython>=1.85,<2", + "biotite>=1.4,<2", + "brotli>=1.1,<2", + "msgpack>=1.1,<2", + "msgpack-numpy>=0.4.8,<1", + "omegaconf>=2.3,<3", # Explicit trusted Boltz Lightning import boundary. + "rdkit>=2025.9,<2027", + "scipy>=1.15,<2", + "zstandard>=0.23,<1", + ] + + +def test_binder_dependencies_are_bounded_and_separate_from_structure() -> None: + binder = _requirements("features/binder.in") + structure = _requirements("features/structure.in") + assert binder == [ + "abnumber==0.4.4", + "anarcii==2.0.8", + "pandas>=3.0,<3.1", + "pyarrow>=25,<26", + ] + assert {_package_name(item) for item in binder}.isdisjoint( + {_package_name(item) for item in structure} + ) + assert _requirements("profiles/binder.in") == [ + "-r ../core.in", + "-r ../features/structure.in", + "-r ../features/binder.in", + ] + + +def test_cueq_dependencies_are_version_aligned_cuda13_and_isolated() -> None: + cueq = _requirements("features/cueq.in") + structure = _requirements("features/structure.in") + assert cueq == [ + 'cuequivariance==0.10.0; platform_system == "Linux"', + 'cuequivariance-torch==0.10.0; platform_system == "Linux"', + 'cuequivariance-ops-torch-cu13==0.10.0; platform_system == "Linux"', + ] + assert not any("cuequivariance" in requirement for requirement in structure) + + source = (ROOT / "src/fastplms/models/esmfold2/modeling_esmfold2_common.py").read_text( + encoding="utf-8" + ) + assert 'find_spec("cuequivariance_ops_torch")' in source + + kernel_sources = [ + source, + (ROOT / "src/fastplms/models/boltz/vb_layers_triangular_mult.py").read_text( + encoding="utf-8" + ), + (ROOT / "src/fastplms/models/boltz/vb_tri_attn_primitives.py").read_text(encoding="utf-8"), + ] + for kernel_source in kernel_sources: + assert "cuequivariance_torch.primitives" not in kernel_source + assert 'find_spec("cuequivariance_ops_torch")' in kernel_source + assert 'import_module("cuequivariance_torch")' in kernel_source + assert "cue_module.triangle_multiplicative_update" in source + assert "cueq.triangle_multiplicative_update" in kernel_sources[1] + assert "cueq.triangle_attention" in kernel_sources[2] + assert _requirements("profiles/candidate-structure.in")[-2:] == [ + "-r ../features/cueq.in", + "-r ../features/train.in", + ] + + +def test_reporting_dependencies_are_separate_from_training_runtime() -> None: + reporting = _requirements("features/reporting.in") + training = _requirements("features/train.in") + assert reporting == [ + "matplotlib>=3.10,<4", + "scikit-learn>=1.7,<2", + "scipy>=1.15,<2", + "seaborn>=0.13,<1", + ] + assert {_package_name(item) for item in training}.isdisjoint( + {_package_name(item) for item in reporting} + ) + + +def test_dependency_instructions_install_the_cpu_validation_profile() -> None: + instructions = (REQUIREMENTS / "README.md").read_text(encoding="utf-8") + + assert "uv pip install" in instructions + assert "-r requirements/profiles/cpu-validation.in" in instructions + assert "-c requirements/constraints/validation.txt" in instructions + + +def test_runtime_import_closure_rejects_undeclared_literal_dynamic_import( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "dynamic.py").write_text( + 'import importlib\nimportlib.import_module("undeclared_dynamic_dependency")\n', + encoding="utf-8", + ) + + with pytest.raises( + RuntimeImportClosureError, + match="undeclared literal dynamic dependencies", + ): + inspect_runtime_import_closure(source_root, ROOT / "requirements") + + +def test_runtime_import_closure_rejects_optional_extra_as_core_import( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "unconditional.py").write_text("import pandas\n", encoding="utf-8") + + with pytest.raises( + RuntimeImportClosureError, + match="Unconditional import dependency scope mismatch", + ): + inspect_runtime_import_closure(source_root, ROOT / "requirements") + + +def test_runtime_import_closure_keeps_top_level_control_flow_import_time( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "conditional.py").write_text( + "enabled = True\n" + "if enabled:\n" + " import pandas\n", + encoding="utf-8", + ) + + with pytest.raises( + RuntimeImportClosureError, + match="Unconditional import dependency scope mismatch", + ): + inspect_runtime_import_closure(source_root, ROOT / "requirements") + + +def test_runtime_import_closure_records_guarded_dependency_intended_extra( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "guarded.py").write_text( + "import importlib\n" + "def load_kernel():\n" + ' return importlib.import_module("kernels")\n', + encoding="utf-8", + ) + + payload = inspect_runtime_import_closure(source_root, ROOT / "requirements") + + assert payload["feature_gated_dynamic_imports"] == [ + { + "declared_scopes": ["extra:flash"], + "kind": "dynamic", + "line": 3, + "module": "kernels", + "required_scope": "extra:flash", + "source": "guarded.py", + "source_scope": "core", + } + ] + + +def test_runtime_import_closure_uses_manifest_scope_for_feature_module( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + module_root = source_root / "models" / "feature" + module_root.mkdir(parents=True) + (source_root / "models.toml").write_text( + "[families.feature]\n" + 'extra = "structure"\n' + 'runtime_paths = ["models/feature"]\n', + encoding="utf-8", + ) + (module_root / "module.py").write_text("import scipy\n", encoding="utf-8") + + payload = inspect_runtime_import_closure(source_root, ROOT / "requirements") + + assert payload["import_time_dependencies"] == [ + { + "declared_scopes": ["extra:reporting", "extra:structure"], + "kind": "static", + "line": 1, + "module": "scipy", + "required_scope": "extra:structure", + "source": "models/feature/module.py", + "source_scope": "extra:structure", + } + ] + + +def test_runtime_import_closure_rejects_escaping_manifest_runtime_path( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "module.py").write_text("import torch\n", encoding="utf-8") + (source_root / "models.toml").write_text( + "[families.feature]\n" + 'extra = "structure"\n' + 'runtime_paths = ["../outside"]\n', + encoding="utf-8", + ) + + with pytest.raises( + RuntimeImportClosureError, + match="non-portable runtime path", + ): + inspect_runtime_import_closure(source_root, ROOT / "requirements") + + +def test_runtime_import_closure_rejects_ambiguous_guarded_extra( + tmp_path: Path, +) -> None: + source_root = tmp_path / "runtime" + source_root.mkdir() + (source_root / "guarded.py").write_text( + "def load_accelerate():\n" + " import accelerate\n" + " return accelerate\n", + encoding="utf-8", + ) + + with pytest.raises( + RuntimeImportClosureError, + match="does not map to one intended dependency scope", + ): + inspect_runtime_import_closure(source_root, ROOT / "requirements") diff --git a/tests/release/test_documentation.py b/tests/release/test_documentation.py new file mode 100644 index 0000000..7b198d1 --- /dev/null +++ b/tests/release/test_documentation.py @@ -0,0 +1,802 @@ +from __future__ import annotations + +import ast +import concurrent.futures +import os +import re +import subprocess +import sys +import pytest +from pathlib import Path +from typing import Any, Self +from urllib.parse import unquote, urlsplit + +from tools.artifacts.generate_docs import ( + CAPABILITY_EVIDENCE_SELECTORS, + EMBEDDING_CAPABILITY_ROWS, + GENERATION_CAPABILITY_ROWS, + STRUCTURE_CAPABILITY_ROWS, + attention_backend_evidence_keys, + autoclass_evidence_keys, + benchmark_autoclass_evidence_pairs, + benchmark_backend_evidence, + synchronize, +) +from tools.debug.check_notation import ( + iter_repository_files, + scan_repository, + violations_in_text, +) + + +ROOT = Path(__file__).resolve().parents[2] +MARKDOWN_ROOTS = ( + ROOT / "AGENTS.md", + ROOT / "CLAUDE.md", + ROOT / "README.md", + ROOT / "THIRD_PARTY_NOTICES.md", + ROOT / "LICENSES", + ROOT / "benchmarks" / "README.md", + ROOT / "docker" / "README.md", + ROOT / "docs", + ROOT / "examples", + ROOT / "model_cards", + ROOT / "tools" / "debug" / "README.md", + ROOT / "tools" / "remote" / "README.md", + ROOT / "vendor" / "README.md", +) +FENCE_PATTERN = re.compile( + r"^```(?P[A-Za-z0-9_+-]*)[^\n]*\n(?P.*?)^```[ \t]*$", + re.MULTILINE | re.DOTALL, +) +LINK_PATTERN = re.compile(r"!?\[[^\]]*\]\((?P[^)]+)\)") +UNBACKED_CLAIM_PATTERNS = ( + re.compile( + r"\b(?:is|are|has been|have been)\s+" + r"(?:fully\s+|exactly\s+)?equivalent\b", + re.I, + ), + re.compile( + r"\b(?:state[- ]of[- ]the[- ]art|outperforms?|" + r"\d+(?:\.\d+)?\s*[x\u00d7]\s+(?:faster|speedup)|" + r"\d+(?:\.\d+)?%\s+faster)\b", + re.I, + ), +) +LICENSE_FILE_PATTERN = re.compile(r"^LICEN[CS]E(?:[._-].*)?$", re.I) +MODEL_CARD_FILE_PATTERN = re.compile(r"^(?:MODEL_CARD|README)\.md$", re.I) +OFFLINE_EXAMPLES = ( + "artifact_loading.py", + "embedding_and_retrieval.py", + "attention_switching.py", + "ankh_embeddings.py", + "generation.py", + "e1_rag.py", + "ttt.py", + "structure_preparation.py", + "task_heads.py", + "fine_tuning.py", + "binder_design_fastplms.py", +) + + +def _markdown_files() -> tuple[Path, ...]: + paths: list[Path] = [] + for candidate in MARKDOWN_ROOTS: + if candidate.is_file(): + paths.append(candidate) + elif candidate.is_dir(): + paths.extend(sorted(candidate.rglob("*.md"))) + return tuple(paths) + + +def _python_snippet(path: Path, marker: str) -> str: + text = path.read_text(encoding="utf-8") + matches = [ + match.group("body") + for match in FENCE_PATTERN.finditer(text) + if match.group("language").lower() in {"python", "py"} and marker in match.group("body") + ] + if len(matches) != 1: + raise AssertionError( + f"Expected one Python snippet containing {marker!r} in {path}, found {len(matches)}." + ) + return matches[0] + + +def _local_link_target(source: Path, raw_target: str) -> Path | None: + target = raw_target.strip() + if target.startswith("<") and target.endswith(">"): + target = target[1:-1] + split = urlsplit(target) + if split.scheme or split.netloc or not split.path: + return None + decoded = unquote(split.path) + destination = ROOT / decoded.lstrip("/") if decoded.startswith("/") else source.parent / decoded + resolved = destination.resolve() + try: + resolved.relative_to(ROOT.resolve()) + except ValueError as error: + raise AssertionError( + f"Documentation link escapes the repository: {source}: {raw_target}" + ) from error + return resolved + + +def test_shape_notation_detector_rejects_square_and_uppercase_dimensions() -> None: + text = ( + "H: [" + "B, L, 81, 2560]\n" + "Z: (" + "B, L, D)\n" + "M: [" + "n_atoms]\n" + "X: [" + "samples, atoms, 3]" + ) + violations = list(violations_in_text(text, path=Path("example.md"))) + assert len(violations) == 4 + + +def test_repository_documentation_uses_canonical_shape_notation() -> None: + violations = scan_repository(ROOT) + assert not violations, "\n" + "\n".join(violation.render(ROOT) for violation in violations) + + +def test_notation_inventory_includes_container_and_provenance_docs() -> None: + paths = {path.relative_to(ROOT).as_posix() for path in iter_repository_files(ROOT)} + assert { + "LICENSES/README.md", + "THIRD_PARTY_NOTICES.md", + "docker/Dockerfile", + "docker/docker-bake.hcl", + "docker/compose.yaml", + "vendor/README.md", + }.issubset(paths) + + +def test_runtime_model_packages_do_not_embed_licenses_or_model_cards() -> None: + model_root = ROOT / "src" / "fastplms" / "models" + misplaced = sorted( + path.relative_to(ROOT).as_posix() + for path in model_root.rglob("*") + if path.is_file() + and ( + LICENSE_FILE_PATTERN.fullmatch(path.name) + or MODEL_CARD_FILE_PATTERN.fullmatch(path.name) + ) + ) + assert not misplaced, ( + "Runtime model packages must not embed license files or model cards; " + "use LICENSES/ and model_cards/:\n" + "\n".join(misplaced) + ) + + +def test_manifest_generated_documentation_is_current() -> None: + failures = synchronize(ROOT, check=True) + assert not failures, "\n" + "\n".join(failures) + + +def test_generated_capability_evidence_covers_manifest() -> None: + from fastplms.registry import load_model_registry + + registry = load_model_registry() + text = (ROOT / "docs" / "generated" / "capability_evidence.md").read_text(encoding="utf-8") + for family in registry.families.values(): + assert f"`{family.id}`" in text + assert f"`{family.tokenizer_mode}`" in text + for auto_class in family.auto_map: + assert f"`{auto_class}`" in text + for backend in family.attention: + assert f"`{backend}`" in text + for heading in ( + "Input, embedding, and storage contracts", + "Generation and adaptation contracts", + "Structure contracts", + ): + assert f"## {heading}" in text + advertised_entries = sum(len(family.auto_map) for family in registry.families.values()) + assert text.count("[runnable AutoClass contract]") == advertised_entries + assert "ANKH embeddings and generation" in text + assert "pocket requests fail closed" in text + + +def test_autoclass_capability_evidence_matches_runtime_and_benchmark_selectors() -> None: + from benchmarks.suite import benchmark_auto_class, benchmark_cases + from fastplms.registry import get_model_registry + from tests.cpu.test_autoclass_evidence_matrix import AUTOCLASS_EVIDENCE + + registry = get_model_registry() + manifest_pairs = { + (family.id, auto_class) + for family in registry.families.values() + for auto_class in family.auto_map + } + assert set(AUTOCLASS_EVIDENCE) == manifest_pairs + + specs_by_checkpoint = { + (spec.fast.repo_id, spec.fast.revision): spec for spec in registry.values() + } + benchmark_pairs: set[tuple[str, str]] = set() + for case in benchmark_cases(family=None, quick=False, local_files_only=True): + if not case.claim_eligible: + continue + spec = specs_by_checkpoint[(case.model, case.revision)] + assert case.auto_class == benchmark_auto_class(spec) + benchmark_pairs.add((spec.family.id, case.auto_class)) + assert benchmark_autoclass_evidence_pairs(registry) == benchmark_pairs + + for family_id, auto_class in manifest_pairs: + family = registry.families[family_id] + evidence = set(autoclass_evidence_keys(registry, family_id, auto_class)) + assert {"cpu:autoclass-runtime", "artifact:checkpoint-autoclasses"}.issubset(evidence) + assert not any(key.startswith("feature:") for key in evidence) + assert ("benchmark:claim-eligible-primary-head" in evidence) == ( + (family_id, auto_class) in benchmark_pairs + ) + if auto_class == "AutoConfig" or "Classification" in auto_class: + assert not any(key.startswith("compliance:") for key in evidence) + assert not any(key.startswith("benchmark:") for key in evidence) + if family.id == "ankh" and auto_class == "AutoModelForSeq2SeqLM": + assert "compliance:ankh-seq2seq" in evidence + assert "benchmark:claim-eligible-primary-head" not in evidence + + +def test_backend_capability_evidence_uses_real_nightly_and_benchmark_scopes() -> None: + from benchmarks.suite import benchmark_cases + from fastplms.registry import get_model_registry + + registry = get_model_registry() + advertised = {backend for family in registry.families.values() for backend in family.attention} + benchmarked = { + case.backend + for case in benchmark_cases(family=None, quick=False, local_files_only=True) + if case.claim_eligible + } + assert benchmark_backend_evidence(registry) == benchmarked + + for backend in advertised: + evidence = set(attention_backend_evidence_keys(registry, backend)) + assert "cpu:attention-contracts" in evidence + assert not any(key.startswith("feature:") for key in evidence) + assert ("benchmark:claim-eligible-backends" in evidence) == (backend in benchmarked) + has_sequence_family = any( + family.tokenizer_mode != "structure" and backend in family.attention + for family in registry.families.values() + ) + has_current_gh200_execution = backend in {"eager", "sdpa", "flex_attention"} + assert ("nightly:sequence-backends" in evidence) == ( + has_sequence_family and has_current_gh200_execution + ) + assert ("historical:fa2-focused" in evidence) == (backend == "flash_attention_2") + assert ("compliance:flash-unavailable-gh200" in evidence) == backend.startswith( + "flash_attention_" + ) + if backend.startswith("flash_attention_"): + assert "compliance:deep-backends" not in evidence + assert "benchmark:claim-eligible-backends" not in evidence + + +def test_capability_rows_fail_closed_against_invalid_tier_inheritance() -> None: + rows = EMBEDDING_CAPABILITY_ROWS + GENERATION_CAPABILITY_ROWS + STRUCTURE_CAPABILITY_ROWS + by_capability = {row.capability: row for row in rows} + assert len(by_capability) == len(rows) + for row in rows: + assert row.evidence + assert set(row.evidence).issubset(CAPABILITY_EVIDENCE_SELECTORS) + + peft = by_capability["Trainer/PEFT LoRA with immutable inputs and verified save/reload"] + assert peft.evidence == ("cpu:peft", "nightly:peft") + + binder = by_capability["Atom-dense binder optimization and critic reporting"] + assert not any(key.startswith("benchmark:") for key in binder.evidence) + + artifact = by_capability["Offline local artifact AutoClass loading"] + assert artifact.evidence == ( + "cpu:artifact-example", + "artifact:checkpoint-autoclasses", + ) + + mapping = by_capability["Ordered mapping or one-shot generator"] + assert "tests/cpu/test_embedding_contracts.py" in mapping.example + extended_poolers = by_capability["Max/norm/median/variance/CLS/PARTI pooling"] + assert "tests/cpu/test_embedding_contracts.py" in extended_poolers.example + for row in EMBEDDING_CAPABILITY_ROWS: + if row.capability == "E1 raw-sequence and MSA-aware ordered embeddings": + continue + assert not any(key.startswith("feature:") for key in row.evidence) + + +def test_capability_evidence_selectors_resolve_to_their_declared_jobs() -> None: + from tools.remote.run import SUITES + + compose = (ROOT / "docker" / "compose.yaml").read_text(encoding="utf-8") + testing = (ROOT / "docs" / "testing.md").read_text(encoding="utf-8") + assert "tests/cpu" in testing and "cpu_contract" in testing + + commands = { + name: tuple(suite.command) + + tuple(part for command in suite.pre_commands for part in command) + for name, suite in SUITES.items() + } + for selector in CAPABILITY_EVIDENCE_SELECTORS.values(): + for target in selector.targets: + relative = target.split("::", maxsplit=1)[0] + path = ROOT / relative + assert path.exists(), f"Evidence selector target does not exist: {relative}" + if selector.tier == "cpu_contract": + assert relative.startswith("tests/cpu/") + elif selector.tier in {"feature", "nightly", "compliance"}: + assert relative in commands[selector.tier] + elif selector.tier == "artifact": + assert relative.startswith("tests/release/") + assert "tests/release" in commands["artifact"] + elif selector.tier == "structure": + assert relative == "tests/structure" or relative.startswith("tests/structure/") + assert "tests/structure" in commands["structure"] + elif selector.tier == "benchmark": + assert relative == "benchmarks/suite.py" + assert "benchmark" in commands["benchmark"] + assert 'entrypoint: ["python", "-m", "benchmarks.suite"]' in compose + elif selector.tier == "historical": + assert relative == "tools/remote/run.py" + assert target.endswith("::_kernel_capability_preflight") + else: + raise AssertionError(f"Unvalidated evidence tier: {selector.tier}") + + +def test_generated_esmc_cards_state_mask_precedence_and_route_hopper_scope_to_docs() -> None: + for name in ("esmc_small.md", "esmc_large.md", "esmc_6b.md"): + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert "When `sequence_id` is supplied" in text + assert "`attention_mask` is ignored" in text + assert "/docs/attention_backends.md" in text + assert "current exact GH200/aarch64" not in text + assert "H100 environment" not in text + + attention_docs = " ".join( + (ROOT / "docs/attention_backends.md").read_text(encoding="utf-8").split() + ) + assert "exact GH200/aarch64 workstation" in attention_docs + assert "H100 and H200" in attention_docs + assert "not current release evidence" in attention_docs + + +def test_generated_cards_publish_canonical_state_commitments() -> None: + from fastplms.registry import get_model_registry + + for spec in get_model_registry().values(): + if spec.canonical_state_sha256 is None: + continue + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + assert f"Canonical transformed state SHA-256: `{spec.canonical_state_sha256}`" in card + assert "Conversion equality attestation: recorded in `provenance.json`" in card + + +def test_curated_offline_examples_expose_executable_help() -> None: + environment = os.environ.copy() + environment.update( + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + OMP_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + MKL_NUM_THREADS="1", + ) + + def run_help(name: str) -> tuple[str, subprocess.CompletedProcess[str]]: + path = ROOT / "examples" / name + return name, subprocess.run( + [sys.executable, str(path), "--help"], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(run_help, OFFLINE_EXAMPLES)) + + for name, result in results: + assert result.returncode == 0, f"{name}: {result.stderr}" + assert "usage:" in result.stdout.lower() + help_by_name = dict(results) + structure_help = " ".join(help_by_name["structure_preparation.py"].stdout.split()) + assert "requires a full 48-block checkpoint" in structure_help + + +def test_routine_setup_avoids_parity_submodules_and_documents_manual_cpu_gate() -> None: + readme = (ROOT / "README.md").read_text(encoding="utf-8") + dependencies = readme.split("## Dependencies", maxsplit=1)[1].split( + "## Quick start", maxsplit=1 + )[0] + validation = readme.split("## Validation and reproducibility", maxsplit=1)[1] + assert "git clone --recurse-submodules" not in dependencies + assert "git submodule update --init --recursive" not in dependencies + assert "Official reference repositories are not runtime" in dependencies + assert "git submodule update --init --recursive" in validation + assert "does not use GitHub Actions" in validation + assert "tests/cpu" in validation + + remote = (ROOT / "tools" / "remote" / "README.md").read_text(encoding="utf-8") + testing = (ROOT / "docs" / "testing.md").read_text(encoding="utf-8") + assert "does not use GitHub Actions" in remote + assert "no GitHub Actions workflows" in testing + + +def test_container_guide_runs_complete_candidate_and_compliance_workflows() -> None: + text = (ROOT / "docker" / "README.md").read_text(encoding="utf-8") + assert "candidate --load" in text + assert "python -m pytest tests/unit tests/integration" in text + assert '-m "not gpu and not slow and not structure"' in text + assert "--suite compliance" in text + assert "Building those images alone does not run\nparity" in text + + +def test_hub_quick_starts_follow_install_and_platform_contracts() -> None: + paths = ( + ROOT / "README.md", + ROOT / "docs" / "attention_backends.md", + ROOT / "docs" / "embedding_api.md", + ROOT / "docs" / "esmfold2.md", + ROOT / "docs" / "finetuning.md", + ROOT / "docs" / "migration.md", + ROOT / "docs" / "models.md", + ) + for path in paths: + text = path.read_text(encoding="utf-8") + loading = text.index(".from_pretrained(") + prefix = text[:loading] + assert "pip install" in prefix, path.relative_to(ROOT) + assert re.search(r"Python 3\.11(?:-| through )3\.14", prefix), path.relative_to(ROOT) + assert "PyTorch 2.13" in prefix, path.relative_to(ROOT) + assert "Transformers 5.13" in prefix, path.relative_to(ROOT) + + +def test_esmfold2_fast_docs_do_not_claim_msa_conditioning() -> None: + for path in ( + ROOT / "README.md", + ROOT / "docs" / "migration.md", + ROOT / "docs" / "models.md", + ROOT / "examples" / "README.md", + ): + text = " ".join(path.read_text(encoding="utf-8").split()) + assert "24 folding blocks" in text, path.relative_to(ROOT) + assert "48" in text and "optional MSA conditioning" in text, path.relative_to(ROOT) + assert "reject MSA-derived inputs" in text, path.relative_to(ROOT) + assert "https://biohub.ai/papers/esm_protein.pdf" in text, path.relative_to(ROOT) + + readme = (ROOT / "README.md").read_text(encoding="utf-8") + quick_start = readme.split( + "### ESMFold2 folding and learned representations", + maxsplit=1, + )[1].split("## Attention backends", maxsplit=1)[0] + normalized_quick_start = " ".join(quick_start.split()) + assert '"Synthyra/ESMFold2-Fast"' in quick_start + assert ( + "quick start below intentionally loads Fast and supplies no MSA" in normalized_quick_start + ) + assert "Protein inputs can also carry an MSA" not in quick_start + assert "Its ESMFold2 MSA branch requires one of the full checkpoints" in normalized_quick_start + + native_preparation = readme.split("### Native biological preparation", maxsplit=1)[1].split( + "### Ordered embedding results", maxsplit=1 + )[0] + normalized_native_preparation = " ".join(native_preparation.split()) + assert "Full ESMFold2 checkpoints additionally retain optional MSA conditioning" in ( + normalized_native_preparation + ) + assert "ESMFold2 Fast checkpoints reject MSA-derived inputs" in normalized_native_preparation + + docs_index = " ".join((ROOT / "docs" / "README.md").read_text(encoding="utf-8").split()) + assert "the distinct full and Fast MSA contracts" in docs_index + + +def test_ankh_seq2seq_docs_describe_live_full_checkpoints() -> None: + paths = ( + ROOT / "README.md", + ROOT / "docs" / "artifacts.md", + ROOT / "docs" / "embedding_api.md", + ROOT / "docs" / "migration.md", + ROOT / "docs" / "models.md", + ROOT / "examples" / "README.md", + ) + for path in paths: + text = path.read_text(encoding="utf-8") + assert "legacy encoder-only" not in " ".join(text.split()), path.relative_to(ROOT) + + readme = (ROOT / "README.md").read_text(encoding="utf-8") + seq2seq_section = readme.split("ANKH selects the encoder final state", maxsplit=1)[1].split( + "### Safetensors output", maxsplit=1 + )[0] + assert '"Synthyra/ANKH_base"' in seq2seq_section + assert "local_files_only=True" not in seq2seq_section + + cards = { + "ankh_base.md": "Synthyra/ANKH_base", + "ankh_large.md": "Synthyra/ANKH_large", + "ankh2_large.md": "Synthyra/ANKH2_large", + "ankh3_large.md": "Synthyra/ANKH3_large", + "ankh3_xl.md": "Synthyra/ANKH3_xl", + } + for filename, repo_id in cards.items(): + path = ROOT / "model_cards" / filename + text = path.read_text(encoding="utf-8") + normalized = " ".join(text.split()) + assert "legacy encoder-only" not in normalized, path.relative_to(ROOT) + assert repo_id in text, path.relative_to(ROOT) + assert "contains the complete ANKH encoder-decoder checkpoint" in normalized + seq2seq_loads = re.findall( + r"AutoModelForSeq2SeqLM\.from_pretrained\((.*?)\)", + text, + flags=re.DOTALL, + ) + assert seq2seq_loads, path.relative_to(ROOT) + for call in seq2seq_loads: + assert repo_id in call or "repo_id" in call, path.relative_to(ROOT) + + +def test_examples_readme_indexes_every_entry_point_and_states_coverage_boundaries() -> None: + text = (ROOT / "examples" / "README.md").read_text(encoding="utf-8") + entry_points = { + path.name + for path in (ROOT / "examples").glob("*.py") + if path.name not in {"__init__.py", "_runtime.py"} + } + assert entry_points + assert not {name for name in entry_points if f"`{name}`" not in text} + for required in ( + "## Embedding coverage matrix", + "base weights + untrained task head", + "LoRA is the demonstrated PEFT method", + "arbitrary `Dataset.save_to_disk()` trees are not", + "--device cpu|cuda[:index]", + "--dtype float32|bfloat16", + ): + assert required in text + + +def test_generated_cards_put_installation_before_hub_quick_start() -> None: + for path in sorted((ROOT / "model_cards").glob("*.md")): + if path.name == "README.md": + continue + text = path.read_text(encoding="utf-8") + assert text.index("## Install and platform requirements") < text.index("## Quick start") + assert "resolve/main/requirements.txt" in text + assert "fastplms @ git+" not in text + assert "implementation itself is embedded in the model repository" in text + + for name in ( + "boltz2.md", + "esmfold.md", + "esmfold2.md", + "esmfold2_fast.md", + "esmfold2_experimental_cutoff2025.md", + "esmfold2_experimental_fast_cutoff2025.md", + ): + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert "exact NVIDIA GH200 on Linux aarch64" in text + assert "validated release target is Linux x86-64" not in text + + +def test_esmfold2_cards_match_checkpoint_specific_msa_contracts() -> None: + generic_msa_claim = "typed interface also supports RNA, protein MSAs" + for name in ( + "esmfold2_fast.md", + "esmfold2_experimental_fast_cutoff2025.md", + ): + path = ROOT / "model_cards" / name + text = path.read_text(encoding="utf-8") + normalized = " ".join(text.split()) + assert "24-block" in normalized or "24 folding blocks" in normalized + assert "trained without MSA conditioning" in normalized + assert "ProteinInput.msa" in normalized + assert "MSA-derived" in normalized and "reject" in normalized + assert "multichain" in normalized and "multimolecule" in normalized + assert "msa=None" in normalized + assert generic_msa_claim not in normalized + + for name in ( + "esmfold2.md", + "esmfold2_experimental_cutoff2025.md", + ): + path = ROOT / "model_cards" / name + normalized = " ".join(path.read_text(encoding="utf-8").split()) + assert "48-block" in normalized or "48 folding blocks" in normalized + assert "optional MSA" in normalized + + +def test_esmc_cards_disclose_supported_divergence_without_fabricated_metrics() -> None: + for name in ("esmc_small.md", "esmc_large.md", "esmc_6b.md"): + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert "Supported, numerically divergent" in text + assert "Pending release measurement" in text + assert "SDPA is the default" in text + + +def test_esmc_release_operations_live_in_docs_not_model_cards() -> None: + card_names = ( + "esmc_small.md", + "esmc_large.md", + "esmc_6b.md", + "esmfold2.md", + "esmfold2_fast.md", + "esmfold2_experimental_cutoff2025.md", + "esmfold2_experimental_fast_cutoff2025.md", + ) + for name in card_names: + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert "Locked oracle package compatibility exception" not in text + assert "nvidia-cusparselt-cu13" not in text + assert "/docs/attention_backends.md" in text + assert "/docs/generated/capability_evidence.md" in text + + evidence = (ROOT / "docs/generated/capability_evidence.md").read_text( + encoding="utf-8" + ) + assert "Locked oracle package compatibility exception" in evidence + assert "nvidia-cusparselt-cu13==0.8.1" in evidence + + +def test_dplm_cards_mark_apache_weights_redistributable() -> None: + for name in ( + "dplm_150m.md", + "dplm_650m.md", + "dplm_3b.md", + "dplm2_150m.md", + "dplm2_650m.md", + "dplm2_3b.md", + ): + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert 'license: "apache-2.0"' in text + assert "Weight license status: `resolved`" in text + assert "Redistributable: `true`" in text + assert "/bytedance/dplm/blob/main/LICENSE" in text + assert "/README.md#overview" in text + + +def test_esmfold2_cards_disclose_race_safe_pickle_boundary() -> None: + for name in ( + "esmfold2.md", + "esmfold2_fast.md", + "esmfold2_experimental_cutoff2025.md", + "esmfold2_experimental_fast_cutoff2025.md", + ): + text = (ROOT / "model_cards" / name).read_text(encoding="utf-8") + assert "`cache_dir`" in text + assert "symlinks are rejected" in text + assert "private temporary snapshot" in text + assert "path-replacement and in-place source-write races" in text + + +def test_artifact_docs_describe_private_verified_runtime_bridge() -> None: + text = (ROOT / "docs" / "artifacts.md").read_text(encoding="utf-8") + assert "loader-owned private `TemporaryDirectory`" in text + assert "re-hashes the exact extracted inventory" in text + assert "rejects symlinks, bytecode, non-file entries" in text + assert "hash-named directory in the Transformers module cache" not in text + + +def test_finetuning_docs_pin_inputs_and_describe_verified_final_artifact() -> None: + readme = (ROOT / "README.md").read_text(encoding="utf-8") + guide = (ROOT / "docs" / "finetuning.md").read_text(encoding="utf-8") + for text in (readme, guide): + assert "--model-revision 185ecbd45665d050a8dae326d91886d330c5f9d0" in text + assert "--classification-dataset-revision " in text + assert "7e18f1b98859b0a3e3da283f63d0a153b774cf1f" in text + for flag, revision in ( + ( + "--regression-train-dataset-revision", + "f4a51e5e9f2c2a0185693f9fbcffc02d9dae08db", + ), + ( + "--regression-validation-dataset-revision", + "826ccfb1488d52b7b361802fbde161373247d084", + ), + ( + "--regression-test-dataset-revision", + "4e22f014745728fca2d9c10f2f2cfd5a29a4981c", + ), + ): + assert flag in guide + assert revision in guide + assert "`ordered_rows_sha256`" in guide + assert "`reload_verified: true`" in guide + assert "first `min(2, len(test_dataset))` rows" in guide + assert "remains on the original\nprepared Trainer" in guide + assert "Before promotion, supplement" not in guide + + +def test_documentation_local_links_resolve() -> None: + failures: list[str] = [] + for path in _markdown_files(): + text = path.read_text(encoding="utf-8") + for match in LINK_PATTERN.finditer(text): + target = _local_link_target(path, match.group("target")) + if target is not None and not target.exists(): + line = text.count("\n", 0, match.start()) + 1 + failures.append( + f"{path.relative_to(ROOT)}:{line}: missing link target " + f"{target.relative_to(ROOT)}" + ) + assert not failures, "\n" + "\n".join(failures) + + +def test_python_documentation_fences_compile() -> None: + failures: list[str] = [] + count = 0 + for path in _markdown_files(): + text = path.read_text(encoding="utf-8") + for match in FENCE_PATTERN.finditer(text): + if match.group("language").lower() not in {"python", "py"}: + continue + count += 1 + line = text.count("\n", 0, match.start("body")) + 1 + try: + ast.parse(match.group("body"), filename=f"{path}:{line}") + except SyntaxError as error: + failures.append(f"{path.relative_to(ROOT)}:{line}: {error.msg}") + assert count > 0, "No executable Python documentation snippets were found." + assert not failures, "\n" + "\n".join(failures) + + +def test_readme_embedding_snippet_executes(monkeypatch: pytest.MonkeyPatch) -> None: + import fastplms + + observed: dict[str, object] = {} + + def fake_embed_dataset(model: object, inputs: object, **kwargs: Any) -> object: + observed.update(model=model, inputs=inputs, kwargs=kwargs) + return object() + + monkeypatch.setattr(fastplms, "embed_dataset", fake_embed_dataset) + namespace = {"model": object()} + snippet = _python_snippet(ROOT / "README.md", "EmbeddingInput, embed_dataset") + exec(compile(snippet, "README.md", "exec"), namespace) + + inputs = observed["inputs"] + assert [record.id for record in inputs] == ["protein-a", "protein-a"] + assert observed["kwargs"] == { + "batch_size": 2, + "pooling": ("mean", "std"), + "output": "embeddings", + } + + +def test_readme_automodel_snippet_executes_without_network( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import transformers + + observed: dict[str, object] = {} + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, model_id: str, **kwargs: Any) -> Self: + observed.update(model_id=model_id, kwargs=kwargs) + return cls() + + def eval(self) -> Self: + return self + + monkeypatch.setattr(transformers, "AutoModel", FakeAutoModel) + snippet = _python_snippet(ROOT / "README.md", 'attn_implementation="sdpa"') + exec(compile(snippet, "README.md", "exec"), {}) + + assert observed == { + "model_id": "Synthyra/ESM2-150M", + "kwargs": { + "trust_remote_code": True, + "attn_implementation": "sdpa", + }, + } + + +def test_documentation_does_not_make_unbacked_equivalence_or_speed_claims() -> None: + failures: list[str] = [] + for path in _markdown_files(): + text = path.read_text(encoding="utf-8") + for pattern in UNBACKED_CLAIM_PATTERNS: + for match in pattern.finditer(text): + line = text.count("\n", 0, match.start()) + 1 + failures.append( + f"{path.relative_to(ROOT)}:{line}: unbacked claim {match.group()!r}" + ) + assert not failures, "\n" + "\n".join(failures) diff --git a/tests/release/test_dplm_source_independence.py b/tests/release/test_dplm_source_independence.py new file mode 100644 index 0000000..00ace09 --- /dev/null +++ b/tests/release/test_dplm_source_independence.py @@ -0,0 +1,119 @@ +"""Fail closed when DPLM runtime units overlap the pinned parity oracle.""" + +from __future__ import annotations + +import ast +import copy +import pytest +from difflib import SequenceMatcher +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +LOCAL_MODEL = ROOT / "src/fastplms/models/dplm/modeling_dplm.py" +UPSTREAM_MODEL = ( + ROOT + / "vendor/upstream/dplm/src/byprot/models/dplm/modules/dplm_modeling_esm.py" +) +MAX_FUNCTION_SIMILARITY = 0.75 +MAX_EXACT_BLOCK_LINES = 10 +SOURCE_PAIRS = ( + ("ModifiedEsmSelfAttention.forward", "ModifiedEsmSelfAttention.forward"), + ("ModifiedEsmAttention.__init__", "ModifiedEsmAttention.__init__"), + ("ModifiedEsmLayer.__init__", "ModifiedEsmLayer.__init__"), + ("ModifiedEsmEncoder.__init__", "ModifiedEsmEncoder.__init__"), + ("FAST_DPLM_ENCODER.forward", "ModifiedEsmModel.forward"), +) + + +def _function(path: Path, qualified_name: str) -> ast.FunctionDef: + body: list[ast.stmt] = ast.parse( + path.read_text(encoding="utf-8"), + filename=str(path), + ).body + selected: ast.AST | None = None + for part in qualified_name.split("."): + selected = next( + ( + node + for node in body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + and node.name == part + ), + None, + ) + assert selected is not None, f"{qualified_name!r} is absent from {path}" + body = selected.body + assert isinstance(selected, ast.FunctionDef) + return selected + + +def _normalized_ast_lines(node: ast.FunctionDef) -> list[str]: + normalized = copy.deepcopy(node) + normalized.name = "function" + normalized.decorator_list = [] + normalized.returns = None + for argument in ( + *normalized.args.posonlyargs, + *normalized.args.args, + *normalized.args.kwonlyargs, + ): + argument.annotation = None + if normalized.args.vararg is not None: + normalized.args.vararg.annotation = None + if normalized.args.kwarg is not None: + normalized.args.kwarg.annotation = None + if ( + normalized.body + and isinstance(normalized.body[0], ast.Expr) + and isinstance(normalized.body[0].value, ast.Constant) + and isinstance(normalized.body[0].value.value, str) + ): + normalized.body.pop(0) + ast.fix_missing_locations(normalized) + return [ + " ".join(line.strip().split()) + for line in ast.unparse(normalized).splitlines() + if line.strip() + ] + + +@pytest.mark.parametrize( + ("local_name", "upstream_name"), + SOURCE_PAIRS, + ids=[local_name for local_name, _ in SOURCE_PAIRS], +) +def test_dplm_functions_are_independently_implemented( + local_name: str, + upstream_name: str, +) -> None: + assert UPSTREAM_MODEL.is_file(), f"pinned DPLM source is missing: {UPSTREAM_MODEL}" + local_lines = _normalized_ast_lines(_function(LOCAL_MODEL, local_name)) + upstream_lines = _normalized_ast_lines(_function(UPSTREAM_MODEL, upstream_name)) + matcher = SequenceMatcher(None, local_lines, upstream_lines, autojunk=False) + similarity = matcher.ratio() + assert similarity < MAX_FUNCTION_SIMILARITY, ( + f"{local_name} has normalized AST similarity {similarity:.3f} to " + f"{UPSTREAM_MODEL.relative_to(ROOT)}::{upstream_name}" + ) + exact_lines = max(block.size for block in matcher.get_matching_blocks()) + assert exact_lines <= MAX_EXACT_BLOCK_LINES, ( + f"{local_name} retains an exact {exact_lines}-line source block from " + f"{UPSTREAM_MODEL.relative_to(ROOT)}::{upstream_name}" + ) + + +def test_dplm_source_does_not_import_the_parity_oracle() -> None: + tree = ast.parse(LOCAL_MODEL.read_text(encoding="utf-8"), filename=str(LOCAL_MODEL)) + imported_roots = { + alias.name.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + } + imported_roots.update( + node.module.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module + ) + assert imported_roots.isdisjoint({"byprot", "vendor"}) diff --git a/tests/release/test_e1_attribution.py b/tests/release/test_e1_attribution.py new file mode 100644 index 0000000..5443318 --- /dev/null +++ b/tests/release/test_e1_attribution.py @@ -0,0 +1,59 @@ +"""Release checks for the E1 agreement's runtime attribution.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] + + +def _run_isolated(code: str) -> subprocess.CompletedProcess[str]: + environment = os.environ.copy() + source_path = str(ROOT / "src") + inherited_path = environment.get("PYTHONPATH") + environment["PYTHONPATH"] = ( + os.pathsep.join((source_path, inherited_path)) if inherited_path else source_path + ) + return subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + env=environment, + check=False, + capture_output=True, + text=True, + ) + + +def test_importing_e1_does_not_display_runtime_attribution() -> None: + result = _run_isolated("import fastplms.models.e1.modeling_e1") + assert result.returncode == 0, result.stderr + assert "Profluent-E1" not in result.stdout + assert "Profluent-E1" not in result.stderr + + +def test_constructing_public_e1_model_displays_attribution_once() -> None: + result = _run_isolated( + """ +from fastplms.models.e1.modeling_e1 import E1Config, E1Model + +config = E1Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + max_num_sequences=8, + max_num_positions_within_seq=64, + max_num_positions_global=256, + dtype="float32", +) +E1Model(config) +""" + ) + assert result.returncode == 0, result.stderr + output = result.stdout + result.stderr + assert output.count("Profluent-E1") == 1, output diff --git a/tests/release/test_e1_source_independence.py b/tests/release/test_e1_source_independence.py new file mode 100644 index 0000000..d73f32e --- /dev/null +++ b/tests/release/test_e1_source_independence.py @@ -0,0 +1,139 @@ +"""Fail closed when E1 runtime functions overlap the pinned parity oracle.""" + +from __future__ import annotations + +import ast +import copy +import pytest +from difflib import SequenceMatcher +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +LOCAL_MODEL = ROOT / "src/fastplms/models/e1/modeling_e1.py" +LOCAL_ATTENTION = ROOT / "src/fastplms/models/e1/attention.py" +LOCAL_PREPARATION = ROOT / "src/fastplms/models/e1/preparation.py" +UPSTREAM = ROOT / "vendor/upstream/e1/src/E1" +MAX_FUNCTION_SIMILARITY = 0.75 + +# Compare only functions that implement the same public or mathematical +# contract. Function-level ASTs cannot be diluted by unrelated model code. +SOURCE_PAIRS = ( + ( + "E1BatchPreparer.prepare_multiseq", + LOCAL_PREPARATION, + UPSTREAM / "batch_preparer.py", + "E1BatchPreparer.prepare_multiseq", + ), + ( + "E1BatchPreparer.prepare_singleseq", + LOCAL_PREPARATION, + UPSTREAM / "batch_preparer.py", + "E1BatchPreparer.prepare_singleseq", + ), + ( + "get_overlapping_blocks", + LOCAL_ATTENTION, + UPSTREAM / "model/varlen_flex_attention.py", + "get_overlapping_blocks", + ), + ( + "direct_block_mask", + LOCAL_ATTENTION, + UPSTREAM / "model/varlen_flex_attention.py", + "direct_block_mask", + ), + ( + "_get_unpad_data", + LOCAL_ATTENTION, + UPSTREAM / "model/flash_attention_utils.py", + "_get_unpad_data", + ), + ( + "E1PreTrainedModel._init_weights", + LOCAL_MODEL, + UPSTREAM / "modeling.py", + "E1PreTrainedModel._init_weights", + ), + ( + "FAST_E1_ENCODER.forward", + LOCAL_MODEL, + UPSTREAM / "modeling.py", + "E1Model.forward", + ), +) + + +def _function(path: Path, qualified_name: str) -> ast.FunctionDef: + body: list[ast.stmt] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)).body + selected: ast.AST | None = None + for part in qualified_name.split("."): + selected = next( + ( + node + for node in body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name == part + ), + None, + ) + assert selected is not None, f"{qualified_name!r} is absent from {path}" + body = selected.body + assert isinstance(selected, ast.FunctionDef) + return selected + + +def _normalized_ast_lines(node: ast.FunctionDef) -> list[str]: + normalized = copy.deepcopy(node) + normalized.name = "function" + normalized.decorator_list = [] + normalized.returns = None + for argument in ( + *normalized.args.posonlyargs, + *normalized.args.args, + *normalized.args.kwonlyargs, + ): + argument.annotation = None + if normalized.args.vararg is not None: + normalized.args.vararg.annotation = None + if normalized.args.kwarg is not None: + normalized.args.kwarg.annotation = None + if ( + normalized.body + and isinstance(normalized.body[0], ast.Expr) + and isinstance(normalized.body[0].value, ast.Constant) + and isinstance(normalized.body[0].value.value, str) + ): + normalized.body.pop(0) + ast.fix_missing_locations(normalized) + return [ + " ".join(line.strip().split()) + for line in ast.unparse(normalized).splitlines() + if line.strip() + ] + + +@pytest.mark.parametrize( + ("local_name", "local_path", "upstream_path", "upstream_name"), + SOURCE_PAIRS, + ids=[local_name for local_name, _, _, _ in SOURCE_PAIRS], +) +def test_e1_functions_are_independently_implemented( + local_name: str, + local_path: Path, + upstream_path: Path, + upstream_name: str, +) -> None: + assert upstream_path.is_file(), f"pinned E1 source is missing: {upstream_path}" + assert local_path.is_file(), f"local E1 source is missing: {local_path}" + local_lines = _normalized_ast_lines(_function(local_path, local_name)) + upstream_lines = _normalized_ast_lines(_function(upstream_path, upstream_name)) + similarity = SequenceMatcher( + None, + local_lines, + upstream_lines, + autojunk=False, + ).ratio() + assert similarity < MAX_FUNCTION_SIMILARITY, ( + f"{local_name} has normalized AST similarity {similarity:.3f} to " + f"{upstream_path.relative_to(ROOT)}::{upstream_name}" + ) diff --git a/tests/release/test_esm3_source_independence.py b/tests/release/test_esm3_source_independence.py new file mode 100644 index 0000000..d2a546d --- /dev/null +++ b/tests/release/test_esm3_source_independence.py @@ -0,0 +1,133 @@ +"""Fail closed when ESM3 runtime functions overlap the Biohub parity oracle.""" + +from __future__ import annotations + +import ast +import copy +import pytest +from difflib import SequenceMatcher +from pathlib import Path + +from fastplms.models.esm3.modeling_esm3 import FastESM3PreTrainedModel + + +ROOT = Path(__file__).resolve().parents[2] +LOCAL_MODEL = ROOT / "src/fastplms/models/esm3/modeling_esm3.py" +BIOHUB = ROOT / "vendor/upstream/biohub-esm/esm" +MAX_FUNCTION_SIMILARITY = 0.75 +SOURCE_PAIRS = ( + ( + "EncodeInputs.__init__", + BIOHUB / "models/esm3.py", + "EncodeInputs.__init__", + ), + ( + "GeometricReasoningOriginalImpl.__init__", + BIOHUB / "layers/geom_attention.py", + "GeometricReasoningOriginalImpl.__init__", + ), + ( + "UnifiedTransformerBlock.forward", + BIOHUB / "layers/blocks.py", + "UnifiedTransformerBlock.forward", + ), +) + + +def _function(path: Path, qualified_name: str) -> ast.FunctionDef: + body: list[ast.stmt] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)).body + selected: ast.AST | None = None + for part in qualified_name.split("."): + selected = next( + ( + node + for node in body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name == part + ), + None, + ) + assert selected is not None, f"{qualified_name!r} is absent from {path}" + body = selected.body + assert isinstance(selected, ast.FunctionDef) + return selected + + +def _normalized_ast_lines(node: ast.FunctionDef) -> list[str]: + normalized = copy.deepcopy(node) + normalized.name = "function" + normalized.decorator_list = [] + normalized.returns = None + for argument in ( + *normalized.args.posonlyargs, + *normalized.args.args, + *normalized.args.kwonlyargs, + ): + argument.annotation = None + if normalized.args.vararg is not None: + normalized.args.vararg.annotation = None + if normalized.args.kwarg is not None: + normalized.args.kwarg.annotation = None + if ( + normalized.body + and isinstance(normalized.body[0], ast.Expr) + and isinstance(normalized.body[0].value, ast.Constant) + and isinstance(normalized.body[0].value.value, str) + ): + normalized.body.pop(0) + ast.fix_missing_locations(normalized) + return [ + " ".join(line.strip().split()) + for line in ast.unparse(normalized).splitlines() + if line.strip() + ] + + +@pytest.mark.parametrize( + ("local_name", "upstream_path", "upstream_name"), + SOURCE_PAIRS, + ids=[local_name for local_name, _, _ in SOURCE_PAIRS], +) +def test_esm3_functions_are_independently_implemented( + local_name: str, + upstream_path: Path, + upstream_name: str, +) -> None: + assert upstream_path.is_file(), f"pinned Biohub source is missing: {upstream_path}" + local_lines = _normalized_ast_lines(_function(LOCAL_MODEL, local_name)) + upstream_lines = _normalized_ast_lines(_function(upstream_path, upstream_name)) + similarity = SequenceMatcher( + None, + local_lines, + upstream_lines, + autojunk=False, + ).ratio() + assert similarity < MAX_FUNCTION_SIMILARITY, ( + f"{local_name} has normalized AST similarity {similarity:.3f} to " + f"{upstream_path.relative_to(ROOT)}::{upstream_name}" + ) + + +def test_esm3_source_does_not_import_upstream_packages() -> None: + tree = ast.parse(LOCAL_MODEL.read_text(encoding="utf-8"), filename=str(LOCAL_MODEL)) + imported_roots = { + alias.name.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + } + imported_roots.update( + node.module.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module + ) + assert imported_roots.isdisjoint({"esm", "vendor"}) + + +def test_esm3_rejects_unavailable_flash_kernels() -> None: + assert FastESM3PreTrainedModel._supports_flash_attn_2 is False + assert FastESM3PreTrainedModel._supports_flash_attn_3 is False + assert FastESM3PreTrainedModel._fastplms_attention_implementations == ( + "eager", + "sdpa", + "flex_attention", + ) diff --git a/tests/release/test_esmc_report_ingestion.py b/tests/release/test_esmc_report_ingestion.py new file mode 100644 index 0000000..190268e --- /dev/null +++ b/tests/release/test_esmc_report_ingestion.py @@ -0,0 +1,647 @@ +"""Fail-closed schema-v3 ESMC documentation evidence contracts.""" + +from __future__ import annotations + +import copy +import json +import os +import shutil +import subprocess +import sys +import pytest +from collections.abc import Callable, Mapping +from pathlib import Path + +from fastplms.registry import ModelRegistry, ModelSpec, get_model_registry +from tests.unit.test_biohub_reference_lock import _reference_environment_payload +from tools.artifacts import generate_docs +from tools.artifacts.generate_docs import ( + ESMC_BACKENDS, + ESMC_MODEL_IDS, + ESMC_PANEL_KINDS, + EsmcReportError, + EsmcReportSet, + EsmcRuntimeIdentity, + load_esmc_report_set, + render_capability_evidence, + render_model_card, +) + + +ROOT = Path(__file__).resolve().parents[2] +REGISTRY = get_model_registry() +RUNTIME_IDENTITY = EsmcRuntimeIdentity( + runtime_revision=f"source-tree-sha256:{'1' * 64}", + source_tree_sha256="1" * 64, + runtime_bundle_sha256="2" * 64, +) +LOCKED_REFERENCE_ENVIRONMENT = _reference_environment_payload() +LOCKED_RUNTIME = LOCKED_REFERENCE_ENVIRONMENT["runtime"] +if not isinstance(LOCKED_RUNTIME, dict) or not isinstance(LOCKED_RUNTIME.get("gpu"), dict): + raise RuntimeError("Synthetic locked Biohub environment fixture is malformed") +GPU = copy.deepcopy(LOCKED_RUNTIME["gpu"]) +REFERENCE_SOURCES: dict[str, dict[str, object]] = { + "biohub-esm": { + "attestation_sha256": "a" * 64, + "file_count": 412, + "import_file": "esm/__init__.py", + "import_name": "esm", + "import_root": "esm", + "package_version": "3.3.0", + "schema_version": 1, + "source_revision": "82ee35553d39169d678f784c8d3f8712ffd7d2c4", + "tree_sha256": "c5489f1fc58de200978803de2c38e1a78f769cb183a2ee90be833f0f4a0212e8", + }, + "biohub-transformers": { + "attestation_sha256": "b" * 64, + "file_count": 5218, + "import_file": "src/transformers/__init__.py", + "import_name": "transformers", + "import_root": "src/transformers", + "package_version": "4.57.6", + "schema_version": 1, + "source_revision": "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf", + "tree_sha256": "28b910cc18b821870db2fb6d1c50376c2d14287ae18485080699e03fa4ba4f43", + }, +} + + +def _candidate_environment() -> dict[str, object]: + return { + "python": LOCKED_RUNTIME["python_version"], + "torch": LOCKED_RUNTIME["torch"], + "transformers": "5.13.0", + "cuda_runtime": LOCKED_RUNTIME["cuda_runtime"], + "cuda_driver": LOCKED_RUNTIME["cuda_driver"], + "gpu": copy.deepcopy(GPU), + "packages": { + "fastplms": "1.0.0", + "huggingface-hub": "1.4.0", + "kernels": "0.12.2", + "tokenizers": "0.22.2", + "transformer-engine": None, + "transformer-engine-torch": None, + }, + } + + +def _reference_environment() -> dict[str, object]: + return { + "cuda_device": GPU["name"], + "cuda_device_capability": copy.deepcopy(GPU["capability"]), + "cuda_total_memory": GPU["total_memory_bytes"], + "cuda_runtime": LOCKED_RUNTIME["cuda_runtime"], + "packages": json.dumps( + { + "python": LOCKED_RUNTIME["python_version"], + "torch": LOCKED_RUNTIME["torch"], + }, + separators=(",", ":"), + sort_keys=True, + ), + "python": LOCKED_RUNTIME["python_version"], + "torch": LOCKED_RUNTIME["torch"], + } + + +def _kernel(registry: ModelRegistry, backend: str) -> dict[str, object]: + kernel = registry.attention_kernels.get(backend) + if kernel is None: + return { + "implementation": backend, + "provider": "torch", + "torch_version": LOCKED_RUNTIME["torch"], + } + return { + "implementation": backend, + "provider": "huggingface_kernels", + "repository": kernel.repository, + "revision": kernel.revision, + "version": kernel.version, + "expected_variant": kernel.expected_variant, + "supported_dtypes": list(kernel.dtypes), + "kernels_package_version": "0.12.2", + } + + +def _tensor_metrics(context: str, base: float) -> list[dict[str, object]]: + result = [] + for output, layer_index, offset in ( + ("hidden_state", 0, 0.0), + ("last_hidden_state", None, 0.00001), + ("logits", None, 0.00002), + ): + value = base + offset + result.append( + { + "context": context, + "output": output, + "layer_index": layer_index, + "relative_l2": value, + "relative_q999": value * 2, + "residue_cosine_p01": 1.0 - value, + "pooled_cosine_min": 1.0 - value / 2, + } + ) + return result + + +def _logits_metrics(base: float) -> dict[str, float]: + return { + "confident_top1_agreement": 1.0 - base, + "mean_jsd": base / 10, + } + + +def _report( + registry: ModelRegistry, + spec: ModelSpec, + backend: str, + panel: Mapping[str, object], +) -> dict[str, object]: + panel_kind = str(panel["kind"]) + is_unavailable = backend in generate_docs.ESMC_UNAVAILABLE_BACKENDS + model_offset = ESMC_MODEL_IDS.index(spec.id) * 0.001 + backend_offset = ESMC_BACKENDS.index(backend) * 0.0001 + panel_offset = ESMC_PANEL_KINDS.index(panel_kind) * 0.00001 + base = 0.001 + model_offset + backend_offset + panel_offset + context = f"{spec.id}:bf16:{backend}:{panel_kind}" + panel_cases = panel["cases"] + assert isinstance(panel_cases, list) + measured_cases = [] + for index, panel_case in enumerate(panel_cases): + assert isinstance(panel_case, Mapping) + case_base = base + index * 0.000001 + case_id = str(panel_case["case_id"]) + measured_cases.append( + { + **panel_case, + "tensor_metrics": _tensor_metrics(f"{context}:case={case_id}", case_base), + "logits_metrics": _logits_metrics(case_base), + } + ) + release_modes = { + "sdpa": "exact", + "eager": "strict_numeric", + "flex_attention": "diagnostic_with_catastrophe_gate", + } + payload: dict[str, object] = { + "schema_version": 3, + "model_id": spec.id, + "candidate": { + "repo_id": spec.fast.repo_id, + "manifest_revision": spec.fast.revision, + "resolved_commit": spec.fast.revision, + "checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "checkpoint_revision": spec.artifact_checkpoint.revision, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": RUNTIME_IDENTITY.runtime_revision, + "source_tree_sha256": RUNTIME_IDENTITY.source_tree_sha256, + "runtime_bundle_sha256": RUNTIME_IDENTITY.runtime_bundle_sha256, + }, + "reference": { + "repo_id": spec.official.repo_id, + "revision": spec.official.revision, + "state_transform": spec.family.state_transform, + "environment": _reference_environment(), + "reference_environment": copy.deepcopy(LOCKED_REFERENCE_ENVIRONMENT), + "reference_sources": copy.deepcopy(REFERENCE_SOURCES), + }, + "record_status": "unavailable" if is_unavailable else "measured", + "unavailability": ( + generate_docs._esmc_unavailability_identity(backend, LOCKED_REFERENCE_ENVIRONMENT) + if is_unavailable + else None + ), + "configured_backend": backend, + "effective_backend": None if is_unavailable else backend, + "dtype": "bfloat16", + "panel": copy.deepcopy(panel), + "environment": _candidate_environment(), + "kernel": _kernel(registry, backend), + "panel_tensor_metrics": None if is_unavailable else _tensor_metrics(context, base), + "panel_logits_metrics": None if is_unavailable else _logits_metrics(base), + "cases": copy.deepcopy(panel_cases) if is_unavailable else measured_cases, + "published_band_violations": [], + "catastrophic_gate": "not_run" if is_unavailable else "passed", + "release_gate": ( + {"mode": "availability", "status": "unavailable"} + if is_unavailable + else {"mode": release_modes[backend], "status": "passed"} + ), + } + payload["report_sha256"] = generate_docs._esmc_report_sha256(payload) + return payload + + +@pytest.fixture +def complete_report_root(tmp_path: Path) -> Path: + report_root = tmp_path / "reports" + report_root.mkdir() + panels = generate_docs._expected_esmc_panels(ROOT) + for model_id in ESMC_MODEL_IDS: + spec = REGISTRY[model_id] + for backend in ESMC_BACKENDS: + for panel_kind in ESMC_PANEL_KINDS: + payload = _report(REGISTRY, spec, backend, panels[panel_kind]) + path = report_root / f"{model_id}-{backend}-{panel_kind}.json" + path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return report_root + + +def _rewrite_report( + path: Path, + mutate: Callable[[dict[str, object]], None], + *, + rehash: bool = True, +) -> None: + payload = json.loads(path.read_text(encoding="utf-8")) + mutate(payload) + if rehash: + payload["report_sha256"] = generate_docs._esmc_report_sha256(payload) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _load(report_root: Path) -> EsmcReportSet: + return load_esmc_report_set( + report_root, + REGISTRY, + source_root=ROOT, + expected_runtime_identity=RUNTIME_IDENTITY, + ) + + +@pytest.mark.parametrize( + ("gpu_name", "architecture", "memory"), + ( + ("NVIDIA H100 80GB HBM3", "x86_64", 80 * 1024**3), + ("NVIDIA H200", "x86_64", 141 * 1024**3), + ("NVIDIA GH200 480GB", "aarch64", 480_000_000_000), + ), +) +def test_generator_dynamic_environment_schema_is_hardware_neutral( + gpu_name: str, + architecture: str, + memory: int, +) -> None: + candidate = _candidate_environment() + gpu = { + "name": gpu_name, + "capability": [9, 0], + "total_memory_bytes": memory, + } + candidate["gpu"] = copy.deepcopy(gpu) + dynamic_reference = _reference_environment() + dynamic_reference["cuda_device"] = gpu_name + dynamic_reference["cuda_device_capability"] = [9, 0] + dynamic_reference["cuda_total_memory"] = memory + locked_reference = { + "runtime": { + "operating_system": "linux", + "architecture": architecture, + "python_version": candidate["python"], + "torch": candidate["torch"], + "cuda_runtime": candidate["cuda_runtime"], + "cuda_driver": candidate["cuda_driver"], + "gpu": copy.deepcopy(gpu), + } + } + + validated_candidate = generate_docs._validate_esmc_candidate_environment(candidate) + validated_reference = generate_docs._validate_esmc_reference_environment(dynamic_reference) + generate_docs._validate_esmc_environment_binding( + validated_candidate, + validated_reference, + locked_reference, + ) + unavailable = generate_docs._esmc_unavailability_identity("flash_attention_2", locked_reference) + assert unavailable["platform"] == f"linux/{architecture}" + assert unavailable["accelerator"] == f"{gpu_name}/SM90" + + +def test_complete_esmc_report_set_renders_checkpoint_specific_measurements( + complete_report_root: Path, +) -> None: + evidence = _load(complete_report_root) + + assert len(evidence.reports) == 30 + assert sum(report["record_status"] == "measured" for report in evidence.reports) == 18 + assert sum(report["record_status"] == "unavailable" for report in evidence.reports) == 12 + assert len(evidence.select("esmc_small")) == 10 + small = render_model_card(REGISTRY["esmc_small"], esmc_evidence=evidence) + large = render_model_card(REGISTRY["esmc_large"], esmc_evidence=evidence) + six_b = render_model_card(REGISTRY["esmc_6b"], esmc_evidence=evidence) + manifest = render_capability_evidence(REGISTRY, esmc_evidence=evidence) + + assert "validated complete set (30/30 records)" in manifest + assert "NVIDIA GH200 480GB" in small + assert REFERENCE_SOURCES["biohub-esm"]["attestation_sha256"] in small + assert REFERENCE_SOURCES["biohub-transformers"]["attestation_sha256"] in small + assert "Locked-platform unavailable backends" in small + assert "no validated artifact" in small + assert "Per-case distributions" in small + assert "0.001 to 0.00102" in small + assert "0.002 to 0.00202" in large + assert "0.003 to 0.00302" in six_b + assert "ESMC-6B Flex Attention exceeds" not in small + assert "ESMC-6B Flex Attention exceeds" not in large + + +def test_default_generation_stays_explicitly_pending_and_ignores_environment( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("FASTPLMS_DIAGNOSTIC_REPORTS", str(tmp_path / "partial")) + + card = render_model_card(REGISTRY["esmc_small"]) + manifest = render_capability_evidence(REGISTRY) + + assert "Pending release measurement" in card + assert "Status: pending" in manifest + assert "ESMC-6B Flex Attention exceeds" not in card + + +def test_missing_esmc_report_fails_closed(complete_report_root: Path) -> None: + (complete_report_root / "esmc_6b-flash_attention_3-real_biological_holdout.json").unlink() + + with pytest.raises(EsmcReportError, match="exactly 30 records"): + _load(complete_report_root) + + +@pytest.mark.parametrize( + ("injected", "message"), + ( + ('"schema_version": 3,', "duplicate key"), + ('"nonfinite": NaN,', "non-finite constant"), + ), +) +def test_non_strict_esmc_json_fails_before_schema_validation( + complete_report_root: Path, + injected: str, + message: str, +) -> None: + path = complete_report_root / "esmc_small-eager-generated_kernel_boundary.json" + encoded = path.read_text(encoding="utf-8") + path.write_text("{\n " + injected + encoded[1:], encoding="utf-8") + + with pytest.raises(EsmcReportError, match=message): + _load(complete_report_root) + + +def test_esmc_schema_validation_remains_fail_closed_under_python_optimized_mode( + complete_report_root: Path, +) -> None: + path = complete_report_root / "esmc_small-sdpa-generated_kernel_boundary.json" + _rewrite_report( + path, + lambda report: report.__setitem__("panel_tensor_metrics", {"invalid": "mapping"}), + ) + script = f""" +from pathlib import Path +from fastplms.registry import get_model_registry +from tools.artifacts.generate_docs import EsmcReportError, EsmcRuntimeIdentity, load_esmc_report_set + +try: + load_esmc_report_set( + Path({str(complete_report_root)!r}), + get_model_registry(), + source_root=Path({str(ROOT)!r}), + expected_runtime_identity=EsmcRuntimeIdentity( + runtime_revision={RUNTIME_IDENTITY.runtime_revision!r}, + source_tree_sha256={RUNTIME_IDENTITY.source_tree_sha256!r}, + runtime_bundle_sha256={RUNTIME_IDENTITY.runtime_bundle_sha256!r}, + ), + ) +except EsmcReportError as error: + if "tensor metrics are missing" not in str(error): + raise +else: + raise SystemExit("optimized-mode validation accepted malformed ESMC evidence") +""" + environment = os.environ.copy() + import_roots = (str(ROOT), str(ROOT / "src")) + environment["PYTHONPATH"] = os.pathsep.join( + (*import_roots, environment.get("PYTHONPATH", "")) + ).rstrip(os.pathsep) + result = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=ROOT, + env=environment, + check=False, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr or result.stdout + + +@pytest.mark.parametrize( + ("mutate", "rehash", "message"), + ( + ( + lambda report: report["candidate"].__setitem__("weights_revision", "0" * 40), + True, + "weights_revision", + ), + ( + lambda report: report["candidate"].__setitem__("source_tree_sha256", "3" * 64), + True, + "source_tree_sha256", + ), + ( + lambda report: report["candidate"].__setitem__("runtime_revision", "4" * 40), + True, + "runtime_revision", + ), + ( + lambda report: report.__setitem__("effective_backend", "sdpa"), + True, + "backend identity", + ), + ( + lambda report: report.__setitem__("dtype", "float32"), + True, + "dtype identity", + ), + ( + lambda report: report["panel"].__setitem__("definition_sha256", "5" * 64), + True, + "immutable definition", + ), + ( + lambda report: report.__setitem__("unexpected", True), + True, + "schema v3", + ), + ( + lambda report: report["reference"]["reference_sources"][ + "biohub-transformers" + ].__setitem__("tree_sha256", "6" * 64), + True, + "reference source biohub-transformers tree_sha256", + ), + ( + lambda report: report["reference"]["reference_environment"]["runtime"][ + "gpu" + ].__setitem__("name", "forged accelerator"), + True, + "locked reference environment is invalid", + ), + ( + lambda report: report["panel_logits_metrics"].__setitem__("mean_jsd", 0.06), + True, + "catastrophe gate", + ), + ( + lambda report: report.__setitem__("catastrophic_gate", "failed"), + True, + "catastrophe gate", + ), + ( + lambda report: report.__setitem__("report_sha256", "f" * 64), + False, + "self-digest", + ), + ), +) +def test_stale_malformed_or_tampered_esmc_report_fails_closed( + complete_report_root: Path, + mutate: Callable[[dict[str, object]], None], + rehash: bool, + message: str, +) -> None: + path = complete_report_root / "esmc_small-flex_attention-generated_kernel_boundary.json" + _rewrite_report(path, mutate, rehash=rehash) + + with pytest.raises(EsmcReportError, match=message): + _load(complete_report_root) + + +@pytest.mark.parametrize( + ("mutate", "message"), + ( + ( + lambda report: report.__setitem__("effective_backend", "flash_attention_3"), + "must not claim effective dispatch", + ), + ( + lambda report: report.__setitem__("panel_tensor_metrics", []), + "must not contain measurements", + ), + ( + lambda report: report["unavailability"].__setitem__( + "dispatch_contract", "silent_fallback" + ), + "unavailability identity", + ), + ), +) +def test_flash_unavailability_records_fail_closed_on_false_execution_claims( + complete_report_root: Path, + mutate: Callable[[dict[str, object]], None], + message: str, +) -> None: + path = complete_report_root / "esmc_small-flash_attention_3-generated_kernel_boundary.json" + _rewrite_report(path, mutate) + + with pytest.raises(EsmcReportError, match=message): + _load(complete_report_root) + + +def test_cross_device_esmc_report_set_fails_closed(complete_report_root: Path) -> None: + path = complete_report_root / "esmc_large-sdpa-real_biological_holdout.json" + + def mutate(report: dict[str, object]) -> None: + environment = report["environment"] + reference = report["reference"] + assert isinstance(environment, dict) + assert isinstance(reference, dict) + gpu = environment["gpu"] + reference_environment = reference["environment"] + assert isinstance(gpu, dict) + assert isinstance(reference_environment, dict) + gpu["name"] = "NVIDIA GH200 96GB" + gpu["total_memory_bytes"] = 95 * 1024**3 + reference_environment["cuda_device"] = gpu["name"] + reference_environment["cuda_total_memory"] = gpu["total_memory_bytes"] + + _rewrite_report(path, mutate) + + with pytest.raises(EsmcReportError, match="locked reference runtime"): + _load(complete_report_root) + + +def test_cli_explicit_report_root_renders_only_after_complete_validation( + complete_report_root: Path, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source_root = tmp_path / "rendered" + fixture = source_root / "tests" / "parity" / "fixtures" / "esmc_biological_holdout.json" + fixture.parent.mkdir(parents=True) + shutil.copyfile( + ROOT / "tests" / "parity" / "fixtures" / "esmc_biological_holdout.json", + fixture, + ) + constraints = source_root / "docker" / "constraints" + constraints.mkdir(parents=True) + for name in ( + "biohub-esm-source.json", + "biohub-transformers-source.json", + "biohub-reference-lock.json", + "biohub-reference.in", + "biohub-reference.lock.txt", + "biohub-biotraj-build.in", + "biohub-biotraj-build.lock.txt", + ): + shutil.copyfile(ROOT / "docker" / "constraints" / name, constraints / name) + shutil.copyfile( + ROOT / "docker" / "biohub-reference-lock.Dockerfile", + source_root / "docker" / "biohub-reference-lock.Dockerfile", + ) + monkeypatch.setattr( + generate_docs, + "_esmc_runtime_identity_from_source", + lambda source, registry: RUNTIME_IDENTITY, + ) + + result = generate_docs.main( + ( + "--source-root", + str(source_root), + "--esmc-report-root", + str(complete_report_root), + ) + ) + + assert result == 0 + card = (source_root / "model_cards" / "esmc_small.md").read_text(encoding="utf-8") + assert "NVIDIA GH200 480GB" in card + assert "Per-case distributions" in card + + +@pytest.mark.parametrize("use_environment", (False, True)) +def test_cli_release_evidence_options_fail_closed_on_missing_set( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + use_environment: bool, +) -> None: + missing = tmp_path / "missing" + monkeypatch.setattr( + generate_docs, + "_esmc_runtime_identity_from_source", + lambda source, registry: RUNTIME_IDENTITY, + ) + arguments = ["--source-root", str(ROOT)] + if use_environment: + monkeypatch.setenv("FASTPLMS_DIAGNOSTIC_REPORTS", str(missing)) + arguments.append("--require-esmc-release-evidence") + else: + arguments.extend(("--esmc-report-root", str(missing))) + + assert generate_docs.main(arguments) == 1 + assert "invalid ESMC release evidence" in capsys.readouterr().out diff --git a/tests/release/test_esmfold2_source_independence.py b/tests/release/test_esmfold2_source_independence.py new file mode 100644 index 0000000..5f95e02 --- /dev/null +++ b/tests/release/test_esmfold2_source_independence.py @@ -0,0 +1,128 @@ +"""Fail closed when ESMFold2 runtime source overlaps its parity oracles.""" + +from __future__ import annotations + +import ast +import pytest +from difflib import SequenceMatcher +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +RUNTIME = ROOT / "src/fastplms/models/esmfold2" +MAX_LINE_SIMILARITY = 0.75 +BIOHUB_ESM = "vendor/upstream/biohub-esm/esm" +BIOHUB_TRANSFORMERS = "vendor/upstream/biohub-transformers/src/transformers/models/esmfold2" + +# Every runtime module must be classified here or in ORIGINAL_RUNTIME_MODULES. This +# inventory makes a new derivative module a release failure instead of silently +# excluding it from the source-boundary check. +SOURCE_COUNTERPARTS = { + "__init__.py": f"{BIOHUB_TRANSFORMERS}/__init__.py", + "configuration_esmfold2.py": f"{BIOHUB_TRANSFORMERS}/configuration_esmfold2.py", + "esmfold2_affine3d.py": f"{BIOHUB_ESM}/utils/structure/affine3d.py", + "esmfold2_aligner.py": f"{BIOHUB_ESM}/utils/structure/aligner.py", + "esmfold2_atom_indexer.py": f"{BIOHUB_ESM}/utils/structure/atom_indexer.py", + "esmfold2_conformers.py": f"{BIOHUB_ESM}/models/esmfold2/conformers.py", + "esmfold2_constants.py": f"{BIOHUB_ESM}/models/esmfold2/constants.py", + "esmfold2_constants_esm3.py": f"{BIOHUB_ESM}/utils/constants/esm3.py", + "esmfold2_input_builder.py": f"{BIOHUB_ESM}/utils/structure/input_builder.py", + "esmfold2_metrics.py": f"{BIOHUB_ESM}/utils/structure/metrics.py", + "esmfold2_misc.py": f"{BIOHUB_ESM}/utils/misc.py", + "esmfold2_mmcif_parsing.py": f"{BIOHUB_ESM}/utils/structure/mmcif_parsing.py", + "esmfold2_molecular_complex.py": f"{BIOHUB_ESM}/utils/structure/molecular_complex.py", + "esmfold2_msa.py": f"{BIOHUB_ESM}/utils/msa/msa.py", + "esmfold2_msa_filter_sequences.py": f"{BIOHUB_ESM}/utils/msa/filter_sequences.py", + "esmfold2_normalize_coordinates.py": (f"{BIOHUB_ESM}/utils/structure/normalize_coordinates.py"), + "esmfold2_output.py": f"{BIOHUB_ESM}/models/esmfold2/output.py", + "esmfold2_paired_msa.py": f"{BIOHUB_ESM}/models/esmfold2/paired_msa.py", + "esmfold2_parsing.py": f"{BIOHUB_ESM}/utils/parsing.py", + "esmfold2_predicted_aligned_error.py": ( + f"{BIOHUB_ESM}/utils/structure/predicted_aligned_error.py" + ), + "esmfold2_prepare_input.py": f"{BIOHUB_ESM}/models/esmfold2/prepare_input.py", + "esmfold2_processor.py": f"{BIOHUB_ESM}/models/esmfold2/processor.py", + "esmfold2_protein_chain.py": f"{BIOHUB_ESM}/utils/structure/protein_chain.py", + "esmfold2_protein_complex.py": f"{BIOHUB_ESM}/utils/structure/protein_complex.py", + "esmfold2_protein_structure.py": (f"{BIOHUB_ESM}/utils/structure/protein_structure.py"), + "esmfold2_residue_constants.py": f"{BIOHUB_ESM}/utils/residue_constants.py", + "esmfold2_sequential_dataclass.py": f"{BIOHUB_ESM}/utils/sequential_dataclass.py", + "esmfold2_system.py": f"{BIOHUB_ESM}/utils/system.py", + "esmfold2_types.py": f"{BIOHUB_ESM}/models/esmfold2/types.py", + "esmfold2_utils_types.py": f"{BIOHUB_ESM}/utils/types.py", + "modeling_esmfold2.py": f"{BIOHUB_TRANSFORMERS}/modeling_esmfold2.py", + "modeling_esmfold2_common.py": (f"{BIOHUB_TRANSFORMERS}/modeling_esmfold2_common.py"), + "modeling_esmfold2_experimental.py": ( + f"{BIOHUB_TRANSFORMERS}/modeling_esmfold2_experimental.py" + ), + "protein_utils.py": f"{BIOHUB_TRANSFORMERS}/protein_utils.py", + "reproducibility.py": f"{BIOHUB_ESM}/models/esmfold2/processor.py", +} +ORIGINAL_RUNTIME_MODULES = frozenset({"attention.py", "embedding.py"}) + + +def _meaningful_lines(text: str) -> list[str]: + return [ + " ".join(line.strip().split()) + for line in text.splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + +def test_esmfold2_runtime_source_inventory_is_complete() -> None: + runtime_modules = {path.name for path in RUNTIME.glob("*.py")} + classified_modules = {*SOURCE_COUNTERPARTS, *ORIGINAL_RUNTIME_MODULES} + assert runtime_modules == classified_modules, ( + "classify every ESMFold2 runtime module as original or bind it to its " + "pinned upstream source counterpart; " + f"missing={sorted(runtime_modules - classified_modules)}, " + f"stale={sorted(classified_modules - runtime_modules)}" + ) + + +@pytest.mark.parametrize( + ("runtime_name", "upstream_relative"), + SOURCE_COUNTERPARTS.items(), + ids=SOURCE_COUNTERPARTS, +) +def test_esmfold2_source_is_independently_organized( + runtime_name: str, upstream_relative: str +) -> None: + runtime_path = RUNTIME / runtime_name + upstream_path = ROOT / upstream_relative + assert upstream_path.is_file(), f"pinned source counterpart is missing: {upstream_relative}" + runtime_text = runtime_path.read_text(encoding="utf-8") + upstream_text = upstream_path.read_text(encoding="utf-8") + assert runtime_text.encode() != upstream_text.encode() + similarity = SequenceMatcher( + None, + _meaningful_lines(runtime_text), + _meaningful_lines(upstream_text), + autojunk=False, + ).ratio() + assert similarity < MAX_LINE_SIMILARITY, ( + f"{runtime_name} has line similarity {similarity:.3f} to {upstream_relative}; " + "reimplement the public behavior instead of copying the parity oracle" + ) + + +def test_esmfold2_runtime_modules_do_not_import_upstream_packages() -> None: + forbidden_roots = {"esm", "vendor"} + for runtime_path in sorted(RUNTIME.glob("*.py")): + source = runtime_path.read_text(encoding="utf-8") + tree = ast.parse(source, filename=runtime_path.name) + imported_roots = { + alias.name.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + } + imported_roots.update( + node.module.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module + ) + assert imported_roots.isdisjoint(forbidden_roots), ( + f"{runtime_path.name} imports an upstream package: " + f"{sorted(imported_roots & forbidden_roots)}" + ) diff --git a/tests/release/test_esmplusplus_source_independence.py b/tests/release/test_esmplusplus_source_independence.py new file mode 100644 index 0000000..4135fff --- /dev/null +++ b/tests/release/test_esmplusplus_source_independence.py @@ -0,0 +1,221 @@ +"""Fail closed when ESM++ runtime functions overlap the Biohub parity oracle.""" + +from __future__ import annotations + +import ast +import copy +import importlib.util +import pytest +import torch +from difflib import SequenceMatcher +from pathlib import Path +from types import ModuleType + +from fastplms.models.esm_plusplus.modeling_esm_plusplus import PreTrainedESMplusplusModel +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + RotaryEmbedding as FastRotaryEmbedding, +) + + +ROOT = Path(__file__).resolve().parents[2] +LOCAL_MODEL = ROOT / "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" +UPSTREAM_ROTARY = ROOT / "vendor/upstream/biohub-esm/esm/layers/rotary.py" +UPSTREAM_TOKENIZER = ROOT / "vendor/upstream/biohub-esm/esm/tokenization/sequence_tokenizer.py" +MAX_FUNCTION_SIMILARITY = 0.75 + +# These functions implement the same public contracts, but the repository source +# must remain independently maintained. Function-level comparisons prevent +# unrelated model code from diluting a copied implementation's similarity. +SOURCE_PAIRS = ( + ( + "EsmSequenceTokenizer.__init__", + UPSTREAM_TOKENIZER, + "EsmSequenceTokenizer.__init__", + ), + ( + "RotaryEmbedding.__init__", + UPSTREAM_ROTARY, + "RotaryEmbedding.__init__", + ), + ( + "RotaryEmbedding._update_cos_sin_cache", + UPSTREAM_ROTARY, + "RotaryEmbedding._update_cos_sin_cache", + ), + ( + "apply_rotary_emb_torch", + UPSTREAM_ROTARY, + "apply_rotary_emb_torch", + ), +) + + +def _function(path: Path, qualified_name: str) -> ast.FunctionDef: + body: list[ast.stmt] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)).body + selected: ast.AST | None = None + for part in qualified_name.split("."): + selected = next( + ( + node + for node in body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name == part + ), + None, + ) + assert selected is not None, f"{qualified_name!r} is absent from {path}" + body = selected.body + assert isinstance(selected, ast.FunctionDef) + return selected + + +def _normalized_ast_lines(node: ast.FunctionDef) -> list[str]: + normalized = copy.deepcopy(node) + normalized.name = "function" + normalized.decorator_list = [] + normalized.returns = None + for argument in ( + *normalized.args.posonlyargs, + *normalized.args.args, + *normalized.args.kwonlyargs, + ): + argument.annotation = None + if normalized.args.vararg is not None: + normalized.args.vararg.annotation = None + if normalized.args.kwarg is not None: + normalized.args.kwarg.annotation = None + if ( + normalized.body + and isinstance(normalized.body[0], ast.Expr) + and isinstance(normalized.body[0].value, ast.Constant) + and isinstance(normalized.body[0].value.value, str) + ): + normalized.body.pop(0) + ast.fix_missing_locations(normalized) + return [ + " ".join(line.strip().split()) + for line in ast.unparse(normalized).splitlines() + if line.strip() + ] + + +def _load_upstream_rotary() -> ModuleType: + spec = importlib.util.spec_from_file_location( + "_fastplms_test_biohub_rotary", + UPSTREAM_ROTARY, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + ("local_name", "upstream_path", "upstream_name"), + SOURCE_PAIRS, + ids=[local_name for local_name, _, _ in SOURCE_PAIRS], +) +def test_esmplusplus_functions_are_independently_implemented( + local_name: str, + upstream_path: Path, + upstream_name: str, +) -> None: + assert upstream_path.is_file(), f"pinned Biohub source is missing: {upstream_path}" + local_lines = _normalized_ast_lines(_function(LOCAL_MODEL, local_name)) + upstream_lines = _normalized_ast_lines(_function(upstream_path, upstream_name)) + similarity = SequenceMatcher( + None, + local_lines, + upstream_lines, + autojunk=False, + ).ratio() + assert similarity < MAX_FUNCTION_SIMILARITY, ( + f"{local_name} has normalized AST similarity {similarity:.3f} to " + f"{upstream_path.relative_to(ROOT)}::{upstream_name}" + ) + + +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) +@pytest.mark.parametrize("interleaved", (False, True)) +def test_reimplemented_rotary_is_exact(dtype: torch.dtype, interleaved: bool) -> None: + upstream_class = _load_upstream_rotary().RotaryEmbedding + local = FastRotaryEmbedding(dim=8, interleaved=interleaved).eval() + upstream = upstream_class(dim=8, interleaved=interleaved).eval() + + generator = torch.Generator().manual_seed(13) + # q: (2, 17, 3, 8) + q = torch.randn((2, 17, 3, 8), generator=generator, dtype=dtype) + # k: (2, 17, 3, 8) + k = torch.randn((2, 17, 3, 8), generator=generator, dtype=dtype) + local_q, local_k = local(q, k) + upstream_q, upstream_k = upstream(q, k) + + assert torch.equal(local.inv_freq, upstream.inv_freq) + assert torch.equal(local._cos_cached, upstream._cos_cached) + assert torch.equal(local._sin_cached, upstream._sin_cached) + assert torch.equal(local_q, upstream_q) + assert torch.equal(local_k, upstream_k) + assert local.state_dict().keys() == upstream.state_dict().keys() + + +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) +def test_reimplemented_scaled_rotary_cache_is_exact(dtype: torch.dtype) -> None: + upstream_class = _load_upstream_rotary().RotaryEmbedding + local = FastRotaryEmbedding(dim=8, scale_base=512).eval() + upstream = upstream_class(dim=8, scale_base=512).eval() + + local._update_cos_sin_cache(19, device=torch.device("cpu"), dtype=dtype) + upstream._update_cos_sin_cache(19, device=torch.device("cpu"), dtype=dtype) + for name in ("_cos_cached", "_sin_cached", "_cos_k_cached", "_sin_k_cached"): + assert torch.equal(getattr(local, name), getattr(upstream, name)) + assert local.state_dict().keys() == upstream.state_dict().keys() + assert torch.equal(local.state_dict()["scale"], upstream.state_dict()["scale"]) + + +@pytest.mark.gpu +def test_reimplemented_rotary_matches_transformers_cuda_policy() -> None: + assert torch.cuda.is_available(), "ESM++ rotary parity requires CUDA" + upstream_class = _load_upstream_rotary().RotaryEmbedding + # local: (...) + local = FastRotaryEmbedding(dim=64).eval().to("cuda") + # upstream: (...) + upstream = upstream_class(dim=64).eval().to("cuda") + + # The original Biohub SDK migrates CPU-computed frequencies. The pinned + # Biohub Transformers oracle instead recomputes them on CUDA after a device + # move. Reproduce that public AutoModel policy on the independent upstream + # rotary implementation before comparing outputs. + # cpu_migrated: (...) + cpu_migrated = upstream.inv_freq.clone() + cuda_native = upstream._compute_inv_freq(torch.device("cuda")) + assert not torch.equal(cpu_migrated, cuda_native) + upstream.register_buffer("inv_freq", cuda_native, persistent=False) + upstream._seq_len_cached = 0 + upstream._cos_cached = None + upstream._sin_cached = None + + generator = torch.Generator(device="cuda").manual_seed(29) + # q: (3, 65, 8, 64) + q = torch.randn((3, 65, 8, 64), generator=generator, device="cuda", dtype=torch.bfloat16) + # k: (3, 65, 8, 64) + k = torch.randn((3, 65, 8, 64), generator=generator, device="cuda", dtype=torch.bfloat16) + local_q, local_k = local(q, k) + upstream_q, upstream_k = upstream(q, k) + + assert torch.equal(local.inv_freq, cuda_native) + assert torch.equal(local._cos_cached, upstream._cos_cached) + assert torch.equal(local._sin_cached, upstream._sin_cached) + assert torch.equal(local_q, upstream_q) + assert torch.equal(local_k, upstream_k) + + +def test_esmplusplus_advertises_only_pinned_flash_kernels() -> None: + assert PreTrainedESMplusplusModel._supports_flash_attn is True + assert PreTrainedESMplusplusModel._supports_flash_attn_2 is True + assert PreTrainedESMplusplusModel._supports_flash_attn_3 is True + assert PreTrainedESMplusplusModel._fastplms_attention_implementations == ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) diff --git a/tests/release/test_flash_source_policy.py b/tests/release/test_flash_source_policy.py new file mode 100644 index 0000000..6a69e1f --- /dev/null +++ b/tests/release/test_flash_source_policy.py @@ -0,0 +1,139 @@ +"""Repository-wide policy against source-built FlashAttention packages.""" + +from __future__ import annotations + +import re +from pathlib import Path + +from fastplms.registry import get_model_registry + + +ROOT = Path(__file__).resolve().parents[2] +_FLASH_PACKAGES = frozenset({"flash-attn", "flash_attn", "flashattention"}) +_FLASH_BACKENDS = frozenset({"flash_attention_2", "flash_attention_3"}) +_SOURCE_FLASH = re.compile(r"flash[-_]?attn|dao-ai(?:lab)?/flash-attention", re.IGNORECASE) +_INSTALL_OR_BUILD = re.compile( + r"(?:^|\s)(?:pip(?:3)?\s+install|python(?:3)?\s+-m\s+pip\s+install|" + r"uv(?:\s+pip)?\s+(?:add|install)|poetry\s+add|conda\s+install|mamba\s+install|" + r"git\s+clone|python(?:3)?\s+setup\.py|cmake(?:\s|$)|ninja(?:\s|$)|make(?:\s|$))", + re.IGNORECASE, +) + + +def _normalized_package(requirement: str) -> str: + name = re.split(r"[<>=!~;\[\s]", requirement.strip(), maxsplit=1)[0] + return re.sub(r"[-_.]+", "-", name).lower() + + +def _requirements(path: Path) -> list[str]: + return [ + line.strip() + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + + +def test_dependency_contract_contains_no_flash_attn_distribution() -> None: + requirement_root = ROOT / "requirements" + dependency_files = sorted(requirement_root.rglob("*.in")) + dependency_files.extend(sorted(requirement_root.rglob("*.txt"))) + requirements = [ + requirement + for path in dependency_files + for requirement in _requirements(path) + ] + + assert not { + requirement + for requirement in requirements + if _normalized_package(requirement) in _FLASH_PACKAGES + } + assert _requirements(requirement_root / "features" / "flash.in") == [ + "kernels>=0.15,<0.16" + ] + + +def test_no_docker_script_or_documentation_command_builds_source_flash_attn() -> None: + roots = ( + ROOT / "requirements", + ROOT / "docker", + ROOT / "tools", + ROOT / "examples", + ROOT / "benchmarks", + ROOT / "docs", + ) + files = [ROOT / "README.md"] + allowed_suffixes = { + ".bat", + ".cmd", + ".hcl", + ".in", + ".md", + ".ps1", + ".py", + ".rst", + ".sh", + ".txt", + ".yml", + ".yaml", + } + for root in roots: + files.extend( + path + for path in root.rglob("*") + if path.is_file() + and (path.suffix.lower() in allowed_suffixes or path.name == "Dockerfile") + ) + + violations: list[str] = [] + for path in sorted(set(files)): + lines = path.read_text(encoding="utf-8").splitlines() + for number, line in enumerate(lines, start=1): + stripped = line.strip() + if ( + path.suffix.lower() in {".in", ".txt"} + and stripped + and not stripped.startswith("#") + and _normalized_package(stripped) in _FLASH_PACKAGES + ): + violations.append(f"{path.relative_to(ROOT)}:{number}: {stripped}") + # Inspect a short logical-command window so shell and Docker line + # continuations cannot hide a source package or repository. + command = " ".join(lines[number - 1 : number + 3]) + if _SOURCE_FLASH.search(command) and _INSTALL_OR_BUILD.search(command): + violations.append(f"{path.relative_to(ROOT)}:{number}: {stripped}") + assert not violations, "Source FlashAttention install/build commands:\n - " + "\n - ".join( + violations + ) + + +def test_fastplms_10_manifest_advertises_only_pinned_flash_backends() -> None: + registry = get_model_registry() + advertised = { + family.id: sorted(set(family.attention).intersection(_FLASH_BACKENDS)) + for family in registry.families.values() + if set(family.attention).intersection(_FLASH_BACKENDS) + } + assert advertised == { + "dplm": ["flash_attention_3"], + "esm2": ["flash_attention_2", "flash_attention_3"], + "esm_plusplus": ["flash_attention_2", "flash_attention_3"], + } + assert { + implementation: kernel.dtypes + for implementation, kernel in registry.attention_kernels.items() + } == { + "flash_attention_2": ("bfloat16",), + "flash_attention_3": ("bfloat16",), + } + for family_id, backends in advertised.items(): + for backend in backends: + assert registry.supported_attention_dtypes(family_id, backend) == ("bfloat16",) + + documentation = (ROOT / "docs" / "attention_backends.md").read_text(encoding="utf-8") + normalized_documentation = " ".join(documentation.split()) + assert "Both pinned FlashAttention kernels are BF16-only." in normalized_documentation + assert ( + "Direct FP32 and FP16 calls raise before kernel loading." + in normalized_documentation + ) diff --git a/tests/release/test_goldens.py b/tests/release/test_goldens.py new file mode 100644 index 0000000..33d5eb4 --- /dev/null +++ b/tests/release/test_goldens.py @@ -0,0 +1,669 @@ +"""Producer and read-only validator contracts for official-generated goldens.""" + +from __future__ import annotations + +import hashlib +import json +import pytest +import torch +from dataclasses import replace +from pathlib import Path +from safetensors.torch import load_file, save_file + +from fastplms.registry import ( + FileDigest, + ModelRegistry, + ModelSpec, + OfficialGolden, + get_model_registry, +) +from tests.parity.support.native_reference import _select_requests +from tests.parity.support.reference_adapters.dplm2 import ( + DPLM2_3B_GENERATION_LIMITATION, +) +from tests.release.test_artifacts import _synthetic_registry +from tools.goldens import ( + GoldenBundleRecord, + GoldenError, + check_tier_specs, + convert_native_result, + golden_generation_matrix, + missing_check_golden_ids, + require_complete_check_goldens, + require_declared_goldens, + validate_golden_bundle, + write_golden_bundle, +) +from tools.goldens.from_native import main as golden_main +from tools.remote.prepare_references import MIXED_LENGTHS, prepare_reference_requests + + +ROOT = Path(__file__).resolve().parents[2] + + +def _with_spec(registry: ModelRegistry, spec: ModelSpec) -> ModelRegistry: + return ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + attention_kernels=registry.attention_kernels, + families={spec.family.id: spec.family}, + models={spec.id: spec}, + legal_files=registry.legal_files, + ) + + +def _official_files(spec: ModelSpec) -> list[dict[str, str]]: + return [ + {"algorithm": item.algorithm, "digest": item.digest, "path": item.path} + for item in spec.official.files + ] + + +def _write_native_sequence_result(path: Path, spec: ModelSpec) -> None: + tensors = { + "input__attention_mask": torch.tensor([[1, 1, 1], [1, 1, 0]]), + "input__input_ids": torch.tensor([[1, 2, 3], [1, 3, 0]]), + "output__hidden_0000": torch.arange(24, dtype=torch.bfloat16).reshape(2, 3, 4), + "output__last_hidden_state": torch.arange( + 24, dtype=torch.bfloat16 + ).reshape(2, 3, 4), + "output__logits": torch.arange(30, dtype=torch.bfloat16).reshape(2, 3, 5), + "residue_mask": torch.tensor( + [[False, True, False], [False, True, False]], dtype=torch.bool + ), + } + path.mkdir(parents=True) + save_file(tensors, path / "bf16.safetensors") + metadata = { + "schema_version": 1, + "model_id": spec.id, + "family": spec.family.id, + "reference_repo_id": spec.official.repo_id, + "reference_revision": spec.official.revision, + "reference_files": _official_files(spec), + "state_transform": spec.family.state_transform, + "environment": { + "cuda_device": "Synthetic H100", + "cuda_runtime": "13.0", + "packages": "{\"torch\":\"2.13.0\"}", + "python": "3.12.12", + "torch": "2.13.0", + }, + "precision_tensor_keys": {"bf16": sorted(tensors)}, + } + (path / "metadata.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _declare_golden( + registry: ModelRegistry, + spec: ModelSpec, + record: GoldenBundleRecord, +) -> tuple[ModelRegistry, ModelSpec]: + declaration = OfficialGolden( + metadata=FileDigest( + "tests/goldens/toy.json", + "sha256", + record.metadata_sha256, + ), + tensors=FileDigest( + "tests/goldens/toy.safetensors", + "sha256", + record.tensors_sha256, + ), + ) + declared_spec = replace(spec, official_golden=declaration) + return ( + ModelRegistry( + schema_version=registry.schema_version, + upstreams=registry.upstreams, + families=registry.families, + models={declared_spec.id: declared_spec}, + legal_files=registry.legal_files, + ), + declared_spec, + ) + + +def test_check_tier_official_golden_matrix_is_complete_and_valid() -> None: + registry = get_model_registry() + require_complete_check_goldens(registry) + records = require_declared_goldens(ROOT, registry, tier="check") + assert len(records) == len(check_tier_specs(registry)) + + +def test_golden_producer_is_deterministic_and_validator_is_read_only(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + root = tmp_path / "repository" + first = root / "tests" / "goldens" + second = tmp_path / "second" + tensors = { + "H": torch.arange(24, dtype=torch.bfloat16).reshape(2, 3, 4), + "logits": torch.arange(10, dtype=torch.float32).reshape(2, 5), + } + input_fingerprint = hashlib.sha256(b"synthetic golden input").hexdigest() + environment = { + "cuda": "13.0", + "python": "3.12.12", + "torch": "2.13.0", + "transformers": "5.13.0", + "upstream_environment": "synthetic-reference", + } + command = ( + "python", + "-m", + "tests.parity.support.generate_golden", + "--model", + "toy", + ) + + first_record = write_golden_bundle( + spec, + registry, + tensors, + metadata_path=first / "toy.json", + tensors_path=first / "toy.safetensors", + generation_command=command, + environment=environment, + input_fingerprint=input_fingerprint, + ) + second_record = write_golden_bundle( + spec, + registry, + tensors, + metadata_path=second / "toy.json", + tensors_path=second / "toy.safetensors", + generation_command=command, + environment=environment, + input_fingerprint=input_fingerprint, + ) + assert first_record == second_record + + declared_registry, declared_spec = _declare_golden(registry, spec, first_record) + metadata = json.loads((first / "toy.json").read_text(encoding="utf-8")) + assert metadata["sources"] == [ + { + "id": "toy", + "revision": "4" * 40, + "url": "https://github.com/example/toy.git", + } + ] + assert metadata["checkpoint"]["repo_id"] == "upstream/ToyModel" + assert metadata["checkpoint"]["revision"] == "2" * 40 + assert metadata["environment"]["fingerprint"] == hashlib.sha256( + json.dumps(environment, separators=(",", ":"), sort_keys=True).encode("utf-8") + ).hexdigest() + assert metadata["generation_command"] == list(command) + assert metadata["input_fingerprint"] == input_fingerprint + assert metadata["source_files"] == {} + assert metadata["tensors"]["H"]["shape"] == [2, 3, 4] + assert metadata["tensors"]["H"]["dtype"] == "bfloat16" + + paths = (first / "toy.json", first / "toy.safetensors") + before = tuple(path.stat().st_mtime_ns for path in paths) + validated = validate_golden_bundle( + declared_spec, + declared_registry, + metadata_path=paths[0], + tensors_path=paths[1], + declaration=declared_spec.official_golden, + ) + check_records = require_declared_goldens(root, declared_registry, tier="check") + after = tuple(path.stat().st_mtime_ns for path in paths) + assert validated == first_record + assert check_records == (first_record,) + assert after == before + + +def test_missing_golden_blocks_only_check_when_manifest_declares_it(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + placeholder = GoldenBundleRecord( + metadata_sha256="a" * 64, + tensors_sha256="b" * 64, + tensor_hashes={}, + ) + declared_registry, _ = _declare_golden(registry, spec, placeholder) + + assert require_declared_goldens(tmp_path, declared_registry, tier="compliance") == () + assert require_declared_goldens(tmp_path, declared_registry, tier="artifact") == () + with pytest.raises(GoldenError, match="Missing required official golden for check tier: toy"): + require_declared_goldens(tmp_path, declared_registry, tier="check") + + +def test_golden_validator_rejects_tampering(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, spec = _synthetic_registry(source_root, checkpoint) + golden_root = tmp_path / "repository" / "tests" / "goldens" + record = write_golden_bundle( + spec, + registry, + {"output": torch.tensor([1.0, 2.0])}, + metadata_path=golden_root / "toy.json", + tensors_path=golden_root / "toy.safetensors", + generation_command=("python", "generate.py"), + environment={"python": "3.12", "reference": "synthetic"}, + input_fingerprint=hashlib.sha256(b"input").hexdigest(), + ) + declared_registry, declared_spec = _declare_golden(registry, spec, record) + tensor_path = golden_root / "toy.safetensors" + tensor_path.write_bytes(tensor_path.read_bytes() + b"tampered") + + with pytest.raises(GoldenError, match="tensor-file digest mismatch"): + validate_golden_bundle( + declared_spec, + declared_registry, + metadata_path=golden_root / "toy.json", + tensors_path=tensor_path, + declaration=declared_spec.official_golden, + ) + + +def test_native_sequence_converter_is_compact_deterministic_and_fail_closed( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, original_spec = _synthetic_registry(source_root, checkpoint) + family = replace(original_spec.family, test_tiers=("check",)) + spec = replace(original_spec, family=family) + registry = _with_spec(registry, spec) + native = tmp_path / "native" / spec.id + _write_native_sequence_result(native, spec) + command = ( + "python", + "-m", + "tools.goldens.from_native", + "--model", + spec.id, + ) + + first_root = tmp_path / "repository-one" + second_root = tmp_path / "repository-two" + first = convert_native_result( + spec, + registry, + native, + first_root / "tests" / "goldens", + generation_command=command, + ) + second = convert_native_result( + spec, + registry, + native, + second_root / "tests" / "goldens", + generation_command=command, + ) + assert first.bundle == second.bundle + assert first.manifest_declaration(first_root).startswith("official_golden = {") + + golden_tensors = load_file(first.tensors_path, device="cpu") + assert sorted(golden_tensors) == [ + "input__attention_mask", + "input__input_ids", + "output__last_hidden_state", + "output__logits", + "residue_mask", + ] + metadata = json.loads(first.metadata_path.read_text(encoding="utf-8")) + assert metadata["source_files"] == { + "native/bf16.safetensors": hashlib.sha256( + (native / "bf16.safetensors").read_bytes() + ).hexdigest(), + "native/metadata.json": hashlib.sha256( + (native / "metadata.json").read_bytes() + ).hexdigest(), + } + assert metadata["environment"]["details"]["cuda_device"] == "Synthetic H100" + assert metadata["input_fingerprint"] != "0" * 64 + + broken_metadata = json.loads((native / "metadata.json").read_text(encoding="utf-8")) + broken_metadata.pop("environment") + (native / "metadata.json").write_text( + json.dumps(broken_metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + with pytest.raises(GoldenError, match="has no environment record"): + convert_native_result( + spec, + registry, + native, + tmp_path / "broken", + generation_command=command, + ) + + +def test_dplm2_3b_official_generation_limitation_is_explicit_and_fail_closed( + tmp_path: Path, +) -> None: + """A broken official sampler is evidence, never successful generation parity.""" + + registry = get_model_registry() + spec = registry["dplm2_3b"] + native = tmp_path / "native" / spec.id + _write_native_sequence_result(native, spec) + metadata_path = native / "metadata.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + metadata["generation_limitation"] = DPLM2_3B_GENERATION_LIMITATION + metadata_path.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + output = tmp_path / "goldens" + record = convert_native_result( + spec, + registry, + native, + output, + generation_command=("python", "-m", "tools.goldens", "--model", spec.id), + ) + compact = json.loads(record.metadata_path.read_text(encoding="utf-8")) + assert compact["limitations"] == [ + {"capability": "generation", **DPLM2_3B_GENERATION_LIMITATION} + ] + validate_golden_bundle( + spec, + registry, + metadata_path=record.metadata_path, + tensors_path=record.tensors_path, + ) + + broken = json.loads(metadata_path.read_text(encoding="utf-8")) + broken["generation_limitation"]["reason"] = "different" + metadata_path.write_text( + json.dumps(broken, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + with pytest.raises(GoldenError, match="generation limitation mismatch"): + convert_native_result( + spec, + registry, + native, + tmp_path / "broken", + generation_command=("python", "generate.py"), + ) + + +def test_check_tier_completeness_reports_exact_undeclared_ids(tmp_path: Path) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, original_spec = _synthetic_registry(source_root, checkpoint) + family = replace(original_spec.family, test_tiers=("check", "compliance")) + spec = replace(original_spec, family=family) + registry = _with_spec(registry, spec) + assert check_tier_specs(registry) == (spec,) + assert missing_check_golden_ids(registry) == ("toy",) + with pytest.raises(GoldenError, match="incomplete: toy"): + require_complete_check_goldens(registry) + + declaration = OfficialGolden( + metadata=FileDigest("tests/goldens/toy.json", "sha256", "a" * 64), + tensors=FileDigest("tests/goldens/toy.safetensors", "sha256", "b" * 64), + ) + declared = replace(spec, official_golden=declaration) + complete_registry = _with_spec(registry, declared) + assert missing_check_golden_ids(complete_registry) == () + require_complete_check_goldens(complete_registry) + + +def test_native_structure_converter_requires_reference_hash_contract( + tmp_path: Path, +) -> None: + source_root = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + registry, original_spec = _synthetic_registry(source_root, checkpoint) + family = replace( + original_spec.family, + tokenizer_mode="structure", + test_tiers=("check", "structure"), + ) + spec = replace(original_spec, family=family, size_category="structure") + registry = _with_spec(registry, spec) + native = tmp_path / "native-structure" / spec.id + native.mkdir(parents=True) + tensors = { + "feature__aatype": torch.tensor([0, 1, 2]), + "output__sample_atom_coords": torch.arange(18, dtype=torch.float32).reshape( + 1, 2, 3, 3 + ), + } + save_file(tensors, native / "bundle.safetensors") + + def raw_hash(T: torch.Tensor) -> str: + # T: (...) + value = T.contiguous().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(value).hexdigest() + + metadata = { + "schema_version": 1, + "producer": "reference", + "model_id": spec.id, + "request_sha256": "c" * 64, + "official": { + "repo_id": spec.official.repo_id, + "revision": spec.official.revision, + "files": _official_files(spec), + }, + "environment": { + "python": "3.12.12", + "torch": "2.13.0", + "packages": {"transformers": "5.13.0"}, + }, + "tensor_keys": sorted(tensors), + "tensor_hashes": {name: raw_hash(T) for name, T in tensors.items()}, + } + metadata_path = native / "metadata.json" + metadata_path.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + record = convert_native_result( + spec, + registry, + native, + tmp_path / "structure-golden", + generation_command=("python", "generate-structure.py"), + ) + golden_metadata = json.loads(record.metadata_path.read_text(encoding="utf-8")) + assert golden_metadata["input_fingerprint"] == "c" * 64 + assert golden_metadata["environment"]["details"]["packages"] == ( + '{"transformers":"5.13.0"}' + ) + + metadata["producer"] = "candidate" + metadata_path.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + with pytest.raises(GoldenError, match="only an official reference bundle"): + convert_native_result( + spec, + registry, + native, + tmp_path / "candidate-not-golden", + generation_command=("python", "generate-structure.py"), + ) + + +def test_native_requests_carry_manifest_checkpoint_file_identities(tmp_path: Path) -> None: + registry = get_model_registry() + paths = prepare_reference_requests(tmp_path) + sequence_specs = tuple( + spec for spec in registry.values() if spec.family.tokenizer_mode != "structure" + ) + assert len(paths) == len(sequence_specs) + by_id = {spec.id: spec for spec in sequence_specs} + for path in paths: + request = json.loads(path.read_text(encoding="utf-8")) + spec = by_id[request["model_id"]] + assert request["reference_repo_id"] == spec.official.repo_id + assert request["reference_revision"] == spec.official.revision + assert request["reference_files"] == _official_files(spec) + assert request["generation_policy"] == spec.generation_contract + if spec.generation_contract == "official_unavailable": + assert request["official_generation_limitation"] == ( + DPLM2_3B_GENERATION_LIMITATION + ) + else: + assert "official_generation_limitation" not in request + assert tuple(map(len, request["sequences"])) == MIXED_LENGTHS + + +def test_manifest_generation_matrix_covers_every_check_checkpoint( + tmp_path: Path, +) -> None: + registry = get_model_registry() + native_root = tmp_path / "native" + output_root = tmp_path / "goldens" + entries = golden_generation_matrix(registry, native_root, output_root) + specs = check_tier_specs(registry) + assert tuple(entry.model_id for entry in entries) == tuple(spec.id for spec in specs) + assert len(entries) == 28 + assert sum(entry.kind == "sequence" for entry in entries) == 23 + assert sum(entry.kind == "structure" for entry in entries) == 5 + for entry, spec in zip(entries, specs, strict=True): + assert entry.reference_container == spec.family.reference_container + assert entry.metadata_path == output_root / f"{spec.id}.json" + assert entry.tensors_path == output_root / f"{spec.id}.safetensors" + assert not entry.native_ready + assert not entry.converted_ready + if entry.kind == "sequence": + assert entry.request_path == ( + native_root + / "requests" + / spec.family.reference_container + / f"{spec.id}.json" + ) + assert entry.native_result_path == native_root / "results" / spec.id + else: + assert entry.request_path == ( + native_root + / "structure" + / "requests" + / spec.family.reference_container + / f"{spec.id}.json" + ) + assert entry.native_result_path == ( + native_root / "structure" / "results" / "reference" / spec.id + ) + + nested_spec = next( + spec for spec in specs if spec.family.tokenizer_mode == "structure" + ) + nested_result = ( + native_root + / "structure" + / "results" + / "reference" + / nested_spec.id + / "bf16" + ) + nested_result.mkdir(parents=True) + (nested_result / "metadata.json").write_text("{}\n", encoding="utf-8") + (nested_result / "bundle.safetensors").write_bytes(b"normalized-bundle") + fp32_result = nested_result.parent / "fp32" + fp32_result.mkdir() + (fp32_result / "metadata.json").write_text("{}\n", encoding="utf-8") + (fp32_result / "bundle.safetensors").write_bytes(b"normalized-fp32-bundle") + refreshed = { + entry.model_id: entry + for entry in golden_generation_matrix(registry, native_root, output_root) + } + assert refreshed[nested_spec.id].native_result_path == nested_result + assert refreshed[nested_spec.id].native_ready + + +def test_native_reference_request_selection_is_explicit_and_fail_closed( + tmp_path: Path, +) -> None: + request_dir = tmp_path / "requests" + request_dir.mkdir() + for model_id in ("alpha", "beta"): + (request_dir / f"{model_id}.json").write_text( + json.dumps( + { + "model_id": model_id, + "deep_reference": model_id == "beta", + } + ), + encoding="utf-8", + ) + assert tuple(path.stem for path in _select_requests(request_dir, None)) == ( + "alpha", + "beta", + ) + assert tuple( + path.stem for path in _select_requests(request_dir, ["beta", "alpha"]) + ) == ("beta", "alpha") + assert tuple( + path.stem for path in _select_requests(request_dir, None, deep_only=True) + ) == ("beta",) + assert tuple( + path.stem + for path in _select_requests( + request_dir, + ["alpha", "beta"], + deep_only=True, + ) + ) == ("beta",) + with pytest.raises(ValueError, match="must be unique"): + _select_requests(request_dir, ["alpha", "alpha"]) + with pytest.raises(FileNotFoundError, match="missing selected models"): + _select_requests(request_dir, ["missing"]) + (request_dir / "alpha.json").write_text( + json.dumps({"model_id": "different"}), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="filename and model ID differ"): + _select_requests(request_dir, ["alpha"]) + + +def test_golden_status_reports_exact_manifest_gap(capsys: pytest.CaptureFixture[str]) -> None: + registry = get_model_registry() + assert golden_main(["--status-only", "--report-missing"]) == 0 + report = json.loads(capsys.readouterr().out) + assert tuple(report["undeclared_check_goldens"]) == missing_check_golden_ids(registry) + if missing_check_golden_ids(registry): + with pytest.raises(GoldenError, match="Check-tier official goldens are incomplete"): + golden_main(["--status-only", "--require-complete"]) + else: + assert golden_main(["--status-only", "--require-complete"]) == 0 + + +def test_golden_status_reports_manifest_wide_generation_matrix( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + registry = get_model_registry() + assert ( + golden_main( + [ + "--status-only", + "--report-matrix", + "--native-root", + str(tmp_path / "native"), + "--output-root", + str(tmp_path / "goldens"), + ] + ) + == 0 + ) + report = json.loads(capsys.readouterr().out) + entries = report["check_golden_matrix"] + assert tuple(entry["model_id"] for entry in entries) == tuple( + spec.id for spec in check_tier_specs(registry) + ) + assert len(entries) == 28 diff --git a/tests/release/test_kernel_lock_sources.py b/tests/release/test_kernel_lock_sources.py new file mode 100644 index 0000000..440e5cf --- /dev/null +++ b/tests/release/test_kernel_lock_sources.py @@ -0,0 +1,39 @@ +"""Source-checkout and Hub-artifact contracts for the immutable kernel lock.""" + +from __future__ import annotations + +import importlib.util +import json +import shutil +from pathlib import Path + +from fastplms.attention import _kernel_lock + + +ROOT = Path(__file__).resolve().parents[2] +LOCK = ROOT / "kernels.lock" + + +def test_source_checkout_resolves_the_tracked_kernel_lock() -> None: + assert _kernel_lock._kernel_lock_path().resolve() == LOCK.resolve() + assert len(json.loads(LOCK.read_text(encoding="utf-8"))) == 2 + + +def test_hub_artifact_resolves_its_embedded_kernel_lock(tmp_path: Path) -> None: + package = tmp_path / "fastplms" + attention = package / "attention" + attention.mkdir(parents=True) + artifact_lock = package / "kernels.lock" + shutil.copyfile(LOCK, artifact_lock) + shutil.copyfile(Path(_kernel_lock.__file__), attention / "_kernel_lock.py") + + module_spec = importlib.util.spec_from_file_location( + "artifact_kernel_lock", + attention / "_kernel_lock.py", + ) + assert module_spec is not None and module_spec.loader is not None + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + + assert module._kernel_lock_path() == artifact_lock + assert artifact_lock.read_bytes() == LOCK.read_bytes() diff --git a/tests/release/test_manifest_readiness.py b/tests/release/test_manifest_readiness.py new file mode 100644 index 0000000..d21a6f9 --- /dev/null +++ b/tests/release/test_manifest_readiness.py @@ -0,0 +1,14 @@ +"""Fail-closed release readiness checks for immutable source identities.""" + +from __future__ import annotations + +import pytest + +from fastplms.registry import get_model_registry + + +@pytest.mark.artifact +def test_release_manifest_has_no_unresolved_files() -> None: + """Require independent hashes for every checkpoint and tokenizer asset.""" + + get_model_registry().require_resolved() diff --git a/tests/release/test_model_card_licenses.py b/tests/release/test_model_card_licenses.py new file mode 100644 index 0000000..8f2c13c --- /dev/null +++ b/tests/release/test_model_card_licenses.py @@ -0,0 +1,240 @@ +"""Fail-closed Hugging Face license metadata contracts for model cards.""" + +from __future__ import annotations + +import pytest +from collections.abc import Callable +from pathlib import Path + +from fastplms.registry import HUB_LICENSE_IDENTIFIERS, ModelSpec, load_model_registry +from tools.artifacts.build import render_model_card as render_artifact_model_card +from tools.artifacts.generate_docs import ( + render_model_card as render_documentation_model_card, +) +from tools.artifacts.generate_docs import ( + render_support, +) +from tools.artifacts.license_metadata import parse_hub_license_metadata + + +ROOT = Path(__file__).resolve().parents[2] + +SEQUENCE_TTT_AUTO_CLASSES = { + "ankh": "AutoModelForMaskedLM", + "dplm": "AutoModelForMaskedLM", + "dplm2": "AutoModelForMaskedLM", + "e1": "AutoModelForMaskedLM", + "esm2": "AutoModelForMaskedLM", + "esm3": "AutoModel", + "esm_plusplus": "AutoModelForMaskedLM", +} + +EMBEDDING_FAMILIES = { + "ankh", + "dplm", + "dplm2", + "e1", + "esm2", + "esm3", + "esm_plusplus", + "esmfold2", +} + + +@pytest.mark.parametrize( + "renderer", + (render_documentation_model_card, render_artifact_model_card), + ids=("documentation", "artifact-fallback"), +) +def test_model_card_renderers_use_typed_hub_license_metadata( + renderer: Callable[[ModelSpec], str], +) -> None: + registry = load_model_registry() + for spec in registry.values(): + metadata = parse_hub_license_metadata(renderer(spec)) + assert metadata == dict(spec.family.hub_license_metadata) + assert metadata["license"] in HUB_LICENSE_IDENTIFIERS + + +def test_checked_in_model_cards_use_typed_hub_license_metadata() -> None: + registry = load_model_registry() + for spec in registry.values(): + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + assert parse_hub_license_metadata(card) == dict(spec.family.hub_license_metadata) + + +def test_model_cards_include_family_appropriate_usage() -> None: + registry = load_model_registry() + expected_sections = { + "ankh": ( + "## Tokenization and forward inference", + "## Encoder and sequence-to-sequence use", + ), + "boltz2": ("## Protein structure prediction",), + "dplm": ("## Tokenization and forward inference", "## Diffusion sequence generation"), + "dplm2": ("## Amino-acid and structure co-generation",), + "e1": ("## Tokenizer-free E1 input",), + "esm2": ( + "## Tokenization and forward inference", + "## Masked language modeling and contacts", + ), + "esm3": ("## Sequence inference and masked-sequence generation",), + "esm_plusplus": ("## Tokenization and forward inference", "## ESMC behavior"), + "esmfold": ("## Protein structure prediction",), + "esmfold2": ("## Protein folding", "## Learned representation and ESMC precision"), + } + stale_fragments = ( + "attn_backend", + "kernels_flash", + "pooling_types", + "save_path=", + 'format="pth"', + ) + for spec in registry.values(): + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + assert "## Quick start" in card + assert spec.fast.repo_id in card + for section in expected_sections[spec.family.id]: + assert section in card + for fragment in stale_fragments: + assert fragment not in card + + if spec.family.id in { + "ankh", + "dplm", + "dplm2", + "e1", + "esm2", + "esm3", + "esm_plusplus", + }: + assert "## Dataset embeddings" in card + + +def test_model_cards_keep_checkpoint_specific_ttt_boundaries() -> None: + esmfold = (ROOT / "model_cards" / "esmfold.md").read_text(encoding="utf-8") + assert "does not expose ProteinTTT" in esmfold + + registry = load_model_registry() + for spec in registry.by_family("esmfold2"): + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + if "experimental" in spec.id: + assert "standard and Fast checkpoints expose" not in card + else: + assert "standard and Fast checkpoints expose" in card + + +def test_every_manifest_model_card_has_the_shared_capability_contract() -> None: + registry = load_model_registry() + specs = tuple(registry.values()) + assert len(specs) == 29 + + for spec in specs: + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + normalized = " ".join(card.split()) + assert "## Capabilities" in card + for feature in ( + "Sequence classification", + "Token classification", + "PEFT fine-tuning", + "Embeddings", + "Test-time training", + "Attention variants", + "Compliance", + ): + assert f"| {feature} |" in card + + for backend in spec.family.attention: + assert f"`{backend}`" in card + assert "An unavailable requested backend raises" in normalized + if "compliance" in spec.family.test_tiers: + assert "This family declares the `compliance` tier." in card + else: + assert "This family does not declare the `compliance` tier." in card + + advertises_heads = { + "AutoModelForSequenceClassification", + "AutoModelForTokenClassification", + }.issubset(spec.auto_map) + assert "## PEFT fine-tuning" in card + assert 'target_modules="all-linear"' in card + assert "ESM2-specific shipped CLI is an example, not a\nsupport boundary" in card + if advertises_heads: + assert "## Downstream classification" in card + assert "base weights with an untrained task head" in card + assert "AutoModelForSequenceClassification.from_pretrained" in card + assert "AutoModelForTokenClassification.from_pretrained" in card + assert "token_labels = torch.full_like(batch[\"input_ids\"], -100)" in card + assert 'modules_to_save=["classifier"]' in card + else: + assert "## Downstream classification" not in card + assert "| PEFT fine-tuning | Supported pattern:" in card + assert "preserve any new head through `modules_to_save`" in card + assert 'modules_to_save=["classifier"]' not in card + + if spec.family.id in EMBEDDING_FAMILIES: + if spec.family.id == "esmfold2": + assert "## Learned representation and ESMC precision" in card + else: + assert "## Dataset embeddings" in card + else: + assert "| Embeddings | Unavailable" in card + + ttt_auto_class = SEQUENCE_TTT_AUTO_CLASSES.get(spec.family.id) + if ttt_auto_class is not None: + section = card.split("## Test-time training", maxsplit=1)[1] + assert f"from transformers import {ttt_auto_class}" in section + assert "updates only injected low-rank" in section + assert 'save_pretrained("adapted", safe_serialization=True)' in section + assert "ttt_model.ttt_reset()" in section + + +def test_esmfold2_cards_publish_embedding_projection_shapes_and_ttt_scope() -> None: + registry = load_model_registry() + for spec in registry.by_family("esmfold2"): + card = (ROOT / "model_cards" / f"{spec.id}.md").read_text(encoding="utf-8") + assert "`H: (b, l, 81, 2560) -> Z: (b, l, 256)`" in card + assert "returns one `(l, 256)` residue" in card + if "experimental" in spec.id: + assert "does not expose folding TTT" in card + else: + assert "## Optional folding TTT" in card + assert "fold_protein_ttt(" in card + assert "not a generic\n`save_pretrained` adapter-persistence path" in card + + +@pytest.mark.parametrize( + "renderer", + (render_documentation_model_card, render_artifact_model_card), + ids=("documentation", "artifact-fallback"), +) +def test_model_card_renderers_surface_typed_limitations( + renderer: Callable[[ModelSpec], str], +) -> None: + registry = load_model_registry() + for spec in registry.values(): + card = renderer(spec) + assert f"Public input: {spec.family.public_input}" in card + assert "Input mode:" not in card + assert "Internal preparation mode:" not in card + assert f"Generation contract: `{spec.generation_contract}`" in card + if spec.notes and spec.family.id != "esm_plusplus": + assert "## Notes and limitations" in card + assert " ".join(spec.notes.split()) in " ".join(card.split()) + elif spec.family.id == "esm_plusplus": + # ESMC measurements are report-derived per checkpoint. The shared + # manifest note contains a historical ESMC-6B-only observation + # and must never be copied into the small or large checkpoint cards. + assert "ESMC-6B Flex Attention exceeds" not in card + + +def test_support_table_exposes_generation_contracts() -> None: + registry = load_model_registry() + support = render_support(registry) + assert "| Public input |" in support + assert "| Input |" not in support + for family in registry.families.values(): + assert family.public_input in support + assert "| Generation contract |" in support + dplm2_3b_row = next(line for line in support.splitlines() if line.startswith("| `dplm2_3b`")) + assert "`official_unavailable`" in dplm2_3b_row diff --git a/tests/release/test_optimized_mode_contracts.py b/tests/release/test_optimized_mode_contracts.py new file mode 100644 index 0000000..1f4943b --- /dev/null +++ b/tests/release/test_optimized_mode_contracts.py @@ -0,0 +1,531 @@ +"""Public validation must survive Python's optimized ``-O`` mode.""" + +from __future__ import annotations + +import ast +import builtins +import os +import subprocess +import sys +import pytest +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] + + +def _public_method(path: str, class_name: str, method_name: str) -> ast.FunctionDef: + module = ast.parse((ROOT / path).read_text(encoding="utf-8")) + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for child in node.body: + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and ( + child.name == method_name + ): + if isinstance(child, ast.AsyncFunctionDef): + raise AssertionError(f"Unexpected async method {class_name}.{method_name}") + return child + raise AssertionError(f"Missing public method {class_name}.{method_name}") + + +def _top_level_function(path: str, function_name: str) -> ast.FunctionDef: + module = ast.parse((ROOT / path).read_text(encoding="utf-8")) + for node in module.body: + if isinstance(node, ast.FunctionDef) and node.name == function_name: + return node + raise AssertionError(f"Missing public function {function_name}") + + +def _assert_no_runtime_asserts(function: ast.FunctionDef, label: str) -> None: + runtime_asserts = [node for node in ast.walk(function) if isinstance(node, ast.Assert)] + assert not runtime_asserts, f"{label} uses validation erased by python -O" + + +def test_boltz_public_boundaries_do_not_use_runtime_asserts() -> None: + model_path = "src/fastplms/models/boltz/modeling_boltz2.py" + config_method = _public_method(model_path, "Boltz2Config", "from_hyperparameters") + _assert_no_runtime_asserts(config_method, "Boltz2Config.from_hyperparameters") + for method_name in ( + "__init__", + "from_boltz_checkpoint", + "predict_structure", + "save_as_cif", + ): + method = _public_method(model_path, "Boltz2Model", method_name) + _assert_no_runtime_asserts(method, f"Boltz2Model.{method_name}") + + for function_name in ("_enforce_pairformer_v2", "_require_key"): + function = _top_level_function(model_path, function_name) + _assert_no_runtime_asserts(function, function_name) + + for function_name in ("_normalize_sequence", "build_boltz2_features"): + function = _top_level_function( + "src/fastplms/models/boltz/minimal_featurizer.py", + function_name, + ) + _assert_no_runtime_asserts(function, function_name) + + for function_name in ("_confidence_per_atom", "write_cif"): + function = _top_level_function( + "src/fastplms/models/boltz/cif_writer.py", + function_name, + ) + _assert_no_runtime_asserts(function, function_name) + + confidence_forward = _public_method( + "src/fastplms/models/boltz/vb_modules_confidencev2.py", + "ConfidenceModule", + "forward", + ) + _assert_no_runtime_asserts(confidence_forward, "ConfidenceModule.forward") + + indexing_matrix = _top_level_function( + "src/fastplms/models/boltz/vb_modules_encodersv2.py", + "get_indexing_matrix", + ) + _assert_no_runtime_asserts(indexing_matrix, "get_indexing_matrix") + for method_name in ("__init__", "compute"): + schedule_method = _public_method( + "src/fastplms/models/boltz/vb_potentials_schedules.py", + "PiecewiseStepFunction", + method_name, + ) + _assert_no_runtime_asserts( + schedule_method, + f"PiecewiseStepFunction.{method_name}", + ) + potential_method = _public_method( + "src/fastplms/models/boltz/vb_potentials_potentials.py", + "FlatBottomPotential", + "compute_function", + ) + _assert_no_runtime_asserts( + potential_method, + "FlatBottomPotential.compute_function", + ) + + +def test_boltz_encoder_schedule_and_potential_validation_survives_optimized_mode() -> None: + script = r""" +import torch + +from fastplms.models.boltz.vb_modules_encodersv2 import get_indexing_matrix +from fastplms.models.boltz.vb_potentials_potentials import FlatBottomPotential +from fastplms.models.boltz.vb_potentials_schedules import PiecewiseStepFunction + + +def must_raise(call, expected): + try: + call() + except expected: + return + raise RuntimeError(f"Expected {expected.__name__}") + + +must_raise( + lambda: get_indexing_matrix(1, 3, 6, torch.device("cpu")), + ValueError, +) +must_raise( + lambda: PiecewiseStepFunction((), (1.0,)), + ValueError, +) +must_raise( + lambda: FlatBottomPotential.compute_function( + object(), + value=torch.tensor([0.5]), + k=torch.tensor(1.0), + lower_bounds=torch.tensor([0.0]), + upper_bounds=torch.tensor([1.0]), + negation_mask=torch.tensor([False]), + ), + ValueError, +) +""" + environment = os.environ.copy() + source_root = str(ROOT / "src") + environment["PYTHONPATH"] = os.pathsep.join( + value for value in (source_root, environment.get("PYTHONPATH", "")) if value + ) + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert completed.returncode == 0, completed.stderr + + +def test_esmfold2_public_boundaries_do_not_use_runtime_asserts() -> None: + model_path = "src/fastplms/models/esmfold2/modeling_esmfold2.py" + for method_name in ( + "_ensure_ttt_lm_head", + "_ttt_tokenize", + "_ttt_predict_logits", + "_fold_protein_no_ttt", + "fold_protein_ttt", + "result_to_cif", + "result_to_pdb", + ): + method = _public_method(model_path, "ESMFold2Model", method_name) + _assert_no_runtime_asserts(method, f"ESMFold2Model.{method_name}") + + for path, class_name, method_name in ( + ( + "src/fastplms/models/esmfold2/esmfold2_protein_chain.py", + "ProteinChain", + "__post_init__", + ), + ( + "src/fastplms/models/esmfold2/esmfold2_protein_complex.py", + "ProteinComplex", + "__post_init__", + ), + ( + "src/fastplms/models/esmfold2/esmfold2_msa.py", + "MSA", + "__post_init__", + ), + ( + "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py", + "MolecularComplex", + "__post_init__", + ), + ): + method = _public_method(path, class_name, method_name) + _assert_no_runtime_asserts(method, f"{class_name}.{method_name}") + + for method_name in ( + "from_atom37", + "from_backbone_atom_coordinates", + "from_mmcif", + "chain_iterable_from_mmcif", + "find_nonpolymer_contacts", + ): + method = _public_method( + "src/fastplms/models/esmfold2/esmfold2_protein_chain.py", + "ProteinChain", + method_name, + ) + _assert_no_runtime_asserts(method, f"ProteinChain.{method_name}") + + experimental_path = "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + for method_name in ("_compute_lm_hidden_states", "result_to_cif", "result_to_pdb"): + method = _public_method( + experimental_path, + "ESMFold2ExperimentalModel", + method_name, + ) + _assert_no_runtime_asserts( + method, + f"ESMFold2ExperimentalModel.{method_name}", + ) + + greedy = _top_level_function( + "src/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py", + "greedy_select_indices", + ) + _assert_no_runtime_asserts(greedy, "greedy_select_indices") + + comparable = _public_method( + "src/fastplms/models/esmfold2/esmfold2_protein_complex.py", + "ProteinComplex", + "_sanity_check_complexes_are_comparable", + ) + _assert_no_runtime_asserts( + comparable, + "ProteinComplex._sanity_check_complexes_are_comparable", + ) + table_validation = _top_level_function( + "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py", + "_assert_table_lengths", + ) + _assert_no_runtime_asserts(table_validation, "MolecularComplex table validation") + + +def test_attention_ankh_and_esmc_contracts_do_not_use_runtime_asserts() -> None: + for class_name, method_names in ( + ("IndexFirstAxis", ("forward", "backward")), + ("IndexPutFirstAxis", ("forward",)), + ): + for method_name in method_names: + method = _public_method( + "src/fastplms/attention/_core.py", + class_name, + method_name, + ) + _assert_no_runtime_asserts(method, f"{class_name}.{method_name}") + + decoder_inputs = _public_method( + "src/fastplms/models/ankh/modeling_ankh.py", + "FastAnkhForConditionalGeneration", + "_prepare_decoder_embedding_inputs", + ) + _assert_no_runtime_asserts( + decoder_inputs, + "FastAnkhForConditionalGeneration._prepare_decoder_embedding_inputs", + ) + + for class_name, method_name in ( + ("RotaryEmbedding", "forward"), + ("TransformerStack", "forward"), + ): + method = _public_method( + "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py", + class_name, + method_name, + ) + _assert_no_runtime_asserts(method, f"{class_name}.{method_name}") + + +def test_embedding_and_state_validation_contracts_do_not_use_runtime_asserts() -> None: + for path, function_name in ( + ("src/fastplms/embeddings/runner.py", "embed_dataset"), + ("src/fastplms/embeddings/storage.py", "load_sqlite_result"), + ("tools/conversion/state_validation.py", "assert_model_parameters_fp32"), + ( + "tools/conversion/state_validation.py", + "assert_state_dict_floating_tensors_fp32", + ), + ("tools/conversion/state_validation.py", "assert_state_dict_equal"), + ): + function = _top_level_function(path, function_name) + _assert_no_runtime_asserts(function, function_name) + + +_COMPOSITE_GUARD_PATHS = ( + "src/fastplms/models/ankh/modeling_ankh.py", + "src/fastplms/models/dplm/modeling_dplm.py", + "src/fastplms/models/dplm2/modeling_dplm2.py", + "src/fastplms/models/e1/modeling_e1.py", + "src/fastplms/models/esm2/modeling_fastesm.py", + "src/fastplms/models/esm3/modeling_esm3.py", + "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py", + "src/fastplms/models/esmfold/modeling_fast_esmfold.py", +) + + +def _fastplms_import_guard(path: str) -> ast.Try: + module = ast.parse((ROOT / path).read_text(encoding="utf-8")) + for node in module.body: + if not isinstance(node, ast.Try): + continue + if any( + isinstance(child, ast.ImportFrom) + and child.module is not None + and child.module.startswith("fastplms") + for child in node.body + ): + return node + raise AssertionError(f"Missing FastPLMs import guard in {path}") + + +def _guard_required_names(guard: ast.Try) -> set[str]: + return { + alias.asname or alias.name + for child in guard.body + if isinstance(child, ast.ImportFrom) + for alias in child.names + } + + +def _execute_guard(guard: ast.Try, *, missing_name: str, shared: set[str]) -> None: + original_import = builtins.__import__ + + def unavailable(name, globals=None, locals=None, fromlist=(), level=0): + if name.startswith("fastplms"): + raise ModuleNotFoundError( + f"synthetic missing dependency: {missing_name}", + name=missing_name, + ) + return original_import(name, globals, locals, fromlist, level) + + namespace = {name: object() for name in shared} + namespace["__builtins__"] = {**vars(builtins), "__import__": unavailable} + code = compile( + ast.fix_missing_locations(ast.Module(body=[guard], type_ignores=[])), + "", + "exec", + ) + exec(code, namespace) + + +@pytest.mark.parametrize("path", _COMPOSITE_GUARD_PATHS) +def test_composite_guards_preserve_transitive_import_errors(path: str) -> None: + guard = _fastplms_import_guard(path) + with pytest.raises(ModuleNotFoundError) as captured: + _execute_guard( + guard, + missing_name="synthetic_runtime_dependency", + shared=_guard_required_names(guard), + ) + assert captured.value.name == "synthetic_runtime_dependency" + + +@pytest.mark.parametrize("path", _COMPOSITE_GUARD_PATHS) +def test_composite_guards_require_every_predefined_shared_symbol(path: str) -> None: + guard = _fastplms_import_guard(path) + required = _guard_required_names(guard) + missing_symbol = sorted(required)[0] + with pytest.raises(ModuleNotFoundError) as captured: + _execute_guard( + guard, + missing_name="fastplms", + shared=required - {missing_symbol}, + ) + assert captured.value.name == "fastplms" + + +@pytest.mark.parametrize("path", _COMPOSITE_GUARD_PATHS) +def test_composite_guards_allow_complete_legacy_flat_context(path: str) -> None: + guard = _fastplms_import_guard(path) + _execute_guard( + guard, + missing_name="fastplms", + shared=_guard_required_names(guard), + ) + + +def test_esmfold2_fallback_does_not_mask_transitive_import_errors() -> None: + guard = _fastplms_import_guard("src/fastplms/models/esmfold2/modeling_esmfold2.py") + handler = guard.handlers[0] + assert isinstance(handler.type, ast.Name) + assert handler.type.id == "ModuleNotFoundError" + assert any( + isinstance(node, ast.Compare) + and any( + isinstance(comparator, ast.Constant) and comparator.value == "fastplms" + for comparator in node.comparators + ) + for node in ast.walk(handler) + ) + + +def test_invalid_attention_ankh_and_esmc_inputs_fail_under_python_optimized_mode() -> None: + script = r""" +import torch + +from fastplms.attention._core import index_first_axis, index_put_first_axis +from fastplms.models.ankh.modeling_ankh import FastAnkhForConditionalGeneration +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + RotaryEmbedding, + TransformerStack, +) + + +def must_raise(call, expected): + try: + call() + except expected: + return + raise RuntimeError(f"Expected {expected.__name__}") + + +must_raise( + lambda: index_first_axis(torch.ones(3), torch.tensor([0])), + ValueError, +) +must_raise( + lambda: index_first_axis(torch.ones(3, 2), torch.tensor([[0]])), + ValueError, +) +must_raise( + lambda: index_put_first_axis(torch.ones(1, 2), torch.tensor([[0]]), 3), + ValueError, +) +must_raise( + lambda: FastAnkhForConditionalGeneration._prepare_decoder_embedding_inputs( + object(), + batch_size=1, + decoder_inputs=None, + decoder_input_ids=None, + decoder_attention_mask=None, + ), + ValueError, +) + +rotary = RotaryEmbedding(8) +rotary._update_cos_sin_cache = lambda *args, **kwargs: None +must_raise( + lambda: rotary(torch.randn(1, 3, 2, 8), torch.randn(1, 3, 2, 8)), + RuntimeError, +) + +stack = TransformerStack(16, 2, 1, attn_backend="sdpa") +must_raise( + lambda: stack( + torch.randn(2, 3, 16), + attention_mask=torch.ones(2, 3, 1), + ), + ValueError, +) +""" + environment = os.environ.copy() + source_root = str(ROOT / "src") + environment["PYTHONPATH"] = os.pathsep.join( + value for value in (source_root, environment.get("PYTHONPATH", "")) if value + ) + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert completed.returncode == 0, completed.stderr + + +def test_state_validation_failures_survive_python_optimized_mode() -> None: + script = r""" +import torch + +from tools.conversion.state_validation import ( + assert_model_parameters_fp32, + assert_state_dict_equal, + assert_state_dict_floating_tensors_fp32, +) + + +def must_raise(call): + try: + call() + except AssertionError: + return + raise RuntimeError("Expected AssertionError") + + +must_raise(lambda: assert_model_parameters_fp32(torch.nn.Identity(), "empty")) +must_raise( + lambda: assert_state_dict_floating_tensors_fp32( + {"weight": torch.ones(1, dtype=torch.float16)}, + "half-state", + ) +) +must_raise( + lambda: assert_state_dict_equal( + {"weight": torch.ones(1)}, + {"weight": torch.zeros(1)}, + "mismatch", + ) +) +""" + environment = os.environ.copy() + source_root = str(ROOT / "src") + environment["PYTHONPATH"] = os.pathsep.join( + value for value in (str(ROOT), source_root, environment.get("PYTHONPATH", "")) if value + ) + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert completed.returncode == 0, completed.stderr diff --git a/tests/release/test_production_source_boundary.py b/tests/release/test_production_source_boundary.py new file mode 100644 index 0000000..b4523bb --- /dev/null +++ b/tests/release/test_production_source_boundary.py @@ -0,0 +1,74 @@ +"""Repository-wide boundary between production code and parity oracles.""" + +from __future__ import annotations + +import ast +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +PACKAGE = ROOT / "src/fastplms" +FORBIDDEN_IMPORT_ROOTS = { + "E1", + "boltz", + "byprot", + "esm", + "openfold", + "vendor", +} + + +def _absolute_import_roots(tree: ast.AST) -> set[str]: + roots = { + alias.name.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + } + roots.update( + node.module.split(".", 1)[0] + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.level == 0 and node.module + ) + return roots + + +def _mutates_python_path(node: ast.Call) -> bool: + function = node.func + return ( + isinstance(function, ast.Attribute) + and function.attr in {"append", "extend", "insert"} + and isinstance(function.value, ast.Attribute) + and function.value.attr == "path" + and isinstance(function.value.value, ast.Name) + and function.value.value.id == "sys" + ) + + +def _dynamic_upstream_import(node: ast.Call) -> bool: + function = node.func + is_import = (isinstance(function, ast.Name) and function.id == "__import__") or ( + isinstance(function, ast.Attribute) and function.attr == "import_module" + ) + if not is_import or not node.args: + return False + module = node.args[0] + if not isinstance(module, ast.Constant) or not isinstance(module.value, str): + return False + return module.value.split(".", 1)[0] in FORBIDDEN_IMPORT_ROOTS + + +def test_production_source_never_imports_parity_oracles() -> None: + failures: list[str] = [] + for path in sorted(PACKAGE.rglob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + forbidden = sorted(_absolute_import_roots(tree) & FORBIDDEN_IMPORT_ROOTS) + if forbidden: + failures.append(f"{path.relative_to(ROOT)} imports upstream roots {forbidden}") + if any(_mutates_python_path(node) for node in ast.walk(tree) if isinstance(node, ast.Call)): + failures.append(f"{path.relative_to(ROOT)} mutates sys.path") + if any( + _dynamic_upstream_import(node) for node in ast.walk(tree) if isinstance(node, ast.Call) + ): + failures.append(f"{path.relative_to(ROOT)} dynamically imports an upstream root") + assert not failures, "\n" + "\n".join(failures) diff --git a/tests/release/test_publish_files_only.py b/tests/release/test_publish_files_only.py new file mode 100644 index 0000000..67118eb --- /dev/null +++ b/tests/release/test_publish_files_only.py @@ -0,0 +1,1527 @@ +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import shutil +import subprocess +import pytest +from collections.abc import Iterator +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo +from huggingface_hub import CommitOperationAdd, CommitOperationDelete + +from fastplms.registry import FileDigest, ModelSpec, get_model_registry +from tools.artifacts import ArtifactError, hash_file +from tools.artifacts import publish as publish_module +from tools.artifacts.build import ( + _RELEASE_TOOL_SCOPE_PATHS, + _artifact_auto_map, + _checkpoint_identity_hash, + _expected_registry_provenance, + _materialize_model_card, + _render_artifact_requirements, + _validated_release_tool_snapshot, + _validated_runtime_snapshot, + _write_bootstrap, + _write_runtime_bundle, + render_model_card, +) +from tools.artifacts.publish import ( + CompletePublishPlan, + _is_weight_path, + _obsolete_registry_pinned_paths, + _run_required_complete_autoclass_probe, + _selected_specs, + _validated_release_text_snapshot, + main, + prepare_complete_plan, + prepare_files_only_plan, + publish_complete, + publish_files_only, +) + + +ROOT = Path(__file__).resolve().parents[2] + + +@pytest.fixture(scope="module", autouse=True) +def _clean_publication_source( + tmp_path_factory: pytest.TempPathFactory, +) -> Iterator[None]: + """Exercise publication against a clean snapshot of the current source bytes. + + Publication must reject dirty or untracked release inputs. The repository under + test is necessarily dirty while a change is being developed, so using it as the + positive fixture prevents the tests from reaching the attack they intend to + exercise. Commit the bounded publication inputs in an isolated repository and + point both the artifact builder and publisher at that immutable snapshot instead. + + This fixture does not bypass the production cleanliness checks: every source byte + consumed below is tracked in the temporary repository, and the production code + still validates status, membership, revision, and Git-archive bytes itself. + """ + + global ROOT + + original_root = ROOT + original_publish_file = publish_module.__file__ + source_root = tmp_path_factory.mktemp("clean-publication-source") + ignored = shutil.ignore_patterns("__pycache__", "*.pyc", "*.pyo") + for relative_name in ( + "src/fastplms", + "model_cards", + "LICENSES", + "requirements", + "tools/artifacts", + "tools/conversion", + ): + source = original_root.joinpath(*relative_name.split("/")) + destination = source_root.joinpath(*relative_name.split("/")) + shutil.copytree(source, destination, ignore=ignored) + for relative_name in ( + "LICENSE", + "THIRD_PARTY_NOTICES.md", + "kernels.lock", + "tools/remote/biohub_reference_environment.py", + "tools/source_provenance.py", + ): + source = original_root / relative_name + destination = source_root / relative_name + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, destination) + + git = ["git", "-c", f"safe.directory={source_root.as_posix()}"] + subprocess.run( + [*git, "init", "--initial-branch=main"], + cwd=source_root, + check=True, + capture_output=True, + ) + subprocess.run( + [*git, "config", "user.email", "tests@example.invalid"], + cwd=source_root, + check=True, + ) + subprocess.run( + [*git, "config", "user.name", "FastPLMs Tests"], + cwd=source_root, + check=True, + ) + subprocess.run( + [*git, "config", "core.autocrlf", "false"], + cwd=source_root, + check=True, + ) + subprocess.run( + [*git, "config", "commit.gpgsign", "false"], + cwd=source_root, + check=True, + ) + subprocess.run( + [*git, "config", "core.hooksPath", ".git/disabled-hooks"], + cwd=source_root, + check=True, + ) + subprocess.run([*git, "add", "."], cwd=source_root, check=True) + subprocess.run( + [*git, "commit", "-m", "immutable publication fixture"], + cwd=source_root, + check=True, + capture_output=True, + ) + + ROOT = source_root + publish_module.__file__ = str(source_root / "tools" / "artifacts" / "publish.py") + try: + yield + finally: + ROOT = original_root + publish_module.__file__ = original_publish_file + + +class FakeApi: + def __init__(self, spec: ModelSpec, *, corrupt_weight: bool = False) -> None: + siblings: list[SimpleNamespace] = [] + for expected in spec.fast.files: + digest = ( + "0" * len(expected.digest) + if corrupt_weight and _is_weight_path(expected.path) + else expected.digest + ) + siblings.append( + SimpleNamespace( + rfilename=expected.path, + blob_id=digest if expected.algorithm == "git-sha1" else None, + lfs={"sha256": digest} if expected.algorithm == "sha256" else None, + ) + ) + self.info = SimpleNamespace(sha="a" * 40, siblings=siblings) + self.model_info_calls: list[dict[str, Any]] = [] + self.create_commit_calls: list[dict[str, Any]] = [] + + def model_info(self, repo_id: str, **kwargs: Any) -> SimpleNamespace: + self.model_info_calls.append({"repo_id": repo_id, **kwargs}) + return self.info + + def create_commit(self, **kwargs: Any) -> SimpleNamespace: + self.create_commit_calls.append(kwargs) + return SimpleNamespace( + oid="b" * 40, + commit_url=f"https://huggingface.co/{kwargs['repo_id']}/commit/{'b' * 40}", + ) + + +def _write_json(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _canonical_release_bytes(path: Path) -> bytes: + return path.read_bytes().replace(b"\r\n", b"\n").replace(b"\r", b"\n") + + +def _self_attest_runtime_mutation(artifact: Path, relative_names: tuple[str, ...]) -> None: + runtime_attestation_path = artifact / "runtime-attestation.json" + runtime_attestation = json.loads(runtime_attestation_path.read_text(encoding="utf-8")) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + for relative_name in relative_names: + path = artifact.joinpath(*relative_name.split("/")) + digest = f"sha256:{hash_file(path)}" + runtime_attestation["files"][relative_name] = digest + manifest[relative_name] = digest + _write_json(runtime_attestation_path, runtime_attestation) + manifest["runtime-attestation.json"] = f"sha256:{hash_file(runtime_attestation_path)}" + _write_json(manifest_path, manifest) + + +def _rewrite_materialized_card(artifact: Path, spec: ModelSpec) -> None: + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + card_source = ROOT / "model_cards" / f"{spec.id}.md" + card_template = ( + card_source.read_text(encoding="utf-8") + if card_source.is_file() + else render_model_card(spec) + ) + materialized = _materialize_model_card( + card_template, + runtime_revision=provenance["runtime_revision"], + source_tree_sha256=provenance["source_tree_sha256"], + runtime_bundle_sha256=provenance["runtime_bundle_sha256"], + ) + (artifact / "README.md").write_text( + materialized, + encoding="utf-8", + newline="\n", + ) + + +def _initialize_release_text_repository(root: Path, spec: ModelSpec) -> Path: + for relative_name in _RELEASE_TOOL_SCOPE_PATHS: + tool_path = root.joinpath(*relative_name.split("/")) + tool_path.parent.mkdir(parents=True, exist_ok=True) + tool_path.write_text(f"# immutable test tool: {relative_name}\n", encoding="utf-8") + (root / "LICENSE").write_text("test license\n", encoding="utf-8") + (root / "THIRD_PARTY_NOTICES.md").write_text( + "test notices\n", + encoding="utf-8", + ) + registry = get_model_registry() + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + for item in source.distribution_files: + legal_path = root.joinpath( + "LICENSES", + source_id, + *item.path.split("/"), + ) + legal_path.parent.mkdir(parents=True, exist_ok=True) + legal_path.write_text("test upstream license\n", encoding="utf-8") + card = root / "model_cards" / f"{spec.id}.md" + card.parent.mkdir(parents=True) + card.write_text("tracked card\n", encoding="utf-8") + subprocess.run(["git", "init", "--initial-branch=main"], cwd=root, check=True) + subprocess.run( + ["git", "config", "user.email", "tests@example.invalid"], + cwd=root, + check=True, + ) + subprocess.run( + ["git", "config", "user.name", "FastPLMs Tests"], + cwd=root, + check=True, + ) + subprocess.run(["git", "add", "."], cwd=root, check=True) + subprocess.run(["git", "commit", "-m", "tracked card"], cwd=root, check=True) + return card + + +def _files_only_artifact(root: Path, spec: ModelSpec) -> Path: + artifact = root / spec.fast.repo_id.split("/", maxsplit=1)[1] + registry = get_model_registry() + card_source = ROOT / "model_cards" / f"{spec.id}.md" + release_tool_revision, release_tool_sha256, release_tool_payloads = ( + _validated_release_tool_snapshot(ROOT) + ) + safe_files: dict[str, str | bytes] = { + "README.md": ( + card_source.read_bytes() + if card_source.is_file() + else render_model_card(spec).encode() + ), + "config.json": "{}\n", + "fastplms_bundle.py": "", + "modeling_fastplms.py": "", + "requirements.txt": _render_artifact_requirements(spec, release_tool_payloads), + "THIRD_PARTY_NOTICES.md": _canonical_release_bytes( + ROOT / "THIRD_PARTY_NOTICES.md" + ), + "LICENSES/FastPLMs-Apache-2.0.txt": _canonical_release_bytes(ROOT / "LICENSE"), + } + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + for item in source.distribution_files: + safe_files[f"LICENSES/{source_id}/{item.path}"] = _canonical_release_bytes( + ROOT.joinpath("LICENSES", source_id, *item.path.split("/")) + ) + for relative_name, contents in safe_files.items(): + path = artifact.joinpath(*relative_name.split("/")) + path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(contents, bytes): + path.write_bytes(contents) + else: + path.write_text(contents, encoding="utf-8", newline="\n") + + source_revision, runtime_payloads, source_tree_sha256 = _validated_runtime_snapshot( + ROOT, + registry, + spec, + ) + packaged_runtime_names: list[str] = [] + for relative_name, payload in runtime_payloads.items(): + packaged_name = f"fastplms/{relative_name}" + path = artifact.joinpath(*packaged_name.split("/")) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload) + packaged_runtime_names.append(packaged_name) + runtime_bundle_sha256 = _write_runtime_bundle( + artifact / "fastplms_bundle.py", + artifact / "fastplms", + ) + card_template = ( + card_source.read_text(encoding="utf-8") + if card_source.is_file() + else render_model_card(spec) + ) + materialized_card = _materialize_model_card( + card_template, + runtime_revision=source_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ) + (artifact / "README.md").write_text( + materialized_card, + encoding="utf-8", + newline="\n", + ) + _write_bootstrap(artifact / "modeling_fastplms.py", spec, runtime_bundle_sha256) + selected_checkpoint = spec.artifact_checkpoint + _write_json( + artifact / "config.json", + { + "auto_map": _artifact_auto_map(spec), + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": selected_checkpoint.repo_id, + "fastplms_checkpoint_revision": selected_checkpoint.revision, + "fastplms_checkpoint_hash": _checkpoint_identity_hash(selected_checkpoint), + "fastplms_weights_revision": selected_checkpoint.revision, + "fastplms_runtime_revision": source_revision, + "fastplms_source_tree_sha256": source_tree_sha256, + "fastplms_runtime_bundle_sha256": runtime_bundle_sha256, + "fastplms_release_tool_revision": release_tool_revision, + "fastplms_release_tool_sha256": release_tool_sha256, + }, + ) + provenance = { + **_expected_registry_provenance(registry, spec), + "model_id": spec.id, + "runtime_revision": source_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "attestations": { + "complete_artifact": { + "scope": "weights+runtime", + "weights_revision": selected_checkpoint.revision, + "runtime_revision": source_revision, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" + if spec.family.weights_publication_allowed + else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + }, + "runtime_update": { + "path": "runtime-attestation.json", + "scope": "runtime-only", + "weights_repo_id": spec.fast.repo_id, + "weights_revision": spec.fast.revision, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" + if spec.family.weights_publication_allowed + else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + }, + }, + "canonical_weights": { + "index": "model.safetensors.index.json", + "shards": { + "model-00001-of-00001.safetensors": "sha256:" + "1" * 64, + }, + }, + } + _write_json(artifact / "provenance.json", provenance) + (artifact / "model.safetensors.index.json").write_bytes(b"test index") + (artifact / "model-00001-of-00001.safetensors").write_bytes(b"test shard") + runtime_files = { + relative_name: f"sha256:{hash_file(artifact.joinpath(*relative_name.split('/')))}" + for relative_name in (*safe_files, *packaged_runtime_names) + } + runtime_attestation = { + "schema_version": 2, + "scope": "runtime-only", + "model_id": spec.id, + "weights": {"repo_id": spec.fast.repo_id, "revision": spec.fast.revision}, + "runtime_revision": source_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" if spec.family.weights_publication_allowed else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + "files": runtime_files, + } + _write_json(artifact / "runtime-attestation.json", runtime_attestation) + manifest = { + relative_name: f"sha256:{hash_file(artifact.joinpath(*relative_name.split('/')))}" + for relative_name in (*safe_files, *packaged_runtime_names) + } + manifest.update( + { + "provenance.json": f"sha256:{hash_file(artifact / 'provenance.json')}", + "runtime-attestation.json": ( + f"sha256:{hash_file(artifact / 'runtime-attestation.json')}" + ), + "model.safetensors.index.json": ( + f"sha256:{hash_file(artifact / 'model.safetensors.index.json')}" + ), + "model-00001-of-00001.safetensors": ( + f"sha256:{hash_file(artifact / 'model-00001-of-00001.safetensors')}" + ), + } + ) + _write_json(artifact / "artifact-manifest.json", manifest) + return artifact + + +def test_files_only_plan_excludes_weights_and_complete_artifact_attestations( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + + assert plan.artifact_path == artifact + assert plan.repo_id == spec.fast.repo_id + assert "README.md" in plan.files + assert "config.json" in plan.files + assert "fastplms/__init__.py" in plan.files + assert "requirements.txt" in plan.files + assert "runtime-attestation.json" in plan.files + assert "artifact-manifest.json" not in plan.files + assert "provenance.json" not in plan.files + assert not any(_is_weight_path(path) for path in plan.files) + assert api.model_info_calls == [ + { + "repo_id": spec.fast.repo_id, + "revision": "main", + "files_metadata": True, + } + ] + + +@pytest.mark.parametrize( + ("relative_name", "error_match"), + ( + ("README.md", "model card differs from the current source tree"), + ( + "LICENSES/FastPLMs-Apache-2.0.txt", + "legal text differs from the current source", + ), + ), +) +def test_files_only_plan_rejects_self_attested_stale_release_text( + tmp_path: Path, + relative_name: str, + error_match: str, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + forged = artifact.joinpath(*relative_name.split("/")) + forged.write_bytes(b"self-attested but not current\n") + + manifest = json.loads( + (artifact / "artifact-manifest.json").read_text(encoding="utf-8") + ) + forged_digest = f"sha256:{hash_file(forged)}" + manifest[relative_name] = forged_digest + runtime_attestation = json.loads( + (artifact / "runtime-attestation.json").read_text(encoding="utf-8") + ) + runtime_attestation["files"][relative_name] = forged_digest + _write_json(artifact / "runtime-attestation.json", runtime_attestation) + manifest["runtime-attestation.json"] = ( + f"sha256:{hash_file(artifact / 'runtime-attestation.json')}" + ) + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match=error_match): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_self_attested_dependency_change(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + requirements = artifact / "requirements.txt" + requirements.write_text("fastplms @ git+https://example.invalid/fastplms.git\n") + _self_attest_runtime_mutation(artifact, ("requirements.txt",)) + + with pytest.raises(ArtifactError, match="direct dependency contract"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_self_attested_card_runtime_placeholder( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + readme = artifact / "README.md" + text = readme.read_text(encoding="utf-8") + readme.write_text( + text.replace("requirements.txt", "requirements.txt@", 1), + encoding="utf-8", + newline="\n", + ) + _self_attest_runtime_mutation(artifact, ("README.md",)) + + with pytest.raises(ArtifactError, match="unresolved runtime-revision placeholder"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_release_snapshot_rejects_tracked_card_deleted_from_worktree( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + source_root = tmp_path / "source" + source_root.mkdir() + card = _initialize_release_text_repository(source_root, spec) + card.unlink() + monkeypatch.setattr( + "tools.artifacts.publish.__file__", + str(source_root / "tools" / "artifacts" / "publish.py"), + ) + + with pytest.raises(ArtifactError, match="regular non-symlink file"): + _validated_release_text_snapshot( + spec, + get_model_registry(), + runtime_revision="a" * 40, + source_tree_sha256="b" * 64, + runtime_bundle_sha256="c" * 64, + ) + + +def test_release_snapshot_rejects_untracked_card_symlink_to_tracked_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + source_root = tmp_path / "source" + source_root.mkdir() + tracked_card = _initialize_release_text_repository(source_root, spec) + tracked_target = source_root / "tracked-target.md" + tracked_target.write_text(tracked_card.read_text(encoding="utf-8"), encoding="utf-8") + subprocess.run(["git", "add", tracked_target.name], cwd=source_root, check=True) + subprocess.run(["git", "commit", "-m", "tracked target"], cwd=source_root, check=True) + tracked_card.unlink() + subprocess.run(["git", "add", "-u"], cwd=source_root, check=True) + subprocess.run(["git", "commit", "-m", "remove declared card"], cwd=source_root, check=True) + try: + tracked_card.symlink_to(tracked_target) + except OSError: + tracked_card.write_text("symlink substitute\n", encoding="utf-8") + original_is_symlink = Path.is_symlink + monkeypatch.setattr( + Path, + "is_symlink", + lambda path: path == tracked_card or original_is_symlink(path), + ) + monkeypatch.setattr( + "tools.artifacts.publish.__file__", + str(source_root / "tools" / "artifacts" / "publish.py"), + ) + + with pytest.raises(ArtifactError, match="untracked at the validated revision"): + _validated_release_text_snapshot( + spec, + get_model_registry(), + runtime_revision="a" * 40, + source_tree_sha256="b" * 64, + runtime_bundle_sha256="c" * 64, + ) + + +def test_files_only_plan_rejects_self_attested_invalid_bundle_data( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + provenance = json.loads((artifact / "provenance.json").read_text(encoding="utf-8")) + (artifact / "fastplms_bundle.py").write_text( + "\n".join( + ( + '"""Forged runtime data."""', + "", + f'RUNTIME_HASH = "{provenance["runtime_bundle_sha256"]}"', + 'RUNTIME_DATA = "☃"', + "", + ) + ), + encoding="utf-8", + newline="\n", + ) + _self_attest_runtime_mutation(artifact, ("fastplms_bundle.py",)) + + with pytest.raises(ArtifactError, match="invalid base85 archive data"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_self_attested_substituted_bundle_member( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + runtime_root = artifact / "fastplms" + files = { + f"fastplms/{path.relative_to(runtime_root).as_posix()}": path.read_bytes() + for path in runtime_root.rglob("*") + if path.is_file() + } + target = sorted(files)[0] + substituted = bytearray(files[target]) + if not substituted: + raise AssertionError(f"Runtime substitution target is unexpectedly empty: {target}") + substituted[0] ^= 0x01 + files[target] = bytes(substituted) + buffer = io.BytesIO() + with ZipFile(buffer, mode="w", compression=ZIP_DEFLATED, compresslevel=9) as archive: + for relative_name, payload in sorted(files.items()): + info = ZipInfo(relative_name, date_time=(1980, 1, 1, 0, 0, 0)) + info.compress_type = ZIP_DEFLATED + info.external_attr = 0o100644 << 16 + archive.writestr(info, payload, compress_type=ZIP_DEFLATED, compresslevel=9) + archive_bytes = buffer.getvalue() + runtime_hash = hashlib.sha256(archive_bytes).hexdigest() + encoded = base64.b85encode(archive_bytes).decode("ascii") + chunks = [encoded[index : index + 100] for index in range(0, len(encoded), 100)] + (artifact / "fastplms_bundle.py").write_text( + "\n".join( + ( + '"""Generated deterministic archive of unchanged FastPLMs runtime sources."""', + "", + f'RUNTIME_HASH = "{runtime_hash}"', + "RUNTIME_DATA = (", + *(f" {chunk!r}" for chunk in chunks), + ")", + "", + ) + ), + encoding="utf-8", + newline="\n", + ) + provenance_path = artifact / "provenance.json" + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + provenance["runtime_bundle_sha256"] = runtime_hash + _write_json(provenance_path, provenance) + config_path = artifact / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["fastplms_runtime_bundle_sha256"] = runtime_hash + _write_json(config_path, config) + runtime_attestation_path = artifact / "runtime-attestation.json" + runtime_attestation = json.loads(runtime_attestation_path.read_text(encoding="utf-8")) + runtime_attestation["runtime_bundle_sha256"] = runtime_hash + _write_json(runtime_attestation_path, runtime_attestation) + _write_bootstrap(artifact / "modeling_fastplms.py", spec, runtime_hash) + _rewrite_materialized_card(artifact, spec) + _self_attest_runtime_mutation( + artifact, + ("README.md", "config.json", "fastplms_bundle.py", "modeling_fastplms.py"), + ) + manifest_path = artifact / "artifact-manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["provenance.json"] = f"sha256:{hash_file(provenance_path)}" + _write_json(manifest_path, manifest) + + with pytest.raises(ArtifactError, match="archive differs at"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_self_attested_modified_bootstrap(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + bootstrap = artifact / "modeling_fastplms.py" + bootstrap.write_text( + bootstrap.read_text(encoding="utf-8") + "\nUNAPPROVED_CODE = True\n", + encoding="utf-8", + newline="\n", + ) + _self_attest_runtime_mutation(artifact, ("modeling_fastplms.py",)) + + with pytest.raises(ArtifactError, match="bootstrap differs"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +@pytest.mark.parametrize("model_id", tuple(get_model_registry())) +def test_files_only_plan_supports_every_manifest_model( + tmp_path: Path, + model_id: str, +) -> None: + spec = get_model_registry()[model_id] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + + if spec.family.requires_complete_weight_publication: + with pytest.raises(ArtifactError, match="complete weights-plus-runtime"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + return + + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + assert plan.model_id == model_id + assert plan.repo_id == spec.fast.repo_id + assert not any(_is_weight_path(path) for path in plan.files) + + +def test_files_only_publish_uses_additions_and_parent_commit_only( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + + results = publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Update runtime files", + ) + + assert results[0].commit_oid == "b" * 40 + assert len(api.create_commit_calls) == 1 + call = api.create_commit_calls[0] + assert call["repo_id"] == spec.fast.repo_id + assert call["repo_type"] == "model" + assert call["revision"] == "main" + assert call["parent_commit"] == "a" * 40 + assert all(isinstance(operation, CommitOperationAdd) for operation in call["operations"]) + assert {operation.path_in_repo for operation in call["operations"]} == set(plan.files) + assert not any(_is_weight_path(operation.path_in_repo) for operation in call["operations"]) + + +def test_files_only_dry_run_performs_no_commit(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + + assert ( + publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Unused", + dry_run=True, + ) + == () + ) + assert api.create_commit_calls == [] + + +def test_files_only_rejects_local_model_identity_mismatch(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + config = {"fastplms_model_id": "esm2_35m"} + _write_json(artifact / "config.json", config) + manifest = json.loads((artifact / "artifact-manifest.json").read_text(encoding="utf-8")) + manifest["config.json"] = f"sha256:{hash_file(artifact / 'config.json')}" + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match="fastplms_model_id"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_rejects_remote_weight_identity_mismatch(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + + with pytest.raises(ArtifactError, match="Hub weight identity differs"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec, corrupt_weight=True), # type: ignore[arg-type] + ) + + +def test_files_only_selection_defaults_to_every_manifest_model() -> None: + registry = get_model_registry() + assert [spec.id for spec in _selected_specs(registry, ("esm2_8m",), all_models=False)] == [ + "esm2_8m" + ] + expected = list(registry) + assert [spec.id for spec in _selected_specs(registry, (), all_models=False)] == expected + assert [spec.id for spec in _selected_specs(registry, (), all_models=True)] == expected + with pytest.raises(ArtifactError, match="not both"): + _selected_specs(registry, ("esm2_8m",), all_models=True) + with pytest.raises(ArtifactError, match="Unknown model IDs"): + _selected_specs(registry, ("not-a-model",), all_models=False) + + +def test_publish_cli_requires_explicit_files_only_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("sys.argv", ["publish", "esm2_8m"]) + with pytest.raises(SystemExit, match="explicit --files-only or --complete"): + main() + + +def test_complete_publish_cli_requires_explicit_model_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("sys.argv", ["publish", "--complete"]) + with pytest.raises(SystemExit, match="requires explicit model IDs"): + main() + + +def test_files_only_plan_rejects_unlisted_and_sensitive_files(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + (artifact / "token.txt").write_text("must never be uploaded", encoding="utf-8") + + with pytest.raises(ArtifactError, match="sensitive publication path"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_unknown_manifest_path(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + unknown = artifact / "helper.exe" + unknown.write_bytes(b"binary") + manifest = json.loads((artifact / "artifact-manifest.json").read_text(encoding="utf-8")) + manifest["helper.exe"] = f"sha256:{hash_file(unknown)}" + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match="outside the publication allowlist"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_undeclared_legal_path(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + undeclared = artifact / "LICENSES" / "UNDECLARED.md" + undeclared.write_text("must not be published", encoding="utf-8") + manifest = json.loads((artifact / "artifact-manifest.json").read_text(encoding="utf-8")) + manifest["LICENSES/UNDECLARED.md"] = f"sha256:{hash_file(undeclared)}" + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match="outside the publication allowlist"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_unlisted_regular_file(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + (artifact / "notes.md").write_text("not in the manifest", encoding="utf-8") + + with pytest.raises(ArtifactError, match="file inventory differs"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_forged_registry_provenance(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + provenance_path = artifact / "provenance.json" + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + provenance["fast_checkpoint"]["revision"] = "f" * 40 + _write_json(provenance_path, provenance) + manifest = json.loads((artifact / "artifact-manifest.json").read_text(encoding="utf-8")) + manifest["provenance.json"] = f"sha256:{hash_file(provenance_path)}" + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match="current registry"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_plan_rejects_stale_runtime_revision(tmp_path: Path) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + stale = "e" * 40 + provenance_path = artifact / "provenance.json" + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + provenance["runtime_revision"] = stale + _write_json(provenance_path, provenance) + config_path = artifact / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["fastplms_runtime_revision"] = stale + _write_json(config_path, config) + attestation_path = artifact / "runtime-attestation.json" + attestation = json.loads(attestation_path.read_text(encoding="utf-8")) + attestation["runtime_revision"] = stale + _write_json(attestation_path, attestation) + _rewrite_materialized_card(artifact, spec) + _self_attest_runtime_mutation(artifact, ("README.md", "config.json")) + manifest = json.loads((artifact / "artifact-manifest.json").read_text(encoding="utf-8")) + manifest["provenance.json"] = f"sha256:{hash_file(provenance_path)}" + _write_json(artifact / "artifact-manifest.json", manifest) + + with pytest.raises(ArtifactError, match="current clean source revision"): + prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=FakeApi(spec), # type: ignore[arg-type] + ) + + +def test_files_only_publish_uses_preflighted_bytes_after_local_mutation( + tmp_path: Path, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + expected = (artifact / "README.md").read_bytes() + (artifact / "README.md").write_text("mutated after preflight", encoding="utf-8") + + publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Runtime update", + ) + + operations = api.create_commit_calls[0]["operations"] + readme = next(operation for operation in operations if operation.path_in_repo == "README.md") + assert readme.path_or_fileobj.getvalue() == expected + + +def test_files_only_publish_rejects_source_change_after_preflight( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + monkeypatch.setattr( + "tools.artifacts.publish._validated_runtime_snapshot", + lambda *_: (plan.runtime_revision, {}, "0" * 64), + ) + + with pytest.raises(ArtifactError, match="changed after publication preflight"): + publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_files_only_publish_rejects_release_text_change_after_preflight( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + monkeypatch.setattr( + "tools.artifacts.publish._validated_release_text_snapshot", + lambda *_, **__: (plan.release_revision, "0" * 64, {}), + ) + + with pytest.raises(ArtifactError, match="release texts changed after publication preflight"): + publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_files_only_publish_rejects_release_tool_change_after_preflight( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + _files_only_artifact(tmp_path, spec) + api = FakeApi(spec) + plan = prepare_files_only_plan( + spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + monkeypatch.setattr( + "tools.artifacts.publish._validated_release_tool_snapshot", + lambda *_: ("f" * 40, "e" * 64, {}), + ) + + with pytest.raises(ArtifactError, match="Release tools changed after publication preflight"): + publish_files_only( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_complete_publish_rehashes_every_file_before_atomic_commit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + first = artifact / "config.json" + second = artifact / "model.safetensors" + first.write_bytes(b"config") + second.write_bytes(b"weights") + plan = CompletePublishPlan( + model_id="toy", + repo_id="Synthyra/Toy", + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=("config.json", "model.safetensors"), + digests=( + ("config.json", f"sha256:{hash_file(first)}"), + ("model.safetensors", f"sha256:{hash_file(second)}"), + ), + ) + api = FakeApi(get_model_registry()["esm2_8m"]) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: None) + + results = publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Complete update", + ) + + assert len(results) == 1 + assert len(api.create_commit_calls) == 1 + call = api.create_commit_calls[0] + assert call["parent_commit"] == "a" * 40 + assert {operation.path_in_repo for operation in call["operations"]} == set(plan.files) + + first.write_bytes(b"changed") + with pytest.raises(ArtifactError, match="digest differs"): + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert len(api.create_commit_calls) == 1 + + +def test_complete_publish_rejects_hand_built_unknown_plan(tmp_path: Path) -> None: + artifact = tmp_path / "forged-artifact" + artifact.mkdir() + plan = CompletePublishPlan( + model_id="unregistered_model", + repo_id="Synthyra/UnregisteredModel", + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=(), + digests=(), + ) + api = FakeApi(get_model_registry()["esm2_8m"]) + + with pytest.raises(ArtifactError, match="absent from the current registry"): + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_complete_plan_deletes_only_superseded_pinned_monolith() -> None: + spec = get_model_registry()["esm2_8m"] + new_inventory = { + item.path for item in spec.fast.files if item.path != "model.safetensors" + } + new_inventory.update( + { + "model.safetensors.index.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + } + ) + + assert _obsolete_registry_pinned_paths( + spec, + FakeApi(spec).info, + new_inventory, + ) == ("model.safetensors",) + + +def test_complete_plan_deletes_superseded_pinned_shards_and_index() -> None: + spec = get_model_registry()["esm2_8m"] + legacy_index = FileDigest.parse( + "model.safetensors.index.json=git-sha1:" + "d" * 40 + ) + legacy_shard = FileDigest.parse( + "model-00001-of-00001.safetensors=sha256:" + "e" * 64 + ) + legacy = replace( + spec, + fast=replace(spec.fast, files=(*spec.fast.files, legacy_index, legacy_shard)), + ) + new_inventory = {item.path for item in spec.fast.files} + + assert _obsolete_registry_pinned_paths( + legacy, + FakeApi(legacy).info, + new_inventory, + ) == (legacy_shard.path, legacy_index.path) + + +@pytest.mark.parametrize( + "relative_name", + ( + "pytorch_model.bin", + "alternate.safetensors", + "stale/model-00001-of-00002.safetensors", + "stale/pytorch_model.bin.index.json", + ), +) +def test_complete_plan_rejects_unpinned_competing_remote_weight( + relative_name: str, +) -> None: + spec = get_model_registry()["esm2_8m"] + api = FakeApi(spec) + api.info.siblings.append( + SimpleNamespace( + rfilename=relative_name, + blob_id=None, + lfs={"sha256": "f" * 64}, + ) + ) + new_inventory = { + "model.safetensors.index.json", + "model-00001-of-00001.safetensors", + } + + with pytest.raises(ArtifactError, match="unpinned competing weight files"): + _obsolete_registry_pinned_paths(spec, api.info, new_inventory) + + +def test_complete_publish_rejects_arbitrary_delete_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + config = artifact / "config.json" + config.write_bytes(b"config") + index = artifact / "model.safetensors.index.json" + shard = artifact / "model-00001-of-00001.safetensors" + index.write_bytes(b"index") + shard.write_bytes(b"replacement") + spec = get_model_registry()["esm2_8m"] + files = (config.name, index.name, shard.name) + plan = CompletePublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=files, + digests=tuple( + (name, f"sha256:{hash_file(artifact / name)}") for name in files + ), + deletes=("unrelated-user-file.txt",), + replacement_weight_paths=(index.name, shard.name), + ) + api = FakeApi(spec) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: spec) + + with pytest.raises(ArtifactError, match="unproven obsolete weight delete"): + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_complete_publish_rejects_config_only_weight_deletion( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + config = artifact / "config.json" + config.write_bytes(b"config") + spec = get_model_registry()["esm2_8m"] + plan = CompletePublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=("config.json",), + digests=(("config.json", f"sha256:{hash_file(config)}"),), + deletes=("model.safetensors",), + ) + api = FakeApi(spec) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: spec) + + with pytest.raises(ArtifactError, match="canonical replacement weight set"): + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_complete_publish_includes_proven_guarded_weight_delete( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + index = artifact / "model.safetensors.index.json" + shard = artifact / "model-00001-of-00001.safetensors" + index.write_bytes(b"index") + shard.write_bytes(b"replacement") + spec = get_model_registry()["esm2_8m"] + files = (index.name, shard.name) + plan = CompletePublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=files, + digests=tuple( + (name, f"sha256:{hash_file(artifact / name)}") for name in files + ), + deletes=("model.safetensors",), + replacement_weight_paths=files, + ) + api = FakeApi(spec) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: spec) + + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Atomic migration", + ) + + operations = api.create_commit_calls[0]["operations"] + assert any(isinstance(operation, CommitOperationAdd) for operation in operations) + assert [ + operation.path_in_repo + for operation in operations + if isinstance(operation, CommitOperationDelete) + ] == ["model.safetensors"] + + +def test_complete_publish_uploads_validated_snapshot_after_source_mutation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact" + artifact.mkdir() + weights = artifact / "model.safetensors" + original = b"validated-weight-payload" + weights.write_bytes(original) + spec = get_model_registry()["esm2_8m"] + plan = CompletePublishPlan( + model_id="toy", + repo_id="Synthyra/Toy", + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=("model.safetensors",), + digests=(("model.safetensors", f"sha256:{hash_file(weights)}"),), + ) + + class MutatingApi(FakeApi): + uploaded: bytes | None = None + + def create_commit(self, **kwargs: Any) -> SimpleNamespace: + weights.write_bytes(b"mutated-in-place-after-validation") + operation = next( + item + for item in kwargs["operations"] + if isinstance(item, CommitOperationAdd) + ) + operation.path_or_fileobj.seek(0) + self.uploaded = operation.path_or_fileobj.read() + return super().create_commit(**kwargs) + + monkeypatch.setattr("tools.artifacts.publish._MAX_RETAINED_COMPLETE_BYTES", 1) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: None) + api = MutatingApi(spec) + + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Frozen snapshot", + ) + + assert api.uploaded == original + + +def test_required_complete_probe_groups_ankh_views( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + current_spec = get_model_registry()["ankh_base"] + spec = replace( + current_spec, + family=replace( + current_spec.family, + requires_complete_weight_publication=True, + ), + ) + artifact = tmp_path / "artifact" + artifact.mkdir() + + def fake_run(command: list[str], **_: object) -> SimpleNamespace: + cases_path = Path(command[command.index("--cases-file") + 1]) + output_path = Path(command[command.index("--output") + 1]) + cases = json.loads(cases_path.read_text(encoding="utf-8")) + assert [case["auto_class"] for case in cases] == [ + "AutoModel", + "AutoModelForSeq2SeqLM", + ] + output_path.write_text( + json.dumps({case["auto_class"]: {"state": "ok"} for case in cases}), + encoding="utf-8", + ) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr("tools.artifacts.publish.subprocess.run", fake_run) + + assert _run_required_complete_autoclass_probe(spec, artifact) == ( + "AutoModel", + "AutoModelForSeq2SeqLM", + ) + + +def test_complete_ankh_publish_rejects_missing_probe_binding( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + current_spec = get_model_registry()["ankh_base"] + spec = replace( + current_spec, + family=replace( + current_spec.family, + requires_complete_weight_publication=True, + ), + ) + artifact = tmp_path / "artifact" + artifact.mkdir() + manifest = artifact / "artifact-manifest.json" + manifest.write_text("{}\n", encoding="utf-8") + index = artifact / "model.safetensors.index.json" + shard = artifact / "model-00001-of-00001.safetensors" + index.write_bytes(b"index") + shard.write_bytes(b"replacement") + files = (manifest.name, index.name, shard.name) + plan = CompletePublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision="main", + parent_commit="a" * 40, + artifact_path=artifact, + files=files, + digests=tuple( + (name, f"sha256:{hash_file(artifact / name)}") for name in files + ), + replacement_weight_paths=(index.name, shard.name), + ) + api = FakeApi(spec) + monkeypatch.setattr(publish_module, "_revalidate_complete_plan", lambda *_: spec) + + with pytest.raises(ArtifactError, match="required AutoClass probe"): + publish_complete( + (plan,), + api=api, # type: ignore[arg-type] + commit_message="Must fail", + ) + assert not api.create_commit_calls + + +def test_complete_publication_rejects_synthetic_unresolved_checkpoint_license( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["dplm_150m"] + unresolved_family = replace( + spec.family, + checkpoint_license="Unresolved synthetic checkpoint terms", + hub_license="other", + hub_license_name="Synthetic checkpoint license unresolved", + hub_license_link="https://example.invalid/checkpoint-license", + weights_publication_allowed=False, + ) + unresolved_spec = replace(spec, family=unresolved_family) + monkeypatch.setattr( + publish_module, + "get_model_registry", + lambda: {unresolved_spec.id: unresolved_spec}, + ) + api = FakeApi(unresolved_spec) + + with pytest.raises(ArtifactError, match="unresolved weight-license"): + prepare_complete_plan( + unresolved_spec, + artifact_root=tmp_path, + revision="main", + api=api, # type: ignore[arg-type] + ) + assert not api.model_info_calls + assert not api.create_commit_calls diff --git a/tests/release/test_published_automodel.py b/tests/release/test_published_automodel.py new file mode 100644 index 0000000..90f893b --- /dev/null +++ b/tests/release/test_published_automodel.py @@ -0,0 +1,742 @@ +"""Fresh-environment validation for every built local Hub artifact.""" + +from __future__ import annotations + +import contextlib +import importlib.util +import json +import os +import subprocess +import sys +import tomllib +import pytest +from collections.abc import Iterator +from pathlib import Path +from types import SimpleNamespace +from typing import Any, ClassVar, Self + +from tools.artifacts.offline_probe import ( + ProbeCase, + _assert_complete_saved_auto_map, + _exercise, + _load_class, + _load_kwargs, + _load_model_exact, + _run_isolated_reload, + _runtime_site_packages, + _save_model_for_probe, + _semantic_config, + probe_many, +) + + +ROOT = Path(__file__).resolve().parents[2] +MANIFEST = tomllib.loads((ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8")) +PROBE = ROOT / "tools" / "artifacts" / "offline_probe.py" + +_INITIAL_WEIGHT_ALLOWANCES: dict[ + tuple[str, str], + tuple[tuple[str, ...], tuple[str, ...]], +] = { + ("ankh", "AutoModel"): ((), ()), + ("ankh", "AutoModelForMaskedLM"): ((), ()), + ("ankh", "AutoModelForSequenceClassification"): ( + ("classifier",), + (), + ), + ("ankh", "AutoModelForTokenClassification"): ( + ("classifier",), + (), + ), + ("dplm", "AutoModel"): ((), ("lm_head",)), + ("dplm", "AutoModelForSequenceClassification"): (("classifier",), ("lm_head",)), + ("dplm", "AutoModelForTokenClassification"): (("classifier",), ("lm_head",)), + ("dplm2", "AutoModel"): ((), ("lm_head",)), + ("dplm2", "AutoModelForSequenceClassification"): (("classifier",), ("lm_head",)), + ("dplm2", "AutoModelForTokenClassification"): (("classifier",), ("lm_head",)), + ("e1", "AutoModel"): ((), ("mlm_head",)), + ("e1", "AutoModelForSequenceClassification"): (("classifier",), ("mlm_head",)), + ("e1", "AutoModelForTokenClassification"): (("classifier",), ("mlm_head",)), + ("esm2", "AutoModel"): ((), ("lm_head",)), + ("esm2", "AutoModelForSequenceClassification"): (("classifier",), ("lm_head",)), + ("esm2", "AutoModelForTokenClassification"): (("classifier",), ("lm_head",)), +} + + +def _checkpoint_cases() -> list[Any]: + cases: list[Any] = [] + families = MANIFEST["families"] + for model in MANIFEST["models"]: + family_id = model["family"] + auto_map = model.get("auto_map", families[family_id]["auto_map"]) + repository_name = model["fast_repo"].split("/", maxsplit=1)[1] + auto_classes: list[dict[str, object]] = [] + for auto_class, class_path in sorted(auto_map.items()): + expected_missing, expected_unexpected = _INITIAL_WEIGHT_ALLOWANCES.get( + (family_id, auto_class), + ((), ()), + ) + auto_classes.append( + { + "auto_class": auto_class, + "class_path": class_path, + "expected_missing_key_prefixes": list(expected_missing), + "expected_unexpected_key_prefixes": list(expected_unexpected), + } + ) + marks = [pytest.mark.artifact, pytest.mark.gpu, pytest.mark.slow] + if model["size_category"] == "xlarge": + marks.append(pytest.mark.large) + cases.append( + pytest.param( + model["id"], + family_id, + repository_name, + tuple(auto_classes), + id=model["id"], + marks=marks, + ) + ) + return cases + + +def _run_probe( + *, + artifact: Path, + family: str, + implementation: str, + output: Path, + auto_class: str | None = None, + class_path: str | None = None, + auto_classes: tuple[dict[str, object], ...] | None = None, + expected_missing_key_prefixes: tuple[str, ...] = (), + expected_unexpected_key_prefixes: tuple[str, ...] = (), + attn_implementation: str | None = None, +) -> subprocess.CompletedProcess[str]: + command = [ + sys.executable, + "-I", + "-S", + str(PROBE), + "--artifact", + str(artifact), + "--family", + family, + "--bf16-execution", + MANIFEST["families"][family]["bf16_execution"], + "--implementation", + implementation, + "--output", + str(output), + ] + for path in _runtime_site_packages(): + command.extend(("--runtime-site-package", str(path))) + if attn_implementation is not None: + command.extend(("--attn-implementation", attn_implementation)) + if auto_classes is not None: + if auto_class is not None or class_path is not None: + raise ValueError("A probe must select one AutoClass or an AutoClass batch") + cases_path = output.with_name(f"{output.stem}-cases.json") + cases_path.write_text( + json.dumps(auto_classes, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + command.extend(("--cases-file", str(cases_path))) + else: + if auto_class is None or class_path is None: + raise ValueError("Single-class probes require auto_class and class_path") + command.extend(("--auto-class", auto_class, "--class-path", class_path)) + for prefix in expected_missing_key_prefixes: + command.extend(("--expected-missing-key-prefix", prefix)) + for prefix in expected_unexpected_key_prefixes: + command.extend(("--expected-unexpected-key-prefix", prefix)) + if implementation == "source": + command.extend(("--source-root", str(ROOT / "src"))) + environment = os.environ.copy() + environment.pop("PYTHONPATH", None) + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + return subprocess.run( + command, + cwd=output.parent, + env=environment, + capture_output=True, + text=True, + check=False, + ) + + +@pytest.mark.artifact +def test_generated_flash_artifacts_resolve_their_embedded_kernel_lock() -> None: + """Remote code must not depend on FastPLMs source outside the artifact.""" + + expected = (ROOT / "kernels.lock").read_bytes() + repository_names = ("ESM2-8M", "ESMplusplus_small", "DPLM-150M") + packages = [ + ROOT / "dist" / "hub" / repository_name / "fastplms" + for repository_name in repository_names + ] + missing = [ + repository_name + for repository_name, package in zip(repository_names, packages, strict=True) + if not package.is_dir() + ] + if missing: + pytest.skip( + "requires locally built Hub artifacts for " + ", ".join(missing) + ) + + for repository_name, package in zip(repository_names, packages, strict=True): + lock = package / "kernels.lock" + module_path = package / "attention" / "_kernel_lock.py" + assert lock.read_bytes() == expected + + module_spec = importlib.util.spec_from_file_location( + f"artifact_kernel_lock_{repository_name}", + module_path, + ) + assert module_spec is not None and module_spec.loader is not None + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + assert module._kernel_lock_path() == lock + + +@pytest.mark.artifact +def test_isolated_reload_rejects_incomplete_saved_remote_code(tmp_path: Path) -> None: + """A fresh interpreter must load every source file from the saved directory.""" + + artifact = tmp_path / "saved" + artifact.mkdir() + (artifact / "config.json").write_text( + json.dumps( + { + "auto_map": {"AutoConfig": "modeling_isolation.IsolationConfig"}, + "model_type": "fastplms-isolation-test", + } + ), + encoding="utf-8", + ) + (artifact / "modeling_isolation.py").write_text( + "from .support_config import IsolationConfig\n", + encoding="utf-8", + ) + support = artifact / "support_config.py" + support.write_text( + "from transformers import PretrainedConfig\n\n" + "class IsolationConfig(PretrainedConfig):\n" + " model_type = 'fastplms-isolation-test'\n", + encoding="utf-8", + ) + + complete = _run_isolated_reload( + artifact=artifact, + family="isolation-test", + bf16_execution="static_parameters", + auto_class="AutoConfig", + class_path="unused.IsolationConfig", + implementation="artifact", + source_root=None, + ) + assert set(complete) == {"config"} + + original_modeling = (artifact / "modeling_isolation.py").read_text(encoding="utf-8") + (artifact / "modeling_isolation.py").write_text( + "from fastplms import __version__\n" + "from .support_config import IsolationConfig\n", + encoding="utf-8", + ) + with pytest.raises(RuntimeError, match="Isolated saved-artifact reload failed"): + _run_isolated_reload( + artifact=artifact, + family="isolation-test", + bf16_execution="static_parameters", + auto_class="AutoConfig", + class_path="unused.IsolationConfig", + implementation="artifact", + source_root=None, + ) + (artifact / "modeling_isolation.py").write_text( + original_modeling, + encoding="utf-8", + ) + + support.unlink() + with pytest.raises(RuntimeError, match="Isolated saved-artifact reload failed"): + _run_isolated_reload( + artifact=artifact, + family="isolation-test", + bf16_execution="static_parameters", + auto_class="AutoConfig", + class_path="unused.IsolationConfig", + implementation="artifact", + source_root=None, + ) + + +def test_offline_probe_load_dtype_follows_manifest_execution_policy() -> None: + torch = SimpleNamespace( + bfloat16="bf16", + float32="fp32", + device=lambda value: value, + ) + static = _load_kwargs("esm2", "static_parameters", torch) + autocast = _load_kwargs("dplm", "fp32_parameters_autocast", torch) + assert static["dtype"] == "bf16" + assert autocast["dtype"] == "fp32" + + +def test_offline_probe_rejects_incomplete_weight_loading(tmp_path: Path) -> None: + class AutoType: + @staticmethod + def from_pretrained(*args: object, **kwargs: object) -> tuple[object, dict[str, object]]: + assert args == (tmp_path,) + assert kwargs["output_loading_info"] is True + return object(), { + "missing_keys": ["encoder.layer.0.weight"], + "unexpected_keys": [], + "mismatched_keys": [], + "error_msgs": [], + } + + with pytest.raises(RuntimeError, match="Validated AutoModel weight loading failed"): + _load_model_exact(AutoType, tmp_path, trust_remote_code=True) + + +def test_offline_probe_accepts_exact_weight_loading(tmp_path: Path) -> None: + model = object() + + class AutoType: + @staticmethod + def from_pretrained(*args: object, **kwargs: object) -> tuple[object, dict[str, object]]: + assert args == (tmp_path,) + assert kwargs["output_loading_info"] is True + return model, { + "missing_keys": [], + "unexpected_keys": [], + "mismatched_keys": [], + "error_msgs": [], + } + + assert _load_model_exact(AutoType, tmp_path, trust_remote_code=False) is model + + +def test_offline_probe_accepts_transformers_set_loading_diagnostics(tmp_path: Path) -> None: + model = object() + + class AutoType: + @staticmethod + def from_pretrained( + *args: object, + **kwargs: object, + ) -> tuple[object, dict[str, object]]: + assert args == (tmp_path,) + assert kwargs["output_loading_info"] is True + return model, { + "missing_keys": set(), + "unexpected_keys": set(), + "mismatched_keys": set(), + "error_msgs": [], + } + + assert _load_model_exact(AutoType, tmp_path) is model + + +def test_offline_probe_allows_only_declared_task_head_keys(tmp_path: Path) -> None: + model = object() + + class AutoType: + loading_info: ClassVar[dict[str, list[str]]] = { + "missing_keys": ["classifier.weight", "classifier.bias"], + "unexpected_keys": ["lm_head.decoder.weight"], + "mismatched_keys": [], + "error_msgs": [], + } + + @classmethod + def from_pretrained( + cls, + *args: object, + **kwargs: object, + ) -> tuple[object, dict[str, object]]: + assert args == (tmp_path,) + assert kwargs["output_loading_info"] is True + return model, cls.loading_info + + assert ( + _load_model_exact( + AutoType, + tmp_path, + expected_missing_key_prefixes=("classifier",), + expected_unexpected_key_prefixes=("lm_head",), + ) + is model + ) + + AutoType.loading_info = { + **AutoType.loading_info, + "missing_keys": ["classifier.weight", "encoder.layer.0.weight"], + } + with pytest.raises(RuntimeError, match=r"encoder\.layer\.0\.weight"): + _load_model_exact( + AutoType, + tmp_path, + expected_missing_key_prefixes=("classifier",), + expected_unexpected_key_prefixes=("lm_head",), + ) + + +def test_offline_probe_semantic_config_excludes_artifact_identity() -> None: + config = SimpleNamespace( + to_dict=lambda: { + "auto_map": {"AutoConfig": "modeling_fastplms.ToyConfig"}, + "hidden_size": 8, + "fastplms_model_id": "toy", + "fastplms_checkpoint_repo_id": "organization/toy", + "fastplms_checkpoint_revision": "a" * 40, + "fastplms_checkpoint_hash": "b" * 64, + "fastplms_weights_revision": "a" * 40, + "fastplms_runtime_revision": "c" * 40, + "fastplms_source_tree_sha256": "d" * 64, + "fastplms_runtime_bundle_sha256": "e" * 64, + } + ) + assert _semantic_config(config) == {"hidden_size": 8} + + +def test_structure_probe_runs_prediction_inside_bf16_autocast(tmp_path: Path) -> None: + state = {"autocast_enabled": False} + + @contextlib.contextmanager + def _autocast(*, device_type: str, dtype: str) -> Iterator[None]: + assert device_type == "cuda" + assert dtype == "bf16" + state["autocast_enabled"] = True + try: + yield + finally: + state["autocast_enabled"] = False + + class _Model: + def predict_structure( + self, + sequence: str, + *, + recycling_steps: int, + num_sampling_steps: int, + diffusion_samples: int, + ) -> str: + del self + assert state["autocast_enabled"] + assert (recycling_steps, num_sampling_steps, diffusion_samples) == (1, 2, 1) + return sequence + + torch = SimpleNamespace( + autocast=_autocast, + bfloat16="bf16", + inference_mode=contextlib.nullcontext, + ) + + output = _exercise( + _Model(), + tmp_path, + "boltz2", + "fp32_parameters_autocast", + torch, + ) + + assert output == "MSTNPKPQ" + assert not state["autocast_enabled"] + + +def test_source_probe_uses_unmodified_remote_code_save( + tmp_path: Path, +) -> None: + class FakeModel: + def __init__(self) -> None: + self.observed: list[bool] = [] + + @classmethod + def is_remote_code(cls) -> bool: + return True + + def save_pretrained(self, _path: Path, *, safe_serialization: bool) -> None: + assert safe_serialization + self.observed.append(self.is_remote_code()) + + model = FakeModel() + _save_model_for_probe(model, tmp_path, "source") + assert model.observed == [True] + assert model.is_remote_code() + + +def test_source_probe_propagates_unmodified_save_failure( + tmp_path: Path, +) -> None: + class FakeModel: + def __init__(self) -> None: + self.observed: list[bool] = [] + + @classmethod + def is_remote_code(cls) -> bool: + return True + + def save_pretrained(self, _path: Path, *, safe_serialization: bool) -> None: + assert safe_serialization + self.observed.append(self.is_remote_code()) + raise RuntimeError("synthetic save failure") + + model = FakeModel() + with pytest.raises(RuntimeError, match="synthetic save failure"): + _save_model_for_probe(model, tmp_path, "source") + + assert model.observed == [True] + assert model.is_remote_code() + + +def test_source_class_uses_transformers_autoclass_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeClass: + _auto_class: ClassVar[str | None] = None + + @classmethod + def register_for_auto_class(cls, auto_class: str) -> None: + cls._auto_class = auto_class + + monkeypatch.setattr( + "tools.artifacts.offline_probe.importlib.import_module", + lambda _name: SimpleNamespace(FakeClass=FakeClass), + ) + assert _load_class("source", "AutoModel", "fake.module.FakeClass") is FakeClass + assert FakeClass._auto_class == "AutoModel" + + +def test_saved_auto_map_must_remain_complete_and_non_null(tmp_path: Path) -> None: + expected = {"AutoConfig", "AutoModel", "AutoModelForMaskedLM"} + config = { + "auto_map": { + "AutoConfig": "modeling_saved.Config", + "AutoModel": "modeling_saved.Model", + "AutoModelForMaskedLM": "modeling_saved.ForMaskedLM", + } + } + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + _assert_complete_saved_auto_map(tmp_path, expected_auto_classes=expected) + + config["auto_map"]["AutoModelForMaskedLM"] = "modeling_saved.None" + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + with pytest.raises(RuntimeError, match="null or invalid"): + _assert_complete_saved_auto_map(tmp_path, expected_auto_classes=expected) + + +def test_encoder_only_exercise_does_not_invent_decoder_inputs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, object] = {} + + class _Tensor: + def to(self, *_args: object, **_kwargs: object) -> Self: + return self + + class _Tokenizer: + def __call__( + self, + _sequences: object, + **_kwargs: object, + ) -> dict[str, _Tensor]: + return { + "input_ids": _Tensor(), + "attention_mask": _Tensor(), + } + + class _EncoderOnly: + config = SimpleNamespace(is_encoder_decoder=True) + + def __call__(self, **kwargs: object) -> object: + observed.update(kwargs) + return kwargs["input_ids"] + + monkeypatch.setattr( + "tools.artifacts.offline_probe._tokenizer", + lambda _artifact, _config: _Tokenizer(), + ) + fake_torch = SimpleNamespace( + inference_mode=contextlib.nullcontext, + is_tensor=lambda value: isinstance(value, _Tensor), + ) + + _exercise( + _EncoderOnly(), + tmp_path, + "ankh", + "static_parameters", + fake_torch, + ) + + assert "decoder_input_ids" not in observed + assert "decoder_attention_mask" not in observed + + +def test_batch_probe_establishes_artifact_isolation_once( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + preparations: list[tuple[str, Path | None, bool]] = [] + observed: list[tuple[str, bool]] = [] + + def prepare( + implementation: str, + source_root: Path | None, + *, + reload_only: bool, + ) -> None: + preparations.append((implementation, source_root, reload_only)) + + def run_case(**kwargs: Any) -> dict[str, str]: + observed.append( + (kwargs["auto_class"], kwargs["_environment_prepared"]) + ) + return {"class_path": kwargs["class_path"]} + + monkeypatch.setattr("tools.artifacts.offline_probe._prepare_probe_environment", prepare) + monkeypatch.setattr("tools.artifacts.offline_probe.probe", run_case) + monkeypatch.setattr("tools.artifacts.offline_probe._release_case_memory", lambda: None) + + result = probe_many( + artifact=tmp_path, + family="toy", + bf16_execution="static_parameters", + cases=( + ProbeCase("AutoConfig", "toy.Config"), + ProbeCase("AutoModel", "toy.Model"), + ), + implementation="artifact", + source_root=None, + ) + + assert preparations == [("artifact", None, False)] + assert observed == [("AutoConfig", True), ("AutoModel", True)] + assert set(result) == {"AutoConfig", "AutoModel"} + + +@pytest.mark.parametrize( + ( + "model_id", + "family", + "repository_name", + "auto_classes", + ), + _checkpoint_cases(), +) +def test_local_artifact_offline_autoclass_parity( + model_id: str, + family: str, + repository_name: str, + auto_classes: tuple[dict[str, object], ...], + tmp_path: Path, +) -> None: + """Exercise one checkpoint's AutoClasses in two isolated processes.""" + + artifact = ROOT / "dist" / "hub" / repository_name + assert artifact.is_dir(), f"Missing required built artifact for {model_id}: {artifact}" + artifact_output = tmp_path / "artifact.json" + source_output = tmp_path / "source.json" + + isolated = _run_probe( + artifact=artifact, + family=family, + auto_classes=auto_classes, + implementation="artifact", + output=artifact_output, + ) + assert isolated.returncode == 0, isolated.stdout + isolated.stderr + source = _run_probe( + artifact=artifact, + family=family, + auto_classes=auto_classes, + implementation="source", + output=source_output, + ) + assert source.returncode == 0, source.stdout + source.stderr + artifact_results = json.loads(artifact_output.read_text(encoding="utf-8")) + source_results = json.loads(source_output.read_text(encoding="utf-8")) + expected_classes = {str(case["auto_class"]) for case in auto_classes} + assert set(artifact_results) == expected_classes + assert artifact_results == source_results + + +@pytest.mark.parametrize( + ("family", "repository_name", "class_path", "attn_implementation"), + ( + ( + "esm2", + "ESM2-8M", + "fastplms.models.esm2.modeling_fastesm.FastEsmModel", + "flash_attention_2", + ), + ( + "esm2", + "ESM2-8M", + "fastplms.models.esm2.modeling_fastesm.FastEsmModel", + "flash_attention_3", + ), + ( + "esm_plusplus", + "ESMplusplus_small", + "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", + "flash_attention_2", + ), + ( + "dplm", + "DPLM-150M", + "fastplms.models.dplm.modeling_dplm.DPLMModel", + "flash_attention_3", + ), + ), +) +@pytest.mark.artifact +@pytest.mark.gpu +@pytest.mark.slow +def test_local_artifact_locked_flash_backend( + family: str, + repository_name: str, + class_path: str, + attn_implementation: str, + tmp_path: Path, +) -> None: + """Exercise every advertised precompiled Flash backend in isolation.""" + + artifact = ROOT / "dist" / "hub" / repository_name + artifact_output = tmp_path / "artifact.json" + source_output = tmp_path / "source.json" + common = { + "artifact": artifact, + "family": family, + "auto_class": "AutoModel", + "class_path": class_path, + "expected_missing_key_prefixes": _INITIAL_WEIGHT_ALLOWANCES.get( + (family, "AutoModel"), + ((), ()), + )[0], + "expected_unexpected_key_prefixes": _INITIAL_WEIGHT_ALLOWANCES.get( + (family, "AutoModel"), + ((), ()), + )[1], + "attn_implementation": attn_implementation, + } + isolated = _run_probe( + **common, + implementation="artifact", + output=artifact_output, + ) + assert isolated.returncode == 0, isolated.stdout + isolated.stderr + source = _run_probe( + **common, + implementation="source", + output=source_output, + ) + assert source.returncode == 0, source.stdout + source.stderr + assert json.loads(artifact_output.read_text(encoding="utf-8")) == json.loads( + source_output.read_text(encoding="utf-8") + ) diff --git a/tests/release/test_python_support_matrix.py b/tests/release/test_python_support_matrix.py new file mode 100644 index 0000000..352886f --- /dev/null +++ b/tests/release/test_python_support_matrix.py @@ -0,0 +1,114 @@ +"""Contracts for the non-canonical Python source-support matrix.""" + +from __future__ import annotations + +from pathlib import Path + +from tools.remote.python_matrix import ( + CANONICAL_GPU_PYTHON, + OFFLINE_SMOKE_ENVIRONMENT, + PYTHON_SUPPORT_VERSIONS, + build_dependency_install_command, + build_smoke_environment, +) +from tools.remote.run import SUITES + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_python_support_versions_match_the_executed_matrix() -> None: + assert CANONICAL_GPU_PYTHON == "3.12" + assert PYTHON_SUPPORT_VERSIONS == ("3.11", "3.13", "3.14") + assert (ROOT / ".python-version").read_text(encoding="utf-8").strip() == "3.12.3" + + +def test_matrix_installs_declared_source_dependencies_with_cpu_constraints() -> None: + command = build_dependency_install_command( + "uv", + Path("/environment/bin/python"), + ROOT, + ) + + assert command == ( + "uv", + "pip", + "install", + "--python", + "/environment/bin/python", + "--torch-backend=cpu", + "-r", + str(ROOT / "requirements/profiles/runtime.in"), + "-c", + str(ROOT / "requirements/constraints/validation.txt"), + ) + + +def test_matrix_smoke_is_offline_cpu_only_and_source_isolated() -> None: + environment = build_smoke_environment( + { + "CUDA_VISIBLE_DEVICES": "0", + "HF_HUB_OFFLINE": "0", + "PYTHONPATH": "/workspace/src", + } + ) + + assert environment["CUDA_VISIBLE_DEVICES"] == "" + assert environment["HF_HUB_OFFLINE"] == "1" + assert environment["TRANSFORMERS_OFFLINE"] == "1" + assert environment["HF_DATASETS_OFFLINE"] == "1" + assert environment["PYTHONPATH"] == "" + assert environment["PYTHONNOUSERSITE"] == "1" + assert environment["UV_TORCH_BACKEND"] == "cpu" + assert set(OFFLINE_SMOKE_ENVIRONMENT).issubset(environment) + + +def test_remote_matrix_runs_source_members_in_parallel_without_raw_logs() -> None: + source = (ROOT / "tools/remote/python_matrix.py").read_text(encoding="utf-8") + + assert "ThreadPoolExecutor" in source + assert "max_workers=min(4, len(versions))" in source + assert "_output_fingerprint" in source + assert '"stdout": _output_fingerprint' in source + assert '"stderr": _output_fingerprint' in source + assert 'stage = "dependency-install"' in source + assert 'stage = "offline-cpu-source-smoke"' in source + assert 'str(python),\n "-I",' in source + assert "uv build" not in source + assert ".whl" not in source + + +def test_remote_matrix_reuses_the_single_candidate_image() -> None: + suite = SUITES["python-matrix"] + + assert suite.bake_targets == ("candidate",) + assert suite.pre_commands == () + assert suite.command == ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "candidate", + "python", + "-m", + "tools.remote.python_matrix", + "--output", + "artifacts/python-matrix.json", + "--junit-output", + "artifacts/junit/python-matrix.xml", + ) + + +def test_python_matrix_is_documented() -> None: + remote = (ROOT / "tools/remote/README.md").read_text(encoding="utf-8") + testing = (ROOT / "docs/testing.md").read_text(encoding="utf-8") + + for document in (remote, testing): + assert "python-matrix" in document + assert "3.11" in document + assert "3.13" in document + assert "3.14" in document + assert "Python 3.12" in document diff --git a/tests/release/test_reference_adapters.py b/tests/release/test_reference_adapters.py new file mode 100644 index 0000000..5968f42 --- /dev/null +++ b/tests/release/test_reference_adapters.py @@ -0,0 +1,327 @@ +"""Manifest-to-native-reference adapter release contracts.""" + +from __future__ import annotations + +import ast +import json +import re +import pytest +from pathlib import Path + +from fastplms.registry import get_model_registry + + +ROOT = Path(__file__).resolve().parents[2] +ADAPTER_ROOT = ROOT / "tests" / "parity" / "support" / "reference_adapters" +_RECONSTRUCTED_FORWARD_ATTRIBUTES = frozenset( + { + "base_z_combine", + "base_z_linear", + "fc1", + "fc2", + "k_proj", + "out_proj", + "q_proj", + "v_proj", + } +) +_FAST_LOADER_NAMES = frozenset( + { + "_load_fast", + "load_fast", + "load_fast_model", + "load_fastplms_model", + } +) + + +def _adapter_path(module: str) -> Path: + prefix = "tests.parity.support.reference_adapters." + assert module.startswith(prefix), f"Reference adapter escapes compliance package: {module}" + return ADAPTER_ROOT / f"{module.removeprefix(prefix)}.py" + + +def _dotted_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + parent = _dotted_name(node.value) + return f"{parent}.{node.attr}" if parent else node.attr + return None + + +def _is_fastplms_module(module: str | None) -> bool: + return module == "fastplms" or bool(module and module.startswith("fastplms.")) + + +def _adapter_violations(source: str, *, filename: str = "") -> list[str]: + """Return statically detectable violations of oracle independence.""" + + tree = ast.parse(source, filename=filename) + violations: list[str] = [] + for node in ast.walk(tree): + line = getattr(node, "lineno", 0) + if isinstance(node, ast.Import): + for alias in node.names: + if _is_fastplms_module(alias.name): + violations.append(f"line {line}: imports FastPLMs ({alias.name})") + if alias.name == "unittest.mock" or alias.name.startswith("unittest.mock."): + violations.append(f"line {line}: imports monkeypatch support ({alias.name})") + elif isinstance(node, ast.ImportFrom): + if _is_fastplms_module(node.module): + violations.append(f"line {line}: imports FastPLMs ({node.module})") + if node.module == "unittest.mock" or bool( + node.module and node.module.startswith("unittest.mock.") + ): + violations.append(f"line {line}: imports monkeypatch support ({node.module})") + if node.module == "tests.parity.test_model_parity": + violations.append(f"line {line}: reuses the FastPLMs parity loader") + elif isinstance(node, ast.Call): + name = _dotted_name(node.func) or "" + leaf = name.rsplit(".", 1)[-1] + if leaf in {"patch", "setattr", "delattr"} or name.endswith("patch.object"): + violations.append(f"line {line}: monkeypatches runtime state ({name})") + if leaf in _FAST_LOADER_NAMES or leaf.startswith(("_load_fastplms", "load_fastplms")): + violations.append(f"line {line}: reuses a FastPLMs loader ({name})") + if name in {"__import__", "importlib.import_module"} and node.args: + module = node.args[0] + if ( + isinstance(module, ast.Constant) + and isinstance(module.value, str) + and _is_fastplms_module(module.value) + ): + violations.append(f"line {line}: dynamically imports FastPLMs ({module.value})") + elif isinstance(node, ast.Attribute): + if node.attr in _RECONSTRUCTED_FORWARD_ATTRIBUTES: + violations.append( + f"line {line}: accesses forward implementation detail ({node.attr})" + ) + elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + raw_targets = list(node.targets) if isinstance(node, ast.Assign) else [node.target] + for target in raw_targets: + for nested in ast.walk(target): + if isinstance(nested, ast.Attribute) and nested.attr in { + "__class__", + "forward", + "from_pretrained", + }: + violations.append( + f"line {line}: replaces upstream behavior ({nested.attr})" + ) + return sorted(set(violations)) + + +def test_every_manifest_reference_adapter_has_a_pinned_loader() -> None: + for family in get_model_registry().families.values(): + path = _adapter_path(family.reference_adapter) + assert path.is_file(), f"Missing reference adapter source: {path}" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + loaders = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "load_official_model" + ] + assert len(loaders) == 1, f"{path} must define exactly one load_official_model" + arguments = {argument.arg for argument in loaders[0].args.args} + assert {"reference_repo_id", "reference_revision", "device", "dtype"}.issubset(arguments), ( + f"{path} does not require the immutable reference loading contract" + ) + + +def test_reference_adapters_are_independent_oracles() -> None: + paths = { + _adapter_path(family.reference_adapter) for family in get_model_registry().families.values() + } + failures: list[str] = [] + for path in sorted(paths): + for violation in _adapter_violations(path.read_text(encoding="utf-8"), filename=str(path)): + failures.append(f"{path.relative_to(ROOT)}: {violation}") + assert not failures, "Reference adapter independence violations:\n - " + "\n - ".join( + failures + ) + + +@pytest.mark.parametrize( + ("source", "expected_fragment"), + [ + ("import fastplms", "imports FastPLMs"), + ("from fastplms.runtime import load_model", "imports FastPLMs"), + ("import importlib\nimportlib.import_module('fastplms.models.e1')", "dynamically imports"), + ("from unittest.mock import patch", "monkeypatch support"), + ("setattr(model, 'forward', replacement)", "monkeypatches runtime state"), + ("from tests.parity.test_model_parity import _load_fast", "parity loader"), + ("model = _load_fast(spec)", "FastPLMs loader"), + ("weights = shim.base_z_combine.softmax(0)", "forward implementation detail"), + ("model.forward = replacement", "replaces upstream behavior"), + ], +) +def test_reference_adapter_static_rules_fail_closed( + source: str, + expected_fragment: str, +) -> None: + violations = _adapter_violations(source) + assert any(expected_fragment in violation for violation in violations), violations + + +def test_every_reference_adapter_maps_to_a_native_container() -> None: + dockerfile = (ROOT / "docker" / "Dockerfile").read_text(encoding="utf-8") + compose = (ROOT / "docker" / "compose.yaml").read_text(encoding="utf-8") + for family in get_model_registry().families.values(): + target = family.reference_container + assert f"AS {target}" in dockerfile + assert f" {target}:" in compose + + +def test_reference_sources_use_target_specific_build_contexts() -> None: + dockerfile = (ROOT / "docker" / "Dockerfile").read_text(encoding="utf-8") + bake = (ROOT / "docker" / "docker-bake.hcl").read_text(encoding="utf-8") + dockerignore = (ROOT / ".dockerignore").read_text(encoding="utf-8") + contexts = { + "upstream_ankh": "vendor/upstream/ankh", + "upstream_biohub_esm": "vendor/upstream/biohub-esm", + "upstream_biohub_transformers": "vendor/upstream/biohub-transformers", + "upstream_boltz": "vendor/upstream/boltz", + "upstream_dplm": "vendor/upstream/dplm", + "upstream_e1": "vendor/upstream/e1", + "upstream_fair_esm": "vendor/upstream/fair-esm", + "upstream_openfold": "vendor/upstream/openfold", + "upstream_protein_ttt": "vendor/upstream/protein-ttt", + } + assert "COPY vendor/upstream" not in dockerfile + assert "!vendor/upstream" not in dockerignore + for name, source in contexts.items(): + assert f"--from={name} --exclude=.git --exclude=.git/**" in dockerfile + assignment = rf"{re.escape(name)}\s*=\s*\"{re.escape(source)}\"" + assert re.search(assignment, bake) + + +def test_biohub_reference_is_locked_attested_and_arm64_native() -> None: + dockerfile = (ROOT / "docker" / "Dockerfile").read_text(encoding="utf-8") + bake = (ROOT / "docker" / "docker-bake.hcl").read_text(encoding="utf-8") + stage = dockerfile.split("FROM python312 AS reference-biohub-esm", maxsplit=1)[1].split( + "FROM python312 AS reference-esm2", maxsplit=1 + )[0] + + assert "git+https://github.com/Biohub/transformers" not in stage + assert "@main" not in stage + assert 'test "${TARGETARCH}" = "arm64"' in stage + assert 'test "$(uname -m)" = "aarch64"' in stage + assert "--from=biohub_biotraj_wheel" in stage + assert "biohub-reference.lock.txt" in stage + assert "biohub-reference-lock.json" in stage + assert "verify-contract" in stage + assert "materialize-wheel-lock" in stage + assert "--require-hashes --only-binary=:all: --no-deps" in stage + assert stage.count("verify-inventory") == 2 + assert stage.count("--profile final") == 2 + assert "--no-deps" in stage + assert "--no-build-isolation" in stage + assert "-e /opt/oracle/vendor" not in stage + assert "--force-reinstall" not in stage + assert "cp -a /opt/oracle/vendor/upstream/biohub-transformers" in stage + assert "cp -a /opt/oracle/vendor/upstream/biohub-esm" in stage + assert stage.count("--exclude=**/__pycache__ ") == 2 + assert stage.count("--exclude=**/__pycache__/**") == 2 + assert stage.count("--exclude=**/*.pyc") == 2 + assert stage.count("--exclude=**/*.pyo") == 2 + assert "verify-pip-check" in stage + assert "dist-info/WHEEL" not in stage + assert "manylinux2014_sbsa" not in stage + assert "test ! -e /opt/oracle/tools/remote/__init__.py" in stage + assert "import tools.remote.biohub_reference_environment" in stage + assert "import tools.remote.biohub_reference_lock" in stage + assert "import tools.remote.reference_source_attestation" in stage + assert "find_spec('tools.remote.run') is None" in stage + assert "reference_source_attestation create" in stage + assert "reference_source_attestation verify" in stage + assert stage.count("--contract /opt/oracle/biohub-esm-source-contract.json") == 2 + assert stage.count("--contract /opt/oracle/biohub-transformers-source-contract.json") == 2 + assert stage.index("pip install --require-hashes") < stage.index( + "COPY --from=reference-protocol" + ) + assert stage.index("pip install --require-hashes") < stage.index("COPY THIRD_PARTY_NOTICES.md") + environment = stage.split("ENV FASTPLMS_BIOHUB_ESM_REVISION", maxsplit=1)[1] + assert "FASTPLMS_BIOHUB_ESM_CONTRACT=" in environment + assert "FASTPLMS_BIOHUB_TRANSFORMERS_CONTRACT=" in environment + assert "FASTPLMS_BIOHUB_LOCK_CONTRACT=" in environment + assert "FASTPLMS_REFERENCE_CONTAINER_IDENTITIES=" in environment + assert "FASTPLMS_REFERENCE_CONTAINER_TARGET=reference-biohub-esm" in environment + assert "PYTHONPATH=/opt/oracle" in environment + assert "biohub-transformers/src" not in environment + assert "vendor/upstream/biohub-esm:/opt/oracle" not in environment + + assert 'target "biohub-biotraj-wheel"' in bake + assert re.search(r'target\s*=\s*"biotraj-wheel-artifact"', bake) + assert 'biohub_biotraj_wheel = "target:biohub-biotraj-wheel"' in bake + + registry = get_model_registry() + expected_sources = { + "biohub-esm": { + "import_name": "esm", + "import_root": "esm", + "package_version": "3.3.0", + "tree_sha256": "c5489f1fc58de200978803de2c38e1a78f769cb183a2ee90be833f0f4a0212e8", + }, + "biohub-transformers": { + "import_name": "transformers", + "import_root": "src/transformers", + "package_version": "4.57.6", + "tree_sha256": "28b910cc18b821870db2fb6d1c50376c2d14287ae18485080699e03fa4ba4f43", + }, + } + for source_id, expected in expected_sources.items(): + source = registry.upstreams[source_id] + contract = json.loads( + (ROOT / f"docker/constraints/{source_id}-source.json").read_text(encoding="utf-8") + ) + assert contract == { + **expected, + "schema_version": 1, + "source_revision": source.revision, + } + environment_name = source_id.removeprefix("biohub-").upper().replace("-", "_") + assert f"FASTPLMS_BIOHUB_{environment_name}_REVISION={source.revision}" in stage + + import_boundaries = { + "esm_plusplus.py": "from transformers import AutoTokenizer", + "esmfold2.py": "from transformers.models.esmfold2.configuration_esmfold2 import", + "esm3.py": "from esm.pretrained import load_local_model", + } + for filename, boundary in import_boundaries.items(): + adapter = (ADAPTER_ROOT / filename).read_text(encoding="utf-8") + load_function = adapter.split("def load_official_model", maxsplit=1)[1] + assert load_function.index("reference_sources()") < load_function.index(boundary) + + shared_gate = (ADAPTER_ROOT / "biohub_source.py").read_text(encoding="utf-8") + for source_id, expected in expected_sources.items(): + assert registry.upstreams[source_id].revision in shared_gate + assert expected["tree_sha256"] in shared_gate + assert "capture_biohub_reference_environment" in shared_gate + assert 'environment_prefix="FASTPLMS_BIOHUB_ESM"' in shared_gate + assert 'environment_prefix="FASTPLMS_BIOHUB_TRANSFORMERS"' in shared_gate + assert 'f"{environment_prefix}_CONTRACT"' in shared_gate + + structure_bundle = (ROOT / "tests/structure/support/esmfold2_bundle.py").read_text( + encoding="utf-8" + ) + reference_producer = structure_bundle.split("def produce_reference", maxsplit=1)[1].split( + "def produce_candidate", maxsplit=1 + )[0] + assert reference_producer.index("reference_sources()") < reference_producer.index( + "_load_reference_model(request, device)" + ) + assert reference_producer.index("reference_environment()") < reference_producer.index( + "_load_reference_model(request, device)" + ) + + native_reference = (ROOT / "tests/parity/support/native_reference.py").read_text( + encoding="utf-8" + ) + assert 'metadata["reference_sources"] = reference_sources' in native_reference + assert 'metadata["reference_environment"] = reference_environment' in native_reference + + esmfold2_stage = dockerfile.split( + "FROM reference-biohub-esm AS reference-esmfold2", maxsplit=1 + )[1].split("FROM python310-reference AS reference-protein-ttt", maxsplit=1)[0] + assert "FASTPLMS_REFERENCE_CONTAINER_TARGET=reference-esmfold2" in esmfold2_stage + assert "reference_sources; reference_sources(); from transformers" in esmfold2_stage diff --git a/tests/release/test_static_typing_scope.py b/tests/release/test_static_typing_scope.py new file mode 100644 index 0000000..96fc432 --- /dev/null +++ b/tests/release/test_static_typing_scope.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from pathlib import Path, PurePosixPath + + +ROOT = Path(__file__).resolve().parents[2] +SCOPE_PATH = ROOT / "tools" / "typing-critical-files.txt" +DIAGNOSTIC_SCOPE_PATH = ROOT / "tools" / "typing-diagnostic-files.txt" +EXPECTED_CRITICAL_TYPING_SCOPE = ( + "src/fastplms/registry.py", + "tools/artifacts/build.py", + "tools/artifacts/publish.py", + "tools/conversion/state_transforms.py", + "tools/source_provenance.py", +) +EXPECTED_DIAGNOSTIC_TYPING_SCOPE = ( + "benchmarks", + "examples", + "src/fastplms", + "tools", +) + + +def _scope_entries(path: Path) -> tuple[str, ...]: + return tuple( + line.strip() + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() + ) + + +def _assert_portable_scope(entries: tuple[str, ...]) -> None: + assert entries == tuple(sorted(set(entries))) + for entry in entries: + path = PurePosixPath(entry) + assert path.as_posix() == entry + assert not path.is_absolute() + assert ".." not in path.parts + assert ROOT.joinpath(*path.parts).exists() + + +def test_critical_typing_scope_is_explicit_complete_and_portable() -> None: + entries = _scope_entries(SCOPE_PATH) + + assert entries == EXPECTED_CRITICAL_TYPING_SCOPE + _assert_portable_scope(entries) + + +def test_diagnostic_typing_scope_retains_the_full_migration_surface() -> None: + entries = _scope_entries(DIAGNOSTIC_SCOPE_PATH) + + assert entries == EXPECTED_DIAGNOSTIC_TYPING_SCOPE + _assert_portable_scope(entries) + discovered = { + path.relative_to(ROOT).as_posix() + for target in entries + for path in ROOT.joinpath(*PurePosixPath(target).parts).rglob("*.py") + if path.is_file() and "__pycache__" not in path.parts + } + covered = { + path.relative_to(ROOT).as_posix() + for path in ( + *ROOT.joinpath("benchmarks").rglob("*.py"), + *ROOT.joinpath("examples").rglob("*.py"), + *ROOT.joinpath("src", "fastplms").rglob("*.py"), + *ROOT.joinpath("tools").rglob("*.py"), + ) + if path.is_file() and "__pycache__" not in path.parts + } + assert discovered == covered + assert not any( + path.startswith(("tests/", "vendor/", "build/", "dist/", "artifacts/")) + for path in discovered + ) + + +def test_manual_validation_documents_the_checked_typing_scope() -> None: + testing = (ROOT / "docs" / "testing.md").read_text(encoding="utf-8") + + assert "tools/typing-critical-files.txt" in testing + assert "--explicit-package-bases" in testing + assert "--follow-imports=silent" in testing diff --git a/tests/release/test_typing_no_regression_gate.py b/tests/release/test_typing_no_regression_gate.py new file mode 100644 index 0000000..535af85 --- /dev/null +++ b/tests/release/test_typing_no_regression_gate.py @@ -0,0 +1,519 @@ +from __future__ import annotations + +import hashlib +import json +import pytest +from collections import Counter +from collections.abc import Callable +from pathlib import Path + +from tools import typing_gate +from tools.typing_gate import ( + BASELINE_CHECKED_SOURCE_FILES, + BASELINE_ERROR_COUNT, + BASELINE_ERROR_FILE_COUNT, + BASELINE_FINGERPRINT_SHA256, + BASELINE_MYPY_VERSION, + BASELINE_PYTHON_VERSION, + BASELINE_RAW_REPORT_SHA256, + BASELINE_REVISION, + BASELINE_SOURCE_INVENTORY_SHA256, + BASELINE_SOURCE_TREE_SHA256, + MYPY_COMMAND, + REQUIRED_SCOPE_TARGETS, + Fingerprint, + TypingGateError, + baseline_payload, + compare_payload, + discover_source_files, + load_baseline, + load_scope_manifest, + main, + parse_mypy_output, +) + + +ROOT = Path(__file__).resolve().parents[2] +BASELINE = ROOT / "tools" / "typing-baselines" / "c240d8a.json" +SCOPE = ROOT / "tools" / "typing-diagnostic-files.txt" + + +def _error(path: str, line: int, message: str, code: str) -> str: + return f"{path}:{line}: error: {message} [{code}]" + + +def _snapshot(*errors: str, checked: int = 4) -> str: + files = len({line.split(":", maxsplit=1)[0] for line in errors}) + return "\n".join( + ( + *errors, + f"Found {len(errors)} errors in {files} files " + f"(checked {checked} source files)", + ) + ) + + +def _success(*, checked: int) -> str: + return f"Success: no issues found in {checked} source files" + + +def _inventory_digest(source_files: tuple[str, ...]) -> str: + payload = json.dumps(list(source_files), separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def _compare( + *, + baseline_fingerprints: Counter[Fingerprint], + candidate_text: str, + source_files: tuple[str, ...], + baseline_source_files: tuple[str, ...] | None = None, + exit_code: int = 1, +) -> dict[str, object]: + candidate = parse_mypy_output(candidate_text, REQUIRED_SCOPE_TARGETS) + baseline_files = baseline_source_files or source_files + return compare_payload( + baseline={ + "baseline_revision": BASELINE_REVISION, + "checked_source_files": len(baseline_files), + "source_files": list(baseline_files), + }, + baseline_fingerprints=baseline_fingerprints, + candidate=candidate, + scope=REQUIRED_SCOPE_TARGETS, + source_files=source_files, + source_inventory_sha256=_inventory_digest(source_files), + source_tree_sha256="a" * 64, + mypy_exit_code=exit_code, + ) + + +def _write_mutated_baseline( + tmp_path: Path, + mutation: Callable[[dict[str, object]], None], +) -> Path: + value = json.loads(BASELINE.read_text(encoding="utf-8")) + assert isinstance(value, dict) + mutation(value) + path = tmp_path / "forged-baseline.json" + path.write_text(json.dumps(value), encoding="utf-8") + return path + + +def test_parser_normalizes_paths_and_preserves_multiplicity() -> None: + text = "\n".join( + ( + "src\\fastplms\\model.py:12:4: error: Invalid value [arg-type]", + "src/fastplms/model.py:999: error: Invalid value [arg-type]", + "Found 2 errors in 1 file (checked 4 source files)", + ) + ) + + parsed = parse_mypy_output(text, REQUIRED_SCOPE_TARGETS) + + assert parsed.fingerprints == Counter( + {Fingerprint("src/fastplms/model.py", "arg-type", "Invalid value"): 2} + ) + + +def test_parser_requires_exactly_one_final_terminal_summary() -> None: + duplicate = "\n".join( + ( + _success(checked=1), + _success(checked=1), + ) + ) + nonterminal = "\n".join( + ( + _success(checked=1), + "tools/a.py: note: late output", + ) + ) + + with pytest.raises(TypingGateError, match="exactly one"): + parse_mypy_output(duplicate, REQUIRED_SCOPE_TARGETS) + with pytest.raises(TypingGateError, match="final nonblank"): + parse_mypy_output(nonterminal, REQUIRED_SCOPE_TARGETS) + + +def test_parser_rejects_unrecognized_or_out_of_scope_errors() -> None: + with pytest.raises(TypingGateError, match="Unrecognized"): + parse_mypy_output( + "tools/a.py:1: error: Missing code\n" + "Found 1 error in 1 file (checked 1 source file)", + REQUIRED_SCOPE_TARGETS, + ) + with pytest.raises(TypingGateError, match="out-of-scope"): + parse_mypy_output( + _snapshot(_error("tests/test_a.py", 1, "Bad", "arg-type"), checked=1), + REQUIRED_SCOPE_TARGETS, + ) + + +def test_payload_fails_for_one_additional_duplicate_error() -> None: + fingerprint = Fingerprint("benchmarks/run.py", "arg-type", "Bad argument") + payload = _compare( + baseline_fingerprints=Counter({fingerprint: 1}), + candidate_text=_snapshot( + _error(fingerprint.path, 1, fingerprint.message, fingerprint.code), + _error(fingerprint.path, 2, fingerprint.message, fingerprint.code), + checked=1, + ), + source_files=(fingerprint.path,), + ) + + assert payload["status"] == "failed" + assert payload["new_error_count"] == 1 + assert payload["retained_error_count"] == 1 + + +def test_payload_allows_resolved_baseline_errors() -> None: + retained = Fingerprint("tools/a.py", "arg-type", "Retained") + resolved = Fingerprint("tools/b.py", "assignment", "Resolved") + payload = _compare( + baseline_fingerprints=Counter({retained: 1, resolved: 2}), + candidate_text=_snapshot( + _error(retained.path, 9, retained.message, retained.code), + checked=2, + ), + source_files=("tools/a.py", "tools/b.py"), + ) + + assert payload["status"] == "passed" + assert payload["new_error_count"] == 0 + assert payload["resolved_error_count"] == 2 + + +def test_payload_fails_closed_on_scope_shrinkage() -> None: + debt = Fingerprint("tools/a.py", "arg-type", "Debt") + payload = _compare( + baseline_fingerprints=Counter({debt: 1}), + candidate_text=_snapshot( + _error(debt.path, 4, debt.message, debt.code), + checked=1, + ), + source_files=("tools/a.py",), + baseline_source_files=("tools/a.py", "tools/b.py"), + ) + + assert payload["status"] == "failed" + assert payload["missing_baseline_source_files"] == ["tools/b.py"] + + +def test_payload_fails_closed_on_infrastructure_exit() -> None: + payload = _compare( + baseline_fingerprints=Counter(), + candidate_text=_success(checked=1), + source_files=("tools/a.py",), + exit_code=2, + ) + + assert payload["status"] == "failed" + assert payload["failure_reasons"] == [ + "mypy exited with infrastructure status 2" + ] + + +def test_payload_rejects_fingerprint_outside_candidate_inventory() -> None: + fingerprint = Fingerprint("tools/b.py", "arg-type", "Bad") + payload = _compare( + baseline_fingerprints=Counter(), + candidate_text=_snapshot( + _error(fingerprint.path, 1, fingerprint.message, fingerprint.code), + checked=1, + ), + source_files=("tools/a.py",), + ) + + assert payload["status"] == "failed" + assert payload["fingerprint_paths_outside_candidate_inventory"] == [ + "tools/b.py" + ] + + +def test_baseline_payload_rejects_infrastructure_exit() -> None: + snapshot = parse_mypy_output( + _snapshot(_error("tools/a.py", 1, "Debt", "arg-type"), checked=1), + REQUIRED_SCOPE_TARGETS, + ) + with pytest.raises(TypingGateError, match="exit status"): + baseline_payload( + snapshot, + scope=REQUIRED_SCOPE_TARGETS, + source_files=("tools/a.py",), + revision=BASELINE_REVISION, + raw_report=b"report", + source_inventory_sha256=_inventory_digest(("tools/a.py",)), + source_tree_sha256="a" * 64, + mypy_exit_code=2, + ) + + +def test_load_baseline_rejects_forged_out_of_scope_source(tmp_path: Path) -> None: + def mutate(value: dict[str, object]) -> None: + source_files = list(value["source_files"]) + source_files.append("tests/forged.py") + value["source_files"] = sorted(source_files) + value["checked_source_files"] = len(source_files) + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="out-of-scope"): + load_baseline(path) + + +def test_load_baseline_rejects_fingerprint_absent_from_inventory( + tmp_path: Path, +) -> None: + def mutate(value: dict[str, object]) -> None: + fingerprints = list(value["fingerprints"]) + fingerprints[0] = {**fingerprints[0], "path": "tools/forged.py"} + value["fingerprints"] = fingerprints + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="not in its source inventory"): + load_baseline(path) + + +def test_load_baseline_rejects_inventory_digest_mutation(tmp_path: Path) -> None: + def mutate(value: dict[str, object]) -> None: + source_files = list(value["source_files"]) + source_files.append("tools/forged_inventory.py") + value["source_files"] = sorted(source_files) + value["checked_source_files"] = len(source_files) + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="inventory digest"): + load_baseline(path) + + +@pytest.mark.parametrize("field", ["python", "mypy"]) +def test_load_baseline_rejects_environment_mutation( + tmp_path: Path, + field: str, +) -> None: + def mutate(value: dict[str, object]) -> None: + environment = dict(value["environment"]) + environment[field] = "forged" + value["environment"] = environment + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="environment"): + load_baseline(path) + + +def test_load_baseline_rejects_environment_field_addition(tmp_path: Path) -> None: + def mutate(value: dict[str, object]) -> None: + environment = dict(value["environment"]) + environment["extra"] = "forged" + value["environment"] = environment + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="environment"): + load_baseline(path) + + +def test_load_baseline_rejects_command_mutation(tmp_path: Path) -> None: + def mutate(value: dict[str, object]) -> None: + command = list(value["mypy_command"]) + command.remove("--strict") + value["mypy_command"] = command + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="command"): + load_baseline(path) + + +def test_load_baseline_rejects_self_consistent_forged_debt( + tmp_path: Path, +) -> None: + def mutate(value: dict[str, object]) -> None: + records = [dict(record) for record in value["fingerprints"]] + records[0]["message"] = f"{records[0]['message']} forged" + value["fingerprints"] = records + counter = Counter( + { + Fingerprint( + path=str(record["path"]), + code=str(record["code"]), + message=str(record["message"]), + ): int(record["count"]) + for record in records + } + ) + value["fingerprint_sha256"] = typing_gate._fingerprint_sha256(counter) + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="immutable c240 identity"): + load_baseline(path) + + +@pytest.mark.parametrize("field", ["source_tree_sha256", "raw_report_sha256"]) +def test_load_baseline_rejects_valid_but_untrusted_attestation_digest( + tmp_path: Path, + field: str, +) -> None: + def mutate(value: dict[str, object]) -> None: + value[field] = "0" * 64 + + path = _write_mutated_baseline(tmp_path, mutate) + + with pytest.raises(TypingGateError, match="immutable c240 identity"): + load_baseline(path) + + +def test_compare_cli_writes_report_for_unparseable_infrastructure_exit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + raw_output = tmp_path / "mypy.txt" + report = tmp_path / "report.json" + monkeypatch.setattr( + typing_gate, + "_run_mypy", + lambda repo_root: (2, b"mypy: error: internal failure\n"), + ) + + exit_code = main( + ( + "compare", + "--baseline", + str(BASELINE), + "--raw-output", + str(raw_output), + "--scope-manifest", + str(SCOPE), + "--repo-root", + str(ROOT), + "--output", + str(report), + ) + ) + + assert exit_code == 1 + payload = json.loads(report.read_text(encoding="utf-8")) + assert payload["status"] == "failed" + assert payload["candidate_output_parsed"] is False + assert "infrastructure status 2" in payload["failure_reasons"][0] + + +def test_baseline_cli_executes_the_owned_mypy_command( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + for relative_name in ( + "benchmarks/a.py", + "examples/a.py", + "src/fastplms/a.py", + "tools/a.py", + ): + path = tmp_path / relative_name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("", encoding="utf-8") + scope_manifest = tmp_path / "scope.txt" + scope_manifest.write_text("\n".join(REQUIRED_SCOPE_TARGETS) + "\n", encoding="utf-8") + raw_output = tmp_path / "raw.txt" + output = tmp_path / "baseline.json" + calls: list[Path] = [] + + monkeypatch.setattr( + typing_gate, + "verified_git_head", + lambda repo_root, scope: BASELINE_REVISION, + ) + monkeypatch.setattr( + typing_gate, + "verify_git_source_identity", + lambda repo_root, scope, source_files: None, + ) + monkeypatch.setattr( + typing_gate, + "_validate_pinned_baseline_identity", + lambda value: None, + ) + + def fake_run(repo_root: Path) -> tuple[int, bytes]: + calls.append(repo_root) + return ( + 1, + _snapshot( + _error("tools/a.py", 1, "Baseline debt", "arg-type"), + checked=4, + ).encode(), + ) + + monkeypatch.setattr(typing_gate, "_run_mypy", fake_run) + + exit_code = main( + ( + "baseline", + "--raw-output", + str(raw_output), + "--scope-manifest", + str(scope_manifest), + "--repo-root", + str(tmp_path), + "--output", + str(output), + ) + ) + + assert exit_code == 0 + assert calls == [tmp_path] + assert raw_output.read_bytes().endswith(b"(checked 4 source files)") + assert json.loads(output.read_text(encoding="utf-8"))["mypy_exit_code"] == 1 + + +def test_checked_baseline_manifest_and_command_are_complete() -> None: + scope = load_scope_manifest(SCOPE) + baseline, fingerprints = load_baseline(BASELINE) + source_files = discover_source_files(ROOT, scope) + + assert scope == REQUIRED_SCOPE_TARGETS + assert MYPY_COMMAND == ( + "python", + "-m", + "mypy", + "--config-file=/dev/null", + "--python-version", + "3.12", + "--strict", + "--warn-unreachable", + "--ignore-missing-imports", + "--no-site-packages", + "--no-incremental", + "--explicit-package-bases", + "--follow-imports=silent", + "--show-error-codes", + "--no-color-output", + "--no-pretty", + "benchmarks", + "examples", + "src/fastplms", + "tools", + ) + assert baseline["baseline_revision"] == BASELINE_REVISION + assert baseline["environment"] == { + "python": BASELINE_PYTHON_VERSION, + "mypy": BASELINE_MYPY_VERSION, + } + assert baseline["mypy_command"] == list(MYPY_COMMAND) + assert baseline["mypy_exit_code"] == 1 + assert baseline["checked_source_files"] == BASELINE_CHECKED_SOURCE_FILES + assert baseline["error_count"] == BASELINE_ERROR_COUNT + assert baseline["error_file_count"] == BASELINE_ERROR_FILE_COUNT + assert baseline["error_count"] == sum(fingerprints.values()) + assert ( + baseline["source_inventory_sha256"] + == BASELINE_SOURCE_INVENTORY_SHA256 + ) + assert baseline["source_tree_sha256"] == BASELINE_SOURCE_TREE_SHA256 + assert baseline["raw_report_sha256"] == BASELINE_RAW_REPORT_SHA256 + assert baseline["fingerprint_sha256"] == BASELINE_FINGERPRINT_SHA256 + assert set(baseline["source_files"]).issubset(source_files) diff --git a/tests/release/test_validation_stack.py b/tests/release/test_validation_stack.py new file mode 100644 index 0000000..8e1463d --- /dev/null +++ b/tests/release/test_validation_stack.py @@ -0,0 +1,68 @@ +"""Exact, low-cost validation-stack checks for Hopper/SM90 release hardware.""" + +from __future__ import annotations + +import sys +import pytest +import torch +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +from tests.structure.support.hardware import hopper_sm90_fingerprint + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_candidate_prefers_the_pytorch_cudnn_runtime() -> None: + dockerfile = (ROOT / "docker" / "Dockerfile").read_text(encoding="utf-8") + wheel_library = "/opt/venv/lib/python3.12/site-packages/nvidia/cudnn/lib" + environment = next( + line.strip() + for line in dockerfile.splitlines() + if line.strip().startswith("LD_LIBRARY_PATH=") + ) + assert environment.split("=", maxsplit=1)[1].split(":", maxsplit=1)[0] == (wheel_library) + + +@pytest.mark.gpu +def test_fp8_validation_stack_uses_the_cuda13_transformer_engine_core() -> None: + expected = "2.12.0" + assert version("transformer-engine") == expected + assert version("transformer-engine-cu13") == expected + assert version("transformer-engine-torch") == expected + + with pytest.raises(PackageNotFoundError): + version("transformer-engine-cu12") + + +@pytest.mark.gpu +def test_gpu_validation_stack_is_exactly_pinned() -> None: + import transformers + + assert sys.version_info[:2] == (3, 12) + assert version("torch").split("+", maxsplit=1)[0] == "2.13.0" + assert torch.__version__.split("+", maxsplit=1)[0] == "2.13.0" + assert version("transformers") == "5.13.0" + assert transformers.__version__ == "5.13.0" + assert version("huggingface-hub") == "1.23.0" + assert version("kernels") == "0.15.2" + assert torch.version.cuda is not None + assert torch.version.cuda.startswith("13.0") + + +@pytest.mark.gpu +def test_release_hopper_sm90_gpu_is_available_without_running_a_model() -> None: + assert torch.cuda.is_available() + assert torch.cuda.device_count() >= 1 + properties = torch.cuda.get_device_properties(0) + hopper_sm90_fingerprint( + { + "cuda_device": properties.name, + "cuda_device_capability": list(torch.cuda.get_device_capability(0)), + "cuda_total_memory": int(properties.total_memory), + } + ) + # probe: (1,) + probe = torch.ones(1, device="cuda") + assert probe.item() == 1.0 diff --git a/tests/release/test_workflow_declared_inputs.py b/tests/release/test_workflow_declared_inputs.py new file mode 100644 index 0000000..f2064df --- /dev/null +++ b/tests/release/test_workflow_declared_inputs.py @@ -0,0 +1,20 @@ +"""Keep repository validation independent of hosted GitHub Actions.""" + +from __future__ import annotations + +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_repository_has_no_github_actions_workflows() -> None: + workflow_root = ROOT / ".github" / "workflows" + workflows = ( + () + if not workflow_root.exists() + else tuple((*workflow_root.glob("*.yml"), *workflow_root.glob("*.yaml"))) + ) + + assert workflows == () + assert not (ROOT / ".github" / "dependabot.yml").exists() diff --git a/tests/structure/__init__.py b/tests/structure/__init__.py new file mode 100644 index 0000000..1d8cd11 --- /dev/null +++ b/tests/structure/__init__.py @@ -0,0 +1 @@ +"""Structure-model tests.""" diff --git a/tests/structure/support/__init__.py b/tests/structure/support/__init__.py new file mode 100644 index 0000000..2f41a98 --- /dev/null +++ b/tests/structure/support/__init__.py @@ -0,0 +1 @@ +"""Shared producers and validators for native structure-compliance bundles.""" diff --git a/tests/structure/support/boltz2_bundle.py b/tests/structure/support/boltz2_bundle.py new file mode 100644 index 0000000..ff40380 --- /dev/null +++ b/tests/structure/support/boltz2_bundle.py @@ -0,0 +1,1115 @@ +"""Produce isolated official and candidate Boltz2 structure bundles. + +The reference producer uses the pinned Boltz public parser, tokenizer, +featurizer, checkpoint constructor, and forward method. The candidate producer +uses the FastPLMs repository source. They communicate only through a +manifest-derived JSON request and hash-verified safetensors bundles. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import importlib +import json +import os +import platform +import shutil +import tarfile +import tempfile +import torch +from collections.abc import Mapping, Sequence +from contextlib import contextmanager +from dataclasses import asdict +from pathlib import Path +from typing import Any, Literal +from safetensors.torch import load_file, save_file + +from tests.structure.support.state_contract import ( + exact_state_contract, + semantic_config_contract, + validate_exact_state_contract, + validate_semantic_config_contract, +) + + +schema_version = 2 +model_id = "boltz2" +reference_container = "reference-boltz2" +fold_sequence = "ACDEFGHIK" +feature_seed = 42 +fold_seed = 17 +fold_recycling_steps = 1 +fold_sampling_steps = 10 +fold_diffusion_samples = 1 +fold_parameter_dtype = "float32" +fold_compute_dtype = "bfloat16" +fold_execution = "fp32_parameters_cuda_bf16_autocast" +fold_dtype = fold_compute_dtype +diffusion_noise_generator = "numpy-pcg64-standard-normal-float32" +conformer_policy = "first-pinned-official-conformer" +_steering_policy = { + "fk_steering": False, + "physical_guidance_update": False, + "contact_guidance_update": False, + "num_gd_steps": 16, +} + +_molecule_archive = "mols.tar" +_feature_names = ( + "affinity_token_mask", + "asym_id", + "atom_backbone_feat", + "atom_pad_mask", + "atom_resolved_mask", + "atom_to_token", + "bfactor", + "contact_conditioning", + "contact_threshold", + "coords", + "cyclic_period", + "deletion_mean", + "deletion_value", + "disto_center", + "disto_coords_ensemble", + "disto_target", + "entity_id", + "frame_resolved_mask", + "frames_idx", + "has_deletion", + "method_feature", + "modified", + "mol_type", + "msa", + "msa_mask", + "msa_paired", + "plddt", + "profile", + "query_to_template", + "r_set_to_rep_atom", + "ref_atom_name_chars", + "ref_charge", + "ref_chirality", + "ref_element", + "ref_pos", + "ref_space_uid", + "res_type", + "residue_index", + "sym_id", + "template_ca", + "template_cb", + "template_frame_rot", + "template_frame_t", + "template_mask", + "template_mask_cb", + "template_mask_frame", + "template_restype", + "token_bonds", + "token_disto_mask", + "token_index", + "token_pad_mask", + "token_resolved_mask", + "token_to_center_atom", + "token_to_rep_atom", + "type_bonds", + "visibility_ids", +) +_required_outputs = ( + "complex_plddt", + "iptm", + "pae", + "pae_logits", + "pde", + "pde_logits", + "pdistogram", + "plddt", + "plddt_logits", + "ptm", + "sample_atom_coords", +) +_exact_features = ( + "affinity_token_mask", + "asym_id", + "atom_backbone_feat", + "atom_pad_mask", + "atom_resolved_mask", + "atom_to_token", + "contact_conditioning", + "contact_threshold", + "cyclic_period", + "deletion_mean", + "deletion_value", + "entity_id", + "frame_resolved_mask", + "frames_idx", + "has_deletion", + "method_feature", + "modified", + "mol_type", + "msa", + "msa_mask", + "msa_paired", + "profile", + "query_to_template", + "r_set_to_rep_atom", + "ref_atom_name_chars", + "ref_charge", + "ref_chirality", + "ref_element", + "ref_space_uid", + "res_type", + "residue_index", + "sym_id", + "template_mask", + "template_mask_cb", + "template_mask_frame", + "template_restype", + "token_bonds", + "token_disto_mask", + "token_index", + "token_pad_mask", + "token_resolved_mask", + "token_to_center_atom", + "token_to_rep_atom", + "type_bonds", + "visibility_ids", +) + + +def _canonical_json(value: Mapping[str, Any]) -> str: + return json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + + +def _request_fingerprint(request: Mapping[str, Any]) -> str: + payload = json.dumps( + request, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + # tensor: (...) + # value: (...) + value = tensor.detach().cpu().contiguous() + return value.reshape(-1).view(torch.uint8).numpy().tobytes() + + +def tensor_sha256(tensor: torch.Tensor) -> str: + """Return the exact byte digest of one tensor.""" + + # tensor: (...) + return hashlib.sha256(_tensor_bytes(tensor)).hexdigest() + + +def tensor_set_sha256(tensors: Mapping[str, torch.Tensor]) -> str: + """Hash tensor names, dtypes, shapes, and values in stable order.""" + + digest = hashlib.sha256() + for name in sorted(tensors): + # tensor: (...) + tensor = tensors[name].detach().cpu().contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(tensor.dtype).encode("ascii")) + digest.update(repr(tuple(tensor.shape)).encode("ascii")) + digest.update(_tensor_bytes(tensor)) + return digest.hexdigest() + + +def _checkpoint_metadata(checkpoint: Any) -> dict[str, Any]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + {"path": item.path, "algorithm": item.algorithm, "digest": item.digest} + for item in checkpoint.files + ], + } + + +def _upstream_metadata(upstream: Any) -> dict[str, Any]: + return { + "id": upstream.id, + "path": upstream.path, + "url": upstream.url, + "revision": upstream.revision, + "license_expression": upstream.license_expression, + } + + +def prepare_request(exchange_root: Path) -> Path: + """Write one manifest-derived request for the isolated Boltz service.""" + + from fastplms.registry import get_model_registry + + registry = get_model_registry() + spec = registry[model_id] + if spec.family.reference_container != reference_container: + raise RuntimeError("Boltz2 reference container disagrees with models.toml.") + if spec.family.upstreams != ("boltz",): + raise RuntimeError("Boltz2 must declare exactly the pinned Boltz upstream.") + if spec.family.bf16_execution != "fp32_parameters_autocast": + raise RuntimeError("Boltz2 must retain FP32 parameters under CUDA BF16 autocast.") + if _molecule_archive not in spec.official.file_map: + raise RuntimeError("Boltz2 official provenance omits its molecule archive.") + request = { + "schema_version": schema_version, + "model_id": model_id, + "architecture": spec.family.architecture, + "reference_container": spec.family.reference_container, + "official": _checkpoint_metadata(spec.official), + "candidate": _checkpoint_metadata(spec.fast), + "candidate_auto_model": spec.auto_map["AutoModel"], + "upstream": _upstream_metadata(registry.upstreams["boltz"]), + "state_transform": spec.family.state_transform, + "sequence": fold_sequence, + "feature_seed": feature_seed, + "seed": fold_seed, + "recycling_steps": fold_recycling_steps, + "sampling_steps": fold_sampling_steps, + "diffusion_samples": fold_diffusion_samples, + "diffusion_noise_generator": diffusion_noise_generator, + "conformer_policy": conformer_policy, + "steering": dict(_steering_policy), + "dtype": fold_dtype, + "parameter_dtype": fold_parameter_dtype, + "compute_dtype": fold_compute_dtype, + "execution": fold_execution, + "feature_names": list(_feature_names), + "output_names": list(_required_outputs), + } + request["request_sha256"] = _request_fingerprint(request) + path = exchange_root / "structure" / "requests" / reference_container / f"{model_id}.json" + _atomic_write_text(path, _canonical_json(request)) + return path + + +def _validate_checkpoint(source: object, *, label: str) -> None: + if not isinstance(source, Mapping): + raise ValueError(f"Boltz2 request omits {label} checkpoint metadata.") + revision = source.get("revision") + if not isinstance(revision, str) or len(revision) != 40: + raise ValueError(f"Boltz2 {label} revision is not immutable.") + files = source.get("files") + if not isinstance(files, list) or not files: + raise ValueError(f"Boltz2 {label} checkpoint has no pinned files.") + for item in files: + if not isinstance(item, Mapping) or item.get("algorithm") not in { + "git-sha1", + "sha256", + }: + raise ValueError(f"Boltz2 {label} contains an invalid file identity.") + + +def _validate_request(request: Mapping[str, Any]) -> None: + if request.get("schema_version") != schema_version: + raise ValueError("Unsupported Boltz2 structure-bundle schema.") + if request.get("model_id") != model_id: + raise ValueError(f"Unsupported Boltz2 model ID: {request.get('model_id')!r}") + if request.get("reference_container") != reference_container: + raise ValueError("Boltz2 request names the wrong reference container.") + if request.get("dtype") != fold_dtype: + raise ValueError("Boltz2 native structure parity requires BF16 mixed precision.") + if request.get("parameter_dtype") != fold_parameter_dtype: + raise ValueError("Boltz2 native structure parity requires FP32 parameters.") + if request.get("compute_dtype") != fold_compute_dtype: + raise ValueError("Boltz2 native structure parity requires BF16 compute.") + if request.get("execution") != fold_execution: + raise ValueError("Boltz2 native structure parity requires CUDA BF16 autocast.") + if request.get("steering") != _steering_policy: + raise ValueError("Boltz2 structure parity requires guidance-disabled sampling.") + if request.get("diffusion_noise_generator") != diffusion_noise_generator: + raise ValueError("Boltz2 structure parity requires portable diffusion noise.") + if request.get("conformer_policy") != conformer_policy: + raise ValueError("Boltz2 structure parity requires its pinned conformer policy.") + if tuple(request.get("feature_names", ())) != _feature_names: + raise ValueError("Boltz2 request feature schema mismatch.") + if tuple(request.get("output_names", ())) != _required_outputs: + raise ValueError("Boltz2 request output schema mismatch.") + expected = dict(request) + observed_fingerprint = expected.pop("request_sha256", None) + if observed_fingerprint != _request_fingerprint(expected): + raise ValueError("Boltz2 request fingerprint mismatch.") + _validate_checkpoint(request.get("official"), label="official") + _validate_checkpoint(request.get("candidate"), label="candidate") + upstream = request.get("upstream") + if not isinstance(upstream, Mapping) or upstream.get("id") != "boltz": + raise ValueError("Boltz2 request omits its pinned upstream.") + if len(str(upstream.get("revision", ""))) != 40: + raise ValueError("Boltz2 upstream revision is not immutable.") + official_files = { + item.get("path"): item for item in request["official"]["files"] if isinstance(item, Mapping) + } + archive = official_files.get(_molecule_archive) + if not isinstance(archive, Mapping) or archive.get("algorithm") != "sha256": + raise ValueError("Boltz2 molecule archive is not SHA-256 pinned.") + + +def load_request(path: Path) -> dict[str, Any]: + """Load and validate one manifest-derived Boltz2 request.""" + + request = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(request, dict): + raise TypeError(f"Boltz2 request must be a JSON object: {path}") + _validate_request(request) + return request + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _download_official_file(request: Mapping[str, Any], filename: str) -> Path: + from huggingface_hub import hf_hub_download + + source = request["official"] + identities = {item["path"]: item for item in source["files"] if isinstance(item, Mapping)} + if filename not in identities: + raise ValueError(f"Boltz2 request does not pin {filename!r}.") + path = Path( + hf_hub_download( + repo_id=source["repo_id"], + filename=filename, + revision=source["revision"], + ) + ) + identity = identities[filename] + if identity["algorithm"] == "sha256" and _file_sha256(path) != identity["digest"]: + raise RuntimeError(f"Boltz2 official asset hash mismatch: {filename}") + return path + + +def _required_residue_names(sequence: str) -> tuple[str, ...]: + from boltz.data import const + + names = set() + for residue in sequence: + token = const.prot_letter_to_token[residue] + names.add(token if isinstance(token, str) else const.tokens[token]) + return tuple(sorted(names)) + + +def _extract_molecules(archive_path: Path, sequence: str) -> Path: + wanted_names = _required_residue_names(sequence) + archive_hash = _file_sha256(archive_path) + cache_root = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + output = cache_root / "fastplms-boltz2" / archive_hash / "mols" + if all((output / f"{name}.pkl").is_file() for name in wanted_names): + return output + + output.mkdir(parents=True, exist_ok=True) + wanted = {f"mols/{name}.pkl": name for name in wanted_names} + found: set[str] = set() + with tarfile.open(archive_path, mode="r") as archive: + for member in archive: + name = wanted.get(member.name) + if name is None: + continue + if not member.isfile() or member.size <= 0: + raise RuntimeError(f"Invalid Boltz2 molecule member: {member.name}") + source = archive.extractfile(member) + if source is None: + raise RuntimeError(f"Cannot read Boltz2 molecule member: {member.name}") + target = output / f"{name}.pkl" + handle, temporary_name = tempfile.mkstemp( + dir=output, + prefix=f".{name}.", + suffix=".tmp", + ) + try: + with os.fdopen(handle, "wb") as stream: + shutil.copyfileobj(source, stream) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, target) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + found.add(name) + if len(found) == len(wanted_names): + break + missing = sorted(set(wanted_names).difference(found)) + if missing: + raise RuntimeError(f"Boltz2 molecule archive omits residues: {missing}") + return output + + +def _normalize_features(features: Mapping[str, object]) -> dict[str, torch.Tensor]: + missing = sorted(set(_feature_names).difference(features)) + if missing: + raise RuntimeError(f"Boltz2 featurizer omitted required inputs: {missing}") + tensors: dict[str, torch.Tensor] = {} + for name in _feature_names: + value = features[name] + if not torch.is_tensor(value): + raise TypeError(f"Boltz2 feature {name!r} is not a tensor.") + # tensors[f'feature__{name}']: (...) + tensors[f"feature__{name}"] = value.detach().cpu().contiguous().clone() + return tensors + + +def _prepare_reference_features( + request: Mapping[str, Any], + molecule_dir: Path, +) -> dict[str, torch.Tensor]: + import numpy as np + from boltz.data import const + from boltz.data.feature.featurizerv2 import Boltz2Featurizer + from boltz.data.module.inferencev2 import collate + from boltz.data.mol import load_molecules + from boltz.data.parse.fasta import parse_fasta + from boltz.data.tokenize.boltz2 import Boltz2Tokenizer + from boltz.data.types import Input + + residue_names = _required_residue_names(str(request["sequence"])) + molecules = load_molecules(molecule_dir, list(residue_names)) + for molecule in molecules.values(): + conformer_ids = sorted(conformer.GetId() for conformer in molecule.GetConformers()) + if not conformer_ids: + raise RuntimeError("Boltz2 molecule has no pinned conformer.") + for conformer_id in conformer_ids[1:]: + molecule.RemoveConformer(conformer_id) + with tempfile.TemporaryDirectory(prefix="boltz2-reference-input-") as directory: + fasta = Path(directory) / "boltz2.fasta" + fasta.write_text( + f">A|protein|empty\n{request['sequence']}\n", + encoding="utf-8", + newline="\n", + ) + target = parse_fasta(fasta, molecules, molecule_dir, boltz2=True) + input_data = Input( + structure=target.structure, + msa={}, + record=target.record, + residue_constraints=target.residue_constraints, + templates=target.templates, + extra_mols=target.extra_mols, + ) + tokenized = Boltz2Tokenizer().tokenize(input_data) + all_molecules = dict(molecules) + all_molecules.update(target.extra_mols or {}) + torch.manual_seed(int(request["feature_seed"])) + features = Boltz2Featurizer().process( + tokenized, + molecules=all_molecules, + random=np.random.default_rng(int(request["feature_seed"])), + training=False, + max_atoms=None, + max_tokens=None, + max_seqs=const.max_msa_seqs, + pad_to_max_seqs=False, + single_sequence_prop=0.0, + compute_frames=True, + inference_pocket_constraints=None, + inference_contact_constraints=None, + compute_constraint_features=True, + override_method=None, + compute_affinity=False, + ) + features = collate([features]) + normalized = _normalize_features(features) + return {name.removeprefix("feature__"): tensor for name, tensor in normalized.items()} + + +def _prepare_candidate_features(request: Mapping[str, Any]) -> dict[str, torch.Tensor]: + from fastplms.models.boltz.minimal_featurizer import build_boltz2_features + + torch.manual_seed(int(request["feature_seed"])) + features, template = build_boltz2_features(str(request["sequence"])) + if template.sequence != request["sequence"]: + raise RuntimeError("Boltz2 candidate normalized the sequence unexpectedly.") + normalized = _normalize_features(features) + return {name.removeprefix("feature__"): tensor for name, tensor in normalized.items()} + + +def _require_floating_parameter_dtype( + model: torch.nn.Module, + expected: torch.dtype, +) -> None: + floating_parameters = { + name: parameter.dtype + for name, parameter in model.named_parameters() + if parameter.is_floating_point() + } + if not floating_parameters: + raise RuntimeError("Boltz2 model has no floating parameters to validate.") + mismatches = {name: dtype for name, dtype in floating_parameters.items() if dtype != expected} + if mismatches: + sample = ", ".join( + f"{name}={dtype}" for name, dtype in list(sorted(mismatches.items()))[:8] + ) + raise RuntimeError( + f"Boltz2 requires {expected} parameter storage before BF16 autocast; found {sample}." + ) + + +def _load_reference_model( + request: Mapping[str, Any], + checkpoint_path: Path, +) -> torch.nn.Module: + from boltz.main import ( + Boltz2DiffusionParams, + BoltzSteeringParams, + MSAModuleArgs, + PairformerArgsV2, + ) + from boltz.model.models.boltz2 import Boltz2 + + predict_args = { + "recycling_steps": int(request["recycling_steps"]), + "sampling_steps": int(request["sampling_steps"]), + "diffusion_samples": int(request["diffusion_samples"]), + "max_parallel_samples": None, + "write_confidence_summary": True, + "write_full_pae": True, + "write_full_pde": True, + } + msa_args = MSAModuleArgs( + subsample_msa=True, + num_subsampled_msa=1024, + use_paired_feature=True, + ) + steering_args = asdict(BoltzSteeringParams()) + steering_args.update(request["steering"]) + model = Boltz2.load_from_checkpoint( + checkpoint_path, + strict=True, + predict_args=predict_args, + map_location="cpu", + diffusion_process_args=asdict(Boltz2DiffusionParams()), + ema=False, + use_kernels=False, + pairformer_args=asdict(PairformerArgsV2()), + msa_args=asdict(msa_args), + steering_args=steering_args, + ) + model = model.eval().to(device="cuda", dtype=torch.float32) + _require_floating_parameter_dtype(model, torch.float32) + return model + + +def _load_candidate_model(request: Mapping[str, Any]) -> torch.nn.Module: + from fastplms.registry import get_model_registry + + spec = get_model_registry()[model_id] + source = request["candidate"] + if source["repo_id"] != spec.fast.repo_id or source["revision"] != spec.fast.revision: + raise RuntimeError("Boltz2 candidate request disagrees with models.toml.") + auto_model = spec.auto_map["AutoModel"] + if request["candidate_auto_model"] != auto_model: + raise RuntimeError("Boltz2 candidate AutoModel request disagrees with models.toml.") + module_name, class_name = auto_model.rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + model = model_class.from_pretrained( + source["repo_id"], + revision=source["revision"], + dtype=torch.float32, + ) + if getattr(model.core, "use_kernels", False): + raise RuntimeError("Boltz2 compliance must use the declared eager implementation.") + for name, value in request["steering"].items(): + if model.core.steering_args.get(name) != value: + raise RuntimeError(f"Boltz2 candidate steering policy mismatch: {name}") + model = model.eval().to(device="cuda", dtype=torch.float32) + _require_floating_parameter_dtype(model, torch.float32) + return model + + +@contextmanager +def _stable_cuda_numerics(): + previous_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + previous_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + previous_benchmark = torch.backends.cudnn.benchmark + previous_deterministic = torch.backends.cudnn.deterministic + try: + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_matmul_tf32 + torch.backends.cudnn.allow_tf32 = previous_cudnn_tf32 + torch.backends.cudnn.benchmark = previous_benchmark + torch.backends.cudnn.deterministic = previous_deterministic + + +@contextmanager +def _portable_random_draws(seed: int): + """Replace every normal draw with a portable, recorded tensor stream. + + NumPy's PCG64 generator produces float32 values on the CPU. Each N tensor + is then copied to the device and dtype requested by the official or local + Boltz2 sampler. Neither wrapper calls PyTorch's random-number generator, so + the injected values are independent of PyTorch, CUDA, and global RNG state. + """ + import numpy as np + + captured: list[torch.Tensor] = [] + original_randn = torch.randn + original_randn_like = torch.randn_like + portable_rng = np.random.default_rng(seed) + + def portable_draw(template: torch.Tensor) -> torch.Tensor: + # N is the portable normal tensor with shape matching the sampler draw. + # template: (...) + N = portable_rng.standard_normal(tuple(template.shape), dtype=np.float32) + # result: (...) + result = torch.from_numpy(N).to(device=template.device, dtype=template.dtype) + result.requires_grad_(template.requires_grad) + captured.append(result.detach().cpu().contiguous().clone()) + return result + + def recording_randn(*args: Any, **kwargs: Any) -> torch.Tensor: + out = kwargs.pop("out", None) + kwargs.pop("generator", None) + # template: (*args,) + template = torch.empty(*args, **kwargs) + result = portable_draw(template) + if out is not None: + out.copy_(result) + # captured[-1]: (...) + captured[-1] = out.detach().cpu().contiguous().clone() + return out + return result + + def recording_randn_like(*args: Any, **kwargs: Any) -> torch.Tensor: + kwargs.pop("generator", None) + template = torch.empty_like(*args, **kwargs) + return portable_draw(template) + + torch.randn = recording_randn + torch.randn_like = recording_randn_like + try: + yield captured + finally: + torch.randn = original_randn + torch.randn_like = original_randn_like + + +def _run_model( + model: torch.nn.Module, + features: Mapping[str, torch.Tensor], + request: Mapping[str, Any], +) -> dict[str, torch.Tensor]: + device_features = { + name: tensor.to(device="cuda", non_blocking=False) for name, tensor in features.items() + } + + torch.manual_seed(int(request["seed"])) + torch.cuda.manual_seed_all(int(request["seed"])) + with ( + _portable_random_draws(int(request["seed"])) as captured_noise, + torch.inference_mode(), + _stable_cuda_numerics(), + torch.autocast("cuda", dtype=torch.bfloat16), + ): + output = model( + feats=device_features, + recycling_steps=int(request["recycling_steps"]), + num_sampling_steps=int(request["sampling_steps"]), + diffusion_samples=int(request["diffusion_samples"]), + max_parallel_samples=None, + run_confidence_sequentially=True, + ) + if not captured_noise: + raise RuntimeError("Boltz2 sampling did not request coordinate noise.") + if not isinstance(output, Mapping): + raise TypeError("Boltz2 forward did not return a tensor mapping.") + missing = sorted(set(_required_outputs).difference(output)) + if missing: + raise RuntimeError(f"Boltz2 forward omitted required outputs: {missing}") + tensors = {f"feature__{name}": tensor for name, tensor in features.items()} + # tensors['noise__initial_standard_normal']: (...) + tensors["noise__initial_standard_normal"] = captured_noise[0] + for index, tensor in enumerate(captured_noise): + tensors[f"noise__draw_{index:03d}"] = tensor + for name in _required_outputs: + value = output[name] + if not torch.is_tensor(value): + raise TypeError(f"Boltz2 output {name!r} is not a tensor.") + # tensors[f'output__{name}']: (...) + tensors[f"output__{name}"] = value.detach().cpu().contiguous().clone() + return tensors + + +def _environment_metadata() -> dict[str, Any]: + versions: dict[str, str | None] = {} + for package in ("boltz", "transformers", "pytorch_lightning", "rdkit"): + try: + module = importlib.import_module(package) + except ImportError: + versions[package] = None + else: + versions[package] = str(getattr(module, "__version__", "unknown")) + cuda_properties = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None + return { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "cuda_device": cuda_properties.name if cuda_properties is not None else None, + "cuda_device_capability": ( + list(torch.cuda.get_device_capability(0)) if cuda_properties is not None else None + ), + "cuda_total_memory": ( + int(cuda_properties.total_memory) if cuda_properties is not None else None + ), + "packages": versions, + } + + +def _metadata( + request: Mapping[str, Any], + *, + producer: Literal["reference", "candidate"], + model: torch.nn.Module, +) -> dict[str, Any]: + if producer == "reference": + raw_config = dict(model.hparams) + state_schema = "official_raw" + else: + raw_config = model.config + state_schema = "canonical" + return { + "schema_version": schema_version, + "producer": producer, + "model_id": model_id, + "request_sha256": request["request_sha256"], + "official": request["official"], + "candidate": request["candidate"], + "upstream": request["upstream"], + "sequence": request["sequence"], + "feature_seed": request["feature_seed"], + "seed": request["seed"], + "recycling_steps": request["recycling_steps"], + "sampling_steps": request["sampling_steps"], + "diffusion_samples": request["diffusion_samples"], + "diffusion_noise_generator": request["diffusion_noise_generator"], + "conformer_policy": request["conformer_policy"], + "steering": request["steering"], + "dtype": request["dtype"], + "parameter_dtype": request["parameter_dtype"], + "compute_dtype": request["compute_dtype"], + "execution": request["execution"], + "attention_backend": "eager", + "state_transform": request["state_transform"], + "state_schema": state_schema, + "semantic_config": semantic_config_contract(raw_config), + "state": exact_state_contract(model), + "environment": _environment_metadata(), + } + + +def canonicalize_reference_state_contract( + contract: Mapping[str, Any], + *, + expected_keys: set[str], +) -> dict[str, Any]: + """Apply the declared Boltz2 inference-core transform to compact metadata.""" + + validate_exact_state_contract(contract) + + def map_name(source_name: str) -> str | None: + if source_name.startswith("ema."): + return None + name = source_name.removeprefix("model.").removeprefix("module.") + canonical = name if name.startswith("core.") else f"core.{name}" + if canonical in expected_keys: + return canonical + bare = canonical.removeprefix("core.") + if bare.startswith(("template_module.", "bfactor_module.")): + return None + raise ValueError(f"Undeclared Boltz2 checkpoint state key: {source_name!r}.") + + tensors: dict[str, Any] = {} + for source_name, record in contract["tensors"].items(): + target = map_name(str(source_name)) + if target is None: + continue + if target in tensors: + raise ValueError(f"Boltz2 state-contract key collision for {target!r}.") + tensors[target] = record + missing = sorted(expected_keys.difference(tensors)) + if missing: + raise ValueError(f"Boltz2 reference state omits canonical keys: {missing[:20]}.") + + aliases: list[list[str]] = [] + for group in contract["aliases"]: + mapped = sorted({target for name in group if (target := map_name(str(name))) is not None}) + if len(mapped) > 1: + aliases.append(mapped) + return {"aliases": sorted(aliases), "tensors": dict(sorted(tensors.items()))} + + +def normalize_inference_config_contract(contract: Mapping[str, Any]) -> dict[str, Any]: + """Remove Boltz training and backend controls that cannot affect eager evaluation.""" + + validate_semantic_config_contract(contract) + fields = json.loads(json.dumps(contract["fields"])) + core_kwargs = fields.get("core_kwargs") + if not isinstance(core_kwargs, dict): + raise ValueError("Boltz2 semantic configuration omits core_kwargs.") + + pairformer_args = core_kwargs.get("pairformer_args") + if not isinstance(pairformer_args, dict): + raise ValueError("Boltz2 semantic configuration omits pairformer_args.") + for name in ( + "activation_checkpointing", + "dropout", + "offload_to_cpu", + "use_trifast", + ): + pairformer_args.pop(name, None) + pairformer_args.setdefault("post_layer_norm", False) + + msa_args = core_kwargs.get("msa_args") + if not isinstance(msa_args, dict): + raise ValueError("Boltz2 semantic configuration omits msa_args.") + for name in ( + "activation_checkpointing", + "msa_dropout", + "num_subsampled_msa", + "offload_to_cpu", + "subsample_msa", + "use_trifast", + "z_dropout", + ): + msa_args.pop(name, None) + msa_args.setdefault("miniformer_blocks", False) + + diffusion_args = core_kwargs.get("diffusion_process_args") + if not isinstance(diffusion_args, dict): + raise ValueError("Boltz2 semantic configuration omits diffusion_process_args.") + # Neither key affects evaluation: the first is not consumed by the + # inference module, and the second is sampled only while training. + diffusion_args.pop("mse_rotational_alignment", None) + diffusion_args.pop("step_scale_random", None) + return semantic_config_contract(fields) + + +def _atomic_write_text(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + handle, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + text=True, + ) + try: + with os.fdopen(handle, "w", encoding="utf-8", newline="\n") as stream: + stream.write(content) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, path) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def write_bundle( + output_dir: Path, + tensors: Mapping[str, torch.Tensor], + metadata: Mapping[str, Any], +) -> None: + """Atomically publish one normalized Boltz2 structure bundle.""" + + output_dir.mkdir(parents=True, exist_ok=True) + normalized = { + name: tensor.detach().cpu().contiguous().clone() for name, tensor in sorted(tensors.items()) + } + features = { + name.removeprefix("feature__"): tensor + for name, tensor in normalized.items() + if name.startswith("feature__") + } + noise_draws = { + name: tensor for name, tensor in normalized.items() if name.startswith("noise__draw_") + } + complete_metadata = dict(metadata) + complete_metadata.update( + { + "feature_sha256": tensor_set_sha256(features), + "diffusion_noise_sha256": tensor_set_sha256(noise_draws), + "diffusion_noise_draw_count": len(noise_draws), + "tensor_hashes": {name: tensor_sha256(tensor) for name, tensor in normalized.items()}, + "tensor_keys": sorted(normalized), + } + ) + handle, temporary_name = tempfile.mkstemp( + dir=output_dir, + prefix=".bundle.", + suffix=".safetensors.tmp", + ) + os.close(handle) + try: + save_file(normalized, temporary_name) + os.replace(temporary_name, output_dir / "bundle.safetensors") + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + _atomic_write_text(output_dir / "metadata.json", _canonical_json(complete_metadata)) + + +def load_bundle(path: Path) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Load one Boltz2 bundle and verify every declared tensor hash.""" + + tensor_path = path / "bundle.safetensors" + metadata_path = path / "metadata.json" + if not tensor_path.is_file() or not metadata_path.is_file(): + raise FileNotFoundError( + f"Missing Boltz2 structure bundle under {path}. Run native producers first." + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if metadata.get("schema_version") != schema_version: + raise ValueError(f"Unsupported Boltz2 bundle schema under {path}.") + tensors = load_file(tensor_path, device="cpu") + if sorted(tensors) != metadata.get("tensor_keys"): + raise ValueError(f"Tensor-key mismatch in Boltz2 bundle {path}.") + observed_hashes = {name: tensor_sha256(tensor) for name, tensor in tensors.items()} + if observed_hashes != metadata.get("tensor_hashes"): + raise ValueError(f"Tensor hash mismatch in Boltz2 bundle {path}.") + features = { + name.removeprefix("feature__"): tensor + for name, tensor in tensors.items() + if name.startswith("feature__") + } + if tensor_set_sha256(features) != metadata.get("feature_sha256"): + raise ValueError(f"Feature hash mismatch in Boltz2 bundle {path}.") + noise_draws = { + name: tensor for name, tensor in tensors.items() if name.startswith("noise__draw_") + } + if tensor_set_sha256(noise_draws) != metadata.get("diffusion_noise_sha256"): + raise ValueError(f"Diffusion-noise hash mismatch in Boltz2 bundle {path}.") + validate_exact_state_contract(metadata.get("state")) + validate_semantic_config_contract(metadata.get("semantic_config")) + return tensors, metadata + + +def produce_reference(request_path: Path, output_dir: Path) -> None: + """Run the pinned upstream Boltz2 public feature and model APIs.""" + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Official Boltz2 structure bundles require CUDA.") + archive = _download_official_file(request, _molecule_archive) + checkpoint = _download_official_file(request, "boltz2_conf.ckpt") + molecule_dir = _extract_molecules(archive, str(request["sequence"])) + features = _prepare_reference_features(request, molecule_dir) + model = _load_reference_model(request, checkpoint) + try: + metadata = _metadata(request, producer="reference", model=model) + tensors = _run_model(model, features, request) + write_bundle( + output_dir, + tensors, + metadata, + ) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def produce_candidate(request_path: Path, output_dir: Path) -> None: + """Run the FastPLMs Boltz2 feature and model APIs from repository source.""" + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Candidate Boltz2 structure bundles require CUDA.") + features = _prepare_candidate_features(request) + model = _load_candidate_model(request) + try: + metadata = _metadata(request, producer="candidate", model=model) + tensors = _run_model(model, features, request) + write_bundle( + output_dir, + tensors, + metadata, + ) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def _default_request(exchange_root: Path) -> Path: + return exchange_root / "structure" / "requests" / reference_container / f"{model_id}.json" + + +def _default_output( + exchange_root: Path, + producer: Literal["reference", "candidate"], +) -> Path: + return exchange_root / "structure" / "results" / producer / model_id / fold_dtype + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + prepare = subparsers.add_parser("prepare") + prepare.add_argument("--exchange-root", type=Path, required=True) + for name in ("produce-reference", "produce-candidate"): + producer = subparsers.add_parser(name) + producer.add_argument("--exchange-root", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Prepare a request or produce one isolated bundle.""" + + arguments = _parser().parse_args(argv) + if arguments.command == "prepare": + output = prepare_request(arguments.exchange_root) + else: + request_path = _default_request(arguments.exchange_root) + producer = "reference" if arguments.command == "produce-reference" else "candidate" + output = _default_output(arguments.exchange_root, producer) + if producer == "reference": + produce_reference(request_path, output) + else: + produce_candidate(request_path, output) + print(output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = [ + "_exact_features", + "_feature_names", + "_required_outputs", + "conformer_policy", + "diffusion_noise_generator", + "feature_seed", + "fold_compute_dtype", + "fold_diffusion_samples", + "fold_dtype", + "fold_execution", + "fold_parameter_dtype", + "fold_recycling_steps", + "fold_sampling_steps", + "fold_seed", + "fold_sequence", + "load_bundle", + "load_request", + "main", + "model_id", + "normalize_inference_config_contract", + "prepare_request", + "produce_candidate", + "produce_reference", + "reference_container", + "schema_version", + "tensor_set_sha256", + "tensor_sha256", + "write_bundle", +] diff --git a/tests/structure/support/esmfold2_bundle.py b/tests/structure/support/esmfold2_bundle.py new file mode 100644 index 0000000..9a1f6b6 --- /dev/null +++ b/tests/structure/support/esmfold2_bundle.py @@ -0,0 +1,802 @@ +"""Produce isolated ESMFold2 folding bundles for release compliance. + +The reference command imports only the pinned Biohub environment. The candidate +command imports FastPLMs only after command dispatch. Both write the same +safetensors and JSON schema so the release gate can compare them without ever +loading both implementations in one Python process. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import importlib +import json +import os +import platform +import tempfile +import torch +from collections.abc import Mapping, Sequence +from contextlib import contextmanager +from dataclasses import asdict, is_dataclass +from pathlib import Path +from typing import Any, Literal +from safetensors.torch import load_file, save_file + +from tests.structure.support.state_contract import ( + exact_state_contract, + semantic_config_contract, + validate_exact_state_contract, + validate_semantic_config_contract, +) + + +schema_version = 1 +reference_container = "reference-esmfold2" +supported_model_ids = ( + "esmfold2", + "esmfold2_fast", + "esmfold2_experimental_cutoff2025", + "esmfold2_experimental_fast_cutoff2025", +) +# Protein G B1 is a compact, experimentally characterized single-chain fold. +fold_sequence = "MQYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE" +fold_seed = 17 +# The checkpoints declare 14 inference steps. Shorter diagnostic schedules can +# leave even the official model outside physically valid C-alpha geometry. +fold_sampling_steps = 14 + +_required_outputs = ( + "atom_pad_mask", + "distogram_logits", + "iptm", + "pae", + "pae_logits", + "plddt", + "plddt_logits", + "ptm", + "sample_atom_coords", +) +_optional_outputs = ("pde_logits",) +_semantic_config_fields = ( + "confidence_head", + "d_pair", + "d_single", + "disable_msa_features", + "folding_trunk", + "force_lm_dropout_during_inference", + "inputs", + "lm_d_model", + "lm_dropout", + "lm_encoder", + "lm_num_layers", + "model_type", + "msa_encoder", + "msa_encoder_overwrite", + "n_relative_chain_bins", + "n_relative_residx_bins", + "num_diffusion_samples", + "num_loops", + "parcae", + "structure_head", + "type", +) + + +def _canonical_json(value: Mapping[str, Any]) -> str: + return json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + + +def _request_fingerprint(request: Mapping[str, Any]) -> str: + payload = json.dumps( + request, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + # tensor: (...) + # value: (...) + value = tensor.detach().cpu().contiguous() + return value.view(torch.uint8).numpy().tobytes() + + +def tensor_sha256(tensor: torch.Tensor) -> str: + """Return the content digest of one tensor without dtype coercion.""" + + # tensor: (...) + return hashlib.sha256(_tensor_bytes(tensor)).hexdigest() + + +def tensor_set_sha256(tensors: Mapping[str, torch.Tensor]) -> str: + """Hash names, dtypes, shapes, and values in deterministic key order.""" + + digest = hashlib.sha256() + for name in sorted(tensors): + # tensor: (...) + tensor = tensors[name].detach().cpu().contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(tensor.dtype).encode("ascii")) + digest.update(repr(tuple(tensor.shape)).encode("ascii")) + digest.update(_tensor_bytes(tensor)) + return digest.hexdigest() + + +def _checkpoint_metadata(checkpoint: Any) -> dict[str, Any]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + { + "path": item.path, + "algorithm": item.algorithm, + "digest": item.digest, + } + for item in checkpoint.files + ], + } + + +def prepare_requests( + exchange_root: Path, + *, + model_ids: Sequence[str] = supported_model_ids, +) -> tuple[Path, ...]: + """Write manifest-derived requests for the isolated reference service.""" + + from fastplms.registry import get_model_registry + + registry = get_model_registry() + actual_ids = tuple(spec.id for spec in registry.by_family("esmfold2")) + if actual_ids != supported_model_ids: + raise RuntimeError( + "The ESMFold2 bundle schema supports exactly the four release variants; " + f"manifest contains {actual_ids}." + ) + + request_root = exchange_root / "structure" / "requests" / reference_container + request_root.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + for model_id in model_ids: + if model_id not in supported_model_ids: + raise ValueError(f"Unsupported ESMFold2 model ID: {model_id!r}") + spec = registry[model_id] + request = { + "schema_version": schema_version, + "model_id": spec.id, + "architecture": spec.family.architecture, + "official": _checkpoint_metadata(spec.official), + "candidate": _checkpoint_metadata(spec.fast), + "candidate_auto_model": spec.auto_map["AutoModel"], + "state_transform": spec.family.state_transform, + "backbone_model": spec.family.backbone_model, + "attention_backend": "sdpa", + "deterministic_algorithms": True, + "sequence": fold_sequence, + "seed": fold_seed, + "sampling_steps": fold_sampling_steps, + } + request["request_sha256"] = _request_fingerprint(request) + path = request_root / f"{model_id}.json" + _atomic_write_text(path, _canonical_json(request)) + paths.append(path) + return tuple(paths) + + +def _validate_request(request: Mapping[str, Any]) -> None: + if request.get("schema_version") != schema_version: + raise ValueError("Unsupported ESMFold2 structure-bundle schema.") + model_id = request.get("model_id") + if model_id not in supported_model_ids: + raise ValueError(f"Unsupported ESMFold2 model ID: {model_id!r}") + expected = dict(request) + observed_fingerprint = expected.pop("request_sha256", None) + expected_fingerprint = _request_fingerprint(expected) + if observed_fingerprint != expected_fingerprint: + raise ValueError( + f"{model_id}: request fingerprint mismatch " + f"{observed_fingerprint!r} != {expected_fingerprint!r}." + ) + for source_name in ("official", "candidate"): + source = request.get(source_name) + if not isinstance(source, Mapping): + raise ValueError(f"{model_id}: missing {source_name} checkpoint metadata.") + revision = source.get("revision") + if not isinstance(revision, str) or len(revision) != 40: + raise ValueError(f"{model_id}: {source_name} revision is not immutable.") + if request.get("state_transform") != "identity": + raise ValueError(f"{model_id}: ESMFold2 requires the identity state transform.") + backbone_model = request.get("backbone_model") + if not isinstance(backbone_model, str) or not backbone_model: + raise ValueError(f"{model_id}: ESMFold2 requires a logical backbone model ID.") + if request.get("attention_backend") != "sdpa": + raise ValueError(f"{model_id}: ESMFold2 folding compliance requires SDPA.") + if request.get("deterministic_algorithms") is not True: + raise ValueError(f"{model_id}: ESMFold2 folding compliance requires determinism.") + + +def load_request(path: Path) -> dict[str, Any]: + """Load and validate one manifest-derived folding request.""" + + request = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(request, dict): + raise TypeError(f"ESMFold2 request must be a JSON object: {path}") + _validate_request(request) + return request + + +def _is_experimental(request: Mapping[str, Any]) -> bool: + return "experimental" in str(request["model_id"]) + + +def _model_package(model: torch.nn.Module) -> str: + return model.__class__.__module__.rsplit(".", maxsplit=1)[0] + + +@contextmanager +def _stable_cuda_numerics(): + """Use deterministic CUDA algorithms with TF32 and autotuning disabled.""" + + previous_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + previous_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + previous_benchmark = torch.backends.cudnn.benchmark + previous_deterministic = torch.backends.cudnn.deterministic + previous_algorithms = torch.are_deterministic_algorithms_enabled() + try: + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.use_deterministic_algorithms(True) + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_matmul_tf32 + torch.backends.cudnn.allow_tf32 = previous_cudnn_tf32 + torch.backends.cudnn.benchmark = previous_benchmark + torch.backends.cudnn.deterministic = previous_deterministic + torch.use_deterministic_algorithms(previous_algorithms) + + +def _run_fold( + model: torch.nn.Module, + request: Mapping[str, Any], +) -> dict[str, torch.Tensor]: + feature_module = importlib.import_module(f"{_model_package(model)}.protein_utils") + common_module = importlib.import_module(f"{_model_package(model)}.modeling_esmfold2_common") + cpu_features = feature_module.prepare_protein_features(request["sequence"]) + if not isinstance(cpu_features, Mapping) or not all( + torch.is_tensor(value) for value in cpu_features.values() + ): + raise TypeError("Official feature preparation did not return a tensor mapping.") + device = next(model.parameters()).device + device_features = {name: tensor.to(device=device) for name, tensor in cpu_features.items()} + + captured_noise: list[torch.Tensor] = [] + original_randn = torch.randn + + def recording_randn(*args: Any, **kwargs: Any) -> torch.Tensor: + tensor = original_randn(*args, **kwargs) + if not captured_noise and tensor.ndim == 3 and tensor.shape[-1] == 3: + captured_noise.append(tensor.detach().cpu().contiguous().clone()) + return tensor + + forward_kwargs: dict[str, Any] = { + "num_loops": 1, + "num_sampling_steps": int(request["sampling_steps"]), + "num_diffusion_samples": 1, + } + if _is_experimental(request): + forward_kwargs.update({"calculate_confidence": True, "seed": int(request["seed"])}) + else: + forward_kwargs.update( + { + "msa_column_mask_rate": 0.0, + "msa_subsample_at_inference": False, + } + ) + + torch.randn = recording_randn + try: + with ( + _stable_cuda_numerics(), + common_module._seed_context(int(request["seed"])), + torch.inference_mode(), + ): + output = model(**device_features, **forward_kwargs) + finally: + torch.randn = original_randn + + if len(captured_noise) != 1: + raise RuntimeError( + f"Expected one initial diffusion-noise tensor, captured {len(captured_noise)}." + ) + if not isinstance(output, Mapping): + raise TypeError("ESMFold2 forward did not return a tensor mapping.") + missing_outputs = sorted(set(_required_outputs).difference(output)) + if missing_outputs: + raise RuntimeError(f"ESMFold2 fold omitted required outputs: {missing_outputs}") + + tensors = { + f"feature__{name}": tensor.detach().cpu().contiguous().clone() + for name, tensor in cpu_features.items() + } + # tensors['noise__initial_standard_normal']: (...) + tensors["noise__initial_standard_normal"] = captured_noise[0] + for name in (*_required_outputs, *_optional_outputs): + value = output.get(name) + if torch.is_tensor(value): + # tensors[f'output__{name}']: (...) + tensors[f"output__{name}"] = value.detach().cpu().contiguous().clone() + return tensors + + +def _environment_metadata() -> dict[str, Any]: + import transformers + + transformer_engine_version = None + try: + import transformer_engine + + transformer_engine_version = transformer_engine.__version__ + except ImportError: + pass + cuda_properties = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None + return { + "python": platform.python_version(), + "torch": torch.__version__, + "transformers": transformers.__version__, + "cuda_runtime": torch.version.cuda, + "cuda_device": cuda_properties.name if cuda_properties is not None else None, + "cuda_device_capability": ( + list(torch.cuda.get_device_capability(0)) if cuda_properties is not None else None + ), + "cuda_total_memory": ( + int(cuda_properties.total_memory) if cuda_properties is not None else None + ), + "transformer_engine": transformer_engine_version, + } + + +def _precision_status(value: object) -> dict[str, Any]: + if is_dataclass(value): + status = asdict(value) + else: + status = { + name: getattr(value, name) + for name in ( + "requested", + "resolved", + "reason", + "device", + "transformer_engine_version", + ) + if hasattr(value, name) + } + return { + name: str(item) if isinstance(item, torch.device) else item for name, item in status.items() + } + + +def _esmfold2_semantic_config( + config: object, + *, + backbone_model: str, +) -> dict[str, Any]: + """Normalize official and mirrored configs to inference semantics. + + Transformers generation defaults, candidate attention/precision policy, + and the official-versus-mirror Hub identifier are packaging or runtime + concerns. The logical backbone identity comes from the typed manifest. + """ + + if isinstance(config, Mapping): + raw = dict(config) + else: + to_dict = getattr(config, "to_dict", None) + if not callable(to_dict): + raise TypeError("ESMFold2 semantic config must be a mapping or expose to_dict().") + raw = to_dict() + missing = [name for name in _semantic_config_fields if name not in raw] + if missing: + raise ValueError(f"ESMFold2 semantic config omits required fields: {missing}.") + fields = {name: raw[name] for name in _semantic_config_fields} + fields["backbone_model"] = backbone_model + return semantic_config_contract(fields) + + +def _load_reference_model( + request: Mapping[str, Any], + device: torch.device, +) -> torch.nn.Module: + from transformers.models.esmfold2.configuration_esmfold2 import ESMFold2Config + from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model + from transformers.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, + ) + + source = request["official"] + config = ESMFold2Config.from_pretrained( + source["repo_id"], + revision=source["revision"], + ) + model_class = ESMFold2ExperimentalModel if config.type == "experimental" else ESMFold2Model + model = model_class.from_pretrained( + source["repo_id"], + revision=source["revision"], + config=config, + load_esmc=False, + dtype=torch.float32, + ) + model = model.eval().to(device=device, dtype=torch.float32) + if config.type == "experimental": + model.load_esmc(model.config.esmc_id) + else: + model.load_esmc(model.config.esmc_id, precision="bf16") + return model + + +def _load_candidate_model( + request: Mapping[str, Any], + device: torch.device, + precision: Literal["bf16", "fp8"], +) -> torch.nn.Module: + from fastplms.registry import get_model_registry + + spec = get_model_registry()[request["model_id"]] + source = request["candidate"] + if source["repo_id"] != spec.fast.repo_id or source["revision"] != spec.fast.revision: + raise RuntimeError(f"{spec.id}: candidate request disagrees with models.toml.") + if request["candidate_auto_model"] != spec.auto_map["AutoModel"]: + raise RuntimeError(f"{spec.id}: candidate AutoModel request disagrees with models.toml.") + module_name, class_name = spec.auto_map["AutoModel"].rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + model = model_class.from_pretrained( + source["repo_id"], + revision=source["revision"], + load_esmc=False, + dtype=torch.float32, + ) + model = model.eval().to(device=device) + model.reload_esmc(precision=precision, device=device) + return model + + +def _base_metadata( + request: Mapping[str, Any], + *, + producer: Literal["reference", "candidate"], + requested_precision: str, + resolved_precision: str, + precision_status: Mapping[str, Any], + model: torch.nn.Module, +) -> dict[str, Any]: + return { + "schema_version": schema_version, + "producer": producer, + "model_id": request["model_id"], + "request_sha256": request["request_sha256"], + "official": request["official"], + "candidate": request["candidate"], + "sequence": request["sequence"], + "seed": request["seed"], + "sampling_steps": request["sampling_steps"], + "state_transform": request["state_transform"], + "attention_backend": request["attention_backend"], + "deterministic_algorithms": request["deterministic_algorithms"], + "semantic_config": _esmfold2_semantic_config( + model.config, + backbone_model=str(request["backbone_model"]), + ), + "state": exact_state_contract(model, excluded_prefixes=("_esmc.",)), + "requested_precision": requested_precision, + "resolved_precision": resolved_precision, + "precision_status": dict(precision_status), + "environment": _environment_metadata(), + } + + +def _atomic_write_text(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + handle, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + text=True, + ) + try: + with os.fdopen(handle, "w", encoding="utf-8", newline="\n") as stream: + stream.write(content) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, path) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def write_bundle( + output_dir: Path, + tensors: Mapping[str, torch.Tensor], + metadata: Mapping[str, Any], +) -> None: + """Atomically publish one normalized structure bundle.""" + + output_dir.mkdir(parents=True, exist_ok=True) + normalized = { + name: tensor.detach().cpu().contiguous().clone() for name, tensor in sorted(tensors.items()) + } + feature_tensors = { + name.removeprefix("feature__"): tensor + for name, tensor in normalized.items() + if name.startswith("feature__") + } + complete_metadata = dict(metadata) + complete_metadata.update( + { + "feature_sha256": tensor_set_sha256(feature_tensors), + "diffusion_noise_sha256": tensor_sha256(normalized["noise__initial_standard_normal"]), + "tensor_hashes": {name: tensor_sha256(tensor) for name, tensor in normalized.items()}, + "tensor_keys": sorted(normalized), + } + ) + + bundle_path = output_dir / "bundle.safetensors" + handle, temporary_name = tempfile.mkstemp( + dir=output_dir, + prefix=".bundle.", + suffix=".safetensors.tmp", + ) + os.close(handle) + try: + save_file(normalized, temporary_name) + os.replace(temporary_name, bundle_path) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + _atomic_write_text(output_dir / "metadata.json", _canonical_json(complete_metadata)) + + +def load_bundle(path: Path) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Load a bundle and verify every declared tensor digest.""" + + tensor_path = path / "bundle.safetensors" + metadata_path = path / "metadata.json" + if not tensor_path.is_file() or not metadata_path.is_file(): + raise FileNotFoundError( + f"Missing ESMFold2 structure bundle under {path}. Run native producers first." + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if metadata.get("schema_version") != schema_version: + raise ValueError(f"Unsupported ESMFold2 bundle schema under {path}.") + tensors = load_file(tensor_path, device="cpu") + if sorted(tensors) != metadata.get("tensor_keys"): + raise ValueError(f"Tensor-key mismatch in ESMFold2 bundle {path}.") + expected_hashes = metadata.get("tensor_hashes") + observed_hashes = {name: tensor_sha256(tensor) for name, tensor in tensors.items()} + if observed_hashes != expected_hashes: + raise ValueError(f"Tensor hash mismatch in ESMFold2 bundle {path}.") + features = { + name.removeprefix("feature__"): tensor + for name, tensor in tensors.items() + if name.startswith("feature__") + } + if tensor_set_sha256(features) != metadata.get("feature_sha256"): + raise ValueError(f"Feature hash mismatch in ESMFold2 bundle {path}.") + if tensor_sha256(tensors["noise__initial_standard_normal"]) != metadata.get( + "diffusion_noise_sha256" + ): + raise ValueError(f"Diffusion-noise hash mismatch in ESMFold2 bundle {path}.") + validate_exact_state_contract(metadata.get("state")) + validate_semantic_config_contract(metadata.get("semantic_config")) + return tensors, metadata + + +def produce_reference(request_path: Path, output_dir: Path) -> None: + """Produce an official BF16 bundle in the native Biohub service.""" + + from tests.parity.support.reference_adapters.biohub_source import ( + reference_environment, + reference_sources, + ) + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Official ESMFold2 structure bundles require CUDA.") + sources = reference_sources() + locked_environment = reference_environment() + device = torch.device("cuda") + model = _load_reference_model(request, device) + try: + metadata = _base_metadata( + request, + producer="reference", + requested_precision="bf16", + resolved_precision="bf16", + precision_status={ + "requested": "bf16", + "resolved": "bf16", + "reason": ( + "Pinned Biohub config retains FP32 folding weights; the public CUDA " + "forward applies internal BF16 autocast and loads ESMC in BF16." + ), + "device": str(device), + "transformer_engine_version": None, + "weight_dtype": "float32", + "compute_dtype": "bfloat16", + "execution": "fp32_parameters_internal_bf16_autocast", + }, + model=model, + ) + metadata["reference_sources"] = sources + metadata["reference_environment"] = locked_environment + tensors = _run_fold(model, request) + write_bundle(output_dir, tensors, metadata) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def produce_candidate( + request_path: Path, + output_dir: Path, + *, + precision: Literal["bf16", "fp8"], +) -> None: + """Produce one FastPLMs BF16 or FP8 bundle in the candidate service.""" + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Candidate ESMFold2 structure bundles require CUDA.") + device = torch.device("cuda") + model = _load_candidate_model(request, device, precision) + try: + status = _precision_status(model.esmc_precision_status) + if status.get("requested") != precision or status.get("resolved") != precision: + raise RuntimeError( + f"{request['model_id']}: requested {precision}, resolved status is {status}." + ) + metadata = _base_metadata( + request, + producer="candidate", + requested_precision=precision, + resolved_precision=str(status["resolved"]), + precision_status=status, + model=model, + ) + tensors = _run_fold(model, request) + write_bundle(output_dir, tensors, metadata) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def _default_request(exchange_root: Path, model_id: str) -> Path: + return exchange_root / "structure" / "requests" / reference_container / f"{model_id}.json" + + +def _all_prepared_requests(exchange_root: Path) -> tuple[Path, ...]: + """Return the exact four prepared requests in the release-schema order.""" + + request_root = exchange_root / "structure" / "requests" / reference_container + available = {path.stem: path for path in sorted(request_root.glob("*.json"))} + missing = sorted(set(supported_model_ids).difference(available)) + unexpected = sorted(set(available).difference(supported_model_ids)) + if missing or unexpected: + raise FileNotFoundError( + "ESMFold2 prepared-request inventory differs from the release schema: " + f"missing={missing}, unexpected={unexpected}." + ) + paths = tuple(available[model_id] for model_id in supported_model_ids) + for model_id, path in zip(supported_model_ids, paths, strict=True): + request = load_request(path) + if request["model_id"] != model_id: + raise ValueError(f"ESMFold2 request filename and model ID differ: {path}") + return paths + + +def _default_output( + exchange_root: Path, + model_id: str, + *, + producer: Literal["reference", "candidate"], + precision: str | None = None, +) -> Path: + path = exchange_root / "structure" / "results" / producer / model_id + return path if precision is None else path / precision + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + + prepare = subparsers.add_parser("prepare") + prepare.add_argument("--exchange-root", type=Path, required=True) + prepare.add_argument("--model-id", action="append", choices=supported_model_ids) + + for command in ("produce-reference", "produce-candidate"): + produce = subparsers.add_parser(command) + produce.add_argument("--exchange-root", type=Path, required=True) + selection = produce.add_mutually_exclusive_group(required=True) + selection.add_argument("--model-id", choices=supported_model_ids) + selection.add_argument( + "--all", + action="store_true", + help="Produce every request in the validated release inventory.", + ) + produce.add_argument("--request", type=Path) + produce.add_argument("--output", type=Path) + if command == "produce-candidate": + produce.add_argument("--precision", choices=("bf16", "fp8"), required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run request preparation or one isolated bundle producer.""" + + args = _parser().parse_args(argv) + if args.command == "prepare": + model_ids = tuple(args.model_id or supported_model_ids) + for path in prepare_requests(args.exchange_root, model_ids=model_ids): + print(path) + return 0 + + if args.all: + if args.request is not None or args.output is not None: + raise ValueError("--request and --output require a single --model-id") + request_paths = _all_prepared_requests(args.exchange_root) + else: + request_paths = (args.request or _default_request(args.exchange_root, args.model_id),) + + for request_path in request_paths: + request = load_request(request_path) + model_id = request["model_id"] + if args.command == "produce-reference": + output_dir = args.output or _default_output( + args.exchange_root, + model_id, + producer="reference", + ) + produce_reference(request_path, output_dir) + else: + output_dir = args.output or _default_output( + args.exchange_root, + model_id, + producer="candidate", + precision=args.precision, + ) + produce_candidate( + request_path, + output_dir, + precision=args.precision, + ) + print(output_dir) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = [ + "fold_sampling_steps", + "fold_seed", + "fold_sequence", + "load_bundle", + "load_request", + "main", + "prepare_requests", + "produce_candidate", + "produce_reference", + "reference_container", + "schema_version", + "supported_model_ids", + "tensor_set_sha256", + "tensor_sha256", + "write_bundle", +] diff --git a/tests/structure/support/esmfold_bundle.py b/tests/structure/support/esmfold_bundle.py new file mode 100644 index 0000000..33c63df --- /dev/null +++ b/tests/structure/support/esmfold_bundle.py @@ -0,0 +1,685 @@ +"""Produce isolated Meta ESMFold v1 bundles for structure compliance. + +The reference path imports only the pinned fair-esm and OpenFold sources through +the manifest adapter. The candidate path imports FastPLMs only after command +dispatch. Both paths use the same raw sequence and normalize the public +``infer`` output into a hash-verified safetensors bundle. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import importlib +import json +import os +import platform +import tempfile +import torch +from collections.abc import Mapping +from contextlib import contextmanager +from dataclasses import asdict +from pathlib import Path +from typing import Any, Literal +from safetensors.torch import load_file, save_file + +from tests.parity.support.state_transforms import transform_parameter_names +from tests.structure.support.state_contract import ( + exact_state_contract, + semantic_config_contract, + validate_exact_state_contract, + validate_semantic_config_contract, +) + + +schema_version = 1 +model_id = "esmfold" +reference_container = "reference-esmfold" +fold_sequence = "MSTNPKPQRKTKRNTNR" +fold_seed = 17 +fold_recycles = 1 +fold_backend = "eager" +fold_deterministic_algorithms = True +supported_precisions = ("fp32", "bf16") +Precision = Literal["fp32", "bf16"] + +_required_outputs = ( + "aatype", + "aligned_confidence_probs", + "atom14_atom_exists", + "atom37_atom_exists", + "chain_index", + "distogram_logits", + "lm_logits", + "mean_plddt", + "plddt", + "positions", + "predicted_aligned_error", + "ptm", + "ptm_logits", + "residue_index", +) +_exact_outputs = ( + "aatype", + "atom14_atom_exists", + "atom37_atom_exists", + "chain_index", + "residue_index", +) +_derived_state_buffers = frozenset( + { + "positional_encoding._float_tensor", + "trunk.structure_module.atom_mask", + "trunk.structure_module.default_frames", + "trunk.structure_module.group_idx", + "trunk.structure_module.lit_positions", + } +) + + +def _canonical_json(value: Mapping[str, Any]) -> str: + return json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + + +def _request_fingerprint(request: Mapping[str, Any]) -> str: + payload = json.dumps( + request, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + # tensor: (...) + # value: (...) + value = tensor.detach().cpu().contiguous() + return value.reshape(-1).view(torch.uint8).numpy().tobytes() + + +def tensor_sha256(tensor: torch.Tensor) -> str: + """Return the exact byte digest for one tensor.""" + + # tensor: (...) + return hashlib.sha256(_tensor_bytes(tensor)).hexdigest() + + +def _checkpoint_metadata(checkpoint: Any) -> dict[str, Any]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + { + "path": item.path, + "algorithm": item.algorithm, + "digest": item.digest, + } + for item in checkpoint.files + ], + } + + +def _upstream_metadata(upstream: Any) -> dict[str, Any]: + return { + "id": upstream.id, + "path": upstream.path, + "url": upstream.url, + "revision": upstream.revision, + "license_expression": upstream.license_expression, + } + + +def _oracle_asset_metadata(asset: Any) -> dict[str, Any]: + return asdict(asset) + + +def prepare_request(exchange_root: Path) -> Path: + """Write the manifest-derived ESMFold v1 reference request.""" + + from fastplms.registry import get_model_registry + + registry = get_model_registry() + spec = registry[model_id] + if spec.family.reference_container != reference_container: + raise RuntimeError("ESMFold reference container disagrees with models.toml.") + expected_upstreams = ("fair-esm", "openfold") + if not set(expected_upstreams).issubset(spec.family.upstreams): + raise RuntimeError("ESMFold is missing its fair-esm or OpenFold provenance.") + request = { + "schema_version": schema_version, + "model_id": model_id, + "architecture": spec.family.architecture, + "adapter": spec.family.reference_adapter, + "reference_container": spec.family.reference_container, + "official": _checkpoint_metadata(spec.official), + "candidate": _checkpoint_metadata(spec.fast), + "candidate_auto_model": spec.auto_map["AutoModel"], + "upstreams": [_upstream_metadata(registry.upstreams[name]) for name in expected_upstreams], + "oracle_assets": [_oracle_asset_metadata(asset) for asset in spec.oracle_assets], + "state_transform": spec.family.state_transform, + "sequence": fold_sequence, + "seed": fold_seed, + "recycles": fold_recycles, + "attention_backend": fold_backend, + "deterministic_algorithms": fold_deterministic_algorithms, + "parameter_dtype": "float32", + "compute_dtypes": list(supported_precisions), + } + request["request_sha256"] = _request_fingerprint(request) + path = exchange_root / "structure" / "requests" / reference_container / f"{model_id}.json" + _atomic_write_text(path, _canonical_json(request)) + return path + + +def _validate_checkpoint(source: object, *, label: str) -> None: + if not isinstance(source, Mapping): + raise ValueError(f"ESMFold request omits {label} checkpoint metadata.") + revision = source.get("revision") + if not isinstance(revision, str) or len(revision) != 40: + raise ValueError(f"ESMFold {label} revision is not immutable.") + files = source.get("files") + if not isinstance(files, list) or not files: + raise ValueError(f"ESMFold {label} checkpoint has no pinned files.") + + +def _validate_request(request: Mapping[str, Any]) -> None: + if request.get("schema_version") != schema_version: + raise ValueError("Unsupported ESMFold structure-bundle schema.") + if request.get("model_id") != model_id: + raise ValueError(f"Unsupported ESMFold model ID: {request.get('model_id')!r}") + if request.get("reference_container") != reference_container: + raise ValueError("ESMFold request names the wrong reference container.") + if request.get("attention_backend") != fold_backend: + raise ValueError("ESMFold official parity requires the eager backend.") + if request.get("deterministic_algorithms") is not fold_deterministic_algorithms: + raise ValueError("ESMFold official parity requires deterministic CUDA algorithms.") + if request.get("parameter_dtype") != "float32": + raise ValueError("ESMFold structure parity requires canonical FP32 parameters.") + if request.get("compute_dtypes") != list(supported_precisions): + raise ValueError("ESMFold structure parity requires FP32 and BF16 compute gates.") + expected = dict(request) + observed_fingerprint = expected.pop("request_sha256", None) + if observed_fingerprint != _request_fingerprint(expected): + raise ValueError("ESMFold request fingerprint mismatch.") + _validate_checkpoint(request.get("official"), label="official") + _validate_checkpoint(request.get("candidate"), label="candidate") + upstreams = request.get("upstreams") + if not isinstance(upstreams, list) or { + item.get("id") for item in upstreams if isinstance(item, Mapping) + } != {"fair-esm", "openfold"}: + raise ValueError("ESMFold request must pin fair-esm and OpenFold.") + for upstream in upstreams: + revision = upstream.get("revision") + if not isinstance(revision, str) or len(revision) != 40: + raise ValueError("ESMFold upstream revision is not immutable.") + assets = request.get("oracle_assets") + if not isinstance(assets, list) or len(assets) != 1: + raise ValueError("ESMFold request must contain its native weights asset.") + asset = assets[0] + if asset.get("role") != "weights" or len(str(asset.get("sha256", ""))) != 64: + raise ValueError("ESMFold native weights asset is not hash-pinned.") + + +def load_request(path: Path) -> dict[str, Any]: + """Load and validate one manifest-derived ESMFold request.""" + + request = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(request, dict): + raise TypeError(f"ESMFold request must be a JSON object: {path}") + _validate_request(request) + return request + + +@contextmanager +def _stable_cuda_numerics(): + """Use deterministic IEEE CUDA numerics for one official or candidate fold.""" + + old_algorithms = torch.are_deterministic_algorithms_enabled() + old_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + + try: + old_fp32_precision = torch.backends.fp32_precision + old_matmul_precision = torch.backends.cuda.matmul.fp32_precision + old_cudnn_precision = torch.backends.cudnn.fp32_precision + except AttributeError: + old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + old_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + old_benchmark = torch.backends.cudnn.benchmark + old_deterministic = torch.backends.cudnn.deterministic + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32 + torch.backends.cudnn.allow_tf32 = old_cudnn_tf32 + torch.backends.cudnn.benchmark = old_benchmark + torch.backends.cudnn.deterministic = old_deterministic + torch.use_deterministic_algorithms(old_algorithms, warn_only=old_warn_only) + return + old_benchmark = torch.backends.cudnn.benchmark + old_deterministic = torch.backends.cudnn.deterministic + torch.backends.fp32_precision = "ieee" + torch.backends.cuda.matmul.fp32_precision = "ieee" + torch.backends.cudnn.fp32_precision = "ieee" + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + try: + yield + finally: + torch.backends.fp32_precision = old_fp32_precision + torch.backends.cuda.matmul.fp32_precision = old_matmul_precision + torch.backends.cudnn.fp32_precision = old_cudnn_precision + torch.backends.cudnn.benchmark = old_benchmark + torch.backends.cudnn.deterministic = old_deterministic + torch.use_deterministic_algorithms(old_algorithms, warn_only=old_warn_only) + + +def _normalize_output(output: object) -> dict[str, torch.Tensor]: + if not isinstance(output, Mapping): + if hasattr(output, "items"): + output = dict(output.items()) + else: + raise TypeError("ESMFold infer did not return a tensor mapping.") + missing = sorted(set(_required_outputs).difference(output)) + if missing: + raise RuntimeError(f"ESMFold infer omitted required outputs: {missing}") + tensors: dict[str, torch.Tensor] = {} + for name in _required_outputs: + value = output[name] + if not torch.is_tensor(value): + raise TypeError(f"ESMFold output {name!r} is not a tensor.") + # tensors[f'output__{name}']: (...) + tensors[f"output__{name}"] = value.detach().cpu().contiguous().clone() + return tensors + + +def _run_infer( + model: torch.nn.Module, + request: Mapping[str, Any], + precision: Precision, +) -> dict[str, torch.Tensor]: + torch.manual_seed(int(request["seed"])) + torch.cuda.manual_seed_all(int(request["seed"])) + with ( + torch.inference_mode(), + _stable_cuda_numerics(), + torch.autocast("cuda", dtype=torch.bfloat16, enabled=precision == "bf16"), + ): + output = model.infer( + request["sequence"], + num_recycles=int(request["recycles"]), + ) + return _normalize_output(output) + + +def _environment_metadata() -> dict[str, Any]: + versions: dict[str, str | None] = {} + for package in ("transformers", "esm", "openfold"): + try: + module = importlib.import_module(package) + except ImportError: + versions[package] = None + else: + versions[package] = str(getattr(module, "__version__", "unknown")) + cuda_properties = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None + return { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "cuda_device": cuda_properties.name if cuda_properties is not None else None, + "cuda_device_capability": ( + list(torch.cuda.get_device_capability(0)) if cuda_properties is not None else None + ), + "cuda_total_memory": ( + int(cuda_properties.total_memory) if cuda_properties is not None else None + ), + "packages": versions, + } + + +def _load_reference_model( + request: Mapping[str, Any], + device: torch.device, +) -> torch.nn.Module: + adapter_name = request["adapter"] + if adapter_name != "tests.parity.support.reference_adapters.esmfold": + raise ValueError(f"Unexpected ESMFold adapter: {adapter_name!r}") + adapter = importlib.import_module(adapter_name) + source = request["official"] + model, tokenizer = adapter.load_official_model( + reference_repo_id=source["repo_id"], + reference_revision=source["revision"], + device=device, + dtype=torch.float32, + oracle_assets=request["oracle_assets"], + ) + if tokenizer is not None: + raise RuntimeError("Meta ESMFold public API must not return a separate tokenizer.") + return model.eval() + + +def _load_candidate_model( + request: Mapping[str, Any], + device: torch.device, +) -> torch.nn.Module: + from fastplms.registry import get_model_registry + + spec = get_model_registry()[model_id] + source = request["candidate"] + if source["repo_id"] != spec.fast.repo_id or source["revision"] != spec.fast.revision: + raise RuntimeError("ESMFold candidate request disagrees with models.toml.") + auto_model = spec.auto_map["AutoModel"] + if request["candidate_auto_model"] != auto_model: + raise RuntimeError("ESMFold candidate AutoModel request disagrees with models.toml.") + module_name, class_name = auto_model.rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + model = model_class.from_pretrained( + source["repo_id"], + revision=source["revision"], + dtype=torch.float32, + attn_implementation=request["attention_backend"], + ) + return model.eval().to(device=device) + + +def _esmfold_semantic_config(model: torch.nn.Module) -> dict[str, Any]: + """Return the shared native/local architecture configuration.""" + + esm = model.esm + esm_config = getattr(esm, "config", None) + hidden_size = getattr(esm, "embed_dim", None) + n_layers = getattr(esm, "num_layers", None) + n_heads = getattr(esm, "attention_heads", None) + if esm_config is not None: + hidden_size = hidden_size or esm_config.hidden_size + n_layers = n_layers or esm_config.num_hidden_layers + n_heads = n_heads or esm_config.num_attention_heads + if hidden_size is None or n_layers is None or n_heads is None: + raise RuntimeError("ESMFold language-model configuration is incomplete.") + token_embedding = getattr(esm, "embed_tokens", None) + if token_embedding is None: + token_embedding = esm.embeddings.word_embeddings + fields = { + "architecture": "ESMFold", + "distogram_bins": int(model.distogram_head.out_features), + "esm_attention_heads": int(n_heads), + "esm_hidden_size": int(hidden_size), + "esm_layers": int(n_layers), + "esm_state_count": int(model.esm_s_combine.numel()), + "folding_blocks": len(model.trunk.blocks), + "esm_vocab_size": int(token_embedding.num_embeddings), + "pairwise_state_dim": int(model.distogram_head.in_features), + "sequence_output_tokens": int(model.lm_head.out_features), + "sequence_state_dim": int(model.lm_head.in_features), + } + return semantic_config_contract(fields) + + +def _metadata( + request: Mapping[str, Any], + *, + producer: Literal["reference", "candidate"], + model: torch.nn.Module, + precision: Precision, +) -> dict[str, Any]: + transform_name = str(request["state_transform"]) + if producer == "reference": + + def name_transform(name: str) -> tuple[str, ...]: + return transform_parameter_names(transform_name, name) + + else: + + def name_transform(name: str) -> tuple[str, ...]: + if name in _derived_state_buffers or name.startswith( + ("mlm_head.", "esm.contact_head.") + ): + return () + return (name,) + + return { + "schema_version": schema_version, + "producer": producer, + "model_id": model_id, + "request_sha256": request["request_sha256"], + "official": request["official"], + "candidate": request["candidate"], + "upstreams": request["upstreams"], + "oracle_assets": request["oracle_assets"], + "sequence": request["sequence"], + "seed": request["seed"], + "recycles": request["recycles"], + "attention_backend": request["attention_backend"], + "deterministic_algorithms": request["deterministic_algorithms"], + "parameter_dtype": request["parameter_dtype"], + "compute_dtype": precision, + "execution": ( + "fp32_parameters_cuda_bf16_autocast" if precision == "bf16" else "fp32_parameters" + ), + "esm_parameter_dtypes": sorted( + { + str(parameter.dtype).removeprefix("torch.") + for parameter in model.esm.parameters() + if parameter.is_floating_point() + } + ), + "state_transform": transform_name, + "semantic_config": _esmfold_semantic_config(model), + "state": exact_state_contract( + model, + name_transform=name_transform, + ), + "environment": _environment_metadata(), + } + + +def _atomic_write_text(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + handle, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + text=True, + ) + try: + with os.fdopen(handle, "w", encoding="utf-8", newline="\n") as stream: + stream.write(content) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, path) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def write_bundle( + output_dir: Path, + tensors: Mapping[str, torch.Tensor], + metadata: Mapping[str, Any], +) -> None: + """Atomically publish one normalized ESMFold structure bundle.""" + + output_dir.mkdir(parents=True, exist_ok=True) + normalized = { + name: tensor.detach().cpu().contiguous().clone() for name, tensor in sorted(tensors.items()) + } + complete_metadata = dict(metadata) + complete_metadata.update( + { + "tensor_hashes": {name: tensor_sha256(tensor) for name, tensor in normalized.items()}, + "tensor_keys": sorted(normalized), + } + ) + bundle_path = output_dir / "bundle.safetensors" + handle, temporary_name = tempfile.mkstemp( + dir=output_dir, + prefix=".bundle.", + suffix=".safetensors.tmp", + ) + os.close(handle) + try: + save_file(normalized, temporary_name) + os.replace(temporary_name, bundle_path) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + _atomic_write_text(output_dir / "metadata.json", _canonical_json(complete_metadata)) + + +def load_bundle(path: Path) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Load one ESMFold bundle and verify all declared tensor hashes.""" + + tensor_path = path / "bundle.safetensors" + metadata_path = path / "metadata.json" + if not tensor_path.is_file() or not metadata_path.is_file(): + raise FileNotFoundError( + f"Missing ESMFold structure bundle under {path}. Run native producers first." + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if metadata.get("schema_version") != schema_version: + raise ValueError(f"Unsupported ESMFold bundle schema under {path}.") + tensors = load_file(tensor_path, device="cpu") + if sorted(tensors) != metadata.get("tensor_keys"): + raise ValueError(f"Tensor-key mismatch in ESMFold bundle {path}.") + observed_hashes = {name: tensor_sha256(tensor) for name, tensor in tensors.items()} + if observed_hashes != metadata.get("tensor_hashes"): + raise ValueError(f"Tensor hash mismatch in ESMFold bundle {path}.") + validate_exact_state_contract(metadata.get("state")) + validate_semantic_config_contract(metadata.get("semantic_config")) + return tensors, metadata + + +def _require_fp32_parameters(model: torch.nn.Module) -> None: + unexpected = sorted( + { + str(parameter.dtype) + for parameter in model.parameters() + if parameter.is_floating_point() and parameter.dtype != torch.float32 + } + ) + if unexpected: + raise RuntimeError(f"ESMFold requires FP32 checkpoint parameters, found {unexpected}.") + + +def produce_reference( + request_path: Path, + output_dir: Path, + *, + precision: Precision, +) -> None: + """Run Meta's pinned public ESMFold v1 constructor and ``infer`` API.""" + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Official ESMFold structure bundles require CUDA.") + model = _load_reference_model(request, torch.device("cuda")) + try: + _require_fp32_parameters(model) + metadata = _metadata(request, producer="reference", model=model, precision=precision) + tensors = _run_infer(model, request, precision) + write_bundle( + output_dir, + tensors, + metadata, + ) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def produce_candidate( + request_path: Path, + output_dir: Path, + *, + precision: Precision, +) -> None: + """Run the pinned FastPLMs ESMFold package class through public ``infer``.""" + + request = load_request(request_path) + if not torch.cuda.is_available(): + raise RuntimeError("Candidate ESMFold structure bundles require CUDA.") + model = _load_candidate_model(request, torch.device("cuda")) + try: + _require_fp32_parameters(model) + metadata = _metadata(request, producer="candidate", model=model, precision=precision) + tensors = _run_infer(model, request, precision) + write_bundle( + output_dir, + tensors, + metadata, + ) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +def _default_request(exchange_root: Path) -> Path: + return exchange_root / "structure" / "requests" / reference_container / f"{model_id}.json" + + +def _default_output( + exchange_root: Path, + producer: Literal["reference", "candidate"], + precision: Precision, +) -> Path: + return exchange_root / "structure" / "results" / producer / model_id / precision + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + prepare = subparsers.add_parser("prepare") + prepare.add_argument("--exchange-root", type=Path, required=True) + for name in ("produce-reference", "produce-candidate"): + producer = subparsers.add_parser(name) + producer.add_argument("--exchange-root", type=Path, required=True) + producer.add_argument("--precision", choices=supported_precisions, required=True) + return parser + + +def main() -> None: + arguments = _parser().parse_args() + if arguments.command == "prepare": + print(prepare_request(arguments.exchange_root)) + return + request_path = _default_request(arguments.exchange_root) + precision = arguments.precision + if arguments.command == "produce-reference": + output = _default_output(arguments.exchange_root, "reference", precision) + produce_reference(request_path, output, precision=precision) + else: + output = _default_output(arguments.exchange_root, "candidate", precision) + produce_candidate(request_path, output, precision=precision) + print(output) + + +if __name__ == "__main__": + main() + + +__all__ = [ + "_exact_outputs", + "fold_backend", + "fold_deterministic_algorithms", + "fold_recycles", + "fold_seed", + "fold_sequence", + "load_bundle", + "load_request", + "model_id", + "prepare_request", + "produce_candidate", + "produce_reference", + "reference_container", + "schema_version", + "supported_precisions", + "tensor_sha256", +] diff --git a/tests/structure/support/hardware.py b/tests/structure/support/hardware.py new file mode 100644 index 0000000..2ee5350 --- /dev/null +++ b/tests/structure/support/hardware.py @@ -0,0 +1,110 @@ +"""Hardware identity contracts for Hopper/SM90 release validation.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + + +HOPPER_SM90_CAPABILITY = (9, 0) +HOPPER_PRODUCT_NAMES = ("H100", "H200", "GH200") +_HOPPER_PRODUCT_PATTERN = re.compile(r"(? HopperDeviceFingerprint: + """Validate and return one allowed NVIDIA Hopper/SM90 device fingerprint.""" + + name = environment.get("cuda_device") + if not isinstance(name, str) or not name.strip(): + raise AssertionError("Hopper validation requires a non-empty CUDA device name.") + match = _HOPPER_PRODUCT_PATTERN.search(name.upper()) + if match is None: + allowed = ", ".join(HOPPER_PRODUCT_NAMES) + raise AssertionError( + f"Release validation requires an NVIDIA Hopper product ({allowed}); got {name!r}." + ) + + raw_capability = environment.get("cuda_device_capability") + if ( + not isinstance(raw_capability, Sequence) + or isinstance(raw_capability, (str, bytes)) + or len(raw_capability) != 2 + or any(not isinstance(value, int) or isinstance(value, bool) for value in raw_capability) + ): + raise AssertionError( + "Hopper validation requires cuda_device_capability as two integer components." + ) + capability = (raw_capability[0], raw_capability[1]) + if capability != HOPPER_SM90_CAPABILITY: + raise AssertionError( + f"Release validation requires compute capability 9.0; got {capability}." + ) + + total_memory = environment.get("cuda_total_memory") + if not isinstance(total_memory, int) or isinstance(total_memory, bool) or total_memory <= 0: + raise AssertionError("Hopper validation requires positive cuda_total_memory bytes.") + + return HopperDeviceFingerprint( + product=match.group(1), + name=name.strip(), + capability=capability, + total_memory=total_memory, + ) + + +def assert_same_hopper_sm90_device( + current: Mapping[str, object], + baseline: Mapping[str, object], +) -> None: + """Reject cross-device comparisons even when both devices are Hopper/SM90.""" + + current_fingerprint = hopper_sm90_fingerprint(current) + baseline_fingerprint = hopper_sm90_fingerprint(baseline) + if current_fingerprint != baseline_fingerprint: + raise AssertionError( + "Cross-device comparison is forbidden: " + f"current={current_fingerprint!r}, baseline={baseline_fingerprint!r}." + ) + + +def assert_recorded_hopper_device_matches( + current: Mapping[str, object], + recorded: Mapping[str, object], +) -> None: + """Match a live Hopper device to a golden's recorded hardware fields. + + Legacy goldens predate capability and memory fields, so their exact device + name remains the strongest available identity. New captures must retain the + additional fields and therefore receive the full comparison. + """ + + current_fingerprint = hopper_sm90_fingerprint(current) + recorded_name = recorded.get("cuda_device") + if recorded_name != current_fingerprint.name: + raise AssertionError( + "Cross-device golden comparison is forbidden: " + f"current device={current_fingerprint.name!r}, recorded device={recorded_name!r}." + ) + optional_fields = ("cuda_device_capability", "cuda_total_memory") + mismatches = [ + field + for field in optional_fields + if field in recorded and recorded[field] != current.get(field) + ] + if mismatches: + details = ", ".join( + f"{field}: current={current.get(field)!r}, recorded={recorded.get(field)!r}" + for field in mismatches + ) + raise AssertionError(f"Cross-device golden comparison is forbidden: {details}.") diff --git a/tests/structure/support/state_contract.py b/tests/structure/support/state_contract.py new file mode 100644 index 0000000..1e98514 --- /dev/null +++ b/tests/structure/support/state_contract.py @@ -0,0 +1,210 @@ +"""Compact exact checkpoint contracts for isolated structure-model oracles.""" + +from __future__ import annotations + +import hashlib +import json +import torch +from collections.abc import Callable, Mapping +from typing import Any + + +NameTransform = Callable[[str], tuple[str, ...]] + +_PACKAGING_CONFIG_FIELDS = frozenset( + { + "_commit_hash", + "_name_or_path", + "architectures", + "auto_map", + "fastplms_checkpoint_hash", + "fastplms_checkpoint_repo_id", + "fastplms_checkpoint_revision", + "fastplms_model_id", + "fastplms_runtime_bundle_sha256", + "fastplms_runtime_revision", + "fastplms_source_tree_sha256", + "fastplms_weights_revision", + "dtype", + "name_or_path", + "torch_dtype", + "transformers_version", + } +) + + +def _canonical_json(value: object) -> bytes: + return json.dumps( + value, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def tensor_sha256(X: torch.Tensor) -> str: + """Hash one tensor exactly, including scalar tensors.""" + + # X: (...) + # value: (-1,) + value = X.detach().to(device="cpu").contiguous().reshape(-1) + return hashlib.sha256(value.view(torch.uint8).numpy().tobytes()).hexdigest() + + +def _included(name: str, excluded_prefixes: tuple[str, ...]) -> bool: + return not any(name.startswith(prefix) for prefix in excluded_prefixes) + + +def exact_state_contract( + model: torch.nn.Module, + *, + name_transform: NameTransform | None = None, + excluded_prefixes: tuple[str, ...] = (), +) -> dict[str, Any]: + """Return exact state tensor and parameter-alias metadata without tensor payloads.""" + + transform = name_transform or (lambda name: (name,)) + tensors: dict[str, dict[str, object]] = {} + for source_name, X in sorted(model.state_dict().items()): + if not _included(source_name, excluded_prefixes): + continue + targets = transform(source_name) + for name in targets: + if name in tensors: + raise RuntimeError(f"State-contract key collision for {name!r}.") + tensors[name] = { + "dtype": str(X.dtype).removeprefix("torch."), + "shape": list(X.shape), + "sha256": tensor_sha256(X), + } + if not tensors: + raise RuntimeError("A structure checkpoint state contract cannot be empty.") + + by_parameter: dict[int, set[str]] = {} + for source_name, parameter in model.named_parameters(remove_duplicate=False): + if not _included(source_name, excluded_prefixes): + continue + by_parameter.setdefault(id(parameter), set()).update(transform(source_name)) + aliases = sorted(sorted(names) for names in by_parameter.values() if len(names) > 1) + payload = {"aliases": aliases, "tensors": dict(sorted(tensors.items()))} + return { + **payload, + "sha256": hashlib.sha256(_canonical_json(payload)).hexdigest(), + } + + +def semantic_config_contract(config: object) -> dict[str, Any]: + """Normalize a Transformers configuration after removing packaging fields.""" + + if hasattr(config, "to_dict"): + raw = config.to_dict() + elif isinstance(config, Mapping): + raw = dict(config) + else: + raise TypeError(f"Unsupported semantic configuration: {type(config)!r}") + + def normalize(value: object) -> object: + if isinstance(value, Mapping): + return { + str(key): normalize(item) + for key, item in sorted(value.items()) + if str(key) not in _PACKAGING_CONFIG_FIELDS + } + if isinstance(value, (list, tuple)): + return [normalize(item) for item in value] + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + normalized = normalize(raw) + assert isinstance(normalized, dict) + return { + "fields": normalized, + "sha256": hashlib.sha256(_canonical_json(normalized)).hexdigest(), + } + + +def validate_exact_state_contract(contract: object) -> None: + """Reject malformed or modified compact state metadata.""" + + if not isinstance(contract, Mapping): + raise ValueError("Structure state contract must be a mapping.") + tensors = contract.get("tensors") + aliases = contract.get("aliases") + if not isinstance(tensors, Mapping) or not tensors or not isinstance(aliases, list): + raise ValueError("Structure state contract is incomplete.") + tensor_names: set[str] = set() + for name, metadata in tensors.items(): + if not isinstance(name, str) or not name or not isinstance(metadata, Mapping): + raise ValueError("Structure state tensor metadata is malformed.") + if set(metadata) != {"dtype", "shape", "sha256"}: + raise ValueError(f"Structure state tensor {name!r} has an invalid schema.") + dtype = metadata["dtype"] + shape = metadata["shape"] + digest = metadata["sha256"] + if not isinstance(dtype, str) or not dtype: + raise ValueError(f"Structure state tensor {name!r} has an invalid dtype.") + if not isinstance(shape, list) or any( + not isinstance(dimension, int) or isinstance(dimension, bool) or dimension < 0 + for dimension in shape + ): + raise ValueError(f"Structure state tensor {name!r} has an invalid shape.") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError(f"Structure state tensor {name!r} has an invalid digest.") + tensor_names.add(name) + + normalized_aliases: list[list[str]] = [] + for group in aliases: + if not isinstance(group, list) or len(group) < 2: + raise ValueError("Structure state aliases are malformed.") + if any(not isinstance(name, str) or name not in tensor_names for name in group): + raise ValueError("Structure state aliases name an unknown tensor.") + normalized = sorted(set(group)) + if len(normalized) != len(group): + raise ValueError("Structure state aliases contain duplicate names.") + normalized_aliases.append(normalized) + if aliases != sorted(normalized_aliases): + raise ValueError("Structure state aliases are not canonical.") + + payload = {"aliases": aliases, "tensors": dict(sorted(tensors.items()))} + expected = hashlib.sha256(_canonical_json(payload)).hexdigest() + if contract.get("sha256") != expected: + raise ValueError("Structure state contract digest mismatch.") + + +def validate_semantic_config_contract(contract: object) -> None: + """Reject malformed or modified compact semantic configuration metadata.""" + + if not isinstance(contract, Mapping) or not isinstance(contract.get("fields"), Mapping): + raise ValueError("Structure semantic configuration contract is incomplete.") + fields = contract["fields"] + + def reject_packaging_fields(value: object) -> None: + if isinstance(value, Mapping): + if any(str(key) in _PACKAGING_CONFIG_FIELDS for key in value): + raise ValueError("Structure semantic configuration contains packaging fields.") + for item in value.values(): + reject_packaging_fields(item) + elif isinstance(value, list): + for item in value: + reject_packaging_fields(item) + + reject_packaging_fields(fields) + expected = hashlib.sha256(_canonical_json(fields)).hexdigest() + if contract.get("sha256") != expected: + raise ValueError("Structure semantic configuration digest mismatch.") + + +__all__ = [ + "exact_state_contract", + "semantic_config_contract", + "tensor_sha256", + "validate_exact_state_contract", + "validate_semantic_config_contract", +] diff --git a/tests/structure/test_boltz2_featurization.py b/tests/structure/test_boltz2_featurization.py new file mode 100644 index 0000000..9f171a1 --- /dev/null +++ b/tests/structure/test_boltz2_featurization.py @@ -0,0 +1,134 @@ +"""Deterministic candidate-side contracts for Boltz2 feature preparation. + +Live parity against the pinned Boltz package runs in the native reference +service. These tests keep inexpensive shape, dtype, padding, and repeatability +failures close to the FastPLMs implementation. +""" + +from __future__ import annotations + +import pytest +import torch + +from fastplms.models.boltz.minimal_featurizer import build_boltz2_features +from fastplms.models.boltz.minimal_structures import ProteinStructureTemplate + + +pytestmark = pytest.mark.structure + +SEQUENCE = "ACDEFGHIKLMNPQRSTVWY" +REQUIRED_FEATURES = { + "atom_pad_mask", + "atom_to_token", + "coords", + "disto_center", + "frames_idx", + "msa", + "msa_mask", + "ref_pos", + "res_type", + "residue_index", + "token_index", + "token_pad_mask", + "token_to_center_atom", + "token_to_rep_atom", +} + + +def _build_with_seed( + seed: int, +) -> tuple[dict[str, torch.Tensor], ProteinStructureTemplate]: + torch.manual_seed(seed) + return build_boltz2_features(SEQUENCE) + + +def test_boltz2_feature_preparation_is_seed_reproducible() -> None: + first, first_template = _build_with_seed(17) + second, second_template = _build_with_seed(17) + + assert first.keys() == second.keys() + for name in first: + assert torch.equal(first[name], second[name]), name + assert first_template == second_template + + +def test_boltz2_single_chain_feature_contract() -> None: + features, template = _build_with_seed(23) + n_residues = len(SEQUENCE) + + assert features.keys() >= REQUIRED_FEATURES + assert template.sequence == SEQUENCE + assert template.num_residues == n_residues + assert features["token_index"].shape == (1, n_residues) + assert features["residue_index"].shape == (1, n_residues) + assert features["token_pad_mask"].shape == (1, n_residues) + assert features["res_type"].shape[:2] == (1, n_residues) + assert features["msa"].shape[:3] == (1, 1, n_residues) + + n_atoms_padded = features["atom_pad_mask"].shape[-1] + assert n_atoms_padded >= template.num_atoms + assert n_atoms_padded % 32 == 0 + assert features["ref_pos"].shape == (1, n_atoms_padded, 3) + assert features["coords"].shape == (1, 1, n_atoms_padded, 3) + assert features["atom_to_token"].shape == (1, n_atoms_padded, n_residues) + assert torch.equal( + features["atom_pad_mask"][0, : template.num_atoms], + torch.ones(template.num_atoms), + ) + assert not features["atom_pad_mask"][0, template.num_atoms :].bool().any() + + +def test_boltz2_feature_dtypes_are_explicit() -> None: + features, _ = _build_with_seed(29) + + integer_features = ( + "atom_backbone_feat", + "atom_to_token", + "contact_conditioning", + "frames_idx", + "msa", + "ref_atom_name_chars", + "ref_chirality", + "ref_element", + "res_type", + "residue_index", + "r_set_to_rep_atom", + "token_index", + "token_to_center_atom", + "token_to_rep_atom", + ) + for name in integer_features: + assert features[name].dtype == torch.long + for name in ("atom_resolved_mask", "frame_resolved_mask", "has_deletion"): + assert features[name].dtype == torch.bool + for name in ("ref_pos", "coords", "token_pad_mask"): + assert features[name].dtype == torch.float32 + + +def test_boltz2_sequence_only_observation_features_are_empty() -> None: + features, _ = _build_with_seed(31) + + assert torch.count_nonzero(features["coords"]) == 0 + assert torch.count_nonzero(features["disto_center"]) == 0 + assert torch.count_nonzero(features["disto_coords_ensemble"]) == 0 + assert torch.all(features["disto_target"][..., 0] == 1) + assert torch.count_nonzero(features["disto_target"][..., 1:]) == 0 + + +def test_boltz2_canonical_charge_and_chirality_tables() -> None: + features, template = _build_with_seed(37) + + def atom_index(residue_name: str, atom_name: str) -> int: + for index, (name, residue_index) in enumerate( + zip(template.atom_names, template.atom_residue_index, strict=True) + ): + if name == atom_name and template.residue_names[residue_index] == residue_name: + return index + raise AssertionError(f"Missing {residue_name} {atom_name} atom.") + + for residue_name, atom_name in (("HIS", "ND1"), ("LYS", "NZ")): + index = atom_index(residue_name, atom_name) + assert features["ref_charge"][0, index].item() == 1.0 + for residue_name, atom_name in (("ALA", "CA"), ("ILE", "CB"), ("THR", "CB")): + index = atom_index(residue_name, atom_name) + assert features["ref_chirality"][0, index].item() == 2 diff --git a/tests/structure/test_boltz2_folding_compliance.py b/tests/structure/test_boltz2_folding_compliance.py new file mode 100644 index 0000000..2562e8a --- /dev/null +++ b/tests/structure/test_boltz2_folding_compliance.py @@ -0,0 +1,552 @@ +"""Release gates over isolated official and candidate Boltz2 bundles.""" + +from __future__ import annotations + +import ast +import copy +import inspect +import os +import pytest +import torch +import torch.nn.functional as F +from collections.abc import Mapping +from pathlib import Path + +from fastplms.models.boltz.modeling_boltz2 import Boltz2Config +from fastplms.registry import get_model_registry +from tests.structure.support import boltz2_bundle +from tests.structure.support.boltz2_bundle import load_bundle, load_request +from tests.structure.support.hardware import ( + assert_same_hopper_sm90_device, + hopper_sm90_fingerprint, +) +from tests.structure.support.state_contract import semantic_config_contract + + +bf16_targets = { + "ca_rmsd": 0.10, + "lddt_ca": 0.995, + "plddt_mae": 0.001, + "pae_mae": 0.10, + "ptm_error": 0.002, + "iptm_error": 0.002, + "mean_probability_jsd": 0.001, +} +bf16_hard_limits = { + "ca_rmsd": 0.25, + "lddt_ca": 0.99, + "plddt_mae": 0.005, + "pae_mae": 0.50, + "ptm_error": 0.005, + "iptm_error": 0.005, + "mean_probability_jsd": 0.005, +} +bf16_relative_l2_target = 1e-2 +bf16_relative_l2_hard_limit = 3e-2 + + +def _exchange_root() -> Path: + return Path(os.environ.get("FASTPLMS_REFERENCE_EXCHANGE", "artifacts/reference")) + + +def _paths() -> tuple[Path, Path, Path]: + root = _exchange_root() + request = ( + root + / "structure" + / "requests" + / boltz2_bundle.reference_container + / f"{boltz2_bundle.model_id}.json" + ) + results = root / "structure" / "results" + reference = results / "reference" / boltz2_bundle.model_id / boltz2_bundle.fold_dtype + candidate = results / "candidate" / boltz2_bundle.model_id / boltz2_bundle.fold_dtype + return request, reference, candidate + + +def _checkpoint_contract(checkpoint: object) -> dict[str, object]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + {"path": item.path, "algorithm": item.algorithm, "digest": item.digest} + for item in checkpoint.files + ], + } + + +def _upstream_contract(upstream: object) -> dict[str, object]: + return { + "id": upstream.id, + "path": upstream.path, + "url": upstream.url, + "revision": upstream.revision, + "license_expression": upstream.license_expression, + } + + +def _features(tensors: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return { + name.removeprefix("feature__"): tensor + for name, tensor in tensors.items() + if name.startswith("feature__") + } + + +def _output(tensors: Mapping[str, torch.Tensor], name: str) -> torch.Tensor: + key = f"output__{name}" + if key not in tensors: + raise KeyError(f"Boltz2 bundle omits required output {name!r}.") + return tensors[key] + + +def _assert_bundle_identity( + metadata: Mapping[str, object], + request: Mapping[str, object], + *, + producer: str, +) -> None: + registry = get_model_registry() + spec = registry[boltz2_bundle.model_id] + assert metadata["producer"] == producer + assert metadata["model_id"] == spec.id + assert metadata["request_sha256"] == request["request_sha256"] + assert metadata["official"] == _checkpoint_contract(spec.official) + assert metadata["candidate"] == _checkpoint_contract(spec.fast) + assert metadata["upstream"] == _upstream_contract(registry.upstreams["boltz"]) + for name in ( + "sequence", + "feature_seed", + "seed", + "recycling_steps", + "sampling_steps", + "diffusion_samples", + "diffusion_noise_generator", + "conformer_policy", + "steering", + "dtype", + "parameter_dtype", + "compute_dtype", + "execution", + ): + assert metadata[name] == request[name] + assert metadata["attention_backend"] == "eager" + environment = metadata["environment"] + assert isinstance(environment, Mapping) + hopper_sm90_fingerprint(environment) + if producer == "candidate": + assert str(environment["torch"]).split("+", maxsplit=1)[0] == "2.13.0" + packages = environment["packages"] + assert isinstance(packages, Mapping) + assert packages["transformers"] == "5.13.0" + assert str(environment["cuda_runtime"]).startswith("13.0") + + +def _assert_exact_features( + actual_tensors: Mapping[str, torch.Tensor], + actual_metadata: Mapping[str, object], + expected_tensors: Mapping[str, torch.Tensor], + expected_metadata: Mapping[str, object], +) -> None: + actual = _features(actual_tensors) + expected = _features(expected_tensors) + assert actual.keys() == expected.keys() == set(boltz2_bundle._feature_names) + for name in boltz2_bundle._exact_features: + # X: (...) + X = actual[name] + X_ref = expected[name] + assert X.dtype == X_ref.dtype, f"{name}: dtype" + assert X.shape == X_ref.shape, f"{name}: shape" + assert torch.equal(X, X_ref), f"{name}: values" + for name in set(actual).difference(boltz2_bundle._exact_features): + # X: (...) + X = actual[name] + X_ref = expected[name] + assert X.dtype == X_ref.dtype, f"{name}: dtype" + assert X.shape == X_ref.shape, f"{name}: shape" + if X.is_floating_point(): + assert torch.isfinite(X).all(), f"{name}: candidate finite values" + assert torch.isfinite(X_ref).all(), f"{name}: reference finite values" + assert torch.equal(actual["ref_pos"], expected["ref_pos"]), "ref_pos: values" + actual_hash = actual_metadata["feature_sha256"] + expected_hash = expected_metadata["feature_sha256"] + assert isinstance(actual_hash, str) + assert isinstance(expected_hash, str) + assert actual_hash == expected_hash, "feature_sha256" + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + difference = torch.linalg.vector_norm(actual.float() - expected.float()) + # scale: (...) + scale = torch.linalg.vector_norm(expected.float()).clamp_min(torch.finfo(torch.float32).tiny) + return (difference / scale).item() + + +def _first_coordinates(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + # X: (...) + X = _output(tensors, "sample_atom_coords").float() + return X.reshape(-1, X.shape[-2], 3)[0] + + +def _ca_mask(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + features = _features(tensors) + # encoded: (4,) + encoded = torch.tensor([ord("C") - 32, ord("A") - 32, 0, 0]) + # atom_names: (...) + atom_names = features["ref_atom_name_chars"][0].argmax(dim=-1) + # atom_mask: (...) + atom_mask = features["atom_pad_mask"][0].bool() + # ca_mask: (...) + ca_mask = atom_names.eq(encoded).all(dim=-1) & atom_mask + assert ca_mask.sum().item() == len(boltz2_bundle.fold_sequence) + return ca_mask + + +def _ca_coordinates(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + return _first_coordinates(tensors)[_ca_mask(tensors)] + + +def _aligned_rmsd(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + # X: (...) + X = actual.float() - actual.float().mean(dim=0, keepdim=True) + X_ref = expected.float() - expected.float().mean(dim=0, keepdim=True) + covariance = X.T @ X_ref + U, _, Vh = torch.linalg.svd(covariance) + # correction: (3, 3) + correction = torch.eye(3) + correction[-1, -1] = torch.sign(torch.det(U @ Vh)) + rotation = U @ correction @ Vh + aligned = X @ rotation + return torch.sqrt(torch.mean(torch.sum((aligned - X_ref) ** 2, dim=-1))).item() + + +def _lddt_ca(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + actual_distances = torch.cdist(actual.float(), actual.float()) + expected_distances = torch.cdist(expected.float(), expected.float()) + # pair_mask: (...) + pair_mask = expected_distances.lt(15.0) + pair_mask.fill_diagonal_(False) + assert pair_mask.any() + errors = (actual_distances - expected_distances).abs() + # scores: (...) + scores = torch.stack([errors.lt(threshold).float() for threshold in (0.5, 1.0, 2.0, 4.0)]).mean( + dim=0 + ) + return scores[pair_mask].mean().item() + + +def _probability_jsd( + actual_logits: torch.Tensor, + expected_logits: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + # actual_logits: (..., c), expected_logits: (..., c), mask: (...) + actual_log_prob = F.log_softmax(actual_logits.float(), dim=-1) + expected_log_prob = F.log_softmax(expected_logits.float(), dim=-1) + actual_prob = actual_log_prob.exp() + expected_prob = expected_log_prob.exp() + mean_prob = 0.5 * (actual_prob + expected_prob) + log_mean_prob = mean_prob.clamp_min(torch.finfo(torch.float32).tiny).log() + divergence = 0.5 * ( + (actual_prob * (actual_log_prob - log_mean_prob)).sum(dim=-1) + + (expected_prob * (expected_log_prob - log_mean_prob)).sum(dim=-1) + ) + while mask.ndim < divergence.ndim: + # mask: (...) + mask = mask.unsqueeze(0) + mask = torch.broadcast_to(mask, divergence.shape) + return divergence[mask].mean() + + +def _metrics( + actual: Mapping[str, torch.Tensor], + expected: Mapping[str, torch.Tensor], +) -> tuple[dict[str, float], dict[str, float]]: + features = _features(actual) + # token_mask: (...) + token_mask = features["token_pad_mask"].bool() + # pair_mask: (...) + pair_mask = token_mask[:, :, None] & token_mask[:, None, :] + plddt_actual = _output(actual, "plddt").float().reshape_as(token_mask) + plddt_expected = _output(expected, "plddt").float().reshape_as(token_mask) + pae_actual = _output(actual, "pae").float().reshape_as(pair_mask) + pae_expected = _output(expected, "pae").float().reshape_as(pair_mask) + probability_values = [ + _probability_jsd( + _output(actual, "pdistogram").squeeze(-2), + _output(expected, "pdistogram").squeeze(-2), + pair_mask, + ), + _probability_jsd( + _output(actual, "pde_logits"), + _output(expected, "pde_logits"), + pair_mask, + ), + _probability_jsd( + _output(actual, "pae_logits"), + _output(expected, "pae_logits"), + pair_mask, + ), + _probability_jsd( + _output(actual, "plddt_logits"), + _output(expected, "plddt_logits"), + token_mask, + ), + ] + structure_metrics = { + "ca_rmsd": _aligned_rmsd(_ca_coordinates(actual), _ca_coordinates(expected)), + "lddt_ca": _lddt_ca(_ca_coordinates(actual), _ca_coordinates(expected)), + "plddt_mae": (plddt_actual[token_mask] - plddt_expected[token_mask]).abs().mean().item(), + "pae_mae": (pae_actual[pair_mask] - pae_expected[pair_mask]).abs().mean().item(), + "ptm_error": ( + _output(actual, "ptm").float().reshape(-1)[0] + - _output(expected, "ptm").float().reshape(-1)[0] + ) + .abs() + .item(), + "iptm_error": ( + _output(actual, "iptm").float().reshape(-1)[0] + - _output(expected, "iptm").float().reshape(-1)[0] + ) + .abs() + .item(), + "mean_probability_jsd": torch.stack(probability_values).mean().item(), + } + relative_l2 = { + name: _relative_l2(_output(actual, name), _output(expected, name)) + for name in ("pdistogram", "pde_logits", "pae_logits", "plddt_logits") + } + return structure_metrics, relative_l2 + + +def _assert_valid_outputs(tensors: Mapping[str, torch.Tensor], *, context: str) -> None: + features = _features(tensors) + # atom_mask: (...) + atom_mask = features["atom_pad_mask"][0].bool() + coordinates = _first_coordinates(tensors) + assert torch.isfinite(coordinates[atom_mask]).all(), f"{context}: coordinates" + for name in boltz2_bundle._required_outputs: + X = _output(tensors, name) + if X.is_floating_point(): + assert torch.isfinite(X).all(), f"{context}: {name}" + + +def test_boltz2_reference_path_has_no_fastplms_dependency() -> None: + tree = ast.parse(inspect.getsource(boltz2_bundle)) + for node in tree.body: + if isinstance(node, ast.Import): + assert all(not alias.name.startswith("fastplms") for alias in node.names) + elif isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith("fastplms") + for function in ( + boltz2_bundle._prepare_reference_features, + boltz2_bundle._load_reference_model, + boltz2_bundle._run_model, + boltz2_bundle.produce_reference, + ): + assert "fastplms" not in inspect.getsource(function).lower() + + +def test_boltz2_request_is_manifest_exact(tmp_path: Path) -> None: + path = boltz2_bundle.prepare_request(tmp_path) + request = load_request(path) + registry = get_model_registry() + spec = registry[boltz2_bundle.model_id] + assert request["official"] == _checkpoint_contract(spec.official) + assert request["candidate"] == _checkpoint_contract(spec.fast) + assert request["upstream"] == _upstream_contract(registry.upstreams["boltz"]) + assert request["sequence"] == boltz2_bundle.fold_sequence + assert spec.family.bf16_execution == "fp32_parameters_autocast" + assert request["parameter_dtype"] == boltz2_bundle.fold_parameter_dtype + assert request["compute_dtype"] == boltz2_bundle.fold_compute_dtype + assert request["execution"] == boltz2_bundle.fold_execution + assert request["feature_names"] == list(boltz2_bundle._feature_names) + assert request["output_names"] == list(boltz2_bundle._required_outputs) + + +def test_boltz2_parameter_storage_contract_is_strict() -> None: + model = torch.nn.Linear(3, 2, bias=False) + boltz2_bundle._require_floating_parameter_dtype(model, torch.float32) + + with pytest.raises(RuntimeError, match=r"requires torch\.float32 parameter storage"): + boltz2_bundle._require_floating_parameter_dtype(model.to(torch.bfloat16), torch.float32) + + +def test_boltz2_portable_noise_never_reads_pytorch_rng( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def forbidden_random_draw(*args: object, **kwargs: object) -> torch.Tensor: + raise AssertionError("Portable Boltz2 noise called PyTorch's RNG.") + + monkeypatch.setattr(torch, "randn", forbidden_random_draw) + monkeypatch.setattr(torch, "randn_like", forbidden_random_draw) + streams: list[list[torch.Tensor]] = [] + for global_seed in (1, 987_654_321): + torch.manual_seed(global_seed) + with boltz2_bundle._portable_random_draws(17) as captured: + # N0 and N1 cover both stochastic interfaces used by Boltz2. + # N0: (2, 3) + N0 = torch.randn((2, 3), dtype=torch.float32) + N1 = torch.randn_like(torch.empty(4, dtype=torch.bfloat16)) + # out: (2,) + out = torch.empty(2, dtype=torch.float64) + # N2: (2,) + N2 = torch.randn(2, dtype=torch.float64, out=out) + assert torch.randn is forbidden_random_draw + assert torch.randn_like is forbidden_random_draw + assert N2 is out + assert len(captured) == 3 + for N_recorded, N_actual in zip(captured, (N0, N1, N2), strict=True): + assert torch.equal(N_recorded, N_actual) + streams.append([N.clone() for N in captured]) + + assert len(streams[0]) == len(streams[1]) == 3 + for N0, N1 in zip(streams[0], streams[1], strict=True): + assert torch.equal(N0, N1) + + +def test_boltz2_inference_config_normalization_is_narrow() -> None: + base = { + "core_kwargs": { + "atom_s": 128, + "pairformer_args": {"v2": True}, + "msa_args": {}, + "diffusion_process_args": {"step_scale": 1.5}, + } + } + training_and_backend_variant = { + "core_kwargs": { + "atom_s": 128, + "pairformer_args": { + "activation_checkpointing": True, + "dropout": 0.25, + "post_layer_norm": False, + "use_trifast": True, + "v2": True, + }, + "msa_args": { + "activation_checkpointing": True, + "miniformer_blocks": False, + "msa_dropout": 0.15, + "subsample_msa": True, + "z_dropout": 0.25, + }, + "diffusion_process_args": { + "mse_rotational_alignment": True, + "step_scale": 1.5, + "step_scale_random": [1.0, 1.5], + }, + }, + "dtype": "float32", + } + normalize = boltz2_bundle.normalize_inference_config_contract + assert normalize(semantic_config_contract(base)) == normalize( + semantic_config_contract(training_and_backend_variant) + ) + + changed_architecture = copy.deepcopy(training_and_backend_variant) + changed_architecture["core_kwargs"]["atom_s"] = 256 + assert normalize(semantic_config_contract(base)) != normalize( + semantic_config_contract(changed_architecture) + ) + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.large +def test_boltz2_live_folding_matches_pinned_official() -> None: + request_path, reference_path, candidate_path = _paths() + request = load_request(request_path) + reference_tensors, reference_metadata = load_bundle(reference_path) + candidate_tensors, candidate_metadata = load_bundle(candidate_path) + _assert_bundle_identity(reference_metadata, request, producer="reference") + _assert_bundle_identity(candidate_metadata, request, producer="candidate") + reference_environment = reference_metadata["environment"] + candidate_environment = candidate_metadata["environment"] + assert isinstance(reference_environment, Mapping) + assert isinstance(candidate_environment, Mapping) + assert_same_hopper_sm90_device(candidate_environment, reference_environment) + expected_keys = set(candidate_metadata["state"]["tensors"]) + canonical_reference = boltz2_bundle.canonicalize_reference_state_contract( + reference_metadata["state"], + expected_keys=expected_keys, + ) + assert canonical_reference["tensors"] == candidate_metadata["state"]["tensors"] + assert canonical_reference["aliases"] == candidate_metadata["state"]["aliases"] + + official_hparams = dict(reference_metadata["semantic_config"]["fields"]) + canonical_config = Boltz2Config.from_hyperparameters( + official_hparams, + use_kernels=False, + ) + canonical_config_contract = boltz2_bundle.normalize_inference_config_contract( + semantic_config_contract(canonical_config) + ) + candidate_config_contract = boltz2_bundle.normalize_inference_config_contract( + candidate_metadata["semantic_config"] + ) + assert canonical_config_contract == candidate_config_contract + _assert_exact_features( + candidate_tensors, + candidate_metadata, + reference_tensors, + reference_metadata, + ) + assert torch.equal( + candidate_tensors["noise__initial_standard_normal"], + reference_tensors["noise__initial_standard_normal"], + ) + candidate_noise = { + name: tensor + for name, tensor in candidate_tensors.items() + if name.startswith("noise__draw_") + } + reference_noise = { + name: tensor + for name, tensor in reference_tensors.items() + if name.startswith("noise__draw_") + } + assert candidate_noise.keys() == reference_noise.keys() + for name in candidate_noise: + assert torch.equal(candidate_noise[name], reference_noise[name]), name + assert ( + candidate_metadata["diffusion_noise_draw_count"] + == reference_metadata["diffusion_noise_draw_count"] + ) + assert ( + candidate_metadata["diffusion_noise_sha256"] == reference_metadata["diffusion_noise_sha256"] + ) + _assert_valid_outputs(reference_tensors, context="official Boltz2") + _assert_valid_outputs(candidate_tensors, context="FastPLMs Boltz2") + metrics, relative_l2 = _metrics(candidate_tensors, reference_tensors) + failures = [] + for name, value in relative_l2.items(): + if value > bf16_relative_l2_hard_limit: + failures.append( + f"{name} relative L2 {value:.6g} exceeds hard limit " + f"{bf16_relative_l2_hard_limit:.6g}" + ) + elif value > bf16_relative_l2_target: + failures.append( + f"{name} relative L2 {value:.6g} misses engineering target " + f"{bf16_relative_l2_target:.6g}" + ) + for name, hard_limit in bf16_hard_limits.items(): + value = metrics[name] + if name == "lddt_ca": + if value < hard_limit: + failures.append(f"{name} {value:.6g} is below hard limit {hard_limit:.6g}") + elif value > hard_limit: + failures.append(f"{name} {value:.6g} exceeds hard limit {hard_limit:.6g}") + for name, target in bf16_targets.items(): + value = metrics[name] + if name == "lddt_ca": + if value < target and value >= bf16_hard_limits[name]: + failures.append(f"{name} {value:.6g} misses target {target:.6g}") + elif value > target and value <= bf16_hard_limits[name]: + failures.append(f"{name} {value:.6g} misses target {target:.6g}") + assert not failures, "Boltz2 compliance failures:\n- " + "\n- ".join(failures) diff --git a/testing/test_esmfold2.py b/tests/structure/test_esmfold2.py similarity index 52% rename from testing/test_esmfold2.py rename to tests/structure/test_esmfold2.py index 6b95cc1..e1a59dd 100644 --- a/testing/test_esmfold2.py +++ b/tests/structure/test_esmfold2.py @@ -1,73 +1,70 @@ -"""ESMFold2 AutoModel and parity tests.""" +"""ESMFold2 package behavior tests. + +Official parity is exercised only through the isolated bundles in +``test_esmfold2_folding_compliance.py``. +""" + from __future__ import annotations import importlib -import os -import subprocess -import sys -import tempfile -from types import SimpleNamespace - import pytest import torch -from transformers import AutoModel +from types import SimpleNamespace +from transformers.modeling_utils import PreTrainedModel -from fastplms.esm_plusplus.modeling_esm_plusplus import ( +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( ESMplusplusConfig, ESMplusplusForMaskedLM, ESMplusplusModel, ) -from fastplms.esmfold2.configuration_esmfold2 import ESMFold2Config -from fastplms.esmfold2.modeling_esmfold2 import ( +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.modeling_esmfold2 import ( _load_fastplms_esmplusplus_for_esmfold2, + _manifest_esmc_checkpoint_contract, ) -from fastplms.esmfold2.modeling_esmfold2_common import ( +from fastplms.models.esmfold2.modeling_esmfold2_common import ( compute_lm_hidden_states, maybe_apply_msa_column_masking, maybe_subsample_msa, ) -from fastplms.esmfold2.modeling_esmc import ( - _PyTorchLayerNormLinear, - _PyTorchLayerNormMLP, -) -from testing.conftest import STRUCTURE_MODEL_REGISTRY +from fastplms.registry import ModelSpec, get_model_registry -ESMFOLD2_MODEL_KEYS = ("esmfold2", "esmfold2_fast") + +REGISTRY = get_model_registry() +ESMFOLD2_MODEL_KEYS = tuple(spec.id for spec in REGISTRY.by_family("esmfold2")) TEST_SEQUENCE = "MSTNPKPQRKTKRNT" -OUTPUT_TOLERANCES = { - "distogram_logits": 0.0, - "plddt": 1e-6, - "pae": 0.0, - "ptm": 0.0, - "iptm": 0.0, -} def test_esmfold2_config_uses_fastplms_esmplusplus_defaults() -> None: config = ESMFold2Config() + backbone = REGISTRY[REGISTRY.families["esmfold2"].backbone_model] - assert config.esmc_id == "Synthyra/ESMplusplus_6B" - assert config.esmc_attn_backend == "flex" + assert config.esmc_id == backbone.fast.repo_id + assert config.esmc_attn_backend is None assert config.lm_mask_pct == 0.0 + legacy = ESMFold2Config(esmc_attn_backend="flex") + assert legacy.esmc_attn_backend == "flex_attention" + def test_esmfold2_config_normalizes_legacy_esmc_ids() -> None: - config = ESMFold2Config(esmc_id="biohub/ESMC-6B") + backbone = REGISTRY[REGISTRY.families["esmfold2"].backbone_model] + config = ESMFold2Config(esmc_id=backbone.official.repo_id) - assert config.esmc_id == "Synthyra/ESMplusplus_6B" + assert config.esmc_id == backbone.fast.repo_id -def test_esmc_pytorch_fallback_accepts_fp32_inputs_with_bf16_weights() -> None: - ln_linear = _PyTorchLayerNormLinear(d_in=8, d_out=12).to(dtype=torch.bfloat16) - ln_mlp = _PyTorchLayerNormMLP(hidden_size=8, ffn_hidden_size=16).to(dtype=torch.bfloat16) - x = torch.randn(2, 4, 8, dtype=torch.float32) +def test_esmfold2_esmc_source_uses_manifest_revision_and_file_identities() -> None: + revision, files = _manifest_esmc_checkpoint_contract("biohub/ESMC-6B") + expected = REGISTRY["esmc_6b"].fast + + assert revision == expected.revision + assert files == {item.path: item.encoded for item in expected.files} - with torch.no_grad(): - linear_out = ln_linear(x) - mlp_out = ln_mlp(x) - assert linear_out.dtype == torch.bfloat16 - assert mlp_out.dtype == torch.bfloat16 +def test_esmfold2_rejects_other_remote_esmc_checkpoints() -> None: + with pytest.raises(ValueError, match="is not the manifest-declared ESMFold2 backbone"): + _manifest_esmc_checkpoint_contract("biohub/ESMC-300M") def test_esmplusplus_sequence_id_masks_cross_chain_attention() -> None: @@ -79,7 +76,9 @@ def test_esmplusplus_sequence_id_masks_cross_chain_attention() -> None: attn_backend="sdpa", ) model = ESMplusplusModel(config).eval() + # input_ids: (1, 4) input_ids = torch.tensor([[0, 3, 4, 2]], dtype=torch.long) + # sequence_id: (1, 4) sequence_id = torch.tensor([[0, 0, 1, 1]], dtype=torch.long) with torch.no_grad(): @@ -105,28 +104,18 @@ def test_esmplusplus_sequence_id_masks_cross_chain_attention() -> None: ) -def test_esmplusplus_flex_sequence_id_masks_run() -> None: - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for flex attention sequence_id regression.") - device = torch.device("cuda") +@pytest.mark.gpu +def test_esmplusplus_rejects_unadvertised_flex_attention() -> None: config = ESMplusplusConfig( vocab_size=16, hidden_size=64, num_attention_heads=4, num_hidden_layers=1, - attn_backend="flex", + attn_backend="flex_attention", ) - model = ESMplusplusModel(config).to(device=device).eval() - input_ids = torch.tensor([[0, 3, 4, 2]], device=device, dtype=torch.long) - sequence_id = torch.tensor([[0, 0, 1, 1]], device=device, dtype=torch.long) - - try: - with torch.no_grad(): - output = model(input_ids=input_ids, sequence_id=sequence_id) - except (AssertionError, RuntimeError) as error: - pytest.skip(f"Flex attention unavailable in this environment: {error}") - assert output.last_hidden_state.shape == (1, 4, 64) + with pytest.raises(ValueError, match=r"does not support 'flex_attention'"): + ESMplusplusModel(config) def test_esmplusplus_esmfold2_hidden_state_layout() -> None: @@ -138,6 +127,7 @@ def test_esmplusplus_esmfold2_hidden_state_layout() -> None: attn_backend="sdpa", ) model = ESMplusplusModel(config).eval() + # input_ids: (1, 4) input_ids = torch.tensor([[0, 3, 4, 2]], dtype=torch.long) with torch.no_grad(): @@ -153,17 +143,47 @@ def test_esmplusplus_esmfold2_hidden_state_layout() -> None: assert len(public_output.hidden_states) == config.num_hidden_layers + 1 assert len(esmfold2_output.hidden_states) == config.num_hidden_layers + 1 torch.testing.assert_close( - esmfold2_output.hidden_states[0], + public_output.hidden_states[0], model.embed(input_ids), rtol=0.0, atol=0.0, ) - torch.testing.assert_close( - esmfold2_output.hidden_states[-1], - public_output.hidden_states[-1], - rtol=0.0, - atol=0.0, + for public_state, esmfold2_state in zip( + public_output.hidden_states, + esmfold2_output.hidden_states, + strict=True, + ): + torch.testing.assert_close(public_state, esmfold2_state, rtol=0.0, atol=0.0) + + +def test_esmplusplus_config_has_official_special_token_ids() -> None: + config = ESMplusplusConfig() + + assert config.pad_token_id == 1 + assert config.mask_token_id == 32 + assert config.classifier_dropout == 0.1 + assert config.initializer_range == 0.02 + assert config.tie_word_embeddings is False + + +def test_esmplusplus_boolean_padding_keeps_eager_outputs_finite() -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="eager", ) + model = ESMplusplusModel(config).eval() + # input_ids: (1, 4) + input_ids = torch.tensor([[0, 3, 1, 1]], dtype=torch.long) + # sequence_id: (1, 4) + sequence_id = torch.tensor([[True, True, False, False]]) + + with torch.no_grad(): + output = model(input_ids=input_ids, sequence_id=sequence_id) + + assert torch.isfinite(output.last_hidden_state).all() def test_esmplusplus_masked_lm_can_skip_logits() -> None: @@ -175,6 +195,7 @@ def test_esmplusplus_masked_lm_can_skip_logits() -> None: attn_backend="sdpa", ) model = ESMplusplusForMaskedLM(config).eval() + # input_ids: (1, 4) input_ids = torch.tensor([[0, 3, 4, 2]], dtype=torch.long) with torch.no_grad(): @@ -186,6 +207,22 @@ def test_esmplusplus_masked_lm_can_skip_logits() -> None: assert with_logits.logits.shape == (1, 4, config.vocab_size) +def test_esmplusplus_auto_classes_share_exact_checkpoint_keys() -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="sdpa", + ) + + base_keys = set(ESMplusplusModel(config).state_dict()) + masked_lm_keys = set(ESMplusplusForMaskedLM(config).state_dict()) + + assert base_keys == masked_lm_keys + assert "sequence_head.0.weight" in base_keys + + def test_esmfold2_loads_shared_esmplusplus_adapter(tmp_path) -> None: config = ESMplusplusConfig( vocab_size=16, @@ -202,7 +239,9 @@ def test_esmfold2_loads_shared_esmplusplus_adapter(tmp_path) -> None: device=torch.device("cpu"), dtype=torch.float32, ) + # input_ids: (1, 4) input_ids = torch.tensor([[0, 3, 4, 2]], dtype=torch.long) + # sequence_id: (1, 4) sequence_id = torch.tensor([[0, 0, 1, 1]], dtype=torch.long) with torch.no_grad(): @@ -216,20 +255,67 @@ def test_esmfold2_loads_shared_esmplusplus_adapter(tmp_path) -> None: assert output.hidden_states.shape == (config.num_hidden_layers + 1, 1, 4, 16) -def test_esmfold2_load_esmc_fp8_requires_transformer_engine(monkeypatch) -> None: - import fastplms.esmfold2.modeling_esmfold2 as esmfold2_module - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model +def test_esmfold2_forwards_manifest_revision_to_esmc_loaders(monkeypatch) -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="sdpa", + ) + calls: list[tuple[str, str, dict[str, object]]] = [] + + def load_config(cls, source: str, **kwargs): + del cls + calls.append(("config", source, kwargs)) + return config + + def load_model(cls, source: str, **kwargs): + del cls + calls.append(("model", source, kwargs)) + return ESMplusplusModel(config) - model = SimpleNamespace() - monkeypatch.setattr(esmfold2_module, "TE_AVAILABLE", False) + monkeypatch.setattr(ESMplusplusConfig, "from_pretrained", classmethod(load_config)) + monkeypatch.setattr(ESMplusplusModel, "from_pretrained", classmethod(load_model)) - with pytest.raises(RuntimeError, match="requires transformer_engine"): + adapter = _load_fastplms_esmplusplus_for_esmfold2( + esmc_model_path="Synthyra/ESMplusplus_6B", + attn_backend="sdpa", + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_revision = REGISTRY["esmc_6b"].fast.revision + assert adapter.config is config + assert calls[0] == ( + "config", + "Synthyra/ESMplusplus_6B", + {"revision": expected_revision}, + ) + assert calls[1][0:2] == ("model", "Synthyra/ESMplusplus_6B") + assert calls[1][2]["revision"] == expected_revision + assert calls[1][2]["config"] is config + assert calls[1][2]["torch_dtype"] == torch.float32 + + +def test_esmfold2_load_esmc_fp8_is_strict_when_unavailable(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as esmfold2_module + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model + + model = SimpleNamespace(device=torch.device("cpu")) + monkeypatch.setattr( + esmfold2_module, + "_te_fp8_capability", + lambda _device: (False, "Transformer Engine reports FP8 unavailable."), + ) + + with pytest.raises(RuntimeError, match="Transformer Engine reports FP8 unavailable"): ESMFold2Model.load_esmc(model, "unused", precision="fp8") -def test_esmfold2_load_esmc_fp8_converts_fastplms_adapter(monkeypatch) -> None: - import fastplms.esmfold2.modeling_esmfold2 as esmfold2_module - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model +def test_esmfold2_load_esmc_auto_selects_bf16_without_probing_fp8(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as esmfold2_module + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model class TinyAdapter(torch.nn.Module): def __init__(self) -> None: @@ -242,10 +328,7 @@ def __init__(self) -> None: def fake_loader(*, esmc_model_path, attn_backend, device, dtype): calls["loader"] = (esmc_model_path, attn_backend, device, dtype) - return adapter.to(device=device, dtype=dtype) - - def fake_converter(module): - calls["converter"] = module + return adapter.to(dtype=dtype) model = SimpleNamespace( config=SimpleNamespace( @@ -253,9 +336,13 @@ def fake_converter(module): lm_d_model=16, lm_num_layers=1, ), - device=torch.device("cpu"), + device=torch.device("cuda"), + ) + monkeypatch.setattr( + esmfold2_module, + "_te_fp8_capability", + lambda _device: pytest.fail("auto must not probe FP8 capability"), ) - monkeypatch.setattr(esmfold2_module, "TE_AVAILABLE", True) monkeypatch.setattr( esmfold2_module, "_load_fastplms_esmplusplus_for_esmfold2", @@ -263,32 +350,113 @@ def fake_converter(module): ) monkeypatch.setattr( esmfold2_module, - "_convert_te_modules_to_fp8_inplace", - fake_converter, + "_convert_esmc_attention_outputs_to_te", + lambda _adapter: pytest.fail("auto must not convert ESMC to FP8"), ) - ESMFold2Model.load_esmc(model, "dummy-esm", precision="fp8") + ESMFold2Model.load_esmc(model, "dummy-esm", precision="auto") + expected_device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cuda") + ) assert calls["loader"] == ( "dummy-esm", "sdpa", - torch.device("cpu"), + expected_device, torch.bfloat16, ) - assert calls["converter"] is adapter assert model._esmc is adapter - assert model._esmc_fp8 is True + assert model._esmc_fp8 is False + assert model._esmc_fp8_module_paths == () + assert model._esmc_precision_status.requested == "auto" + assert model._esmc_precision_status.resolved == "bf16" + assert "defaults to BF16" in model._esmc_precision_status.reason assert model._ttt_lm_head is None assert all(not parameter.requires_grad for parameter in adapter.parameters()) -def test_esmfold2_ttt_rejects_fp8_adapter() -> None: - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model +@pytest.mark.parametrize("experimental", [False, True], ids=("released", "experimental")) +def test_esmfold2_from_pretrained_preserves_loading_info(monkeypatch, experimental) -> None: + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model + from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, + ) + + model_class = ESMFold2ExperimentalModel if experimental else ESMFold2Model + config = SimpleNamespace(esmc_id="dummy-esmc", esmc_precision="auto") + calls: list[tuple[str, str]] = [] + model = SimpleNamespace( + config=config, + load_esmc=lambda source, *, precision: calls.append((source, precision)), + ) + loading_info = { + "missing_keys": [], + "unexpected_keys": [], + "mismatched_keys": [], + "error_msgs": [], + } + + def fake_from_pretrained(cls, source, *args, **kwargs): + del args + assert cls is model_class + assert source == "dummy-fold" + assert kwargs["config"] is config + assert kwargs["output_loading_info"] is True + return model, loading_info + + monkeypatch.setattr( + PreTrainedModel, + "from_pretrained", + classmethod(fake_from_pretrained), + ) + + loaded = model_class.from_pretrained( + "dummy-fold", + config=config, + output_loading_info=True, + esmc_precision="bf16", + ) + + assert loaded == (model, loading_info) + assert calls == [("dummy-esmc", "bf16")] + - model = SimpleNamespace(_esmc=torch.nn.Linear(1, 1), _esmc_fp8=True) +def test_esmfold2_ttt_reloads_canonical_bf16_adapter() -> None: + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model - with pytest.raises(RuntimeError, match="TTT is not supported with fp8"): - ESMFold2Model._ttt_get_trainable_modules(model) + class DummyModel: + _ensure_ttt_bf16 = ESMFold2Model._ensure_ttt_bf16 + _ttt_get_trainable_modules = ESMFold2Model._ttt_get_trainable_modules + + def __init__(self) -> None: + self._esmc = torch.nn.Linear(1, 1) + self._esmc_fp8 = True + self._esmc_precision_policy = "auto" + self._esmc_precision_status = SimpleNamespace( + device="cpu", + transformer_engine_version=None, + ) + self.config = SimpleNamespace(esmc_precision="auto") + self.device = torch.device("cpu") + self.reload_precision = None + + def reload_esmc(self, precision="auto", device=None) -> None: + self.reload_precision = (precision, device) + self._esmc = torch.nn.Linear(1, 1, dtype=torch.bfloat16) + self._esmc_fp8 = False + + model = DummyModel() + trainable = model._ttt_get_trainable_modules() + + assert model.reload_precision == ("bf16", torch.device("cpu")) + assert model._esmc_fp8 is False + assert model.config.esmc_precision == "auto" + assert model._esmc_precision_status.requested == "auto" + assert model._esmc_precision_status.resolved == "bf16" + assert trainable == [model._esmc] + assert next(model._esmc.parameters()).dtype == torch.bfloat16 def test_compute_lm_hidden_states_pads_and_masks_non_special_tokens() -> None: @@ -300,10 +468,13 @@ def __init__(self) -> None: def forward(self, input_ids, sequence_id, output_hidden_states): assert output_hidden_states is True + # input_ids: (b, l) self.input_ids = input_ids.detach().clone() + # sequence_id: (b, l) self.sequence_id = sequence_id.detach().clone() num_layers = 2 hidden_size = 3 + # hidden_states: (num_layers, *input_ids.shape, hidden_size) hidden_states = torch.arange( num_layers * input_ids.numel() * hidden_size, dtype=torch.float32, @@ -311,8 +482,11 @@ def forward(self, input_ids, sequence_id, output_hidden_states): return SimpleNamespace(hidden_states=hidden_states) esmc = CapturingEsmc() + # input_ids: (1, 3) input_ids = torch.tensor([[5, 6, 7]], dtype=torch.long) + # asym_id: (1, 3) asym_id = torch.tensor([[0, 0, 0]], dtype=torch.long) + # residue_index: (1, 3) residue_index = torch.tensor([[0, 1, 2]], dtype=torch.long) mol_type = torch.zeros_like(input_ids) token_mask = torch.ones_like(input_ids, dtype=torch.bool) @@ -337,9 +511,13 @@ def forward(self, input_ids, sequence_id, output_hidden_states): def test_msa_subsample_keeps_query_row() -> None: + # msa: (1, 5, 4) msa = torch.arange(20, dtype=torch.long).reshape(1, 5, 4) + # msa_attention_mask: (1, 5, 4) msa_attention_mask = torch.ones(1, 5, 4, dtype=torch.bool) + # has_deletion: (1, 5, 4) has_deletion = torch.zeros(1, 5, 4, dtype=torch.bool) + # deletion_value: (1, 5, 4) deletion_value = torch.zeros(1, 5, 4) torch.manual_seed(0) @@ -363,6 +541,7 @@ def test_msa_subsample_keeps_query_row() -> None: def test_msa_column_masking_keeps_query_row() -> None: + # msa_attention_mask: (2, 3, 4) msa_attention_mask = torch.ones(2, 3, 4, dtype=torch.bool) masked = maybe_apply_msa_column_masking(msa_attention_mask, rate=1.0) @@ -372,34 +551,26 @@ def test_msa_column_masking_keeps_query_row() -> None: assert not masked[:, 1:, :].any() -def _enable_deterministic_forward() -> None: - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = True - torch.use_deterministic_algorithms(True) - - -def _load_official_model(model_key: str) -> torch.nn.Module: - config = STRUCTURE_MODEL_REGISTRY[model_key] - module = pytest.importorskip("transformers.models.esmfold2.modeling_esmfold2") - official_cls = module.ESMFold2Model - return ( - official_cls.from_pretrained( - config["official_path"], - load_esmc=False, - dtype=torch.float32, - ) - .eval() - .cuda() - ) +def _esmfold2_spec(model_key: str) -> ModelSpec: + spec = REGISTRY[model_key] + assert spec.family.id == "esmfold2" + return spec def _load_fast_model(model_key: str) -> torch.nn.Module: - config = STRUCTURE_MODEL_REGISTRY[model_key] + spec = _esmfold2_spec(model_key) + is_experimental = "experimental" in spec.id + module_name = ( + "fastplms.models.esmfold2.modeling_esmfold2_experimental" + if is_experimental + else "fastplms.models.esmfold2.modeling_esmfold2" + ) + class_name = "ESMFold2ExperimentalModel" if is_experimental else "ESMFold2Model" + model_class = getattr(importlib.import_module(module_name), class_name) return ( - AutoModel.from_pretrained( - config["fast_path"], - trust_remote_code=True, + model_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, load_esmc=False, dtype=torch.float32, ) @@ -408,23 +579,8 @@ def _load_fast_model(model_key: str) -> torch.nn.Module: ) -def _run_short_fold(model: torch.nn.Module) -> dict[str, torch.Tensor]: - common_module_name = ( - model.__class__.__module__.rsplit(".", 1)[0] - + ".modeling_esmfold2_common" - ) - common_module = importlib.import_module(common_module_name) - with common_module._seed_context(0), torch.no_grad(): - return model.infer_protein( - TEST_SEQUENCE, - num_loops=1, - num_sampling_steps=2, - num_diffusion_samples=1, - ) - - def test_esmfold2_fold_protein_accepts_msa_path(tmp_path, monkeypatch) -> None: - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model captured = {} @@ -456,7 +612,7 @@ def fake_fold(self, input_value, **kwargs): def test_esmfold2_fold_protein_rejects_msa_query_mismatch(tmp_path, monkeypatch) -> None: - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model def fake_fold(self, input_value, **kwargs): del self, input_value, kwargs @@ -467,12 +623,12 @@ def fake_fold(self, input_value, **kwargs): msa_path.write_text(">query\nAAAA\n", encoding="utf-8") model = object.__new__(ESMFold2Model) - with pytest.raises(AssertionError, match="MSA query does not match sequence"): + with pytest.raises(ValueError, match="MSA query does not match sequence"): ESMFold2Model.fold_protein(model, "MSTN", msa_path=msa_path) def test_esmfold2_fold_protein_without_msa_preserves_single_sequence(monkeypatch) -> None: - from fastplms.esmfold2.modeling_esmfold2 import ESMFold2Model + from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model captured = {} @@ -492,55 +648,6 @@ def fake_fold(self, input_value, **kwargs): assert protein_input.msa is None -def _aligned_rmsd( - actual: torch.Tensor, - expected: torch.Tensor, - atom_mask: torch.Tensor, -) -> torch.Tensor: - mask = atom_mask[0].bool() if atom_mask.ndim == 2 else atom_mask.bool() - actual_coords = actual[0, mask].float() - expected_coords = expected[0, mask].float() - - actual_centered = actual_coords - actual_coords.mean(dim=0, keepdim=True) - expected_centered = expected_coords - expected_coords.mean(dim=0, keepdim=True) - cov = actual_centered.T @ expected_centered - u, _, vh = torch.linalg.svd(cov) - det = torch.det(u @ vh) - correction = torch.eye(3, device=actual.device, dtype=torch.float32) - correction[2, 2] = torch.sign(det) - rotation = u @ correction @ vh - aligned = actual_centered @ rotation - return torch.sqrt(torch.mean(torch.sum((aligned - expected_centered) ** 2, dim=-1))) - - -def _assert_forward_parity(model_key: str) -> None: - _enable_deterministic_forward() - official_model = _load_official_model(model_key) - fast_model = _load_fast_model(model_key) - - official_output = _run_short_fold(official_model) - fast_output = _run_short_fold(fast_model) - - for key, atol in OUTPUT_TOLERANCES.items(): - torch.testing.assert_close( - fast_output[key], - official_output[key], - rtol=0.0, - atol=atol, - msg=f"ESMFold2 output mismatch: {key}", - ) - - rmsd = _aligned_rmsd( - fast_output["sample_atom_coords"], - official_output["sample_atom_coords"], - official_output["atom_pad_mask"], - ) - assert rmsd.item() < 1e-2, f"Aligned coordinate RMSD too high: {rmsd.item()}" - - del official_model, fast_model, official_output, fast_output - torch.cuda.empty_cache() - - @pytest.mark.structure @pytest.mark.gpu @pytest.mark.slow @@ -560,58 +667,6 @@ def test_esmfold2_automodel_loads(model_key: str) -> None: torch.cuda.empty_cache() -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -@pytest.mark.parametrize("model_key", ESMFOLD2_MODEL_KEYS) -def test_esmfold2_weight_parity(model_key: str) -> None: - official_model = _load_official_model(model_key) - fast_model = _load_fast_model(model_key) - - official_state = official_model.state_dict() - fast_state = fast_model.state_dict() - assert official_state.keys() == fast_state.keys() - - for name, official_tensor in official_state.items(): - fast_tensor = fast_state[name] - torch.testing.assert_close( - fast_tensor, - official_tensor, - rtol=0.0, - atol=0.0, - msg=f"ESMFold2 parameter mismatch: {name}", - ) - - del official_model, fast_model - torch.cuda.empty_cache() - - -@pytest.mark.structure -@pytest.mark.gpu -@pytest.mark.slow -@pytest.mark.parametrize("model_key", ESMFOLD2_MODEL_KEYS) -def test_esmfold2_forward_parity(model_key: str) -> None: - env = os.environ.copy() - with tempfile.TemporaryDirectory() as module_cache: - env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - env["HF_MODULES_CACHE"] = module_cache - result = subprocess.run( - [ - sys.executable, - __file__, - "--esmfold2-forward-parity", - model_key, - ], - capture_output=True, - text=True, - check=False, - env=env, - ) - if result.returncode != 0 and "Skipped:" in result.stderr: - pytest.skip(result.stderr.split("Skipped:", 1)[1].strip()) - assert result.returncode == 0, result.stdout + result.stderr - - @pytest.mark.structure @pytest.mark.gpu @pytest.mark.slow @@ -649,9 +704,3 @@ def test_esmfold2_input_builder_complex_and_exports() -> None: del model, features, result torch.cuda.empty_cache() - - -if __name__ == "__main__": - assert len(sys.argv) == 3 - assert sys.argv[1] == "--esmfold2-forward-parity" - _assert_forward_parity(sys.argv[2]) diff --git a/tests/structure/test_esmfold2_complex_identity.py b/tests/structure/test_esmfold2_complex_identity.py new file mode 100644 index 0000000..47e8668 --- /dev/null +++ b/tests/structure/test_esmfold2_complex_identity.py @@ -0,0 +1,85 @@ +"""Fast, source-local contracts for ESMFold2 complex identity and storage.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from fastplms.models.esmfold2.esmfold2_molecular_complex import MolecularComplex +from fastplms.models.esmfold2.esmfold2_protein_complex import ( + ProteinComplex, + ProteinComplexMetadata, +) + + +pytestmark = pytest.mark.structure + + +def _homomer_with_repeated_chain_label() -> ProteinComplex: + sequence = "AC|AC|GG" + length = len(sequence) + # atom37_positions: (length, 37, 3) + atom37_positions = np.full((length, 37, 3), np.nan, dtype=np.float32) + # atom37_mask: (length, 37) + atom37_mask = np.zeros((length, 37), dtype=bool) + for index, residue in enumerate(sequence): + if residue == "|": + continue + atom37_positions[index, 0] = (index, 0.0, 0.0) + atom37_positions[index, 4] = (index, 1.0, 0.0) + atom37_mask[index, (0, 4)] = True + + return ProteinComplex( + id="repeated-chain-homomer", + sequence=sequence, + entity_id=np.asarray([7, 7, -1, 7, 7, -1, 9, 9], dtype=np.int64), + chain_id=np.asarray([0, 0, -1, 0, 0, -1, 1, 1], dtype=np.int64), + sym_id=np.asarray([0, 0, 0, 1, 1, 0, 0, 0], dtype=np.int64), + residue_index=np.asarray([1, 2, -1, 1, 2, -1, 1, 2], dtype=np.int64), + insertion_code=np.asarray([""] * length, dtype=object), + atom37_positions=atom37_positions, + atom37_mask=atom37_mask, + confidence=np.asarray([0.9, 0.8, 0.0, 0.7, 0.6, 0.0, 0.5, 0.4]), + metadata=ProteinComplexMetadata( + entity_lookup={7: 101, 9: 202}, + chain_lookup={0: "A", 1: "B"}, + assembly_composition={"1": ["A", "A", "B"]}, + ), + ) + + +def test_molecular_round_trip_preserves_identity_and_repeated_chain_boundaries() -> None: + original = _homomer_with_repeated_chain_label() + + molecular = MolecularComplex.from_protein_complex(original) + restored_molecular = MolecularComplex.from_blob(molecular.to_blob()) + restored = restored_molecular.to_protein_complex() + + residue_rows = np.asarray(list(original.sequence)) != "|" + assert molecular.entity_id is not None + assert molecular.sym_id is not None + np.testing.assert_array_equal(restored_molecular.entity_id, molecular.entity_id) + np.testing.assert_array_equal(restored_molecular.sym_id, molecular.sym_id) + assert restored.sequence == original.sequence + np.testing.assert_array_equal(restored.chain_id, original.chain_id) + np.testing.assert_array_equal(restored.entity_id, original.entity_id) + np.testing.assert_array_equal( + restored.sym_id[residue_rows], original.sym_id[residue_rows] + ) + assert len(list(restored.chain_iter())) == 3 + assert restored.metadata.chain_lookup == original.metadata.chain_lookup + assert restored.metadata.entity_lookup == original.metadata.entity_lookup + assert restored.metadata.assembly_composition == original.metadata.assembly_composition + + +def test_backbone_state_dict_does_not_mutate_source_atom_mask() -> None: + complex_value = _homomer_with_repeated_chain_label() + original_mask = complex_value.atom37_mask.copy() + + backbone_state = complex_value.state_dict(backbone_only=True) + + np.testing.assert_array_equal(complex_value.atom37_mask, original_mask) + assert original_mask[:, 4].any() + assert not backbone_state["atom37_mask"][:, 3:].any() + full_state = complex_value.state_dict() + assert len(full_state["atom37_positions"]) == int(original_mask.sum()) diff --git a/tests/structure/test_esmfold2_embeddings_fp8.py b/tests/structure/test_esmfold2_embeddings_fp8.py new file mode 100644 index 0000000..143236a --- /dev/null +++ b/tests/structure/test_esmfold2_embeddings_fp8.py @@ -0,0 +1,518 @@ +from __future__ import annotations + +import inspect +import subprocess +import sys +import pytest +import torch +from types import SimpleNamespace +from torch import nn + +from fastplms.embeddings import EmbeddingBatch +from fastplms.models.esmfold2.attention import ESMFold2AttentionMixin +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.embedding import ESMFold2EmbeddingMixin +from fastplms.models.esmfold2.modeling_esmfold2 import ( + ESMFold2Model, + _convert_esmc_attention_outputs_to_te, + _drop_transient_esmc_state, + _install_esmc_backbone, + _resolve_esmc_precision, +) +from fastplms.models.esmfold2.modeling_esmfold2_common import ( + LanguageModelShim, + compute_lm_hidden_states, +) +from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, +) +from fastplms.registry import get_model_registry + + +def test_language_model_projection_matches_original_operation() -> None: + torch.manual_seed(0) + shim = LanguageModelShim(d_z=7, d_model=11, num_layers=3) + # H: (2, 5, 4, 11) + H = torch.randn(2, 5, 4, 11) + # M: (2, 5) + M = torch.tensor([[True, True, True, False, False], [True, True, True, True, True]]) + + projected_states = shim.base_z_linear(H) + expected = shim.base_z_combine.softmax(dim=0) @ projected_states + expected = expected * M.unsqueeze(-1) + Z = shim.project_sequence(H, M) + + assert torch.equal(Z, expected) + assert tuple(Z.shape) == (2, 5, 7) + assert not any(key.startswith("project") for key in shim.state_dict()) + assert set(shim.state_dict()) == { + "base_z_combine", + "base_z_mlp.0.downproject.weight", + "base_z_mlp.0.downproject.bias", + "base_z_mlp.0.output_mlp.0.weight", + "base_z_mlp.0.output_mlp.0.bias", + "base_z_mlp.0.output_mlp.2.weight", + "base_z_mlp.0.output_mlp.2.bias", + "base_z_mlp.1.weight", + "base_z_mlp.1.bias", + "base_z_linear.0.weight", + "base_z_linear.0.bias", + "base_z_linear.1.weight", + } + + +def test_language_model_projection_preserves_single_residue_axis() -> None: + shim = LanguageModelShim(d_z=7, d_model=11, num_layers=3) + # H: (2, 1, 4, 11) + H = torch.randn(2, 1, 4, 11) + # M: (2, 1) + M = torch.ones((2, 1), dtype=torch.bool) + + Z = shim.project_sequence(H, M) + + assert Z.shape == (2, 1, 7) + + +def test_language_model_projection_matches_checkpoint_dtype() -> None: + shim = LanguageModelShim(d_z=7, d_model=11, num_layers=3).to(dtype=torch.bfloat16) + # H: (2, 5, 4, 11) + H = torch.randn(2, 5, 4, 11, dtype=torch.float32) + + Z = shim.project_sequence(H) + + assert Z.dtype == torch.bfloat16 + assert torch.isfinite(Z).all() + + +def test_projection_validates_official_state_count() -> None: + shim = LanguageModelShim() + # H: (1, 2, 80, 2560) + H = torch.zeros(1, 2, 80, 2560) + with pytest.raises(ValueError, match="expected 81"): + shim.project_sequence(H) + + model = SyntheticESMFold2() + with pytest.raises(ValueError, match="official ordered 81-state"): + model.project_esmc_hidden_states(torch.zeros(1, 2, 80, 4)) + + +class SyntheticESMFold2(ESMFold2EmbeddingMixin, nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self._esmc = object() + self.language_model = LanguageModelShim(d_z=3, d_model=4, num_layers=80) + + @property + def device(self) -> torch.device: + return self.anchor.device + + def _compute_lm_hidden_states( + self, + input_ids: torch.Tensor, + asym_id: torch.Tensor, + residue_index: torch.Tensor, + mol_type: torch.Tensor, + residue_mask: torch.Tensor, + ) -> torch.Tensor: + # input_ids: (b, l); asym_id, residue_index, mol_type: (...) + # residue_mask: (b, l) + del asym_id, residue_index, mol_type + # H: (*input_ids.shape, 81, 4) + H = torch.zeros(*input_ids.shape, 81, 4) + # H[..., 0]: (...) + H[..., 0] = input_ids.unsqueeze(-1) + return H * residue_mask[..., None, None] + + +def test_esmfold2_embedding_is_residue_only_and_rejects_complexes() -> None: + model = SyntheticESMFold2() + batch = model._embedding_batch(["ACD", "GG"]) + assert isinstance(batch, EmbeddingBatch) + assert tuple(batch.X.shape) == (2, 3, 3) + assert batch.residue_mask.tolist() == [[True, True, True], [True, True, False]] + assert torch.equal(batch.X[1, 2], torch.zeros(3)) + result = model.embed_dataset(["ACD"], full_embeddings=True) + assert result.metadata["layer"] == "all_81_esmc_states" + assert result.metadata["projection"] == "esmfold2_learned_sequence_summary" + assert "BOS" in result.metadata["token_policy"]["exclude"] + with pytest.raises(ValueError, match="one ungapped protein chain"): + model._embedding_batch(["ACD|GG"]) + with pytest.raises(ValueError, match="at least one protein residue"): + model._embedding_batch([""]) + + +def test_auto_precision_uses_bf16_without_probing_fp8(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + monkeypatch.setattr( + module, + "_te_fp8_capability", + lambda _device: pytest.fail("auto must not probe FP8 capability"), + ) + status = _resolve_esmc_precision("auto", torch.device("cuda")) + assert status.requested == "auto" + assert status.resolved == "bf16" + assert "defaults to BF16" in status.reason + + +def test_auto_precision_remains_bf16_when_transformer_engine_is_installed( + monkeypatch, +) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + monkeypatch.setattr( + module, + "_transformer_engine_version", + lambda: "2.12.0", + ) + status = _resolve_esmc_precision("auto", torch.device("cuda")) + assert status.requested == "auto" + assert status.resolved == "bf16" + assert "defaults to BF16" in status.reason + assert status.transformer_engine_version == "2.12.0" + + +def test_explicit_fp8_is_strict(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + monkeypatch.setattr( + module, + "_te_fp8_capability", + lambda device: (False, f"FP8 unavailable on {device}"), + ) + with pytest.raises(RuntimeError, match="FP8 unavailable on cuda"): + _resolve_esmc_precision("fp8", torch.device("cuda")) + + +def test_fp8_capability_rejects_non_cuda_devices() -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + available, reason = module._te_fp8_capability(torch.device("cpu")) + assert available is False + assert reason == "FP8 requires direct ESMC loading onto a CUDA device." + + +def test_fp8_converter_replaces_only_80_attention_output_projections( + monkeypatch, +) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + class FakeTELinear(nn.Linear): + def __init__( + self, + in_features, + out_features, + *, + bias, + params_dtype, + device, + ) -> None: + super().__init__( + in_features, + out_features, + bias=bias, + device=device, + dtype=params_dtype, + ) + + class Block(nn.Module): + def __init__(self) -> None: + super().__init__() + self.attn = nn.Module() + self.attn.out_proj = nn.Linear(4, 4, bias=False) + self.ffn = nn.Linear(4, 4, bias=False) + + backbone = nn.Module() + backbone.layers = nn.ModuleList([Block() for _ in range(80)]) + expected_weights = [block.attn.out_proj.weight.detach().clone() for block in backbone.layers] + monkeypatch.setattr( + module, + "_load_transformer_engine", + lambda: (SimpleNamespace(Linear=FakeTELinear), SimpleNamespace()), + ) + + paths = _convert_esmc_attention_outputs_to_te(backbone) + + assert len(paths) == 80 + assert all(path.endswith(".attn.out_proj") for path in paths) + for block, expected in zip(backbone.layers, expected_weights, strict=True): + assert isinstance(block.attn.out_proj, FakeTELinear) + assert isinstance(block.ffn, nn.Linear) + assert torch.equal(block.attn.out_proj.weight, expected) + + +class FakeBackbone(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + self.config = SimpleNamespace(hidden_size=4, num_hidden_layers=2) + + +class PrecisionOwner(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + esmc_attn_backend="flex_attention", + _attn_implementation="flex_attention", + lm_d_model=4, + lm_num_layers=2, + esmc_id="canonical-esmc", + esmc_precision="auto", + ) + self._esmc = None + self._esmc_fp8 = False + self._ttt_lm_head = None + + @property + def device(self) -> torch.device: + return self.anchor.device + + +class _AttentionOwnerBase(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.config = config + self._esmc = None + + def _check_and_adjust_attn_implementation( + self, + attn_implementation: str | None, + is_init_check: bool = False, + allow_all_kernels: bool = False, + ) -> str: + """Model the Transformers base-class hook used by the production owner.""" + + del is_init_check, allow_all_kernels + return attn_implementation or "sdpa" + + +class AttentionOwner(ESMFold2AttentionMixin, _AttentionOwnerBase): + pass + + +class RecordingAttentionBackbone(nn.Module): + def __init__(self) -> None: + super().__init__() + self.implementations: list[str] = [] + + def set_attn_implementation(self, implementation: str) -> None: + self.implementations.append(implementation) + + +def test_attention_config_defaults_to_unspecified_and_normalizes_legacy_flex() -> None: + default = ESMFold2Config() + assert default.esmc_attn_backend is None + assert default._attn_implementation is None + + legacy = ESMFold2Config(esmc_attn_backend="flex") + assert legacy.esmc_attn_backend == "flex_attention" + assert legacy._attn_implementation == "flex_attention" + + +def test_public_attention_setting_overrides_legacy_config_value() -> None: + config = ESMFold2Config( + esmc_attn_backend="flex", + attn_implementation="eager", + ) + assert config.esmc_attn_backend == "eager" + assert config._attn_implementation == "eager" + + +def test_transformers_config_reload_applies_public_attention_override(tmp_path) -> None: + ESMFold2Config(esmc_attn_backend="flex").save_pretrained(tmp_path) + + config = ESMFold2Config.from_pretrained( + tmp_path, + attn_implementation="eager", + ) + + assert config._attn_implementation == "eager" + assert config.esmc_attn_backend == "eager" + + +def test_runtime_attention_setting_updates_loaded_esmc() -> None: + config = ESMFold2Config(attn_implementation="sdpa") + owner = AttentionOwner(config) + esmc = RecordingAttentionBackbone() + owner._esmc = esmc + + owner.set_attn_implementation("eager") + + assert owner.config._attn_implementation == "eager" + assert owner.config.esmc_attn_backend == "eager" + assert esmc.implementations == ["eager"] + + +@pytest.mark.parametrize("implementation", ["flex", "flash_attention_2", "unknown"]) +def test_runtime_attention_rejects_unadvertised_names(implementation: str) -> None: + owner = AttentionOwner(ESMFold2Config(attn_implementation="sdpa")) + with pytest.raises(ValueError, match="does not support"): + owner.set_attn_implementation(implementation) + + +def test_compile_helpers_do_not_mutate_global_dynamo_configuration() -> None: + for model_class in (ESMFold2Model, ESMFold2ExperimentalModel): + source = inspect.getsource(model_class.apply_torch_compile) + assert "torch._dynamo.config" not in source + + +def test_fresh_esmfold2_import_does_not_mutate_global_torch_settings() -> None: + script = """ +import torch +import torch._dynamo + +def snapshot(): + return ( + torch.get_default_dtype(), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.backends.cudnn.benchmark, + torch.backends.cudnn.deterministic, + torch._dynamo.config.cache_size_limit, + torch._dynamo.config.accumulated_cache_size_limit, + torch._dynamo.config.capture_scalar_outputs, + ) + +before = snapshot() +import fastplms.embeddings +import fastplms.models.esmfold2.modeling_esmfold2 +after = snapshot() +assert after == before, (before, after) +""" + subprocess.run([sys.executable, "-c", script], check=True) + + +def test_install_records_bf16_auto_default_and_policy(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + owner = PrecisionOwner() + load_options: list[dict[str, object]] = [] + + def load_backbone(**kwargs): + load_options.append(kwargs) + return FakeBackbone() + + monkeypatch.setattr( + module, + "_load_fastplms_esmplusplus_for_esmfold2", + load_backbone, + ) + _install_esmc_backbone(owner, "canonical-esmc", precision="auto", device=torch.device("cpu")) + assert owner._esmc_fp8 is False + assert owner._esmc_precision_status.resolved == "bf16" + assert "defaults to BF16" in owner._esmc_precision_status.reason + assert owner.config.esmc_precision == "auto" + assert load_options[0]["attn_backend"] == "flex_attention" + + +def test_install_records_validated_fp8_projection_set(monkeypatch) -> None: + import fastplms.models.esmfold2.modeling_esmfold2 as module + + owner = PrecisionOwner().to("meta") + owner.anchor = nn.Parameter(torch.zeros((), device="meta")) + backbone = FakeBackbone().to("meta") + monkeypatch.setattr( + module, + "_te_fp8_capability", + lambda device: (True, f"FP8 available on {device}"), + ) + monkeypatch.setattr( + module, + "_load_fastplms_esmplusplus_for_esmfold2", + lambda **kwargs: backbone, + ) + paths = tuple(f"model.layers.{index}.attn.out_proj" for index in range(80)) + monkeypatch.setattr( + module, + "_convert_esmc_attention_outputs_to_te", + lambda esmc: paths, + ) + + _install_esmc_backbone(owner, "canonical-esmc", precision="fp8", device="meta") + + assert owner._esmc_fp8 is True + assert owner._esmc_fp8_module_paths == paths + assert owner._esmc_precision_status.resolved == "fp8" + assert "Converted 80 projections" in owner._esmc_precision_status.reason + + +def test_ttt_switches_runtime_fp8_back_to_bf16() -> None: + calls: list[tuple[str, torch.device]] = [] + owner = SimpleNamespace( + _esmc_fp8=True, + _esmc_precision_policy="auto", + _esmc_precision_status=SimpleNamespace(device="cuda:0", transformer_engine_version="2.16"), + config=SimpleNamespace(esmc_precision="auto"), + device=torch.device("cuda"), + reload_esmc=lambda precision, device: calls.append((precision, device)), + ) + ESMFold2Model._ensure_ttt_bf16(owner) + assert calls == [("bf16", torch.device("cuda"))] + assert owner.config.esmc_precision == "auto" + assert owner._esmc_precision_status.requested == "auto" + assert owner._esmc_precision_status.resolved == "bf16" + + +def test_runtime_esmc_state_is_not_persisted() -> None: + state = { + "language_model.base_z_combine": torch.ones(1), + "_esmc.layer.weight": torch.ones(1), + "_ttt_lm_head.weight": torch.ones(1), + } + _drop_transient_esmc_state(nn.Identity(), state, "", {}) + assert list(state) == ["language_model.base_z_combine"] + + +class RecordingESMC(nn.Module): + def __init__(self) -> None: + super().__init__() + self.seen_l = 0 + + def forward(self, input_ids, sequence_id, output_hidden_states): + del sequence_id, output_hidden_states + b, sequence_length = input_ids.shape + self.seen_l = sequence_length + # H: (81, b, sequence_length, 4) + H = torch.zeros(81, b, sequence_length, 4) + H[:] = torch.arange(81).view(81, 1, 1, 1) # RHS broadcasts from (81, 1, 1, 1) + return SimpleNamespace(hidden_states=H) + + +def test_fp8_language_model_input_is_padded_to_multiple_of_16() -> None: + esmc = RecordingESMC() + sequence_length = 15 + # input_ids: (1, sequence_length) + input_ids = torch.full((1, sequence_length), 4, dtype=torch.long) + asym_id = torch.zeros_like(input_ids) + # residue_index: (...) + residue_index = torch.arange(sequence_length).unsqueeze(0) + mol_type = torch.zeros_like(input_ids) + M = torch.ones_like(input_ids, dtype=torch.bool) + H = compute_lm_hidden_states( + esmc, + input_ids, + asym_id, + residue_index, + mol_type, + M, + pad_to_multiple=16, + ) + assert esmc.seen_l == 32 + assert tuple(H.shape) == (1, sequence_length, 81, 4) + assert torch.equal(H[0, 0, :, 0], torch.arange(81, dtype=H.dtype)) + + +def test_supported_variants_are_exactly_the_four_approved_sources() -> None: + official_repositories = { + spec.official.repo_id for spec in get_model_registry().by_family("esmfold2") + } + assert official_repositories == { + "biohub/ESMFold2", + "biohub/ESMFold2-Fast", + "biohub/ESMFold2-Experimental-Cutoff2025", + "biohub/ESMFold2-Experimental-Fast-Cutoff2025", + } + assert hasattr(ESMFold2Model, "project_esmc_hidden_states") + assert hasattr(ESMFold2ExperimentalModel, "project_esmc_hidden_states") diff --git a/tests/structure/test_esmfold2_experimental.py b/tests/structure/test_esmfold2_experimental.py new file mode 100644 index 0000000..25a5bac --- /dev/null +++ b/tests/structure/test_esmfold2_experimental.py @@ -0,0 +1,101 @@ +"""ESMFold2 experimental model tests.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from pathlib import Path + +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.modeling_esmfold2_common import NUM_RES_TYPES +from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, +) +from fastplms.models.esmfold2.protein_utils import prepare_protein_features +from fastplms.registry import ModelSpec, get_model_registry + + +TEST_SEQUENCE = "MSTNPKPQRKTKRNT" +REGISTRY = get_model_registry() +EXPERIMENTAL_SPECS = tuple( + spec for spec in REGISTRY.by_family("esmfold2") if "experimental" in spec.id +) +DEFAULT_EXPERIMENTAL_SPEC = REGISTRY["esmfold2_experimental_fast_cutoff2025"] +EXPERIMENTAL_AUTO_MAP = { + auto_class: class_path.removeprefix("fastplms.models.esmfold2.") + for auto_class, class_path in REGISTRY["esmfold2_experimental_cutoff2025"].auto_map.items() +} + + +def _load_fast_model(spec: ModelSpec) -> ESMFold2ExperimentalModel: + return ( + ESMFold2ExperimentalModel.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + load_esmc=False, + dtype=torch.float32, + ) + .eval() + .cuda() + ) +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +def test_esmfold2_experimental_res_type_soft_gradients() -> None: + model = _load_fast_model(DEFAULT_EXPERIMENTAL_SPEC) + features = { + name: tensor.cuda() for name, tensor in prepare_protein_features(TEST_SEQUENCE).items() + } + # res_type_soft: (...) + res_type_soft = F.one_hot(features["res_type"].long(), num_classes=NUM_RES_TYPES).float() + res_type_soft.requires_grad_(True) + + output = model( + **features, + res_type_soft=res_type_soft, + num_loops=0, + num_sampling_steps=1, + num_diffusion_samples=1, + calculate_confidence=False, + seed=0, + ) + # loss: () + loss = output["distogram_logits"].float().mean() + loss.backward() + + assert "representative_atom_coords" in output + assert output["representative_atom_coords"].shape[-1] == 3 + assert output["representative_atom_coords"].shape[-2] == features["res_type"].shape[1] + assert res_type_soft.grad is not None + assert torch.isfinite(res_type_soft.grad).all() + assert res_type_soft.grad.abs().sum().item() > 0 + + del model, output, features + torch.cuda.empty_cache() + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.parametrize("spec", EXPERIMENTAL_SPECS, ids=lambda spec: spec.id) +def test_esmfold2_experimental_model_loads(spec: ModelSpec) -> None: + model = _load_fast_model(spec) + + assert callable(model.infer_protein_as_pdb) + assert callable(model.fold) + assert callable(model.prepare_structure_input) + + del model + torch.cuda.empty_cache() + + +def test_esmfold2_experimental_export_config(tmp_path: Path) -> None: + config = ESMFold2Config(type="experimental") + config.auto_map = EXPERIMENTAL_AUTO_MAP + config.architectures = ["ESMFold2ExperimentalModel"] + config.save_pretrained(tmp_path) + + loaded = ESMFold2Config.from_pretrained(tmp_path) + assert loaded.auto_map == EXPERIMENTAL_AUTO_MAP + assert loaded.architectures == ["ESMFold2ExperimentalModel"] diff --git a/tests/structure/test_esmfold2_folding_compliance.py b/tests/structure/test_esmfold2_folding_compliance.py new file mode 100644 index 0000000..bde4b13 --- /dev/null +++ b/tests/structure/test_esmfold2_folding_compliance.py @@ -0,0 +1,575 @@ +"""Release gates over isolated official and candidate ESMFold2 fold bundles.""" + +from __future__ import annotations + +import ast +import inspect +import os +import pytest +import torch +import torch.nn.functional as F +from collections.abc import Callable, Mapping +from pathlib import Path + +from fastplms.registry import ModelSpec, get_model_registry +from tests.parity.support.reference_adapters.biohub_source import ( + BIOHUB_ESM_REVISION, + BIOHUB_ESM_TREE_SHA256, + BIOHUB_REFERENCE_SOURCE_NAMES, + BIOHUB_TRANSFORMERS_REVISION, + BIOHUB_TRANSFORMERS_TREE_SHA256, +) +from tests.structure.support import esmfold2_bundle +from tests.structure.support.esmfold2_bundle import load_bundle, load_request +from tests.structure.support.hardware import ( + assert_same_hopper_sm90_device, + hopper_sm90_fingerprint, +) +from tools.remote.biohub_reference_environment import ( + validate_biohub_reference_environment_evidence, +) +from tools.remote.reference_source_attestation import validate_reference_sources_evidence + + +ROOT = Path(__file__).resolve().parents[2] + +bf16_targets = { + "ca_rmsd": 0.10, + "lddt_ca": 0.995, + "plddt_mae": 0.001, + "pae_mae": 0.10, + "ptm_error": 0.002, + "iptm_error": 0.002, +} +bf16_hard_limits = { + "ca_rmsd": 0.25, + "lddt_ca": 0.99, + "plddt_mae": 0.005, + "pae_mae": 0.50, + "ptm_error": 0.005, + "iptm_error": 0.005, +} + + +def _exchange_root() -> Path: + return Path(os.environ.get("FASTPLMS_REFERENCE_EXCHANGE", "artifacts/reference")) + + +def _bundle_paths(spec: ModelSpec) -> tuple[Path, Path, Path]: + root = _exchange_root() + request = ( + root / "structure" / "requests" / esmfold2_bundle.reference_container / f"{spec.id}.json" + ) + reference = root / "structure" / "results" / "reference" / spec.id + candidate = root / "structure" / "results" / "candidate" / spec.id + return request, reference, candidate / "bf16" + + +def _feature_tensors(tensors: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return { + name.removeprefix("feature__"): tensor + for name, tensor in tensors.items() + if name.startswith("feature__") + } + + +def _output(tensors: Mapping[str, torch.Tensor], name: str) -> torch.Tensor: + key = f"output__{name}" + if key not in tensors: + raise KeyError(f"Structure bundle omits required output {name!r}.") + return tensors[key] + + +def _checkpoint_contract(checkpoint: object) -> dict[str, object]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + { + "path": item.path, + "algorithm": item.algorithm, + "digest": item.digest, + } + for item in checkpoint.files + ], + } + + +def _assert_bundle_identity( + metadata: Mapping[str, object], + request: Mapping[str, object], + spec: ModelSpec, + *, + producer: str, + precision: str, +) -> None: + assert metadata["producer"] == producer + assert metadata["model_id"] == spec.id + assert metadata["request_sha256"] == request["request_sha256"] + assert metadata["official"] == _checkpoint_contract(spec.official) + assert metadata["candidate"] == _checkpoint_contract(spec.fast) + assert metadata["sequence"] == request["sequence"] + assert metadata["seed"] == request["seed"] + assert metadata["sampling_steps"] == request["sampling_steps"] + assert metadata["attention_backend"] == request["attention_backend"] == "sdpa" + assert metadata["deterministic_algorithms"] is True + assert metadata["requested_precision"] == precision + assert metadata["resolved_precision"] == precision + status = metadata["precision_status"] + assert isinstance(status, Mapping) + assert status["requested"] == precision + assert status["resolved"] == precision + environment = metadata["environment"] + assert isinstance(environment, Mapping) + hopper_sm90_fingerprint(environment) + if producer == "candidate": + assert str(environment["torch"]).split("+", maxsplit=1)[0] == "2.13.0" + assert environment["transformers"] == "5.13.0" + assert str(environment["cuda_runtime"]).startswith("13.0") + else: + locked_environment = validate_biohub_reference_environment_evidence( + metadata.get("reference_environment"), + repository_root=ROOT, + contract_path=ROOT / "docker/constraints/biohub-reference-lock.json", + ) + assert locked_environment["reference_container_target"] == "reference-esmfold2" + sources = validate_reference_sources_evidence( + metadata.get("reference_sources"), + required_sources=BIOHUB_REFERENCE_SOURCE_NAMES, + ) + esm_source = sources["biohub-esm"] + assert esm_source["source_revision"] == BIOHUB_ESM_REVISION + assert esm_source["tree_sha256"] == BIOHUB_ESM_TREE_SHA256 + assert esm_source["package_version"] == "3.3.0" + assert esm_source["import_file"] == "esm/__init__.py" + transformers_source = sources["biohub-transformers"] + assert transformers_source["source_revision"] == BIOHUB_TRANSFORMERS_REVISION + assert transformers_source["tree_sha256"] == BIOHUB_TRANSFORMERS_TREE_SHA256 + assert transformers_source["package_version"] == "4.57.6" + assert transformers_source["import_file"] == "src/transformers/__init__.py" + if precision == "fp8": + assert status["transformer_engine_version"] + assert environment["transformer_engine"] + + +def _assert_exact_inputs( + actual_tensors: Mapping[str, torch.Tensor], + actual_metadata: Mapping[str, object], + expected_tensors: Mapping[str, torch.Tensor], + expected_metadata: Mapping[str, object], + *, + context: str, +) -> None: + actual_features = _feature_tensors(actual_tensors) + expected_features = _feature_tensors(expected_tensors) + assert actual_features.keys() == expected_features.keys(), context + for name in actual_features: + actual = actual_features[name] + expected = expected_features[name] + assert actual.dtype == expected.dtype, f"{context}: {name} dtype" + assert actual.shape == expected.shape, f"{context}: {name} shape" + assert torch.equal(actual, expected), f"{context}: {name} values" + assert actual_metadata["feature_sha256"] == expected_metadata["feature_sha256"] + assert torch.equal( + actual_tensors["noise__initial_standard_normal"], + expected_tensors["noise__initial_standard_normal"], + ), f"{context}: initial diffusion noise" + assert actual_metadata["diffusion_noise_sha256"] == expected_metadata["diffusion_noise_sha256"] + + +def _first_coordinate_sample(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + # coordinates: (..., 3) + coordinates = _output(tensors, "sample_atom_coords").float() + if coordinates.ndim == 4: + # coordinates: (-1, coordinates.shape[-2], 3) + coordinates = coordinates.reshape(-1, coordinates.shape[-2], 3) + assert coordinates.ndim == 3 and coordinates.shape[-1] == 3 + return coordinates[0] + + +def _ca_mask(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + features = _feature_tensors(tensors) + # encoded_ca: (4,) + encoded_ca = torch.tensor([ord("C") - 32, ord("A") - 32, 0, 0]) + atom_names = features["ref_atom_name_chars"][0] + # atom_mask: (...) + atom_mask = features["atom_attention_mask"][0].bool() + # mask: (...) + mask = atom_names.eq(encoded_ca).all(dim=-1) & atom_mask + token_ids = features["atom_to_token"][0, mask] + valid_token_ids = features["token_attention_mask"][0].nonzero(as_tuple=True)[0] + assert torch.equal(token_ids, valid_token_ids), ( + "Each biological residue must have exactly one C-alpha atom." + ) + return mask + + +def _ca_coordinates(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + return _first_coordinate_sample(tensors)[_ca_mask(tensors)] + + +def _aligned_ca_rmsd(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + actual_centered = actual.float() - actual.float().mean(dim=0, keepdim=True) + expected_centered = expected.float() - expected.float().mean(dim=0, keepdim=True) + covariance = actual_centered.T @ expected_centered + left, _, right = torch.linalg.svd(covariance) + # correction: (3, 3) + correction = torch.eye(3, dtype=torch.float32) + correction[-1, -1] = torch.sign(torch.det(left @ right)) + rotation = left @ correction @ right + aligned = actual_centered @ rotation + return torch.sqrt(torch.mean(torch.sum((aligned - expected_centered) ** 2, dim=-1))).item() + + +def _lddt_ca(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + actual_distances = torch.cdist(actual.float(), actual.float()) + expected_distances = torch.cdist(expected.float(), expected.float()) + # pair_mask: (...) + pair_mask = expected_distances.lt(15.0) + pair_mask.fill_diagonal_(False) + assert pair_mask.any(), "No valid C-alpha pairs for lDDT." + errors = (actual_distances - expected_distances).abs() + # score: (...) + score = torch.stack([errors.lt(threshold).float() for threshold in (0.5, 1.0, 2.0, 4.0)]).mean( + dim=0 + ) + return score[pair_mask].mean().item() + + +def _token_vector( + tensors: Mapping[str, torch.Tensor], + name: str, + sequence_length: int, +) -> torch.Tensor: + return _output(tensors, name).float().reshape(-1, sequence_length)[0] + + +def _token_pair( + tensors: Mapping[str, torch.Tensor], + name: str, + sequence_length: int, +) -> torch.Tensor: + return ( + _output(tensors, name) + .float() + .reshape( + -1, + sequence_length, + sequence_length, + )[0] + ) + + +def _structure_metrics( + actual: Mapping[str, torch.Tensor], + expected: Mapping[str, torch.Tensor], +) -> dict[str, float]: + actual_features = _feature_tensors(actual) + # token_mask: (...) + token_mask = actual_features["token_attention_mask"][0].bool() + sequence_length = token_mask.numel() + # pair_mask: (...) + pair_mask = token_mask[:, None] & token_mask[None, :] + return { + "ca_rmsd": _aligned_ca_rmsd( + _ca_coordinates(actual), + _ca_coordinates(expected), + ), + "lddt_ca": _lddt_ca( + _ca_coordinates(actual), + _ca_coordinates(expected), + ), + "plddt_mae": ( + _token_vector(actual, "plddt", sequence_length)[token_mask] + - _token_vector(expected, "plddt", sequence_length)[token_mask] + ) + .abs() + .mean() + .item(), + "pae_mae": ( + _token_pair(actual, "pae", sequence_length)[pair_mask] + - _token_pair(expected, "pae", sequence_length)[pair_mask] + ) + .abs() + .mean() + .item(), + "ptm_error": ( + _output(actual, "ptm").float().reshape(-1)[0] + - _output(expected, "ptm").float().reshape(-1)[0] + ) + .abs() + .item(), + "iptm_error": ( + _output(actual, "iptm").float().reshape(-1)[0] + - _output(expected, "iptm").float().reshape(-1)[0] + ) + .abs() + .item(), + } + + +def _probability_jsd( + actual_logits: torch.Tensor, + expected_logits: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + # actual_logits: (..., c), expected_logits: (..., c), mask: (...) + actual_log_prob = F.log_softmax(actual_logits.float(), dim=-1) + expected_log_prob = F.log_softmax(expected_logits.float(), dim=-1) + actual_prob = actual_log_prob.exp() + expected_prob = expected_log_prob.exp() + mean_prob = 0.5 * (actual_prob + expected_prob) + log_mean_prob = mean_prob.clamp_min(torch.finfo(torch.float32).tiny).log() + jsd = 0.5 * ( + (actual_prob * (actual_log_prob - log_mean_prob)).sum(dim=-1) + + (expected_prob * (expected_log_prob - log_mean_prob)).sum(dim=-1) + ) + while mask.ndim < jsd.ndim: + # mask: (...) + mask = mask.unsqueeze(0) + mask = torch.broadcast_to(mask, jsd.shape) + assert mask.any() + return jsd[mask].mean() + + +def _mean_probability_jsd( + actual: Mapping[str, torch.Tensor], + expected: Mapping[str, torch.Tensor], +) -> float: + features = _feature_tensors(actual) + # atom_mask: (...) + atom_mask = features["atom_attention_mask"].bool() + # token_mask: (...) + token_mask = features["token_attention_mask"].bool() + # pair_mask: (...) + pair_mask = token_mask[:, :, None] & token_mask[:, None, :] + values = [] + for name in ("distogram_logits", "plddt_logits", "pae_logits", "pde_logits"): + key = f"output__{name}" + if key not in actual or key not in expected: + continue + mask = atom_mask if name == "plddt_logits" else pair_mask + values.append(_probability_jsd(actual[key], expected[key], mask)) + assert values, "No probability tensors were returned for JSD compliance." + return torch.stack(values).mean().item() + + +def _assert_valid_geometry( + tensors: Mapping[str, torch.Tensor], + *, + context: str, +) -> None: + features = _feature_tensors(tensors) + coordinates = _first_coordinate_sample(tensors) + # atom_mask: (...) + atom_mask = features["atom_attention_mask"][0].bool() + assert torch.equal( + _output(tensors, "atom_pad_mask").bool().reshape_as(atom_mask), + atom_mask, + ) + assert torch.isfinite(coordinates[atom_mask]).all(), f"{context}: non-finite coordinates" + # ca_coordinates: (..., 3) + ca_coordinates = coordinates[_ca_mask(tensors)] + ca_steps = torch.linalg.vector_norm(ca_coordinates[1:] - ca_coordinates[:-1], dim=-1) + assert torch.isfinite(ca_steps).all(), f"{context}: non-finite C-alpha distances" + assert ca_steps.gt(2.0).all() and ca_steps.lt(5.0).all(), ( + f"{context}: invalid consecutive C-alpha distances {ca_steps.tolist()}" + ) + for name, tensor in tensors.items(): + if name.startswith("output__") and tensor.is_floating_point(): + assert torch.isfinite(tensor).all(), f"{context}: {name} contains NaN or inf" + + +def _assert_thresholds( + metrics: Mapping[str, float], + *, + targets: Mapping[str, float], + hard_limits: Mapping[str, float], + context: str, +) -> None: + for name, hard_limit in hard_limits.items(): + value = metrics[name] + if name == "lddt_ca": + assert value >= hard_limit, ( + f"{context}: {name}={value:.6g} is below hard limit {hard_limit:.6g}" + ) + else: + assert value <= hard_limit, ( + f"{context}: {name}={value:.6g} exceeds hard limit {hard_limit:.6g}" + ) + for name, target in targets.items(): + value = metrics[name] + if name == "lddt_ca": + assert value >= target, ( + f"{context}: {name}={value:.6g} misses engineering target {target:.6g}" + ) + else: + assert value <= target, ( + f"{context}: {name}={value:.6g} misses engineering target {target:.6g}" + ) + + +def _spec_parameter(spec: ModelSpec) -> object: + return pytest.param(spec, id=spec.id, marks=pytest.mark.large) + + +def test_structure_bundle_reference_path_has_no_fastplms_dependency() -> None: + """Keep the module copyable into a native service with FastPLMs absent.""" + + tree = ast.parse(inspect.getsource(esmfold2_bundle)) + for node in tree.body: + if isinstance(node, ast.Import): + assert all(not alias.name.startswith("fastplms") for alias in node.names) + elif isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith("fastplms") + for function in ( + esmfold2_bundle._load_reference_model, + esmfold2_bundle._run_fold, + esmfold2_bundle.produce_reference, + ): + assert "fastplms" not in inspect.getsource(function).lower() + + +def test_structure_metric_helpers_are_exact_for_rigid_identity() -> None: + # expected: (4, 3) + expected = torch.tensor([[0.0, 0.0, 0.0], [3.8, 0.0, 0.0], [7.2, 1.0, 0.0], [9.0, 4.0, 1.0]]) + # rotation: (3, 3) + rotation = torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + actual = expected @ rotation + torch.tensor([4.0, -2.0, 7.0]) + assert _aligned_ca_rmsd(actual, expected) == pytest.approx(0.0, abs=1e-5) + assert _lddt_ca(actual, expected) == pytest.approx(1.0) + # logits: (1, 2, 2) + logits = torch.tensor([[[1.0, 2.0], [0.0, -1.0]]]) + # mask: (1, 2) + mask = torch.tensor([[True, True]]) + assert _probability_jsd(logits, logits, mask).item() == pytest.approx(0.0) + + +def test_prepare_structure_requests_is_manifest_exact(tmp_path: Path) -> None: + paths = esmfold2_bundle.prepare_requests(tmp_path) + assert tuple(path.stem for path in paths) == esmfold2_bundle.supported_model_ids + registry = get_model_registry() + for path in paths: + request = load_request(path) + spec = registry[request["model_id"]] + assert request["official"] == _checkpoint_contract(spec.official) + assert request["candidate"] == _checkpoint_contract(spec.fast) + assert request["candidate_auto_model"] == spec.auto_map["AutoModel"] + assert request["backbone_model"] == spec.family.backbone_model + assert request["attention_backend"] == "sdpa" + assert request["deterministic_algorithms"] is True + + +def test_all_prepared_requests_requires_the_exact_release_inventory(tmp_path: Path) -> None: + paths = esmfold2_bundle.prepare_requests(tmp_path) + selected = esmfold2_bundle._all_prepared_requests(tmp_path) + assert selected == paths + + paths[0].unlink() + with pytest.raises(FileNotFoundError, match=r"missing=.*esmfold2"): + esmfold2_bundle._all_prepared_requests(tmp_path) + + +def test_esmfold2_semantic_config_ignores_only_packaging_and_runtime_policy() -> None: + fields = {name: 1 for name in esmfold2_bundle._semantic_config_fields} + reference = { + **fields, + "esmc_id": "biohub/ESMC-6B", + "max_length": 20, + } + candidate = { + **fields, + "esmc_id": "Synthyra/ESMplusplus_6B", + "attn_backend": "sdpa", + "esmc_precision": "bf16", + } + expected = esmfold2_bundle._esmfold2_semantic_config( + reference, + backbone_model="esmc_6b", + ) + assert ( + esmfold2_bundle._esmfold2_semantic_config( + candidate, + backbone_model="esmc_6b", + ) + == expected + ) + + changed = {**candidate, "d_pair": 2} + assert ( + esmfold2_bundle._esmfold2_semantic_config( + changed, + backbone_model="esmc_6b", + ) + != expected + ) + + +@pytest.mark.structure +@pytest.mark.compliance +@pytest.mark.slow +@pytest.mark.parametrize( + "spec", + [_spec_parameter(spec) for spec in get_model_registry().by_family("esmfold2")], +) +def test_esmfold2_isolated_bf16_folding_compliance( + spec: ModelSpec, + record_property: Callable[[str, object], None], +) -> None: + """Gate native BF16 parity for one compact stable request per variant.""" + + request_path, reference_path, bf16_path = _bundle_paths(spec) + request = load_request(request_path) + reference_tensors, reference_metadata = load_bundle(reference_path) + bf16_tensors, bf16_metadata = load_bundle(bf16_path) + + _assert_bundle_identity( + reference_metadata, + request, + spec, + producer="reference", + precision="bf16", + ) + _assert_bundle_identity( + bf16_metadata, + request, + spec, + producer="candidate", + precision="bf16", + ) + reference_environment = reference_metadata["environment"] + candidate_environment = bf16_metadata["environment"] + assert isinstance(reference_environment, Mapping) + assert isinstance(candidate_environment, Mapping) + assert_same_hopper_sm90_device(candidate_environment, reference_environment) + assert bf16_metadata["semantic_config"] == reference_metadata["semantic_config"] + assert bf16_metadata["state"] == reference_metadata["state"] + _assert_exact_inputs( + bf16_tensors, + bf16_metadata, + reference_tensors, + reference_metadata, + context=f"{spec.id} BF16 official parity", + ) + _assert_valid_geometry(reference_tensors, context=f"{spec.id} official BF16") + _assert_valid_geometry(bf16_tensors, context=f"{spec.id} FastPLMs BF16") + + bf16_metrics = _structure_metrics(bf16_tensors, reference_tensors) + _assert_thresholds( + bf16_metrics, + targets=bf16_targets, + hard_limits=bf16_hard_limits, + context=f"{spec.id} BF16 official parity", + ) + + record_property("fast_checkpoint_revision", spec.fast.revision) + record_property("official_checkpoint_revision", spec.official.revision) + record_property("feature_sha256", bf16_metadata["feature_sha256"]) + record_property( + "diffusion_noise_sha256", + bf16_metadata["diffusion_noise_sha256"], + ) + for name, value in bf16_metrics.items(): + record_property(f"bf16_{name}", value) diff --git a/tests/structure/test_esmfold2_fp8_compliance.py b/tests/structure/test_esmfold2_fp8_compliance.py new file mode 100644 index 0000000..c1b4e69 --- /dev/null +++ b/tests/structure/test_esmfold2_fp8_compliance.py @@ -0,0 +1,134 @@ +"""Experimental H100 smoke coverage for explicit ESMFold2 FP8 reloads.""" + +from __future__ import annotations + +import gc +import importlib +import pytest +import torch + +from fastplms.registry import ModelSpec, get_model_registry + + +SEQUENCES = ( + "MSTNPKPQRKTKRNTNR", + "ACDEFGHIK", +) +EXPECTED_FP8_PROJECTIONS = 80 + + +def _esmfold2_specs() -> tuple[ModelSpec, ...]: + return get_model_registry().by_family("esmfold2") + + +def _parameter(spec: ModelSpec) -> object: + return pytest.param(spec, id=spec.id, marks=pytest.mark.large) + + +def _base_spec() -> ModelSpec: + return get_model_registry()["esmfold2"] + + +def _load_current_model(spec: ModelSpec, device: torch.device) -> torch.nn.Module: + module_name, class_name = spec.auto_map["AutoModel"].rsplit(".", maxsplit=1) + model_class = getattr(importlib.import_module(module_name), class_name) + model = model_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + load_esmc=False, + dtype=torch.bfloat16, + ) + return model.eval().to(device=device) + + +def _assert_fp8_smoke(model: torch.nn.Module) -> None: + status = model.esmc_precision_status + assert status.requested == "fp8" + assert status.resolved == "fp8" + assert str(status.device).startswith("cuda") + assert status.transformer_engine_version + paths = model._esmc_fp8_module_paths + assert len(paths) == len(set(paths)) == EXPECTED_FP8_PROJECTIONS + assert all(path.endswith(".attn.out_proj") for path in paths) + + result = model.embed_dataset( + list(SEQUENCES), + batch_size=2, + full_embeddings=True, + dtype=torch.float32, + ) + embeddings = tuple(record.load_tensor() for record in result) + assert tuple(tensor.shape for tensor in embeddings) == ((17, 256), (9, 256)) + assert all(torch.isfinite(tensor).all() for tensor in embeddings) + assert not any(key.startswith("_esmc.") for key in model.state_dict()) + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.parametrize("spec", [_parameter(spec) for spec in _esmfold2_specs()]) +def test_explicit_fp8_smoke_on_each_esmfold2_variant(spec: ModelSpec) -> None: + """Exercise the experimental FP8 opt-in once on every supported variant.""" + + device = torch.device("cuda") + model = _load_current_model(spec, device) + try: + model.reload_esmc(precision="fp8", device=device) + _assert_fp8_smoke(model) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.large +def test_standard_esmfold2_rebuilds_fp8_from_bf16_three_times() -> None: + """Rebuild transient FP8 modules from canonical BF16 state on every reload.""" + + device = torch.device("cuda") + model = _load_current_model(_base_spec(), device) + try: + for _cycle in range(3): + model.reload_esmc(precision="bf16", device=device) + assert model.esmc_precision_status.resolved == "bf16" + assert model._esmc_fp8_module_paths == () + + model.reload_esmc(precision="fp8", device=device) + _assert_fp8_smoke(model) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.large +def test_auto_precision_selects_runtime_bf16_on_the_locked_h100() -> None: + """Exercise the stable BF16 auto policy on a directly loaded CUDA model.""" + + device = torch.device("cuda") + model = _load_current_model(_base_spec(), device) + try: + model.reload_esmc(precision="auto", device=device) + status = model.esmc_precision_status + assert status.requested == "auto" + assert status.resolved == "bf16" + assert str(status.device).startswith("cuda") + assert model._esmc_fp8_module_paths == () + result = model.embed_dataset( + list(SEQUENCES), + batch_size=2, + full_embeddings=True, + dtype=torch.float32, + ) + assert all(torch.isfinite(record.load_tensor()).all() for record in result) + assert not any(key.startswith("_esmc.") for key in model.state_dict()) + finally: + del model + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/structure/test_esmfold2_leaf_outputs.py b/tests/structure/test_esmfold2_leaf_outputs.py new file mode 100644 index 0000000..2f55d4c --- /dev/null +++ b/tests/structure/test_esmfold2_leaf_outputs.py @@ -0,0 +1,122 @@ +"""Deterministic output-contract tests for ESMFold2 leaf utilities.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from types import SimpleNamespace + +from fastplms.models.esmfold2.esmfold2_constants import ( + MOL_TYPE_NONPOLYMER, + MOL_TYPE_PROTEIN, +) +from fastplms.models.esmfold2.esmfold2_output import ( + build_molecular_complex_from_features, +) +from fastplms.models.esmfold2.modeling_esmfold2_common import LanguageModelShim + + +pytestmark = pytest.mark.structure + + +def _encoded_name(name: str) -> list[int]: + return [ord(character) - 32 if character != " " else 0 for character in name.ljust(4)] + + +def _token( + token_index: int, + residue_index: int, + residue_name: str, + atom_start: int, + atom_count: int, +) -> SimpleNamespace: + return SimpleNamespace( + token_index=token_index, + residue_index=residue_index, + residue_name=residue_name, + atom_start=atom_start, + atom_count=atom_count, + ) + + +def test_feature_output_groups_modified_residues_and_ligand_atoms() -> None: + polymer_tokens = [ + _token(0, 0, "ALA", 0, 2), + _token(1, 1, "GLY", 2, 1), + _token(2, 1, "GLY", 3, 1), + ] + ligand_tokens = [ + _token(3, 0, "LIG", 4, 1), + _token(4, 0, "LIG", 5, 1), + ] + chain_infos = [ + SimpleNamespace( + asym_id=0, + entity_id=0, + chain_id="A", + mol_type=MOL_TYPE_PROTEIN, + tokens=polymer_tokens, + ), + SimpleNamespace( + asym_id=1, + entity_id=1, + chain_id="B", + mol_type=MOL_TYPE_NONPOLYMER, + tokens=ligand_tokens, + ), + ] + # X: (6, 3) + X = torch.arange(18, dtype=torch.float32).reshape(6, 3) + # atom_names: (...) + atom_names = torch.tensor([_encoded_name(name) for name in ("N", "CA", "C", "O", "C1", "N1")]) + + complex_record = build_molecular_complex_from_features( + coords=X, + plddt=torch.tensor([0.1, 0.3, 0.5, 0.7, 0.9]), + atom_mask=torch.ones(6, dtype=torch.bool), + ref_element=torch.tensor([7, 6, 6, 8, 6, 7]), + ref_atom_name_chars=atom_names, + chain_infos=chain_infos, + complex_id="fixture", + ) + + assert complex_record.id == "fixture" + assert complex_record.sequence == ["ALA", "GLY", "LIG"] + np.testing.assert_array_equal( + complex_record.token_to_atoms, + np.array([[0, 2], [2, 4], [4, 6]], dtype=np.int32), + ) + np.testing.assert_array_equal(complex_record.chain_id, np.array([0, 0, 1])) + np.testing.assert_allclose(complex_record.plddt, np.array([0.1, 0.4, 0.8])) + np.testing.assert_array_equal( + complex_record.atom_hetero, + np.array([False, False, False, False, True, True]), + ) + np.testing.assert_array_equal(complex_record.atom_positions, X.numpy()) + assert complex_record.atom_names.tolist() == ["N", "CA", "C", "O", "C1", "N1"] + assert complex_record.metadata.chain_lookup == {0: "A", 1: "B"} + assert complex_record.metadata.entity_lookup == {0: "polymer", 1: "non-polymer"} + + +@pytest.mark.gpu +def test_learned_sequence_projection_matches_explicit_cuda_operation() -> None: + assert torch.cuda.is_available() + torch.manual_seed(7) + device = torch.device("cuda") + shim = LanguageModelShim(d_z=7, d_model=11, num_layers=3).to( + device=device, dtype=torch.bfloat16 + ) + # H: (2, 17, 4, 11) + H = torch.randn((2, 17, 4, 11), device=device, dtype=torch.bfloat16) + # M: (2, 17) + M = torch.tensor([[True] * 13 + [False] * 4, [True] * 17], device=device, dtype=torch.bool) + + projected_states = shim.base_z_linear(H) + expected = shim.base_z_combine.softmax(dim=0) @ projected_states + expected = expected * M.unsqueeze(-1) + Z = shim.project_sequence(H, M) + + assert torch.equal(Z, expected) + assert Z.shape == (2, 17, 7) + assert Z.dtype == torch.bfloat16 diff --git a/tests/structure/test_esmfold_folding_compliance.py b/tests/structure/test_esmfold_folding_compliance.py new file mode 100644 index 0000000..2d4a51a --- /dev/null +++ b/tests/structure/test_esmfold_folding_compliance.py @@ -0,0 +1,392 @@ +"""Release gates over isolated Meta and FastPLMs ESMFold v1 bundles.""" + +from __future__ import annotations + +import ast +import inspect +import os +import pytest +import torch +from collections.abc import Mapping +from pathlib import Path + +from fastplms.registry import get_model_registry +from tests.structure.support import esmfold_bundle +from tests.structure.support.esmfold_bundle import load_bundle, load_request +from tests.structure.support.hardware import ( + assert_same_hopper_sm90_device, + hopper_sm90_fingerprint, +) + + +relative_l2_targets = {"fp32": 2e-6, "bf16": 1e-2} +relative_l2_hard_limits = {"fp32": 2e-5, "bf16": 3e-2} +structure_targets = { + "ca_rmsd": 0.10, + "lddt_ca": 0.995, + "plddt_mae": 0.0012, + "pae_mae": 0.10, + "ptm_error": 0.002, +} +structure_hard_limits = { + "ca_rmsd": 0.25, + "lddt_ca": 0.99, + "plddt_mae": 0.005, + "pae_mae": 0.50, + "ptm_error": 0.005, +} + + +def _exchange_root() -> Path: + return Path(os.environ.get("FASTPLMS_REFERENCE_EXCHANGE", "artifacts/reference")) + + +def _paths(precision: str) -> tuple[Path, Path, Path]: + root = _exchange_root() + request = ( + root + / "structure" + / "requests" + / esmfold_bundle.reference_container + / f"{esmfold_bundle.model_id}.json" + ) + results = root / "structure" / "results" + reference = results / "reference" / esmfold_bundle.model_id / precision + candidate = results / "candidate" / esmfold_bundle.model_id / precision + return request, reference, candidate + + +def _checkpoint_contract(checkpoint: object) -> dict[str, object]: + return { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": [ + { + "path": item.path, + "algorithm": item.algorithm, + "digest": item.digest, + } + for item in checkpoint.files + ], + } + + +def _upstream_contract(upstream: object) -> dict[str, object]: + return { + "id": upstream.id, + "path": upstream.path, + "url": upstream.url, + "revision": upstream.revision, + "license_expression": upstream.license_expression, + } + + +def _output(tensors: Mapping[str, torch.Tensor], name: str) -> torch.Tensor: + key = f"output__{name}" + if key not in tensors: + raise KeyError(f"ESMFold bundle omits required output {name!r}.") + return tensors[key] + + +def _residue_mask(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + # atom37_mask: (...) + atom37_mask = _output(tensors, "atom37_atom_exists").bool() + assert atom37_mask.ndim == 3 and atom37_mask.shape[-1] == 37 + return atom37_mask[0, :, 1] + + +def _ca_coordinates(tensors: Mapping[str, torch.Tensor]) -> torch.Tensor: + # P is the atom14 position tensor with shape (n_blocks, b, l, 14, 3). + P = _output(tensors, "positions").float() + assert P.ndim == 5 and P.shape[-2:] == (14, 3) + # coordinates: (..., 3) + coordinates = P[-1, 0, :, 1] + return coordinates[_residue_mask(tensors)] + + +def _aligned_ca_rmsd(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + actual_centered = actual.float() - actual.float().mean(dim=0, keepdim=True) + expected_centered = expected.float() - expected.float().mean(dim=0, keepdim=True) + covariance = actual_centered.T @ expected_centered + left, _, right = torch.linalg.svd(covariance) + # correction: (3, 3) + correction = torch.eye(3, dtype=torch.float32) + correction[-1, -1] = torch.sign(torch.det(left @ right)) + rotation = left @ correction @ right + aligned = actual_centered @ rotation + return torch.sqrt(torch.mean(torch.sum((aligned - expected_centered) ** 2, dim=-1))).item() + + +def _lddt_ca(actual: torch.Tensor, expected: torch.Tensor) -> float: + # actual: (...), expected: (...) + actual_distances = torch.cdist(actual.float(), actual.float()) + expected_distances = torch.cdist(expected.float(), expected.float()) + # pair_mask: (...) + pair_mask = expected_distances.lt(15.0) + pair_mask.fill_diagonal_(False) + assert pair_mask.any(), "No valid C-alpha pairs for ESMFold lDDT." + errors = (actual_distances - expected_distances).abs() + # scores: (...) + scores = torch.stack([errors.lt(threshold).float() for threshold in (0.5, 1.0, 2.0, 4.0)]).mean( + dim=0 + ) + return scores[pair_mask].mean().item() + + +def _structure_metrics( + actual: Mapping[str, torch.Tensor], + expected: Mapping[str, torch.Tensor], +) -> dict[str, float]: + residue_mask = _residue_mask(actual) + # pair_mask: (...) + pair_mask = residue_mask[:, None] & residue_mask[None, :] + # Meta ESMFold reports pLDDT on (0, 100); compliance uses (0, 1). + # actual_plddt: (...) + actual_plddt = _output(actual, "plddt").float()[0, :, 1] / 100.0 + # expected_plddt: (...) + expected_plddt = _output(expected, "plddt").float()[0, :, 1] / 100.0 + # actual_pae: (...) + actual_pae = _output(actual, "predicted_aligned_error").float()[0] + # expected_pae: (...) + expected_pae = _output(expected, "predicted_aligned_error").float()[0] + return { + "ca_rmsd": _aligned_ca_rmsd( + _ca_coordinates(actual), + _ca_coordinates(expected), + ), + "lddt_ca": _lddt_ca( + _ca_coordinates(actual), + _ca_coordinates(expected), + ), + "plddt_mae": (actual_plddt[residue_mask] - expected_plddt[residue_mask]) + .abs() + .mean() + .item(), + "pae_mae": (actual_pae[pair_mask] - expected_pae[pair_mask]).abs().mean().item(), + "ptm_error": ( + _output(actual, "ptm").float().reshape(-1)[0] + - _output(expected, "ptm").float().reshape(-1)[0] + ) + .abs() + .item(), + } + + +def _relative_l2( + actual: torch.Tensor, + expected: torch.Tensor, + mask: torch.Tensor, +) -> float: + # actual: (...), expected: (...), mask: (...) + while mask.ndim < actual.ndim: + # mask: (...) + mask = mask.unsqueeze(-1) + mask = torch.broadcast_to(mask, actual.shape) + difference = (actual.float() - expected.float())[mask] + reference = expected.float()[mask] + return ( + torch.linalg.vector_norm(difference) + / torch.linalg.vector_norm(reference).clamp_min(torch.finfo(torch.float32).tiny) + ).item() + + +def _logit_metrics( + actual: Mapping[str, torch.Tensor], + expected: Mapping[str, torch.Tensor], +) -> dict[str, float]: + residue_mask = _residue_mask(actual) + # pair_mask: (...) + pair_mask = residue_mask[:, None] & residue_mask[None, :] + return { + "distogram_logits": _relative_l2( + _output(actual, "distogram_logits"), + _output(expected, "distogram_logits"), + pair_mask.unsqueeze(0), + ), + "ptm_logits": _relative_l2( + _output(actual, "ptm_logits"), + _output(expected, "ptm_logits"), + pair_mask.unsqueeze(0), + ), + "lm_logits": _relative_l2( + _output(actual, "lm_logits"), + _output(expected, "lm_logits"), + residue_mask.unsqueeze(0), + ), + } + + +def _assert_valid_bundle( + tensors: Mapping[str, torch.Tensor], + *, + context: str, +) -> None: + residue_mask = _residue_mask(tensors) + assert residue_mask.sum().item() == len(esmfold_bundle.fold_sequence) + for name, tensor in tensors.items(): + if tensor.is_floating_point(): + assert torch.isfinite(tensor).all(), f"{context}: {name} contains NaN or inf" + ca_coordinates = _ca_coordinates(tensors) + ca_steps = torch.linalg.vector_norm(ca_coordinates[1:] - ca_coordinates[:-1], dim=-1) + assert ca_steps.gt(2.0).all() and ca_steps.lt(5.0).all(), ( + f"{context}: invalid consecutive C-alpha distances {ca_steps.tolist()}" + ) + + +def _assert_bundle_identity( + metadata: Mapping[str, object], + request: Mapping[str, object], + *, + producer: str, + precision: str, +) -> None: + registry = get_model_registry() + spec = registry[esmfold_bundle.model_id] + assert metadata["producer"] == producer + assert metadata["model_id"] == spec.id + assert metadata["request_sha256"] == request["request_sha256"] + assert metadata["official"] == _checkpoint_contract(spec.official) + assert metadata["candidate"] == _checkpoint_contract(spec.fast) + assert metadata["upstreams"] == [ + _upstream_contract(registry.upstreams[name]) for name in ("fair-esm", "openfold") + ] + assert metadata["sequence"] == esmfold_bundle.fold_sequence + assert metadata["seed"] == esmfold_bundle.fold_seed + assert metadata["recycles"] == esmfold_bundle.fold_recycles + assert metadata["attention_backend"] == esmfold_bundle.fold_backend + assert metadata["deterministic_algorithms"] is True + assert metadata["parameter_dtype"] == "float32" + assert metadata["compute_dtype"] == precision + assert metadata["esm_parameter_dtypes"] == ["float32"] + expected_execution = ( + "fp32_parameters_cuda_bf16_autocast" if precision == "bf16" else "fp32_parameters" + ) + assert metadata["execution"] == expected_execution + environment = metadata["environment"] + assert isinstance(environment, Mapping) + hopper_sm90_fingerprint(environment) + if producer == "candidate": + assert str(environment["torch"]).split("+", maxsplit=1)[0] == "2.13.0" + assert str(environment["cuda_runtime"]).startswith("13.0") + packages = environment["packages"] + assert isinstance(packages, Mapping) + assert packages["transformers"] == "5.13.0" + + +def test_esmfold_reference_path_has_no_fastplms_dependency() -> None: + """Keep the Meta oracle copyable into a native image without FastPLMs.""" + + tree = ast.parse(inspect.getsource(esmfold_bundle)) + for node in tree.body: + if isinstance(node, ast.Import): + assert all(not alias.name.startswith("fastplms") for alias in node.names) + elif isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith("fastplms") + for function in ( + esmfold_bundle._load_reference_model, + esmfold_bundle._run_infer, + esmfold_bundle.produce_reference, + ): + assert "fastplms" not in inspect.getsource(function).lower() + + +def test_prepare_esmfold_request_is_manifest_exact(tmp_path: Path) -> None: + path = esmfold_bundle.prepare_request(tmp_path) + request = load_request(path) + registry = get_model_registry() + spec = registry[esmfold_bundle.model_id] + assert request["official"] == _checkpoint_contract(spec.official) + assert request["candidate"] == _checkpoint_contract(spec.fast) + assert request["candidate_auto_model"] == spec.auto_map["AutoModel"] + assert request["adapter"] == spec.family.reference_adapter + assert request["deterministic_algorithms"] is True + + +def test_esmfold_metric_helpers_are_exact_for_rigid_identity() -> None: + # expected: (4, 3) + expected = torch.tensor([[0.0, 0.0, 0.0], [3.8, 0.0, 0.0], [7.2, 1.0, 0.0], [9.0, 4.0, 1.0]]) + # rotation: (3, 3) + rotation = torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + actual = expected @ rotation + torch.tensor([4.0, -2.0, 7.0]) + assert _aligned_ca_rmsd(actual, expected) == pytest.approx(0.0, abs=1e-5) + assert _lddt_ca(actual, expected) == pytest.approx(1.0) + + +def test_esmfold_tensor_hash_accepts_scalar_outputs() -> None: + assert esmfold_bundle.tensor_sha256(torch.tensor(0.5)) == esmfold_bundle.tensor_sha256( + torch.tensor([0.5]) + ) + + +@pytest.mark.parametrize("precision", esmfold_bundle.supported_precisions) +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.large +def test_esmfold_live_official_structure_parity(precision: str) -> None: + """The pinned candidate matches Meta ESMFold v1 under FP32 and BF16 compute.""" + + request_path, reference_path, candidate_path = _paths(precision) + request = load_request(request_path) + reference_tensors, reference_metadata = load_bundle(reference_path) + candidate_tensors, candidate_metadata = load_bundle(candidate_path) + _assert_bundle_identity( + reference_metadata, + request, + producer="reference", + precision=precision, + ) + _assert_bundle_identity( + candidate_metadata, + request, + producer="candidate", + precision=precision, + ) + reference_environment = reference_metadata["environment"] + candidate_environment = candidate_metadata["environment"] + assert isinstance(reference_environment, Mapping) + assert isinstance(candidate_environment, Mapping) + assert_same_hopper_sm90_device(candidate_environment, reference_environment) + assert candidate_metadata["semantic_config"] == reference_metadata["semantic_config"] + assert candidate_metadata["state"] == reference_metadata["state"] + assert reference_tensors.keys() == candidate_tensors.keys() + for name in esmfold_bundle._exact_outputs: + actual = _output(candidate_tensors, name) + expected = _output(reference_tensors, name) + assert actual.dtype == expected.dtype, name + assert actual.shape == expected.shape, name + assert torch.equal(actual, expected), name + _assert_valid_bundle(reference_tensors, context="Meta ESMFold v1") + _assert_valid_bundle(candidate_tensors, context="FastPLMs ESMFold v1") + + logits = _logit_metrics(candidate_tensors, reference_tensors) + relative_l2_target = relative_l2_targets[precision] + relative_l2_hard_limit = relative_l2_hard_limits[precision] + for name, value in logits.items(): + assert value <= relative_l2_hard_limit, ( + f"ESMFold {name} relative L2 {value:.6g} exceeds hard limit " + f"{relative_l2_hard_limit:.6g} under {precision} compute." + ) + assert value <= relative_l2_target, ( + f"ESMFold {name} relative L2 {value:.6g} misses engineering target " + f"{relative_l2_target:.6g} under {precision} compute." + ) + + metrics = _structure_metrics(candidate_tensors, reference_tensors) + for name, hard_limit in structure_hard_limits.items(): + value = metrics[name] + if name == "lddt_ca": + assert value >= hard_limit, ( + f"ESMFold {name} {value:.6g} is below hard limit {hard_limit:.6g}." + ) + else: + assert value <= hard_limit, ( + f"ESMFold {name} {value:.6g} exceeds hard limit {hard_limit:.6g}." + ) + for name, target in structure_targets.items(): + value = metrics[name] + if name == "lddt_ca": + assert value >= target, f"ESMFold {name} {value:.6g} misses target {target:.6g}." + else: + assert value <= target, f"ESMFold {name} {value:.6g} misses target {target:.6g}." diff --git a/testing/test_structure_models.py b/tests/structure/test_structure_models.py similarity index 55% rename from testing/test_structure_models.py rename to tests/structure/test_structure_models.py index 9bf5a94..d915e71 100644 --- a/testing/test_structure_models.py +++ b/tests/structure/test_structure_models.py @@ -6,7 +6,14 @@ import pytest import torch -from transformers import AutoModel +from transformers import AutoConfig, AutoModel + +from fastplms.registry import get_model_registry + + +REGISTRY = get_model_registry() +BOLTZ2 = REGISTRY["boltz2"] +ESMFOLD = REGISTRY["esmfold"] @pytest.mark.structure @@ -14,11 +21,16 @@ @pytest.mark.slow def test_boltz2_loads() -> None: """Boltz2 loads via AutoModel with trust_remote_code=True.""" - model = AutoModel.from_pretrained( - "Synthyra/Boltz2", - trust_remote_code=True, - dtype=torch.float32, - ).eval().cuda() + model = ( + AutoModel.from_pretrained( + BOLTZ2.fast.repo_id, + revision=BOLTZ2.fast.revision, + trust_remote_code=True, + dtype=torch.float32, + ) + .eval() + .cuda() + ) assert model is not None assert hasattr(model, "predict_structure") @@ -33,11 +45,16 @@ def test_boltz2_loads() -> None: @pytest.mark.slow def test_boltz2_forward() -> None: """Boltz2 predict_structure returns valid coordinates and confidence scores.""" - model = AutoModel.from_pretrained( - "Synthyra/Boltz2", - trust_remote_code=True, - dtype=torch.float32, - ).eval().cuda() + model = ( + AutoModel.from_pretrained( + BOLTZ2.fast.repo_id, + revision=BOLTZ2.fast.revision, + trust_remote_code=True, + dtype=torch.float32, + ) + .eval() + .cuda() + ) sequence = "MSTNPKPQRKTKRNTNRRPQDVKFPGG" output = model.predict_structure( @@ -48,7 +65,9 @@ def test_boltz2_forward() -> None: ) assert output.sample_atom_coords is not None - assert output.sample_atom_coords.ndim in (2, 3), f"Expected 2D or 3D coords, got {output.sample_atom_coords.ndim}D" + assert output.sample_atom_coords.ndim in (2, 3), ( + f"Expected 2D or 3D coords, got {output.sample_atom_coords.ndim}D" + ) assert output.sample_atom_coords.shape[-1] == 3 assert not torch.isnan(output.sample_atom_coords).any(), "NaN in predicted coordinates" assert output.plddt is not None @@ -63,15 +82,22 @@ def test_boltz2_forward() -> None: @pytest.mark.slow def test_esmfold_loads() -> None: """ESMFold loads via AutoModel with trust_remote_code=True.""" - model = AutoModel.from_pretrained( - "Synthyra/FastESMFold", - trust_remote_code=True, - dtype=torch.float32, - ).eval().cuda() + model = ( + AutoModel.from_pretrained( + ESMFOLD.fast.repo_id, + revision=ESMFOLD.fast.revision, + trust_remote_code=True, + dtype=torch.float32, + ) + .eval() + .cuda() + ) assert model is not None assert hasattr(model, "infer") assert hasattr(model, "fold_protein") + assert not any(name.startswith("mlm_head.") for name in model.state_dict()) + assert not any(name.startswith("esm.contact_head.") for name in model.state_dict()) del model torch.cuda.empty_cache() @@ -82,17 +108,22 @@ def test_esmfold_loads() -> None: @pytest.mark.slow def test_esmfold_forward() -> None: """ESMFold infer produces valid pLDDT and structure output.""" - from transformers import AutoConfig - - config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True) - config.ttt_config = {"steps": 0} - - model = AutoModel.from_pretrained( - "Synthyra/FastESMFold", - config=config, + config = AutoConfig.from_pretrained( + ESMFOLD.fast.repo_id, + revision=ESMFOLD.fast.revision, trust_remote_code=True, - dtype=torch.float32, - ).eval().cuda() + ) + model = ( + AutoModel.from_pretrained( + ESMFOLD.fast.repo_id, + revision=ESMFOLD.fast.revision, + config=config, + trust_remote_code=True, + dtype=torch.float32, + ) + .eval() + .cuda() + ) sequence = "MKTLLILAVVAAALA" @@ -100,6 +131,7 @@ def test_esmfold_forward() -> None: output = model.infer(sequence) assert "plddt" in output + # plddt: (...) plddt = output["plddt"] assert not torch.isnan(plddt).any(), "NaN in pLDDT" diff --git a/tests/structure/test_structure_official_goldens.py b/tests/structure/test_structure_official_goldens.py new file mode 100644 index 0000000..cef4ee5 --- /dev/null +++ b/tests/structure/test_structure_official_goldens.py @@ -0,0 +1,154 @@ +"""Candidate-only, hardware-bound regression against immutable structure goldens. + +Unlike compliance tests, this module never imports or executes an official +implementation. It consumes the compact, hash-verified outputs already checked +into ``tests/goldens`` and is intended for the conditional structure GPU lane. +""" + +from __future__ import annotations + +import gc +import json +import pytest +import torch +from pathlib import Path +from safetensors.torch import load_file + +from fastplms.registry import ModelSpec, get_model_registry +from tests.structure import test_esmfold2_folding_compliance as esmfold2_metrics +from tests.structure import test_esmfold_folding_compliance as esmfold_metrics +from tests.structure.support import esmfold2_bundle, esmfold_bundle +from tests.structure.support.hardware import assert_recorded_hopper_device_matches +from tools.goldens import validate_golden_bundle + + +ROOT = Path(__file__).resolve().parents[2] +REGISTRY = get_model_registry() + + +def _golden(spec: ModelSpec) -> tuple[dict[str, torch.Tensor], dict[str, object]]: + declaration = spec.official_golden + if declaration is None: + raise AssertionError( + f"{spec.id}: no measured immutable structure golden is declared; " + "candidate-only structure validation must fail closed." + ) + metadata_path = ROOT / declaration.metadata.path + tensors_path = ROOT / declaration.tensors.path + validate_golden_bundle( + spec, + REGISTRY, + metadata_path=metadata_path, + tensors_path=tensors_path, + declaration=declaration, + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + return load_file(tensors_path, device="cpu"), metadata + + +def _release_parameter(spec: ModelSpec) -> object: + return pytest.param(spec, id=spec.id, marks=pytest.mark.large) + + +def _assert_golden_device_matches_current(metadata: dict[str, object]) -> None: + environment = metadata["environment"] + assert isinstance(environment, dict) + recorded = environment["details"] + assert isinstance(recorded, dict) + properties = torch.cuda.get_device_properties(0) + assert_recorded_hopper_device_matches( + { + "cuda_device": properties.name, + "cuda_device_capability": list(torch.cuda.get_device_capability(0)), + "cuda_total_memory": int(properties.total_memory), + }, + recorded, + ) + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.checkpoint +@pytest.mark.network +def test_esmfold_candidate_matches_checked_structure_golden(tmp_path: Path) -> None: + spec = REGISTRY[esmfold_bundle.model_id] + golden, metadata = _golden(spec) + _assert_golden_device_matches_current(metadata) + request_path = esmfold_bundle.prepare_request(tmp_path) + request = esmfold_bundle.load_request(request_path) + assert metadata["input_fingerprint"] == request["request_sha256"] + + model = esmfold_bundle._load_candidate_model(request, torch.device("cuda")) + candidate = esmfold_bundle._run_infer(model, request, "bf16") + + for name in esmfold_bundle._exact_outputs: + assert torch.equal(candidate[f"output__{name}"], golden[f"output__{name}"]), name + esmfold_metrics._assert_valid_bundle(golden, context="checked Meta ESMFold golden") + esmfold_metrics._assert_valid_bundle(candidate, context="FastPLMs ESMFold candidate") + for name, value in esmfold_metrics._logit_metrics(candidate, golden).items(): + assert value <= esmfold_metrics.relative_l2_hard_limits["bf16"], ( + f"{name} relative L2 {value:.6g} exceeds the measured BF16 hard limit" + ) + structure = esmfold_metrics._structure_metrics(candidate, golden) + for name, hard_limit in esmfold_metrics.structure_hard_limits.items(): + if name == "lddt_ca": + assert structure[name] >= hard_limit + else: + assert structure[name] <= hard_limit + + del model, candidate, golden + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.structure +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.checkpoint +@pytest.mark.network +@pytest.mark.parametrize( + "spec", + [_release_parameter(REGISTRY[model_id]) for model_id in esmfold2_bundle.supported_model_ids], +) +def test_esmfold2_candidate_matches_checked_structure_golden( + spec: ModelSpec, + tmp_path: Path, +) -> None: + golden, metadata = _golden(spec) + _assert_golden_device_matches_current(metadata) + request_path = esmfold2_bundle.prepare_requests(tmp_path, model_ids=(spec.id,))[0] + request = esmfold2_bundle.load_request(request_path) + assert metadata["input_fingerprint"] == request["request_sha256"] + + model = esmfold2_bundle._load_candidate_model(request, torch.device("cuda"), "bf16") + candidate = esmfold2_bundle._run_fold(model, request) + + immutable_inputs = [ + name + for name in golden + if name.startswith("feature__") or name == "noise__initial_standard_normal" + ] + assert immutable_inputs + for name in immutable_inputs: + assert name in candidate, f"{spec.id}: candidate omitted golden input {name}" + assert torch.equal(candidate[name], golden[name]), f"{spec.id}: {name}" + esmfold2_metrics._assert_valid_geometry( + golden, + context=f"{spec.id} checked official golden", + ) + esmfold2_metrics._assert_valid_geometry( + candidate, + context=f"{spec.id} FastPLMs candidate", + ) + structure = esmfold2_metrics._structure_metrics(candidate, golden) + esmfold2_metrics._assert_thresholds( + structure, + targets=esmfold2_metrics.bf16_targets, + hard_limits=esmfold2_metrics.bf16_hard_limits, + context=f"{spec.id} checked BF16 golden", + ) + + del model, candidate, golden + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/structure/test_structure_public_helpers.py b/tests/structure/test_structure_public_helpers.py new file mode 100644 index 0000000..8ab541c --- /dev/null +++ b/tests/structure/test_structure_public_helpers.py @@ -0,0 +1,592 @@ +"""CPU contracts for public structure convenience helpers.""" + +from __future__ import annotations + +import random +import numpy as np +import pytest +import torch +from collections.abc import Iterator +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any + +from examples.binder_design_fastplms import compute_structure_losses +from fastplms.models.boltz import modeling_boltz2, vb_modules_confidencev2 +from fastplms.models.boltz.minimal_featurizer import build_boltz2_features +from fastplms.models.boltz.vb_loss_diffusionv2 import smooth_lddt_loss +from fastplms.models.boltz.vb_modules_encodersv2 import get_indexing_matrix +from fastplms.models.boltz.vb_potentials_potentials import FlatBottomPotential +from fastplms.models.boltz.vb_potentials_schedules import PiecewiseStepFunction +from fastplms.models.esmfold.modeling_fast_esmfold import FastEsmForProteinFolding +from fastplms.models.esmfold2.esmfold2_affine3d import Affine3D +from fastplms.models.esmfold2.esmfold2_predicted_aligned_error import tm_loss +from fastplms.models.esmfold2.protein_utils import prepare_protein_features +from fastplms.models.esmfold2.reproducibility import seed_context + + +pytestmark = pytest.mark.structure + + +class _FakeBoltz: + def __init__(self, *, device: str = "cpu", dtype: torch.dtype = torch.float32) -> None: + self.device = torch.device(device) + self.config = SimpleNamespace(num_bins=64) + self.core = SimpleNamespace( + input_embedder=SimpleNamespace( + atom_encoder=SimpleNamespace(atoms_per_window_queries=32) + ) + ) + self.parameter = torch.nn.Parameter(torch.ones(1, dtype=dtype)) + + def parameters(self) -> Iterator[torch.nn.Parameter]: + yield self.parameter + + def _to_model_device( + self, + feats: dict[str, torch.Tensor], + float_dtype: torch.dtype, + ) -> dict[str, torch.Tensor]: + assert float_dtype == torch.float32 + return feats + + def forward(self, **kwargs: Any) -> dict[str, torch.Tensor]: + del kwargs + return { + "sample_atom_coords": torch.randn((1, 1, 3)), + "plddt": torch.rand((1, 1)), + "complex_plddt": torch.rand((1,)), + "iptm": torch.rand((1,)), + "ptm": torch.rand((1,)), + } + + +def _fake_boltz_features( + amino_acid_sequence: str, + num_bins: int, + atoms_per_window_queries: int, +) -> tuple[dict[str, torch.Tensor], SimpleNamespace]: + assert num_bins == 64 + assert atoms_per_window_queries == 32 + # Exercise every ambient stream that the public helper promises to scope. + random.random() + np.random.random() + return { + "atom_pad_mask": torch.ones((1, 1)), + "ref_pos": torch.randn((1, 1, 3)), + }, SimpleNamespace(sequence=amino_acid_sequence) + + +def test_boltz_public_helper_is_seeded_and_restores_ambient_rng( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(modeling_boltz2, "build_boltz2_features", _fake_boltz_features) + model = _FakeBoltz() + + random.seed(91) + np.random.seed(91) + torch.manual_seed(91) + expected_next = (random.random(), float(np.random.random()), torch.rand(1)) + random.seed(91) + np.random.seed(91) + torch.manual_seed(91) + + first = modeling_boltz2.Boltz2Model.predict_structure(model, "ACD", seed=17) + observed_next = (random.random(), float(np.random.random()), torch.rand(1)) + second = modeling_boltz2.Boltz2Model.predict_structure(model, "ACD", seed=17) + + assert observed_next[0] == expected_next[0] + assert observed_next[1] == expected_next[1] + torch.testing.assert_close(observed_next[2], expected_next[2], rtol=0.0, atol=0.0) + torch.testing.assert_close( + first.sample_atom_coords, + second.sample_atom_coords, + rtol=0.0, + atol=0.0, + ) + assert first.seed == second.seed == 17 + + +@pytest.mark.parametrize( + "invalid_seed", + (True, False, 1.0, "17", b"17", np.int64(17)), +) +def test_boltz_public_helper_rejects_coerced_seed_types_before_rng_mutation( + monkeypatch: pytest.MonkeyPatch, + invalid_seed: object, +) -> None: + monkeypatch.setattr( + modeling_boltz2, + "build_boltz2_features", + lambda *_args, **_kwargs: pytest.fail("feature preparation was reached"), + ) + model = _FakeBoltz() + + random.seed(91) + np.random.seed(91) + torch.manual_seed(91) + expected_next = (random.random(), float(np.random.random()), torch.rand(1)) + random.seed(91) + np.random.seed(91) + torch.manual_seed(91) + + with pytest.raises(TypeError, match="seed must be an int or None"): + modeling_boltz2.Boltz2Model.predict_structure( + model, + "ACD", + seed=invalid_seed, # type: ignore[arg-type] + ) + + observed_next = (random.random(), float(np.random.random()), torch.rand(1)) + assert observed_next[0] == expected_next[0] + assert observed_next[1] == expected_next[1] + torch.testing.assert_close(observed_next[2], expected_next[2], rtol=0.0, atol=0.0) + + +def test_boltz_public_helper_owns_cuda_bf16_autocast_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(modeling_boltz2, "build_boltz2_features", _fake_boltz_features) + observed: list[tuple[str, torch.dtype]] = [] + + @contextmanager + def fake_autocast(*, device_type: str, dtype: torch.dtype) -> Iterator[None]: + observed.append((device_type, dtype)) + yield + + monkeypatch.setattr(torch, "autocast", fake_autocast) + modeling_boltz2.Boltz2Model.predict_structure(_FakeBoltz(device="cuda"), "ACD", seed=3) + + assert observed == [("cuda", torch.bfloat16)] + with pytest.raises(ValueError, match="requires FP32 parameter storage"): + modeling_boltz2.Boltz2Model.predict_structure( + _FakeBoltz(dtype=torch.bfloat16), + "ACD", + seed=3, + ) + + +def test_boltz_public_helper_rejects_non_finite_coordinates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(modeling_boltz2, "build_boltz2_features", _fake_boltz_features) + model = _FakeBoltz() + model.forward = lambda **_kwargs: { # type: ignore[method-assign] + "sample_atom_coords": torch.tensor([[[float("nan"), 0.0, 0.0]]]) + } + + with pytest.raises(RuntimeError, match="sample_atom_coords contains non-finite"): + modeling_boltz2.Boltz2Model.predict_structure(model, "ACD", seed=3) + + +@pytest.mark.parametrize( + ("k", "w", "h", "error_type", "message"), + ( + (True, 4, 8, TypeError, "k must be an int"), + (0, 4, 8, ValueError, "k must be positive"), + (1, 3, 6, ValueError, "w must be even"), + (1, 4, 7, ValueError, "h must be divisible"), + (1, 4, 6, ValueError, "even number of half-window key blocks"), + ), +) +def test_boltz_indexing_matrix_rejects_invalid_public_dimensions( + k: object, + w: int, + h: int, + error_type: type[Exception], + message: str, +) -> None: + with pytest.raises(error_type, match=message): + get_indexing_matrix( # type: ignore[arg-type] + k, + w, + h, + torch.device("cpu"), + ) + + +def test_boltz_piecewise_schedule_validates_and_owns_its_configuration() -> None: + with pytest.raises(ValueError, match="at least one threshold"): + PiecewiseStepFunction((), (1.0,)) + with pytest.raises(ValueError, match="exactly one more value"): + PiecewiseStepFunction((0.5,), (1.0,)) + with pytest.raises(ValueError, match="strictly increasing"): + PiecewiseStepFunction((0.5, 0.5), (1.0, 2.0, 3.0)) + + thresholds = [0.5] + values = [1.0, 2.0] + schedule = PiecewiseStepFunction(thresholds, values) + thresholds.clear() + values.clear() + assert schedule.compute(0.5) == 1.0 + assert schedule.compute(0.5001) == 2.0 + + +@pytest.mark.parametrize( + ("negation_mask", "error_type", "message"), + ( + (torch.tensor([0]), TypeError, "must be a boolean tensor"), + (torch.tensor([False, False]), ValueError, "broadcastable to value shape"), + (torch.tensor([False]), ValueError, "at least one bound is infinite"), + ), +) +def test_boltz_flat_bottom_potential_rejects_invalid_negation_masks( + negation_mask: torch.Tensor, + error_type: type[Exception], + message: str, +) -> None: + # negation_mask: (...) + with pytest.raises(error_type, match=message): + FlatBottomPotential.compute_function( + object(), + value=torch.tensor([0.5]), + k=torch.tensor(1.0), + lower_bounds=torch.tensor([0.0]), + upper_bounds=torch.tensor([1.0]), + negation_mask=negation_mask, + ) + + +def _assert_boltz_atom_confidence_mapping(device: torch.device) -> None: + batch_size = 2 + multiplicity = 2 + token_count = 2 + slots_per_token = 3 + # token_logits: (batch_size * multiplicity, token_count, slots_per_token, 1) + token_logits = torch.empty( + batch_size * multiplicity, + token_count, + slots_per_token, + 1, + device=device, + ) + for batch_sample in range(batch_size * multiplicity): + for token in range(token_count): + for slot in range(slots_per_token): + # A fully indexed token logit is scalar. + token_logits[batch_sample, token, slot, 0] = 100 * batch_sample + 10 * token + slot + # atom_to_token: (2, 4, 2) + atom_to_token = torch.tensor( + ( + ((0, 1), (1, 0), (0, 1), (0, 0)), + ((1, 0), (1, 0), (0, 1), (0, 0)), + ), + dtype=torch.bool, + device=device, + ) + # atom_pad_mask: (2, 4) + atom_pad_mask = torch.tensor( + ((1, 1, 1, 0), (1, 1, 1, 0)), + dtype=torch.bool, + device=device, + ) + + atom_logits = vb_modules_confidencev2._token_slot_logits_to_atom_logits( + token_logits, + atom_to_token, + atom_pad_mask, + multiplicity=multiplicity, + ) + + # expected: (4, 4) + expected = torch.tensor( + ( + (10, 0, 11, 0), + (110, 100, 111, 0), + (200, 201, 210, 0), + (300, 301, 310, 0), + ), + dtype=token_logits.dtype, + device=device, + ) + torch.testing.assert_close(atom_logits.squeeze(-1), expected, rtol=0.0, atol=0.0) + + +def test_boltz_atom_confidence_mapping_preserves_batch_multiplicity_and_atom_order() -> None: + _assert_boltz_atom_confidence_mapping(torch.device("cpu")) + + +@pytest.mark.gpu +def test_boltz_atom_confidence_mapping_preserves_batch_multiplicity_on_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the GH200 atom-confidence contract.") + _assert_boltz_atom_confidence_mapping(torch.device("cuda")) + + +def test_boltz_atom_confidence_is_finite_for_short_uneven_batches_and_multiplicity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + batch_size = 2 + multiplicity = 2 + token_count = 2 + atom_count = 4 + hidden_size = 4 + head = vb_modules_confidencev2.ConfidenceHeads( + token_s=hidden_size, + token_z=hidden_size, + num_plddt_bins=4, + num_pde_bins=4, + num_pae_bins=4, + token_level_confidence=False, + ) + # s: (batch_size * multiplicity, token_count, hidden_size) + s = torch.randn( + batch_size * multiplicity, + token_count, + hidden_size, + requires_grad=True, + ) + # z: (batch_size * multiplicity, token_count, token_count, hidden_size) + z = torch.randn( + batch_size * multiplicity, + token_count, + token_count, + hidden_size, + requires_grad=True, + ) + # atom_to_token: (2, 4, 2) + atom_to_token = torch.tensor( + ( + ((1, 0), (0, 0), (0, 0), (0, 0)), + ((1, 0), (0, 1), (0, 1), (0, 0)), + ), + dtype=torch.bool, + ) + # atom_pad_mask: (2, 4) + atom_pad_mask = torch.tensor( + ((1, 0, 0, 0), (1, 1, 1, 0)), + dtype=torch.bool, + ) + feats = { + "atom_to_token": atom_to_token, + "atom_pad_mask": atom_pad_mask, + "mol_type": torch.zeros((batch_size, token_count), dtype=torch.long), + "asym_id": torch.tensor(((0, 0), (0, 1)), dtype=torch.long), + "token_pad_mask": torch.tensor(((1, 0), (1, 1)), dtype=torch.float32), + } + # zeros: (batch_size * multiplicity,) + zeros = torch.zeros(batch_size * multiplicity) + monkeypatch.setattr( + vb_modules_confidencev2, + "compute_ptms", + lambda *_args, **_kwargs: (zeros, zeros, zeros, zeros, {}), + ) + + output = head( + s=s, + z=z, + x_pred=torch.randn(batch_size * multiplicity, atom_count, 3), + d=torch.zeros(batch_size * multiplicity, token_count, token_count), + feats=feats, + multiplicity=multiplicity, + pred_distogram_logits=torch.zeros( + batch_size, + token_count, + token_count, + 64, + ), + ) + + assert output["plddt_logits"].shape == ( + batch_size * multiplicity, + atom_count, + 4, + ) + assert output["resolved_logits"].shape == ( + batch_size * multiplicity, + atom_count, + 2, + ) + assert torch.isfinite(output["complex_pde"]).all() + torch.testing.assert_close( + output["complex_pde"][:multiplicity], + torch.zeros(multiplicity), + rtol=0.0, + atol=0.0, + ) + assert torch.isfinite(output["complex_plddt"]).all() + assert torch.isfinite(output["complex_iplddt"]).all() + assert not output["plddt_logits"][:multiplicity, 1:].any() + assert not output["plddt_logits"][multiplicity:, 3:].any() + + (output["plddt"].mean() + output["pde"].mean()).backward() + assert s.grad is not None and torch.isfinite(s.grad).all() + assert z.grad is not None and torch.isfinite(z.grad).all() + + +def test_esmfold_fold_single_uses_linker_masked_mean_plddt() -> None: + class FakeESMFold: + def infer(self, sequence: str) -> dict[str, torch.Tensor]: + assert sequence == "AC:DE" + return { + "plddt": torch.full((1, 29, 37), 1.0), + "mean_plddt": torch.tensor([87.5]), + "ptm": torch.tensor([0.75]), + } + + result = FastEsmForProteinFolding._fold_single( + FakeESMFold(), + "AC:DE", + return_pdb_string=False, + ) + + assert result["plddt"] == 87.5 + assert result["ptm"] == 0.75 + assert "pdb_string" not in result + + +def test_boltz_real_features_flow_through_tiny_core_and_structure_loss() -> None: + with seed_context(19): + features, template = build_boltz2_features("ACDE") + tiny_core = torch.nn.Linear(3, 3, bias=False) + + # reference_positions: (..., 3) + reference_positions = features["ref_pos"] + predicted_positions = tiny_core(reference_positions) + # atom_mask: (...) + atom_mask = features["atom_pad_mask"].bool() + loss = smooth_lddt_loss( + predicted_positions, + reference_positions, + is_nucleotide=torch.zeros_like(atom_mask), + coords_mask=atom_mask, + ) + loss.backward() + + assert len(template.atom_names) == int(atom_mask.sum()) + assert reference_positions.shape[1] % 32 == 0 + assert torch.isfinite(loss) + assert tiny_core.weight.grad is not None + assert torch.isfinite(tiny_core.weight.grad).all() + assert torch.count_nonzero(tiny_core.weight.grad) > 0 + + +@pytest.mark.parametrize("invalid_seed", (True, 1.5, "19")) +def test_esmfold2_seed_context_rejects_non_integer_seeds_without_rng_mutation( + invalid_seed: object, +) -> None: + random.seed(101) + np.random.seed(102) + torch.manual_seed(103) + python_state = random.getstate() + numpy_state = np.random.get_state() + torch_state = torch.random.get_rng_state() + + with ( + pytest.raises(TypeError, match="excluding bool"), + seed_context(invalid_seed), # type: ignore[arg-type] + ): + raise AssertionError("invalid seeds must fail before entering the context") + + assert random.getstate() == python_state + assert np.array_equal(np.random.get_state()[1], numpy_state[1]) + assert np.random.get_state()[0] == numpy_state[0] + assert np.random.get_state()[2:] == numpy_state[2:] + assert torch.equal(torch.random.get_rng_state(), torch_state) + + +@pytest.mark.parametrize("raise_inside", (False, True)) +def test_esmfold2_seed_context_restores_all_available_rng_streams( + raise_inside: bool, +) -> None: + random.seed(201) + np.random.seed(202) + torch.manual_seed(203) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(204) + python_state = random.getstate() + numpy_state = np.random.get_state() + torch_state = torch.random.get_rng_state() + cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + + def exercise() -> None: + with seed_context(29): + random.random() + np.random.random() + torch.rand(3) + if torch.cuda.is_available(): + torch.rand(3, device="cuda") + if raise_inside: + raise RuntimeError("seed-context-test") + + if raise_inside: + with pytest.raises(RuntimeError, match="seed-context-test"): + exercise() + else: + exercise() + + assert random.getstate() == python_state + assert np.array_equal(np.random.get_state()[1], numpy_state[1]) + assert np.random.get_state()[0] == numpy_state[0] + assert np.random.get_state()[2:] == numpy_state[2:] + assert torch.equal(torch.random.get_rng_state(), torch_state) + if cuda_state is not None: + current_cuda_state = torch.cuda.get_rng_state_all() + assert len(current_cuda_state) == len(cuda_state) + assert all( + torch.equal(actual, expected) + for actual, expected in zip(current_cuda_state, cuda_state, strict=True) + ) + + +def test_esmfold2_real_features_flow_through_tiny_core_and_tm_loss() -> None: + features = prepare_protein_features("ACDE") + # input_ids: (b, l) + input_ids = features["input_ids"] + sequence_length = input_ids.shape[1] + + with seed_context(23): + embedding = torch.nn.Embedding(int(input_ids.max()) + 1, 8) + pair_projection = torch.nn.Linear(8, 16) + + token_embeddings = embedding(input_ids) + pair_features = token_embeddings.unsqueeze(2) + token_embeddings.unsqueeze(1) + pae_logits = pair_projection(pair_features) + target_frames = Affine3D.identity( + (input_ids.shape[0], sequence_length), + dtype=pae_logits.dtype, + device=pae_logits.device, + ).tensor + loss = tm_loss( + pae_logits, + pred_affine=target_frames, + targ_affine=target_frames, + targ_mask=features["token_attention_mask"], + ) + loss.backward() + + assert features["ref_pos"].shape[1] % 32 == 0 + assert int(features["atom_attention_mask"].sum()) > sequence_length + assert torch.isfinite(loss) + for parameter in (*embedding.parameters(), *pair_projection.parameters()): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + assert ( + sum( + int(torch.count_nonzero(parameter.grad)) + for parameter in (*embedding.parameters(), *pair_projection.parameters()) + if parameter.grad is not None + ) + > 0 + ) + + +def test_binder_structure_loss_is_finite_and_differentiable() -> None: + with seed_context(29): + # distogram_logits: (2, 16, 16, 128) + distogram_logits = torch.randn((2, 16, 16, 128), requires_grad=True) + + losses = compute_structure_losses(distogram_logits, binder_length=12) + losses["total_loss"].mean().backward() + + assert set(losses) == { + "glob_loss", + "inter_contact_loss", + "intra_contact_loss", + "total_loss", + } + assert all(loss.shape == (2,) for loss in losses.values()) + assert all(torch.isfinite(loss).all() for loss in losses.values()) + assert distogram_logits.grad is not None + assert torch.isfinite(distogram_logits.grad).all() + assert torch.count_nonzero(distogram_logits.grad) > 0 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..c966081 --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1 @@ +"""Fast, offline unit tests.""" diff --git a/tests/unit/test_ankh_config.py b/tests/unit/test_ankh_config.py new file mode 100644 index 0000000..0e634b7 --- /dev/null +++ b/tests/unit/test_ankh_config.py @@ -0,0 +1,66 @@ +"""ANKH configuration round-trip contracts.""" + +from __future__ import annotations + +import pytest +import torch.nn as nn +from pathlib import Path + +from fastplms.models.ankh.modeling_ankh import ( + FastAnkhConfig, + FastAnkhForMaskedLMExtension, + FastAnkhModel, +) + + +def _small_config() -> FastAnkhConfig: + return FastAnkhConfig( + vocab_size=32, + d_model=32, + d_kv=8, + d_ff=64, + num_heads=4, + num_layers=2, + is_encoder_decoder=True, + ) + + +def test_ankh_config_consumes_serialized_encoder_decoder_field() -> None: + config = _small_config() + + restored = FastAnkhConfig.from_dict(config.to_dict()) + + assert restored.is_encoder_decoder is True + assert restored.to_dict() == config.to_dict() + + +def test_ankh_config_exposes_generic_t5_dimension_aliases() -> None: + config = _small_config() + + assert config.hidden_size == config.d_model == 32 + assert config.head_dim == config.d_kv == 8 + assert config.num_attention_heads == config.num_heads == 4 + assert config.num_hidden_layers == config.num_layers == 2 + + +@pytest.mark.parametrize("model_class", (FastAnkhModel, FastAnkhForMaskedLMExtension)) +def test_ankh_encoder_embedding_alias_is_declared_and_safetensors_safe( + model_class: type[FastAnkhModel] | type[FastAnkhForMaskedLMExtension], + tmp_path: Path, +) -> None: + model = model_class(_small_config()) + + assert model.encoder.embed_tokens.weight is model.shared.weight + assert model._tied_weights_keys == {"encoder.embed_tokens.weight": "shared.weight"} + + replacement = nn.Embedding(model.config.vocab_size, model.config.d_model) + model.set_input_embeddings(replacement) + assert model.get_input_embeddings() is replacement + assert model.shared is replacement + + model.save_pretrained(tmp_path, safe_serialization=True) + + +def test_ankh_config_rejects_encoder_only_serialized_contract() -> None: + with pytest.raises(ValueError, match="requires is_encoder_decoder=true"): + FastAnkhConfig(is_encoder_decoder=False) diff --git a/tests/unit/test_ankh_cpu_contract.py b/tests/unit/test_ankh_cpu_contract.py new file mode 100644 index 0000000..2236035 --- /dev/null +++ b/tests/unit/test_ankh_cpu_contract.py @@ -0,0 +1,891 @@ +"""Hermetic CPU contracts for ANKH's encoder and complete T5 views.""" + +from __future__ import annotations + +import json +import pytest +import torch +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier, Lock +from types import SimpleNamespace +from typing import ClassVar +from tokenizers import Tokenizer, decoders, models, pre_tokenizers, processors +from transformers import AutoTokenizer, GenerationConfig + +import fastplms.models.ankh.modeling_ankh as ankh_module +from fastplms.models.ankh.modeling_ankh import ( + FAST_ANKH_ENCODER, + FastAnkhConfig, + FastAnkhForConditionalGeneration, + FastAnkhForMaskedLMExtension, + FastAnkhForSequenceClassification, + FastAnkhForTokenClassification, + FastAnkhModel, + configure_ankh_tokenizer, + normalize_ankh_decoder_prompt, + normalize_ankh_sequence, + tokenize_ankh_decoder_prompts, + tokenize_ankh_sequences, +) + + +def _config(**overrides: object) -> FastAnkhConfig: + values = { + "vocab_size": 16, + "d_model": 8, + "d_kv": 4, + "d_ff": 16, + "num_heads": 2, + "num_layers": 2, + "num_decoder_layers": 2, + "dropout_rate": 0.0, + "pad_token_id": 0, + "eos_token_id": 1, + "decoder_start_token_id": 0, + "attn_backend": "eager", + "use_cache": False, + } + values.update(overrides) + return FastAnkhConfig(**values) + + +class _TinyTokenizer: + all_special_ids = (0, 1, 7) + name_or_path = "tiny-ankh" + vocab_size = 16 + special_tokens_map: ClassVar[dict[str, str]] = { + "pad_token": "", + "eos_token": "", + } + model_max_length = 64 + padding_side = "right" + truncation_side = "right" + + def __call__(self, sequences, **_kwargs): + alphabet = {"A": 2, "C": 3, "D": 4, "E": 5, "F": 6, "S": 7} + rows = [[alphabet.get(character, 8) for character in value] + [1] for value in sequences] + width = max(len(row) for row in rows) + input_ids = torch.tensor( # (b, l) + [row + [0] * (width - len(row)) for row in rows] + ) + return { + "input_ids": input_ids, + "attention_mask": input_ids.ne(0).to(dtype=torch.long), + } + + def get_vocab(self): + return {"": 0, "": 1, "A": 2, "C": 3, "D": 4, "E": 5, "F": 6} + + def get_added_vocab(self): + return {"": 7} + + +class _RecordingTokenizer(_TinyTokenizer): + is_fast = True + + def __init__(self) -> None: + self.backend_tokenizer = SimpleNamespace(pre_tokenizer=None) + self.calls: list[tuple[list[str], dict[str, object]]] = [] + + def __call__(self, sequences, **kwargs): + values = [sequences] if isinstance(sequences, str) else list(sequences) + self.calls.append((values, dict(kwargs))) + return super().__call__(values, **kwargs) + + +def test_ankh_tokenization_normalizes_raw_sequences_and_tight_sentinel_prompts() -> None: + tokenizer = _RecordingTokenizer() + + source = tokenize_ankh_sequences( + tokenizer, + ["M S T N P K", "AC\nDE"], + return_tensors="pt", + padding=True, + ) + prompt = tokenize_ankh_decoder_prompts( + tokenizer, + ["M ", "A\t"], + return_tensors="pt", + add_special_tokens=False, + ) + + assert tokenizer.calls[0][0] == ["MSTNPK", "ACDE"] + assert tokenizer.calls[1][0] == ["M", "A"] + assert source["input_ids"].ndim == 2 + assert prompt["input_ids"].ndim == 2 + assert type(tokenizer.backend_tokenizer.pre_tokenizer).__name__ == "Metaspace" + assert normalize_ankh_sequence(" M S T \n") == "MST" + assert normalize_ankh_decoder_prompt(" M ") == "M" + + +def test_ankh_tokenization_rejects_empty_inputs_and_real_slow_tokenizers() -> None: + with pytest.raises(ValueError, match="protein sequences must not be empty"): + normalize_ankh_sequence(" \n\t") + with pytest.raises(ValueError, match="decoder prompts must not be empty"): + normalize_ankh_decoder_prompt(" ") + with pytest.raises(TypeError, match="requires a fast tokenizer"): + configure_ankh_tokenizer(SimpleNamespace(is_fast=False)) + + +def test_offline_auto_tokenizer_flags_and_seq2seq_generation_config( + tmp_path: Path, +) -> None: + """Transformers 5.13 must resolve both tokenizer flags from local artifact bytes.""" + + vocabulary = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("A", -1.0), + ("C", -1.0), + ("D", -1.0), + ("X", -1.0), + ("", 0.0), + ] + backend = Tokenizer(models.Unigram(vocabulary, unk_id=2)) + replacement = "\N{LOWER ONE EIGHTH BLOCK}" + backend.pre_tokenizer = pre_tokenizers.Metaspace( + replacement=replacement, + prepend_scheme="never", + split=False, + ) + backend.decoder = decoders.Metaspace( + replacement=replacement, + prepend_scheme="never", + split=False, + ) + backend.post_processor = processors.TemplateProcessing( + single="$A ", + pair="$A $B ", + special_tokens=[("", 1)], + ) + backend.add_special_tokens(["", "", "", ""]) + backend.save(str(tmp_path / "tokenizer.json")) + tokenizer_config = { + "tokenizer_class": "T5Tokenizer", + "extra_ids": 1, + "pad_token": "", + "eos_token": "", + "unk_token": "", + "additional_special_tokens": [""], + } + special_tokens = { + "pad_token": "", + "eos_token": "", + "unk_token": "", + "additional_special_tokens": [""], + } + generation_config = { + "decoder_start_token_id": 0, + "pad_token_id": 0, + "eos_token_id": 1, + } + for name, payload in ( + ("tokenizer_config.json", tokenizer_config), + ("special_tokens_map.json", special_tokens), + ("generation_config.json", generation_config), + ): + (tmp_path / name).write_text( + json.dumps(payload, sort_keys=True), + encoding="utf-8", + ) + + encoded = [] + for use_fast in (True, False): + tokenizer = AutoTokenizer.from_pretrained( + tmp_path, + use_fast=use_fast, + local_files_only=True, + ) + assert tokenizer.is_fast + configure_ankh_tokenizer(tokenizer) + encoded.append( + tokenizer( + ["ACD", "AX"], + padding=True, + return_tensors="pt", + ) + ) + + assert torch.equal(encoded[0]["input_ids"], encoded[1]["input_ids"]) + assert torch.equal(encoded[0]["attention_mask"], encoded[1]["attention_mask"]) + loaded_generation = GenerationConfig.from_pretrained( + tmp_path, + local_files_only=True, + ) + assert loaded_generation.decoder_start_token_id == 0 + assert loaded_generation.pad_token_id == 0 + assert loaded_generation.eos_token_id == 1 + + +def test_ankh_explicit_and_model_owned_tokenizers_share_the_raw_sequence_contract() -> None: + explicit = _RecordingTokenizer() + owned = _RecordingTokenizer() + model = FastAnkhModel(_config()) + model.tokenizer = owned + + explicit_ids = model._tokenize_sequence_batch( + ["A C D", "EF"], + tokenizer=explicit, + return_tensors="pt", + padding=True, + )["input_ids"] + owned_ids = model._tokenize_sequence_batch( + ["A C D", "EF"], + return_tensors="pt", + padding=True, + )["input_ids"] + + assert explicit.calls[-1][0] == ["ACD", "EF"] + assert owned.calls[-1][0] == ["ACD", "EF"] + torch.testing.assert_close(explicit_ids, owned_ids) + + result = model.embed_dataset( + ["A C D"], + tokenizer=explicit, + full_embeddings=True, + ) + assert explicit.calls[-1][0] == ["ACD"] + assert result[0].load_tensor().shape == (3, model.config.d_model) + + +def test_ankh_ttt_uses_raw_residue_tokenization() -> None: + tokenizer = _RecordingTokenizer() + model = FastAnkhForMaskedLMExtension(_config()) + model.tokenizer = tokenizer + + input_ids = model._ttt_tokenize(seq=["A C D", "EF"]) + + assert tokenizer.calls[-1][0] == ["ACD", "EF"] + assert input_ids.ndim == 2 + + +def test_ankh_decoder_embedding_inputs_use_tight_sentinel_tokenization() -> None: + tokenizer = _RecordingTokenizer() + model = FastAnkhForConditionalGeneration(_config()).eval() + + decoder_input_ids, decoder_attention_mask, resolved = model._prepare_decoder_embedding_inputs( + batch_size=2, + decoder_inputs=["M ", "A\t"], + decoder_input_ids=None, + decoder_attention_mask=None, + tokenizer=tokenizer, + ) + + assert resolved is tokenizer + assert tokenizer.calls[-1][0] == ["M", "A"] + assert decoder_input_ids.shape == decoder_attention_mask.shape + + +@pytest.mark.parametrize("dropout_rate", (True, False, None, "0.1")) +def test_ankh_config_rejects_non_real_or_boolean_dropout(dropout_rate: object) -> None: + with pytest.raises(TypeError, match="dropout_rate must be a real number"): + _config(dropout_rate=dropout_rate) + + +@pytest.mark.parametrize("dropout_rate", (-0.01, 1.0, float("inf"), float("nan"))) +def test_ankh_config_rejects_dropout_outside_half_open_probability_range( + dropout_rate: float, +) -> None: + with pytest.raises(ValueError, match=r"dropout_rate must be in \[0, 1\)"): + _config(dropout_rate=dropout_rate) + + +def test_ankh_custom_encoder_invokes_every_t5_stack_dropout_site() -> None: + model = FastAnkhModel(_config(num_layers=1, dropout_rate=0.4)).train() + layer = model.encoder.block[0] + dropout_sites = { + "stack_input_and_final": model.encoder.dropout, + "attention_residual": layer.layer[0].dropout, + "ff_internal": layer.layer[1].DenseReluDense.dropout, + "ff_residual": layer.layer[1].dropout, + } + observed_calls = {name: 0 for name in dropout_sites} + handles = [] + for name, module in dropout_sites.items(): + handles.append( + module.register_forward_hook( + lambda _module, _inputs, _output, site=name: observed_calls.__setitem__( + site, observed_calls[site] + 1 + ) + ) + ) + + try: + model( + input_ids=torch.tensor([[2, 3, 4, 1]]), + attention_mask=torch.ones(1, 4, dtype=torch.long), + ) + finally: + for handle in handles: + handle.remove() + + assert observed_calls == { + "stack_input_and_final": 2, + "attention_residual": 1, + "ff_internal": 1, + "ff_residual": 1, + } + assert all(module.p == pytest.approx(0.4) for module in dropout_sites.values()) + + +def test_ankh_dropout_is_eval_exact_and_seeded_in_training() -> None: + input_ids = torch.tensor([[2, 3, 4, 1], [5, 6, 1, 0]]) # (b=2, l=4) + attention_mask = input_ids.ne(0) # (b, l) + + torch.manual_seed(31) + zero_dropout = FastAnkhModel(_config(dropout_rate=0.0)).eval() + torch.manual_seed(31) + configured_dropout = FastAnkhModel(_config(dropout_rate=0.5)).eval() + + zero_output = zero_dropout( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + configured_output = configured_dropout( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + torch.testing.assert_close(configured_output, zero_output, rtol=0.0, atol=0.0) + + configured_dropout.train() + torch.manual_seed(47) + first = configured_dropout( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + torch.manual_seed(47) + repeated = configured_dropout( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + torch.manual_seed(53) + different_seed = configured_dropout( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + + torch.testing.assert_close(first, repeated, rtol=0.0, atol=0.0) + assert not torch.equal(first, configured_output) + assert not torch.equal(first, different_seed) + + +@pytest.mark.parametrize( + "model_class", + ( + FAST_ANKH_ENCODER, + FastAnkhModel, + FastAnkhForMaskedLMExtension, + FastAnkhForSequenceClassification, + FastAnkhForTokenClassification, + ), +) +@pytest.mark.parametrize("argument", ("use_cache", "decoder_input_ids", "misspelled_option")) +def test_ankh_encoder_views_reject_every_unexpected_forward_argument( + model_class: ( + type[FastAnkhModel] + | type[FastAnkhForMaskedLMExtension] + | type[FastAnkhForSequenceClassification] + | type[FastAnkhForTokenClassification] + ), + argument: str, +) -> None: + model = model_class(_config(num_labels=3)).eval() + + with pytest.raises(TypeError, match=argument): + model(input_ids=torch.tensor([[2, 3, 1]]), **{argument: True}) + + +def test_seq2seq_embedding_selects_encoder_and_explicit_decoder_layers() -> None: + torch.manual_seed(7) + model = FastAnkhForConditionalGeneration(_config()).eval() + model.tokenizer = _TinyTokenizer() + input_ids = torch.tensor([[2, 3, 4, 1], [5, 6, 1, 0]]) # (b=2, l_enc=4) + attention_mask = input_ids.ne(0) # (b, l_enc) + decoder_input_ids = torch.tensor([[2, 7, 1, 0], [3, 1, 0, 0]]) # (b, l_dec=4) + decoder_attention_mask = decoder_input_ids.ne(0) # (b, l_dec) + + encoder_states = model._embed( + input_ids, + attention_mask, + hidden_state_source="encoder", + store_all_hidden_states=True, + ) + decoder_states = model._embed( + input_ids, + attention_mask, + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + store_all_hidden_states=True, + ) + + assert encoder_states.shape == (2, 3, 4, 8) + assert decoder_states.shape == (2, 3, 4, 8) + for layer_index in range(encoder_states.shape[1]): + assert torch.equal( + model._embed( + input_ids, + attention_mask, + hidden_state_source="encoder", + hidden_state_index=layer_index, + ), + encoder_states[:, layer_index], + ) + for layer_index in range(decoder_states.shape[1]): + assert torch.equal( + model._embed( + input_ids, + attention_mask, + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + hidden_state_index=layer_index, + ), + decoder_states[:, layer_index], + ) + assert torch.equal( + model._embed(input_ids, attention_mask), + encoder_states[:, -1], + ) + + +def test_decoder_default_attention_mask_preserves_t5_start_and_masks_padding() -> None: + model = FastAnkhForConditionalGeneration(_config()).eval() + decoder_input_ids = torch.tensor([[0, 5, 1, 0], [0, 6, 7, 1]]) # (b=2, l_dec=4) + + prepared_ids, prepared_mask, _ = model._prepare_decoder_embedding_inputs( + batch_size=2, + decoder_inputs=None, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=None, + tokenizer=_TinyTokenizer(), + ) + + assert torch.equal(prepared_ids, decoder_input_ids) + assert torch.equal( + prepared_mask, + torch.tensor([[True, True, True, False], [True, True, True, True]]), + ) + + +def test_seq2seq_view_forces_eager_without_changing_encoder_backend_contract() -> None: + seq2seq_config = _config(attn_backend=None) + seq2seq = FastAnkhForConditionalGeneration(seq2seq_config) + encoder = FastAnkhModel(_config(attn_backend=None)) + + assert seq2seq.config.attn_backend == "eager" + assert seq2seq.config._attn_implementation == "eager" + assert encoder.attn_backend == "sdpa" + + +def test_decoder_embeddings_require_explicit_aligned_inputs() -> None: + model = FastAnkhForConditionalGeneration(_config()).eval() + model.tokenizer = _TinyTokenizer() + input_ids = torch.tensor([[2, 3, 1], [4, 1, 0]]) # (b=2, l=3) + + with pytest.raises(ValueError, match="requires exactly one"): + model._embed(input_ids, hidden_state_source="decoder") + with pytest.raises(ValueError, match="exactly one"): + model._embed( + input_ids, + hidden_state_source="decoder", + decoder_inputs=["AC", "D"], + decoder_input_ids=input_ids, + ) + with pytest.raises(ValueError, match="align one-to-one"): + model._embed( + input_ids, + hidden_state_source="decoder", + decoder_inputs=["AC"], + ) + with pytest.raises(ValueError, match="only valid"): + model._embed( + input_ids, + hidden_state_source="encoder", + decoder_input_ids=input_ids, + ) + + +def test_decoder_embedding_batch_masks_start_eos_padding_and_sentinels() -> None: + model = FastAnkhForConditionalGeneration(_config()).eval() + tokenizer = _TinyTokenizer() + decoder_input_ids = torch.tensor([[2, 7, 1, 0], [3, 1, 0, 0]]) # (b=2, l_dec=4) + decoder_attention_mask = decoder_input_ids.ne(0) # (b, l_dec) + + batch = model._embedding_batch( + ["ACD", "EF"], + tokenizer=tokenizer, + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + ) + + assert batch.X.shape == (2, 4, 8) + assert torch.equal( + batch.residue_mask, + torch.tensor([[True, False, False, False], [True, False, False, False]]), + ) + + string_batch = model._embedding_batch( + ["ACD", "EF"], + tokenizer=tokenizer, + hidden_state_source="decoder", + decoder_inputs=["AS", "C"], + ) + assert string_batch.X.shape == (2, 3, 8) + assert torch.equal( + string_batch.residue_mask, + torch.tensor([[True, False, False], [True, False, False]]), + ) + + +def test_decoder_embed_dataset_slices_aligned_inputs_and_records_provenance() -> None: + model = FastAnkhForConditionalGeneration(_config()).eval() + tokenizer = _TinyTokenizer() + decoder_input_ids = torch.tensor([[2, 7, 1, 0], [3, 1, 0, 0]]) # (b=2, l_dec=4) + decoder_attention_mask = decoder_input_ids.ne(0) # (b, l_dec) + + result = model.embed_dataset( + ["ACD", "EF"], + tokenizer=tokenizer, + batch_size=1, + full_embeddings=True, + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + ) + + assert [tuple(record.load_tensor().shape) for record in result] == [(1, 8), (1, 8)] + assert result.metadata["hidden_state_source"] == "decoder" + assert result.metadata["hidden_state_index"] == -1 + assert result.metadata["store_all_hidden_states"] is False + assert result.metadata["decoder_alignment"] == "input-position" + assert len(result.metadata["decoder_input_fingerprint"]) == 64 + assert len(result.metadata["decoder_attention_mask_fingerprint"]) == 64 + assert result.metadata["model_embedding"]["hidden_state_stack"] == "decoder" + assert ( + result.metadata["model_embedding"]["decoder_residue_mask"] + == "attention-mask-minus-tokenizer-specials" + ) + + +def test_encoder_only_view_rejects_decoder_hidden_states() -> None: + model = FastAnkhModel(_config()).eval() + with pytest.raises(ValueError, match="AutoModelForSeq2SeqLM"): + model._embed( + torch.tensor([[2, 3, 1]]), + hidden_state_source="decoder", + decoder_input_ids=torch.tensor([[2, 1]]), + ) + + +def test_sdpa_output_attentions_fallback_keeps_padding_mask_and_backend() -> None: + model = FastAnkhModel(_config(attn_backend="sdpa")).eval() + input_ids = torch.tensor([[2, 3, 1, 0], [4, 1, 0, 0]]) # (b=2, l=4) + attention_mask = input_ids.ne(0) # (b, l) + + with pytest.warns( + RuntimeWarning, + match="requested 'sdpa'.*using 'eager'.*call only", + ) as fallback_warnings: + output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + return_dict=True, + ) + + assert len(fallback_warnings) == 1 + assert model.attn_backend == "sdpa" + assert output.attentions is not None + for layer_attention in output.attentions: + assert torch.count_nonzero(layer_attention[0, :, :, 3]) == 0 + assert torch.count_nonzero(layer_attention[1, :, :, 2:]) == 0 + + +@pytest.mark.parametrize( + "model_class", + ( + FastAnkhModel, + FastAnkhForMaskedLMExtension, + FastAnkhForSequenceClassification, + FastAnkhForTokenClassification, + ), +) +def test_encoder_auto_classes_honor_tuple_and_dict_outputs( + model_class: ( + type[FastAnkhModel] + | type[FastAnkhForMaskedLMExtension] + | type[FastAnkhForSequenceClassification] + | type[FastAnkhForTokenClassification] + ), +) -> None: + model = model_class(_config(num_labels=3)).eval() + input_ids = torch.tensor([[2, 3, 1], [4, 1, 0]]) # (b=2, l=3) + attention_mask = input_ids.ne(0) # (b, l) + + dictionary_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + tuple_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=False, + ) + + assert isinstance(tuple_output, tuple) + expected = ( + dictionary_output.last_hidden_state + if model_class is FastAnkhModel + else dictionary_output.logits + ) + assert torch.equal(tuple_output[0], expected) + assert tuple_output[1] is not None + + +@pytest.mark.parametrize( + "model_class", + ( + FastAnkhModel, + FastAnkhForMaskedLMExtension, + FastAnkhForSequenceClassification, + FastAnkhForTokenClassification, + ), +) +def test_encoder_auto_classes_resize_shared_input_embeddings( + model_class: ( + type[FastAnkhModel] + | type[FastAnkhForMaskedLMExtension] + | type[FastAnkhForSequenceClassification] + | type[FastAnkhForTokenClassification] + ), +) -> None: + model = model_class(_config(num_labels=3)) + assert model.shared is model.get_input_embeddings() + assert model.encoder.embed_tokens is model.shared + + resized = model.resize_token_embeddings(19) + + assert resized is model.get_input_embeddings() + assert model.config.vocab_size == 19 + assert model.shared is model.get_input_embeddings() + assert model.encoder.embed_tokens is model.shared + if model_class is FastAnkhForMaskedLMExtension: + assert model.get_output_embeddings().out_features == 19 + + +@pytest.mark.parametrize( + ("model_class", "labels"), + ( + ( + FastAnkhForMaskedLMExtension, + torch.tensor([[2, 3, 1], [4, 1, -100]]), + ), + ( + FastAnkhForSequenceClassification, + torch.tensor([1, 2]), + ), + ( + FastAnkhForTokenClassification, + torch.tensor([[1, 2, 0], [2, 0, -100]]), + ), + ), +) +def test_encoder_task_heads_produce_finite_loss_and_gradients( + model_class: ( + type[FastAnkhForMaskedLMExtension] + | type[FastAnkhForSequenceClassification] + | type[FastAnkhForTokenClassification] + ), + labels: torch.Tensor, +) -> None: + model = model_class(_config(num_labels=3)).train() + input_ids = torch.tensor([[2, 3, 1], [4, 1, 0]]) # (b=2, l=3) + output = model( + input_ids=input_ids, + attention_mask=input_ids.ne(0), + labels=labels, + return_dict=True, + ) + + assert output.loss is not None and torch.isfinite(output.loss) + output.loss.backward() + assert model.shared.weight.grad is not None + assert torch.isfinite(model.shared.weight.grad).all() + + +def test_seq2seq_head_produces_finite_loss_and_gradients() -> None: + model = FastAnkhForConditionalGeneration(_config()).train() + input_ids = torch.tensor([[2, 3, 1], [4, 1, 0]]) # (b=2, l=3) + output = model( + input_ids=input_ids, + attention_mask=input_ids.ne(0), + labels=torch.tensor([[2, 1, -100], [3, 1, -100]]), + return_dict=True, + ) + + assert output.loss is not None and torch.isfinite(output.loss) + output.loss.backward() + assert model.shared.weight.grad is not None + assert torch.isfinite(model.shared.weight.grad).all() + + +def test_complete_t5_checkpoint_loads_clean_encoder_and_seq2seq_views( + tmp_path: Path, +) -> None: + torch.manual_seed(11) + full = FastAnkhForConditionalGeneration(_config()).eval() + full.save_pretrained(tmp_path, safe_serialization=True) + + encoder, encoder_info = FastAnkhModel.from_pretrained( + tmp_path, + local_files_only=True, + output_loading_info=True, + ) + reloaded, seq2seq_info = FastAnkhForConditionalGeneration.from_pretrained( + tmp_path, + local_files_only=True, + output_loading_info=True, + ) + + assert not hasattr(encoder, "decoder") + assert not encoder_info["missing_keys"] + assert not encoder_info["unexpected_keys"] + assert not seq2seq_info["missing_keys"] + assert not seq2seq_info["unexpected_keys"] + expected_state = full.state_dict() + observed_state = reloaded.state_dict() + assert set(observed_state) == set(expected_state) + assert all(torch.equal(observed_state[key], value) for key, value in expected_state.items()) + + +def test_tokenizer_load_context_is_per_instance_and_offline_scoped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[tuple[str, dict[str, object]]] = [] + + class Tokenizer(_TinyTokenizer): + backend_tokenizer = SimpleNamespace(pre_tokenizer=None) + + def load_tokenizer(source: str, **kwargs): + requests.append((source, kwargs)) + return Tokenizer() + + monkeypatch.setattr( + ankh_module.AutoTokenizer, + "from_pretrained", + staticmethod(load_tokenizer), + ) + first = FastAnkhModel(_config()) + second = FastAnkhModel(_config()) + first.config._name_or_path = "ankh-first" + first.config._commit_hash = "a" * 40 + second.config._name_or_path = "ankh-second" + second.config._commit_hash = "b" * 40 + first_token = object() + second_token = object() + first.__dict__["_fastplms_tokenizer_load_context"] = { + "cache_dir": "cache-one", + "local_files_only": True, + "token": first_token, + } + second.__dict__["_fastplms_tokenizer_load_context"] = { + "cache_dir": "cache-two", + "local_files_only": False, + "token": second_token, + } + + assert first.tokenizer is first.tokenizer + assert second.tokenizer is second.tokenizer + assert requests[0][0] == "ankh-first" + assert requests[0][1] == { + "cache_dir": "cache-one", + "local_files_only": True, + "token": first_token, + "revision": "a" * 40, + } + assert requests[1][0] == "ankh-second" + assert requests[1][1] == { + "cache_dir": "cache-two", + "local_files_only": False, + "token": second_token, + "revision": "b" * 40, + } + + +def test_tokenizer_load_context_is_isolated_during_concurrent_first_access( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Concurrent models must never exchange revision, cache, or token context.""" + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + rendezvous = Barrier(2) + request_lock = Lock() + requests: dict[str, dict[str, object]] = {} + + class Tokenizer(_TinyTokenizer): + def __init__(self, source: str) -> None: + self.source = source + self.backend_tokenizer = SimpleNamespace(pre_tokenizer=None) + + def load_tokenizer(source: str, **kwargs): + rendezvous.wait(timeout=3.0) + with request_lock: + requests[source] = dict(kwargs) + return Tokenizer(source) + + monkeypatch.setattr( + ankh_module.AutoTokenizer, + "from_pretrained", + staticmethod(load_tokenizer), + ) + first = FastAnkhModel(_config(num_layers=1, num_decoder_layers=1)) + second = FastAnkhModel(_config(num_layers=1, num_decoder_layers=1)) + first.config._name_or_path = "ankh-first" + first.config._commit_hash = "a" * 40 + second.config._name_or_path = "ankh-second" + second.config._commit_hash = "b" * 40 + first_token = object() + second_token = object() + first.__dict__["_fastplms_tokenizer_load_context"] = { + "cache_dir": "cache-one", + "local_files_only": True, + "token": first_token, + } + second.__dict__["_fastplms_tokenizer_load_context"] = { + "cache_dir": "cache-two", + "local_files_only": True, + "token": second_token, + } + + with ThreadPoolExecutor(max_workers=2) as executor: + first_future = executor.submit(getattr, first, "tokenizer") + second_future = executor.submit(getattr, second, "tokenizer") + first_tokenizer = first_future.result(timeout=5.0) + second_tokenizer = second_future.result(timeout=5.0) + + assert first_tokenizer.source == "ankh-first" + assert second_tokenizer.source == "ankh-second" + assert requests == { + "ankh-first": { + "cache_dir": "cache-one", + "local_files_only": True, + "token": first_token, + "revision": "a" * 40, + }, + "ankh-second": { + "cache_dir": "cache-two", + "local_files_only": True, + "token": second_token, + "revision": "b" * 40, + }, + } + assert first.tokenizer is first_tokenizer + assert second.tokenizer is second_tokenizer diff --git a/tests/unit/test_attention_interfaces.py b/tests/unit/test_attention_interfaces.py new file mode 100644 index 0000000..55bbc8f --- /dev/null +++ b/tests/unit/test_attention_interfaces.py @@ -0,0 +1,1552 @@ +"""Transformers registry and Flex cache contracts.""" + +from __future__ import annotations + +import ast +import inspect +import json +import sys +import warnings +import pytest +import torch +from concurrent.futures import ThreadPoolExecutor +from importlib.metadata import version +from pathlib import Path +from threading import Barrier +from types import SimpleNamespace +from transformers import AttentionInterface + + +_TRANSFORMERS_FLASH_HANDLERS = { + name: AttentionInterface()[name] for name in ("flash_attention_2", "flash_attention_3") +} + +import fastplms.attention.interfaces as attention_interfaces # noqa: E402 +import fastplms.models.ankh.modeling_ankh as ankh_module # noqa: E402 +import fastplms.models.dplm.modeling_dplm as dplm_module # noqa: E402 +import fastplms.models.esm_plusplus.modeling_esm_plusplus as esmpp_module # noqa: E402 +from fastplms.attention import ( # noqa: E402 + FASTPLMS_ATTENTION_FUNCTIONS, + FASTPLMS_ATTENTION_MASKS, + AttentionBackend, + FastPLMsAttentionMixin, + _core, + _kernel_lock, + validate_transformers_attention_interfaces, +) +from fastplms.embeddings.runner import _attention_kernel_metadata # noqa: E402 +from fastplms.models.ankh.modeling_ankh import ( # noqa: E402 + AnkhSelfAttention, + FastAnkhConfig, + FastAnkhModel, +) +from fastplms.models.dplm.modeling_dplm import ( # noqa: E402 + DPLMConfig, + DPLMModel, + ModifiedEsmSelfAttention, +) +from fastplms.models.dplm2.modeling_dplm2 import ( # noqa: E402 + DPLM2Config, + DPLM2Model, +) +from fastplms.models.esm2.modeling_fastesm import ( # noqa: E402 + EsmSelfAttention, + FastEsmConfig, + FastEsmModel, +) +from fastplms.registry import get_model_registry # noqa: E402 + + +FUNCTION_BACKENDS = ( + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", +) +MASK_BACKENDS = ("eager", *FUNCTION_BACKENDS) + + +class _RejectingTransformersBase: + def _check_and_adjust_attn_implementation(self, *args, **kwargs) -> str: + raise AssertionError("Transformers' source-Flash resolver must not run.") + + +class _SupportedFlashMixin(FastPLMsAttentionMixin, _RejectingTransformersBase): + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_flash_attn_3 = True + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + + +@pytest.mark.parametrize( + ( + "implementation", + "repository", + "revision", + "kernel", + "variant", + ), + ( + ( + "flash_attention_2", + "kernels-community/flash-attn2", + "db6b51744f0cd7061386442c09df890fc6d9f47e", + SimpleNamespace( + fwd=lambda **kwargs: kwargs["q"], + varlen_fwd=lambda **kwargs: kwargs["q"], + flash_attn_func=lambda **kwargs: kwargs["q"], + flash_attn_varlen_func=lambda **kwargs: kwargs["q"], + ), + "flash_attn2", + ), + ( + "flash_attention_3", + "kernels-community/flash-attn3", + "43f0bd269777115d94ff826e0d113ce9c1c9087b", + SimpleNamespace( + flash_attn_func=lambda **kwargs: kwargs["q"], + flash_attn_varlen_func=lambda **kwargs: kwargs["q"], + ), + "flash_attn3", + ), + ), +) +def test_flash_backend_loads_only_its_hugging_face_kernel( + monkeypatch: pytest.MonkeyPatch, + implementation: str, + repository: str, + revision: str, + kernel: object, + variant: str, +) -> None: + requested: list[tuple[str, str]] = [] + + def locked_kernel(repo_id: str, revision: str) -> object: + requested.append((repo_id, revision)) + return kernel + + monkeypatch.setattr(_core, "load_locked_kernel", locked_kernel) + _core._FLASH_KERNELS.clear() + + assert _core._ensure_flash_kernels_loaded(implementation) == (kernel, variant) + assert _core._ensure_flash_kernels_loaded(implementation) == (kernel, variant) + assert requested == [(repository, revision)] + + +def test_flash_kernel_variant_mismatch_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requested: list[tuple[str, str]] = [] + flash_attention_3_kernel = SimpleNamespace( + flash_attn_func=object(), + flash_attn_varlen_func=object(), + ) + + def locked_kernel(repo_id: str, revision: str) -> object: + requested.append((repo_id, revision)) + return flash_attention_3_kernel + + monkeypatch.setattr(_core, "load_locked_kernel", locked_kernel) + _core._FLASH_KERNELS.clear() + + with pytest.raises( + RuntimeError, + match=( + "kernels-community/flash-attn2@" + "db6b51744f0cd7061386442c09df890fc6d9f47e exposed 'flash_attn3'; " + "expected 'flash_attn2'" + ), + ): + _core._ensure_flash_kernels_loaded("flash_attention_2") + assert requested == [ + ( + "kernels-community/flash-attn2", + "db6b51744f0cd7061386442c09df890fc6d9f47e", + ) + ] + + +def test_flash_kernel_load_error_retains_pinned_identity_and_cause( + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure = OSError("incompatible precompiled kernel") + + def locked_kernel(repo_id: str, revision: str) -> object: + assert repo_id == "kernels-community/flash-attn3" + assert revision == "43f0bd269777115d94ff826e0d113ce9c1c9087b" + raise failure + + monkeypatch.setattr(_core, "load_locked_kernel", locked_kernel) + _core._FLASH_KERNELS.clear() + + with pytest.raises( + RuntimeError, + match=( + "Unable to load the manifest-pinned kernel " + "kernels-community/flash-attn3@" + "43f0bd269777115d94ff826e0d113ce9c1c9087b" + ), + ) as captured: + _core._ensure_flash_kernels_loaded("flash_attention_3") + assert captured.value.__cause__ is failure + + +def test_locked_kernel_is_hash_validated_before_import( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + revision = "d" * 40 + lock_path = tmp_path / "kernels.lock" + lock_path.write_text( + json.dumps( + [ + { + "repo_id": "kernels-community/flash-attn2", + "sha": revision, + "variants": { + "torch213-cxx11-cu130-x86_64-linux": { + "hash": f"sha256-{'a' * 64}", + "hash_type": "git_lfs_concat", + } + }, + } + ] + ), + encoding="utf-8", + ) + events: list[str] = [] + kernel = object() + validated_path = tmp_path / "validated-variant" + + class KernelLock: + @classmethod + def from_json(cls, entry: dict[str, object]) -> SimpleNamespace: + return SimpleNamespace(sha=entry["sha"], variants=entry["variants"]) + + def install_kernel( + repository: str, + *, + revision: str, + variant_locks: dict[str, object], + ) -> Path: + assert repository == "kernels-community/flash-attn2" + assert revision == "d" * 40 + assert set(variant_locks) == {"torch213-cxx11-cu130-x86_64-linux"} + events.append("validate") + return validated_path + + def get_local_kernel(path: Path) -> object: + assert path == validated_path + events.append("import") + return kernel + + monkeypatch.setattr(_kernel_lock, "_kernel_lock_path", lambda: lock_path) + monkeypatch.setitem( + sys.modules, + "kernels", + SimpleNamespace( + get_local_kernel=get_local_kernel, + install_kernel=install_kernel, + ), + ) + monkeypatch.setitem(sys.modules, "kernels.lockfile", SimpleNamespace(KernelLock=KernelLock)) + + assert _kernel_lock.load_locked_kernel("kernels-community/flash-attn2", revision) is kernel + assert events == ["validate", "import"] + + +def test_locked_kernel_offline_resolves_sparse_snapshot_without_hub_api( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + snapshot = tmp_path / "snapshot" + variant_path = snapshot / "build" / "torch213-cxx11-cu130-x86_64-linux" + variant_path.mkdir(parents=True) + variant = SimpleNamespace(variant_str=variant_path.name) + expected_hash = f"sha256-{'a' * 64}" + events: list[str] = [] + kernel = object() + + def validate_kernel(*, repo_path: Path, variant: str, hash: str) -> None: + assert repo_path == snapshot + assert variant == variant_path.name + assert hash == expected_hash + events.append("validate") + + def get_local_kernel(path: Path) -> object: + assert path == variant_path + events.append("import") + return kernel + + monkeypatch.setattr(_kernel_lock, "_offline_snapshot_path", lambda *_: snapshot) + monkeypatch.setitem(sys.modules, "kernels", SimpleNamespace(get_local_kernel=get_local_kernel)) + monkeypatch.setitem( + sys.modules, + "kernels.utils", + SimpleNamespace(validate_kernel=validate_kernel), + ) + monkeypatch.setitem( + sys.modules, + "kernels.variants", + SimpleNamespace( + get_variants_local=lambda path: [variant], + resolve_variants=lambda variants: (variants, []), + ), + ) + + assert ( + _kernel_lock._load_offline_locked_kernel( + "kernels-community/flash-attn2", + "d" * 40, + {variant_path.name: SimpleNamespace(hash=expected_hash)}, + ) + is kernel + ) + assert events == ["validate", "import"] + + +def test_locked_kernel_offline_rejects_unlocked_cached_variant( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + snapshot = tmp_path / "snapshot" + (snapshot / "build" / "unexpected-variant").mkdir(parents=True) + monkeypatch.setattr(_kernel_lock, "_offline_snapshot_path", lambda *_: snapshot) + + with pytest.raises(RuntimeError, match="contains unlocked variants"): + _kernel_lock._load_offline_locked_kernel( + "kernels-community/flash-attn2", + "d" * 40, + {"expected-variant": SimpleNamespace(hash=f"sha256-{'a' * 64}")}, + ) + + +def test_transformers_flash_hook_defers_binary_loading_until_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dependency_checks: list[None] = [] + monkeypatch.setattr( + attention_interfaces, + "require_kernels_package", + lambda: dependency_checks.append(None), + ) + model = object.__new__(_SupportedFlashMixin) + + for implementation in ("flash_attention_2", "flash_attention_3"): + assert ( + model._check_and_adjust_attn_implementation( + implementation, + is_init_check=True, + ) + == implementation + ) + + assert dependency_checks == [None, None] + + +@pytest.mark.parametrize( + "implementation", + ( + "kernels-community/flash-attn2", + "another-owner/custom-kernel", + "paged|flash_attention_2", + "flash_attention_4", + ), +) +def test_transformers_flash_hook_rejects_external_and_unadvertised_kernels( + implementation: str, +) -> None: + model = object.__new__(_SupportedFlashMixin) + with pytest.raises(ValueError, match="does not support"): + model._check_and_adjust_attn_implementation(implementation) + + with pytest.raises(ValueError, match="does not load external"): + model._check_and_adjust_attn_implementation( + "flash_attention_2", + allow_all_kernels=True, + ) + + +def test_transformers_flash_hook_rejects_an_unadvertised_family() -> None: + model = object.__new__(FastPLMsAttentionMixin) + with pytest.raises(ValueError, match="does not support"): + model._check_and_adjust_attn_implementation("flash_attention_2") + + +def test_public_attention_setter_matches_transformers_513_kernel_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + signature = inspect.signature(FastPLMsAttentionMixin.set_attn_implementation) + assert tuple(signature.parameters) == ( + "self", + "attn_implementation", + "allow_all_kernels", + ) + assert signature.parameters["allow_all_kernels"].default is False + + model = object.__new__(_SupportedFlashMixin) + with pytest.raises(ValueError, match="does not load external"): + model.set_attn_implementation( + "flash_attention_2", + allow_all_kernels=True, + ) + + def unavailable_kernel_runtime() -> None: + raise RuntimeError("precompiled kernel runtime is unavailable") + + monkeypatch.setattr( + attention_interfaces, + "require_kernels_package", + unavailable_kernel_runtime, + ) + with pytest.raises(RuntimeError, match="kernel runtime is unavailable"): + model.set_attn_implementation("flash_attention_2") + + +def test_flash_kernel_identity_is_typed_and_manifest_owned() -> None: + registry = get_model_registry() + assert { + implementation: ( + spec.repository, + spec.revision, + spec.version, + spec.expected_variant, + spec.dtypes, + ) + for implementation, spec in registry.attention_kernels.items() + } == { + "flash_attention_2": ( + "kernels-community/flash-attn2", + "db6b51744f0cd7061386442c09df890fc6d9f47e", + 2, + "flash_attn2", + ("bfloat16",), + ), + "flash_attention_3": ( + "kernels-community/flash-attn3", + "43f0bd269777115d94ff826e0d113ce9c1c9087b", + 1, + "flash_attn3", + ("bfloat16",), + ), + } + source = Path(_core.__file__).read_text(encoding="utf-8") + assert "kernels-community/flash-attn2" not in source + assert "kernels-community/flash-attn3" not in source + for family in registry.families.values(): + if any(name.startswith("flash_attention_") for name in family.attention): + assert "registry.py" in family.runtime_paths + + assert _attention_kernel_metadata("flash_attention_3") == { + "repository": "kernels-community/flash-attn3", + "revision": "43f0bd269777115d94ff826e0d113ce9c1c9087b", + "version": 1, + "expected_variant": "flash_attn3", + "dtypes": ["bfloat16"], + } + + +def test_kernels_flash_rejects_fp32_before_loading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_load(_implementation: str) -> None: + raise AssertionError("unsupported dtypes must fail before kernel loading") + + monkeypatch.setattr(_core, "_ensure_flash_kernels_loaded", unexpected_load) + monkeypatch.setattr( + _core, + "_validate_kernels_flash_device", + lambda *_args: torch.device("cuda"), + ) + X = torch.zeros(1, 4, 2, 8, dtype=torch.float32) + with pytest.raises(RuntimeError, match=r"bfloat16.*received float32"): + _core.kernels_flash_attention_func(X, X, X, implementation="flash_attention_3") + + +def test_kernels_flash_rejects_mixed_qkv_dtypes_before_loading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_load(_implementation: str) -> None: + raise AssertionError("mismatched dtypes must fail before kernel loading") + + monkeypatch.setattr(_core, "_ensure_flash_kernels_loaded", unexpected_load) + monkeypatch.setattr( + _core, + "_validate_kernels_flash_device", + lambda *_args: torch.device("cuda"), + ) + Q = torch.zeros(1, 4, 2, 8, dtype=torch.bfloat16) + K = torch.zeros(1, 4, 2, 8, dtype=torch.float32) + V = torch.zeros(1, 4, 2, 8, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="Q, K, and V to share one dtype"): + _core.kernels_flash_attention_func(Q, K, V, implementation="flash_attention_3") + + +def test_kernels_flash_rejects_cpu_bf16_before_loading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_load(_implementation: str) -> None: + raise AssertionError("CPU tensors must fail before kernel loading") + + monkeypatch.setattr(_core, "_ensure_flash_kernels_loaded", unexpected_load) + X = torch.zeros(1, 4, 2, 8, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match=r"requires CUDA Q, K, and V.*cpu"): + _core.kernels_flash_attention_func(X, X, X, implementation="flash_attention_2") + + +def test_kernels_flash_rejects_mixed_devices_before_loading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_load(_implementation: str) -> None: + raise AssertionError("mixed devices must fail before kernel loading") + + monkeypatch.setattr(_core, "_ensure_flash_kernels_loaded", unexpected_load) + Q = torch.zeros(1, 4, 2, 8, dtype=torch.bfloat16) + K = torch.empty(1, 4, 2, 8, dtype=torch.bfloat16, device="meta") + with pytest.raises(RuntimeError, match=r"on one device.*cpu, meta, cpu"): + _core.kernels_flash_attention_func(Q, K, Q, implementation="flash_attention_3") + + +def test_causal_masked_flash_uses_varlen_and_zeroes_padding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + _core, + "_validate_kernels_flash_device", + lambda *_args: torch.device("cpu"), + ) + monkeypatch.setattr( + _core, + "_ensure_flash_kernels_loaded", + lambda _implementation: (object(), "flash_attn3"), + ) + observed: dict[str, object] = {} + + def varlen_forward(**kwargs): + observed.update(kwargs) + return kwargs["query_states"] + 1 + + monkeypatch.setattr(_core, "_kernels_flash_varlen_forward", varlen_forward) + monkeypatch.setattr( + _core, + "_kernels_flash_forward", + lambda **_kwargs: (_ for _ in ()).throw( + AssertionError("a masked causal call must not use dense FlashAttention") + ), + ) + X = torch.zeros(2, 4, 2, 8, dtype=torch.bfloat16) + mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]], dtype=torch.long) + + output = _core.kernels_flash_attention_func( + X, + X, + X, + attention_mask_2d=mask, + causal=True, + implementation="flash_attention_3", + ) + + assert observed["causal"] is True + assert observed["query_states"].shape == (5, 2, 8) + assert torch.equal(output[mask.bool()], torch.ones(5, 2, 8, dtype=torch.bfloat16)) + assert torch.count_nonzero(output[~mask.bool()]) == 0 + + +def test_masked_flash_validates_padding_mask_shape_before_kernel_loading( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + _core, + "_validate_kernels_flash_device", + lambda *_args: torch.device("cpu"), + ) + monkeypatch.setattr( + _core, + "_ensure_flash_kernels_loaded", + lambda _implementation: (_ for _ in ()).throw( + AssertionError("invalid masks must fail before kernel loading") + ), + ) + X = torch.zeros(2, 4, 2, 8, dtype=torch.bfloat16) + with pytest.raises(ValueError, match=r"expected \(2, 4\), received \(2, 3\)"): + _core.kernels_flash_attention_func( + X, + X, + X, + attention_mask_2d=torch.ones(2, 3, dtype=torch.bool), + causal=True, + implementation="flash_attention_2", + ) + + +def test_flash_dependency_profile_has_no_source_build_path() -> None: + root = Path(__file__).resolve().parents[2] + dependency_files = [ + path + for path in (root / "requirements").rglob("*") + if path.is_file() + ] + dependency_files.append(root / "docker" / "Dockerfile") + + for path in dependency_files: + text = path.read_text(encoding="utf-8") + assert "flash-attn" not in text + assert "flash_attn" not in text + assert "no-build-isolation-package" not in text + assert "extra-build-dependencies" not in text + + for path in (root / "src").rglob("*.py"): + text = path.read_text(encoding="utf-8") + assert 'import_module("flash_attn' not in text + assert "from flash_attn" not in text + assert "import flash_attn" not in text + + +def test_transformers_513_exposes_every_advertised_handler() -> None: + assert version("transformers") == "5.13.0" + validate_transformers_attention_interfaces() + for name in FUNCTION_BACKENDS: + assert name in FASTPLMS_ATTENTION_FUNCTIONS + assert callable(FASTPLMS_ATTENTION_FUNCTIONS[name]) + for name in MASK_BACKENDS: + assert name in FASTPLMS_ATTENTION_MASKS + assert callable(FASTPLMS_ATTENTION_MASKS[name]) + transformers_registry = AttentionInterface() + for name in ("flash_attention_2", "flash_attention_3"): + assert transformers_registry[name] is _TRANSFORMERS_FLASH_HANDLERS[name] + assert FASTPLMS_ATTENTION_FUNCTIONS[name] is not transformers_registry[name] + + +@pytest.mark.parametrize( + ("family_id", "relative_path", "class_name", "flash_backends"), + ( + ( + "esm2", + "src/fastplms/models/esm2/modeling_fastesm.py", + "FastEsmPreTrainedModel", + ("flash_attention_2", "flash_attention_3"), + ), + ( + "esm_plusplus", + "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py", + "PreTrainedESMplusplusModel", + ("flash_attention_2", "flash_attention_3"), + ), + ( + "dplm", + "src/fastplms/models/dplm/modeling_dplm.py", + "DPLMPreTrainedModel", + ("flash_attention_3",), + ), + ( + "dplm2", + "src/fastplms/models/dplm2/modeling_dplm2.py", + "DPLM2PreTrainedModel", + (), + ), + ("e1", "src/fastplms/models/e1/modeling_e1.py", "E1PreTrainedModel", ()), + ( + "ankh", + "src/fastplms/models/ankh/modeling_ankh.py", + "AnkhPreTrainedModel", + (), + ), + ( + "esm3", + "src/fastplms/models/esm3/modeling_esm3.py", + "FastESM3PreTrainedModel", + (), + ), + ( + "esmfold", + "src/fastplms/models/esmfold/modeling_fast_esmfold.py", + "FastEsmForProteinFolding", + (), + ), + ( + "esmfold2", + "src/fastplms/models/esmfold2/attention.py", + "ESMFold2AttentionMixin", + (), + ), + ), +) +def test_model_flash_flags_match_the_manifest( + family_id: str, + relative_path: str, + class_name: str, + flash_backends: tuple[str, ...], +) -> None: + root = Path(__file__).resolve().parents[2] + module = ast.parse((root / relative_path).read_text(encoding="utf-8")) + module_constants = { + target.id: ast.literal_eval(statement.value) + for statement in module.body + if isinstance(statement, ast.Assign) + for target in statement.targets + if isinstance(target, ast.Name) + and isinstance(statement.value, (ast.Constant, ast.Tuple, ast.List)) + } + class_node = next( + node for node in module.body if isinstance(node, ast.ClassDef) and node.name == class_name + ) + + def assignment_value(value: ast.expr) -> object: + if isinstance(value, ast.Name): + return module_constants[value.id] + return ast.literal_eval(value) + + assignments = { + target.id: assignment_value(statement.value) + for statement in class_node.body + if isinstance(statement, ast.Assign) + for target in statement.targets + if isinstance(target, ast.Name) + and target.id + in { + "_supports_flash_attn", + "_supports_flash_attn_2", + "_supports_flash_attn_3", + "_fastplms_attention_implementations", + } + } + assert assignments["_supports_flash_attn_2"] is ("flash_attention_2" in flash_backends) + assert assignments["_supports_flash_attn_3"] is ("flash_attention_3" in flash_backends) + assert assignments.get("_supports_flash_attn", False) is bool(flash_backends) + expected = get_model_registry().families[family_id].attention + assert assignments["_fastplms_attention_implementations"] == expected + assert tuple(name for name in expected if name.startswith("flash_")) == flash_backends + + +@pytest.mark.parametrize( + ("model_class", "config_class", "vocab_size", "unsupported"), + ( + ( + DPLMModel, + DPLMConfig, + 33, + ("flash_attention_2",), + ), + ( + DPLM2Model, + DPLM2Config, + 64, + ("eager", "flex_attention", "flash_attention_2", "flash_attention_3"), + ), + ), +) +def test_dplm_families_reject_unadvertised_attention( + model_class: type, + config_class: type, + vocab_size: int, + unsupported: tuple[str, ...], +) -> None: + config = config_class( + vocab_size=vocab_size, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=64, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=32, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + attn_backend="sdpa", + ) + model = model_class(config) + expected_flex_support = "flex_attention" not in unsupported + assert model._supports_flex_attn is expected_flex_support + + for implementation in unsupported: + with pytest.raises(ValueError, match="does not support"): + model.set_attn_implementation(implementation) + with pytest.raises(ValueError, match="does not support"): + model.attn_backend = implementation + + +@pytest.mark.parametrize( + "implementation", + ("flex_attention", "flash_attention_2", "flash_attention_3"), +) +def test_esm2_training_rejects_unsupported_attention_dropout( + implementation: str, +) -> None: + attention = EsmSelfAttention( + FastEsmConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.1, + attn_backend=implementation, + ) + ) + heads = torch.randn(1, 2, 3, 4) + + with pytest.raises(RuntimeError, match=r"inference-only.*dropout.*eager or SDPA"): + attention._attn(heads, heads, heads) + + +@pytest.mark.parametrize( + "implementation", + ("sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"), +) +def test_output_attentions_warns_when_falling_back_to_eager( + implementation: str, +) -> None: + with pytest.warns( + RuntimeWarning, + match=rf"output_attentions=True.*{implementation!r}.*using 'eager'", + ): + effective = _core.resolve_attention_backend_for_call( + implementation, + output_attentions=True, + ) + + assert effective == _core.AttentionBackend.EAGER + + +def test_output_attentions_does_not_warn_for_configured_eager() -> None: + with warnings.catch_warnings(): + warnings.simplefilter("error") + effective = _core.resolve_attention_backend_for_call( + "eager", + output_attentions=True, + ) + + assert effective == _core.AttentionBackend.EAGER + + +@pytest.mark.parametrize("training", (False, True)) +def test_dplm_sdpa_uses_attention_dropout_only_during_training( + monkeypatch: pytest.MonkeyPatch, + training: bool, +) -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.25, + position_embedding_type="absolute", + attn_backend="sdpa", + ) + ) + attention.train(training) + heads = torch.randn(1, 2, 3, 4) + observed: dict[str, float] = {} + + def fake_sdpa(*args, **kwargs): + observed["dropout_p"] = kwargs["dropout_p"] + return args[0] + + monkeypatch.setattr(dplm_module.F, "scaled_dot_product_attention", fake_sdpa) + attention._sdpa_attn(heads, heads, heads) + + assert observed["dropout_p"] == (0.25 if training else 0.0) + + +def test_dplm_manual_attention_applies_configured_training_dropout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.25, + position_embedding_type="absolute", + attn_backend="eager", + ) + ) + attention.train() + heads = torch.randn(1, 2, 3, 4) + observed: dict[str, object] = {} + + def fake_dropout(tensor, *, p, training): + observed.update(p=p, training=training) + return tensor + + monkeypatch.setattr(dplm_module.F, "dropout", fake_dropout) + attention._manual_attn(heads, heads, heads) + + assert observed == {"p": 0.25, "training": True} + + +@pytest.mark.parametrize( + "implementation", + ("flex_attention", "flash_attention_3"), +) +def test_dplm_training_rejects_unsupported_attention_dropout( + implementation: str, +) -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.1, + position_embedding_type="absolute", + attn_backend=implementation, + ) + ) + heads = torch.randn(1, 2, 3, 4) + + with pytest.raises(RuntimeError, match=r"inference-only.*dropout.*eager or SDPA"): + attention._attn(heads, heads, heads) + + +@pytest.mark.parametrize("implementation", ("eager", "sdpa")) +def test_dplm_cross_attention_executes_the_requested_supported_backend( + implementation: str, +) -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend=implementation, + ) + ).eval() + hidden_states = torch.randn(1, 3, 8) + encoder_hidden_states = torch.randn(1, 4, 8) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + output, weights, _ = attention( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + ) + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert weights is None + + +def test_dplm_eager_cross_attention_applies_additive_encoder_mask() -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend="eager", + ) + ).eval() + additive_mask = torch.tensor([[[[0.0, 0.0, -10_000.0, -10_000.0]]]]) + + output, weights, _ = attention( + torch.randn(1, 3, 8), + encoder_hidden_states=torch.randn(1, 4, 8), + encoder_attention_mask=additive_mask, + output_attentions=True, + ) + + assert torch.isfinite(output).all() + assert weights is not None + assert torch.isfinite(weights).all() + assert torch.equal(weights[..., 2:], torch.zeros_like(weights[..., 2:])) + + +@pytest.mark.parametrize("implementation", ("flex_attention", "flash_attention_3")) +def test_dplm_cross_attention_rejects_unimplemented_backends( + implementation: str, +) -> None: + attention = ModifiedEsmSelfAttention( + DPLMConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend=implementation, + ) + ).eval() + + with pytest.raises(RuntimeError, match=r"cross-attention.*Use eager or SDPA"): + attention( + torch.randn(1, 3, 8), + encoder_hidden_states=torch.randn(1, 4, 8), + ) + with pytest.raises(RuntimeError, match=r"cross-attention.*Use eager or SDPA"): + attention( + torch.randn(1, 3, 8), + encoder_hidden_states=torch.randn(1, 4, 8), + output_attentions=True, + ) + + +def test_esm2_config_normalizes_null_boundary_token_ids(tmp_path: Path) -> None: + config = FastEsmConfig(bos_token_id=None, eos_token_id=None) + + assert config.bos_token_id == 0 + assert config.eos_token_id == 2 + + (tmp_path / "config.json").write_text( + json.dumps({"model_type": "fast_esm", "bos_token_id": None, "eos_token_id": None}), + encoding="utf-8", + ) + reloaded = FastEsmConfig.from_pretrained(tmp_path, local_files_only=True) + assert reloaded.bos_token_id == 0 + assert reloaded.eos_token_id == 2 + + +def test_esm_family_base_automodels_do_not_create_untrained_poolers() -> None: + esm2_config = FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + position_embedding_type="absolute", + attn_backend="sdpa", + ) + dplm_config = DPLMConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + position_embedding_type="absolute", + attn_backend="eager", + ) + dplm2_config = DPLM2Config( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + position_embedding_type="absolute", + attn_backend="sdpa", + ) + + assert FastEsmModel(esm2_config).pooler is None + assert DPLMModel(dplm_config).pooler is None + assert DPLM2Model(dplm2_config).pooler is None + assert FastEsmModel(esm2_config, add_pooling_layer=True).pooler is not None + assert DPLMModel(dplm_config, add_pooling_layer=True).pooler is not None + assert DPLM2Model(dplm2_config, add_pooling_layer=True).pooler is not None + + +@pytest.mark.parametrize( + ("model_class", "config"), + [ + ( + FastEsmModel, + FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + pad_token_id=1, + mask_token_id=5, + position_embedding_type="absolute", + attn_backend="eager", + ), + ), + ( + DPLMModel, + DPLMConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + pad_token_id=1, + position_embedding_type="absolute", + attn_backend="eager", + ), + ), + ( + DPLM2Model, + DPLM2Config( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + pad_token_id=1, + position_embedding_type="absolute", + attn_backend="sdpa", + ), + ), + ], +) +def test_optional_esm_pooler_is_returned_and_round_trips( + model_class: type, + config: object, + tmp_path: Path, +) -> None: + model = model_class(config, add_pooling_layer=True).eval() + model.save_pretrained(tmp_path) + reloaded = model_class.from_pretrained(tmp_path).eval() + output = reloaded(input_ids=torch.tensor([[0, 3, 4, 2]])) + + assert reloaded.config.add_pooling_layer is True + assert reloaded.pooler is not None + assert output.pooler_output is not None + assert output.pooler_output.shape == (1, 8) + + +@pytest.mark.parametrize("should_raise", (False, True)) +def test_ankh_sdpa_never_mutates_process_global_reduction_policy( + monkeypatch: pytest.MonkeyPatch, + should_raise: bool, +) -> None: + attention = AnkhSelfAttention( + FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + attn_backend="sdpa", + ) + ) + query = torch.randn(1, 2, 3, 4) + mutations: list[bool] = [] + + monkeypatch.setattr( + torch.backends.cuda, + "allow_fp16_bf16_reduction_math_sdp", + lambda enabled: mutations.append(enabled), + ) + + def fake_sdpa(*args, **_kwargs): + if should_raise: + raise RuntimeError("forced SDPA failure") + return args[0] + + monkeypatch.setattr(ankh_module.F, "scaled_dot_product_attention", fake_sdpa) + if should_raise: + with pytest.raises(RuntimeError, match="forced SDPA failure"): + attention._sdpa_attn(query, query, query, None) + else: + output = attention._sdpa_attn(query, query, query, None) + assert output.shape == (1, 3, 8) + assert mutations == [] + + +def test_ankh_concurrent_fallback_and_sdpa_keep_backend_and_global_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = FastAnkhModel( + FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + attn_backend="sdpa", + ) + ).eval() + attention = model.encoder.block[0].layer[0].SelfAttention + configured_backend = attention.attn_backend + input_ids = torch.tensor(((2, 3, 1, 0), (4, 1, 0, 0))) + attention_mask = input_ids.ne(0) + rendezvous = Barrier(2) + global_policy_mutations: list[tuple[str, bool]] = [] + sdpa_calls: list[torch.Tensor] = [] + original_manual_attention = attention._manual_attn + + def synchronized_manual_attention(query, key, value, position_bias): + rendezvous.wait(timeout=3.0) + return original_manual_attention(query, key, value, position_bias) + + def synchronized_sdpa(query, _key, _value, **_kwargs): + sdpa_calls.append(query) + if len(sdpa_calls) == 1: + rendezvous.wait(timeout=3.0) + return query + + monkeypatch.setattr(attention, "_manual_attn", synchronized_manual_attention) + monkeypatch.setattr(ankh_module.F, "scaled_dot_product_attention", synchronized_sdpa) + monkeypatch.setattr( + torch.backends.cuda, + "allow_fp16_bf16_reduction_math_sdp", + lambda enabled: global_policy_mutations.append(("reduction_math", enabled)), + ) + for setter_name in ( + "_set_sdp_use_math", + "_set_sdp_use_flash", + "_set_sdp_use_mem_efficient", + "_set_sdp_use_cudnn", + "_set_sdp_use_overrideable", + ): + monkeypatch.setattr( + torch._C, + setter_name, + lambda enabled, name=setter_name: global_policy_mutations.append((name, enabled)), + raising=False, + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + with ThreadPoolExecutor(max_workers=2) as executor: + fallback_future = executor.submit( + model, + input_ids, + attention_mask=attention_mask, + output_attentions=True, + ) + sdpa_future = executor.submit( + model, + input_ids, + attention_mask=attention_mask, + ) + fallback_output = fallback_future.result(timeout=5.0) + sdpa_output = sdpa_future.result(timeout=5.0) + + fallback_warnings = [ + warning + for warning in captured + if issubclass(warning.category, RuntimeWarning) + and "output_attentions=True" in str(warning.message) + ] + assert len(fallback_warnings) == 1 + warning_message = str(fallback_warnings[0].message) + assert "requested 'sdpa'" in warning_message + assert "using 'eager'" in warning_message + assert "call only" in warning_message + assert model.attn_backend == "sdpa" + assert model.config.attn_backend == "sdpa" + assert model.config._attn_implementation == "sdpa" + assert model.encoder.attention_backend == configured_backend + assert attention.attn_backend == configured_backend + assert global_policy_mutations == [] + assert fallback_output.last_hidden_state.shape == (2, 4, 8) + assert fallback_output.attentions is not None + assert sdpa_output.last_hidden_state.shape == (2, 4, 8) + assert sdpa_output.attentions is None + + with warnings.catch_warnings(): + warnings.simplefilter("error") + post_fallback_output = model( + input_ids, + attention_mask=attention_mask, + output_attentions=False, + ) + assert post_fallback_output.attentions is None + assert len(sdpa_calls) == 2 + assert global_policy_mutations == [] + + +def test_ankh_rejects_unadvertised_flex_attention() -> None: + model = FastAnkhModel( + FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + attn_backend="sdpa", + ) + ) + with pytest.raises(ValueError, match="does not support 'flex_attention'"): + model.set_attn_implementation("flex_attention") + + +def test_attention_mixin_leaves_unspecified_backend_to_transformers() -> None: + observed: list[str | None] = [] + + class Base: + def __init__(self, config) -> None: + observed.append(config._attn_implementation) + config._attn_implementation_internal = "sdpa" + + class Model(FastPLMsAttentionMixin, Base): + _fastplms_attention_implementations = ("eager", "sdpa") + + config = SimpleNamespace( + _attn_implementation=None, + attn_backend=None, + ) + Model(config) + + assert observed == [None] + assert config._attn_implementation_internal == "sdpa" + assert config.attn_backend == "sdpa" + + +def test_attention_mixin_forwards_explicit_legacy_backend_to_transformers() -> None: + observed: list[str | None] = [] + + class Base: + def __init__(self, config) -> None: + observed.append(config._attn_implementation) + config._attn_implementation_internal = config._attn_implementation + + class Model(FastPLMsAttentionMixin, Base): + _fastplms_attention_implementations = ("eager", "sdpa", "flex_attention") + + config = SimpleNamespace( + _attn_implementation=None, + attn_backend="flex_attention", + ) + Model(config) + + assert observed == ["flex_attention"] + assert config._attn_implementation_internal == "flex_attention" + assert config.attn_backend == "flex_attention" + + +def test_attention_mixin_preserves_explicit_transformers_override() -> None: + observed: list[str | None] = [] + + class Base: + def __init__(self, config) -> None: + observed.append(config._attn_implementation) + + class Model(FastPLMsAttentionMixin, Base): + _fastplms_attention_implementations = ("eager", "sdpa", "flex_attention") + + config = SimpleNamespace( + _attn_implementation="eager", + _attn_implementation_internal="eager", + attn_backend="flex_attention", + ) + Model(config) + + assert observed == ["eager"] + assert config._attn_implementation_internal == "eager" + assert config.attn_backend == "eager" + + +def test_compiled_flex_cache_key_covers_execution_not_batch_contents( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = object() + compiled: list[object] = [] + + def fake_compile(function, *, dynamic): + assert function is source + assert dynamic is False + result = object() + compiled.append(result) + return result + + monkeypatch.setattr(_core, "flex_attention", source) + monkeypatch.setattr(torch, "compile", fake_compile) + monkeypatch.setattr( + torch.nn.attention.flex_attention, + "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", + False, + raising=False, + ) + _core._compiled_flex_attention.clear() + + base = { + "device": torch.device("cuda:0"), + "dtype": torch.bfloat16, + "shape": (2, 8, 64, 64), + "sequence_lengths": (64, 31), + "mask_semantics": "padding", + } + first = _core._get_flex_attention_fn(**base) + assert _core._get_flex_attention_fn(**base) is first + assert ( + _core._get_flex_attention_fn( + **{**base, "sequence_lengths": (64, 30)}, + ) + is first + ) + + variants = ( + {**base, "device": torch.device("cuda:1")}, + {**base, "dtype": torch.float32}, + {**base, "shape": (2, 8, 128, 64)}, + {**base, "mask_semantics": "chain_and_padding"}, + ) + values = [_core._get_flex_attention_fn(**variant) for variant in variants] + assert all(value is not first for value in values) + assert len({id(value) for value in values}) == len(values) + assert len(compiled) == 1 + len(variants) + + +def test_flex_block_mask_supports_disjoint_valid_spans_and_exact_cache_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created: list[tuple[object, object]] = [] + + def fake_create_block_mask(mask_mod, *args, **kwargs): + block_mask = object() + created.append((mask_mod, block_mask)) + return block_mask + + monkeypatch.setattr(_core, "flex_attention", object()) + monkeypatch.setattr(_core, "create_block_mask", fake_create_block_mask) + _core._flex_block_masks.clear() + first_pattern = torch.tensor(((1, 1, 0, 1), (1, 0, 1, 0)), dtype=torch.bool) + + _, _, first = _core.get_attention_mask( + _core.AttentionBackend.FLEX_ATTENTION, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + attention_mask=first_pattern, + dtype=torch.bfloat16, + ) + _, _, repeated = _core.get_attention_mask( + _core.AttentionBackend.FLEX_ATTENTION, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + attention_mask=first_pattern.clone(), + dtype=torch.bfloat16, + ) + assert repeated is first + assert len(created) == 1 + mask_mod = created[0][0] + for batch_index in range(2): + for query_index in range(4): + for key_index in range(4): + assert bool(mask_mod(batch_index, 0, query_index, key_index)) is bool( + first_pattern[batch_index, key_index] + ) + + second_pattern = torch.tensor(((1, 0, 1, 1), (0, 1, 0, 1)), dtype=torch.bool) + _, _, second = _core.get_attention_mask( + _core.AttentionBackend.FLEX_ATTENTION, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + attention_mask=second_pattern, + dtype=torch.bfloat16, + ) + assert second is not first + assert len(created) == 2 + + +def test_flex_block_mask_key_separates_equal_bytes_with_different_pattern_dtypes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created: list[object] = [] + + def fake_create_block_mask(mask_mod, *args, **kwargs): + del mask_mod, args, kwargs + block_mask = object() + created.append(block_mask) + return block_mask + + monkeypatch.setattr(_core, "create_block_mask", fake_create_block_mask) + _core.clear_flex_attention_caches() + boolean_pattern = torch.tensor(((True, False), (False, True))) + byte_pattern = boolean_pattern.to(dtype=torch.uint8) + assert boolean_pattern.view(torch.uint8).numpy().tobytes() == byte_pattern.numpy().tobytes() + + common = { + "batch_size": 2, + "query_length": 2, + "key_value_length": 2, + "device": torch.device("cpu"), + "dtype": torch.bfloat16, + "mask_semantics": "dtype-collision-contract", + "mask_mod": lambda batch_idx, head_idx, query_idx, key_idx: True, + } + boolean_mask = _core._get_flex_block_mask( + mask_pattern=boolean_pattern, + **common, + ) + repeated_boolean_mask = _core._get_flex_block_mask( + mask_pattern=boolean_pattern.clone(), + **common, + ) + byte_mask = _core._get_flex_block_mask( + mask_pattern=byte_pattern, + **common, + ) + + assert repeated_boolean_mask is boolean_mask + assert byte_mask is not boolean_mask + assert len(created) == 2 + + +def test_esmplusplus_flex_sequence_masks_share_exact_bounded_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created: list[object] = [] + + def fake_create_block_mask(mask_mod, *args, **kwargs): + del args, kwargs + block_mask = object() + created.append((mask_mod, block_mask)) + return block_mask + + monkeypatch.setattr(_core, "create_block_mask", fake_create_block_mask) + monkeypatch.setattr(_core, "_MAX_FLEX_CACHE_ENTRIES", 2) + _core.clear_flex_attention_caches() + stack = esmpp_module.TransformerStack( + d_model=8, + n_heads=2, + n_layers=1, + attn_backend="flex_attention", + ) + boolean_pattern = torch.tensor( + ((True, True, False, False), (True, False, True, False)) + ) + + *_, first = stack._sequence_id_attention_masks( + boolean_pattern, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + dtype=torch.bfloat16, + effective_backend=AttentionBackend.FLEX, + ) + *_, repeated = stack._sequence_id_attention_masks( + boolean_pattern.clone(), + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + dtype=torch.bfloat16, + effective_backend=AttentionBackend.FLEX, + ) + assert repeated is first + assert len(created) == 1 + + *_, different_dtype = stack._sequence_id_attention_masks( + boolean_pattern, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + dtype=torch.float32, + effective_backend=AttentionBackend.FLEX, + ) + assert different_dtype is not first + + chain_pattern = torch.tensor(((0, 0, -1, -1), (0, 1, 1, -1))) + *_, chain_mask = stack._sequence_id_attention_masks( + chain_pattern, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + dtype=torch.bfloat16, + effective_backend=AttentionBackend.FLEX, + ) + assert chain_mask is not first + assert len(created) == 3 + assert len(_core._flex_block_masks) == 2 + + *_, rebuilt = stack._sequence_id_attention_masks( + boolean_pattern, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + dtype=torch.bfloat16, + effective_backend=AttentionBackend.FLEX, + ) + assert rebuilt is not first + assert len(created) == 4 + assert len(_core._flex_block_masks) == 2 + + _core.clear_flex_attention_caches() + assert not _core._flex_block_masks diff --git a/tests/unit/test_attention_regressions.py b/tests/unit/test_attention_regressions.py new file mode 100644 index 0000000..ef7926d --- /dev/null +++ b/tests/unit/test_attention_regressions.py @@ -0,0 +1,1074 @@ +"""Focused CPU regressions for attention dispatch, masks, and caches.""" + +from __future__ import annotations + +import warnings +import pytest +import torch +from types import SimpleNamespace +from transformers.models.esm.configuration_esm import EsmConfig + +import fastplms.models.ankh.modeling_ankh as ankh_module +import fastplms.models.e1.modeling_e1 as e1_module +import fastplms.models.esm2.modeling_fastesm as esm2_module +import fastplms.models.esm3.modeling_esm3 as esm3_module +import fastplms.models.esm_plusplus.modeling_esm_plusplus as esmpp_module +from fastplms.attention import AttentionBackend, _core, _kernel_lock +from fastplms.models.ankh.modeling_ankh import ( + AnkhSelfAttention, + FastAnkhConfig, + FastAnkhForMaskedLMExtension, +) +from fastplms.models.dplm.modeling_dplm import ( + DPLMConfig, + ModifiedEsmEncoder, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + DPLM2Config, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + ModifiedEsmEncoder as DPLM2ModifiedEsmEncoder, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + ModifiedEsmSelfAttention as DPLM2ModifiedEsmSelfAttention, +) +from fastplms.models.e1.attention import build_block_causal_mask_4d +from fastplms.models.e1.modeling_e1 import ( + FAST_E1_ENCODER, + E1Config, +) +from fastplms.models.e1.modeling_e1 import ( + Attention as E1Attention, +) +from fastplms.models.e1.modeling_e1 import ( + AttentionArgs as E1AttentionArgs, +) +from fastplms.models.e1.modeling_e1 import ( + AttentionLayerType as E1AttentionLayerType, +) +from fastplms.models.esm2.modeling_fastesm import ( + EsmEncoder, + FastEsmConfig, + FastEsmPreTrainedModel, +) +from fastplms.models.esm2.modeling_fastesm import ( + EsmSelfAttention as FastEsmSelfAttention, +) +from fastplms.models.esm3.modeling_esm3 import FastESM3Config, FastESM3Model +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + MultiHeadAttention as ESMplusplusMultiHeadAttention, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + TransformerStack, +) +from fastplms.models.esmfold.modeling_fast_esmfold import ( + EsmSelfAttention as EsmFoldSelfAttention, +) +from fastplms.models.esmfold.modeling_fast_esmfold import ( + FastEsmEncoder as EsmFoldEncoder, +) +from fastplms.models.ttt import LoraInjectedLinear + + +@pytest.mark.parametrize( + "invalid_mask", + ( + torch.ones(2, 4, 1, dtype=torch.bool), + torch.ones(2, 3, dtype=torch.bool), + torch.ones(1, 4, dtype=torch.bool), + ), +) +def test_attention_masks_require_exact_batch_sequence_shape( + invalid_mask: torch.Tensor, +) -> None: + with pytest.raises(ValueError, match=r"attention_mask.*shape"): + _core.get_attention_mask( + AttentionBackend.EAGER, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + attention_mask=invalid_mask, + ) + + +@pytest.mark.parametrize("backend", tuple(AttentionBackend)) +def test_attention_masks_reject_rows_without_valid_keys_before_dispatch( + backend: AttentionBackend, +) -> None: + with pytest.raises( + ValueError, + match="attention_mask must keep at least one valid key per batch row", + ): + _core.get_attention_mask( + backend, + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + attention_mask=torch.tensor(((1, 1, 0, 0), (0, 0, 0, 0))), + ) + + +def test_esmplusplus_rejects_empty_attention_rows_without_fallback_or_mutation() -> None: + stack = TransformerStack(d_model=8, n_heads=2, n_layers=1, attn_backend="sdpa").eval() + configured_backend = stack.attention_backend + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + with pytest.raises( + ValueError, + match="attention_mask must keep at least one valid key per batch row", + ): + stack( + torch.randn(2, 4, 8), + attention_mask=torch.tensor(((1, 1, 0, 0), (0, 0, 0, 0))), + output_attentions=True, + ) + + assert captured == [] + assert stack.attention_backend == configured_backend + + +def _assert_masked_output_attentions_fallback(encoder: torch.nn.Module) -> None: + hidden_states = torch.randn(1, 4, 8) # (b=1, l=4, d=8) + attention_mask = torch.tensor([[1, 1, 0, 0]], dtype=torch.long) # (b, l) + configured_backend = encoder.attention_backend + backend_name = configured_backend.value + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + output = encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=True, + ) + + assert len(captured) == 1 + assert issubclass(captured[0].category, RuntimeWarning) + warning = str(captured[0].message) + assert "output_attentions=True" in warning + assert repr(backend_name) in warning + assert "using 'eager'" in warning + assert encoder.attention_backend == configured_backend + assert output.attentions is not None + assert len(output.attentions) == len(encoder.layer) + for attention_weights in output.attentions: + assert attention_weights is not None + assert torch.count_nonzero(attention_weights[..., 2:]) == 0 + + +def test_esm2_output_attentions_fallback_preserves_padding_mask() -> None: + encoder = EsmEncoder( + FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=2, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend="flex_attention", + ) + ).eval() + _assert_masked_output_attentions_fallback(encoder) + + +def test_dplm_output_attentions_fallback_preserves_padding_mask() -> None: + encoder = ModifiedEsmEncoder( + DPLMConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend="flash_attention_3", + ) + ).eval() + _assert_masked_output_attentions_fallback(encoder) + + +def test_dplm2_output_attentions_fallback_is_call_scoped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = DPLM2Config( + vocab_size=64, + hidden_size=8, + num_hidden_layers=2, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend="sdpa", + ) + encoder = DPLM2ModifiedEsmEncoder(config).eval() + hidden_states = torch.randn(2, 4, 8) # (b=2, l=4, d=8) + attention_mask = torch.tensor( # (b, l) + ((1, 1, 1, 0), (1, 1, 0, 0)), + dtype=torch.long, + ) + configured_backend = encoder.attention_backend + configured_layer_backends = tuple(layer.attention.self.attn_backend for layer in encoder.layer) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + fallback_output = encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=True, + ) + + assert len(captured) == 1 + assert issubclass(captured[0].category, RuntimeWarning) + warning = str(captured[0].message) + assert "output_attentions=True" in warning + assert "requested 'sdpa'" in warning + assert "using 'eager'" in warning + assert "call only" in warning + assert config.attn_backend == "sdpa" + assert encoder.attention_backend == configured_backend + assert ( + tuple(layer.attention.self.attn_backend for layer in encoder.layer) + == configured_layer_backends + ) + assert torch.isfinite(fallback_output.last_hidden_state).all() + assert fallback_output.attentions is not None + assert len(fallback_output.attentions) == 2 + invalid_keys = attention_mask[:, None, None, :].logical_not() + for attention_weights in fallback_output.attentions: + assert attention_weights is not None + assert attention_weights.shape == (2, 2, 4, 4) + assert torch.isfinite(attention_weights).all() + masked_weights = attention_weights.masked_select(invalid_keys.expand_as(attention_weights)) + assert torch.equal(masked_weights, torch.zeros_like(masked_weights)) + + dispatches: list[DPLM2ModifiedEsmSelfAttention] = [] + original_sdpa = DPLM2ModifiedEsmSelfAttention._sdpa_attn + + def record_sdpa(self, *args, **kwargs): + dispatches.append(self) + return original_sdpa(self, *args, **kwargs) + + monkeypatch.setattr(DPLM2ModifiedEsmSelfAttention, "_sdpa_attn", record_sdpa) + with warnings.catch_warnings(): + warnings.simplefilter("error") + subsequent_output = encoder( + hidden_states, + attention_mask=attention_mask, + output_attentions=False, + ) + + assert dispatches == [layer.attention.self for layer in encoder.layer] + assert torch.isfinite(subsequent_output.last_hidden_state).all() + assert subsequent_output.attentions is None + assert config.attn_backend == "sdpa" + assert encoder.attention_backend == configured_backend + assert ( + tuple(layer.attention.self.attn_backend for layer in encoder.layer) + == configured_layer_backends + ) + + +def test_e1_output_attentions_fallback_preserves_block_causal_mask_and_backend() -> None: + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_num_sequences=4, + max_num_positions_within_seq=16, + max_num_positions_global=32, + global_attention_every_n_layers=1, + attn_backend="flex_attention", + dtype="float32", + ) + attention = E1Attention(config, layer_idx=0).eval() + assert attention.layer_type == E1AttentionLayerType.GLOBAL + configured_backend = attention.attn_backend + sequence_ids = torch.tensor(((0, 0, 1, 1),), dtype=torch.long) # (b=1, l=4) + attention_args = E1AttentionArgs(block_causal_mask_4d=build_block_causal_mask_4d(sequence_ids)) + query = torch.randn(1, 4, 2, 4) # (b, l, h=2, d_h=4) + key = torch.randn(1, 4, 2, 4) # (b, l, h, d_h) + value = torch.randn(1, 4, 2, 4) # (b, l, h, d_h) + + with pytest.warns( + RuntimeWarning, + match=r"requested 'flex_attention'.*using 'eager'.*call only", + ) as fallback_warnings: + output, weights, s_max = attention._attn( + query, + key, + value, + sequence_ids=sequence_ids, + attention_args=attention_args, + output_attentions=True, + ) + + assert len(fallback_warnings) == 1 + assert attention.attn_backend == configured_backend + assert output.shape == (1, 4, 8) + assert torch.isfinite(output).all() + assert weights is not None + assert weights.shape == (1, 2, 4, 4) + assert torch.isfinite(weights).all() + assert torch.count_nonzero(weights[:, :, :2, 2:]) == 0 + assert torch.count_nonzero(weights[:, :, 2:, :2]) == weights[:, :, 2:, :2].numel() + assert s_max is None + + +def test_e1_public_fallback_warns_once_masks_padding_and_preserves_flex( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_num_sequences=4, + max_num_positions_within_seq=16, + max_num_positions_global=32, + global_attention_every_n_layers=2, + attn_backend="flex_attention", + dtype="float32", + ) + model = FAST_E1_ENCODER(config).eval() + configured_backend = model._attn_backend + configured_layer_backends = tuple( + layer.norm_attn_norm.self_attn.attn_backend for layer in model.layers + ) + inputs_embeds = torch.randn(2, 5, 8) # (b=2, l=5, d=8) + within_positions = torch.tensor( # (b, l) + ((0, 1, 0, 1, -1), (0, 1, 2, -1, -1)) + ) + global_positions = torch.tensor( # (b, l) + ((0, 1, 2, 3, -1), (0, 1, 2, -1, -1)) + ) + sequence_ids = torch.tensor( # (b, l) + ((0, 0, 1, 1, -1), (0, 0, 0, -1, -1)) + ) + + with pytest.warns( + RuntimeWarning, + match=r"requested 'flex_attention'.*using 'eager'.*call only", + ) as fallback_warnings: + fallback_output = model( + inputs_embeds=inputs_embeds, + within_seq_position_ids=within_positions, + global_position_ids=global_positions, + sequence_ids=sequence_ids, + output_attentions=True, + ) + + assert len(fallback_warnings) == 1 + assert model._attn_backend == configured_backend + assert config.attn_backend == "flex_attention" + assert ( + tuple(layer.norm_attn_norm.self_attn.attn_backend for layer in model.layers) + == configured_layer_backends + ) + assert torch.isfinite(fallback_output.last_hidden_state).all() + assert fallback_output.attentions is not None + expected_masks = ( + e1_module.build_within_seq_mask_4d(sequence_ids), + e1_module.build_block_causal_mask_4d(sequence_ids), + ) + for attention_weights, expected_mask in zip( + fallback_output.attentions, + expected_masks, + strict=True, + ): + assert torch.isfinite(attention_weights).all() + expanded_mask = expected_mask.expand_as(attention_weights) + masked_weights = attention_weights.masked_select(~expanded_mask) + assert torch.equal(masked_weights, torch.zeros_like(masked_weights)) + + within_block_mask = object() + global_block_mask = object() + flex_dispatches: list[dict[str, object]] = [] + monkeypatch.setattr( + e1_module, + "create_within_seq_block_mask", + lambda _sequence_ids: within_block_mask, + ) + monkeypatch.setattr( + e1_module, + "create_block_causal_mask_optimized", + lambda _sequence_ids: global_block_mask, + ) + + def fake_flex_attention(query, _key, _value, **kwargs): + flex_dispatches.append(kwargs) + return query + + monkeypatch.setattr(e1_module, "flex_attention_func", fake_flex_attention) + with warnings.catch_warnings(): + warnings.simplefilter("error") + optimized_output = model( + inputs_embeds=inputs_embeds, + within_seq_position_ids=within_positions, + global_position_ids=global_positions, + sequence_ids=sequence_ids, + output_attentions=False, + ) + + assert torch.isfinite(optimized_output.last_hidden_state).all() + assert [dispatch["block_mask"] for dispatch in flex_dispatches] == [ + within_block_mask, + global_block_mask, + ] + assert [dispatch["mask_semantics"] for dispatch in flex_dispatches] == [ + E1AttentionLayerType.WITHIN_SEQ.value, + E1AttentionLayerType.GLOBAL.value, + ] + assert model._attn_backend == configured_backend + assert ( + tuple(layer.norm_attn_norm.self_attn.attn_backend for layer in model.layers) + == configured_layer_backends + ) + + +def test_esm3_sdpa_fallback_is_call_scoped_and_preserves_padding_mask( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=2, + attn_backend="sdpa", + ) + ).eval() + input_ids = torch.tensor(((0, 3, 4, 2, 1), (0, 6, 2, 1, 1))) # (b=2, l=5) + attention_mask = input_ids.ne(1) # (b, l) + configured_stack_backend = model.esm3.transformer.attention_backend + configured_layer_backends = tuple( + block.attn.attn_backend for block in model.esm3.transformer.blocks + ) + + with pytest.warns( + RuntimeWarning, + match=r"requested 'sdpa'.*using 'eager'.*call only", + ) as fallback_warnings: + fallback_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=True, + ) + + assert len(fallback_warnings) == 1 + assert model.attn_backend == "sdpa" + assert model.config._attn_implementation == "sdpa" + assert model.esm3.transformer.attention_backend == configured_stack_backend + assert ( + tuple(block.attn.attn_backend for block in model.esm3.transformer.blocks) + == configured_layer_backends + ) + assert torch.isfinite(fallback_output.last_hidden_state).all() + assert fallback_output.attentions is not None + invalid_keys = ~attention_mask[:, None, None, :] + for attention_weights in fallback_output.attentions: + assert torch.isfinite(attention_weights).all() + masked_weights = attention_weights.masked_select(invalid_keys.expand_as(attention_weights)) + assert torch.equal(masked_weights, torch.zeros_like(masked_weights)) + + sdpa_masks: list[torch.Tensor] = [] + + def fake_sdpa(query, _key, _value, **kwargs): + sdpa_masks.append(kwargs["attn_mask"].detach().clone()) + return query + + monkeypatch.setattr(esm3_module.F, "scaled_dot_product_attention", fake_sdpa) + with warnings.catch_warnings(): + warnings.simplefilter("error") + optimized_output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=False, + ) + + assert torch.isfinite(optimized_output.last_hidden_state).all() + expected_padding_mask = attention_mask[:, None, None, :] + assert len(sdpa_masks) == 2 + for observed_mask in sdpa_masks: + assert torch.equal(observed_mask, expected_padding_mask) + assert model.attn_backend == "sdpa" + assert model.esm3.transformer.attention_backend == configured_stack_backend + assert ( + tuple(block.attn.attn_backend for block in model.esm3.transformer.blocks) + == configured_layer_backends + ) + + +def test_esm3_sequence_id_grouping_combines_with_public_padding_mask() -> None: + model = FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + ).eval() + input_ids = torch.tensor(((0, 3, 4, 2, 1), (0, 6, 2, 1, 1))) # (b=2, l=5) + attention_mask = input_ids.ne(1) # (b, l) + sequence_id = torch.tensor( # (b, l) + ((0, 0, 1, 1, -1), (0, 1, 1, -1, -1)) + ) + + output = model( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_attentions=True, + ) + + assert output.attentions is not None + expected_mask = ( + sequence_id.unsqueeze(-1).eq(sequence_id.unsqueeze(-2)).unsqueeze(1) + & attention_mask[:, None, None, :] + ) + attention_weights = output.attentions[0] + assert torch.isfinite(attention_weights).all() + masked_weights = attention_weights.masked_select(~expected_mask.expand_as(attention_weights)) + assert torch.equal(masked_weights, torch.zeros_like(masked_weights)) + assert torch.count_nonzero(attention_weights[0, :, :2, 2:4]) == 0 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + ( + ("attention_mask", torch.ones(2, 4), r"attention_mask must have shape"), + ("attention_mask", torch.tensor(((1, 1, 2), (1, 0, 0))), r"only boolean or 0/1"), + ("sequence_id", torch.zeros(1, 3, dtype=torch.long), r"sequence_id must have shape"), + ), +) +def test_esm3_rejects_malformed_padding_and_sequence_masks( + field: str, + value: torch.Tensor, + message: str, +) -> None: + model = FastESM3Model( + FastESM3Config( + hidden_size=8, + num_attention_heads=2, + num_vector_heads=2, + num_hidden_layers=1, + attn_backend="eager", + ) + ).eval() + inputs = { + "input_ids": torch.tensor(((0, 3, 2), (0, 4, 2))), + field: value, + } + + with pytest.raises(ValueError, match=message): + model(**inputs) + + +def test_dplm_sdpa_output_attentions_fallback_preserves_cross_attention_mask_and_backend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + encoder = ModifiedEsmEncoder( + DPLMConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + is_decoder=True, + add_cross_attention=True, + attn_backend="sdpa", + ) + ).eval() + cross_attention = encoder.layer[0].crossattention.self + configured_backend = encoder.attention_backend + additive_encoder_mask = torch.tensor( # (b=1, 1, 1, l_k=4) + [[[[0.0, 0.0, -10_000.0, -10_000.0]]]] + ) + observed: dict[str, torch.Tensor] = {} + original_manual_attention = cross_attention._manual_attn + + def record_cross_attention(query, key, value, attention_mask_4d=None, output_s_max=False): + result = original_manual_attention( + query, + key, + value, + attention_mask_4d, + output_s_max, + ) + observed["mask"] = attention_mask_4d.detach().clone() + observed["weights"] = result[1].detach().clone() + return result + + monkeypatch.setattr(cross_attention, "_manual_attn", record_cross_attention) + with pytest.warns( + RuntimeWarning, + match=r"requested 'sdpa'.*using 'eager'.*call only", + ) as fallback_warnings: + output = encoder( + torch.randn(1, 3, 8), + attention_mask=torch.ones(1, 3, dtype=torch.long), + encoder_hidden_states=torch.randn(1, 4, 8), + encoder_attention_mask=additive_encoder_mask, + output_attentions=True, + ) + + assert len(fallback_warnings) == 1 + assert encoder.attention_backend == configured_backend + assert encoder.layer[0].attention.self.attn_backend == configured_backend + assert cross_attention.attn_backend == configured_backend + assert torch.equal(observed["mask"], additive_encoder_mask) + assert torch.equal(observed["weights"][..., 2:], torch.zeros_like(observed["weights"][..., 2:])) + assert torch.isfinite(output.last_hidden_state).all() + + +def test_esmfold_output_attentions_fallback_preserves_padding_mask() -> None: + encoder = EsmFoldEncoder( + EsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + position_embedding_type="absolute", + attn_backend="flex_attention", + ) + ).eval() + _assert_masked_output_attentions_fallback(encoder) + + +def test_esmfold_flex_training_rejects_unimplemented_attention_dropout() -> None: + attention = EsmFoldSelfAttention( + EsmConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.1, + position_embedding_type="absolute", + attn_backend="flex_attention", + ) + ).train() + heads = torch.randn(1, 2, 3, 4) # (b=1, h=2, l=3, d_h=4) + + with pytest.raises(RuntimeError, match=r"inference-only.*dropout.*eager or SDPA"): + attention._attn(heads, heads, heads) + + +def test_esmplusplus_chain_masks_fail_closed_without_assertions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stack = TransformerStack(d_model=8, n_heads=2, n_layers=1, attn_backend="sdpa") + + with pytest.raises(ValueError, match=r"sequence_id must have shape \(2, 4\)"): + stack._sequence_id_attention_masks( + sequence_id=torch.ones(2, 4, 1, dtype=torch.bool), + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + ) + + stack.attention_backend = AttentionBackend.FLASH_ATTENTION_3 + with pytest.raises(ValueError, match="only supports boolean sequence_id"): + stack._sequence_id_attention_masks( + sequence_id=torch.zeros(2, 4, dtype=torch.long), + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + ) + + stack.attention_backend = AttentionBackend.FLEX_ATTENTION + monkeypatch.setattr(_core, "create_block_mask", None) + with pytest.raises(RuntimeError, match="create_block_mask is unavailable"): + stack._sequence_id_attention_masks( + sequence_id=torch.ones(2, 4, dtype=torch.bool), + batch_size=2, + seq_len=4, + device=torch.device("cpu"), + ) + + +def test_flash_attention_2_dense_and_varlen_use_autograd_wrappers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: list[dict[str, object]] = [] + + def low_level(*_args, **_kwargs): + raise AssertionError("low-level FlashAttention entry points must not be called") + + def dense(**kwargs): + observed.append(kwargs) + return kwargs["q"] + kwargs["k"] + kwargs["v"] + + def varlen(**kwargs): + observed.append(kwargs) + return kwargs["q"] + kwargs["k"] + kwargs["v"] + + kernel = SimpleNamespace( + fwd=low_level, + varlen_fwd=low_level, + flash_attn_func=dense, + flash_attn_varlen_func=varlen, + ) + monkeypatch.setattr( + _core, + "_ensure_flash_kernels_loaded", + lambda _implementation: (kernel, "flash_attn2"), + ) + query = torch.randn(1, 3, 2, 4, requires_grad=True) # (b=1, l=3, h=2, d_h=4) + key = torch.randn(1, 3, 2, 4, requires_grad=True) # (b, l, h, d_h) + value = torch.randn(1, 3, 2, 4, requires_grad=True) # (b, l, h, d_h) + + dense_output = _core._kernels_flash_forward( + query, + key, + value, + implementation="flash_attention_2", + ) + dense_output.sum().backward() + assert all(tensor.grad is not None for tensor in (query, key, value)) + assert observed[-1]["dropout_p"] == 0.0 + + flat_query = query.detach().reshape(3, 2, 4).requires_grad_() + flat_key = key.detach().reshape(3, 2, 4).requires_grad_() + flat_value = value.detach().reshape(3, 2, 4).requires_grad_() + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) # (b + 1,) + varlen_output = _core._kernels_flash_varlen_forward( + flat_query, + flat_key, + flat_value, + cu_seqlens, + cu_seqlens, + 3, + 3, + implementation="flash_attention_2", + ) + varlen_output.sum().backward() + assert all(tensor.grad is not None for tensor in (flat_query, flat_key, flat_value)) + assert observed[-1]["dropout_p"] == 0.0 + + +def test_flash_attention_2_dense_and_varlen_preserve_lora_and_input_gradients( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dispatch_count = {"dense": 0, "varlen": 0} + + def dense(**kwargs): + dispatch_count["dense"] += 1 + return kwargs["q"] + kwargs["k"] + kwargs["v"] + + def varlen(**kwargs): + dispatch_count["varlen"] += 1 + return kwargs["q"] + kwargs["k"] + kwargs["v"] + + kernel = SimpleNamespace( + flash_attn_func=dense, + flash_attn_varlen_func=varlen, + ) + monkeypatch.setattr( + _core, + "_ensure_flash_kernels_loaded", + lambda _implementation: (kernel, "flash_attn2"), + ) + monkeypatch.setattr( + _core, + "_validate_kernels_flash_device", + lambda query, _key, _value, _implementation: query.device, + ) + monkeypatch.setattr( + _core, + "_validate_kernels_flash_dtype", + lambda query, _key, _value, _implementation: query.dtype, + ) + + projections = torch.nn.ModuleList( + [ + LoraInjectedLinear( + torch.nn.Linear(8, 8, bias=False), + rank=2, + alpha=1.0, + ) + for _ in range(3) + ] + ) + down_weights = torch.arange(1, 17, dtype=torch.float32).reshape(2, 8) / 16 # (r=2, d=8) + with torch.no_grad(): + for projection in projections: + projection.linear.weight.copy_(torch.eye(8)) + projection.lora_down.weight.copy_(down_weights) + projection.lora_up.weight.fill_(0.05) + + def project(hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]: + return tuple(projection(hidden_states).reshape(2, 4, 2, 4) for projection in projections) + + dense_input = (torch.arange(1, 65, dtype=torch.float32).reshape(2, 4, 8) / 64).requires_grad_() + padded_input = ( + torch.arange(65, 129, dtype=torch.float32).reshape(2, 4, 8) / 128 + ).requires_grad_() + dense_output = _core.kernels_flash_attention_func( + *project(dense_input), + implementation="flash_attention_2", + ) + attention_mask = torch.tensor( + [[True, True, True, True], [True, True, False, False]], + ) + padded_output = _core.kernels_flash_attention_func( + *project(padded_input), + attention_mask_2d=attention_mask, + implementation="flash_attention_2", + ) + + (dense_output.sum() + padded_output.sum()).backward() + + assert dispatch_count == {"dense": 1, "varlen": 1} + for hidden_states in (dense_input, padded_input): + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() + assert torch.count_nonzero(hidden_states.grad) > 0 + assert torch.count_nonzero(padded_input.grad[1, 2:]) == 0 + for projection in projections: + assert projection.linear.weight.grad is None + for parameter in (projection.lora_down.weight, projection.lora_up.weight): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + assert torch.count_nonzero(parameter.grad) > 0 + + +def test_flash_attention_2_rejects_low_level_only_kernel_artifact( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = SimpleNamespace( + fwd=lambda **kwargs: kwargs["q"], + varlen_fwd=lambda **kwargs: kwargs["q"], + ) + monkeypatch.setattr(_core, "load_locked_kernel", lambda *_args: kernel) + + with pytest.raises(RuntimeError, match="autograd-enabled flash_attn_func"): + _core._load_kernels_flash("flash_attention_2") + + +@pytest.mark.parametrize("varlen", (False, True)) +def test_flash_attention_3_preserves_internal_type_errors( + monkeypatch: pytest.MonkeyPatch, + varlen: bool, +) -> None: + failure = TypeError("internal FlashAttention kernel failure") + calls = 0 + + def fail(**_kwargs): + nonlocal calls + calls += 1 + raise failure + + kernel = SimpleNamespace( + flash_attn_func=fail, + flash_attn_varlen_func=fail, + ) + monkeypatch.setattr( + _core, + "_ensure_flash_kernels_loaded", + lambda _implementation: (kernel, "flash_attn3"), + ) + query = torch.randn(1, 3, 2, 4) # (b=1, l=3, h=2, d_h=4) + + with pytest.raises(TypeError) as captured: + if varlen: + flat_query = query.reshape(3, 2, 4) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) # (b + 1,) + _core._kernels_flash_varlen_forward( + flat_query, + flat_query, + flat_query, + cu_seqlens, + cu_seqlens, + 3, + 3, + implementation="flash_attention_3", + ) + else: + _core._kernels_flash_forward( + query, + query, + query, + implementation="flash_attention_3", + ) + + assert captured.value is failure + assert calls == 1 + + +@pytest.mark.parametrize("variable", ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")) +def test_kernel_loader_honors_all_offline_environment_variables( + monkeypatch: pytest.MonkeyPatch, + variable: str, +) -> None: + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.setenv(variable, "yes") + + assert _kernel_lock._offline_mode() + + +def test_public_flex_cache_cleanup_is_scoped_and_complete() -> None: + _core._compiled_flex_attention[("compiled",)] = object() + _core._flex_block_masks[("mask",)] = object() + + _core.clear_flex_attention_caches() + + assert not _core._compiled_flex_attention + assert not _core._flex_block_masks + + +def _ankh_attention(dropout_rate: float) -> AnkhSelfAttention: + config = FastAnkhConfig( + vocab_size=16, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + dropout_rate=dropout_rate, + attn_backend="sdpa", + ) + return AnkhSelfAttention(config) + + +@pytest.mark.parametrize( + ("training", "expected_calls"), + ((True, [(0.25, True)]), (False, [])), +) +def test_ankh_output_attentions_eager_fallback_honors_attention_dropout( + monkeypatch: pytest.MonkeyPatch, + training: bool, + expected_calls: list[tuple[float, bool]], +) -> None: + attention = _ankh_attention(dropout_rate=0.25).train(training) + attention.attn_backend = AttentionBackend.SDPA + dropout_calls: list[tuple[float, bool]] = [] + + def record_dropout( + tensor: torch.Tensor, + *, + p: float, + training: bool, + ) -> torch.Tensor: + dropout_calls.append((p, training)) + return tensor + + monkeypatch.setattr(ankh_module.F, "dropout", record_dropout) + hidden_states = torch.randn(1, 3, 8) # (b=1, l=3, d=8) + + with pytest.warns(RuntimeWarning, match="output_attentions=True"): + output, attention_weights, _ = attention( + hidden_states, + output_attentions=True, + ) + + assert output.shape == hidden_states.shape + assert attention_weights is not None + assert dropout_calls == expected_calls + + +@pytest.mark.parametrize( + ("training", "expected_dropout"), + ((True, 0.25), (False, 0.0)), +) +def test_ankh_sdpa_receives_training_attention_dropout( + monkeypatch: pytest.MonkeyPatch, + training: bool, + expected_dropout: float, +) -> None: + attention = _ankh_attention(dropout_rate=0.25).train(training) + observed_dropout: list[float] = [] + + def record_sdpa( + query: torch.Tensor, + _key: torch.Tensor, + _value: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + observed_dropout.append(kwargs["dropout_p"]) + return query + + monkeypatch.setattr(ankh_module.F, "scaled_dot_product_attention", record_sdpa) + heads = torch.randn(1, 2, 3, 4) # (b=1, h=2, l=3, d_h=4) + + output = attention._sdpa_attn(heads, heads, heads, None) + + assert output.shape == (1, 3, 8) + assert observed_dropout == [expected_dropout] + + +def test_ankh_ttt_missing_input_uses_optimization_safe_validation() -> None: + model = object.__new__(FastAnkhForMaskedLMExtension) + + with pytest.raises(ValueError, match="either seq or input_ids"): + model._ttt_tokenize() + + +def test_esm2_attention_validates_head_divisibility_without_assertions() -> None: + config = FastEsmConfig( + hidden_size=10, + num_attention_heads=3, + attention_probs_dropout_prob=0.0, + attn_backend="eager", + ) + + with pytest.raises(ValueError, match="not a multiple"): + FastEsmSelfAttention(config) + + +def test_esm2_legacy_backend_setter_uses_explicit_validation() -> None: + model = FastEsmPreTrainedModel(FastEsmConfig(attn_backend="sdpa")) + + with pytest.raises(ValueError, match="does not support"): + model.attn_backend = "not_an_attention_backend" + + +def test_bool_to_additive_mask_rejects_non_boolean_input_explicitly() -> None: + with pytest.raises(TypeError, match="requires a bool tensor"): + _core.bool_to_additive_mask(torch.ones(1, 2), torch.float32) + + +def test_esm2_flex_unavailability_raises_explicit_runtime_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attention = FastEsmSelfAttention( + FastEsmConfig( + hidden_size=8, + num_attention_heads=2, + attention_probs_dropout_prob=0.0, + attn_backend="sdpa", + ) + ) + attention.attn_backend = AttentionBackend.FLEX_ATTENTION + monkeypatch.setattr(esm2_module, "flex_attention", None) + heads = torch.randn(1, 2, 3, 4) # (b=1, h=2, l=3, d_h=4) + + with pytest.raises(RuntimeError, match="Flex attention is not available"): + attention._flex_attn(heads, heads, heads) + + +def test_esmplusplus_flex_unavailability_raises_explicit_runtime_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attention = ESMplusplusMultiHeadAttention( + d_model=8, + n_heads=2, + attn_backend="sdpa", + ) + attention.attn_backend = AttentionBackend.FLEX_ATTENTION + monkeypatch.setattr(esmpp_module, "flex_attention", None) + heads = torch.randn(1, 2, 3, 4) # (b=1, h=2, l=3, d_h=4) + + with pytest.raises(RuntimeError, match="Flex attention is not available"): + attention._flex_attn(heads, heads, heads) diff --git a/tests/unit/test_benchmark_execution.py b/tests/unit/test_benchmark_execution.py new file mode 100644 index 0000000..8b7c180 --- /dev/null +++ b/tests/unit/test_benchmark_execution.py @@ -0,0 +1,322 @@ +"""Short H100 execution smoke for the external benchmark harness.""" + +from __future__ import annotations + +import pytest +import torch +import transformers +from pathlib import Path +from types import SimpleNamespace + +from benchmarks.run import ( + _load_model, + _prepare_esmfold2_inputs, + _run_esmfold2_esmc_projection, + cuda_sample_ms, + measure_blocks, + prepare_inputs, + warm_until_stable, +) +from fastplms.registry import get_model_registry + + +def test_prepare_inputs_counts_residues_not_special_tokens() -> None: + """Keep logical throughput independent of tokenizer control tokens.""" + + class FakeTokenizer: + def __call__( + self, + sequences: list[str], + *, + return_tensors: str, + padding: str, + max_length: int, + truncation: bool, + ) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + assert padding == "max_length" + assert max_length == 8 + assert truncation + assert [len(sequence) for sequence in sequences] == [6, 3] + + input_ids = torch.zeros((2, max_length), dtype=torch.long) # (b=2, l=8) + attention_mask = torch.zeros_like(input_ids) # (b=2, l=8) + # Both sequences receive BOS and EOS control tokens. + attention_mask[0, :8] = 1 # (l=8,) + attention_mask[1, :5] = 1 # (l=5,) + return {"input_ids": input_ids, "attention_mask": attention_mask} + + class FakeModel: + tokenizer = FakeTokenizer() + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + # input_ids: (b, l); attention_mask: (b, l) + del attention_mask + return input_ids # (b, l) + + model_inputs, logical_tokens, padded_tokens, sequences = prepare_inputs( + FakeModel(), + "unused-local-model", + (8, 5), + torch.device("cpu"), + revision=None, + local_files_only=True, + ) + + assert [len(sequence) for sequence in sequences] == [6, 3] + assert model_inputs["attention_mask"].sum().item() == 13 + assert logical_tokens == 9 + assert padded_tokens == 16 + + +def test_local_artifact_model_load_omits_hub_revision_and_keeps_registry_policy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esm2_8m"] + artifact = tmp_path / "ESM2-8M" + artifact.mkdir() + calls: list[tuple[object, dict[str, object]]] = [] + + class FakeModel: + def eval(self) -> FakeModel: + return self + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, source: object, **kwargs: object) -> FakeModel: + calls.append((source, kwargs)) + return FakeModel() + + monkeypatch.setattr(transformers, "AutoModelForMaskedLM", FakeAutoModel) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + arguments = SimpleNamespace( + model=spec.fast.repo_id, + revision=spec.fast.revision, + load_model=artifact, + load_revision=None, + auto_class="AutoModelForMaskedLM", + backend="sdpa", + precision="bf16", + bf16_execution=spec.family.bf16_execution, + mode="steady", + local_files_only=True, + esmc_load_model=None, + ) + + _load_model(arguments, torch) + + assert calls[0][0] == artifact + assert "revision" not in calls[0][1] + assert calls[0][1]["local_files_only"] is True + expected_dtype = ( + torch.float32 + if spec.family.bf16_execution == "fp32_parameters_autocast" + else torch.bfloat16 + ) + assert calls[0][1]["dtype"] == expected_dtype + + +def test_local_artifact_tokenizer_load_omits_hub_revision( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "ESM2-8M" + artifact.mkdir() + calls: list[tuple[object, dict[str, object]]] = [] + + class FakeTokenizer: + def __call__(self, sequences: list[str], **kwargs: object) -> dict[str, torch.Tensor]: + del sequences, kwargs + return { + "input_ids": torch.zeros((1, 4), dtype=torch.long), # (b=1, l=4) + "attention_mask": torch.ones((1, 4), dtype=torch.long), # (b=1, l=4) + } + + class FakeAutoTokenizer: + @classmethod + def from_pretrained(cls, source: object, **kwargs: object) -> FakeTokenizer: + calls.append((source, kwargs)) + return FakeTokenizer() + + class FakeModel: + tokenizer = None + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + # input_ids: (b, l); attention_mask: (b, l) + del attention_mask + return input_ids # (b, l) + + monkeypatch.setattr(transformers, "AutoTokenizer", FakeAutoTokenizer) + prepare_inputs( + FakeModel(), + artifact, + (4,), + torch.device("cpu"), + revision=None, + local_files_only=True, + ) + + assert calls == [ + ( + artifact, + {"trust_remote_code": True, "local_files_only": True}, + ) + ] + + +def test_local_esmfold2_load_uses_validated_local_esmc_dependency( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = get_model_registry()["esmfold2"] + artifact = tmp_path / "ESMFold2" + esmc_artifact = tmp_path / "ESMC-6B" + artifact.mkdir() + esmc_artifact.mkdir() + top_level_calls: list[tuple[object, dict[str, object]]] = [] + esmc_calls: list[tuple[str, dict[str, object]]] = [] + + class FakeModel: + def eval(self) -> FakeModel: + return self + + def load_esmc(self, source: str, **kwargs: object) -> None: + esmc_calls.append((source, kwargs)) + + class FakeAutoModel: + @classmethod + def from_pretrained(cls, source: object, **kwargs: object) -> FakeModel: + top_level_calls.append((source, kwargs)) + return FakeModel() + + monkeypatch.setattr(transformers, "AutoModel", FakeAutoModel) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + arguments = SimpleNamespace( + model=spec.fast.repo_id, + revision=spec.fast.revision, + load_model=artifact, + load_revision=None, + auto_class="AutoModel", + backend="sdpa", + precision="bf16", + bf16_execution=spec.family.bf16_execution, + mode="esmfold2_embed", + local_files_only=True, + esmc_load_model=esmc_artifact, + ) + + _load_model(arguments, torch) + + assert top_level_calls[0][0] == artifact + assert top_level_calls[0][1]["load_esmc"] is False + assert "revision" not in top_level_calls[0][1] + assert esmc_calls == [ + ( + str(esmc_artifact), + { + "precision": "bf16", + "device": torch.device("cuda"), + "local_files_only": True, + }, + ) + ] + + +@pytest.mark.benchmark +@pytest.mark.gpu +def test_cuda_event_benchmark_smoke() -> None: + """Exercise event timing and block accounting without benchmarking a model.""" + + assert torch.cuda.is_available() + X = torch.randn((64, 64), device="cuda", dtype=torch.bfloat16) # (n=64, d=64) + W = torch.randn((64, 64), device="cuda", dtype=torch.bfloat16) # (d=64, d=64) + + def operation() -> torch.Tensor: + return X @ W # (n=64, d=64) + + warmup = warm_until_stable( + torch, + operation, + window=2, + tolerance=1.0, + minimum_samples=4, + maximum_samples=20, + ) + blocks = measure_blocks( + torch, + operation, + logical_tokens_per_forward=64, + padded_tokens_per_forward=64, + blocks=1, + minimum_block_ms=1.0, + minimum_forwards=2, + ) + + assert len(warmup) >= 4 + assert len(blocks) == 1 + assert blocks[0].forwards >= 2 + assert blocks[0].logical_tokens_per_second > 0 + assert blocks[0].padded_tokens_per_second > 0 + assert all(sample > 0 for sample in blocks[0].samples_ms) + + +@pytest.mark.benchmark +@pytest.mark.gpu +def test_esmfold2_esmc_projection_path_smoke() -> None: + """Exercise residue preparation and the complete representation operation.""" + + assert torch.cuda.is_available() + + class FakeESMFold2: + def _compute_lm_hidden_states( + self, + input_ids: torch.Tensor, + asym_id: torch.Tensor, + residue_index: torch.Tensor, + mol_type: torch.Tensor, + residue_mask: torch.Tensor, + ) -> torch.Tensor: + # Each input tensor has shape (b, l). + del asym_id, residue_index, mol_type + H = input_ids.to(torch.bfloat16)[..., None, None].expand( # (b, l, 81, d=4) + -1, -1, 81, 4 + ) + return H * residue_mask[..., None, None] # (b, l, 81, d=4) + + def project_esmc_hidden_states( + self, + hidden_states: torch.Tensor, + residue_mask: torch.Tensor, + ) -> torch.Tensor: + # hidden_states: (b, l, 81, d); residue_mask: (b, l) + Z = hidden_states.mean(dim=2) # (b, l, d) + return Z * residue_mask[..., None] # (b, l, d) + + model_inputs, logical_tokens, padded_tokens = _prepare_esmfold2_inputs( + torch, (7, 3) + ) # tensor values: (b=2, l=7) + assert logical_tokens == 10 + assert padded_tokens == 14 + assert model_inputs["residue_mask"].sum().item() == 10 + assert not model_inputs["residue_mask"][1, 3:].any() + + model = FakeESMFold2() + + def operation() -> torch.Tensor: + with torch.inference_mode(): + return _run_esmfold2_esmc_projection(model, model_inputs) # (b=2, l=7, d=4) + + elapsed_ms = cuda_sample_ms(torch, operation) + Z = operation() # (b=2, l=7, d=4) + assert elapsed_ms >= 0.0 + assert Z.shape == (2, 7, 4) + assert torch.count_nonzero(Z[1, 3:]) == 0 diff --git a/tests/unit/test_benchmark_matrix.py b/tests/unit/test_benchmark_matrix.py new file mode 100644 index 0000000..f9f62f5 --- /dev/null +++ b/tests/unit/test_benchmark_matrix.py @@ -0,0 +1,540 @@ +"""Manifest-derived benchmark matrix tests.""" + +from __future__ import annotations + +import json +import pytest +from pathlib import Path +from types import SimpleNamespace + +import benchmarks.suite as benchmark_suite +from benchmarks.run import ( + _benchmark_load_dtype, + _resolve_bf16_execution, + _uses_bf16_autocast, + run_case, + validate_hopper_sm90_environment, +) +from benchmarks.suite import ( + ESMFOLD2_DEDICATED_MODE, + ESMFOLD2_REPRESENTATION_PROFILE, + FIXED_SHAPES, + FLASH_BACKEND_CLAIM_INELIGIBILITY_REASON, + FLASH_BACKEND_HISTORICAL_EVIDENCE, + SEQUENCE_FORWARD_PROFILE, + STRUCTURE_DEDICATED_MODE, + STRUCTURE_STARTUP_PROFILE, + benchmark_auto_class, + benchmark_cases, + benchmark_model_key, + bind_local_artifacts, + build_parser, + exhaustive_benchmark_cases, +) +from fastplms.registry import ModelSpec, get_model_registry + + +_RUNTIME_REVISION = "a" * 40 +_SOURCE_SHA256 = "b" * 64 +_RUNTIME_BUNDLE_SHA256 = "c" * 64 +_STATE_SHA256 = "d" * 64 + + +def _write_benchmark_artifact( + root: Path, + spec: ModelSpec, + *, + runtime_revision: str = _RUNTIME_REVISION, + source_tree_sha256: str = _SOURCE_SHA256, + config_updates: dict[str, object] | None = None, +) -> Path: + path = root / spec.fast.repo_id.rsplit("/", maxsplit=1)[1] + path.mkdir(parents=True) + config: dict[str, object] = { + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "fastplms_checkpoint_revision": spec.artifact_checkpoint.revision, + "fastplms_weights_revision": spec.artifact_checkpoint.revision, + "fastplms_runtime_revision": runtime_revision, + "fastplms_source_tree_sha256": source_tree_sha256, + "fastplms_runtime_bundle_sha256": _RUNTIME_BUNDLE_SHA256, + } + if config_updates: + config.update(config_updates) + provenance = { + "model_id": spec.id, + "artifact_checkpoint": { + "repo_id": spec.artifact_checkpoint.repo_id, + "revision": spec.artifact_checkpoint.revision, + }, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": runtime_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": _RUNTIME_BUNDLE_SHA256, + "canonical_weights": { + "state_digest": { + "schema_version": 1, + "algorithm": "sha256", + "sha256": _STATE_SHA256, + } + }, + } + (path / "config.json").write_text(json.dumps(config), encoding="utf-8") + (path / "provenance.json").write_text(json.dumps(provenance), encoding="utf-8") + (path / "artifact-manifest.json").write_text( + json.dumps({"config.json": "sha256:" + "e" * 64}), + encoding="utf-8", + ) + return path + + +def _stub_artifact_validation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(benchmark_suite, "_validate_built_artifact", lambda *_args: None) + monkeypatch.setattr( + benchmark_suite, + "_frozen_runtime_identity", + lambda *_args: (_RUNTIME_REVISION, _SOURCE_SHA256), + ) + + +def _shape(case: SimpleNamespace) -> tuple[int, int, tuple[int, ...]]: + return case.batch_size, case.sequence_length, tuple(case.lengths) + + +def test_full_matrix_covers_fixed_shapes_for_each_sequence_backend() -> None: + cases = list(benchmark_cases(family=None, quick=False, local_files_only=True)) + registry = get_model_registry() + representatives = [ + spec + for spec in registry.values() + if spec.is_deep_reference + and spec.family.tokenizer_mode != "structure" + and "benchmark" in spec.family.test_tiers + ] + assert {spec.family.id for spec in representatives} == { + "ankh", + "dplm", + "dplm2", + "e1", + "esm2", + "esm3", + "esm_plusplus", + } + for spec in representatives: + for backend in spec.family.attention: + matching = [ + case + for case in cases + if case.model == spec.fast.repo_id + and case.backend == backend + and case.mode == "steady" + ] + assert {_shape(case) for case in matching} == set(FIXED_SHAPES) + assert all(case.bf16_execution == spec.family.bf16_execution for case in matching) + compile_cases = [ + case + for case in cases + if case.model == spec.fast.repo_id + and case.backend == backend + and case.mode == "compile" + ] + assert len(compile_cases) == 1 + assert _shape(compile_cases[0]) == (1, 512, ()) + startup = [ + case for case in cases if case.model == spec.fast.repo_id and case.mode == "startup" + ] + embedding = [ + case for case in cases if case.model == spec.fast.repo_id and case.mode == "embed" + ] + assert len(startup) == len(embedding) == 1 + assert startup[0].suite_profile == SEQUENCE_FORWARD_PROFILE + assert embedding[0].suite_profile == SEQUENCE_FORWARD_PROFILE + + +def test_gh200_matrix_explicitly_selects_eager_sdpa_and_flex_only() -> None: + selected = ("eager", "sdpa", "flex_attention") + cases = list( + benchmark_cases( + family=None, + quick=False, + local_files_only=True, + backends=selected, + ) + ) + + assert {case.backend for case in cases}.issubset(set(selected)) + assert "flash_attention_2" not in {case.backend for case in cases} + assert "flash_attention_3" not in {case.backend for case in cases} + for spec in get_model_registry().values(): + if not spec.is_deep_reference or "benchmark" not in spec.family.test_tiers: + continue + expected = set(spec.family.attention).intersection(selected) + measured = { + case.backend + for case in cases + if case.model == spec.fast.repo_id and case.mode in {"compile", "steady"} + } + if spec.family.tokenizer_mode != "structure": + assert measured == expected + + +def test_flash_benchmark_cases_are_never_claim_eligible() -> None: + cases = list(benchmark_cases(family=None, quick=False, local_files_only=True)) + flash_cases = [case for case in cases if case.backend in FLASH_BACKEND_HISTORICAL_EVIDENCE] + + assert {case.backend for case in flash_cases} == set(FLASH_BACKEND_HISTORICAL_EVIDENCE) + assert all(not case.claim_eligible for case in flash_cases) + assert all( + case.claim_eligibility_reason == FLASH_BACKEND_CLAIM_INELIGIBILITY_REASON + for case in flash_cases + ) + for case in flash_cases: + assert case.historical_evidence == FLASH_BACKEND_HISTORICAL_EVIDENCE[case.backend] + + current_cases = [ + case + for case in cases + if case.backend in {"eager", "sdpa", "flex_attention"} and case.claim_eligible + ] + assert current_cases + assert all( + case.historical_evidence == "current_release_execution" for case in current_cases + ) + + +def test_esmfold2_matrix_separates_projection_from_esmc_precision() -> None: + cases = list(benchmark_cases(family="esmfold2", quick=False, local_files_only=True)) + expected = {spec.fast.repo_id for spec in get_model_registry().by_family("esmfold2")} + assert {case.model for case in cases} == expected + for model_id in expected: + model_cases = [case for case in cases if case.model == model_id] + projection = [case for case in model_cases if case.mode == "projection"] + assert {case.precision for case in projection} == {"bf16"} + assert {_shape(case) for case in projection} == set(FIXED_SHAPES) + assert all(case.backend == "sdpa" for case in projection) + + esmc_projection = [case for case in model_cases if case.mode == "esmc_projection"] + for precision in ("bf16", "fp8"): + for backend in ("eager", "sdpa", "flex_attention"): + matching = [ + case + for case in esmc_projection + if case.precision == precision and case.backend == backend + ] + assert {_shape(case) for case in matching} == set(FIXED_SHAPES) + assert all(case.claim_eligible for case in matching) + + full_embedding = [case for case in model_cases if case.mode == "esmfold2_embed"] + assert {case.precision for case in full_embedding} == {"bf16", "fp8"} + assert all(_shape(case) == (1, 512, ()) for case in full_embedding) + assert all(not case.claim_eligible for case in full_embedding) + assert all( + case.suite_profile == ESMFOLD2_REPRESENTATION_PROFILE + and case.dedicated_mode == ESMFOLD2_DEDICATED_MODE + for case in model_cases + ) + + +def test_structure_models_use_dedicated_startup_records() -> None: + for family in ("esmfold", "boltz2"): + cases = list(benchmark_cases(family=family, quick=False, local_files_only=True)) + assert len(cases) == 1 + case = cases[0] + assert case.mode == "startup" + assert case.suite_profile == STRUCTURE_STARTUP_PROFILE + assert case.dedicated_mode == STRUCTURE_DEDICATED_MODE + assert not case.claim_eligible + + +def test_every_manifest_family_declares_a_benchmark_tier() -> None: + registry = get_model_registry() + assert all("benchmark" in family.test_tiers for family in registry.families.values()) + + +def test_exhaustive_matrix_is_all_checkpoint_and_descriptive() -> None: + cases = list( + exhaustive_benchmark_cases( + family=None, + batch_sizes=(1, 2), + sequence_lengths=(128, 256), + local_files_only=True, + ) + ) + registry = get_model_registry() + expected = { + spec.fast.repo_id for spec in registry.values() if "benchmark" in spec.family.test_tiers + } + assert {case.model for case in cases} == expected + assert all(case.matrix_kind == "exhaustive" for case in cases) + assert all(not case.claim_eligible for case in cases) + + for spec in registry.values(): + if spec.family.tokenizer_mode == "structure": + continue + model_cases = [case for case in cases if case.model == spec.fast.repo_id] + expected_axes = { + (backend, batch_size, sequence_length) + for backend in spec.family.attention + for batch_size in (1, 2) + for sequence_length in (128, 256) + } + assert { + (case.backend, case.batch_size, case.sequence_length) for case in model_cases + } == expected_axes + + for family in ("esmfold", "boltz2"): + model_id = registry[registry.families[family].representative].fast.repo_id + matching = [case for case in cases if case.model == model_id] + assert len(matching) == 1 + assert matching[0].mode == "startup" + + for spec in registry.by_family("esmfold2"): + model_cases = [case for case in cases if case.model == spec.fast.repo_id] + projections = [case for case in model_cases if case.mode == "projection"] + assert { + (case.batch_size, case.sequence_length, case.precision) for case in projections + } == { + (batch_size, sequence_length, "bf16") + for batch_size in (1, 2) + for sequence_length in (128, 256) + } + esmc_projections = [case for case in model_cases if case.mode == "esmc_projection"] + assert { + ( + case.backend, + case.batch_size, + case.sequence_length, + case.precision, + ) + for case in esmc_projections + } == { + (backend, batch_size, sequence_length, precision) + for backend in spec.family.attention + for batch_size in (1, 2) + for sequence_length in (128, 256) + for precision in ("bf16", "fp8") + } + + +def test_projection_mode_rejects_a_misleading_fp8_label_before_gpu_setup() -> None: + arguments = SimpleNamespace(mode="projection", precision="fp8") + with pytest.raises(ValueError, match="esmc_projection"): + run_case(arguments) + + +def test_quick_matrix_is_one_short_case() -> None: + cases = list(benchmark_cases(family=None, quick=True, local_files_only=True)) + assert len(cases) == 1 + assert cases[0].sequence_length <= 128 + + +@pytest.mark.parametrize( + "gpu", + ("NVIDIA H100 PCIe", "NVIDIA H200 NVL", "NVIDIA GH200 480GB"), +) +def test_release_benchmark_accepts_named_hopper_sm90_products(gpu: str) -> None: + validate_hopper_sm90_environment({"gpu": gpu, "gpu_capability": [9, 0]}) + + +@pytest.mark.parametrize( + "environment", + ( + {"gpu": "NVIDIA A100-SXM4-80GB", "gpu_capability": [8, 0]}, + {"gpu": "NVIDIA B200", "gpu_capability": [10, 0]}, + {"gpu": "NVIDIA H100 PCIe", "gpu_capability": [8, 0]}, + ), +) +def test_release_benchmark_rejects_non_hopper_sm90_hardware( + environment: dict[str, object], +) -> None: + with pytest.raises(RuntimeError): + validate_hopper_sm90_environment(environment) + + +def test_benchmark_load_class_is_manifest_advertised() -> None: + registry = get_model_registry() + for spec in registry.values(): + if "benchmark" not in spec.family.test_tiers: + continue + selected = benchmark_auto_class(spec) + assert selected in spec.auto_map + + +def test_benchmark_parameter_dtype_follows_manifest_bf16_execution() -> None: + import torch + + static = SimpleNamespace( + precision="bf16", + bf16_execution="static_parameters", + ) + autocast = SimpleNamespace( + precision="bf16", + bf16_execution="fp32_parameters_autocast", + ) + assert _benchmark_load_dtype(static, torch) == torch.bfloat16 + assert not _uses_bf16_autocast(static) + assert _benchmark_load_dtype(autocast, torch) == torch.float32 + assert _uses_bf16_autocast(autocast) + + +def test_benchmark_derives_registered_bf16_execution_from_manifest() -> None: + registry = get_model_registry() + for model_id in ( + "esm2_8m", + "dplm_150m", + "dplm2_150m", + "boltz2", + "esmfold", + "esmfold2", + ): + spec = registry[model_id] + arguments = SimpleNamespace(model=spec.fast.repo_id, bf16_execution=None) + assert _resolve_bf16_execution(arguments) == spec.family.bf16_execution + + +def test_benchmark_rejects_precision_override_that_conflicts_with_manifest() -> None: + spec = get_model_registry()["dplm_150m"] + arguments = SimpleNamespace( + model=spec.fast.repo_id, + bf16_execution="static_parameters", + ) + with pytest.raises(ValueError, match="conflicts with the manifest policy"): + _resolve_bf16_execution(arguments) + + +def test_architecture_specific_benchmark_heads() -> None: + registry = get_model_registry() + assert benchmark_auto_class(registry["ankh_base"]) == "AutoModel" + assert benchmark_auto_class(registry["esm3_small"]) == "AutoModel" + assert benchmark_auto_class(registry["e1_150m"]) == "AutoModelForMaskedLM" + assert benchmark_auto_class(registry["esmfold"]) == "AutoModel" + assert benchmark_auto_class(registry["boltz2"]) == "AutoModel" + assert benchmark_auto_class(registry["esmfold2"]) == "AutoModel" + + +def test_model_cache_key_reuses_backends_and_shapes_but_not_precision() -> None: + cases = list(benchmark_cases(family="esm2", quick=False, local_files_only=True)) + assert len({benchmark_model_key(case) for case in cases}) == 1 + + structure_cases = list(benchmark_cases(family="esmfold2", quick=False, local_files_only=True)) + by_model = { + model_id: {benchmark_model_key(case) for case in structure_cases if case.model == model_id} + for model_id in {case.model for case in structure_cases} + } + assert all(len(keys) == 2 for keys in by_model.values()) + + +def test_suite_parser_accepts_local_artifact_root() -> None: + arguments = build_parser().parse_args( + [ + "--output", + "report.json", + "--artifact-root", + "dist/hub", + "--backends", + "eager", + "sdpa", + "flex_attention", + ] + ) + + assert arguments.artifact_root == Path("dist/hub") + assert arguments.backends == ["eager", "sdpa", "flex_attention"] + + +def test_local_artifact_binding_preserves_registry_report_identity( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_artifact_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + artifact = _write_benchmark_artifact(tmp_path, spec) + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=False)) + + identities = bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + case = cases[0] + assert case.model == spec.fast.repo_id + assert case.revision == spec.fast.revision + assert case.load_model == artifact.resolve() + assert case.load_revision is None + assert case.local_files_only is True + assert identities[spec.id] == case.artifact_identity + assert case.artifact_identity["weights_revision"] == spec.artifact_checkpoint.revision + assert str(tmp_path) not in json.dumps(identities, sort_keys=True) + + +def test_local_artifact_binding_rejects_missing_and_stale_artifacts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_artifact_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=True)) + + with pytest.raises(ValueError, match="Missing or invalid selected benchmark artifacts"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + _write_benchmark_artifact( + tmp_path, + spec, + runtime_revision="f" * 40, + source_tree_sha256="0" * 64, + ) + with pytest.raises(ValueError, match="registry/frozen source"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + +def test_local_artifact_binding_rejects_swapped_or_forged_artifacts( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_artifact_validation(monkeypatch) + spec = get_model_registry()["esm2_8m"] + _write_benchmark_artifact( + tmp_path, + spec, + config_updates={"fastplms_model_id": "esm2_35m"}, + ) + cases = list(benchmark_cases(family="esm2", quick=True, local_files_only=True)) + + with pytest.raises(ValueError, match="fastplms_model_id"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + monkeypatch.setattr( + benchmark_suite, + "_validate_built_artifact", + lambda *_args: (_ for _ in ()).throw(ValueError("forged manifest")), + ) + config_path = tmp_path / spec.fast.repo_id.rsplit("/", maxsplit=1)[1] / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["fastplms_model_id"] = spec.id + config_path.write_text(json.dumps(config), encoding="utf-8") + with pytest.raises(ValueError, match="forged manifest"): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + +def test_esmfold2_local_artifact_binding_requires_and_records_esmc_backbone( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _stub_artifact_validation(monkeypatch) + registry = get_model_registry() + spec = registry["esmfold2"] + backbone_id = spec.family.backbone_model + assert backbone_id is not None + backbone = registry[backbone_id] + primary_path = _write_benchmark_artifact(tmp_path, spec) + cases = list(benchmark_cases(family="esmfold2", quick=True, local_files_only=False)) + + with pytest.raises(ValueError, match=backbone_id): + bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + backbone_path = _write_benchmark_artifact(tmp_path, backbone) + identities = bind_local_artifacts(cases, tmp_path, source_root=tmp_path) + + case = cases[0] + assert case.load_model == primary_path.resolve() + assert case.esmc_load_model == backbone_path.resolve() + assert case.artifact_dependencies == {"esmc": identities[backbone_id]} + assert str(tmp_path) not in json.dumps(case.artifact_dependencies, sort_keys=True) diff --git a/tests/unit/test_benchmark_regression.py b/tests/unit/test_benchmark_regression.py new file mode 100644 index 0000000..b5ea666 --- /dev/null +++ b/tests/unit/test_benchmark_regression.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import pytest + +from benchmarks.regression import ( + GateThresholds, + bootstrap_ratio_interval, + compare_reports, +) + + +ENVIRONMENT = { + "python": "3.12.3", + "platform": "Linux-test", + "machine": "aarch64", + "torch": "2.13.0+cu130", + "cuda_runtime": "13.0", + "cudnn": 92000, + "transformers": "5.13.0", + "fastplms": "1.0.0", + "transformer_engine": "2.12.0", + "kernels": "0.15.2", + "kernels_data": "0.16.0", + "gpu": "NVIDIA H100 80GB HBM3", + "gpu_capability": [9, 0], + "nvidia_smi": { + "name": "NVIDIA H100 80GB HBM3", + "driver_version": "999.0", + "memory.total": "81559", + "temperature.gpu": "31", + "clocks.sm": "345", + }, +} + + +def _report(values: list[float], *, memory: int = 1_000_000_000) -> dict: + return { + "schema_version": 3, + "status": "complete", + "environment": ENVIRONMENT, + "matrix_kind": "fixed", + "claim_scope": "validated_hopper_sm90_exact_device", + "backend_policy": { + "requested": ["eager", "sdpa", "flex_attention"], + "selection": "explicit_subset", + "external_kernel_downloads": False, + "external_kernel_builds": False, + }, + "timing_contract": { + "cold_compile_field": "results[].compile_ms", + "warm_throughput_field": "results[].blocks", + }, + "baseline_promotion_contract": { + "requires_exact_environment_match": True, + "requires_exact_artifact_inventory_match": True, + }, + "expected_case_count": 1, + "completed_case_count": 1, + "results": [ + { + "model": "example/model", + "revision": "abc123", + "auto_class": "AutoModel", + "backend": "sdpa", + "precision": "bf16", + "mode": "steady", + "batch_size": 1, + "sequence_length": 512, + "lengths": [512], + "blocks": [{"logical_tokens_per_second": value} for value in values], + "memory": {"peak_allocated_bytes": memory}, + } + ], + } + + +def _with_artifact_inventory(report: dict, *, runtime_revision: str = "a" * 40) -> dict: + report["artifact_load_mode"] = "validated_local_build" + report["artifacts"] = { + "esm2_8m": { + "model_id": "esm2_8m", + "registry_repo_id": "Synthyra/ESM2-8M", + "registry_revision": "b" * 40, + "weights_revision": "b" * 40, + "runtime_revision": runtime_revision, + "source_tree_sha256": "c" * 64, + "runtime_bundle_sha256": "d" * 64, + "canonical_state_sha256": "e" * 64, + "artifact_manifest_sha256": "f" * 64, + } + } + return report + + +def test_bootstrap_interval_is_deterministic() -> None: + first = bootstrap_ratio_interval([105, 106, 104], [100, 100, 100], samples=500) + second = bootstrap_ratio_interval([105, 106, 104], [100, 100, 100], samples=500) + assert first == second + assert first[0] == 1.05 + + +def test_gate_accepts_equivalent_results() -> None: + thresholds = GateThresholds(bootstrap_samples=500) + result = compare_reports( + _report([101, 100, 99, 101, 100, 99, 100]), + _report([100] * 7), + thresholds, + ) + assert result.passed + assert result.cases[0].passed + + +def test_gate_rejects_hard_throughput_regression() -> None: + thresholds = GateThresholds(bootstrap_samples=500) + result = compare_reports(_report([85] * 7), _report([100] * 7), thresholds) + assert not result.passed + assert any("hard limit" in reason for reason in result.cases[0].reasons) + + +def test_gate_rejects_large_memory_growth() -> None: + thresholds = GateThresholds(bootstrap_samples=100) + result = compare_reports( + _report([100] * 7, memory=1_400_000_000), + _report([100] * 7, memory=1_000_000_000), + thresholds, + ) + assert not result.passed + assert any("memory" in reason for reason in result.cases[0].reasons) + + +def test_gate_requires_every_baseline_case() -> None: + current = _report([100] * 7) + current["results"] = [] + current["expected_case_count"] = 0 + current["completed_case_count"] = 0 + result = compare_reports(current, _report([100] * 7)) + assert not result.passed + assert result.unmatched_baseline + + +def test_gate_rejects_current_case_without_a_baseline() -> None: + baseline = _report([100] * 7) + current = _report([100] * 7) + extra = dict(current["results"][0]) + extra["backend"] = "flex_attention" + current["results"].append(extra) + current["expected_case_count"] = 2 + current["completed_case_count"] = 2 + + result = compare_reports(current, baseline, GateThresholds(bootstrap_samples=100)) + + assert not result.passed + assert result.unmatched_current + + +def test_gate_rejects_environment_drift() -> None: + current = _report([100] * 7) + current["environment"] = {**ENVIRONMENT, "torch": "2.13.1+cu130"} + + result = compare_reports( + current, + _report([100] * 7), + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert result.environment_mismatches == ( + "environment.torch: current='2.13.1+cu130', baseline='2.13.0+cu130'", + ) + + +def test_gate_rejects_cross_device_hopper_comparison() -> None: + current = _report([100] * 7) + current["environment"] = { + **ENVIRONMENT, + "gpu": "NVIDIA GH200 480GB", + "nvidia_smi": { + **ENVIRONMENT["nvidia_smi"], + "name": "NVIDIA GH200 480GB", + "memory.total": "97871", + }, + } + + result = compare_reports( + current, + _report([100] * 7), + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert any( + mismatch.startswith("environment.gpu:") for mismatch in result.environment_mismatches + ) + assert any( + mismatch.startswith("environment.nvidia_smi.name:") + for mismatch in result.environment_mismatches + ) + assert any( + mismatch.startswith("environment.nvidia_smi.memory.total:") + for mismatch in result.environment_mismatches + ) + + +def test_gate_rejects_architecture_and_driver_drift() -> None: + current = _report([100] * 7) + current["environment"] = { + **ENVIRONMENT, + "machine": "x86_64", + "nvidia_smi": { + **ENVIRONMENT["nvidia_smi"], + "driver_version": "580.105.08", + }, + } + + result = compare_reports( + current, + _report([100] * 7), + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert any( + mismatch.startswith("environment.machine:") + for mismatch in result.environment_mismatches + ) + assert any( + mismatch.startswith("environment.nvidia_smi.driver_version:") + for mismatch in result.environment_mismatches + ) + + +def test_gate_rejects_missing_required_environment_identity() -> None: + baseline = _report([100] * 7) + baseline["environment"] = { + key: value for key, value in ENVIRONMENT.items() if key != "machine" + } + + result = compare_reports( + _report([100] * 7), + baseline, + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert ( + "environment.machine: baseline report is missing the field" + in result.environment_mismatches + ) + + +def test_gate_rejects_incomplete_or_different_promotion_contracts() -> None: + current = _report([100] * 7) + current["status"] = "running" + baseline = _report([100] * 7) + baseline["backend_policy"] = { + **baseline["backend_policy"], + "requested": ["sdpa"], + } + + result = compare_reports( + current, + baseline, + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert "current report status is not complete" in result.report_mismatches + assert any( + mismatch.startswith("backend_policy:") for mismatch in result.report_mismatches + ) + + +def test_gate_rejects_duplicate_and_nonfinite_measurements() -> None: + duplicate = _report([100] * 7) + duplicate["results"].append(dict(duplicate["results"][0])) + duplicate["expected_case_count"] = 2 + duplicate["completed_case_count"] = 2 + with pytest.raises(ValueError, match="duplicate case"): + compare_reports(duplicate, _report([100] * 7)) + + nonfinite = _report([100] * 7) + nonfinite["results"][0]["blocks"][0]["logical_tokens_per_second"] = float("nan") + with pytest.raises(ValueError, match="non-positive/non-finite"): + compare_reports(nonfinite, _report([100] * 7)) + + +def test_gate_rejects_artifact_identity_drift() -> None: + current = _with_artifact_inventory(_report([100] * 7), runtime_revision="1" * 40) + baseline = _with_artifact_inventory(_report([100] * 7), runtime_revision="2" * 40) + + result = compare_reports( + current, + baseline, + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert any( + mismatch.startswith("artifacts.esm2_8m:") for mismatch in result.artifact_mismatches + ) + + +def test_gate_rejects_missing_artifact_inventory() -> None: + current = _with_artifact_inventory(_report([100] * 7)) + + result = compare_reports( + current, + _report([100] * 7), + GateThresholds(bootstrap_samples=100), + ) + + assert not result.passed + assert "baseline report has no artifact inventory mapping" in result.artifact_mismatches + + +def test_gate_ignores_telemetry_that_is_not_environment_identity() -> None: + current = _report([100] * 7) + current_smi = dict(ENVIRONMENT["nvidia_smi"]) + current_smi["temperature.gpu"] = "79" + current_smi["clocks.sm"] = "1980" + current["environment"] = {**ENVIRONMENT, "nvidia_smi": current_smi} + + result = compare_reports( + current, + _report([100] * 7), + GateThresholds(bootstrap_samples=100), + ) + + assert result.passed + assert not result.environment_mismatches + + +def test_descriptive_records_do_not_enter_throughput_gate() -> None: + current = _report([100] * 7) + baseline = _report([100] * 7) + descriptive = { + "model": "example/model", + "revision": "abc123", + "backend": "sdpa", + "mode": "embed", + "batch_size": 1, + "sequence_length": 512, + "lengths": [512], + "blocks": [], + "embedding_ms": 12.0, + "memory": {"peak_allocated_bytes": 1_000_000_000}, + } + current["results"].append({**descriptive, "embedding_ms": 11.0}) + baseline["results"].append(descriptive) + current["expected_case_count"] = baseline["expected_case_count"] = 2 + current["completed_case_count"] = baseline["completed_case_count"] = 2 + + result = compare_reports( + current, + baseline, + GateThresholds(bootstrap_samples=100), + ) + + assert result.passed + assert len(result.cases) == 1 diff --git a/tests/unit/test_binder_example_contracts.py b/tests/unit/test_binder_example_contracts.py new file mode 100644 index 0000000..f52557d --- /dev/null +++ b/tests/unit/test_binder_example_contracts.py @@ -0,0 +1,74 @@ +"""Unit contracts for the public binder-design example helpers.""" + +from __future__ import annotations + +import sys +import pytest +import torch +from types import ModuleType +from typing import Any + +from examples import binder_design_fastplms as binder + + +class _Position: + def __init__(self, in_cdr: bool) -> None: + self._in_cdr = in_cdr + + def is_in_cdr(self) -> bool: + return self._in_cdr + + +class _Chain: + def __init__(self, sequence: str, cdr_offsets: set[int]) -> None: + self.seq = sequence + self._cdr_offsets = cdr_offsets + + def __iter__(self): + for offset, residue in enumerate(self.seq): + yield _Position(offset in self._cdr_offsets), residue + + +def test_cdr_indices_use_public_abnumber_multiple_domain_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, Any] = {} + + class PublicChain: + @classmethod + def multiple_domains(cls, sequence: str, **kwargs: Any) -> list[_Chain]: + del cls + observed.update(sequence=sequence, kwargs=kwargs) + return [ + _Chain("AAACCC", {1, 4}), + _Chain("GGGTTT", {0, 5}), + ] + + abnumber = ModuleType("abnumber") + abnumber.Chain = PublicChain # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "abnumber", abnumber) + + sequence = "AAACCCLINKGGGTTT" + assert binder._cdr_indices(sequence) == [1, 4, 10, 15] + assert observed == { + "sequence": sequence, + "kwargs": { + "scheme": "chothia", + "allowed_species": None, + "use_anarcii": True, + }, + } + + +def test_binder_helper_validation_uses_explicit_exceptions() -> None: + with pytest.raises(ValueError, match="Unsupported fixed binder residue"): + binder.build_initial_soft_sequence_logits("A?", batch_size=1) + with pytest.raises(ValueError, match="Distogram logits must have shape"): + binder.compute_distogram_iptm_proxy( + torch.zeros(3, 3, 128), + target_length=2, + binder_sequence="AA", + is_antibody=False, + ) + with pytest.raises(ValueError, match="separated by one"): + binder._binder_sequence_from_designed_sequence("missing-separator") diff --git a/tests/unit/test_biohub_reference_environment.py b/tests/unit/test_biohub_reference_environment.py new file mode 100644 index 0000000..bde8b81 --- /dev/null +++ b/tests/unit/test_biohub_reference_environment.py @@ -0,0 +1,239 @@ +"""Fail-closed contracts for the pinned Biohub reference environment.""" + +from __future__ import annotations + +import hashlib +import json +import sys +import pytest +from pathlib import Path +from types import ModuleType + +from tools.remote.biohub_reference_requirements import ( + BiohubReferenceRequirementsError, + extract_biohub_reference_requirements, + write_biohub_reference_requirements, +) +from tools.remote.reference_source_attestation import ( + ReferenceSourceAttestationError, + create_reference_source_attestation, + validate_reference_source_evidence, + validate_reference_sources_evidence, + verify_reference_source, +) +from tools.source_provenance import actual_tree_paths, tracked_tree_digest + + +_TRANSFORMERS_MAIN = "transformers @ git+https://github.com/Biohub/transformers.git@main" + + +def _write_biohub_pyproject(path: Path, dependencies: list[str]) -> None: + rendered = "\n".join(f' "{requirement}",' for requirement in dependencies) + path.write_text( + f"[project]\nname = \"esm\"\ndependencies = [\n{rendered}\n]\n", + encoding="utf-8", + ) + + +def test_biohub_requirements_remove_only_the_known_transformers_main_url( + tmp_path: Path, +) -> None: + pyproject = tmp_path / "pyproject.toml" + output = tmp_path / "requirements.txt" + _write_biohub_pyproject( + pyproject, + ["torch>=2.2", _TRANSFORMERS_MAIN, "accelerate", "biotite>=1"], + ) + + assert extract_biohub_reference_requirements(pyproject) == ( + "torch>=2.2", + "accelerate", + "biotite>=1", + ) + assert write_biohub_reference_requirements(pyproject, output) == ( + "torch>=2.2", + "accelerate", + "biotite>=1", + ) + assert output.read_text(encoding="utf-8") == "torch>=2.2\naccelerate\nbiotite>=1\n" + + +@pytest.mark.parametrize( + "dependencies", + ( + ["torch", "transformers>=4.57"], + ["torch", _TRANSFORMERS_MAIN, "other @ https://example.invalid/archive.whl"], + ["torch", _TRANSFORMERS_MAIN, _TRANSFORMERS_MAIN], + ["torch", _TRANSFORMERS_MAIN, "-r injected.txt"], + ["torch", _TRANSFORMERS_MAIN, "--extra-index-url=https://example.invalid"], + ["torch", _TRANSFORMERS_MAIN, "../local-wheel.whl"], + ["torch", _TRANSFORMERS_MAIN, "package; python_version >= '3.12'"], + ), +) +def test_biohub_requirements_reject_dependency_contract_drift( + tmp_path: Path, + dependencies: list[str], +) -> None: + pyproject = tmp_path / "pyproject.toml" + _write_biohub_pyproject(pyproject, dependencies) + + with pytest.raises(BiohubReferenceRequirementsError): + extract_biohub_reference_requirements(pyproject) + + +def _reference_source_fixture(tmp_path: Path) -> tuple[Path, Path, str]: + source_root = tmp_path / "source" + package = source_root / "src" / "pinned_reference_probe" + package.mkdir(parents=True) + (package / "__init__.py").write_text('__version__ = "1.2.3"\n', encoding="utf-8") + (package / "model.py").write_text("VALUE = 7\n", encoding="utf-8") + digest = tracked_tree_digest(source_root, actual_tree_paths(source_root)) + revision = "a" * 40 + contract = tmp_path / "contract.json" + contract.write_text( + json.dumps( + { + "schema_version": 1, + "source_revision": revision, + "tree_sha256": digest, + "import_name": "pinned_reference_probe", + "import_root": "src/pinned_reference_probe", + "package_version": "1.2.3", + } + ) + + "\n", + encoding="utf-8", + ) + return source_root, contract, revision + + +def test_reference_source_attestation_rehashes_and_proves_import_origin( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + created = create_reference_source_attestation(source_root, contract, attestation) + assert created["source_revision"] == revision + assert created["file_count"] == 2 + + monkeypatch.syspath_prepend(str(source_root / "src")) + monkeypatch.setattr(sys, "dont_write_bytecode", False) + sys.modules.pop("pinned_reference_probe", None) + evidence = verify_reference_source( + source_root, + attestation, + contract, + expected_revision=revision, + ) + assert evidence == { + "attestation_sha256": hashlib.sha256(attestation.read_bytes()).hexdigest(), + "file_count": 2, + "import_file": "src/pinned_reference_probe/__init__.py", + "import_name": "pinned_reference_probe", + "import_root": "src/pinned_reference_probe", + "package_version": "1.2.3", + "schema_version": 1, + "source_revision": revision, + "tree_sha256": created["tree_sha256"], + } + assert validate_reference_source_evidence(evidence) == evidence + sources = { + "biohub-esm": dict(evidence), + "biohub-transformers": dict(evidence), + } + assert validate_reference_sources_evidence( + sources, + required_sources=("biohub-esm", "biohub-transformers"), + ) == sources + with pytest.raises(ReferenceSourceAttestationError, match="names differ"): + validate_reference_sources_evidence( + {"biohub-transformers": evidence}, + required_sources=("biohub-esm", "biohub-transformers"), + ) + assert sys.dont_write_bytecode is False + assert not any(path.name == "__pycache__" for path in source_root.rglob("*")) + + (source_root / "src/pinned_reference_probe/model.py").write_text( + "VALUE = 8\n", + encoding="utf-8", + ) + with pytest.raises(ReferenceSourceAttestationError, match="changed after image construction"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) + + +def test_reference_source_attestation_rejects_untracked_import_code(tmp_path: Path) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + create_reference_source_attestation(source_root, contract, attestation) + (source_root / "src/pinned_reference_probe/injected.py").write_text( + "VALUE = 'untracked'\n", + encoding="utf-8", + ) + + with pytest.raises(ReferenceSourceAttestationError, match="inventory differs"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) + + +def test_reference_source_attestation_rejects_untracked_source_root_code( + tmp_path: Path, +) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + create_reference_source_attestation(source_root, contract, attestation) + (source_root / "src/sitecustomize.py").write_text( + "raise RuntimeError('must never execute')\n", + encoding="utf-8", + ) + + with pytest.raises(ReferenceSourceAttestationError, match="inventory differs"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) + + +def test_reference_source_attestation_rejects_self_authored_tree_identity( + tmp_path: Path, +) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + create_reference_source_attestation(source_root, contract, attestation) + (source_root / "src/pinned_reference_probe/model.py").write_text( + "VALUE = 9\n", + encoding="utf-8", + ) + payload = json.loads(attestation.read_text(encoding="utf-8")) + payload["tree_sha256"] = tracked_tree_digest(source_root, payload["tracked_files"]) + attestation.write_text(json.dumps(payload) + "\n", encoding="utf-8") + + with pytest.raises(ReferenceSourceAttestationError, match="checked-in contract"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) + + +def test_reference_source_attestation_rejects_cached_external_import( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + create_reference_source_attestation(source_root, contract, attestation) + rogue = ModuleType("pinned_reference_probe") + rogue.__file__ = str(tmp_path / "rogue/pinned_reference_probe/__init__.py") + monkeypatch.setitem(sys.modules, "pinned_reference_probe", rogue) + + with pytest.raises(ReferenceSourceAttestationError, match="outside the pinned source"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) + + +def test_reference_source_attestation_rejects_cached_external_submodule( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source_root, contract, revision = _reference_source_fixture(tmp_path) + attestation = tmp_path / "attestation.json" + create_reference_source_attestation(source_root, contract, attestation) + monkeypatch.delitem(sys.modules, "pinned_reference_probe", raising=False) + rogue = ModuleType("pinned_reference_probe.poison") + rogue.__file__ = str(tmp_path / "rogue/pinned_reference_probe/poison.py") + monkeypatch.setitem(sys.modules, "pinned_reference_probe.poison", rogue) + + with pytest.raises(ReferenceSourceAttestationError, match="outside the pinned source"): + verify_reference_source(source_root, attestation, contract, expected_revision=revision) diff --git a/tests/unit/test_biohub_reference_lock.py b/tests/unit/test_biohub_reference_lock.py new file mode 100644 index 0000000..659ee3d --- /dev/null +++ b/tests/unit/test_biohub_reference_lock.py @@ -0,0 +1,423 @@ +"""Fail-closed tests for the native GH200 Biohub dependency lock.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import shutil +import subprocess +import pytest +from dataclasses import asdict +from pathlib import Path + +import tools.remote.biohub_reference_lock as biohub_reference_lock +from tools.remote.biohub_reference_environment import ( + BiohubReferenceEnvironmentError, + validate_biohub_reference_environment_evidence, +) +from tools.remote.biohub_reference_lock import ( + BiohubReferenceLockError, + assert_exact_installed_inventory, + expected_installed_inventory, + load_biohub_reference_lock_contract, + materialize_biotraj_wheel_lock, + verify_biohub_reference_lock_contract, + verify_current_pip_check, +) + + +_ROOT = Path(__file__).parents[2] +_CONTRACT = _ROOT / "docker/constraints/biohub-reference-lock.json" +_CONTRACT_FILES = ( + "docker/biohub-reference-lock.Dockerfile", + "docker/constraints/biohub-reference.in", + "docker/constraints/biohub-reference.lock.txt", + "docker/constraints/biohub-biotraj-build.in", + "docker/constraints/biohub-biotraj-build.lock.txt", + "docker/constraints/biohub-reference-lock.json", +) + + +def _copy_contract_root(tmp_path: Path) -> tuple[Path, Path]: + for relative in _CONTRACT_FILES: + destination = tmp_path / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(_ROOT / relative, destination) + return tmp_path, tmp_path / "docker/constraints/biohub-reference-lock.json" + + +def test_checked_in_biohub_reference_lock_is_exact_gh200_contract() -> None: + evidence = verify_biohub_reference_lock_contract(_ROOT, _CONTRACT) + + assert evidence == { + "schema_version": 2, + "target": { + "hardware": "NVIDIA GH200 480GB", + "operating_system": "linux", + "architecture": "aarch64", + "container_platform": "linux/arm64", + "python_implementation": "CPython", + "python_version": "3.12", + "cuda_version": "13.0", + "torch_backend": "cu130", + }, + "runtime_lock_sha256": ("f87033dffffe953478b482dae82f91603fa705a68e92ee7683c1831586c94ca0"), + "runtime_package_count": 108, + "build_lock_sha256": ("c7864daa96028aba35081110c563b8b08c968fd312f7827449d355638f18079d"), + "build_package_count": 7, + "biotraj_sdist_sha256": ( + "4bcba92101ed50f369cc1487fb5dfcfe1d8402ad47adaa9232b080553271663a" + ), + "biotraj_wheel_sha256": ( + "253c1354c401e97d6e951f29e0d768deb5263de6662001281870425c37719f6b" + ), + "pip_check_platform_exceptions": [ + { + "distribution": "nvidia-cusparselt-cu13", + "version": "0.8.1", + "wheel_filename": ( + "nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl" + ), + "wheel_sha256": ( + "4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f" + ), + "filename_platform_tag": "py3-none-manylinux2014_aarch64", + "wheel_metadata_platform_tag": "py3-none-manylinux2014_sbsa", + "target_hardware": "NVIDIA GH200 480GB", + "target_operating_system": "linux", + "target_architecture": "aarch64", + "accepted_diagnostic": ( + "nvidia-cusparselt-cu13 0.8.1 is not supported on this platform" + ), + "resolution": "validated-vendor-metadata-exception-no-wheel-rewrite", + } + ], + } + + +def test_biohub_lock_contract_rejects_any_target_broadening(tmp_path: Path) -> None: + payload = json.loads(_CONTRACT.read_text(encoding="utf-8")) + payload["target"]["architecture"] = "x86_64" + changed = tmp_path / "contract.json" + changed.write_text(json.dumps(payload) + "\n", encoding="utf-8") + + with pytest.raises(BiohubReferenceLockError, match="Unsupported Biohub lock target"): + load_biohub_reference_lock_contract(changed) + + +def test_biohub_lock_rejects_byte_mutation_and_forged_digest(tmp_path: Path) -> None: + root, contract_path = _copy_contract_root(tmp_path) + lock = root / "docker/constraints/biohub-reference.lock.txt" + lock.write_bytes(lock.read_bytes() + b"# injected\n") + + with pytest.raises(BiohubReferenceLockError, match="file digest differs"): + verify_biohub_reference_lock_contract(root, contract_path) + + payload = json.loads(contract_path.read_text(encoding="utf-8")) + payload["runtime"]["lock_sha256"] = hashlib.sha256(lock.read_bytes()).hexdigest() + contract_path.write_text(json.dumps(payload) + "\n", encoding="utf-8") + with pytest.raises(BiohubReferenceLockError, match="Unexpected comment"): + verify_biohub_reference_lock_contract(root, contract_path) + + +def test_materialized_runtime_lock_uses_only_attested_native_wheel( + tmp_path: Path, +) -> None: + root, contract_path = _copy_contract_root(tmp_path) + wheel = tmp_path / "biotraj-1.2.2-cp312-cp312-linux_aarch64.whl" + wheel.write_bytes(b"deterministic-native-wheel") + wheel_sha256 = hashlib.sha256(wheel.read_bytes()).hexdigest() + payload = json.loads(contract_path.read_text(encoding="utf-8")) + payload["biotraj"]["wheel_sha256"] = wheel_sha256 + payload["biotraj"]["wheel_size"] = wheel.stat().st_size + contract_path.write_text(json.dumps(payload) + "\n", encoding="utf-8") + output = tmp_path / "materialized.lock.txt" + wheel_uri = f"file:///opt/wheels/{wheel.name}" + + parsed = materialize_biotraj_wheel_lock( + root, + contract_path, + wheel, + output, + wheel_uri=wheel_uri, + ) + + rendered = output.read_text(encoding="utf-8") + assert len(parsed.inventory) == 108 + assert parsed.inventory["biotraj"] == "1.2.2" + assert "--no-binary biotraj" not in rendered + assert f"biotraj @ {wheel_uri}" in rendered + assert f"--hash=sha256:{wheel_sha256}" in rendered + assert "biotraj-1.2.2.tar.gz" not in rendered + + wheel.write_bytes(wheel.read_bytes() + b"mutation") + with pytest.raises(BiohubReferenceLockError, match="size differs"): + materialize_biotraj_wheel_lock( + root, + contract_path, + wheel, + output, + wheel_uri=wheel_uri, + ) + + +def test_exact_inventory_profiles_include_only_declared_overlays() -> None: + build = expected_installed_inventory(_ROOT, _CONTRACT, profile="build") + runtime = expected_installed_inventory(_ROOT, _CONTRACT, profile="runtime") + final = expected_installed_inventory(_ROOT, _CONTRACT, profile="final") + + assert len(build) == 7 + assert len(runtime) == 109 + assert len(final) == 112 + assert runtime["pip"] == "26.1.1" + assert "esm" not in runtime and "transformers" not in runtime and "uv" not in runtime + assert final["esm"] == "3.3.0" + assert final["transformers"] == "4.57.6" + assert final["uv"] == "0.10.12" + + +def test_inventory_comparison_normalizes_names_and_torch_cuda_local_version() -> None: + expected = {"huggingface-hub": "0.36.2", "torch": "2.13.0+cu130"} + observed = {"huggingface_hub": "0.36.2", "Torch": "2.13.0+cu13_0"} + + assert assert_exact_installed_inventory(expected, observed) == { + "huggingface-hub": "0.36.2", + "torch": "2.13.0+cu130", + } + + +@pytest.mark.parametrize( + "observed", + ( + {"torch": "2.13.0+cu130"}, + {"torch": "2.13.0+cu130", "numpy": "1.26.4", "rogue": "1"}, + {"torch": "2.13.0+cpu", "numpy": "1.26.4"}, + ), +) +def test_inventory_comparison_rejects_missing_extra_and_changed( + observed: dict[str, str], +) -> None: + expected = {"torch": "2.13.0+cu130", "numpy": "1.26.4"} + + with pytest.raises(BiohubReferenceLockError, match="inventory differs"): + assert_exact_installed_inventory(expected, observed) + + +def _patch_pip_check_runtime( + monkeypatch: pytest.MonkeyPatch, + *, + wheel_tag: str = "py3-none-manylinux2014_sbsa", + stdout: str = "nvidia-cusparselt-cu13 0.8.1 is not supported on this platform\n", + returncode: int = 1, +) -> None: + class Distribution: + @staticmethod + def read_text(filename: str) -> str: + assert filename == "WHEEL" + return f"Wheel-Version: 1.0\nTag: {wheel_tag}\n" + + monkeypatch.setattr( + biohub_reference_lock, + "verify_current_installed_inventory", + lambda *_args, **_kwargs: {}, + ) + monkeypatch.setattr(biohub_reference_lock.platform, "system", lambda: "Linux") + monkeypatch.setattr(biohub_reference_lock.platform, "machine", lambda: "aarch64") + monkeypatch.setattr( + biohub_reference_lock.importlib.metadata, + "distribution", + lambda _name: Distribution(), + ) + monkeypatch.setattr( + biohub_reference_lock.subprocess, + "run", + lambda *_args, **_kwargs: subprocess.CompletedProcess( + args=[], returncode=returncode, stdout=stdout, stderr="" + ), + ) + + +def test_pip_check_accepts_only_attested_nvidia_sbsa_tag_defect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_pip_check_runtime(monkeypatch) + contract = load_biohub_reference_lock_contract(_CONTRACT) + + assert verify_current_pip_check(_ROOT, _CONTRACT) == { + "status": "accepted-platform-exception", + "returncode": 1, + "diagnostics": ["nvidia-cusparselt-cu13 0.8.1 is not supported on this platform"], + "accepted_platform_exceptions": [ + asdict(exception) for exception in contract.pip_check_platform_exceptions + ], + } + + +def test_pip_check_rejects_any_additional_diagnostic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_pip_check_runtime( + monkeypatch, + stdout=( + "nvidia-cusparselt-cu13 0.8.1 is not supported on this platform\n" + "rogue 1.0 requires missing-package\n" + ), + ) + + with pytest.raises(BiohubReferenceLockError, match="differs from the one accepted"): + verify_current_pip_check(_ROOT, _CONTRACT) + + +def test_pip_check_rejects_changed_vendor_wheel_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_pip_check_runtime(monkeypatch, wheel_tag="py3-none-manylinux2014_aarch64") + + with pytest.raises(BiohubReferenceLockError, match="WHEEL tag differs"): + verify_current_pip_check(_ROOT, _CONTRACT) + + +def _canonical_digest(value: object) -> str: + serialized = (json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n").encode() + return hashlib.sha256(serialized).hexdigest() + + +def _reference_environment_payload() -> dict[str, object]: + contract = load_biohub_reference_lock_contract(_CONTRACT) + inventory = expected_installed_inventory(_ROOT, _CONTRACT, profile="final") + digest = "sha256:" + "a" * 64 + image = { + "content_digest": digest, + "image_id": digest, + "os": "linux", + "architecture": "arm64", + "resolved_platform": "linux/arm64", + } + container_identity = { + "schema_version": 1, + "resolved_platform": "linux/arm64", + "docker_server": { + "Version": "28.0.0", + "ApiVersion": "1.48", + "Os": "linux", + "Arch": "arm64", + }, + "docker_buildx": "github.com/docker/buildx v0.25.0 deadbeef", + "images": { + "biohub-biotraj-wheel": dict(image), + "reference-biohub-esm": dict(image), + }, + } + return { + "schema_version": 2, + "contract": contract.contract, + "contract_sha256": hashlib.sha256(_CONTRACT.read_bytes()).hexdigest(), + "target": asdict(contract.target), + "build_container": asdict(contract.container), + "locks": verify_biohub_reference_lock_contract(_ROOT, _CONTRACT), + "biotraj": asdict(contract.biotraj), + "installed_inventory": inventory, + "installed_inventory_sha256": _canonical_digest(inventory), + "pip_check": { + "status": "accepted-platform-exception", + "returncode": 1, + "diagnostics": [ + exception.accepted_diagnostic + for exception in contract.pip_check_platform_exceptions + ], + "accepted_platform_exceptions": [ + asdict(exception) for exception in contract.pip_check_platform_exceptions + ], + }, + "reference_container_target": "reference-biohub-esm", + "container_identity": container_identity, + "container_identity_sha256": _canonical_digest(container_identity), + "runtime": { + "operating_system": "linux", + "architecture": "aarch64", + "python_implementation": "CPython", + "python_version": "3.12.11", + "torch": inventory["torch"], + "cuda_runtime": "13.0", + "cuda_driver": "580.65.06", + "gpu": { + "name": "NVIDIA GH200 480GB", + "capability": [9, 0], + "total_memory_bytes": 480_000_000_000, + }, + "uname": { + "system": "Linux", + "release": "6.8.0", + "version": "#1 SMP PREEMPT_DYNAMIC", + "machine": "aarch64", + }, + }, + } + + +def test_reference_environment_evidence_is_exact_and_deterministic() -> None: + payload = _reference_environment_payload() + + assert validate_biohub_reference_environment_evidence( + payload, + repository_root=_ROOT, + contract_path=_CONTRACT, + ) == {field: payload[field] for field in sorted(payload)} + + ephemeral = copy.deepcopy(payload) + assert isinstance(ephemeral["runtime"], dict) + assert isinstance(ephemeral["runtime"]["uname"], dict) + ephemeral["runtime"]["uname"]["node"] = "container-id" + with pytest.raises(BiohubReferenceEnvironmentError, match="uname identity"): + validate_biohub_reference_environment_evidence( + ephemeral, + repository_root=_ROOT, + contract_path=_CONTRACT, + ) + + missing_build_image = copy.deepcopy(payload) + assert isinstance(missing_build_image["container_identity"], dict) + images = missing_build_image["container_identity"]["images"] + assert isinstance(images, dict) + del images["biohub-biotraj-wheel"] + missing_build_image["container_identity_sha256"] = _canonical_digest( + missing_build_image["container_identity"] + ) + with pytest.raises(BiohubReferenceEnvironmentError, match="Required Biohub"): + validate_biohub_reference_environment_evidence( + missing_build_image, + repository_root=_ROOT, + contract_path=_CONTRACT, + ) + + drifted_inventory = copy.deepcopy(payload) + assert isinstance(drifted_inventory["installed_inventory"], dict) + drifted_inventory["installed_inventory"]["torch"] = "2.13.1+cu130" + drifted_inventory["installed_inventory_sha256"] = _canonical_digest( + drifted_inventory["installed_inventory"] + ) + with pytest.raises(BiohubReferenceEnvironmentError, match="installed_inventory"): + validate_biohub_reference_environment_evidence( + drifted_inventory, + repository_root=_ROOT, + contract_path=_CONTRACT, + ) + + broadened_pip_check = copy.deepcopy(payload) + assert isinstance(broadened_pip_check["pip_check"], dict) + assert isinstance(broadened_pip_check["pip_check"]["diagnostics"], list) + broadened_pip_check["pip_check"]["diagnostics"].append("rogue dependency failure") + with pytest.raises(BiohubReferenceEnvironmentError, match="pip_check differs"): + validate_biohub_reference_environment_evidence( + broadened_pip_check, + repository_root=_ROOT, + contract_path=_CONTRACT, + ) + + +def test_biohub_environment_image_validation_survives_python_optimized_mode() -> None: + source = (_ROOT / "tools/remote/biohub_reference_environment.py").read_text(encoding="utf-8") + + assert "assert isinstance(images, Mapping)" not in source + assert source.count("Reference container identity images must be a mapping.") == 2 diff --git a/tests/unit/test_boltz_checkpoint_io.py b/tests/unit/test_boltz_checkpoint_io.py new file mode 100644 index 0000000..83c5874 --- /dev/null +++ b/tests/unit/test_boltz_checkpoint_io.py @@ -0,0 +1,173 @@ +"""Safe persistence contracts for the FastPLMs Boltz2 model.""" + +from __future__ import annotations + +import pytest +import torch +from pathlib import Path +from typing import Any +from torch import nn + +from fastplms.models.boltz import modeling_boltz2 +from fastplms.models.boltz.modeling_boltz2 import Boltz2Config, Boltz2Model + + +class _TinyCore(nn.Module): + """Minimal checkpoint-facing core used to exercise model persistence.""" + + def __init__(self, width: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(width)) # (d=width,) + + +def _install_tiny_core(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(modeling_boltz2, "Boltz2InferenceCore", _TinyCore) + + +def test_lightning_checkpoint_rejects_pickle_without_explicit_opt_in( + monkeypatch: pytest.MonkeyPatch, +) -> None: + load_called = False + + def _unexpected_load(*args: Any, **kwargs: Any) -> None: + nonlocal load_called + load_called = True + + monkeypatch.setattr(torch, "load", _unexpected_load) + + with pytest.raises( + ValueError, + match=r"allow_unsafe_pickle=True only for a trusted, hash-verified checkpoint", + ): + Boltz2Model.from_boltz_checkpoint("untrusted.ckpt") + + assert not load_called + + +def test_lightning_checkpoint_rejects_truthy_non_boolean_opt_in( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + torch, + "load", + lambda *args, **kwargs: pytest.fail("unsafe deserializer was reached"), + ) + + with pytest.raises(ValueError, match=r"allow_unsafe_pickle=True"): + Boltz2Model.from_boltz_checkpoint( + "untrusted.ckpt", + allow_unsafe_pickle=1, # type: ignore[arg-type] + ) + + +def test_lightning_checkpoint_load_requires_and_honors_explicit_opt_in( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_tiny_core(monkeypatch) + expected = torch.tensor([1.25, -2.5]) # (d=2,) + load_call: dict[str, Any] = {} + + def _load( + path: str, + *, + map_location: str | torch.device, + weights_only: bool, + ) -> dict[str, Any]: + load_call.update( + path=path, + map_location=map_location, + weights_only=weights_only, + ) + return { + "hyper_parameters": {}, + "state_dict": {"model.weight": expected.clone()}, # (d=2,) + } + + def _config_from_hyperparameters( + cls: type[Boltz2Config], + hparams: dict[str, Any], + **kwargs: Any, + ) -> Boltz2Config: + assert hparams == {} + assert kwargs["use_kernels"] is False + return cls(core_kwargs={"width": 2}) + + monkeypatch.setattr(torch, "load", _load) + monkeypatch.setattr( + Boltz2Config, + "from_hyperparameters", + classmethod(_config_from_hyperparameters), + ) + + model = Boltz2Model.from_boltz_checkpoint( + "trusted.ckpt", + allow_unsafe_pickle=True, + ) + + assert load_call == { + "path": "trusted.ckpt", + "map_location": "cpu", + "weights_only": False, + } + assert torch.equal(model.core.weight, expected) + assert not model.training + + +def test_lightning_checkpoint_missing_state_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_tiny_core(monkeypatch) + monkeypatch.setattr( + torch, + "load", + lambda *_args, **_kwargs: { + "hyper_parameters": {}, + "state_dict": {}, + }, + ) + monkeypatch.setattr( + Boltz2Config, + "from_hyperparameters", + classmethod(lambda cls, _hparams, **_kwargs: cls(core_kwargs={"width": 2})), + ) + + with pytest.raises(RuntimeError, match="missing required parameters"): + Boltz2Model.from_boltz_checkpoint( + "trusted-but-incomplete.ckpt", + allow_unsafe_pickle=True, + ) + + +def test_save_pretrained_defaults_to_safetensors_and_round_trips( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _install_tiny_core(monkeypatch) + source = Boltz2Model(Boltz2Config(core_kwargs={"width": 3})) + source.core.weight.data.copy_(torch.tensor([0.25, -1.5, 3.0])) # (d=3,) + + source.save_pretrained(tmp_path) + + assert (tmp_path / "model.safetensors").is_file() + assert not (tmp_path / "pytorch_model.bin").exists() + reloaded = Boltz2Model.from_pretrained(tmp_path, local_files_only=True) + assert torch.equal(reloaded.core.weight, source.core.weight) + + +def test_floating_features_default_to_fp32_parameter_storage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_tiny_core(monkeypatch) + model = Boltz2Model(Boltz2Config(core_kwargs={"width": 3})) + features = { + "positions": torch.randn(2, 3, dtype=torch.bfloat16), # (n=2, xyz=3) + "indices": torch.tensor((1, 2), dtype=torch.int64), # (n=2,) + "mask": torch.tensor((True, False)), # (n=2,) + } + + moved = model._to_model_device(features) # shapes unchanged + + assert moved["positions"].dtype == torch.float32 + assert moved["indices"].dtype == torch.int64 + assert moved["mask"].dtype == torch.bool + assert {tensor.device for tensor in moved.values()} == {model.device} diff --git a/tests/unit/test_boltz_opm_diagnostic.py b/tests/unit/test_boltz_opm_diagnostic.py new file mode 100644 index 0000000..7cdf7be --- /dev/null +++ b/tests/unit/test_boltz_opm_diagnostic.py @@ -0,0 +1,58 @@ +"""Contracts for the bounded Boltz2 OPM runtime diagnostic.""" + +from __future__ import annotations + +import json +import torch +import torch.nn.functional as F +from pathlib import Path +from safetensors.torch import save_file + +from tools.debug.analyze_boltz_opm_projection import _PREFIX, main + + +def test_opm_report_localizes_an_output_kernel_difference(tmp_path: Path) -> None: + X = torch.arange(8, dtype=torch.float32).reshape(1, 2, 4) / 10 # (b=1, l=2, d=4) + W = torch.arange(12, dtype=torch.float32).reshape(3, 4) / 20 # (d_out=3, d=4) + bias = torch.tensor([0.1, -0.2, 0.3], dtype=torch.float32) # (d_out=3,) + candidate_output = F.linear(X, W, bias).to(torch.bfloat16) # (b=1, l=2, d_out=3) + reference_output = candidate_output.clone() # (b=1, l=2, d_out=3) + reference_output.reshape(-1)[0] += torch.tensor(0.125, dtype=torch.bfloat16) # () + + input_key = f"{_PREFIX}__call_000__args__0" + output_key = f"{_PREFIX}__call_000__output" + weight_key = f"{_PREFIX}__parameter__weight" + bias_key = f"{_PREFIX}__parameter__bias" + common = {input_key: X, weight_key: W, bias_key: bias} + candidate_path = tmp_path / "candidate.safetensors" + reference_path = tmp_path / "reference.safetensors" + report_path = tmp_path / "report.json" + save_file({**common, output_key: candidate_output}, candidate_path) + save_file({**common, output_key: reference_output}, reference_path) + + assert ( + main( + [ + str(candidate_path), + str(reference_path), + "--output", + str(report_path), + "--device", + "cpu", + ] + ) + == 0 + ) + report = json.loads(report_path.read_text(encoding="utf-8")) + localization = report["localization"] + assert localization == { + "operation": "msa_module.layers.0.outer_product_mean.proj_o", + "first_differing_kernel": "autocast_bf16_linear_output", + "input_equal": True, + "weight_equal": True, + "bias_equal": True, + "recorded_output_equal": False, + } + recorded = report["variants"]["recorded_autocast_bf16"] + assert recorded["exact"] is False + assert recorded["unequal_values"] == 1 diff --git a/tests/unit/test_build_all_artifacts.py b/tests/unit/test_build_all_artifacts.py new file mode 100644 index 0000000..a9b398c --- /dev/null +++ b/tests/unit/test_build_all_artifacts.py @@ -0,0 +1,244 @@ +"""Contracts for manifest-wide local artifact orchestration.""" + +from __future__ import annotations + +import pytest +from pathlib import Path +from types import SimpleNamespace + +from benchmarks.suite import benchmark_artifact_model_ids +from fastplms.registry import ModelRegistry, ModelSpec, get_model_registry +from tools.artifacts import build as build_module +from tools.artifacts import build_all as build_all_module + + +def test_build_all_registry_binds_official_source_artifact_validation( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """ANKH and DPLM2 must not be revalidated from self-attestation alone.""" + + registry = get_model_registry() + model_ids = ("ankh_base", "dplm2_150m") + assert all(registry[model_id].artifact_source == "official" for model_id in model_ids) + + downloaded: dict[Path, tuple[str, str]] = {} + + def fake_snapshot_download( + *, + repo_id: str, + revision: str, + allow_patterns: list[str], + ) -> str: + assert allow_patterns + destination = tmp_path / "snapshots" / str(len(downloaded)) + destination.mkdir(parents=True) + downloaded[destination] = (repo_id, revision) + return str(destination) + + built: dict[str, Path] = {} + + def fake_build_local_artifact( + *, + model_id: str, + checkpoint_dir: Path, + output_root: Path, + source_root: Path, + tokenizer_dir: Path | None, + replace: bool, + ) -> Path: + del source_root + assert not replace + spec = registry[model_id] + selected = spec.artifact_checkpoint + assert downloaded[checkpoint_dir] == (selected.repo_id, selected.revision) + if tokenizer_dir is not None: + assert tokenizer_dir in downloaded + destination = output_root / spec.fast.repo_id.split("/", maxsplit=1)[1] + destination.mkdir(parents=True) + built[model_id] = destination + return destination + + validated: list[tuple[Path, ModelSpec, ModelRegistry]] = [] + + def fake_validate_artifact( + path: Path, + *, + spec: ModelSpec, + registry: ModelRegistry, + ) -> None: + validated.append((path, spec, registry)) + + monkeypatch.setattr(build_all_module, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr( + build_all_module, + "build_local_artifact", + fake_build_local_artifact, + ) + monkeypatch.setattr(build_all_module, "validate_artifact", fake_validate_artifact) + + output_root = tmp_path / "artifacts" + destinations = build_all_module.build_all_artifacts( + output_root=output_root, + source_root=tmp_path / "source", + model_ids=model_ids, + ) + + assert destinations == tuple(built[model_id] for model_id in model_ids) + assert validated == [ + (built[model_id], registry[model_id], registry) for model_id in model_ids + ] + + +def test_ankh_build_all_downloads_and_provenances_every_declared_source_asset( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Required tokenizer/generation assets travel with the selected official state.""" + + registry = get_model_registry() + model_ids = ( + "ankh_base", + "ankh_large", + "ankh2_large", + "ankh3_large", + "ankh3_xl", + ) + downloads: list[tuple[str, str, tuple[str, ...]]] = [] + + def fake_snapshot_download( + *, + repo_id: str, + revision: str, + allow_patterns: list[str], + ) -> str: + downloads.append((repo_id, revision, tuple(allow_patterns))) + destination = tmp_path / "snapshots" / str(len(downloads)) + destination.mkdir(parents=True) + return str(destination) + + def fake_build_local_artifact( + *, + model_id: str, + checkpoint_dir: Path, + output_root: Path, + source_root: Path, + tokenizer_dir: Path | None, + replace: bool, + ) -> Path: + del checkpoint_dir, source_root + assert tokenizer_dir is not None + assert not replace + destination = output_root / model_id + destination.mkdir(parents=True) + return destination + + monkeypatch.setattr(build_all_module, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr( + build_all_module, + "build_local_artifact", + fake_build_local_artifact, + ) + monkeypatch.setattr(build_all_module, "validate_artifact", lambda *_args, **_kwargs: None) + + build_all_module.build_all_artifacts( + output_root=tmp_path / "artifacts", + source_root=tmp_path / "source", + model_ids=model_ids, + ) + + assert len(downloads) == len(model_ids) + for model_id, (repo_id, revision, allow_patterns) in zip( + model_ids, + downloads, + strict=True, + ): + spec = registry[model_id] + assert (repo_id, revision) == ( + spec.artifact_checkpoint.repo_id, + spec.artifact_checkpoint.revision, + ) + assert allow_patterns == tuple(item.path for item in spec.official.files) + provenance = build_module._expected_registry_provenance(registry, spec) + assert provenance["official_checkpoint"]["files"] == { + item.path: item.encoded for item in spec.official.files + } + + assert "generation_config.json" in downloads[2][2] + assert "spiece.model" in downloads[3][2] + assert "generation_config.json" in downloads[4][2] + assert "spiece.model" in downloads[4][2] + assert "pytorch_model.bin.index.json" not in downloads[4][2] + + +def test_benchmark_artifact_selection_is_manifest_derived_and_includes_nested_backbone() -> None: + registry = get_model_registry() + selected = benchmark_artifact_model_ids() + expected = { + spec.id + for spec in registry.values() + if (spec.is_deep_reference or spec.family.id == "esmfold2") + and "benchmark" in spec.family.test_tiers + } + backbone_id = registry.families["esmfold2"].backbone_model + assert backbone_id is not None + expected.add(backbone_id) + + assert set(selected) == expected + assert "esmc_6b" in selected + assert selected == tuple(model_id for model_id in registry if model_id in expected) + + +def test_build_all_rejects_mixed_explicit_and_benchmark_selection(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + build_all_module.build_all_artifacts( + output_root=tmp_path / "artifacts", + source_root=tmp_path / "source", + model_ids=("esm2_8m",), + benchmark_suite=True, + ) + + +def test_single_artifact_cli_revalidates_against_current_registry( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """The normal CLI must not fall back to self-attested artifact validation.""" + + registry = get_model_registry() + spec = registry["ankh_base"] + destination = tmp_path / "ANKH-Base" + destination.mkdir() + monkeypatch.setattr( + build_module, + "_parse_args", + lambda: SimpleNamespace( + model_id=spec.id, + checkpoint_dir=tmp_path / "checkpoint", + output_root=tmp_path / "dist", + source_root=tmp_path, + tokenizer_dir=None, + replace=False, + ), + ) + monkeypatch.setattr(build_module, "get_model_registry", lambda: registry) + monkeypatch.setattr( + build_module, + "build_local_artifact", + lambda **_kwargs: destination, + ) + validated: list[tuple[Path, ModelSpec, ModelRegistry]] = [] + + def validate( + path: Path, + *, + spec: ModelSpec, + registry: ModelRegistry, + ) -> None: + validated.append((path, spec, registry)) + + monkeypatch.setattr(build_module, "validate_artifact", validate) + + build_module.main() + + assert validated == [(destination, spec, registry)] diff --git a/tests/unit/test_container_contract.py b/tests/unit/test_container_contract.py new file mode 100644 index 0000000..daea557 --- /dev/null +++ b/tests/unit/test_container_contract.py @@ -0,0 +1,407 @@ +"""Static container-boundary checks that do not require a Docker daemon.""" + +from __future__ import annotations + +import re +from pathlib import Path + +from fastplms.registry import get_model_registry + + +ROOT = Path(__file__).resolve().parents[2] +DOCKERFILE = ROOT / "docker" / "Dockerfile" +BAKE_FILE = ROOT / "docker" / "docker-bake.hcl" +COMPOSE_FILE = ROOT / "docker" / "compose.yaml" +DOCKERIGNORE = ROOT / ".dockerignore" +REQUIREMENTS = ROOT / "requirements" + + +def _stages(text: str) -> dict[str, str]: + return { + name: base + for base, name in re.findall( + r"^FROM\s+(\S+)\s+AS\s+(\S+)\s*$", + text, + flags=re.MULTILINE | re.IGNORECASE, + ) + } + + +def _stage_section(text: str, stage: str) -> str: + marker = re.compile( + rf"^FROM\s+\S+\s+AS\s+{re.escape(stage)}\s*$", + flags=re.MULTILINE | re.IGNORECASE, + ) + match = marker.search(text) + assert match is not None, f"Missing Docker stage {stage!r}" + next_stage = re.search(r"^FROM\s+", text[match.end() :], flags=re.MULTILINE) + end = match.end() + next_stage.start() if next_stage is not None else len(text) + return text[match.start() : end] + + +def test_manifest_reference_containers_are_build_targets() -> None: + bake = BAKE_FILE.read_text(encoding="utf-8") + targets = set(re.findall(r'^target\s+"([^"]+)"', bake, flags=re.MULTILINE)) + expected = {spec.family.reference_container for spec in get_model_registry().values()} + assert expected.issubset(targets), f"Missing reference targets: {sorted(expected - targets)}" + + +def test_bake_defaults_to_native_host_platform_without_an_amd64_override() -> None: + bake = BAKE_FILE.read_text(encoding="utf-8") + common = bake.split('target "common" {', maxsplit=1)[1].split("\n}", maxsplit=1)[0] + assert "platforms" not in common + assert "linux/amd64" not in bake + + +def test_reference_stages_do_not_inherit_candidate_or_runtime_layers() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + stages = _stages(dockerfile) + reference_stages = {name for name in stages if name.startswith("reference-")} + assert reference_stages + for stage in reference_stages: + ancestry: list[str] = [] + current = stage + while current in stages: + base = stages[current] + ancestry.append(base) + if base not in stages: + break + current = base + candidate_stages = { + "source", + "runtime", + "candidate", + "candidate-structure", + "candidate-fp8", + } + assert not candidate_stages.intersection(ancestry), ( + f"{stage} inherits candidate/runtime layers: {ancestry}" + ) + + +def test_reference_stages_copy_notices_and_no_checkpoint_assets() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + for section in re.split(r"(?=^FROM\s+)", dockerfile, flags=re.MULTILINE): + match = re.match(r"FROM\s+\S+\s+AS\s+(reference-\S+)", section) + if match is None: + continue + stage = match.group(1) + if stage in {"reference-protocol", "reference-esmfold2"}: + # The protocol is copied into notice-bearing final stages; ESMFold2 + # inherits the notice-bearing Biohub reference stage. + continue + assert "THIRD_PARTY_NOTICES.md" in section, f"{stage} omits required notices" + copied_sources = re.findall(r"^COPY\s+(\S+)", section, flags=re.MULTILINE) + weight_suffixes = (".bin", ".ckpt", ".pt", ".pth", ".safetensors") + assert not any(source.endswith(weight_suffixes) for source in copied_sources), ( + f"{stage} copies checkpoint weights" + ) + + +def test_reference_stages_copy_distribution_licenses_for_each_source_context() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + registry = get_model_registry() + context_to_source = { + source.id.replace("-", "_"): source.id for source in registry.upstreams.values() + } + checked: set[str] = set() + for section in re.split(r"(?=^FROM\s+)", dockerfile, flags=re.MULTILINE): + match = re.match(r"FROM\s+\S+\s+AS\s+(reference-\S+)", section) + if match is None: + continue + for context in re.findall(r"--from=upstream_([a-z0-9_]+)", section): + source_id = context_to_source[context] + assert f"COPY LICENSES/{source_id} /licenses/{source_id}" in section, ( + f"{match.group(1)} omits distribution licenses for {source_id}" + ) + checked.add(source_id) + assert checked == set(registry.upstreams) + + +def test_compose_centralizes_gpu_and_ipc_configuration() -> None: + compose = COMPOSE_FILE.read_text(encoding="utf-8") + assert "ipc: host" in compose + assert "driver: nvidia" in compose + assert "count: all" in compose + assert "capabilities: [gpu]" in compose + assert "HF_HOME: /cache/huggingface" in compose + assert "TORCH_HOME: /cache/torch" in compose + assert compose.count('CUBLAS_WORKSPACE_CONFIG: ":4096:8"') == 2 + for volume in ("hf", "torch", "xdg"): + assert f"name: fastplms-{volume}-cache" in compose + + +def test_dependency_tool_caches_match_their_buildkit_mounts() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + assert "UV_CACHE_DIR=/root/.cache/uv" in dockerfile + assert "PIP_CACHE_DIR=/root/.cache/pip" in dockerfile + assert "--mount=type=cache,target=/root/.cache/uv" in dockerfile + assert "--mount=type=cache,target=/root/.cache/pip" in dockerfile + + +def test_reference_esm2_installs_oracle_runtime_dependencies() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + section = dockerfile.split("FROM python312 AS reference-esm2", maxsplit=1)[1].split( + "FROM python310-reference AS reference-boltz2", maxsplit=1 + )[0] + assert "huggingface-hub==0.36.2" in section + assert "numpy==1.26.4" in section + + +def test_reference_services_receive_only_the_exchange_and_cache_mounts() -> None: + compose = COMPOSE_FILE.read_text(encoding="utf-8") + reference_anchor = compose.split("x-reference: &reference", maxsplit=1)[1].split( + "services:", maxsplit=1 + )[0] + assert "../artifacts/reference:/exchange" in reference_anchor + assert "..:/workspace" not in reference_anchor + for service in ( + "reference-ankh", + "reference-biohub-esm", + "reference-boltz2", + "reference-dplm", + "reference-e1", + "reference-esm2", + "reference-esmfold", + "reference-esmfold2", + "reference-protein-ttt", + ): + start = compose.index(f" {service}:") + section = compose[start : start + 240] + assert "<<: *reference" in section + + +def test_compose_and_bake_reference_contexts_are_synchronized() -> None: + compose = COMPOSE_FILE.read_text(encoding="utf-8") + bake = BAKE_FILE.read_text(encoding="utf-8") + expected = { + "reference-ankh": {"upstream_ankh": "vendor/upstream/ankh"}, + "reference-biohub-esm": { + "upstream_biohub_esm": "vendor/upstream/biohub-esm", + "upstream_biohub_transformers": "vendor/upstream/biohub-transformers", + }, + "reference-boltz2": {"upstream_boltz": "vendor/upstream/boltz"}, + "reference-dplm": {"upstream_dplm": "vendor/upstream/dplm"}, + "reference-e1": {"upstream_e1": "vendor/upstream/e1"}, + "reference-esm2": {"upstream_fair_esm": "vendor/upstream/fair-esm"}, + "reference-esmfold": { + "upstream_fair_esm": "vendor/upstream/fair-esm", + "upstream_openfold": "vendor/upstream/openfold", + }, + "reference-esmfold2": { + "upstream_biohub_esm": "vendor/upstream/biohub-esm", + "upstream_biohub_transformers": "vendor/upstream/biohub-transformers", + }, + "reference-protein-ttt": { + "upstream_protein_ttt": "vendor/upstream/protein-ttt" + }, + } + + for service, contexts in expected.items(): + compose_tail = compose.split(f" {service}:\n", maxsplit=1)[1] + next_service = re.search(r"\n [^\s][^:\n]*:\n", compose_tail) + compose_section = ( + compose_tail[: next_service.start()] if next_service is not None else compose_tail + ) + bake_section = bake.split(f'target "{service}" {{', maxsplit=1)[1].split( + "\n}", maxsplit=1 + )[0] + assert "additional_contexts:" in compose_section + assert "contexts = {" in bake_section + for name, relative_path in contexts.items(): + assert f"{name}: ../{relative_path}" in compose_section + assert re.search( + rf"{re.escape(name)}\s*=\s*\"{re.escape(relative_path)}\"", + bake_section, + ) + + +def test_reference_protocol_contains_the_isolated_esmfold2_bundle_producer() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + protocol = dockerfile.split("FROM scratch AS reference-protocol", maxsplit=1)[1].split( + "FROM python310-reference AS reference-ankh", + maxsplit=1, + )[0] + assert "tests/structure/support/esmfold2_bundle.py" in protocol + + +def test_reference_protocol_only_copies_existing_repository_paths() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + protocol = _stage_section(dockerfile, "reference-protocol") + copied_sources = re.findall(r"^COPY\s+(\S+)\s+\S+\s*$", protocol, flags=re.MULTILINE) + assert copied_sources + missing = [source for source in copied_sources if not (ROOT / source).exists()] + assert missing == [] + assert "COPY tools/remote/__init__.py" not in protocol + assert "tests/parity/support/semantic_config.py" in protocol + assert "src/fastplms" not in protocol + + +def test_candidate_dependencies_are_cached_before_source() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + dependency_files = _stage_section(dockerfile, "dependency-files") + assert "COPY requirements ./requirements" in dependency_files + assert "COPY kernels.lock ./" in dependency_files + assert "COPY src" not in dependency_files + assert "COPY tests" not in dependency_files + + expected_profiles = { + "source-dependencies": "runtime.in", + "candidate-dependencies": "candidate.in", + "candidate-structure-dependencies": "candidate-structure.in", + "candidate-fp8-dependencies": "candidate-fp8.in", + } + for dependencies, profile in expected_profiles.items(): + dependency_section = _stage_section(dockerfile, dependencies) + assert "uv pip install --python /opt/venv/bin/python" in dependency_section + assert f"--requirement requirements/profiles/{profile}" in dependency_section + assert "--constraint requirements/constraints/validation.txt" in dependency_section + assert "COPY src" not in dependency_section + assert "COPY tests" not in dependency_section + + for final in ("source", "runtime", "candidate", "candidate-structure", "candidate-fp8"): + final_section = _stage_section(dockerfile, final) + assert "COPY src ./src" in final_section + assert "uv pip install" not in final_section + assert "PYTHONPATH=" in final_section + + +def test_reference_protocol_and_legal_text_do_not_invalidate_dependency_layers() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + for stage in _stages(dockerfile): + if not stage.startswith("reference-") or stage in { + "reference-protocol", + "reference-esmfold2", + }: + continue + section = _stage_section(dockerfile, stage) + last_install = section.rfind("pip install") + if last_install < 0: + continue + for late_copy in ( + "COPY --from=reference-protocol", + "COPY THIRD_PARTY_NOTICES.md", + "COPY LICENSES/", + ): + if late_copy in section: + assert section.index(late_copy) > last_install, ( + f"{stage} copies {late_copy!r} before its dependency layer" + ) + + +def test_runtime_is_one_fail_closed_parameterized_stage() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + sections = { + name: section + for name, section in re.findall( + r"^FROM\s+\S+\s+AS\s+(\S+)\s*$([\s\S]*?)(?=^FROM\s+|\Z)", + dockerfile, + flags=re.MULTILINE | re.IGNORECASE, + ) + } + assert set(name for name in sections if name.startswith("runtime")) == { + "runtime-dependencies", + "runtime", + } + runtime_dependencies = sections["runtime-dependencies"] + assert "ARG FASTPLMS_RUNTIME_PROFILE=core" in runtime_dependencies + assert "core)" in runtime_dependencies + assert "requirements/profiles/runtime.in" in runtime_dependencies + assert "esmfold2-fp8)" in runtime_dependencies + assert "requirements/profiles/runtime-fp8.in" in runtime_dependencies + assert "Unsupported FastPLMs runtime profile" in runtime_dependencies + assert "exit 64" in runtime_dependencies + + runtime = sections["runtime"] + assert "uv pip install" not in runtime + assert "ENV PYTHONPATH=/opt/fastplms/src" in runtime + + bake = BAKE_FILE.read_text(encoding="utf-8") + runtime_targets = { + name: section + for name, section in re.findall( + r'^target\s+"([^"]+)"\s*\{([\s\S]*?)^\}', + bake, + flags=re.MULTILINE, + ) + if name in {"runtime", "runtime-fp8"} + } + assert set(runtime_targets) == {"runtime", "runtime-fp8"} + assert all('target = "runtime"' in section for section in runtime_targets.values()) + assert 'FASTPLMS_RUNTIME_PROFILE = "core"' in runtime_targets["runtime"] + assert ( + 'FASTPLMS_RUNTIME_PROFILE = "esmfold2-fp8"' + in runtime_targets["runtime-fp8"] + ) + + +def test_fp8_dependency_is_confined_to_fp8_container_targets() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + sections = { + name: section + for name, section in re.findall( + r"^FROM\s+\S+\s+AS\s+(\S+)\s*$([\s\S]*?)(?=^FROM\s+|\Z)", + dockerfile, + flags=re.MULTILINE | re.IGNORECASE, + ) + } + assert "requirements/profiles/runtime-fp8.in" in sections["runtime-dependencies"] + assert "requirements/profiles/candidate-fp8.in" in sections[ + "candidate-fp8-dependencies" + ] + assert "requirements/overrides/cuda.txt" in sections["runtime-dependencies"] + assert "requirements/overrides/cuda.txt" in sections["candidate-fp8-dependencies"] + for stage in ( + "source-dependencies", + "candidate-dependencies", + "candidate-structure-dependencies", + "candidate-artifact", + ): + assert "profiles/runtime-fp8.in" not in sections[stage] + assert "profiles/candidate-fp8.in" not in sections[stage] + assert "requirements/overrides/cuda.txt" not in sections[stage] + + +def test_cueq_dependency_is_confined_to_structure_validation_targets() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + assert "requirements/profiles/candidate-structure.in" in _stage_section( + dockerfile, "candidate-structure-dependencies" + ) + assert "requirements/profiles/candidate-fp8.in" in _stage_section( + dockerfile, "candidate-fp8-dependencies" + ) + + structure_profile = ( + REQUIREMENTS / "profiles" / "candidate-structure.in" + ).read_text(encoding="utf-8") + fp8_profile = (REQUIREMENTS / "profiles" / "candidate-fp8.in").read_text( + encoding="utf-8" + ) + assert "-r ../features/cueq.in" in structure_profile + assert "-r candidate-structure.in" in fp8_profile + + for profile in ("runtime.in", "candidate.in", "artifact.in"): + profile_text = (REQUIREMENTS / "profiles" / profile).read_text(encoding="utf-8") + assert "cueq.in" not in profile_text + + +def test_kernel_lock_is_available_to_source_and_artifact_images() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + assert "COPY kernels.lock ./" in _stage_section(dockerfile, "dependency-files") + assert "FROM dependency-files AS candidate-artifact" in dockerfile + assert "!kernels.lock" in DOCKERIGNORE.read_text(encoding="utf-8").splitlines() + assert "!requirements/" in DOCKERIGNORE.read_text(encoding="utf-8").splitlines() + + +def test_candidate_profile_supports_transformers_device_map() -> None: + dev = (REQUIREMENTS / "features" / "dev.in").read_text(encoding="utf-8").splitlines() + assert "accelerate>=1.10,<2" in dev + + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + candidate_profile = (REQUIREMENTS / "profiles" / "candidate.in").read_text( + encoding="utf-8" + ) + assert "-r ../features/dev.in" in candidate_profile + assert "requirements/profiles/candidate.in" in _stage_section( + dockerfile, "candidate-dependencies" + ) diff --git a/tests/unit/test_dplm_rotary.py b/tests/unit/test_dplm_rotary.py new file mode 100644 index 0000000..534ed4c --- /dev/null +++ b/tests/unit/test_dplm_rotary.py @@ -0,0 +1,55 @@ +"""DPLM rotary compatibility with the pinned Transformers release.""" + +from __future__ import annotations + +import torch + +from fastplms.models._esm_rotary import RotaryEmbedding +from fastplms.models.dplm.modeling_dplm import DPLMConfig, ModifiedEsmSelfAttention + + +def test_dplm_initializes_the_checkpoint_compatible_rotary_buffer() -> None: + config = DPLMConfig( + vocab_size=33, + hidden_size=64, + num_attention_heads=4, + num_hidden_layers=1, + intermediate_size=128, + position_embedding_type="rotary", + attn_backend="sdpa", + ) + attention = ModifiedEsmSelfAttention(config) + + assert isinstance(attention.rotary_embeddings, RotaryEmbedding) + assert set(attention.rotary_embeddings.state_dict()) == {"inv_freq"} + expected = 1.0 / ( # (d_h / 2=8,) + 10_000 + ** ( + torch.arange(0, attention.attention_head_size, 2).float() + / attention.attention_head_size + ) + ) + assert torch.equal(attention.rotary_embeddings.inv_freq, expected) + + +def test_dplm_rotary_forward_is_finite() -> None: + config = DPLMConfig( + vocab_size=33, + hidden_size=64, + num_attention_heads=4, + num_hidden_layers=1, + intermediate_size=128, + position_embedding_type="rotary", + attn_backend="sdpa", + ) + attention = ModifiedEsmSelfAttention(config).eval() + hidden_states = torch.randn(2, 11, config.hidden_size) # (b=2, l=11, d=64) + output, weights, s_max = attention( + hidden_states, + attention_mask_2d=torch.ones(2, 11, dtype=torch.bool), # (b=2, l=11) + ) # output: (b=2, l=11, d=64); weights: None; s_max: None + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert weights is None + assert s_max is None diff --git a/tests/unit/test_e1_cache_contract.py b/tests/unit/test_e1_cache_contract.py new file mode 100644 index 0000000..2975b6f --- /dev/null +++ b/tests/unit/test_e1_cache_contract.py @@ -0,0 +1,527 @@ +"""Fast CPU contracts for E1 retrieval caching and load context.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from threading import Barrier +from typing import Any +from transformers import PreTrainedModel +from transformers.modeling_outputs import ModelOutput + +from fastplms.models.e1 import modeling_e1 as e1_modeling +from fastplms.models.e1.cache import DynamicCache, KVCache +from fastplms.models.e1.modeling_e1 import ( + Attention, + AttentionLayerType, + E1Config, + E1ForMaskedLM, + E1ForSequenceClassification, + E1ForTokenClassification, + E1Model, +) + + +@dataclass +class _CacheableE1Output(ModelOutput): + logits: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + embeddings: torch.Tensor | None = None + token_embeddings: torch.Tensor | None = None + past_key_values: DynamicCache | None = None + hidden_states: tuple[torch.Tensor, ...] | None = None + + +def _dynamic_cache(sequence_length: int) -> DynamicCache: + cache = DynamicCache() + states = torch.arange(sequence_length * 2, dtype=torch.float32).reshape( # (b=1, l, h=1, d_h=2) + 1, sequence_length, 1, 2 + ) + cache.update(states, states + 1, layer_idx=0) + return cache + + +def _cache_batch(sequence_length: int = 5, context_length: int = 3) -> dict[str, Any]: + token_values = torch.arange(sequence_length).unsqueeze(0) + return { + "context": ["ACD"], + "context_len": [context_length], + "use_cache": True, + "input_ids": token_values.clone(), + "within_seq_position_ids": token_values.clone(), + "global_position_ids": token_values.clone(), + "sequence_ids": token_values.clone(), + "labels": token_values.clone(), + } + + +def test_e1_cache_miss_slices_every_sequence_aligned_output_alias() -> None: + cache = KVCache(cache_size=1) + batch = _cache_batch() + output_values = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1) + outputs = _CacheableE1Output( + logits=output_values + 10, + last_hidden_state=output_values + 20, + embeddings=output_values + 30, + token_embeddings=output_values + 40, + past_key_values=_dynamic_cache(sequence_length=5), + hidden_states=(output_values + 50, output_values + 60), + ) + + cache.after_forward(batch, outputs) + + for field_name in cache.tensor_input_field_names: + assert batch[field_name].shape[1] == 2 + for field_name in cache.tensor_output_field_names: + value = outputs[field_name] + assert value is not None + assert value.shape[1] == 2 + assert isinstance(outputs.hidden_states, tuple) + assert all(hidden_state.shape[1] == 2 for hidden_state in outputs.hidden_states) + assert cache.cache_dict["ACD"].get_seq_length() == 3 + + +def test_e1_cache_hit_does_not_slice_target_outputs_twice() -> None: + cache = KVCache(cache_size=1) + miss_batch = _cache_batch() + miss_outputs = _CacheableE1Output( + last_hidden_state=torch.zeros(1, 5, 2), + past_key_values=_dynamic_cache(sequence_length=5), + ) + cache.after_forward(miss_batch, miss_outputs) + + hit_batch = _cache_batch() + cache.before_forward(hit_batch) + assert hit_batch["input_ids"].shape[1] == 2 + assert hit_batch["past_key_values"] is cache.cache_dict["ACD"] + + hit_outputs = _CacheableE1Output( + last_hidden_state=torch.ones(1, 2, 2), + past_key_values=cache.cache_dict["ACD"], + ) + cache.after_forward(hit_batch, hit_outputs) + + assert hit_outputs.last_hidden_state is not None + assert hit_outputs.last_hidden_state.shape[1] == 2 + + +def test_e1_from_pretrained_tokenizer_context_is_thread_local( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Concurrent loads must not exchange source, revision, cache, or token settings.""" + + barrier = Barrier(2) + + def fake_from_pretrained( + cls: type[E1ForMaskedLM], + pretrained_model_name_or_path: str, + *model_args: Any, + **kwargs: Any, + ) -> dict[str, Any]: + del pretrained_model_name_or_path, model_args, kwargs + barrier.wait(timeout=5) + observed = cls._tokenizer_kwargs_from_config(E1Config()) + barrier.wait(timeout=5) + return observed + + monkeypatch.setattr( + PreTrainedModel, + "from_pretrained", + classmethod(fake_from_pretrained), + ) + + load_specs = ( + ("model-a", "cache-a", "revision-a", "token-a"), + ("model-b", "cache-b", "revision-b", "token-b"), + ) + + def load(spec: tuple[str, str, str, str]) -> dict[str, Any]: + source, cache_dir, revision, token = spec + return E1ForMaskedLM.from_pretrained( + source, + local_files_only=True, + cache_dir=cache_dir, + revision=revision, + token=token, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + observed = list(executor.map(load, load_specs)) + + for load_spec, tokenizer_kwargs in zip(load_specs, observed, strict=True): + source, cache_dir, revision, token = load_spec + assert tokenizer_kwargs == { + "tokenizer_source": source, + "local_files_only": True, + "cache_dir": cache_dir, + "revision": revision, + "token": token, + } + + +def test_e1_lazy_tokenizer_uses_resolved_weight_commit_per_instance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Moving weight revisions must resolve to the tokenizer's immutable commit.""" + + load_barrier = Barrier(2) + access_barrier = Barrier(2) + commits = { + "model-a": "a" * 40, + "model-b": "b" * 40, + } + + def fake_from_pretrained( + cls: type[E1ForMaskedLM], + pretrained_model_name_or_path: str, + *model_args: Any, + **kwargs: Any, + ) -> E1ForMaskedLM: + del model_args, kwargs + config = _tiny_e1_config() + config._commit_hash = commits[pretrained_model_name_or_path] + load_barrier.wait(timeout=5) + return cls(config) + + requests: list[dict[str, Any]] = [] + + class RecordingPreparer: + def __init__(self, *, data_prep_config: Any, **kwargs: Any) -> None: + del data_prep_config + access_barrier.wait(timeout=5) + requests.append(kwargs) + + monkeypatch.setattr( + PreTrainedModel, + "from_pretrained", + classmethod(fake_from_pretrained), + ) + monkeypatch.setattr(e1_modeling, "E1BatchPreparer", RecordingPreparer) + + load_specs = ( + ("model-a", "cache-a", None, "token-a"), + ("model-b", "cache-b", "main", "token-b"), + ) + + def load(spec: tuple[str, str, str | None, str]) -> E1ForMaskedLM: + source, cache_dir, revision, token = spec + return E1ForMaskedLM.from_pretrained( + source, + local_files_only=True, + cache_dir=cache_dir, + revision=revision, + token=token, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + models = list(executor.map(load, load_specs)) + with ThreadPoolExecutor(max_workers=2) as executor: + preparers = list(executor.map(lambda model: model.prep_tokens, models)) + + assert all(isinstance(preparer, RecordingPreparer) for preparer in preparers) + observed = {request["tokenizer_source"]: request for request in requests} + for source, cache_dir, _requested_revision, token in load_specs: + assert observed[source] == { + "tokenizer_source": source, + "local_files_only": True, + "cache_dir": cache_dir, + "revision": commits[source], + "token": token, + } + + +def _tiny_e1_config() -> E1Config: + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + max_num_sequences=4, + max_num_positions_within_seq=32, + max_num_positions_global=64, + attn_backend="sdpa", + dtype="float32", + num_labels=3, + ) + config.output_hidden_states = True + config.use_cache = True + return config + + +def _tiny_e1_batch() -> dict[str, torch.Tensor]: + return { + "input_ids": torch.tensor([[1, 5, 6, 2]], dtype=torch.long), + "within_seq_position_ids": torch.arange(4).unsqueeze(0), + "global_position_ids": torch.arange(4).unsqueeze(0), + "sequence_ids": torch.zeros(1, 4, dtype=torch.long), + } + + +def test_e1_config_round_trip_preserves_cache_policy(tmp_path: Path) -> None: + config = _tiny_e1_config() + + config.save_pretrained(tmp_path) + restored = E1Config.from_pretrained(tmp_path, local_files_only=True) + + assert restored.use_cache is True + + +def test_e1_encoder_embedding_filters_training_only_preparer_fields() -> None: + model = E1Model(_tiny_e1_config()).eval() + + with torch.inference_mode(): + hidden, token_mask = model._embed(["ACD", "G"], return_attention_mask=True) + + assert hidden.shape[:2] == token_mask.shape + assert token_mask.sum(dim=-1).tolist() == [7, 5] + + +def _assert_nested_output_close(actual: Any, expected: Any) -> None: + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor) + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, DynamicCache): + assert isinstance(actual, DynamicCache) + assert len(actual.key_cache) == len(expected.key_cache) + assert len(actual.value_cache) == len(expected.value_cache) + for actual_tensor, expected_tensor in zip( + actual.key_cache + actual.value_cache, + expected.key_cache + expected.value_cache, + strict=True, + ): + torch.testing.assert_close(actual_tensor, expected_tensor) + return + if isinstance(expected, (tuple, list)): + assert isinstance(actual, type(expected)) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_nested_output_close(actual_item, expected_item) + return + assert actual == expected + + +@pytest.mark.parametrize( + "model_class", + [E1Model, E1ForMaskedLM, E1ForSequenceClassification, E1ForTokenClassification], +) +def test_e1_public_models_honor_config_output_flags_and_return_dict( + model_class: type[PreTrainedModel], +) -> None: + model = model_class(_tiny_e1_config()).eval() + batch = _tiny_e1_batch() + + with torch.inference_mode(): + structured = model(**batch, return_dict=True) + tuple_output = model(**batch, return_dict=False) + + assert isinstance(structured, ModelOutput) + assert structured.hidden_states is not None + assert structured.past_key_values is not None + assert isinstance(tuple_output, tuple) + assert len(tuple_output) == len(structured.to_tuple()) + assert isinstance(tuple_output[0], torch.Tensor) + assert tuple_output[0].shape == structured.to_tuple()[0].shape + + +@pytest.mark.parametrize( + "model_class", + [E1ForMaskedLM, E1ForSequenceClassification, E1ForTokenClassification], +) +def test_e1_loss_bearing_head_tuples_start_with_loss_then_logits( + model_class: type[PreTrainedModel], +) -> None: + model = model_class(_tiny_e1_config()).eval() + batch = _tiny_e1_batch() + if model_class is E1ForSequenceClassification: + labels = torch.tensor([1], dtype=torch.long) + elif model_class is E1ForTokenClassification: + labels = batch["input_ids"].remainder(model.config.num_labels) + else: + labels = batch["input_ids"].clone() + + structured = model(**batch, labels=labels, return_dict=True) + tuple_output = model(**batch, labels=labels, return_dict=False) + + assert structured.loss is not None + torch.testing.assert_close(tuple_output[0], structured.loss) + torch.testing.assert_close(tuple_output[1], structured.logits) + if model_class is E1ForMaskedLM: + assert structured.mlm_loss is not None + assert structured.to_tuple()[2] is structured.hidden_states + elif model_class is E1ForSequenceClassification: + assert structured.to_tuple()[2] is structured.past_key_values + else: + assert structured.to_tuple()[2] is structured.hidden_states + assert len(tuple_output) == len(structured.to_tuple()) + for actual, expected in zip(tuple_output, structured.to_tuple(), strict=True): + _assert_nested_output_close(actual, expected) + + +@pytest.mark.parametrize( + "model_class", + [E1Model, E1ForMaskedLM, E1ForSequenceClassification, E1ForTokenClassification], +) +def test_e1_public_forwards_reject_unknown_arguments( + model_class: type[PreTrainedModel], +) -> None: + model = model_class(_tiny_e1_config()).eval() + with pytest.raises(TypeError, match="unexpected_argument"): + model(**_tiny_e1_batch(), unexpected_argument=True) + + +def test_e1_base_model_rejects_misaligned_biological_indices() -> None: + model = E1Model(_tiny_e1_config()).eval() + batch = _tiny_e1_batch() + + with pytest.raises(ValueError, match="Cannot specify both"): + model(**batch, inputs_embeds=torch.zeros(1, 4, model.config.hidden_size)) + with pytest.raises(ValueError, match="sequence_ids must have shape"): + model(**{**batch, "sequence_ids": torch.zeros(2, 2, dtype=torch.long)}) + with pytest.raises(ValueError, match="Global position ids must be in the range"): + model( + **{ + **batch, + "global_position_ids": torch.full((1, 4), model.config.max_num_positions_global), + } + ) + with pytest.raises(ValueError, match="Sequence ids must be in the range"): + model( + **{ + **batch, + "sequence_ids": torch.full((1, 4), model.config.max_num_sequences), + } + ) + + +def test_e1_masked_lm_resizes_input_and_output_embeddings_together() -> None: + model = E1ForMaskedLM(_tiny_e1_config()).eval() + + resized_input = model.resize_token_embeddings(39) + + assert resized_input.num_embeddings == 39 + assert model.get_input_embeddings().num_embeddings == 39 + assert model.get_output_embeddings().out_features == 39 + assert model.config.vocab_size == 39 + assert model.vocab_size == 39 + + batch = _tiny_e1_batch() + batch["input_ids"][0, 1] = 38 + with torch.inference_mode(): + output = model(**batch) + assert output.logits.shape == (1, 4, 39) + + +def test_e1_legacy_backend_setter_rejects_unadvertised_backends() -> None: + model = E1Model(_tiny_e1_config()).eval() + + with pytest.raises(ValueError, match="E1 does not support 'eager'"): + model.attn_backend = "eager" + + +def _tiny_attention(layer_type: AttentionLayerType) -> Attention: + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_num_sequences=8, + max_num_positions_within_seq=32, + max_num_positions_global=32, + global_attention_every_n_layers=(1 if layer_type == AttentionLayerType.GLOBAL else 0), + dtype="float32", + attn_backend="sdpa", + ) + attention = Attention(config, layer_idx=0).eval() + assert attention.layer_type == layer_type + return attention + + +@pytest.mark.parametrize( + "layer_type", + (AttentionLayerType.WITHIN_SEQ, AttentionLayerType.GLOBAL), +) +def test_e1_cached_sdpa_preserves_layer_attention_semantics( + layer_type: AttentionLayerType, +) -> None: + torch.manual_seed(11) + attention = _tiny_attention(layer_type) + query = torch.randn(1, 2, 2, 4) + key = torch.randn(1, 5, 2, 4) + value = torch.randn(1, 5, 2, 4) + sequence_ids = torch.tensor([[1, 1]]) + + actual, _ = attention._sdpa_attn( + query, + key, + value, + sequence_ids=sequence_ids, + effective_layer_type=layer_type, + is_cache_prefilled=True, + ) + + expected_key = key[:, -2:] if layer_type == AttentionLayerType.WITHIN_SEQ else key + expected_value = value[:, -2:] if layer_type == AttentionLayerType.WITHIN_SEQ else value + mask = attention._cached_attention_mask_4d( + sequence_ids, + key.shape[1], + layer_type, + ) + expected_heads = F.scaled_dot_product_attention( + query.transpose(1, 2), + expected_key.transpose(1, 2), + expected_value.transpose(1, 2), + attn_mask=mask, + ) + expected = expected_heads.transpose(1, 2).reshape(1, 2, 8) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize( + ("layer_type", "expected_path", "expected_kv_length"), + ( + (AttentionLayerType.WITHIN_SEQ, "dense", 2), + (AttentionLayerType.GLOBAL, "packed", 5), + ), +) +def test_e1_cached_flex_dispatch_keeps_or_discards_context_by_layer( + monkeypatch: pytest.MonkeyPatch, + layer_type: AttentionLayerType, + expected_path: str, + expected_kv_length: int, +) -> None: + attention = _tiny_attention(layer_type) + query = torch.randn(1, 2, 2, 4) + key = torch.randn(1, 5, 2, 4) + value = torch.randn(1, 5, 2, 4) + observed: list[tuple[str, int]] = [] + + def fake_dense(q, k, v, **kwargs): + del v, kwargs + observed.append(("dense", k.shape[1])) + return torch.zeros_like(q) + + def fake_packed(q, k, v, **kwargs): + del v, kwargs + observed.append(("packed", k.shape[1])) + return torch.zeros_like(q) + + monkeypatch.setattr(e1_modeling, "flex_attention_func", fake_dense) + monkeypatch.setattr(e1_modeling, "varlen_flex_attention_func", fake_packed) + + attention._flex_attn( + query, + key, + value, + sequence_ids=torch.tensor([[1, 1]]), + effective_layer_type=layer_type, + is_cache_prefilled=True, + ) + + assert observed == [(expected_path, expected_kv_length)] diff --git a/tests/unit/test_e1_rotary.py b/tests/unit/test_e1_rotary.py new file mode 100644 index 0000000..107fb8d --- /dev/null +++ b/tests/unit/test_e1_rotary.py @@ -0,0 +1,30 @@ +"""Contracts for E1 rotary-position buffers.""" + +from __future__ import annotations + +import torch + +from fastplms.models.e1.modeling_e1 import RotaryPositionalEmbedding + + +def test_e1_rotary_lazily_initializes_after_meta_materialization() -> None: + """Meta-device construction cannot leave uninitialized trigonometric caches.""" + + with torch.device("meta"): + rotary = RotaryPositionalEmbedding(dim=8, max_position_embeddings=16) + rotary = rotary.to_empty(device="cpu") + + Q = torch.randn(2, 5, 3, 8) # (b=2, l=5, h_q=3, d_h=8) + K = torch.randn(2, 5, 1, 8) # (b=2, l=5, h_k=1, d_h=8) + position_ids = torch.tensor( # (b=2, l=5) + [[0, 1, 2, 3, 4], [0, 1, 2, 3, -1]] + ) + Q_rotated, K_rotated = rotary(Q, K, position_ids) # Q: (2, 5, 3, 8); K: (2, 5, 1, 8) + + assert torch.isfinite(Q_rotated).all() + assert torch.isfinite(K_rotated).all() + assert rotary.max_seq_len_cached == 5 + assert rotary.cos_cached.shape == (5, 8) + assert rotary.sin_cached.shape == (5, 8) + assert rotary.cos_cached.abs().max() <= 1 + assert rotary.sin_cached.abs().max() <= 1 diff --git a/tests/unit/test_embeddings_api.py b/tests/unit/test_embeddings_api.py new file mode 100644 index 0000000..e57528c --- /dev/null +++ b/tests/unit/test_embeddings_api.py @@ -0,0 +1,1957 @@ +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import struct +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace +from torch import nn + +from fastplms.embeddings import ( + EmbeddingBatch, + EmbeddingInput, + EmbeddingRecord, + EmbeddingResult, + LazyTensorReference, + Pooler, + convert_legacy_sqlite, + embed_dataset, + garbage_collect_safetensors_generations, + iter_fasta, + load_legacy_pth, + load_safetensors_result, + load_sqlite_result, + pagerank_weights, + parse_fasta, + save_safetensors_result, + save_sqlite_result, +) +from fastplms.embeddings.storage import SafetensorsStreamWriter + + +class SyntheticEmbeddingModel(nn.Module): + def __init__(self, backend: str = "eager") -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + model_type="synthetic", + _name_or_path="synthetic/checkpoint", + _commit_hash="abc123", + _attn_implementation=backend, + ) + + def _embedding_batch(self, sequences: list[str]) -> EmbeddingBatch: + b = len(sequences) + sequence_length = max(map(len, sequences)) + 2 + X = torch.zeros(b, sequence_length, 2) + M = torch.zeros(b, sequence_length, dtype=torch.bool) + for batch_index, sequence in enumerate(sequences): + for residue_index, residue in enumerate(sequence, start=1): + X[batch_index, residue_index] = torch.tensor( + [float(ord(residue)), float(residue_index)] + ) + M[batch_index, residue_index] = True + # A uses two heads and includes BOS/EOS rows that M removes. + A = torch.ones(b, 2, sequence_length, sequence_length) + return EmbeddingBatch(X=X, residue_mask=M, attentions=(A, A * 2)) + + +class InterruptibleEmbeddingModel(SyntheticEmbeddingModel): + def __init__(self, fail_on_call: int | None) -> None: + super().__init__() + self.fail_on_call = fail_on_call + self.calls = 0 + + def _embedding_batch(self, sequences: list[str]) -> EmbeddingBatch: + self.calls += 1 + if self.calls == self.fail_on_call: + raise RuntimeError("simulated interruption") + return super()._embedding_batch(sequences) + + +class TrainingAwareEmbeddingModel(SyntheticEmbeddingModel): + def __init__(self) -> None: + super().__init__() + self.observed_training: list[bool] = [] + + def _embedding_batch(self, sequences: list[str]) -> EmbeddingBatch: + self.observed_training.append(self.training) + return super()._embedding_batch(sequences) + + +class SyntheticAllStatesModel(SyntheticEmbeddingModel): + def _embedding_batch( + self, + sequences: list[str], + *, + store_all_hidden_states: bool = False, + ) -> EmbeddingBatch: + assert store_all_hidden_states is True + batch = super()._embedding_batch(sequences) + return EmbeddingBatch( + X=torch.stack((batch.X, batch.X + 100), dim=1), + residue_mask=batch.residue_mask, + ) + + +class SyntheticE1Preparer: + boundary_token_ids = torch.tensor([0, 1, 2, 3]) + + def get_batch_kwargs(self, sequences, device): + assert sequences == ["AC"] + return {"input_ids": torch.tensor([[0, 1, 5, 6, 2, 3]], device=device)} + + +class SyntheticE1Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace(model_type="e1") + self.prep_tokens = SyntheticE1Preparer() + + def _embed(self, sequences, return_attention_mask, **kwargs): + del kwargs + prepared = self.prep_tokens.get_batch_kwargs(sequences, self.anchor.device) + X = prepared["input_ids"].float().unsqueeze(-1) + M = torch.ones(X.shape[:2], dtype=torch.bool) + assert return_attention_mask is True + return X, M + + +class SyntheticDecoderEmbeddingModel(SyntheticEmbeddingModel): + def __init__(self) -> None: + super().__init__() + self.config.model_type = "fast_ankh" + self.observed_decoder_rows: list[list[int]] = [] + + def _embedding_metadata(self, **context): + return { + "hidden_state_stack": context["hidden_state_source"], + "source": context["hidden_state_source"], + } + + def _embedding_batch( + self, + sequences: list[str], + *, + tokenizer, + max_length, + truncate, + need_attentions, + hidden_state_source, + decoder_input_ids=None, + decoder_attention_mask=None, + **kwargs, + ) -> EmbeddingBatch: + del tokenizer, max_length, truncate, need_attentions, kwargs + assert hidden_state_source == "decoder" + assert decoder_input_ids is not None + assert decoder_attention_mask is not None + self.observed_decoder_rows.append(decoder_input_ids[:, 0].tolist()) + X = decoder_input_ids.to(dtype=torch.float32).unsqueeze(-1) + return EmbeddingBatch(X=X, residue_mask=decoder_attention_mask.to(dtype=torch.bool)) + + +def test_result_preserves_order_and_duplicates() -> None: + model = SyntheticEmbeddingModel() + inputs = [ + EmbeddingInput("first", "ACD"), + EmbeddingInput("second", "GG"), + EmbeddingInput("first", "ACD"), + ] + result = embed_dataset(model, inputs, batch_size=2, pooling="mean") + + assert [record.id for record in result] == ["first", "second", "first"] + assert [record.sequence for record in result] == ["ACD", "GG", "ACD"] + assert torch.equal(result[0].load_tensor(), result[2].load_tensor()) + assert len(result.metadata["tensor_hashes"]) == 3 + assert result.metadata["outputs"][0]["sha256"] == result.metadata["tensor_hashes"][0] + assert result.metadata["token_policy"]["unit"] == "residue" + assert result.metadata["layer"] == -1 + with pytest.raises(ValueError, match="Duplicate id"): + result.as_dict() + assert list(result.as_dict(duplicates="first")) == ["first", "second"] + + +def test_bounded_length_bucketing_restores_input_order() -> None: + model = SyntheticEmbeddingModel() + observed: list[list[str]] = [] + original = model._embedding_batch + + def recording_batch(sequences: list[str]) -> EmbeddingBatch: + observed.append(list(sequences)) + return original(sequences) + + model._embedding_batch = recording_batch # type: ignore[method-assign] + inputs = ["A", "BBBB", "CC", "DDD"] + result = embed_dataset( + model, + inputs, + batch_size=2, + max_tokens_per_batch=8, + ) + + assert observed == [["BBBB", "DDD"], ["CC", "A"]] + assert [record.sequence for record in result] == inputs + assert result.metadata["batching"]["batch_window_size"] == 32 + assert result.metadata["batching"]["ordering"] == ("bounded-length-bucketed-stable-output") + + +def test_invalid_storage_and_pooling_fail_before_input_consumption(tmp_path: Path) -> None: + consumed = False + + def inputs(): + nonlocal consumed + consumed = True + yield "ACD" + + with pytest.raises(ValueError, match="format must"): + embed_dataset( + SyntheticEmbeddingModel(), + inputs(), + output=tmp_path / "output", + format="unknown", + ) + assert consumed is False + + with pytest.raises(ValueError, match="cannot be combined"): + embed_dataset( + SyntheticEmbeddingModel(), + ["ACD"], + full_embeddings=True, + pooling="mean", + ) + + +def test_decoder_companions_are_fingerprinted_and_bucket_aligned() -> None: + model = SyntheticDecoderEmbeddingModel() + inputs = ["A", "BBBB", "CC"] + decoder_input_ids = torch.tensor([[11, 11], [22, 22], [33, 33]]) + decoder_attention_mask = torch.ones_like(decoder_input_ids) + + result = embed_dataset( + model, + inputs, + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + batch_size=2, + batch_window_size=3, + ) + + assert model.observed_decoder_rows == [[22, 33], [11]] + assert [record.load_tensor().item() for record in result] == [11.0, 22.0, 33.0] + assert result.metadata["hidden_state_source"] == "decoder" + assert result.metadata["decoder_input_fingerprint"] + assert result.metadata["decoder_attention_mask_fingerprint"] + assert result.metadata["decoder_alignment"] == "input-position" + assert result.metadata["model_embedding"] == { + "hidden_state_stack": "decoder", + "source": "decoder", + } + + with pytest.raises(ValueError, match="exactly one"): + embed_dataset(model, inputs, hidden_state_source="decoder") + + +def test_decoder_embeddings_require_an_explicit_model_capability() -> None: + with pytest.raises(ValueError, match="does not declare decoder embedding support"): + embed_dataset( + SyntheticEmbeddingModel(), + ["AC"], + hidden_state_source="decoder", + decoder_input_ids=torch.tensor([[3, 4]]), + ) + + +def test_mapping_inputs_embed_values_with_mapping_keys_as_ids() -> None: + inputs = { + "protein-a": "ACD", + "protein-b": "GG", + } + + result = embed_dataset(SyntheticEmbeddingModel(), inputs, pooling="mean") + + assert [record.id for record in result] == ["protein-a", "protein-b"] + assert [record.sequence for record in result] == ["ACD", "GG"] + assert result[0].load_tensor()[0].item() == pytest.approx(sum(map(ord, "ACD")) / len("ACD")) + with pytest.raises(ValueError, match="at least one sequence"): + embed_dataset(SyntheticEmbeddingModel(), {}, pooling="mean") + + +def test_embedding_temporarily_uses_eval_and_restores_training_state() -> None: + model = TrainingAwareEmbeddingModel() + model.train() + + embed_dataset(model, ["ACD"]) + + assert model.observed_training == [False] + assert model.training is True + + interrupted = InterruptibleEmbeddingModel(fail_on_call=1) + interrupted.train() + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset(interrupted, ["ACD"]) + assert interrupted.training is True + + +def test_full_embeddings_contain_biological_residues_only() -> None: + result = embed_dataset( + SyntheticEmbeddingModel(), + ["ACD", "GG"], + full_embeddings=True, + ) + assert tuple(result[0].load_tensor().shape) == (3, 2) + assert tuple(result[1].load_tensor().shape) == (2, 2) + assert result.metadata["residue_mask_policy"] == "biological-residues-only" + + +@pytest.mark.parametrize("format", ("safetensors", "sqlite")) +def test_all_hidden_state_embeddings_trim_token_axis_and_round_trip( + tmp_path: Path, + format: str, +) -> None: + output = tmp_path / ("all-states.sqlite" if format == "sqlite" else "all-states") + result = embed_dataset( + SyntheticAllStatesModel(), + ["ACD", "GG"], + full_embeddings=True, + store_all_hidden_states=True, + output=output, + format=format, + ) + + assert tuple(result[0].load_tensor().shape) == (2, 3, 2) + assert tuple(result[1].load_tensor().shape) == (2, 2, 2) + assert torch.equal(result[0].load_tensor()[1], result[0].load_tensor()[0] + 100) + loaded = load_sqlite_result(output) if format == "sqlite" else load_safetensors_result(output) + assert torch.equal(loaded[0].load_tensor(), result[0].load_tensor()) + assert loaded.metadata["record_count"] == 2 + assert loaded.metadata["descriptor_index"] in { + "sqlite-records", + "safetensors-generation-index", + } + assert "outputs" not in loaded.metadata + assert "tensor_hashes" not in loaded.metadata + + +def test_all_hidden_states_require_full_embeddings() -> None: + with pytest.raises( + ValueError, + match="store_all_hidden_states=True requires full_embeddings=True", + ): + embed_dataset( + SyntheticAllStatesModel(), + ["ACD"], + store_all_hidden_states=True, + ) + + +def test_embedding_fingerprint_records_loaded_esmc_identity() -> None: + model = SyntheticEmbeddingModel() + model._esmc_source = "Synthyra/ESMplusplus_6B" + model._esmc_source_revision = "a" * 40 + model._esmc_source_files = {"model.safetensors": "sha256:" + "b" * 64} + + result = embed_dataset(model, ["ACD"], pooling="mean") + + assert result.metadata["esmc_source"] == model._esmc_source + assert result.metadata["esmc_revision"] == model._esmc_source_revision + assert result.metadata["esmc_files"] == model._esmc_source_files + + model._esmc_source_revision = "c" * 40 + changed = embed_dataset(model, ["ACD"], pooling="mean") + assert changed.metadata["run_fingerprint"] != result.metadata["run_fingerprint"] + + +def test_embedding_fingerprint_binds_persisted_parameters_and_buffers( + tmp_path: Path, +) -> None: + model = SyntheticEmbeddingModel() + model.register_buffer("running_value", torch.tensor([3.0])) + initial = embed_dataset(model, ["ACD"], output=tmp_path / "initial") + + with torch.no_grad(): + model.anchor.fill_(1) + changed_parameter = embed_dataset(model, ["ACD"], output=tmp_path / "parameter") + assert ( + changed_parameter.metadata["model_state_fingerprint"] + != (initial.metadata["model_state_fingerprint"]) + ) + assert changed_parameter.metadata["run_fingerprint"] != initial.metadata["run_fingerprint"] + + model.running_value.add_(1) + changed_buffer = embed_dataset(model, ["ACD"], output=tmp_path / "buffer") + assert ( + changed_buffer.metadata["model_state_fingerprint"] + != (changed_parameter.metadata["model_state_fingerprint"]) + ) + assert initial.metadata["fingerprint_schema_version"] == 3 + assert initial.metadata["model_state_fingerprint_source"] == "computed" + + +def test_model_state_fingerprint_rehashes_data_and_storage_alias_mutations( + tmp_path: Path, +) -> None: + model = SyntheticEmbeddingModel() + original_path = tmp_path / "original" + original = embed_dataset(model, ["ACD"], output=original_path) + + # ``Parameter.data`` mutation bypasses autograd's version counter. A state + # fingerprint must still derive from current bytes rather than object/version + # metadata retained by a cache. + model.anchor.data.fill_(1) + data_mutated = embed_dataset(model, ["ACD"], output=tmp_path / "data-mutated") + assert ( + data_mutated.metadata["model_state_fingerprint"] + != (original.metadata["model_state_fingerprint"]) + ) + assert data_mutated.metadata["run_fingerprint"] != original.metadata["run_fingerprint"] + with pytest.raises(ValueError, match="different run fingerprint"): + embed_dataset(model, ["ACD"], output=original_path) + + storage_alias = model.anchor.data + storage_alias.fill_(2) + alias_mutated = embed_dataset(model, ["ACD"], output=tmp_path / "alias-mutated") + assert ( + alias_mutated.metadata["model_state_fingerprint"] + != (data_mutated.metadata["model_state_fingerprint"]) + ) + assert alias_mutated.metadata["run_fingerprint"] != (data_mutated.metadata["run_fingerprint"]) + + +def test_in_memory_embedding_skips_model_state_hash() -> None: + class NoStateDictEmbeddingModel(SyntheticEmbeddingModel): + def state_dict( + self, + *args: object, + **kwargs: object, + ) -> dict[str, torch.Tensor]: + del args, kwargs + raise AssertionError("in-memory embeddings must not hash the full model state") + + result = embed_dataset(NoStateDictEmbeddingModel(), ["ACD"]) + + assert result.metadata["model_state_fingerprint"] is None + assert result.metadata["model_state_fingerprint_source"] == "not-computed" + + +def test_caller_owned_model_state_fingerprint_overrides_state_hash() -> None: + model = SyntheticEmbeddingModel() + first = embed_dataset(model, ["ACD"], model_state_fingerprint="external-state-v1") + with torch.no_grad(): + model.anchor.fill_(9) + second = embed_dataset(model, ["ACD"], model_state_fingerprint="external-state-v1") + + assert second.metadata["run_fingerprint"] == first.metadata["run_fingerprint"] + assert second.metadata["model_state_fingerprint"] == "external-state-v1" + assert second.metadata["model_state_fingerprint_source"] == "caller" + with pytest.raises(ValueError, match="must not be empty"): + embed_dataset(model, ["ACD"], model_state_fingerprint=" ") + + +def test_runtime_versions_are_part_of_resume_identity(monkeypatch) -> None: + import fastplms.embeddings.runner as runner + + model = SyntheticEmbeddingModel() + first = embed_dataset(model, ["ACD"]) + versions = runner._software_versions() + monkeypatch.setattr( + runner, + "_software_versions", + lambda: {**versions, "torch": "different-runtime"}, + ) + changed = embed_dataset(model, ["ACD"]) + + assert changed.metadata["run_fingerprint"] != first.metadata["run_fingerprint"] + + +def test_tokenizer_content_changes_run_fingerprint() -> None: + class Tokenizer: + all_special_ids = (0, 2) + name_or_path = "synthetic/tokenizer" + vocab_size = 4 + model_max_length = 32 + padding_side = "right" + truncation_side = "right" + + def __init__(self, vocab, *, mode: str = "first") -> None: + self._vocab = vocab + self.init_kwargs = {"mode": mode} + self.special_tokens_map = {"bos_token": "", "eos_token": ""} + + def get_vocab(self): + return self._vocab + + model = SyntheticEmbeddingModel() + first = embed_dataset( + model, + ["ACD"], + tokenizer=Tokenizer({"": 0, "A": 1, "": 2, "D": 3}), + ) + changed_vocab = embed_dataset( + model, + ["ACD"], + tokenizer=Tokenizer({"": 0, "D": 1, "": 2, "A": 3}), + ) + changed_config = embed_dataset( + model, + ["ACD"], + tokenizer=Tokenizer( + {"": 0, "A": 1, "": 2, "D": 3}, + mode="second", + ), + ) + + assert ( + first.metadata["tokenizer"]["content_sha256"] + != (changed_vocab.metadata["tokenizer"]["content_sha256"]) + ) + assert first.metadata["run_fingerprint"] != changed_vocab.metadata["run_fingerprint"] + assert first.metadata["run_fingerprint"] != changed_config.metadata["run_fingerprint"] + + +def test_native_sequence_tokenizer_loader_context_is_bound_without_secret_values() -> None: + model = SyntheticEmbeddingModel() + model.__dict__["_fastplms_tokenizer_kwargs"] = { + "tokenizer_source": "Synthyra/Profluent-E1-150M", + "revision": "tokenizer-revision-a", + "cache_dir": "/immutable/cache", + "local_files_only": True, + "token": "do-not-persist-this-token", + } + first = embed_dataset(model, ["ACD"]) + model.__dict__["_fastplms_tokenizer_kwargs"]["revision"] = "tokenizer-revision-b" + changed = embed_dataset(model, ["ACD"]) + + assert first.metadata["tokenizer"] == { + "mode": "native-sequence", + "source": "Synthyra/Profluent-E1-150M", + "revision": "tokenizer-revision-a", + "cache_dir": "/immutable/cache", + "local_files_only": True, + "token_policy": "provided", + } + assert "do-not-persist-this-token" not in json.dumps(first.metadata) + assert first.metadata["run_fingerprint"] != changed.metadata["run_fingerprint"] + + +def test_local_artifact_identity_fills_embedding_provenance() -> None: + model = SyntheticEmbeddingModel() + model.config._name_or_path = "dist/hub/ESM2-8M" + model.config._commit_hash = None + model.config.fastplms_model_id = "esm2_8m" + model.config.fastplms_checkpoint_repo_id = "Synthyra/ESM2-8M" + model.config.fastplms_checkpoint_revision = "d" * 40 + model.config.fastplms_checkpoint_hash = "e" * 64 + + result = embed_dataset(model, ["ACD"], pooling="mean") + + assert result.metadata["model_id"] == "esm2_8m" + assert result.metadata["model_revision"] == "d" * 40 + assert result.metadata["checkpoint_repo_id"] == "Synthyra/ESM2-8M" + assert result.metadata["checkpoint_revision"] == "d" * 40 + assert result.metadata["checkpoint_hash"] == "e" * 64 + + model.config.fastplms_checkpoint_hash = "f" * 64 + changed = embed_dataset(model, ["ACD"], pooling="mean") + assert changed.metadata["run_fingerprint"] != result.metadata["run_fingerprint"] + + +def test_e1_native_path_removes_all_boundary_tokens() -> None: + result = embed_dataset(SyntheticE1Model(), ["AC"], pooling="mean") + assert torch.equal(result[0].load_tensor(), torch.tensor([5.5])) + + +def test_generic_embedding_uses_a_model_sequence_tokenizer_adapter() -> None: + class Tokenizer: + all_special_ids = (0, 2) + name_or_path = "synthetic/multimodal" + vocab_size = 8 + + def __call__(self, *args, **kwargs): + raise AssertionError("The generic tokenizer path must not run.") + + class AdaptedModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace(model_type="synthetic-multimodal") + self.tokenizer = Tokenizer() + self.sequences: list[str] | None = None + + def _tokenize_sequence_batch(self, sequences, *, tokenizer, **kwargs): + self.sequences = list(sequences) + assert tokenizer is self.tokenizer + assert kwargs == {"return_tensors": "pt", "padding": True, "truncation": True} + return { + "input_ids": torch.tensor([[0, 4, 5, 2]]), + "attention_mask": torch.ones(1, 4, dtype=torch.long), + } + + def _embed(self, input_ids, attention_mask, **kwargs): + del attention_mask, kwargs + return input_ids.float().unsqueeze(-1) + + model = AdaptedModel() + result = embed_dataset(model, ["AC"], full_embeddings=True) + + assert model.sequences == ["AC"] + assert torch.equal(result[0].load_tensor(), torch.tensor([[4.0], [5.0]])) + + +def test_max_length_counts_biological_residues_not_special_tokens() -> None: + class Tokenizer: + all_special_ids = (0, 2) + name_or_path = "synthetic/residue-limit" + vocab_size = 8 + + def num_special_tokens_to_add(self, *, pair: bool) -> int: + assert pair is False + return 2 + + def __call__(self, sequences, **kwargs): + assert sequences == ["ACD"] + assert kwargs["max_length"] == 5 + return { + "input_ids": torch.tensor([[0, 3, 4, 5, 2]]), + "attention_mask": torch.ones(1, 5, dtype=torch.long), + } + + class Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace(model_type="synthetic") + + def _embed(self, input_ids, attention_mask, **kwargs): + del attention_mask, kwargs + return input_ids.float().unsqueeze(-1) + + result = embed_dataset( + Model(), + ["ACDE"], + tokenizer=Tokenizer(), + max_length=3, + full_embeddings=True, + ) + + assert result[0].load_tensor().shape == (3, 1) + + +@pytest.mark.parametrize("input_kind", ("list", "mapping", "generator", "fasta")) +def test_truncate_false_rejects_overlength_inputs_before_custom_adapter_inference( + tmp_path: Path, + input_kind: str, +) -> None: + model = SyntheticEmbeddingModel() + inference_called = False + + def fail_if_called(sequences: list[str]) -> EmbeddingBatch: + del sequences + nonlocal inference_called + inference_called = True + raise AssertionError("over-length input reached inference") + + model._embedding_batch = fail_if_called # type: ignore[method-assign] + records = [ + EmbeddingInput("short", "AC"), + EmbeddingInput("too-long", "ACDE"), + ] + if input_kind == "list": + inputs = records + elif input_kind == "mapping": + inputs = {record.id: record.sequence for record in records} + elif input_kind == "generator": + inputs = (record for record in records) + else: + inputs = tmp_path / "proteins.fasta" + inputs.write_text(">short\nAC\n>too-long\nACDE\n", encoding="utf-8") + + with pytest.raises(ValueError) as error: + embed_dataset( + model, + inputs, + max_length=3, + truncate=False, + ) + + message = str(error.value) + assert "position 1" in message + assert "id 'too-long'" in message + assert "4 biological residues" in message + assert "max_length=3" in message + assert inference_called is False + + +def test_truncate_false_rejects_overlength_inputs_before_raw_adapter_inference() -> None: + model = SyntheticE1Model() + + def fail_if_called(*args, **kwargs): + del args, kwargs + raise AssertionError("over-length input reached raw adapter inference") + + model._embed = fail_if_called # type: ignore[method-assign] + with pytest.raises(ValueError, match=r"position 0.*id '0'.*max_length=3"): + embed_dataset( + model, + ["ACDE"], + max_length=3, + truncate=False, + ) + + +def test_all_poolers_and_output_slices() -> None: + names = ("mean", "max", "norm", "median", "std", "var", "cls", "parti") + result = embed_dataset(SyntheticEmbeddingModel(), ["ACD", "GG"], pooling=names) + assert tuple(result[0].load_tensor().shape) == (16,) + assert torch.isfinite(result[0].load_tensor()).all() + assert result.metadata["pool_slices"]["mean"] == (0, 2) + assert result.metadata["pool_slices"]["parti"] == (14, 16) + + +def test_poolers_ignore_nonfinite_excluded_positions_and_reject_nonfinite_output() -> None: + X = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [torch.nan, torch.inf]]]) + M = torch.tensor([[True, True, False]]) + + pooled = Pooler(("mean", "norm", "std", "var"))(X, M) + + assert torch.isfinite(pooled).all() + torch.testing.assert_close( + pooled, + torch.tensor( + [ + [ + 2.0, + 3.0, + 10.0**0.5, + 20.0**0.5, + 1.0, + 1.0, + 1.0, + 1.0, + ] + ] + ), + ) + with pytest.raises(ValueError, match="produced non-finite output"): + Pooler("mean")(torch.tensor([[[torch.nan], [1.0]]]), torch.tensor([[True, False]])) + + +def test_parti_requires_eager_attention() -> None: + with pytest.raises(ValueError, match="requires attn_implementation='eager'"): + embed_dataset(SyntheticEmbeddingModel("sdpa"), ["ACD"], pooling="parti") + + +def test_parti_rejects_overlength_input_before_model_inference() -> None: + class Tokenizer: + all_special_ids = (0, 2) + vocab_size = 3 + name_or_path = "synthetic/tokenizer" + + def __call__(self, sequences, **kwargs): + del kwargs + assert sequences == ["A" * 2_049] + input_ids = torch.tensor([[0, *([1] * 2_049), 2]]) + return { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + } + + class Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + model_type="synthetic", + _attn_implementation="eager", + ) + + def _embed(self, *args, **kwargs): + raise AssertionError("parti length validation must run before inference") + + with pytest.raises(ValueError, match="at most 2,048 biological residues"): + embed_dataset( + Model(), + ["A" * 2_049], + pooling="parti", + tokenizer=Tokenizer(), + ) + + +def test_torch_pagerank_handles_dangling_rows() -> None: + A = torch.tensor([[0.0, 1.0], [0.0, 0.0]]) + w = pagerank_weights(A) + assert torch.isclose(w.sum(), torch.tensor(1.0)) + assert bool((w > 0).all()) + + +def test_fasta_preserves_headers_order_and_duplicates(tmp_path: Path) -> None: + path = tmp_path / "proteins.fasta" + path.write_text(">a description\nACD\n>a\nGG\n", encoding="utf-8") + records = parse_fasta(path) + assert records == [EmbeddingInput("a", "ACD"), EmbeddingInput("a", "GG")] + + +def test_fasta_parser_streams_without_path_read_text( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + path = tmp_path / "stream.fasta" + path.write_text(">p1\nAC\nD\n>p2\nGG\n", encoding="utf-8") + + def reject_read_text(*args, **kwargs): + del args, kwargs + raise AssertionError("FASTA parsing must stream from the file handle") + + monkeypatch.setattr(Path, "read_text", reject_read_text) + assert [(record.id, record.sequence) for record in iter_fasta(path)] == [ + ("p1", "ACD"), + ("p2", "GG"), + ] + + +@pytest.mark.parametrize("source_kind", ("generator", "fasta")) +def test_large_streaming_inputs_use_bounded_disk_windows( + monkeypatch, + tmp_path, + source_kind: str, +) -> None: + import fastplms.embeddings.runner as runner + + count = 2_050 + iterations = 0 + expected_input_digest = hashlib.sha256() + source_liveness = {"live": 0, "peak": 0} + + class TrackedText: + def __init__(self, value: str) -> None: + self.value = value + source_liveness["live"] += 1 + source_liveness["peak"] = max(source_liveness["peak"], source_liveness["live"]) + if source_liveness["live"] > 32: + raise AssertionError( + "The source generator was fully materialized before disk spooling." + ) + + def __str__(self) -> str: + return self.value + + def __del__(self) -> None: + source_liveness["live"] -= 1 + + def update_expected_fingerprint(input_id: str, sequence: str) -> None: + for value in (input_id, sequence): + encoded = value.encode("utf-8") + expected_input_digest.update(len(encoded).to_bytes(8, "big")) + expected_input_digest.update(encoded) + + class SinglePassInputs: + def __iter__(self): + nonlocal iterations + iterations += 1 + if iterations > 1: + raise AssertionError("The source iterable was consumed more than once.") + for position in range(count): + input_id = f"protein-{position}" + sequence = "A" * (position % 11 + 1) + update_expected_fingerprint(input_id, sequence) + yield (TrackedText(input_id), TrackedText(sequence)) + + if source_kind == "generator": + inputs = SinglePassInputs() + else: + fasta = tmp_path / "large.fasta" + with fasta.open("w", encoding="utf-8") as handle: + for position in range(count): + input_id = f"protein-{position}" + sequence = "A" * (position % 11 + 1) + update_expected_fingerprint(input_id, sequence) + handle.write(f">{input_id}\n{sequence}\n") + inputs = fasta + + observed_slice_widths: list[int] = [] + original_getitem = runner._InputSpool.__getitem__ + + def tracking_getitem(self, index): + if isinstance(index, slice): + start, stop, step = index.indices(len(self)) + if step == 1: + observed_slice_widths.append(stop - start) + return original_getitem(self, index) + + def reject_full_spool_iteration(self): + del self + raise AssertionError("The disk spool must be consumed through bounded slices.") + + monkeypatch.setattr(runner._InputSpool, "__getitem__", tracking_getitem) + monkeypatch.setattr(runner._InputSpool, "__iter__", reject_full_spool_iteration) + output = tmp_path / f"{source_kind}.sqlite" + result = embed_dataset( + SyntheticEmbeddingModel(), + inputs, + batch_size=64, + batch_window_size=127, + output=output, + format="sqlite", + ) + + assert len(result) == count + for position, record in enumerate(result): + assert record.id == f"protein-{position}" + assert record.sequence == "A" * (position % 11 + 1) + assert observed_slice_widths + assert max(observed_slice_widths) <= 127 + assert result.metadata["batching"]["input_storage"] == "disk-spool" + expected_input_digest.update(count.to_bytes(8, "big")) + assert result.metadata["input_fingerprint"] == expected_input_digest.hexdigest() + reopened = load_sqlite_result(output) + assert reopened.metadata["input_fingerprint"] == result.metadata["input_fingerprint"] + assert reopened.metadata["run_fingerprint"] == result.metadata["run_fingerprint"] + if source_kind == "generator": + assert iterations == 1 + assert source_liveness["peak"] <= 32 + assert source_liveness["live"] == 0 + + +def test_sqlite_round_trip_is_lazy_and_bf16_lossless(tmp_path: Path) -> None: + source = embed_dataset(SyntheticEmbeddingModel(), ["ACD", "GG"]) + source = EmbeddingResult( + [ + type(record)(record.id, record.sequence, record.load_tensor().to(torch.bfloat16)) + for record in source + ], + source.metadata, + ) + path = tmp_path / "embeddings.sqlite" + saved = save_sqlite_result(source, path) + loaded = load_sqlite_result(path) + assert isinstance(saved[0].tensor, LazyTensorReference) + assert torch.equal(loaded[0].load_tensor(), source[0].load_tensor()) + assert loaded[0].load_tensor().dtype == torch.bfloat16 + + +def test_sqlite_tensor_corruption_is_detected_on_materialization(tmp_path: Path) -> None: + path = tmp_path / "corrupt.sqlite" + source = EmbeddingResult( + [EmbeddingRecord("protein", "AC", torch.tensor([1.0, 2.0]))], + {"run_fingerprint": "corrupt-sqlite", "complete": True}, + ) + save_sqlite_result(source, path) + with sqlite3.connect(path) as connection: + connection.execute( + "UPDATE tensors SET data = ?", + (sqlite3.Binary(torch.tensor([3.0, 4.0]).numpy().tobytes()),), + ) + connection.commit() + + with pytest.raises(ValueError, match="failed SHA-256 verification"): + load_sqlite_result(path)[0].load_tensor() + + +def test_sqlite_loading_and_lazy_tensor_reads_use_read_only_connections( + monkeypatch, + tmp_path, +) -> None: + import fastplms.embeddings.storage as storage + + path = tmp_path / "readonly.sqlite" + save_sqlite_result(embed_dataset(SyntheticEmbeddingModel(), ["ACD"]), path) + original_connect = storage.sqlite3.connect + observed: list[tuple[object, dict[str, object]]] = [] + + def tracking_connect(database, *args, **kwargs): + observed.append((database, dict(kwargs))) + return original_connect(database, *args, **kwargs) + + monkeypatch.setattr(storage.sqlite3, "connect", tracking_connect) + loaded = load_sqlite_result(path) + loaded[0].load_tensor() + + # Loading metadata, resolving the lazy descriptor, and loading its tensor + # are independent read-only operations. + assert len(observed) == 3 + assert all("mode=ro" in str(database) for database, _ in observed) + assert all(kwargs.get("uri") is True for _, kwargs in observed) + + +def test_sqlite_filtered_retrieval_preserves_selector_order_and_duplicates( + tmp_path: Path, +) -> None: + path = tmp_path / "selection.sqlite" + result = EmbeddingResult( + [ + EmbeddingRecord("x", "AA", torch.tensor([0.0])), + EmbeddingRecord("y", "BB", torch.tensor([1.0])), + EmbeddingRecord("x", "AA", torch.tensor([2.0])), + ], + {"run_fingerprint": "selection-run", "complete": True}, + ) + save_sqlite_result(result, path) + + by_position = load_sqlite_result(path, positions=[2, 0, 2]) + assert [record.load_tensor().item() for record in by_position] == [2.0, 0.0, 2.0] + by_id = load_sqlite_result(path, record_ids=["x", "y", "x"]) + assert [record.load_tensor().item() for record in by_id] == [ + 0.0, + 2.0, + 1.0, + 0.0, + 2.0, + ] + by_sequence = load_sqlite_result(path, sequences=["BB", "AA"]) + assert [record.id for record in by_sequence] == ["y", "x", "x"] + + +def test_legacy_sqlite_converter_accepts_compact_blobs_without_pickle( + tmp_path: Path, +) -> None: + source = tmp_path / "legacy.sqlite" + output = tmp_path / "converted.sqlite" + tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) + shape = tuple(tensor.shape) + blob = ( + struct.pack( + f" None: + path = tmp_path / "stream.sqlite" + inputs = ["ACD", "GG", "M"] + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=2), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + ) + partial = load_sqlite_result(path) + assert len(partial) == 1 + assert partial.metadata["complete"] is False + + resumed = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + ) + assert [record.sequence for record in resumed] == inputs + assert resumed.metadata["complete"] is True + assert resumed.metadata["record_count"] == 3 + assert resumed.metadata["descriptor_index"] == "sqlite-records" + assert "outputs" not in resumed.metadata + + +def test_sqlite_successful_overwrite_becomes_default_and_retains_prior_run( + tmp_path, +) -> None: + path = tmp_path / "overwrite.sqlite" + model = SyntheticEmbeddingModel() + original = embed_dataset( + model, + ["AC"], + output=path, + format="sqlite", + ) + original_run_id = original.metadata["run_fingerprint"] + original_tensor = original[0].load_tensor().clone() + + replacement = embed_dataset( + model, + ["GG", "M"], + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + resume=False, + ) + + current = load_sqlite_result(path) + retained = load_sqlite_result(path, run_id=original_run_id) + assert current.metadata["run_fingerprint"] == replacement.metadata["run_fingerprint"] + assert current.metadata["run_fingerprint"] != original_run_id + assert [record.sequence for record in current] == ["GG", "M"] + assert current.metadata["complete"] is True + assert [record.sequence for record in retained] == ["AC"] + torch.testing.assert_close(retained[0].load_tensor(), original_tensor, rtol=0.0, atol=0.0) + # Readers opened before the replacement remain bound to their explicit run. + torch.testing.assert_close(original[0].load_tensor(), original_tensor, rtol=0.0, atol=0.0) + + +def test_interrupted_sqlite_overwrite_retains_prior_run_and_resumable_prefix( + tmp_path, +) -> None: + path = tmp_path / "interrupted-overwrite.sqlite" + original = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + ["AC"], + output=path, + format="sqlite", + ) + original_run_id = original.metadata["run_fingerprint"] + original_tensor = original[0].load_tensor().clone() + replacement_inputs = ["GG", "M"] + + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=2), + replacement_inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + resume=False, + ) + + partial = load_sqlite_result(path) + retained = load_sqlite_result(path, run_id=original_run_id) + assert partial.metadata["run_fingerprint"] != original_run_id + assert partial.metadata["complete"] is False + assert [record.sequence for record in partial] == replacement_inputs[:1] + assert [record.sequence for record in retained] == ["AC"] + torch.testing.assert_close(retained[0].load_tensor(), original_tensor, rtol=0.0, atol=0.0) + + resumed = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + replacement_inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + ) + assert resumed.metadata["run_fingerprint"] == partial.metadata["run_fingerprint"] + assert resumed.metadata["complete"] is True + assert [record.sequence for record in resumed] == replacement_inputs + assert [record.sequence for record in load_sqlite_result(path, run_id=original_run_id)] == [ + "AC" + ] + + +def test_sqlite_first_batch_publication_is_atomic_and_hidden_run_resumes( + tmp_path: Path, +) -> None: + path = tmp_path / "first-batch-atomic.sqlite" + original = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + ["AC"], + output=path, + format="sqlite", + ) + original_run_id = original.metadata["run_fingerprint"] + replacement_inputs = ["GG", "M"] + + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=1), + replacement_inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + resume=False, + ) + + current = load_sqlite_result(path) + assert current.metadata["run_fingerprint"] == original_run_id + assert [record.sequence for record in current] == ["AC"] + + resumed = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + replacement_inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + ) + assert resumed.metadata["run_fingerprint"] != original_run_id + assert resumed.metadata["complete"] is True + assert [record.sequence for record in resumed] == replacement_inputs + assert ( + load_sqlite_result(path).metadata["run_fingerprint"] == resumed.metadata["run_fingerprint"] + ) + + +def test_sqlite_same_run_replacement_is_deferred_until_first_batch_commit( + tmp_path: Path, +) -> None: + path = tmp_path / "same-run-atomic.sqlite" + inputs = ["AC", "GG"] + original = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + ) + original_run_id = original.metadata["run_fingerprint"] + original_tensors = [record.load_tensor().clone() for record in original] + + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=1), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + resume=False, + ) + + retained = load_sqlite_result(path) + assert retained.metadata["run_fingerprint"] == original_run_id + assert retained.metadata["complete"] is True + for observed, expected in zip(retained, original_tensors, strict=True): + torch.testing.assert_close(observed.load_tensor(), expected, rtol=0.0, atol=0.0) + + replacement = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="sqlite", + resume=False, + ) + assert replacement.metadata["run_fingerprint"] == original_run_id + assert replacement.metadata["complete"] is True + assert [record.sequence for record in replacement] == inputs + + +def test_sqlite_prepublication_schema_remains_readable_and_migrates( + tmp_path: Path, +) -> None: + path = tmp_path / "pre-publication-schema.sqlite" + original = embed_dataset( + SyntheticEmbeddingModel(), + ["AC"], + output=path, + format="sqlite", + ) + original_run_id = original.metadata["run_fingerprint"] + with sqlite3.connect(path) as connection: + connection.execute("DROP INDEX runs_published_order_idx") + connection.execute("ALTER TABLE runs DROP COLUMN published_order") + connection.commit() + + legacy_view = load_sqlite_result(path) + assert legacy_view.metadata["run_fingerprint"] == original_run_id + assert [record.sequence for record in legacy_view] == ["AC"] + + replacement = embed_dataset( + SyntheticEmbeddingModel(), + ["GG"], + output=path, + format="sqlite", + resume=False, + ) + with sqlite3.connect(path) as connection: + columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(runs)").fetchall()} + assert "published_order" in columns + assert ( + load_sqlite_result(path).metadata["run_fingerprint"] + == replacement.metadata["run_fingerprint"] + ) + assert [record.sequence for record in load_sqlite_result(path, run_id=original_run_id)] == [ + "AC" + ] + + +def test_safetensors_round_trip_is_lazy(tmp_path: Path) -> None: + source = embed_dataset(SyntheticEmbeddingModel(), ["ACD", "GG"]) + output = tmp_path / "safe" + with pytest.raises(ValueError, match="cannot fit"): + save_safetensors_result(source, output, shard_size=4) + saved = save_safetensors_result(source, output, shard_size=8) + loaded = load_safetensors_result(output) + assert isinstance(saved[0].tensor, LazyTensorReference) + assert len(list(output.glob("*.safetensors"))) >= 2 + index_path = output / "index.json" + run_manifest = json.loads((output / "run.json").read_text(encoding="utf-8")) + generation_path = output / run_manifest["index"]["file"] + assert run_manifest["index"] == { + "file": generation_path.name, + "sha256": hashlib.sha256(generation_path.read_bytes()).hexdigest(), + } + assert run_manifest["version"] == 2 + assert run_manifest["record_count"] == len(source) + pointer = json.loads(index_path.read_text(encoding="utf-8")) + assert pointer["index"] == run_manifest["index"] + generation = json.loads(generation_path.read_text(encoding="utf-8")) + assert "records" not in generation + assert generation["record_count"] == len(source) + assert generation["metadata"]["descriptor_index"] == "safetensors-generation-index" + assert "outputs" not in generation["metadata"] + descriptors = [] + for shard in generation["descriptor_shards"]: + descriptors.extend( + json.loads(line) + for line in (output / shard["file"]).read_text(encoding="utf-8").splitlines() + if line + ) + assert len(descriptors) == len(source) + assert len({item["tensor"]["key"] for item in descriptors}) == len(source) + assert torch.equal(loaded[1].load_tensor(), source[1].load_tensor()) + + +def test_safetensors_descriptor_shards_have_a_bounded_record_count( + tmp_path: Path, +) -> None: + output = tmp_path / "bounded-descriptors" + source = EmbeddingResult( + [ + EmbeddingRecord(str(position), "A", torch.tensor([float(position)])) + for position in range(1_025) + ], + {"run_fingerprint": "bounded-descriptors", "complete": True}, + ) + + save_safetensors_result(source, output, shard_size=1024**2) + run = json.loads((output / "run.json").read_text(encoding="utf-8")) + generation = json.loads((output / run["index"]["file"]).read_text(encoding="utf-8")) + + assert [item["count"] for item in generation["descriptor_shards"]] == [1_024, 1] + assert sum(item["count"] for item in generation["descriptor_shards"]) == len(source) + assert "records" not in generation + assert "outputs" not in generation["metadata"] + + +def test_safetensors_tensor_corruption_is_detected_on_materialization( + tmp_path: Path, +) -> None: + from safetensors.torch import save_file + + output = tmp_path / "corrupt-safe" + source = EmbeddingResult( + [EmbeddingRecord("protein", "AC", torch.tensor([1.0, 2.0]))], + {"run_fingerprint": "corrupt-safe", "complete": True}, + ) + save_safetensors_result(source, output) + run = json.loads((output / "run.json").read_text(encoding="utf-8")) + generation = json.loads((output / run["index"]["file"]).read_text(encoding="utf-8")) + descriptor_path = output / generation["descriptor_shards"][0]["file"] + descriptor = json.loads(descriptor_path.read_text(encoding="utf-8").splitlines()[0]) + tensor_path = output / descriptor["tensor"]["file"] + save_file({descriptor["tensor"]["key"]: torch.tensor([3.0, 4.0])}, tensor_path) + + with pytest.raises(ValueError, match="failed SHA-256 verification"): + load_safetensors_result(output)[0].load_tensor() + + +def test_safetensors_streaming_resumes_an_ordered_prefix(tmp_path: Path) -> None: + path = tmp_path / "stream-safe" + inputs = ["ACD", "GG", "M"] + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=3), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="safetensors", + shard_size=8, + ) + partial = load_safetensors_result(path) + assert len(partial) == 1 + assert partial.metadata["complete"] is False + + resumed = embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=None), + inputs, + batch_size=1, + batch_window_size=1, + output=path, + format="safetensors", + shard_size=8, + ) + assert [record.sequence for record in resumed] == inputs + assert resumed.metadata["complete"] is True + assert resumed.metadata["record_count"] == 3 + assert resumed.metadata["descriptor_index"] == "safetensors-generation-index" + assert "outputs" not in resumed.metadata + + +@pytest.mark.parametrize( + ("format", "expected_granularity"), + (("sqlite", "batch-window"), ("safetensors", "shard-flush")), +) +def test_persistent_resume_metadata_records_true_commit_granularity( + tmp_path: Path, + format: str, + expected_granularity: str, +) -> None: + output = tmp_path / ("embeddings.sqlite" if format == "sqlite" else "embeddings") + + result = embed_dataset( + SyntheticEmbeddingModel(), + ["AC", "GG"], + batch_size=1, + batch_window_size=1, + output=output, + format=format, + ) + + assert result.metadata["batching"]["resume_commit_granularity"] == (expected_granularity) + + +def test_safetensors_streaming_packs_batches_into_shards(tmp_path: Path) -> None: + output = tmp_path / "packed" + embed_dataset( + SyntheticEmbeddingModel(), + ["AC", "GG", "MM"], + batch_size=1, + output=output, + format="safetensors", + shard_size=24, + ) + + assert len(list(output.glob("*.safetensors"))) == 1 + + +def test_safetensors_manifest_rejects_shard_path_traversal(tmp_path: Path) -> None: + output = tmp_path / "safe" + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("protein", "AC", torch.tensor([1.0, 2.0]))], + {"complete": True}, + ), + output, + ) + run_path = output / "run.json" + run = json.loads(run_path.read_text(encoding="utf-8")) + generation_path = output / run["index"]["file"] + generation = json.loads(generation_path.read_text(encoding="utf-8")) + descriptor_reference = generation["descriptor_shards"][0] + descriptor_path = output / descriptor_reference["file"] + descriptors = [ + json.loads(line) + for line in descriptor_path.read_text(encoding="utf-8").splitlines() + if line + ] + descriptors[0]["tensor"]["file"] = "../outside.safetensors" + descriptor_bytes = b"".join( + json.dumps(item, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n" + for item in descriptors + ) + descriptor_path.write_bytes(descriptor_bytes) + descriptor_reference["sha256"] = hashlib.sha256(descriptor_bytes).hexdigest() + generation_bytes = (json.dumps(generation, indent=2, sort_keys=True) + "\n").encode() + generation_path.write_bytes(generation_bytes) + run["index"]["sha256"] = hashlib.sha256(generation_bytes).hexdigest() + run_path.write_text(json.dumps(run, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + with pytest.raises(ValueError, match="outside its output directory"): + load_safetensors_result(output)[0] + + +def test_pooler_rejects_duplicate_operations() -> None: + with pytest.raises(ValueError, match="Duplicate pooling operations"): + Pooler(("mean", "mean")) + + +def test_failed_safetensors_overwrite_preserves_previous_valid_generation( + tmp_path: Path, +) -> None: + output = tmp_path / "safe" + original = EmbeddingResult( + [EmbeddingRecord("old", "AC", torch.tensor([1.0, 2.0]))], + {"complete": True}, + ) + save_safetensors_result(original, output, shard_size=8) + replacement = EmbeddingResult( + [ + EmbeddingRecord("new-1", "GG", torch.tensor([3.0, 4.0])), + EmbeddingRecord("new-2", "MM", torch.tensor([5.0, 6.0])), + EmbeddingRecord("too-large", "M", torch.arange(3, dtype=torch.float32)), + ], + {"complete": True}, + ) + + original_manifest = (output / "run.json").read_bytes() + with pytest.raises(ValueError, match="cannot fit"): + save_safetensors_result(replacement, output, shard_size=8) + + assert (output / "run.json").read_bytes() == original_manifest + loaded = load_safetensors_result(output) + assert [(record.id, record.sequence) for record in loaded] == [("old", "AC")] + assert torch.equal(loaded[0].load_tensor(), torch.tensor([1.0, 2.0])) + + +def test_open_safetensors_reader_survives_successful_overwrite(tmp_path: Path) -> None: + output = tmp_path / "retained" + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("old", "AC", torch.tensor([1.0, 2.0]))], + {"complete": True}, + ), + output, + ) + old_reader = load_safetensors_result(output) + old_manifest = json.loads((output / "run.json").read_text(encoding="utf-8")) + old_index_path = output / old_manifest["index"]["file"] + old_index = json.loads(old_index_path.read_text(encoding="utf-8")) + old_descriptor_path = output / old_index["descriptor_shards"][0]["file"] + old_tensor_path = output / old_index["descriptor_shards"][0]["tensor_file"] + + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("middle", "GG", torch.tensor([3.0, 4.0]))], + {"complete": True}, + ), + output, + ) + # A third writer starts while the first generation is already + # non-authoritative. Writer initialization must not treat it as an orphan. + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("new", "MM", torch.tensor([5.0, 6.0]))], + {"complete": True}, + ), + output, + ) + current = load_safetensors_result(output) + + assert [(record.id, record.sequence) for record in current] == [("new", "MM")] + assert torch.equal(current[0].load_tensor(), torch.tensor([5.0, 6.0])) + assert old_index_path.is_file() + assert old_descriptor_path.is_file() + assert old_tensor_path.is_file() + assert [(record.id, record.sequence) for record in old_reader] == [("old", "AC")] + assert torch.equal(old_reader[0].load_tensor(), torch.tensor([1.0, 2.0])) + + +def test_safetensors_generation_gc_is_dry_run_and_explicitly_exclusive( + tmp_path: Path, +) -> None: + output = tmp_path / "retained" + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("old", "AC", torch.tensor([1.0, 2.0]))], + {"complete": True}, + ), + output, + ) + old_manifest = json.loads((output / "run.json").read_text(encoding="utf-8")) + old_index_path = output / old_manifest["index"]["file"] + old_index = json.loads(old_index_path.read_text(encoding="utf-8")) + old_generation_paths = { + old_index_path, + output / old_index["descriptor_shards"][0]["file"], + output / old_index["descriptor_shards"][0]["tensor_file"], + } + save_safetensors_result( + EmbeddingResult( + [EmbeddingRecord("new", "GG", torch.tensor([3.0, 4.0]))], + {"complete": True}, + ), + output, + ) + + preview = set(garbage_collect_safetensors_generations(output)) + assert old_generation_paths.issubset(preview) + assert all(path.is_file() for path in old_generation_paths) + with pytest.raises(ValueError, match="confirm_no_active_readers_or_writers"): + garbage_collect_safetensors_generations(output, dry_run=False) + assert all(path.is_file() for path in old_generation_paths) + + deleted = set( + garbage_collect_safetensors_generations( + output, + dry_run=False, + confirm_no_active_readers_or_writers=True, + ) + ) + assert deleted == preview + assert not any(path.exists() for path in old_generation_paths) + current = load_safetensors_result(output) + assert [(record.id, record.sequence) for record in current] == [("new", "GG")] + assert torch.equal(current[0].load_tensor(), torch.tensor([3.0, 4.0])) + + +def test_interrupted_embedding_overwrite_preserves_previous_generation( + tmp_path: Path, +) -> None: + output = tmp_path / "safe" + original = embed_dataset( + SyntheticEmbeddingModel(), + ["AC"], + output=output, + format="safetensors", + ) + + with pytest.raises(RuntimeError, match="simulated interruption"): + embed_dataset( + InterruptibleEmbeddingModel(fail_on_call=2), + ["GG", "M"], + batch_size=1, + output=output, + format="safetensors", + resume=False, + ) + + loaded = load_safetensors_result(output) + assert [(record.id, record.sequence) for record in loaded] == [ + (record.id, record.sequence) for record in original + ] + assert torch.equal(loaded[0].load_tensor(), original[0].load_tensor()) + + +def test_interrupted_metadata_publish_recovers_last_committed_generation( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + output = tmp_path / "safe" + original = EmbeddingResult( + [EmbeddingRecord("old", "AC", torch.tensor([1.0, 2.0]))], + {"complete": True}, + ) + save_safetensors_result(original, output, shard_size=8) + + writer = SafetensorsStreamWriter( + output, + {"complete": False}, + shard_size=8, + publish_initial=False, + publish_incremental=False, + ) + writer.append( + [EmbeddingRecord("new", "GG", torch.tensor([3.0, 4.0]))], + publish=False, + ) + run_manifest_path = output / "run.json" + original_replace = Path.replace + + def interrupt_manifest_replace(path: Path, target: Path) -> Path: + if Path(target) == run_manifest_path: + raise OSError("simulated metadata publication interruption") + return original_replace(path, target) + + monkeypatch.setattr(Path, "replace", interrupt_manifest_replace) + with pytest.raises(OSError, match="metadata publication interruption"): + writer.publish(complete=True) + + loaded = load_safetensors_result(output) + assert [(record.id, record.sequence) for record in loaded] == [("old", "AC")] + assert torch.equal(loaded[0].load_tensor(), torch.tensor([1.0, 2.0])) + + +def test_named_safetensors_outputs_do_not_share_shards(tmp_path: Path) -> None: + first_source = embed_dataset(SyntheticEmbeddingModel(), ["ACD"]) + second_source = embed_dataset(SyntheticEmbeddingModel(), ["GGG"]) + first_path = tmp_path / "first.safetensors" + second_path = tmp_path / "second.safetensors" + + save_safetensors_result(first_source, first_path) + save_safetensors_result(second_source, second_path) + + first_loaded = load_safetensors_result(first_path) + second_loaded = load_safetensors_result(second_path) + assert torch.equal(first_loaded[0].load_tensor(), first_source[0].load_tensor()) + assert torch.equal(second_loaded[0].load_tensor(), second_source[0].load_tensor()) + assert list(tmp_path.glob("first-embeddings-*.safetensors")) + assert list(tmp_path.glob("second-embeddings-*.safetensors")) + + +def test_safetensors_manifest_snapshot_recovers_mismatched_standalone_index( + tmp_path: Path, +) -> None: + output = tmp_path / "safe" + source = embed_dataset(SyntheticEmbeddingModel(), ["ACD"]) + save_safetensors_result(source, output) + index_path = output / "index.json" + index_path.write_text(index_path.read_text(encoding="utf-8") + "\n", encoding="utf-8") + + recovered = load_safetensors_result(output) + assert len(recovered) == 1 + assert recovered[0].sequence == "ACD" + + run_path = output / "run.json" + manifest = json.loads(run_path.read_text(encoding="utf-8")) + generation_path = output / manifest["index"]["file"] + generation_path.write_text( + generation_path.read_text(encoding="utf-8") + "\n", + encoding="utf-8", + ) + with pytest.raises(ValueError, match="does not match its index"): + load_safetensors_result(output) + + +def test_resume_recovers_from_authoritative_manifest_when_index_is_missing( + tmp_path: Path, +) -> None: + output = tmp_path / "safe" + model = SyntheticEmbeddingModel() + original = embed_dataset(model, ["ACD"], output=output) + (output / "index.json").unlink() + + recovered = load_safetensors_result(output) + resumed = embed_dataset(model, ["ACD"], output=output) + + assert not (output / "index.json").exists() + assert torch.equal(recovered[0].load_tensor(), original[0].load_tensor()) + assert resumed.metadata["run_fingerprint"] == original.metadata["run_fingerprint"] + + +def test_resume_requires_matching_fingerprint(tmp_path: Path) -> None: + output = tmp_path / "resume" + first = embed_dataset(SyntheticEmbeddingModel(), ["ACD"], output=output) + resumed = embed_dataset(SyntheticEmbeddingModel(), ["ACD"], output=output) + assert isinstance(first[0].tensor, LazyTensorReference) + assert resumed.metadata["run_fingerprint"] == first.metadata["run_fingerprint"] + with pytest.raises(ValueError, match="different run fingerprint"): + embed_dataset(SyntheticEmbeddingModel(), ["GG"], output=output) + + +def test_resume_rejects_legacy_fingerprint_schema(tmp_path: Path) -> None: + output = tmp_path / "legacy-schema.sqlite" + embed_dataset( + SyntheticEmbeddingModel(), + ["ACD"], + output=output, + format="sqlite", + ) + with sqlite3.connect(output) as connection: + run_id, metadata_json = connection.execute( + "SELECT run_id, metadata_json FROM runs" + ).fetchone() + metadata = json.loads(metadata_json) + metadata["fingerprint_schema_version"] = 1 + connection.execute( + "UPDATE runs SET metadata_json = ? WHERE run_id = ?", + (json.dumps(metadata, sort_keys=True), run_id), + ) + + with pytest.raises(ValueError, match="incompatible run fingerprint schema"): + embed_dataset( + SyntheticEmbeddingModel(), + ["ACD"], + output=output, + format="sqlite", + ) + + +def test_legacy_pth_import_requires_explicit_unsafe_opt_in(tmp_path: Path) -> None: + path = tmp_path / "legacy.pth" + torch.save({"ACD": torch.ones(3)}, path) + with pytest.raises(ValueError, match="allow_unsafe_pickle=True"): + load_legacy_pth(path) + loaded = load_legacy_pth(path, allow_unsafe_pickle=True) + assert torch.equal(loaded[0].load_tensor(), torch.ones(3)) + + +def test_new_api_never_writes_pth(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="not supported"): + embed_dataset( + SyntheticEmbeddingModel(), + ["ACD"], + output=tmp_path / "embeddings.pth", + format="pth", + ) + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("batch_size", True), + ("batch_size", 1.5), + ("max_length", True), + ("max_tokens_per_batch", True), + ("batch_window_size", True), + ("shard_size", True), + ("full_embeddings", 1), + ("resume", 1), + ("truncate", 1), + ("model_state_fingerprint", ""), + ], +) +def test_strict_embedding_controls_fail_before_consuming_inputs(name, value) -> None: + consumed = False + + def source(): + nonlocal consumed + consumed = True + yield "AC" + + with pytest.raises((TypeError, ValueError)): + embed_dataset(SyntheticEmbeddingModel(), source(), **{name: value}) + assert consumed is False + + +@pytest.mark.parametrize( + "decoder_input_ids", + [ + torch.tensor([1], dtype=torch.int64), + torch.empty((1, 0), dtype=torch.int64), + torch.tensor([[1]], dtype=torch.int16), + torch.tensor([[1.0]], dtype=torch.float32), + ], +) +def test_decoder_input_ids_require_nonempty_2d_int32_or_int64(decoder_input_ids) -> None: + with pytest.raises((TypeError, ValueError)): + embed_dataset( + SyntheticDecoderEmbeddingModel(), + ["AC"], + hidden_state_source="decoder", + decoder_input_ids=decoder_input_ids, + ) + + +@pytest.mark.parametrize( + "decoder_attention_mask", + [ + torch.tensor([[0, 2]], dtype=torch.int64), + torch.tensor([[0.0, float("nan")]]), + torch.tensor([[1]], dtype=torch.int64), + ], +) +def test_decoder_attention_masks_are_exact_finite_binary_shapes( + decoder_attention_mask, +) -> None: + with pytest.raises(ValueError): + embed_dataset( + SyntheticDecoderEmbeddingModel(), + ["AC"], + hidden_state_source="decoder", + decoder_input_ids=torch.tensor([[1, 2]], dtype=torch.int64), + decoder_attention_mask=decoder_attention_mask, + ) + + +class InvalidEmbeddingBatchModel(SyntheticEmbeddingModel): + def __init__(self, case: str) -> None: + super().__init__() + self.case = case + + def _embedding_batch(self, sequences: list[str]) -> EmbeddingBatch: + batch = super()._embedding_batch(sequences) + if self.case == "x_type": + return EmbeddingBatch(X="bad", residue_mask=batch.residue_mask) # type: ignore[arg-type] + if self.case == "mask_type": + return EmbeddingBatch(X=batch.X, residue_mask="bad") # type: ignore[arg-type] + if self.case == "x_integer": + return EmbeddingBatch(X=batch.X.to(torch.int64), residue_mask=batch.residue_mask) + if self.case == "mask_nonbinary": + return EmbeddingBatch(X=batch.X, residue_mask=batch.residue_mask.float() * 0.5) + if self.case == "wrong_batch": + return EmbeddingBatch(X=batch.X[:1], residue_mask=batch.residue_mask[:1]) + X = batch.X.clone() + X[0, 1, 0] = torch.inf + return EmbeddingBatch(X=X, residue_mask=batch.residue_mask) + + +@pytest.mark.parametrize( + "case", + ["x_type", "mask_type", "x_integer", "mask_nonbinary", "wrong_batch", "nonfinite"], +) +def test_embedding_batch_adapter_outputs_are_validated(case) -> None: + with pytest.raises((TypeError, ValueError)): + embed_dataset(InvalidEmbeddingBatchModel(case), ["AC", "GG"]) + + +@pytest.mark.parametrize( + "case", + ["loader_type", "dtype", "verify", "record_tensor"], +) +def test_embedding_value_types_fail_closed(case) -> None: + if case == "record_tensor": + with pytest.raises(TypeError): + EmbeddingRecord("id", "AC", object()) # type: ignore[arg-type] + return + reference = LazyTensorReference( + source="memory", + key="x", + dtype="float32", + shape=(1,), + sha256="0" * 64, + _loader=(lambda: object()) if case == "loader_type" else (lambda: torch.ones(1)), + ) + if case == "loader_type": + with pytest.raises(TypeError): + reference.load(verify=False) + elif case == "dtype": + wrong_dtype = LazyTensorReference( + source="memory", + key="x", + dtype="float64", + shape=(1,), + sha256="0" * 64, + _loader=lambda: torch.ones(1), + ) + with pytest.raises(ValueError, match="dtype"): + wrong_dtype.load(verify=False) + else: + with pytest.raises(TypeError, match="verify"): + reference.load(verify=1) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "kwargs", + [ + {"damping": float("nan")}, + {"tolerance": 0.0}, + {"max_iterations": True}, + ], +) +def test_pagerank_controls_require_finite_valid_values(kwargs) -> None: + with pytest.raises((TypeError, ValueError)): + pagerank_weights(torch.ones(2, 2), **kwargs) + + +def test_tensor_sha256_uses_bounded_chunks_and_preserves_legacy_digest(monkeypatch) -> None: + from fastplms.embeddings import storage + + X = torch.arange(12, dtype=torch.float32).reshape(3, 4).transpose(0, 1) + expected = hashlib.sha256() + expected.update(b"float32") + expected.update(json.dumps(tuple(X.shape)).encode()) + expected.update(X.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()) + + monkeypatch.setattr(storage, "_TENSOR_HASH_CHUNK_BYTES", 7) + monkeypatch.setattr( + storage, + "_tensor_bytes", + lambda _: (_ for _ in ()).throw(AssertionError("full byte copy used")), + ) + assert storage.tensor_sha256(X) == expected.hexdigest() + + +def test_sqlite_lazy_references_are_absolute_after_cwd_change(tmp_path: Path, monkeypatch) -> None: + output = tmp_path / "embeddings.sqlite" + saved = save_sqlite_result( + EmbeddingResult( + [EmbeddingRecord("id", "AC", torch.arange(4, dtype=torch.float32))], + {"run_fingerprint": "absolute-path"}, + ), + output, + ) + monkeypatch.chdir(tmp_path.parent) + assert Path(saved[0].tensor.source).is_absolute() + assert torch.equal(saved[0].load_tensor(), torch.arange(4, dtype=torch.float32)) diff --git a/tests/unit/test_esm3_rotary.py b/tests/unit/test_esm3_rotary.py new file mode 100644 index 0000000..c2f176e --- /dev/null +++ b/tests/unit/test_esm3_rotary.py @@ -0,0 +1,28 @@ +"""Contracts for ESM3 rotary-position buffers.""" + +from __future__ import annotations + +import torch + +from fastplms.models.esm3.modeling_esm3 import RotaryEmbedding + + +def test_esm3_rotary_rebuilds_frequency_after_meta_materialization() -> None: + """A meta-loaded non-persistent frequency buffer must not contain garbage.""" + + with torch.device("meta"): + rotary = RotaryEmbedding(dim=8) + rotary = rotary.to_empty(device="cpu") + + Q = torch.randn(2, 5, 3, 8) # (b=2, l=5, h=3, d_h=8) + K = torch.randn(2, 5, 3, 8) # (b=2, l=5, h=3, d_h=8) + Q_rotated, K_rotated = rotary(Q, K) # each: (b=2, l=5, h=3, d_h=8) + + expected_inv_freq = rotary._compute_inv_freq(torch.device("cpu")) # (d_h / 2=4,) + assert torch.equal(rotary.inv_freq, expected_inv_freq) + assert torch.isfinite(Q_rotated).all() + assert torch.isfinite(K_rotated).all() + assert rotary._cos_cached is not None + assert rotary._sin_cached is not None + assert rotary._cos_cached.abs().max() <= 1 + assert rotary._sin_cached.abs().max() <= 1 diff --git a/tests/unit/test_esmc_diagnostics.py b/tests/unit/test_esmc_diagnostics.py new file mode 100644 index 0000000..530b8f2 --- /dev/null +++ b/tests/unit/test_esmc_diagnostics.py @@ -0,0 +1,690 @@ +"""Strict schema-v3 ESMC evidence with warning-only backend bands.""" + +from __future__ import annotations + +import copy +import json +import pytest +import torch +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace + +from fastplms.registry import get_model_registry +from tests.parity import test_native_results as diagnostics +from tests.parity.support.esmc_calibration import ( + ESMC_CALIBRATION_SEED, + esmc_calibration_batches, + validate_esmc_calibration_batch, +) +from tests.unit.test_biohub_reference_lock import _reference_environment_payload +from tools.remote.prepare_references import _esmc_calibration_batches + + +SPEC = get_model_registry()["esmc_small"] +ADVERTISED_BF16_BACKENDS = tuple( + backend + for backend in SPEC.family.attention + if "bfloat16" + in get_model_registry().supported_attention_dtypes( + SPEC.family.id, + backend, + ) +) +PANEL_KINDS = ("generated_kernel_boundary", "real_biological_holdout") +SOURCE_TREE_SHA256 = "1" * 64 +RUNTIME_BUNDLE_SHA256 = "2" * 64 +REFERENCE_SOURCES: dict[str, dict[str, object]] = { + "biohub-esm": { + "attestation_sha256": "a" * 64, + "file_count": 412, + "import_file": "esm/__init__.py", + "import_name": "esm", + "import_root": "esm", + "package_version": "3.3.0", + "schema_version": 1, + "source_revision": diagnostics.BIOHUB_ESM_REVISION, + "tree_sha256": diagnostics.BIOHUB_ESM_TREE_SHA256, + }, + "biohub-transformers": { + "attestation_sha256": "b" * 64, + "file_count": 5218, + "import_file": "src/transformers/__init__.py", + "import_name": "transformers", + "import_root": "src/transformers", + "package_version": "4.57.6", + "schema_version": 1, + "source_revision": diagnostics.BIOHUB_TRANSFORMERS_REVISION, + "tree_sha256": diagnostics.BIOHUB_TRANSFORMERS_TREE_SHA256, + }, +} + + +def _output(hidden: torch.Tensor) -> SimpleNamespace: + signal = hidden[..., :1] * 8 + logits = torch.cat((signal, -signal), dim=-1) + return SimpleNamespace( + hidden_states=(hidden * 0.5, hidden), + last_hidden_state=hidden, + logits=logits, + ) + + +def _batch(kind: str) -> dict[str, object]: + return copy.deepcopy( + next(batch for batch in esmc_calibration_batches() if batch["kind"] == kind) + ) + + +def _panel_tensors( + kind: str, + *, + candidate_scale: float, +) -> tuple[dict[str, object], torch.Tensor, torch.Tensor, torch.Tensor]: + batch = _batch(kind) + cases = batch["cases"] + assert isinstance(cases, list) + lengths = torch.tensor([int(case["sequence_length"]) for case in cases]) # (b,) + maximum = int(lengths.max().item()) + residue_mask = torch.arange(maximum).unsqueeze(0) < lengths.unsqueeze(1) # (b, l) + reference = torch.ones(len(cases), maximum, 4) # (b, l, d=4) + candidate = reference * candidate_scale + return batch, candidate, reference, residue_mask + + +def _candidate_model( + runtime_revision: str = f"source-tree-sha256:{SOURCE_TREE_SHA256}", +) -> SimpleNamespace: + config = SimpleNamespace( + fastplms_model_id=SPEC.id, + fastplms_checkpoint_repo_id=SPEC.artifact_checkpoint.repo_id, + fastplms_checkpoint_revision=SPEC.artifact_checkpoint.revision, + fastplms_weights_revision=SPEC.artifact_checkpoint.revision, + fastplms_runtime_revision=runtime_revision, + fastplms_source_tree_sha256=SOURCE_TREE_SHA256, + fastplms_runtime_bundle_sha256=RUNTIME_BUNDLE_SHA256, + _commit_hash=SPEC.fast.revision, + ) + return SimpleNamespace(config=config) + + +def _reference_metadata() -> dict[str, object]: + locked_environment = _reference_environment_payload() + runtime = locked_environment["runtime"] + if not isinstance(runtime, dict): + raise AssertionError("Synthetic locked reference runtime is malformed") + gpu = runtime["gpu"] + if not isinstance(gpu, dict): + raise AssertionError("Synthetic locked reference GPU is malformed") + return { + "reference_repo_id": SPEC.official.repo_id, + "reference_revision": SPEC.official.revision, + "state_transform": SPEC.family.state_transform, + "reference_sources": copy.deepcopy(REFERENCE_SOURCES), + "reference_environment": locked_environment, + "environment": { + "cuda_device": gpu["name"], + "cuda_device_capability": copy.deepcopy(gpu["capability"]), + "cuda_total_memory": gpu["total_memory_bytes"], + "cuda_runtime": runtime["cuda_runtime"], + "packages": json.dumps( + {"python": runtime["python_version"], "torch": runtime["torch"]}, + separators=(",", ":"), + sort_keys=True, + ), + "python": runtime["python_version"], + "torch": runtime["torch"], + }, + } + + +def _candidate_environment() -> dict[str, object]: + locked_environment = _reference_environment_payload() + runtime = locked_environment["runtime"] + if not isinstance(runtime, dict): + raise AssertionError("Synthetic locked reference runtime is malformed") + return { + "python": runtime["python_version"], + "torch": runtime["torch"], + "transformers": "5.13.0", + "cuda_runtime": runtime["cuda_runtime"], + "cuda_driver": runtime["cuda_driver"], + "gpu": copy.deepcopy(runtime["gpu"]), + "packages": { + "fastplms": "1.0.0", + "huggingface-hub": "1.4.0", + "kernels": "0.12.2", + "tokenizers": "0.22.2", + "transformer-engine": None, + "transformer-engine-torch": None, + }, + } + + +def _patch_candidate_environment(monkeypatch: pytest.MonkeyPatch) -> None: + environment = _candidate_environment() + monkeypatch.setattr( + diagnostics, + "_candidate_environment_identity", + lambda: copy.deepcopy(environment), + ) + + +def _build_report( + *, + backend: str, + kind: str, + candidate_scale: float = 1.0, + runtime_revision: str = f"source-tree-sha256:{SOURCE_TREE_SHA256}", +) -> dict[str, object]: + batch, candidate, reference, residue_mask = _panel_tensors( + kind, + candidate_scale=candidate_scale, + ) + return diagnostics._build_esmc_diagnostic_report( + SPEC, + _output(candidate), + _output(reference), + residue_mask, + backend=backend, + effective_backend=backend, + context=f"{SPEC.id}:bf16:{backend}:{kind}", + calibration_batch=batch, + model=_candidate_model(runtime_revision), + reference_metadata=_reference_metadata(), + ) + + +def test_esmc_schema_v3_partitions_every_advertised_bf16_backend() -> None: + assert (*diagnostics.ESMC_MEASURED_BACKENDS, *diagnostics.ESMC_UNAVAILABLE_BACKENDS) == ( + ADVERTISED_BF16_BACKENDS + ) + + +@pytest.mark.parametrize( + ("gpu_name", "architecture", "memory"), + ( + ("NVIDIA H100 80GB HBM3", "x86_64", 80 * 1024**3), + ("NVIDIA H200", "x86_64", 141 * 1024**3), + ("NVIDIA GH200 480GB", "aarch64", 480_000_000_000), + ), +) +def test_esmc_dynamic_environment_schema_is_hardware_neutral_and_exactly_bound( + monkeypatch: pytest.MonkeyPatch, + gpu_name: str, + architecture: str, + memory: int, +) -> None: + candidate = _candidate_environment() + gpu = { + "name": gpu_name, + "capability": [9, 0], + "total_memory_bytes": memory, + } + candidate["gpu"] = copy.deepcopy(gpu) + monkeypatch.setattr( + diagnostics, + "_candidate_environment_identity", + lambda: copy.deepcopy(candidate), + ) + diagnostics._validate_candidate_environment(candidate) + + dynamic_reference = { + "cuda_device": gpu_name, + "cuda_device_capability": [9, 0], + "cuda_total_memory": memory, + "cuda_runtime": candidate["cuda_runtime"], + "packages": json.dumps({"torch": candidate["torch"]}), + "python": candidate["python"], + "torch": candidate["torch"], + } + diagnostics._validate_reference_environment(dynamic_reference) + locked_reference = { + "runtime": { + "operating_system": "linux", + "architecture": architecture, + "python_version": candidate["python"], + "torch": candidate["torch"], + "cuda_runtime": candidate["cuda_runtime"], + "cuda_driver": candidate["cuda_driver"], + "gpu": copy.deepcopy(gpu), + } + } + diagnostics._validate_esmc_environment_binding( + candidate, + dynamic_reference, + locked_reference, + ) + unavailable = diagnostics._esmc_unavailability_identity("flash_attention_3", locked_reference) + assert unavailable["platform"] == f"linux/{architecture}" + assert unavailable["accelerator"] == f"{gpu_name}/SM90" + + +def test_esmc_environment_binding_rejects_mismatch_and_malformed_capability() -> None: + candidate = _candidate_environment() + dynamic_reference = _reference_metadata()["environment"] + locked_reference = _reference_environment_payload() + if not isinstance(dynamic_reference, dict): + raise AssertionError("Synthetic dynamic reference environment is malformed") + diagnostics._validate_esmc_environment_binding( + candidate, + dynamic_reference, + locked_reference, + ) + + mismatched = copy.deepcopy(dynamic_reference) + mismatched["cuda_device"] = "NVIDIA H100 80GB HBM3" + with pytest.raises(ValueError, match="native reference environments differ"): + diagnostics._validate_esmc_environment_binding( + candidate, + mismatched, + locked_reference, + ) + + malformed = copy.deepcopy(candidate) + malformed_gpu = malformed["gpu"] + if not isinstance(malformed_gpu, dict): + raise AssertionError("Synthetic candidate GPU is malformed") + malformed_gpu["capability"] = [True, 0] + with pytest.raises(ValueError, match="GPU identity is malformed"): + diagnostics._validate_candidate_environment(malformed) + + +@pytest.mark.parametrize("backend", diagnostics.ESMC_MEASURED_BACKENDS) +@pytest.mark.parametrize("kind", PANEL_KINDS) +def test_esmc_schema_v3_covers_every_measured_bf16_backend_and_panel( + monkeypatch: pytest.MonkeyPatch, + backend: str, + kind: str, +) -> None: + _patch_candidate_environment(monkeypatch) + report = _build_report(backend=backend, kind=kind) + + assert report["schema_version"] == 3 + assert report["model_id"] == SPEC.id + assert report["configured_backend"] == backend + assert report["effective_backend"] == backend + assert report["record_status"] == "measured" + assert report["unavailability"] is None + assert report["dtype"] == "bfloat16" + assert report["catastrophic_gate"] == "passed" + assert report["release_gate"] == { + "mode": { + "sdpa": "exact", + "eager": "strict_numeric", + "flex_attention": "diagnostic_with_catastrophe_gate", + }[backend], + "status": "passed", + } + assert report["environment"] == _candidate_environment() + assert report["candidate"] == { + "repo_id": SPEC.fast.repo_id, + "manifest_revision": SPEC.fast.revision, + "resolved_commit": SPEC.fast.revision, + "checkpoint_repo_id": SPEC.artifact_checkpoint.repo_id, + "checkpoint_revision": SPEC.artifact_checkpoint.revision, + "weights_revision": SPEC.artifact_checkpoint.revision, + "runtime_revision": f"source-tree-sha256:{SOURCE_TREE_SHA256}", + "source_tree_sha256": SOURCE_TREE_SHA256, + "runtime_bundle_sha256": RUNTIME_BUNDLE_SHA256, + } + assert report["reference"] == diagnostics._reference_identity( + SPEC, + _reference_metadata(), + ) + assert report["reference"]["reference_sources"] == REFERENCE_SOURCES + panel = report["panel"] + assert isinstance(panel, Mapping) + assert panel["kind"] == kind + assert panel["seed"] == ESMC_CALIBRATION_SEED + assert len(str(panel["definition_sha256"])) == 64 + cases = report["cases"] + panel_cases = panel["cases"] + assert isinstance(cases, list) + assert isinstance(panel_cases, list) + assert len(cases) == len(panel_cases) + assert report["published_band_violations"] == [] + assert report["report_sha256"] == diagnostics._report_sha256(report) + diagnostics._validate_esmc_diagnostic_report(report, SPEC) + + panel_metrics = report["panel_tensor_metrics"] + assert isinstance(panel_metrics, list) + assert [(metric["output"], metric["layer_index"]) for metric in panel_metrics] == [ + ("hidden_state", 0), + ("hidden_state", 1), + ("last_hidden_state", None), + ("logits", None), + ] + for case, panel_case in zip(cases, panel_cases, strict=True): + assert { + name: case[name] + for name in ( + "case_id", + "sequence_length", + "sequence_sha256", + "source", + "source_sha256", + ) + } == panel_case + assert len(case["sequence_sha256"]) == 64 + assert [(metric["output"], metric["layer_index"]) for metric in case["tensor_metrics"]] == [ + ("hidden_state", 0), + ("hidden_state", 1), + ("last_hidden_state", None), + ("logits", None), + ] + assert set(case["logits_metrics"]) == { + "confident_top1_agreement", + "mean_jsd", + } + + +@pytest.mark.parametrize("backend", diagnostics.ESMC_UNAVAILABLE_BACKENDS) +@pytest.mark.parametrize("kind", PANEL_KINDS) +def test_esmc_schema_v3_records_locked_flash_unavailability_without_metrics( + monkeypatch: pytest.MonkeyPatch, + backend: str, + kind: str, +) -> None: + _patch_candidate_environment(monkeypatch) + report = diagnostics._build_esmc_unavailable_report( + SPEC, + backend=backend, + calibration_batch=_batch(kind), + model=_candidate_model(), + reference_metadata=_reference_metadata(), + ) + + assert report["record_status"] == "unavailable" + assert report["configured_backend"] == backend + assert report["effective_backend"] is None + assert report["panel_tensor_metrics"] is None + assert report["panel_logits_metrics"] is None + assert report["catastrophic_gate"] == "not_run" + assert report["release_gate"] == {"mode": "availability", "status": "unavailable"} + reference = report["reference"] + assert isinstance(reference, Mapping) + locked_environment = reference["reference_environment"] + assert isinstance(locked_environment, Mapping) + assert report["unavailability"] == diagnostics._esmc_unavailability_identity( + backend, locked_environment + ) + assert report["cases"] == report["panel"]["cases"] + diagnostics._validate_esmc_diagnostic_report(report, SPEC) + + +@pytest.mark.parametrize( + "runtime_revision", + ( + "a" * 40, + f"source-tree-sha256:{SOURCE_TREE_SHA256}", + ), + ids=("clean-git-revision", "content-addressed-fallback"), +) +def test_esmc_accepts_both_artifact_runtime_revision_forms( + monkeypatch: pytest.MonkeyPatch, + runtime_revision: str, +) -> None: + _patch_candidate_environment(monkeypatch) + report = _build_report( + backend="sdpa", + kind="generated_kernel_boundary", + runtime_revision=runtime_revision, + ) + + candidate = report["candidate"] + assert isinstance(candidate, Mapping) + assert candidate["runtime_revision"] == runtime_revision + diagnostics._validate_esmc_diagnostic_report( + report, + SPEC, + expected_candidate=candidate, + ) + + +@pytest.mark.parametrize( + "runtime_revision", + ( + "main", + "A" * 40, + "source-tree-sha256:" + "f" * 64, + ), + ids=("symbolic-ref", "noncanonical-git", "wrong-source-digest"), +) +def test_esmc_rejects_runtime_revision_not_emitted_by_artifact_builder( + runtime_revision: str, +) -> None: + with pytest.raises(ValueError, match="clean Git revision or the exact source-tree"): + diagnostics._candidate_identity( + SPEC, + _candidate_model(runtime_revision), + ) + + +def test_esmc_report_must_match_validated_artifact_runtime_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_candidate_environment(monkeypatch) + report = _build_report( + backend="sdpa", + kind="real_biological_holdout", + runtime_revision="a" * 40, + ) + expected_candidate = copy.deepcopy(report["candidate"]) + report["candidate"]["runtime_revision"] = "b" * 40 + report["report_sha256"] = diagnostics._report_sha256(report) + + with pytest.raises(ValueError, match="validated artifact identity"): + diagnostics._validate_esmc_diagnostic_report( + report, + SPEC, + expected_candidate=expected_candidate, + ) + + +def test_native_result_reference_identity_raises_explicitly( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + result_dir = tmp_path / SPEC.id + result_dir.mkdir() + metadata = _reference_metadata() + (result_dir / "metadata.json").write_text( + json.dumps(metadata, sort_keys=True), + encoding="utf-8", + ) + monkeypatch.setenv("FASTPLMS_REFERENCE_RESULTS", str(tmp_path)) + + observed, directory = diagnostics._result(SPEC) + assert observed == metadata + assert directory == result_dir + + metadata["reference_repo_id"] = "untrusted/reference" + (result_dir / "metadata.json").write_text( + json.dumps(metadata, sort_keys=True), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="reference_repo_id"): + diagnostics._result(SPEC) + + metadata["reference_repo_id"] = SPEC.official.repo_id + metadata["reference_sources"]["biohub-transformers"]["tree_sha256"] = "0" * 64 + (result_dir / "metadata.json").write_text( + json.dumps(metadata, sort_keys=True), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="reference source evidence tree_sha256"): + diagnostics._result(SPEC) + + +@pytest.mark.parametrize( + ("backend", "kind"), + (("flex_attention", "generated_kernel_boundary"),), +) +def test_esmc_supported_backend_deviation_warns_and_writes_complete_metrics( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + backend: str, + kind: str, +) -> None: + _patch_candidate_environment(monkeypatch) + monkeypatch.setenv("FASTPLMS_DIAGNOSTIC_REPORTS", str(tmp_path)) + batch, candidate, reference, residue_mask = _panel_tensors( + kind, + candidate_scale=1.04, + ) + + with pytest.warns( + UserWarning, + match=( + rf"configured backend={backend}, effective backend={backend}.*" + r"outside the published ESMC backend bands" + ), + ) as diagnostic_warnings: + diagnostics._assert_and_record_esmc_diagnostic( + SPEC, + _output(candidate), + _output(reference), + residue_mask, + backend=backend, + effective_backend=backend, + context=f"{SPEC.id}:bf16:{backend}:{kind}", + calibration_batch=batch, + model=_candidate_model(), + reference_metadata=_reference_metadata(), + ) + + assert len(diagnostic_warnings) == 1 + report_path = tmp_path / f"{SPEC.id}-{backend}-{kind}.json" + report = json.loads(report_path.read_text(encoding="utf-8")) + diagnostics._validate_esmc_diagnostic_report(report, SPEC) + assert report["published_band_violations"] + assert any("case=" in violation for violation in report["published_band_violations"]) + assert report["kernel"]["implementation"] == backend + assert report["kernel"]["provider"] == ( + "torch" if backend == "flex_attention" else "huggingface_kernels" + ) + + +def test_esmc_catastrophic_disagreement_remains_a_hard_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _patch_candidate_environment(monkeypatch) + monkeypatch.setenv("FASTPLMS_DIAGNOSTIC_REPORTS", str(tmp_path)) + kind = "real_biological_holdout" + batch, _, reference, residue_mask = _panel_tensors( + kind, + candidate_scale=1.0, + ) + candidate = torch.zeros_like(reference) # (b, l, d) + + with pytest.raises(AssertionError, match="relative_l2"): + diagnostics._assert_and_record_esmc_diagnostic( + SPEC, + _output(candidate), + _output(reference), + residue_mask, + backend="flex_attention", + effective_backend="flex_attention", + context=f"{SPEC.id}:bf16:flex_attention:{kind}", + calibration_batch=batch, + model=_candidate_model(), + reference_metadata=_reference_metadata(), + ) + assert not tuple(tmp_path.glob("*.json")) + + +def test_esmc_immutable_panels_fail_closed_on_drift() -> None: + seed_drift = _batch("generated_kernel_boundary") + seed_drift["seed"] = ESMC_CALIBRATION_SEED + 1 + with pytest.raises(ValueError, match="seed differs"): + validate_esmc_calibration_batch(seed_drift) + + sequence_drift = _batch("real_biological_holdout") + cases = sequence_drift["cases"] + assert isinstance(cases, list) + cases[0]["sequence_sha256"] = "0" * 64 + with pytest.raises(ValueError, match="differs from the release contract"): + validate_esmc_calibration_batch(sequence_drift) + + +def test_reference_request_uses_the_shared_immutable_panels() -> None: + assert _esmc_calibration_batches() == list(esmc_calibration_batches()) + for batch in _esmc_calibration_batches(): + assert validate_esmc_calibration_batch(batch)["definition_sha256"] + + +def test_esmc_schema_rejects_stale_identity_and_tampering( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_candidate_environment(monkeypatch) + report = _build_report( + backend="flex_attention", + kind="generated_kernel_boundary", + ) + + stale = copy.deepcopy(report) + stale["candidate"]["weights_revision"] = "f" * 40 + stale["report_sha256"] = diagnostics._report_sha256(stale) + with pytest.raises(ValueError, match="weights_revision differs"): + diagnostics._validate_esmc_diagnostic_report(stale, SPEC) + + fallback = copy.deepcopy(report) + fallback["effective_backend"] = "sdpa" + fallback["report_sha256"] = diagnostics._report_sha256(fallback) + with pytest.raises(ValueError, match="fallback"): + diagnostics._validate_esmc_diagnostic_report(fallback, SPEC) + + panel_drift = copy.deepcopy(report) + panel_drift["panel"]["definition_sha256"] = "0" * 64 + panel_drift["report_sha256"] = diagnostics._report_sha256(panel_drift) + with pytest.raises(ValueError, match="immutable definition"): + diagnostics._validate_esmc_diagnostic_report(panel_drift, SPEC) + + source_drift = copy.deepcopy(report) + source_drift["reference"]["reference_sources"]["biohub-esm"]["tree_sha256"] = "0" * 64 + source_drift["report_sha256"] = diagnostics._report_sha256(source_drift) + with pytest.raises(ValueError, match="reference source evidence tree_sha256"): + diagnostics._validate_esmc_diagnostic_report(source_drift, SPEC) + + nonfinite = copy.deepcopy(report) + nonfinite["panel_tensor_metrics"][0]["relative_l2"] = float("nan") + nonfinite["report_sha256"] = diagnostics._report_sha256(nonfinite) + with pytest.raises(ValueError, match="finite number"): + diagnostics._validate_esmc_diagnostic_report(nonfinite, SPEC) + + digest_mismatch = copy.deepcopy(report) + digest_mismatch["panel_tensor_metrics"][0]["context"] += ":tampered" + with pytest.raises(ValueError, match="digest does not match"): + diagnostics._validate_esmc_diagnostic_report(digest_mismatch, SPEC) + + +def test_esmc_report_write_is_atomic_idempotent_and_no_clobber( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _patch_candidate_environment(monkeypatch) + monkeypatch.setenv("FASTPLMS_DIAGNOSTIC_REPORTS", str(tmp_path)) + report = _build_report( + backend="sdpa", + kind="generated_kernel_boundary", + ) + + path = diagnostics._write_esmc_diagnostic_report(SPEC, report) + assert diagnostics._write_esmc_diagnostic_report(SPEC, report) == path + + different = copy.deepcopy(report) + different["panel_tensor_metrics"][0]["context"] += ":second-run" + different["report_sha256"] = diagnostics._report_sha256(different) + with pytest.raises(RuntimeError, match="Refusing to replace different ESMC evidence"): + diagnostics._write_esmc_diagnostic_report(SPEC, different) + assert len(tuple(tmp_path.glob("*.json"))) == 1 + assert not tuple(tmp_path.glob("*.tmp")) + + +def test_esmc_calibration_contains_no_expected_failures() -> None: + source = (Path(__file__).resolve().parents[1] / "parity" / "test_native_results.py").read_text( + encoding="utf-8" + ) + assert "pytest.mark.xfail" not in source diff --git a/tests/unit/test_esmfold2_leaf_contracts.py b/tests/unit/test_esmfold2_leaf_contracts.py new file mode 100644 index 0000000..c84be11 --- /dev/null +++ b/tests/unit/test_esmfold2_leaf_contracts.py @@ -0,0 +1,176 @@ +"""Fast ESMFold2 schema and protein-feature regression gates.""" + +from __future__ import annotations + +import hashlib +import io +import json +import subprocess +import sys +import torch + +from fastplms.models.esmfold2 import esmfold2_constants as molecular_schema +from fastplms.models.esmfold2 import esmfold2_constants_esm3 as token_schema +from fastplms.models.esmfold2.esmfold2_parsing import parse_fasta, read_sequences +from fastplms.models.esmfold2.protein_utils import prepare_protein_features + + +_OFFICIAL_FEATURE_DIGEST = "255bab0048c0c8984e03eb878b7a9c88dfa0ad083983e2641d6a4a0eb8e39fc5" +_OFFICIAL_SCHEMA_DIGEST = "34b14af14c0f640034eec8174f80081234a98d617a5dbf9682a06a04465dc024" +_DIGEST_SEQUENCES = ("ACGX", "ARNDCQEGHILKMFPSTWYV", "M" * 31, "M" * 33) + +_MOLECULAR_SCHEMA_NAMES = ( + "MOL_TYPE_PROTEIN", + "MOL_TYPE_DNA", + "MOL_TYPE_RNA", + "MOL_TYPE_NONPOLYMER", + "PROTEIN_RESIDUE_TO_RES_TYPE", + "PROTEIN_UNK_RES_TYPE", + "RNA_RESIDUE_TO_RES_TYPE", + "RNA_UNK_RES_TYPE", + "DNA_RESIDUE_TO_RES_TYPE", + "DNA_UNK_RES_TYPE", + "GAP_RES_TYPE", + "PROTEIN_3TO1", + "PROTEIN_1TO3", + "DNA_1TO3", + "RNA_1TO3", + "ESM_PROTEIN_VOCAB", + "DNA_RNA_LIGAND_INPUT_ID", + "MSA_PAD_TOKEN_ID", + "MSA_GAP_TOKEN_ID", + "RES_TYPE_TO_CCD", + "ELEMENT_TO_ATOMIC_NUM", + "ELEMENT_NUMBER_TO_SYMBOL", + "PROTEIN_HEAVY_ATOMS", + "DNA_HEAVY_ATOMS", + "RNA_HEAVY_ATOMS", + "DNA_BACKBONE_ATOMS", + "RNA_BACKBONE_ATOMS", +) +_TOKEN_SCHEMA_NAMES = ( + "SEQUENCE_BOS_TOKEN", + "SEQUENCE_PAD_TOKEN", + "SEQUENCE_EOS_TOKEN", + "SEQUENCE_CHAINBREAK_TOKEN", + "SEQUENCE_MASK_TOKEN", + "VQVAE_CODEBOOK_SIZE", + "VQVAE_SPECIAL_TOKENS", + "VQVAE_DIRECTION_LOSS_BINS", + "VQVAE_PAE_BINS", + "VQVAE_MAX_PAE_BIN", + "VQVAE_PLDDT_BINS", + "STRUCTURE_MASK_TOKEN", + "STRUCTURE_BOS_TOKEN", + "STRUCTURE_EOS_TOKEN", + "STRUCTURE_PAD_TOKEN", + "STRUCTURE_CHAINBREAK_TOKEN", + "STRUCTURE_UNDEFINED_TOKEN", + "SASA_PAD_TOKEN", + "SS8_PAD_TOKEN", + "INTERPRO_PAD_TOKEN", + "RESIDUE_PAD_TOKEN", + "CHAIN_BREAK_STR", + "SEQUENCE_BOS_STR", + "SEQUENCE_EOS_STR", + "MASK_STR_SHORT", + "SEQUENCE_MASK_STR", + "SASA_MASK_STR", + "SS8_MASK_STR", + "SEQUENCE_VOCAB", + "SEQUENCE_STANDARD_AA_MIN_TOKEN", + "SEQUENCE_STANDARD_AA_MAX_TOKEN", + "SSE_8CLASS_VOCAB", + "SSE_3CLASS_VOCAB", + "SSE_8CLASS_TO_3CLASS_MAP", + "SASA_DISCRETIZATION_BOUNDARIES", + "MAX_RESIDUE_ANNOTATIONS", + "TFIDF_VECTOR_SIZE", + "FUNCTION_TOKENS_DEPTH", +) + + +def _feature_digest() -> str: + digest = hashlib.sha256() + for sequence in _DIGEST_SEQUENCES: + for name, tensor in sorted(prepare_protein_features(sequence).items()): + value = tensor.detach().cpu().contiguous() + digest.update(sequence.encode()) + digest.update(name.encode()) + digest.update(str(value.dtype).encode()) + digest.update(json.dumps(list(value.shape)).encode()) + digest.update(value.numpy().tobytes()) + return digest.hexdigest() + + +def _schema_digest() -> str: + payload = {name: getattr(molecular_schema, name) for name in _MOLECULAR_SCHEMA_NAMES} + payload["CHARGED_ATOMS"] = sorted( + [*key, value] for key, value in molecular_schema.CHARGED_ATOMS.items() + ) + payload["tokens"] = {name: getattr(token_schema, name) for name in _TOKEN_SCHEMA_NAMES} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode()).hexdigest() + + +def test_protein_features_match_pinned_official_tensor_digest() -> None: + """The digest was generated from Biohub Transformers at its manifest revision.""" + + assert _feature_digest() == _OFFICIAL_FEATURE_DIGEST + + +def test_generated_schemas_match_pinned_official_semantic_digest() -> None: + assert _schema_digest() == _OFFICIAL_SCHEMA_DIGEST + + +def test_protein_features_preserve_padding_and_unknown_residues() -> None: + # Feature dimensions are b=1, l=2 residues, a=32 reference atoms, xyz=3. + features = prepare_protein_features("GX") + assert features["res_type"].tolist() == [[9, 22]] + assert features["input_ids"].tolist() == [[6, 3]] + assert features["ref_pos"].shape == (1, 32, 3) + assert features["atom_attention_mask"].sum().item() == 8 + assert torch.equal(features["msa"], features["res_type"].unsqueeze(1)) + + +def test_generated_token_and_molecular_schemas_keep_checkpoint_indices() -> None: + assert len(token_schema.SEQUENCE_VOCAB) == 33 + assert token_schema.SEQUENCE_VOCAB[4:24] == list("LAGVSERTIDPKQNFYMHWC") + assert token_schema.VQVAE_SPECIAL_TOKENS == { + "MASK": 4096, + "EOS": 4097, + "BOS": 4098, + "PAD": 4099, + "CHAINBREAK": 4100, + } + assert molecular_schema.PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] == 14 + assert molecular_schema.RES_TYPE_TO_CCD[32] == "DN" + assert molecular_schema.ELEMENT_TO_ATOMIC_NUM["U"] == 92 + assert 2 not in molecular_schema.ELEMENT_NUMBER_TO_SYMBOL + + +def test_fasta_stream_order_and_ownership() -> None: + source = io.StringIO(">first\nAC\n>second\nGX\n") + assert list(read_sequences(source)) == [("first", "AC"), ("second", "GX")] + assert not source.closed + assert list(parse_fasta("# note\n>x\nAA\n")) == [("x", "AA")] + + +def test_esmfold2_package_init_is_lazy() -> None: + probe = """ +import sys +import fastplms.models.esmfold2 as package + +assert package.__all__ == [ + "ESMFold2Config", + "ESMFold2ExperimentalModel", + "ESMFold2Model", +] +assert "fastplms.models.esmfold2.modeling_esmfold2" not in sys.modules +""" + subprocess.run( + [sys.executable, "-c", probe], + check=True, + capture_output=True, + text=True, + ) diff --git a/tests/unit/test_esmfold2_public_contracts.py b/tests/unit/test_esmfold2_public_contracts.py new file mode 100644 index 0000000..e84c2a7 --- /dev/null +++ b/tests/unit/test_esmfold2_public_contracts.py @@ -0,0 +1,311 @@ +from __future__ import annotations + +import os +import subprocess +import sys +import pytest +import torch +from pathlib import Path +from types import MethodType, SimpleNamespace + +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.esmfold2_msa import MSA +from fastplms.models.esmfold2.esmfold2_processor import ESMFold2InputBuilder +from fastplms.models.esmfold2.esmfold2_types import ProteinInput, StructurePredictionInput +from fastplms.models.esmfold2.modeling_esmfold2 import ESMFold2Model +from fastplms.models.esmfold2.modeling_esmfold2_common import ( + MSA_CONDITIONING_INPUT_NAMES, + validate_msa_conditioning_inputs, +) +from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, +) +from fastplms.registry import RegistryError, load_model_registry +from tools.artifacts.build import ( + ArtifactError, + _apply_artifact_config_contract, + _expected_registry_provenance, +) + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_esmfold2_msa_conditioning_is_manifest_typed() -> None: + registry = load_model_registry() + assert { + spec.id: spec.msa_conditioning for spec in registry.by_family("esmfold2") + } == { + "esmfold2": True, + "esmfold2_fast": False, + "esmfold2_experimental_cutoff2025": True, + "esmfold2_experimental_fast_cutoff2025": False, + } + assert all( + spec.msa_conditioning is None + for spec in registry.values() + if spec.family.id != "esmfold2" + ) + + +def test_esmfold2_msa_conditioning_is_required_and_family_scoped(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + candidate = tmp_path / "models.toml" + candidate.write_text(manifest.replace("msa_conditioning = true\n", "", 1), encoding="utf-8") + with pytest.raises(RegistryError, match="must be an explicit boolean"): + load_model_registry(candidate) + + candidate.write_text( + manifest.replace('id = "esm2_8m"\n', 'id = "esm2_8m"\nmsa_conditioning = false\n', 1), + encoding="utf-8", + ) + with pytest.raises(RegistryError, match="only valid for ESMFold2"): + load_model_registry(candidate) + + +@pytest.mark.parametrize("value", [None, "false", 0, 1]) +def test_esmfold2_config_rejects_invalid_msa_contract_values(value: object) -> None: + with pytest.raises(TypeError, match="msa_conditioning must be a boolean"): + ESMFold2Config(msa_encoder={"enabled": False}, msa_conditioning=value) + + +def test_esmfold2_config_requires_encoder_contract_agreement() -> None: + with pytest.raises(ValueError, match=r"must match msa_encoder\.enabled"): + ESMFold2Config(msa_encoder={"enabled": True}, msa_conditioning=False) + config = ESMFold2Config(msa_encoder={"enabled": True}) + assert config.msa_conditioning is True + assert config.to_dict()["msa_conditioning"] is True + + +@pytest.mark.parametrize("msa_conditioning", [False, True]) +def test_esmfold2_config_msa_contract_survives_save_reload( + tmp_path: Path, + msa_conditioning: bool, +) -> None: + config = ESMFold2Config( + msa_encoder={"enabled": msa_conditioning}, + msa_conditioning=msa_conditioning, + ) + config.save_pretrained(tmp_path) + reloaded = ESMFold2Config.from_pretrained(tmp_path) + assert reloaded.msa_conditioning is msa_conditioning + assert reloaded.msa_encoder.enabled is msa_conditioning + assert reloaded.to_dict()["msa_conditioning"] is msa_conditioning + + +def test_artifact_config_materializes_manifest_msa_contract() -> None: + registry = load_model_registry() + for model_id, expected in ( + ("esmfold2", True), + ("esmfold2_fast", False), + ("esmfold2_experimental_cutoff2025", True), + ("esmfold2_experimental_fast_cutoff2025", False), + ): + config = {"msa_encoder": {"enabled": expected}} + _apply_artifact_config_contract(registry[model_id], config) + assert config["msa_conditioning"] is expected + assert config["msa_encoder"]["enabled"] is expected + + with pytest.raises(ArtifactError, match="msa_encoder must be an object"): + _apply_artifact_config_contract(registry["esmfold2_fast"], {}) + + +def test_artifact_provenance_materializes_manifest_msa_contract() -> None: + registry = load_model_registry() + for spec in registry.by_family("esmfold2"): + provenance = _expected_registry_provenance(registry, spec) + assert provenance["msa_conditioning"] is spec.msa_conditioning + assert "msa_conditioning" not in _expected_registry_provenance( + registry, + registry["esm2_8m"], + ) + + +@pytest.mark.parametrize( + ("model_id", "config_enabled"), + (("esmfold2", False), ("esmfold2_fast", True)), +) +def test_artifact_config_rejects_manifest_msa_disagreement( + model_id: str, + config_enabled: bool, +) -> None: + registry = load_model_registry() + with pytest.raises(ArtifactError, match=r"differs from models\.toml"): + _apply_artifact_config_contract( + registry[model_id], + {"msa_encoder": {"enabled": config_enabled}}, + ) + + expected = registry[model_id].msa_conditioning + with pytest.raises(ArtifactError, match=r"config\.msa_conditioning differs"): + _apply_artifact_config_contract( + registry[model_id], + { + "msa_encoder": {"enabled": expected}, + "msa_conditioning": not expected, + }, + ) + + +@pytest.mark.parametrize("provided_name", MSA_CONDITIONING_INPUT_NAMES) +def test_fast_checkpoint_rejects_every_low_level_msa_input(provided_name: str) -> None: + config = ESMFold2Config(msa_encoder={"enabled": False}, msa_conditioning=False) + values = {name: None for name in MSA_CONDITIONING_INPUT_NAMES} + values[provided_name] = torch.zeros(1) # (n=1,) + with pytest.raises(ValueError, match=provided_name): + validate_msa_conditioning_inputs(config, **values) + + +def test_full_checkpoint_accepts_low_level_msa_inputs() -> None: + config = ESMFold2Config(msa_encoder={"enabled": True}, msa_conditioning=True) + tensor = torch.zeros(1) # (n=1,) + validate_msa_conditioning_inputs( + config, + msa=tensor, + msa_attention_mask=tensor, + has_deletion=tensor, + deletion_value=tensor, + deletion_mean=tensor, + ) + + +@pytest.mark.parametrize("model_type", [ESMFold2Model, ESMFold2ExperimentalModel]) +@pytest.mark.parametrize("provided_name", MSA_CONDITIONING_INPUT_NAMES) +def test_fast_model_forward_rejects_every_msa_input_before_computation( + model_type: type[ESMFold2Model] | type[ESMFold2ExperimentalModel], + provided_name: str, +) -> None: + config = ESMFold2Config(msa_encoder={"enabled": False}, msa_conditioning=False) + model = SimpleNamespace(config=config) + tensor = torch.zeros(1) # (n=1,) + required = { + name: tensor + for name in ( + "token_index", + "residue_index", + "asym_id", + "sym_id", + "entity_id", + "mol_type", + "res_type", + "token_bonds", + "token_attention_mask", + "ref_pos", + "ref_element", + "ref_charge", + "ref_atom_name_chars", + "ref_space_uid", + "atom_attention_mask", + "atom_to_token", + "distogram_atom_idx", + ) + } + required[provided_name] = tensor + with pytest.raises(ValueError, match=provided_name): + model_type.forward(model, **required) + + +def _builder_with_features(features: dict[str, torch.Tensor]) -> ESMFold2InputBuilder: + builder = object.__new__(ESMFold2InputBuilder) + + def prepare_input( + self: ESMFold2InputBuilder, + input: object, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, torch.Tensor], list[object]]: + del self, input, seed, device + return dict(features), [] + + builder.prepare_input = MethodType(prepare_input, builder) + return builder + + +def test_fast_high_level_input_rejects_explicit_msa_and_strips_synthetic_features() -> None: + features = { + name: torch.zeros(1) # (n=1,) + for name in MSA_CONDITIONING_INPUT_NAMES + } + features["token_index"] = torch.zeros(1) # (n=1,) + builder = _builder_with_features(features) + model = SimpleNamespace(config=SimpleNamespace(msa_conditioning=False)) + sequence_only = StructurePredictionInput( + sequences=[ProteinInput(id="A", sequence="AC")] + ) + prepared, _ = builder.prepare_model_input(model, sequence_only) + assert set(prepared) == {"token_index"} + + explicit = StructurePredictionInput( + sequences=[ + ProteinInput(id="A", sequence="AC", msa=MSA.from_sequences(["AC", "AC"])) + ] + ) + with pytest.raises(ValueError, match="rejects explicit MSAs"): + builder.prepare_model_input(model, explicit) + + +def test_full_high_level_input_preserves_msa_features() -> None: + features = { + name: torch.zeros(1) # (n=1,) + for name in MSA_CONDITIONING_INPUT_NAMES + } + builder = _builder_with_features(features) + model = SimpleNamespace(config=SimpleNamespace(msa_conditioning=True)) + explicit = StructurePredictionInput( + sequences=[ + ProteinInput(id="A", sequence="AC", msa=MSA.from_sequences(["AC", "AC"])) + ] + ) + prepared, _ = builder.prepare_model_input(model, explicit) + assert set(prepared) == set(MSA_CONDITIONING_INPUT_NAMES) + + +def test_public_validation_survives_python_optimized_mode() -> None: + script = r''' +import numpy as np +import torch + +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.esmfold2_affine3d import ( + Affine3D, + RotationMatrix, + RotationQuat, + build_affine3d_from_coordinates, +) +from fastplms.models.esmfold2.esmfold2_misc import concat_objects, merge_ranges +from fastplms.models.esmfold2.esmfold2_protein_chain import ProteinChain +from fastplms.models.esmfold2.modeling_esmfold2_common import DropoutResidual + +def expect(error_type, fn): + try: + fn() + except error_type: + return + raise RuntimeError(f"expected {error_type.__name__}") + +expect(TypeError, lambda: RotationQuat([1, 0, 0, 0])) +expect(ValueError, lambda: RotationQuat(torch.zeros(3))) +expect(ValueError, lambda: RotationMatrix(torch.zeros(2, 2))) +expect(ValueError, lambda: Affine3D(torch.zeros(2, 2), RotationMatrix.identity((2,)))) +expect(ValueError, lambda: build_affine3d_from_coordinates(torch.zeros(2, 3, 3))) +expect(ValueError, lambda: DropoutResidual(0.1, batch_dim=0)) +expect(ValueError, lambda: DropoutResidual(0.1, batch_dim=True)) +expect(TypeError, lambda: concat_objects(["A", "B"], separator=1)) +expect(ValueError, lambda: merge_ranges([range(0, 1)], merge_gap_max=-1)) +expect(ValueError, lambda: ProteinChain.from_atom37(np.zeros((2, 36, 3)))) +expect( + ValueError, + lambda: ESMFold2Config(msa_encoder={"enabled": True}, msa_conditioning=False), +) +''' + env = dict(os.environ) + env["PYTHONPATH"] = str(ROOT / "src") + completed = subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + check=False, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr diff --git a/tests/unit/test_esmfold2_reimplemented_leaves.py b/tests/unit/test_esmfold2_reimplemented_leaves.py new file mode 100644 index 0000000..de4ca17 --- /dev/null +++ b/tests/unit/test_esmfold2_reimplemented_leaves.py @@ -0,0 +1,295 @@ +"""Behavioral contracts for independently organized ESMFold2 leaf utilities.""" + +from __future__ import annotations + +import dataclasses +import sys +import numpy as np +import pytest +import torch +from pathlib import Path +from types import SimpleNamespace + +from fastplms.models.esmfold2 import esmfold2_input_builder +from fastplms.models.esmfold2.esmfold2_aligner import Aligner +from fastplms.models.esmfold2.esmfold2_atom_indexer import AtomIndexer +from fastplms.models.esmfold2.esmfold2_msa_filter_sequences import ( + greedy_select_indices, + hhfilter, +) +from fastplms.models.esmfold2.esmfold2_normalize_coordinates import index_by_atom_name +from fastplms.models.esmfold2.esmfold2_predicted_aligned_error import ( + compute_predicted_aligned_error, + compute_tm, +) +from fastplms.models.esmfold2.esmfold2_protein_chain import ProteinChain +from fastplms.models.esmfold2.esmfold2_system import run_subprocess_with_errorcheck +from fastplms.models.esmfold2.esmfold2_types import ( + PocketConditioning, + ProteinInput, + StructurePredictionInput, +) + + +pytestmark = pytest.mark.structure + + +def test_greedy_msa_selection_preserves_official_tie_order() -> None: + sequences = np.asarray( # (n=5, l=4) + [list(row) for row in ("AAAA", "AAAT", "AATT", "TTTT", "ATAT")], + dtype="S1", + ) + assert greedy_select_indices(sequences, 3, mode="max") == [0, 1, 3] + assert greedy_select_indices(sequences, 3, mode="min") == [0, 1, 2] + assert greedy_select_indices(sequences, 10) == [0, 1, 2, 3, 4] + with pytest.raises(ValueError, match="unsupported selection mode"): + greedy_select_indices(sequences, 2, mode="median") + with pytest.raises(ValueError, match="greater than zero"): + greedy_select_indices(sequences, 0) + with pytest.raises(ValueError, match="non-empty shape"): + greedy_select_indices(np.empty((0, 4), dtype="S1"), 1) + + +def test_fast_msa_stack_preserves_headerless_and_mixed_inputs() -> None: + from fastplms.models.esmfold2.esmfold2_msa import FastMSA + + first = FastMSA(np.asarray([list("AAA"), list("AAT")], dtype="S1")) # (n=2, l=3) + second = FastMSA(np.asarray([list("AAA"), list("ATT")], dtype="S1")) # (n=2, l=3) + headerless = FastMSA.stack([first, second]) + + assert headerless.depth == 3 + assert headerless.headers is None + + with_headers = FastMSA( + np.asarray([list("AAA"), list("ATA")], dtype="S1"), # (n=2, l=3) + ["query", "named"], + ) + mixed = FastMSA.stack([first, with_headers]) + assert mixed.depth == 3 + assert mixed.headers == ["", "", "named"] + + +def test_hhfilter_passes_paths_as_distinct_arguments(tmp_path: Path) -> None: + executable = tmp_path / "fake_hhfilter.py" + executable.write_text( + "#!/usr/bin/env python3\n" + "import pathlib, sys\n" + "output = pathlib.Path(sys.argv[sys.argv.index('-o') + 1])\n" + "output.write_text('>2\\nCCC\\n>0\\nAAA\\n', encoding='utf-8')\n", + encoding="utf-8", + ) + executable.chmod(0o755) + assert hhfilter(["AAA", "BBB", "CCC"], binary=str(executable)) == [2, 0] + + +def test_subprocess_failure_includes_standard_error() -> None: + with pytest.raises(RuntimeError, match="intentional failure"): + run_subprocess_with_errorcheck( + [ + sys.executable, + "-c", + "import sys; sys.stderr.write('intentional failure'); sys.exit(4)", + ], + capture_output=True, + ) + + +def test_schema_namespace_preserves_type_identity() -> None: + assert ProteinInput is esmfold2_input_builder.ProteinInput + assert PocketConditioning is esmfold2_input_builder.PocketConditioning + assert StructurePredictionInput is esmfold2_input_builder.StructurePredictionInput + + +def test_msa_rejects_unequal_biological_rows() -> None: + from fastplms.models.esmfold2.esmfold2_types import MSA + + with pytest.raises(ValueError, match="MSA row length mismatch"): + MSA.from_sequences(["ACDE", "ACD"]) + + +def test_a3m_dot_insertions_and_raw_rows_have_consistent_metadata() -> None: + import io + + from fastplms.models.esmfold2.esmfold2_msa import MSA + + source = ">query\nA.CD\n>hit\nAaCD\n" + match_columns = MSA.from_a3m(io.StringIO(source)) + assert match_columns.sequences == ["ACD", "ACD"] + assert match_columns.deletions is not None + # n=2 aligned sequences, l=3 retained match columns. + assert match_columns.deletions.shape == (2, 3) + + raw_rows = MSA.from_a3m(io.StringIO(source), remove_insertions=False) + assert raw_rows.sequences == ["A.CD", "AaCD"] + assert raw_rows.deletions is None + + +def test_protein_chain_rejects_misaligned_atom37_tables() -> None: + with pytest.raises(ValueError, match=r"shape \(length, 37, 3\)"): + ProteinChain.from_atom37( + np.zeros((2, 36, 3), dtype=np.float32) # (l=2, a=36, xyz=3) + ) + + +def test_protein_chain_contacts_require_retained_mmcif_source() -> None: + chain = ProteinChain.from_atom37( + np.zeros((2, 37, 3), dtype=np.float32), # (l=2, a=37, xyz=3) + sequence="AC", + ) + with pytest.raises(ValueError, match="keep_source=True"): + chain.find_nonpolymer_contacts() + + +def test_unavailable_structure_kernel_backend_fails_before_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastplms.models.esmfold2 import modeling_esmfold2_common as common + + with pytest.raises(RuntimeError, match="does not bundle"): + common.validate_kernel_backend("fused") + monkeypatch.setattr(common, "CUE_AVAILABLE", False) + with pytest.raises( + RuntimeError, + match=r"requires cuequivariance_torch.*cuequivariance_ops_torch", + ): + common.validate_kernel_backend("cuequivariance") + with pytest.raises(ValueError, match="backend must be one of"): + common.validate_kernel_backend("silent-fallback") + + +def test_experimental_top_level_kernel_backend_validates_before_zero_layer_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastplms.models.esmfold2 import modeling_esmfold2_common as common + from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, + ) + + class ZeroLayerRecorder: + def __init__(self) -> None: + self.calls: list[str | None] = [] + + def set_kernel_backend(self, backend: str | None) -> None: + self.calls.append(backend) + + folding_trunk = ZeroLayerRecorder() + structure_head = ZeroLayerRecorder() + model = SimpleNamespace( + folding_trunk=folding_trunk, + confidence_head=None, + structure_head=structure_head, + _kernel_backend=None, + ) + monkeypatch.setattr(common, "CUE_AVAILABLE", False) + + with pytest.raises(RuntimeError, match="requires cuequivariance_torch"): + ESMFold2ExperimentalModel.set_kernel_backend(model, "cuequivariance") + assert folding_trunk.calls == [] + assert structure_head.calls == [] + assert model._kernel_backend is None + + ESMFold2ExperimentalModel.set_kernel_backend(model, None) + assert folding_trunk.calls == [None] + assert structure_head.calls == [None] + assert model._kernel_backend is None + + +def test_pocket_conditioning_is_rejected_instead_of_silently_dropped() -> None: + from fastplms.models.esmfold2.esmfold2_processor import clean_esmfold2_input + + request = StructurePredictionInput( + sequences=[ProteinInput(id="A", sequence="ACD")], + pocket=PocketConditioning(binder_chain_id="A", contacts=[("A", 0)]), + ) + with pytest.raises(NotImplementedError, match="refuses this input"): + clean_esmfold2_input(request) + + +def test_multiple_delimited_proteins_keep_their_own_split_identity() -> None: + from fastplms.models.esmfold2.esmfold2_processor import clean_esmfold2_input + + request = StructurePredictionInput( + sequences=[ + ProteinInput(id="first", sequence="AA:CC"), + ProteinInput(id="second", sequence="GG:TT"), + ] + ) + cleaned = clean_esmfold2_input(request) + + assert [protein.sequence for protein in cleaned.sequences] == ["AA", "CC", "GG", "TT"] + assert [protein.id for protein in cleaned.sequences] == [ + ["first_0"], + ["first_1"], + ["second_0"], + ["second_1"], + ] + + +@dataclasses.dataclass +class _Structure: + atom37_positions: np.ndarray + atom37_mask: np.ndarray + + def __len__(self) -> int: + return self.atom37_positions.shape[0] + + +def test_atom_indexer_selects_the_declared_property_and_axis() -> None: + positions = np.arange(2 * 37 * 3, dtype=np.float32).reshape( # (l=2, a=37, xyz=3) + 2, 37, 3 + ) + structure = _Structure(positions, np.ones((2, 37), dtype=bool)) # mask: (l=2, a=37) + indexer = AtomIndexer(structure, "atom37_positions", dim=1) + np.testing.assert_array_equal(indexer["CA"], positions[:, 1]) + np.testing.assert_array_equal(indexer[["N", "C"]], positions[:, [0, 2]]) + + +def test_atom_name_selection_matches_for_numpy_and_torch() -> None: + positions = np.arange(2 * 37 * 3, dtype=np.float32).reshape( # (l=2, a=37, xyz=3) + 2, 37, 3 + ) + expected = positions[:, [0, 1, 2]] # (l=2, a_selected=3, xyz=3) + np.testing.assert_array_equal( + index_by_atom_name(positions, ["N", "CA", "C"]), + expected, + ) + assert torch.equal( + index_by_atom_name(torch.from_numpy(positions), ["N", "CA", "C"]), + torch.from_numpy(expected), + ) + + +def test_pae_bin_expectation_and_tm_score_are_exact_for_uniform_logits() -> None: + logits = torch.zeros((1, 2, 2, 4), dtype=torch.float64) # (b=1, l=2, l=2, c=4) + mask = torch.ones((1, 2), dtype=torch.bool) # (b=1, l=2) + pae = compute_predicted_aligned_error(logits, mask) # (b=1, l=2, l=2) + assert torch.equal( + pae, + torch.full((1, 2, 2), 31.0, dtype=torch.float64), # (b=1, l=2, l=2) + ) + + centers = torch.tensor([7.75, 23.25, 38.75, 54.25], dtype=torch.float64) # (c=4,) + d0 = 1.24 * 4 ** (1 / 3) - 1.8 + expected_tm = (1 / (1 + (centers / d0) ** 2)).mean().reshape(1) # (b=1,) + torch.testing.assert_close(compute_tm(logits, mask), expected_tm, rtol=2e-6, atol=1e-12) + + +def test_aligner_recovers_a_rigid_translation() -> None: + mobile_positions = np.full((1, 37, 3), np.nan, dtype=np.float32) # (l=1, a=37, xyz=3) + mobile_positions[0, :3] = np.asarray( # (a_selected=3, xyz=3) + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=np.float32, + ) + mask = np.zeros((1, 37), dtype=bool) # (l=1, a=37) + mask[0, :3] = True # (a_selected=3,) + target_positions = mobile_positions.copy() # (l=1, a=37, xyz=3) + target_positions[mask] += np.asarray( # (xyz=3,), broadcast over n_selected=3 + [2.0, -1.0, 3.0], dtype=np.float32 + ) + mobile = _Structure(mobile_positions, mask) + target = _Structure(target_positions, mask) + + aligner = Aligner(mobile, target, only_use_backbone=True) + aligned = aligner.apply(mobile) # atom37_positions: (l=1, a=37, xyz=3) + assert aligner.rmsd < 1e-7 + np.testing.assert_allclose(aligned.atom37_positions[mask], target_positions[mask], atol=1e-6) diff --git a/tests/unit/test_esmfold2_runtime_assets.py b/tests/unit/test_esmfold2_runtime_assets.py new file mode 100644 index 0000000..2c19525 --- /dev/null +++ b/tests/unit/test_esmfold2_runtime_assets.py @@ -0,0 +1,225 @@ +"""Security contracts for ESMFold2 runtime assets and tensor payloads.""" + +from __future__ import annotations + +import hashlib +import io +import os +import pickle +import pytest +import torch +import zstandard +from pathlib import Path +from types import SimpleNamespace +from typing import Any, BinaryIO + +from fastplms.models.esmfold2 import esmfold2_conformers as conformers +from fastplms.models.esmfold2.esmfold2_misc import deserialize_tensors + + +def _contract(payload: bytes) -> SimpleNamespace: + return SimpleNamespace( + repository="biohub/ESMFold2", + revision="1ebf0e3481a5184eb6171d40615c79e384b48796", + path="ccd.pkl", + sha256=hashlib.sha256(payload).hexdigest(), + size=len(payload), + trust_kind="hash_pinned_pickle", + ) + + +def _install_contract(monkeypatch: pytest.MonkeyPatch, payload: bytes) -> SimpleNamespace: + contract = _contract(payload) + registry = SimpleNamespace(runtime_assets={"esmfold2_ccd": contract}) + monkeypatch.setattr(conformers, "get_model_registry", lambda: registry) + return contract + + +def test_ccd_local_asset_is_verified_before_pickle_load( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + trusted_payload = pickle.dumps({"ALA": "fixture"}) + _install_contract(monkeypatch, trusted_payload) + (tmp_path / "ccd.pkl").write_bytes(b"not-the-approved-pickle") + + def fail_if_loaded(_handle: object) -> object: + raise AssertionError("pickle.load must not run before identity verification") + + monkeypatch.setattr(conformers.pickle, "load", fail_if_loaded) + store = conformers._ChemicalComponentStore() + with pytest.raises(ValueError, match=r"size mismatch|SHA256 mismatch"): + store.load(tmp_path) + + +def test_ccd_hub_download_uses_manifest_revision( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + payload = pickle.dumps({}) + contract = _install_contract(monkeypatch, payload) + asset = tmp_path / "ccd.pkl" + asset.write_bytes(payload) + monkeypatch.delenv("ESMCFOLD_CCD_PATH", raising=False) + observed: dict[str, str] = {} + + def fake_download(**kwargs: str) -> str: + observed.update(kwargs) + return str(asset) + + monkeypatch.setattr(conformers, "hf_hub_download", fake_download) + resolved = conformers._ChemicalComponentStore()._resolve_asset(None) + + assert resolved == asset + assert observed == { + "repo_id": contract.repository, + "filename": contract.path, + "revision": contract.revision, + } + + +def test_ccd_verified_pickle_loads(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + payload = pickle.dumps({"ALA": "fixture"}) + _install_contract(monkeypatch, payload) + (tmp_path / "ccd.pkl").write_bytes(payload) + + assert conformers._ChemicalComponentStore().load(tmp_path) == {"ALA": "fixture"} + + +def test_ccd_path_replacement_after_hashing_cannot_change_loaded_bytes( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + trusted_payload = pickle.dumps({"ALA": "trusted"}) + replacement_payload = pickle.dumps({"ALA": "replacement"}) + _install_contract(monkeypatch, trusted_payload) + asset = tmp_path / "ccd.pkl" + asset.write_bytes(trusted_payload) + replacement = tmp_path / "replacement.pkl" + replacement.write_bytes(replacement_payload) + real_file_digest = conformers.file_digest + + def replace_path_after_hash(handle: BinaryIO, algorithm: str) -> Any: + digest = real_file_digest(handle, algorithm) + os.replace(replacement, asset) + return digest + + monkeypatch.setattr(conformers, "file_digest", replace_path_after_hash) + loaded = conformers._ChemicalComponentStore().load(tmp_path) + + assert loaded == {"ALA": "trusted"} + assert pickle.loads(asset.read_bytes()) == {"ALA": "replacement"} + + +def test_ccd_in_place_mutation_after_hashing_cannot_change_loaded_bytes( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + trusted_payload = pickle.dumps({"ALA": "trusted"}) + replacement_payload = pickle.dumps({"ALA": "replacement"}) + _install_contract(monkeypatch, trusted_payload) + asset = tmp_path / "ccd.pkl" + asset.write_bytes(trusted_payload) + real_file_digest = conformers.file_digest + + def mutate_source_after_hash(handle: BinaryIO, algorithm: str) -> Any: + digest = real_file_digest(handle, algorithm) + asset.write_bytes(replacement_payload) + return digest + + monkeypatch.setattr(conformers, "file_digest", mutate_source_after_hash) + loaded = conformers._ChemicalComponentStore().load(tmp_path) + + assert loaded == {"ALA": "trusted"} + assert pickle.loads(asset.read_bytes()) == {"ALA": "replacement"} + + +def test_ccd_loader_allows_manifest_owned_hub_snapshot_symlink( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + payload = pickle.dumps({"ALA": "trusted"}) + contract = _install_contract(monkeypatch, payload) + hub_root = tmp_path / "hub" + repository_cache = hub_root / "models--biohub--ESMFold2" + blob = repository_cache / "blobs" / contract.sha256 + blob.parent.mkdir(parents=True) + blob.write_bytes(payload) + snapshot = repository_cache / "snapshots" / contract.revision / contract.path + snapshot.parent.mkdir(parents=True) + snapshot.symlink_to(blob) + monkeypatch.delenv("ESMCFOLD_CCD_PATH", raising=False) + monkeypatch.setattr(conformers, "HF_HUB_CACHE", str(hub_root)) + monkeypatch.setattr(conformers, "hf_hub_download", lambda **_kwargs: str(snapshot)) + + assert conformers._ChemicalComponentStore().load() == {"ALA": "trusted"} + + +def test_ccd_loader_rejects_hub_snapshot_link_outside_repo_blob_cache( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + payload = pickle.dumps({"ALA": "trusted"}) + contract = _install_contract(monkeypatch, payload) + hub_root = tmp_path / "hub" + repository_cache = hub_root / "models--biohub--ESMFold2" + (repository_cache / "blobs").mkdir(parents=True) + outside = tmp_path / "outside.pkl" + outside.write_bytes(payload) + snapshot = repository_cache / "snapshots" / contract.revision / contract.path + snapshot.parent.mkdir(parents=True) + snapshot.symlink_to(outside) + monkeypatch.delenv("ESMCFOLD_CCD_PATH", raising=False) + monkeypatch.setattr(conformers, "HF_HUB_CACHE", str(hub_root)) + monkeypatch.setattr(conformers, "hf_hub_download", lambda **_kwargs: str(snapshot)) + + with pytest.raises(ValueError, match="escapes its repository blob cache"): + conformers._ChemicalComponentStore().load() + + +def test_ccd_loader_rejects_configured_cache_symlinks( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + payload = pickle.dumps({"ALA": "fixture"}) + _install_contract(monkeypatch, payload) + target = tmp_path / "trusted.pkl" + target.write_bytes(payload) + (tmp_path / "ccd.pkl").symlink_to(target) + + with pytest.raises(ValueError, match="must not be a symlink"): + conformers._ChemicalComponentStore().load(tmp_path) + + +def test_ccd_loader_rejects_non_regular_files( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + payload = pickle.dumps({"ALA": "fixture"}) + _install_contract(monkeypatch, payload) + (tmp_path / "ccd.pkl").mkdir() + + with pytest.raises(ValueError, match="must be a regular file"): + conformers._ChemicalComponentStore().load(tmp_path) + + +def test_tensor_deserialization_rejects_arbitrary_pickle_globals() -> None: + class UnsafePayload: + def __reduce__(self) -> tuple[object, tuple[str]]: + return eval, ("40 + 2",) + + buffer = io.BytesIO() + torch.save(UnsafePayload(), buffer) + compressed = zstandard.ZstdCompressor().compress(buffer.getvalue()) + + with pytest.raises((pickle.UnpicklingError, RuntimeError)): + deserialize_tensors(compressed) + + +def test_tensor_deserialization_accepts_tensor_mappings() -> None: + buffer = io.BytesIO() + expected = {"X": torch.arange(6).reshape(2, 3)} # (n=2, d=3) + torch.save(expected, buffer) + compressed = zstandard.ZstdCompressor().compress(buffer.getvalue()) + + actual = deserialize_tensors(compressed) # actual["X"]: (n=2, d=3) + assert actual.keys() == expected.keys() + assert torch.equal(actual["X"], expected["X"]) diff --git a/tests/unit/test_esmfold_api.py b/tests/unit/test_esmfold_api.py new file mode 100644 index 0000000..1d2359a --- /dev/null +++ b/tests/unit/test_esmfold_api.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import inspect +import pytest +import torch +from types import MethodType +from transformers.models.esm.modeling_esmfold import EsmForProteinFolding + +from fastplms.models.esmfold.modeling_fast_esmfold import FastEsmForProteinFolding + + +def test_forward_uses_official_plddt_scale(monkeypatch: pytest.MonkeyPatch) -> None: + def forward(*args: object, **kwargs: object) -> dict[str, torch.Tensor]: + return {"plddt": torch.tensor([0.75])} # (b=1,) + + monkeypatch.setattr(EsmForProteinFolding, "forward", forward) + model = FastEsmForProteinFolding.__new__(FastEsmForProteinFolding) + torch.nn.Module.__init__(model) + + output = model.forward( # output["plddt"]: (b=1,) + torch.zeros(1, 1, dtype=torch.int64) # (b=1, l=1) + ) + + assert output["plddt"].equal(torch.tensor([75.0])) + + +def test_infer_preserves_official_multimer_contract() -> None: + assert list(inspect.signature(FastEsmForProteinFolding.infer).parameters) == [ + "self", + "sequences", + "residx", + "masking_pattern", + "num_recycles", + "residue_index_offset", + "chain_linker", + ] + model = FastEsmForProteinFolding.__new__(FastEsmForProteinFolding) + torch.nn.Module.__init__(model) + model.register_parameter( + "device_anchor", + torch.nn.Parameter(torch.empty(0), requires_grad=False), + ) + observed: dict[str, object] = {} + + def forward( + self: FastEsmForProteinFolding, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: object, + ) -> dict[str, torch.Tensor]: + # input_ids: (b, l); attention_mask: (b, l) + observed.update( + input_ids=input_ids, + attention_mask=attention_mask, + **kwargs, + ) + atom_mask = torch.ones(input_ids.shape[0], input_ids.shape[1], 37) # (b, l, a=37) + return { + "aatype": input_ids, # (b, l) + "atom37_atom_exists": atom_mask, # (b, l, a=37) + "plddt": torch.full_like(atom_mask, 0.75), # (b, l, a=37) + } + + model.forward = MethodType(forward, model) + output = model.infer( # per-residue outputs use (b=1, l=6, ...) + "AC:DE", + num_recycles=2, + residue_index_offset=32, + chain_linker="GG", + ) + + assert output["aatype"].shape == (1, 6) + assert observed["attention_mask"].equal( + torch.ones(1, 6, dtype=torch.int64) # (b=1, l=6) + ) + assert observed["position_ids"].equal( + torch.tensor([[0, 1, 2, 3, 36, 37]], dtype=torch.int64) # (b=1, l=6) + ) + assert observed["masking_pattern"] is None + assert observed["num_recycles"] == 2 + assert output["chain_index"].equal( + torch.tensor([[0, 0, 0, 0, 1, 1]], dtype=torch.int64) # (b=1, l=6) + ) + assert output["atom37_atom_exists"][0, :, 0].equal( + torch.tensor([1, 1, 0, 0, 1, 1], dtype=torch.float32) # (l=6,) + ) + assert output["mean_plddt"].equal(torch.tensor([0.75])) # (b=1,) diff --git a/tests/unit/test_esmplusplus_masking.py b/tests/unit/test_esmplusplus_masking.py new file mode 100644 index 0000000..4f55c6d --- /dev/null +++ b/tests/unit/test_esmplusplus_masking.py @@ -0,0 +1,292 @@ +"""Focused ESM++ input-mask contracts.""" + +from __future__ import annotations + +import ast +import inspect +import textwrap +import pytest +import torch +from pathlib import Path + +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ESMplusplusForSequenceClassification, + ESMplusplusForTokenClassification, + ESMplusplusModel, + TransformerStack, +) + + +@pytest.mark.parametrize("model_class", (ESMplusplusModel, ESMplusplusForMaskedLM)) +def test_esmplusplus_infers_padding_mask_from_input_ids( + model_class: type[ESMplusplusModel] | type[ESMplusplusForMaskedLM], +) -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="eager", + pad_token_id=1, + ) + model = model_class(config).eval() + input_ids = torch.tensor([[0, 3, 4, 1, 1]], dtype=torch.long) # (b=1, l=5) + attention_mask = input_ids.ne(config.pad_token_id) # (b=1, l=5) + + kwargs = {"compute_logits": False} if model_class is ESMplusplusForMaskedLM else {} + with torch.inference_mode(): + inferred = model(input_ids=input_ids, **kwargs).last_hidden_state # (b=1, l=5, d=16) + explicit = model( + input_ids=input_ids, + attention_mask=attention_mask, + **kwargs, + ).last_hidden_state # (b=1, l=5, d=16) + + torch.testing.assert_close(inferred, explicit, rtol=0.0, atol=0.0) + + +def test_esmplusplus_boolean_sequence_id_matches_biohub_equality_mask() -> None: + stack = TransformerStack( + d_model=16, + n_heads=4, + n_layers=1, + attn_backend="eager", + ) + sequence_id = torch.tensor([[True, True, True, False, False]]) # (b=1, l=5) + + mask_2d, mask_4d, block_mask = stack._sequence_id_attention_masks( + sequence_id=sequence_id, + batch_size=1, + seq_len=5, + device=torch.device("cpu"), + ) # mask_2d: (b=1, l=5); mask_4d: (b=1, 1, l=5, l=5); block_mask: None + + expected = sequence_id[:, None, :, None] == sequence_id[:, None, None, :] # (1, 1, 5, 5) + assert torch.equal(mask_2d, sequence_id) + assert torch.equal(mask_4d, expected) + assert block_mask is None + + +@pytest.mark.parametrize("model_class", (ESMplusplusModel, ESMplusplusForMaskedLM)) +def test_esmplusplus_embedding_helper_infers_padding_mask( + model_class: type[ESMplusplusModel] | type[ESMplusplusForMaskedLM], +) -> None: + config = ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="eager", + pad_token_id=1, + ) + model = model_class(config).eval() + input_ids = torch.tensor([[0, 3, 4, 1, 1]], dtype=torch.long) # (b=1, l=5) + attention_mask = input_ids.ne(config.pad_token_id) # (b=1, l=5) + + with torch.inference_mode(): + inferred = model._embed(input_ids) # (b=1, l=5, d=16) + explicit = model._embed(input_ids, attention_mask=attention_mask) # (1, 5, 16) + + torch.testing.assert_close(inferred, explicit, rtol=0.0, atol=0.0) + + +def _sequence_classifier_config() -> ESMplusplusConfig: + return ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + num_labels=3, + attn_backend="eager", + pad_token_id=1, + ) + + +def test_esmplusplus_sequence_classifier_accepts_explicit_pooling_types() -> None: + model = ESMplusplusForSequenceClassification( + _sequence_classifier_config(), + pooling_types=["mean"], + ).eval() + + assert model.pooler.names == ("mean",) + assert model.classifier[0].in_features == 16 + + with torch.inference_mode(): + output = model( + input_ids=torch.tensor([[0, 3, 4, 2]], dtype=torch.long) # (b=1, l=4) + ) + + assert output.logits is not None + # b=1 sequences, c=3 classes. + assert output.logits.shape == (1, 3) + + +def test_esmplusplus_sequence_classifier_pooling_round_trips(tmp_path: Path) -> None: + model = ESMplusplusForSequenceClassification( + _sequence_classifier_config(), + pooling_types=["mean"], + ).eval() + model.save_pretrained(tmp_path) + + reloaded = ESMplusplusForSequenceClassification.from_pretrained(tmp_path).eval() + + assert reloaded.config.classifier_pooling_types == ["mean"] + assert reloaded.pooler.names == ("mean",) + assert reloaded.classifier[0].in_features == 16 + for name, value in model.state_dict().items(): + torch.testing.assert_close(reloaded.state_dict()[name], value, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize( + ("model_class", "pooling_types", "labels", "expected_shape"), + ( + ( + ESMplusplusForSequenceClassification, + ["mean", "var"], + torch.tensor([1, 2]), # (b=2,) + (2, 3), + ), + ( + ESMplusplusForTokenClassification, + None, + torch.tensor(((0, 1, 2, 1), (2, 1, 0, 1))), # (b=2, l=4) + (2, 4, 3), + ), + ), +) +def test_esmplusplus_wide_classifier_forward_backward_and_reload( + model_class: ( + type[ESMplusplusForSequenceClassification] + | type[ESMplusplusForTokenClassification] + ), + pooling_types: list[str] | None, + labels: torch.Tensor, + expected_shape: tuple[int, ...], + tmp_path: Path, +) -> None: + config = _sequence_classifier_config() + kwargs = {} if pooling_types is None else {"pooling_types": pooling_types} + model = model_class(config, **kwargs).train() + assert model.classifier[0].out_features == config.hidden_size * 4 + assert model.classifier[3].in_features == config.hidden_size * 4 + input_ids = torch.tensor(((0, 3, 4, 2), (0, 5, 6, 2))) # (b=2, l=4) + + output = model(input_ids=input_ids, labels=labels) + assert output.logits.shape == expected_shape + assert output.loss is not None + assert torch.isfinite(output.loss) + output.loss.backward() + classifier_gradients = [ + parameter.grad for parameter in model.classifier.parameters() if parameter.requires_grad + ] + assert classifier_gradients + assert all(gradient is not None for gradient in classifier_gradients) + assert all(torch.isfinite(gradient).all() for gradient in classifier_gradients) + + model.eval() + with torch.inference_mode(): + expected_logits = model(input_ids=input_ids).logits # (b, c) or (b, l, c) + save_path = tmp_path / model_class.__name__ + model.save_pretrained(save_path, safe_serialization=True) + reloaded = model_class.from_pretrained(save_path, local_files_only=True).eval() + with torch.inference_mode(): + actual_logits = reloaded(input_ids=input_ids).logits # same shape as expected_logits + torch.testing.assert_close(actual_logits, expected_logits, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize( + ("pooling_types", "exception"), + [ + ([], ValueError), + (("mean",), TypeError), + (["mean", 1], TypeError), + (["parti"], ValueError), + ], +) +def test_esmplusplus_sequence_classifier_validates_pooling_types( + pooling_types: object, + exception: type[Exception], +) -> None: + with pytest.raises(exception, match="pooling_types"): + ESMplusplusForSequenceClassification( + _sequence_classifier_config(), + pooling_types=pooling_types, + ) + + +def test_esmplusplus_sequence_classifier_is_right_padding_invariant_without_mask() -> None: + model = ESMplusplusForSequenceClassification(_sequence_classifier_config()).eval() + unpadded = torch.tensor([[0, 3, 4, 2]], dtype=torch.long) # (b=1, l=4) + right_padded = torch.tensor([[0, 3, 4, 2, 1, 1]], dtype=torch.long) # (b=1, l=6) + + with torch.inference_mode(): + unpadded_logits = model(input_ids=unpadded).logits # (b=1, c=3) + padded_logits = model(input_ids=right_padded).logits # (b=1, c=3) + + assert unpadded_logits is not None + assert padded_logits is not None + torch.testing.assert_close(padded_logits, unpadded_logits, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + "model_class", + (ESMplusplusModel, ESMplusplusForMaskedLM, ESMplusplusForSequenceClassification), +) +def test_esmplusplus_public_models_require_exactly_one_input_form( + model_class: ( + type[ESMplusplusModel] + | type[ESMplusplusForMaskedLM] + | type[ESMplusplusForSequenceClassification] + ), +) -> None: + model = model_class( + ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="eager", + ) + ) + input_ids = torch.ones(1, 2, dtype=torch.long) # (b=1, l=2) + inputs_embeds = torch.zeros(1, 2, 16) # (b=1, l=2, d=16) + + with pytest.raises(ValueError, match="either input_ids or inputs_embeds"): + model() + with pytest.raises(ValueError, match="both input_ids and inputs_embeds"): + model(input_ids=input_ids, inputs_embeds=inputs_embeds) + + +def test_esmplusplus_masked_lm_labels_require_logits() -> None: + model = ESMplusplusForMaskedLM( + ESMplusplusConfig( + vocab_size=16, + hidden_size=16, + num_attention_heads=4, + num_hidden_layers=1, + attn_backend="eager", + ) + ) + input_ids = torch.ones(1, 2, dtype=torch.long) # (b=1, l=2) + + with pytest.raises(ValueError, match="labels require compute_logits=True"): + model(input_ids=input_ids, labels=input_ids, compute_logits=False) + + +@pytest.mark.parametrize( + "model_class", + (ESMplusplusModel, ESMplusplusForMaskedLM, ESMplusplusForSequenceClassification), +) +def test_esmplusplus_public_input_validation_survives_python_optimization( + model_class: ( + type[ESMplusplusModel] + | type[ESMplusplusForMaskedLM] + | type[ESMplusplusForSequenceClassification] + ), +) -> None: + forward_source = textwrap.dedent(inspect.getsource(model_class.forward)) + forward_tree = ast.parse(forward_source) + + assert not any(isinstance(node, ast.Assert) for node in ast.walk(forward_tree)) diff --git a/tests/unit/test_fine_tuning_example.py b/tests/unit/test_fine_tuning_example.py new file mode 100644 index 0000000..83ee244 --- /dev/null +++ b/tests/unit/test_fine_tuning_example.py @@ -0,0 +1,722 @@ +"""Lightweight source-contract tests for the optional fine-tuning example.""" + +from __future__ import annotations + +import ast +import hashlib +import json +import os +import platform +import re +import shutil +import sys +import tempfile +import numpy as np +import pytest +from collections.abc import Mapping +from importlib import metadata +from pathlib import Path +from types import SimpleNamespace +from typing import Any, ClassVar + + +ROOT = Path(__file__).resolve().parents[2] +EXAMPLE = ROOT / "examples" / "fine_tuning.py" + + +def _tree() -> ast.Module: + return ast.parse(EXAMPLE.read_text(encoding="utf-8"), filename=str(EXAMPLE)) + + +def _assignment_value(tree: ast.Module, name: str) -> Any: + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name for target in node.targets + ): + return ast.literal_eval(node.value) + raise AssertionError(f"Missing assignment for {name}.") + + +def _function(tree: ast.Module, name: str) -> ast.FunctionDef: + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"Missing function {name}.") + + +def _definition(tree: ast.Module, name: str) -> ast.FunctionDef | ast.ClassDef: + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name == name: + return node + raise AssertionError(f"Missing definition {name}.") + + +def test_lora_configuration_persists_classifier_and_preserves_custom_modules() -> None: + tree = _tree() + classifier_name = _assignment_value(tree, "CLASSIFIER_MODULE_NAME") + helper = _function(tree, "_ensure_classifier_persistence") + namespace = { + "Any": Any, + "CLASSIFIER_MODULE_NAME": classifier_name, + } + exec(compile(ast.Module(body=[helper], type_ignores=[]), str(EXAMPLE), "exec"), namespace) + ensure_persistence = namespace["_ensure_classifier_persistence"] + + config = SimpleNamespace(modules_to_save=["contact_head"]) + assert ensure_persistence(config) is config + assert config.modules_to_save == ["contact_head", "classifier"] + + ensure_persistence(config) + assert config.modules_to_save == ["contact_head", "classifier"] + + empty_config = SimpleNamespace(modules_to_save=None) + ensure_persistence(empty_config) + assert empty_config.modules_to_save == ["classifier"] + + +def test_defaults_only_advertise_sequence_classification_artifacts() -> None: + tree = _tree() + default_model = _assignment_value(tree, "DEFAULT_MODEL") + assert default_model == "Synthyra/ESM2-8M" + + for function_name in ("train_regression_model", "train_classification_model"): + function = _function(tree, function_name) + model_default = function.args.defaults[0] + assert isinstance(model_default, ast.Name) + assert model_default.id == "DEFAULT_MODEL" + + source = EXAMPLE.read_text(encoding="utf-8") + assert "Synthyra/ESMplusplus_small" not in source + assert "Synthyra/ESMplusplus_large" not in source + + +def test_cli_exposes_symmetric_lora_switch() -> None: + tree = _tree() + lora_argument: ast.Call | None = None + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + continue + if node.func.attr != "add_argument": + continue + options = [ + ast.literal_eval(argument) + for argument in node.args + if isinstance(argument, ast.Constant) and isinstance(argument.value, str) + ] + if "--use-lora" in options: + lora_argument = node + assert {"--use-lora", "--use_lora"}.issubset(options) + break + + assert lora_argument is not None + keywords = {keyword.arg: keyword.value for keyword in lora_argument.keywords} + action = keywords["action"] + assert isinstance(action, ast.Attribute) + assert isinstance(action.value, ast.Name) + assert (action.value.id, action.attr) == ("argparse", "BooleanOptionalAction") + assert ast.literal_eval(keywords["default"]) is True + + +def test_reporting_is_opt_in_for_the_minimal_training_install() -> None: + tree = _tree() + plot_argument: ast.Call | None = None + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + continue + if node.func.attr != "add_argument": + continue + options = [ + ast.literal_eval(argument) + for argument in node.args + if isinstance(argument, ast.Constant) and isinstance(argument.value, str) + ] + if "--plot-results" in options: + plot_argument = node + break + + assert plot_argument is not None + keywords = {keyword.arg: keyword.value for keyword in plot_argument.keywords} + assert ast.literal_eval(keywords["default"]) is False + for function_name in ("train_regression_model", "train_classification_model"): + function = _function(tree, function_name) + defaults = dict( + zip( + (argument.arg for argument in function.args.args[-len(function.args.defaults) :]), + function.args.defaults, + strict=True, + ) + ) + assert ast.literal_eval(defaults["plot_results"]) is False + + +def test_max_length_contract_is_an_encoded_budget_including_added_tokens() -> None: + tree = _tree() + for function_name in ( + "train_regression_model", + "train_classification_model", + "PairDatasetHF", + "SequenceDatasetHF", + ): + definition = _definition(tree, function_name) + docstring = ast.get_docstring(definition) or "" + assert "Encoded token budget" in docstring + assert "special tokens" in docstring + + max_length_argument: ast.Call | None = None + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + continue + if node.func.attr != "add_argument": + continue + options = [ + ast.literal_eval(argument) + for argument in node.args + if isinstance(argument, ast.Constant) and isinstance(argument.value, str) + ] + if "--max_length" in options: + max_length_argument = node + break + assert max_length_argument is not None + keywords = {keyword.arg: keyword.value for keyword in max_length_argument.keywords} + help_text = ast.literal_eval(keywords["help"]) + assert "encoded token count" in help_text + assert "special and pair separator tokens" in help_text + + +def test_plot_contract_uses_task_output_paths_and_has_no_interactive_or_overwrite_path() -> None: + source = EXAMPLE.read_text(encoding="utf-8") + assert 'Path(output_dir) / "regression_results.png"' in source + assert 'Path(output_dir) / "classification_results.png"' in source + assert "plt.show(" not in source + + for function_name in ("plot_regression_results", "plot_classification_results"): + function = _function(_tree(), function_name) + assert "output_path" in [argument.arg for argument in function.args.args] + called_names = { + node.func.id + for node in ast.walk(function) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "_save_figure_exclusive" in called_names + + +def test_pair_token_budget_includes_special_tokens_at_the_exact_boundary() -> None: + torch = pytest.importorskip("torch") + tree = _tree() + namespace = {"Any": Any, "torch": torch} + nodes = [ + _definition(tree, "_encoded_length"), + _definition(tree, "_fits_token_budget"), + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + + class FakeTokenizer: + def __call__( + self, + sequence: str, + pair: str | None, + **kwargs: Any, + ) -> dict[str, list[int]]: + assert kwargs["add_special_tokens"] is True + assert kwargs["truncation"] is False + special_tokens = 3 if pair is not None else 2 + return {"input_ids": [0] * (len(sequence) + len(pair or "") + special_tokens)} + + fits = namespace["_fits_token_budget"] + assert fits(FakeTokenizer(), "AC", "DEF", 8) + assert not fits(FakeTokenizer(), "AC", "DEF", 7) + + +def test_pair_collator_enforces_longest_first_tokenizer_limit() -> None: + torch = pytest.importorskip("torch") + tree = _tree() + namespace = {"Any": Any, "torch": torch} + nodes = [ + _definition(tree, "_tokenization_kwargs"), + _definition(tree, "PairCollator"), + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + + class FakeTokenizer: + def __init__(self) -> None: + self.kwargs: dict[str, Any] | None = None + + def __call__(self, seqs_a: Any, seqs_b: Any, **kwargs: Any) -> dict[str, Any]: + del seqs_a, seqs_b + self.kwargs = kwargs + return { + "input_ids": torch.zeros( # (b=2, l=max_length) + (2, kwargs["max_length"]), + dtype=torch.long, + ), + "attention_mask": torch.ones((2, kwargs["max_length"]), dtype=torch.long), + } + + tokenizer = FakeTokenizer() + collator = namespace["PairCollator"](tokenizer, regression=True, max_length=16) + batch = collator([("AC", "DE", 1.0), ("FG", "HI", 2.0)]) + + assert tokenizer.kwargs is not None + assert tokenizer.kwargs["truncation"] == "longest_first" + assert tokenizer.kwargs["max_length"] == 16 + assert tokenizer.kwargs["pad_to_multiple_of"] == 8 + assert batch["input_ids"].shape == (2, 16) + assert batch["labels"].dtype == torch.float32 + + +def test_training_manifest_records_reproducible_model_data_and_tokenizer_identity( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + torch = pytest.importorskip("torch") + tree = _tree() + namespace = { + "Any": Any, + "Mapping": Mapping, + "Path": Path, + "TrainingArguments": Any, + "hashlib": hashlib, + "json": json, + "metadata": metadata, + "platform": platform, + "_IMMUTABLE_REVISION": re.compile(r"^[0-9a-f]{40}$", re.IGNORECASE), + "_PINNED_DEFAULT_REVISIONS": {}, + "sys": sys, + "torch": torch, + } + nodes = [ + _definition(tree, name) + for name in ( + "_package_version", + "_json_safe", + "_sha256_json", + "_tokenizer_identity", + "_tree_sha256", + "_immutable_source_identity", + "_verify_training_source_unchanged", + "_effective_attention_backend", + "_ordered_rows_sha256", + "_dataset_identity", + "_write_training_manifest", + ) + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + + class FakeConfig: + _name_or_path = "Synthyra/tiny" + _commit_hash = "a" * 40 + attn_backend = "sdpa" + + def to_dict(self) -> dict[str, Any]: + return {"hidden_size": 8, "attn_backend": self.attn_backend} + + class FakeAdapterConfig: + def to_dict(self) -> dict[str, Any]: + return {"r": 2, "target_modules": {"query", "value"}} + + class FakeModel: + config = FakeConfig() + + def __init__(self) -> None: + self.peft_config = {"default": FakeAdapterConfig()} + self.parameter = torch.nn.Parameter(torch.ones(1)) + self._fastplms_training_source_identity = { + "kind": "hub", + "source_kind": "model", + "repo_id": "Synthyra/tiny", + "revision": "a" * 40, + } + + def parameters(self) -> Any: + yield self.parameter + + class FakeTokenizer: + name_or_path = "Synthyra/tiny" + cls_token_id = 0 + eos_token_id = 2 + pad_token_id = 1 + + def __init__(self) -> None: + self.init_kwargs = {"revision": "a" * 40} + + def get_vocab(self) -> dict[str, int]: + return {"": 0, "": 1, "": 2, "A": 3} + + class FakeDataset: + _fingerprint = "dataset-fingerprint" + info = SimpleNamespace(builder_name="builder", config_name="config", version="1.0") + rows: ClassVar[list[dict[str, object]]] = [ + {"sequence": "AC", "label": 0}, + {"sequence": "DEF", "label": 1}, + {"sequence": "GHIK", "label": 0}, + ] + + def __len__(self) -> int: + return len(self.rows) + + def __iter__(self) -> Any: + return iter(self.rows) + + class FakeTrainingArguments: + device = torch.device("cpu") + bf16 = False + fp16 = False + optim = "adamw_torch" + lr_scheduler_type = "linear" + warmup_steps = 5 + weight_decay = 0.01 + eval_strategy = "steps" + eval_steps = 4 + save_strategy = "steps" + save_steps = 4 + logging_strategy = "steps" + logging_steps = 2 + load_best_model_at_end = True + metric_for_best_model = "eval_loss" + greater_is_better = False + label_names: ClassVar[list[str]] = ["labels"] + report_to: ClassVar[list[str]] = [] + + monkeypatch.setattr(sys, "argv", ["fine_tuning.py", "--task", "regression"]) + namespace["_write_training_manifest"]( + str(tmp_path), + task="test", + model=FakeModel(), + tokenizer=FakeTokenizer(), + model_name="Synthyra/tiny", + model_revision="a" * 40, + seed=7, + max_length=16, + use_lora=True, + batch_size=2, + gradient_accumulation_steps=3, + learning_rate=1e-4, + num_epochs=1, + full_determinism=True, + datasets={"train": FakeDataset()}, + dataset_contracts={ + "train": { + "source": { + "kind": "hub", + "source_kind": "dataset", + "repo_id": "Synthyra/tiny-data", + "revision": "b" * 40, + }, + "split": "train", + "columns": ("sequence", "label"), + } + }, + training_arguments=FakeTrainingArguments(), + patience=2, + final_artifact={"reload_verified": True, "tree_sha256": "c" * 64}, + ) + + manifest = json.loads((tmp_path / "run_manifest.json").read_text(encoding="utf-8")) + assert manifest["command"] == ["fine_tuning.py", "--task", "regression"] + assert manifest["model"]["revision"] == "a" * 40 + assert manifest["model"]["attention_backend"] == "sdpa" + assert manifest["model"]["requested_attention_backend"] == "sdpa" + assert manifest["model"]["effective_attention_backend"] == "sdpa" + assert manifest["model"]["parameter_dtype"] == "torch.float32" + assert len(manifest["model"]["configuration_sha256"]) == 64 + assert manifest["model"]["adapters"]["default"]["target_modules"] == [ + "query", + "value", + ] + assert manifest["tokenizer"]["revision"] == "a" * 40 + assert len(manifest["tokenizer"]["vocab_sha256"]) == 64 + assert manifest["datasets"]["train"]["library_fingerprint_advisory"] == "dataset-fingerprint" + assert len(manifest["datasets"]["train"]["ordered_rows_sha256"]) == 64 + assert manifest["datasets"]["train"]["columns"] == ["sequence", "label"] + assert manifest["datasets"]["train"]["source"]["revision"] == "b" * 40 + assert manifest["datasets"]["train"]["rows"] == 3 + assert manifest["training"]["compute_dtype"] == "torch.float32" + assert manifest["training"]["optimizer"] == "adamw_torch" + assert manifest["training"]["early_stopping_patience"] == 2 + assert manifest["training"]["eval_steps"] == 4 + assert manifest["training"]["save_steps"] == 4 + assert manifest["training"]["logging_steps"] == 2 + assert manifest["training"]["max_length_semantics"] == ( + "encoded token budget including tokenizer-added special and pair tokens" + ) + assert manifest["final_artifact"]["reload_verified"] is True + + +def test_immutable_sources_reject_moving_refs_pin_only_shipped_defaults_and_detect_drift( + tmp_path: Path, +) -> None: + tree = _tree() + default_model = _assignment_value(tree, "DEFAULT_MODEL") + default_revision = _assignment_value(tree, "DEFAULT_MODEL_REVISION") + namespace = { + "Any": Any, + "Mapping": Mapping, + "Path": Path, + "hashlib": hashlib, + "_IMMUTABLE_REVISION": re.compile(r"^[0-9a-f]{40}$", re.IGNORECASE), + "_PINNED_DEFAULT_REVISIONS": {("model", default_model): default_revision}, + } + nodes = [ + _definition(tree, name) + for name in ( + "_tree_sha256", + "_immutable_source_identity", + "_verify_training_source_unchanged", + ) + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + identity = namespace["_immutable_source_identity"] + + pinned = identity(default_model, None, source_kind="model") + assert pinned["revision"] == default_revision + with pytest.raises(ValueError, match="40-character"): + identity("Custom/model", None, source_kind="model") + with pytest.raises(ValueError, match="branches, tags"): + identity("Custom/model", "main", source_kind="model") + + local_model = tmp_path / "local-model" + local_model.mkdir() + (local_model / "config.json").write_text('{"hidden_size": 8}\n', encoding="utf-8") + loaded_identity = identity(str(local_model), None, source_kind="model") + model = SimpleNamespace(_fastplms_training_source_identity=loaded_identity) + assert ( + namespace["_verify_training_source_unchanged"]( + model, + str(local_model), + None, + ) + == loaded_identity + ) + (local_model / "config.json").write_text('{"hidden_size": 16}\n', encoding="utf-8") + with pytest.raises(RuntimeError, match="changed after initialization"): + namespace["_verify_training_source_unchanged"]( + model, + str(local_model), + None, + ) + + +def test_ordered_training_row_hash_is_content_and_order_sensitive() -> None: + tree = _tree() + namespace = {"Any": Any, "Mapping": Mapping, "hashlib": hashlib, "json": json} + nodes = [ + _definition(tree, name) for name in ("_json_safe", "_sha256_json", "_ordered_rows_sha256") + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + hash_rows = namespace["_ordered_rows_sha256"] + rows = [ + {"sequence": "AC", "label": 0, "ignored": "x"}, + {"sequence": "DEF", "label": 1, "ignored": "y"}, + ] + columns = ("sequence", "label") + + baseline = hash_rows(rows, columns) + assert hash_rows([dict(row) for row in rows], columns) == baseline + assert hash_rows(list(reversed(rows)), columns) != baseline + changed = [dict(row) for row in rows] + changed[1]["sequence"] = "DEG" + assert hash_rows(changed, columns) != baseline + ignored_change = [dict(row) for row in rows] + ignored_change[0]["ignored"] = "changed" + assert hash_rows(ignored_change, columns) == baseline + + +def test_persisted_hash_scope_covers_full_state_or_only_lora_payload() -> None: + torch = pytest.importorskip("torch") + tree = _tree() + namespace = {"Any": Any, "hashlib": hashlib, "torch": torch} + nodes = [_definition(tree, name) for name in ("_tensor_sha256", "_persisted_parameter_hashes")] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + + class FakeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.backbone = torch.nn.Linear(2, 2) + self.classifier = torch.nn.Linear(2, 1) + self.lora_A = torch.nn.Parameter(torch.ones(1, 2)) + self.modules_to_save = torch.nn.Linear(2, 1) + self.register_buffer("running_state", torch.tensor([3.0])) # (1,) + + model = FakeModel() + hashes = namespace["_persisted_parameter_hashes"] + full_hashes = hashes(model, use_lora=False) + lora_hashes = hashes(model, use_lora=True) + assert set(full_hashes) == set(model.state_dict()) + assert "running_state" in full_hashes + assert lora_hashes + assert all("lora_" in name or "modules_to_save" in name for name in lora_hashes) + assert not any("backbone" in name for name in lora_hashes) + assert not any(name.startswith("classifier") for name in lora_hashes) + + +def test_atomic_final_artifact_reload_preserves_trainer_and_held_out_logits( + tmp_path: Path, +) -> None: + torch = pytest.importorskip("torch") + tree = _tree() + namespace = { + "Any": Any, + "Mapping": Mapping, + "Path": Path, + "Trainer": Any, + "hashlib": hashlib, + "json": json, + "np": np, + "os": os, + "shutil": shutil, + "tempfile": tempfile, + "torch": torch, + } + nodes = [ + _definition(tree, name) + for name in ( + "_tree_sha256", + "_tensor_sha256", + "_persisted_parameter_hashes", + "_primary_prediction_tensor", + "_held_out_reload_verification", + "_save_reload_verify_final_artifact", + ) + ] + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(EXAMPLE), "exec"), namespace) + + class TinyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.projection = torch.nn.Linear(2, 1, bias=False) + self._fastplms_training_source_identity = { + "kind": "hub", + "source_kind": "model", + "repo_id": "offline/tiny", + "revision": "a" * 40, + } + + def forward(self, input_ids: Any, labels: Any = None) -> Any: + del labels + return SimpleNamespace(logits=self.projection(input_ids.float())) + + model = TinyModel().eval() + with torch.no_grad(): + model.projection.weight.copy_(torch.tensor([[0.25, -0.5]])) # (c=1, d=2) + + def collate(rows: list[tuple[torch.Tensor, float]]) -> dict[str, torch.Tensor]: + inputs, labels = zip(*rows, strict=True) + return { + "input_ids": torch.stack(inputs), + "labels": torch.tensor(labels), # (b,) + } + + class TinyTrainer: + def __init__(self, trained_model: TinyModel) -> None: + self.model = trained_model + self.model_wrapped = trained_model + self.args = SimpleNamespace(device=torch.device("cpu"), bf16=False, fp16=False) + + def predict(self, rows: Any) -> Any: + batch = collate(rows) + with torch.inference_mode(): + logits = self.model(**batch).logits.numpy() # (b, c) + return SimpleNamespace(predictions=logits) + + def save_model(self, directory: str | Path) -> None: + Path(directory, "model.safetensors").write_bytes(b"safe-test-weights") + + class TinyTokenizer: + def save_pretrained(self, directory: str | Path) -> None: + Path(directory, "tokenizer.json").write_text("{}\n", encoding="utf-8") + + trainer = TinyTrainer(model) + original_model = trainer.model + reloaded = TinyModel() + reloaded.load_state_dict(model.state_dict()) + namespace["_verify_training_source_unchanged"] = ( + lambda current_model, model_name, model_revision: dict( + current_model._fastplms_training_source_identity + ) + ) + namespace["_reload_final_model"] = lambda *args, **kwargs: reloaded + verification_rows = [ + (torch.tensor([1.0, 2.0]), 0.0), + (torch.tensor([3.0, 4.0]), 1.0), + ] + + artifact = namespace["_save_reload_verify_final_artifact"]( + trainer, + TinyTokenizer(), + output_dir=str(tmp_path / "output"), + model_name="offline/tiny", + model_revision="a" * 40, + num_labels=1, + use_lora=False, + verification_dataset=verification_rows, + data_collator=collate, + ) + final_dir = tmp_path / "output" / "final_model" + assert trainer.model is original_model + assert trainer.model_wrapped is original_model + assert final_dir.is_dir() + assert not list((tmp_path / "output").glob(".final-model-*")) + assert artifact["reload_verified"] is True + assert artifact["held_out_inference"]["rows"] == 2 + assert artifact["held_out_inference"]["max_absolute_error"] == 0.0 + metadata_payload = json.loads( + (final_dir / "artifact_metadata.json").read_text(encoding="utf-8") + ) + assert metadata_payload["held_out_inference"] == artifact["held_out_inference"] + + +def test_lora_adapter_save_reload_preserves_trained_classifier_logits( + tmp_path: Path, +) -> None: + torch = pytest.importorskip("torch") + peft = pytest.importorskip("peft") + from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForSequenceClassification, + ) + + torch.manual_seed(7) + config = FastEsmConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + num_labels=2, + pad_token_id=1, + mask_token_id=5, + position_embedding_type="absolute", + attn_backend="eager", + ) + base = FastEsmForSequenceClassification(config) + base_state = {name: tensor.detach().clone() for name, tensor in base.state_dict().items()} + adapter = peft.get_peft_model( + base, + peft.LoraConfig( + task_type=peft.TaskType.SEQ_CLS, + r=2, + lora_alpha=4, + target_modules=["query", "value"], + modules_to_save=["classifier"], + ), + ) + with torch.no_grad(): + for parameter in adapter.base_model.model.classifier.parameters(): + parameter.add_(0.25) + + inputs = { + "input_ids": torch.tensor([[0, 3, 4, 2]], dtype=torch.long), # (b=1, l=4) + "attention_mask": torch.ones((1, 4), dtype=torch.long), + } + adapter.eval() + with torch.inference_mode(): + expected = adapter(**inputs).logits # (b, c) + adapter.save_pretrained(tmp_path) + + restored_base = FastEsmForSequenceClassification(config) + restored_base.load_state_dict(base_state) + restored = peft.PeftModel.from_pretrained(restored_base, tmp_path).eval() + with torch.inference_mode(): + observed = restored(**inputs).logits # (b, c) + + torch.testing.assert_close(observed, expected, rtol=0.0, atol=0.0) diff --git a/tests/unit/test_hopper_hardware_contract.py b/tests/unit/test_hopper_hardware_contract.py new file mode 100644 index 0000000..d1fe881 --- /dev/null +++ b/tests/unit/test_hopper_hardware_contract.py @@ -0,0 +1,97 @@ +"""CPU contracts for Hopper/SM90 release-hardware identification.""" + +from __future__ import annotations + +import pytest + +from tests.structure.support.hardware import ( + HOPPER_SM90_CAPABILITY, + assert_recorded_hopper_device_matches, + assert_same_hopper_sm90_device, + hopper_sm90_fingerprint, +) + + +def _environment( + name: str, + *, + capability: tuple[int, int] = HOPPER_SM90_CAPABILITY, + total_memory: int = 96 * 1024**3, +) -> dict[str, object]: + return { + "cuda_device": name, + "cuda_device_capability": list(capability), + "cuda_total_memory": total_memory, + } + + +@pytest.mark.parametrize( + ("name", "product"), + ( + ("NVIDIA H100 PCIe", "H100"), + ("NVIDIA H200 NVL", "H200"), + ("NVIDIA GH200 480GB", "GH200"), + ), +) +def test_release_hardware_accepts_named_hopper_sm90_products( + name: str, + product: str, +) -> None: + fingerprint = hopper_sm90_fingerprint(_environment(name)) + + assert fingerprint.product == product + assert fingerprint.capability == (9, 0) + + +@pytest.mark.parametrize( + "environment", + ( + _environment("NVIDIA A100-SXM4-80GB", capability=(8, 0)), + _environment("NVIDIA B200", capability=(10, 0)), + _environment("NVIDIA H100 PCIe", capability=(8, 0)), + { + "cuda_device": "NVIDIA H200", + "cuda_device_capability": [9, 0], + "cuda_total_memory": 0, + }, + ), +) +def test_release_hardware_rejects_non_hopper_or_incomplete_identity( + environment: dict[str, object], +) -> None: + with pytest.raises(AssertionError): + hopper_sm90_fingerprint(environment) + + +def test_comparisons_require_the_exact_same_hopper_device_fingerprint() -> None: + h100 = _environment("NVIDIA H100 PCIe", total_memory=80 * 1024**3) + assert_same_hopper_sm90_device(h100, dict(h100)) + + with pytest.raises(AssertionError, match="Cross-device comparison is forbidden"): + assert_same_hopper_sm90_device( + _environment("NVIDIA GH200 480GB", total_memory=96 * 1024**3), + h100, + ) + with pytest.raises(AssertionError, match="Cross-device comparison is forbidden"): + assert_same_hopper_sm90_device( + _environment("NVIDIA H100 PCIe", total_memory=94 * 1024**3), + h100, + ) + + +def test_golden_comparison_rejects_cross_device_and_honors_new_identity_fields() -> None: + current = _environment("NVIDIA GH200 480GB", total_memory=96 * 1024**3) + legacy_record = {"cuda_device": "NVIDIA GH200 480GB"} + assert_recorded_hopper_device_matches(current, legacy_record) + assert_recorded_hopper_device_matches(current, dict(current)) + + with pytest.raises(AssertionError, match="Cross-device golden comparison is forbidden"): + assert_recorded_hopper_device_matches( + current, + {"cuda_device": "NVIDIA H100 PCIe"}, + ) + with pytest.raises(AssertionError, match="cuda_total_memory"): + assert_recorded_hopper_device_matches( + current, + {**current, "cuda_total_memory": 80 * 1024**3}, + ) diff --git a/tests/unit/test_import_hygiene.py b/tests/unit/test_import_hygiene.py new file mode 100644 index 0000000..ebe9419 --- /dev/null +++ b/tests/unit/test_import_hygiene.py @@ -0,0 +1,104 @@ +"""Fresh-process import contracts for the package and model modules.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +MODEL_MODULES = ( + "fastplms.models.ankh.modeling_ankh", + "fastplms.models.boltz.modeling_boltz2", + "fastplms.models.dplm.modeling_dplm", + "fastplms.models.dplm2.modeling_dplm2", + "fastplms.models.e1.modeling_e1", + "fastplms.models.esm2.modeling_fastesm", + "fastplms.models.esm3.modeling_esm3", + "fastplms.models.esm_plusplus.modeling_esm_plusplus", + "fastplms.models.esmfold.modeling_fast_esmfold", + "fastplms.models.esmfold2.modeling_esmfold2", + "fastplms.models.esmfold2.modeling_esmfold2_experimental", +) + + +def _run_fresh(script: str) -> subprocess.CompletedProcess[str]: + environment = { + **os.environ, + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "TOKENIZERS_PARALLELISM": "false", + } + return subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + env=environment, + timeout=120, + ) + + +def test_top_level_package_and_models_namespace_are_lazy() -> None: + completed = _run_fresh( + "import sys\n" + "import fastplms\n" + "import fastplms.models\n" + "assert 'torch' not in sys.modules\n" + "assert 'transformers' not in sys.modules\n" + "assert 'huggingface_hub' not in sys.modules\n" + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout == "" + + +def test_model_imports_do_not_download_compile_log_or_mutate_torch_globals() -> None: + modules = repr(MODEL_MODULES) + completed = _run_fresh( + f"MODULES = {modules}\n" + "import importlib\n" + "import logging\n" + "import socket\n" + "import torch\n" + "import torch._dynamo.config as dynamo_config\n" + "import torch._inductor.config as inductor_config\n" + "import transformers\n" + "import huggingface_hub\n" + "def forbidden(*args, **kwargs):\n" + " raise AssertionError('import attempted a forbidden side effect')\n" + "torch.compile = forbidden\n" + "logging.basicConfig = forbidden\n" + "socket.create_connection = forbidden\n" + "huggingface_hub.hf_hub_download = forbidden\n" + "huggingface_hub.snapshot_download = forbidden\n" + "transformers.AutoTokenizer.from_pretrained = forbidden\n" + "def snapshot():\n" + " return (\n" + " torch.get_float32_matmul_precision(),\n" + " torch.backends.cuda.matmul.allow_tf32,\n" + " torch.backends.cudnn.allow_tf32,\n" + " torch.backends.cudnn.benchmark,\n" + " torch.backends.cudnn.deterministic,\n" + " repr(getattr(dynamo_config, '_config', None)),\n" + " repr(getattr(inductor_config, '_config', None)),\n" + " )\n" + "before = snapshot()\n" + "for module in MODULES:\n" + " importlib.import_module(module)\n" + " assert snapshot() == before, module\n" + "assert not torch.cuda.is_initialized()\n" + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout == "" + + +def test_production_models_do_not_embed_unpinned_checkpoint_downloaders() -> None: + root = Path(__file__).resolve().parents[2] / "src" / "fastplms" / "models" + for relative_path in ( + "esm_plusplus/modeling_esm_plusplus.py", + "esm3/modeling_esm3.py", + ): + source = (root / relative_path).read_text(encoding="utf-8") + assert "snapshot_download" not in source + assert "from_pretrained_esm" not in source diff --git a/tests/unit/test_parity_hidden_states.py b/tests/unit/test_parity_hidden_states.py new file mode 100644 index 0000000..af7f3af --- /dev/null +++ b/tests/unit/test_parity_hidden_states.py @@ -0,0 +1,33 @@ +"""Focused contracts for live parity output normalization.""" + +import torch +from types import SimpleNamespace + +from tests.parity.test_model_parity import _hidden_state_tuple, _last_hidden, tensor_metrics + + +def test_live_parity_accepts_layer_stacked_hidden_states() -> None: + """Official ESMC returns H stacked across the leading layer axis.""" + + H = torch.arange(3 * 2 * 4 * 5).reshape(3, 2, 4, 5) # (n_layers=3, b=2, l=4, d=5) + output = SimpleNamespace(hidden_states=H, last_hidden_state=None) + + layers = _hidden_state_tuple(output) # 3 * (b=2, l=4, d=5) + assert len(layers) == 3 + assert all(torch.equal(layer, H[index]) for index, layer in enumerate(layers)) + assert torch.equal(_last_hidden(output), H[-1]) + + +def test_live_parity_pooling_ignores_nonfinite_padding() -> None: + candidate = torch.tensor( # (b=1, l=2, d=2) + [[[1.0, 2.0], [float("nan"), float("nan")]]] + ) + official = torch.tensor( # (b=1, l=2, d=2) + [[[1.0, 2.0], [float("nan"), float("nan")]]] + ) + residue_mask = torch.tensor([[True, False]]) # (b=1, l=2) + + metrics = tensor_metrics(candidate, official, residue_mask) + + assert metrics.relative_l2 == 0.0 + assert metrics.pooled_cosine_min > 0.999999 diff --git a/tests/unit/test_reference_adapter_contracts.py b/tests/unit/test_reference_adapter_contracts.py new file mode 100644 index 0000000..d20c173 --- /dev/null +++ b/tests/unit/test_reference_adapter_contracts.py @@ -0,0 +1,406 @@ +"""Small contracts for isolated official-reference adapters.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from types import SimpleNamespace + +from tests.parity.support.native_reference import ( + _adapter_reference_sources, + _generation_contract, + _record_generation_contract, + _validated_generation_limitation, +) +from tests.parity.support.reference_adapters import OfficialGenerationUnavailable +from tests.parity.support.reference_adapters.dplm2 import ( + DPLM2_3B_GENERATION_LIMITATION, + _accepts_type_ids, + _call_checkpoint_forward, + _call_checkpoint_generate, +) +from tools.remote.reference_source_attestation import ReferenceSourceAttestationError + + +class _AcceptsTypeIds(nn.Module): + def forward( + self, + input_ids: torch.Tensor, + type_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + # input_ids: (b, l); type_ids: (b, l) or None + del type_ids + return input_ids # (b, l) + + +class _AcceptsKeywordArguments(nn.Module): + def forward(self, input_ids: torch.Tensor, **kwargs: object) -> torch.Tensor: + # input_ids: (b, l) + del kwargs + return input_ids # (b, l) + + +class _RejectsTypeIds(nn.Module): + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids # (b, l) + + +class _OfficialWrapper(nn.Module): + def __init__(self) -> None: + super().__init__() + self.generation_calls: list[dict[str, object]] = [] + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids + 1 # (b, l) + + def generate(self, input_tokens: torch.Tensor, **kwargs: object) -> torch.Tensor: + # input_tokens: (b, l) + self.generation_calls.append({"input_tokens": input_tokens, **kwargs}) + return input_tokens + 2 # (b, l) + + +class _CheckpointNetwork(_RejectsTypeIds): + def __init__(self) -> None: + super().__init__() + self.generation_calls: list[dict[str, object]] = [] + + def generate( + self, + batch: dict[str, torch.Tensor], + max_iter: int, + sampling_strategy: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + # batch["input_ids"]: (b, l) + self.generation_calls.append( + { + "batch": batch, + "max_iter": max_iter, + "sampling_strategy": sampling_strategy, + } + ) + return ( # each: (b, l) + batch["input_ids"] + 3, + torch.zeros_like(batch["input_ids"]), + ) + + +class EsmForDPLM(_RejectsTypeIds): + """Minimal reproduction of the pinned 3B sampler's missing BOS token.""" + + bos_id = None + + def generate( + self, + batch: dict[str, torch.Tensor], + max_iter: int, + sampling_strategy: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + del max_iter, sampling_strategy + tokens = batch["input_ids"] # (b, l) + tokens.ne(self.bos_id) + raise AssertionError("unreachable") + + +class _AnkhGenerationTokenizer: + def __call__( + self, + text: str, + *, + return_tensors: str, + add_special_tokens: bool = True, + ) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + if add_special_tokens: + assert text == "M S T N P K" + input_ids = torch.tensor([[4, 5, 1]]) # (b=1, l=3) + else: + assert text == "A C" + input_ids = torch.tensor([[2, 3]]) # (b=1, l=2) + return { + "input_ids": input_ids, # (b=1, l) + "attention_mask": torch.ones_like(input_ids), # (b=1, l) + } + + +class _AnkhGenerationModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(decoder_start_token_id=0) + self.generation_calls: list[dict[str, object]] = [] + + def generate(self, **kwargs: object) -> torch.Tensor: + self.generation_calls.append(kwargs) + decoder_input_ids = kwargs["decoder_input_ids"] + assert torch.is_tensor(decoder_input_ids) + return torch.cat( # (b=1, l + 1) + (decoder_input_ids, decoder_input_ids.new_tensor([[9]])), # (1, l); (1, 1) + dim=1, + ) + + +class _AnkhGenerationAdapter: + def __init__(self) -> None: + self.model = _AnkhGenerationModel() + self.loads: list[dict[str, object]] = [] + + def load_official_seq2seq( + self, + **kwargs: object, + ) -> tuple[_AnkhGenerationModel, _AnkhGenerationTokenizer]: + self.loads.append(kwargs) + return self.model, _AnkhGenerationTokenizer() + + +_VALID_SOURCE_ATTESTATION: dict[str, object] = { + "attestation_sha256": "b" * 64, + "file_count": 5218, + "import_file": "src/transformers/__init__.py", + "import_name": "transformers", + "import_root": "src/transformers", + "package_version": "4.57.6", + "schema_version": 1, + "source_revision": "a" * 40, + "tree_sha256": "c" * 64, +} +_VALID_REFERENCE_SOURCES = { + "biohub-esm": { + **_VALID_SOURCE_ATTESTATION, + "file_count": 157, + "import_file": "esm/__init__.py", + "import_name": "esm", + "import_root": "esm", + "package_version": "3.3.0", + "source_revision": "d" * 40, + "tree_sha256": "e" * 64, + }, + "biohub-transformers": _VALID_SOURCE_ATTESTATION, +} + + +@pytest.mark.parametrize("family", ("esm_plusplus", "esm3", "esmfold2")) +def test_biohub_native_adapter_requires_both_stable_source_attestations(family: str) -> None: + adapter = SimpleNamespace( + reference_sources=lambda: { + name: dict(evidence) for name, evidence in _VALID_REFERENCE_SOURCES.items() + } + ) + request = {"model_id": "biohub-probe", "family": family} + + assert _adapter_reference_sources(adapter, request) == _VALID_REFERENCE_SOURCES + with pytest.raises(RuntimeError, match="omits source attestations"): + _adapter_reference_sources(object(), request) + + incomplete = SimpleNamespace( + reference_sources=lambda: {"biohub-transformers": _VALID_SOURCE_ATTESTATION} + ) + with pytest.raises(ReferenceSourceAttestationError, match="names differ"): + _adapter_reference_sources(incomplete, request) + + +def test_native_adapter_rejects_malformed_source_attestation() -> None: + malformed = { + **_VALID_SOURCE_ATTESTATION, + "import_file": "/tmp/untrusted/transformers/__init__.py", + } + adapter = SimpleNamespace( + reference_sources=lambda: { + **_VALID_REFERENCE_SOURCES, + "biohub-transformers": malformed, + } + ) + + with pytest.raises(ReferenceSourceAttestationError, match="portable relative"): + _adapter_reference_sources( + adapter, + {"model_id": "esmc_small", "family": "esm_plusplus"}, + ) + + +def test_dplm2_checkpoint_forward_selection_is_signature_gated() -> None: + """Unsupported modality keywords select the public checkpoint network.""" + + assert _accepts_type_ids(_AcceptsTypeIds()) + assert _accepts_type_ids(_AcceptsKeywordArguments()) + + model = _RejectsTypeIds() + input_tensor = torch.tensor([1]) # (l=1,) + type_ids = torch.tensor([0]) # (l=1,) + assert not _accepts_type_ids(model) + with pytest.raises(TypeError, match="type_ids"): + model(input_tensor, type_ids=type_ids) + assert torch.equal( + _call_checkpoint_forward(_OfficialWrapper(), model, input_tensor, {}), + input_tensor, + ) + assert torch.equal( + _call_checkpoint_forward( + _OfficialWrapper(), + _AcceptsTypeIds(), + input_tensor, + {}, + ), + input_tensor + 1, + ) + + +def test_dplm2_checkpoint_generation_uses_public_selected_network() -> None: + """The 3B architecture bypasses only the broken multimodal generator.""" + + oracle = _OfficialWrapper() + network = _CheckpointNetwork() + input_tokens = torch.tensor([[1, 2]]) # (b=1, l=2) + generated = _call_checkpoint_generate( + oracle, + network, + input_tokens, + { + "max_iter": 4, + "sampling_strategy": "argmax", + "unmasking_strategy": "deterministic", + }, + ) + + assert torch.equal(generated, input_tokens + 3) + assert not oracle.generation_calls + assert network.generation_calls == [ + { + "batch": {"input_ids": input_tokens}, + "max_iter": 4, + "sampling_strategy": "argmax", + } + ] + + +def test_dplm2_multimodal_generation_is_retained_when_supported() -> None: + """Native DPLM2 networks continue through the official outer sampler.""" + + oracle = _OfficialWrapper() + input_tokens = torch.tensor([[1, 2]]) # (b=1, l=2) + generated = _call_checkpoint_generate( + oracle, + _AcceptsTypeIds(), + input_tokens, + {"max_iter": 4}, + ) + + assert torch.equal(generated, input_tokens + 2) + assert oracle.generation_calls == [ + {"input_tokens": input_tokens, "max_iter": 4} + ] + + +def test_dplm2_checkpoint_generation_normalizes_exact_public_failure() -> None: + """The unusable 3B sampler records evidence without patching the oracle.""" + + with pytest.raises(OfficialGenerationUnavailable) as captured: + _call_checkpoint_generate( + _OfficialWrapper(), + EsmForDPLM(), + torch.tensor([[1, 2]]), + {"max_iter": 4, "sampling_strategy": "argmax"}, + ) + + assert captured.value.as_record() == DPLM2_3B_GENERATION_LIMITATION + assert isinstance(captured.value.__cause__, TypeError) + + +def test_native_generation_limitation_policy_is_fail_closed() -> None: + """Only an exact official_unavailable request may publish the limitation.""" + + error = OfficialGenerationUnavailable( + public_method=DPLM2_3B_GENERATION_LIMITATION["public_method"], + exception_type=DPLM2_3B_GENERATION_LIMITATION["exception_type"], + reason=DPLM2_3B_GENERATION_LIMITATION["reason"], + ) + request = { + "model_id": "dplm2_3b", + "generation_policy": "official_unavailable", + "official_generation_limitation": DPLM2_3B_GENERATION_LIMITATION, + } + assert _validated_generation_limitation(request, error) == ( + DPLM2_3B_GENERATION_LIMITATION + ) + + with pytest.raises(RuntimeError, match="official generation is required"): + _validated_generation_limitation( + {"model_id": "dplm2_3b", "generation_policy": "required"}, + error, + ) + mutated = dict(request) + mutated["official_generation_limitation"] = { + **DPLM2_3B_GENERATION_LIMITATION, + "reason": "different", + } + with pytest.raises(RuntimeError, match="differs from the manifest-derived request"): + _validated_generation_limitation(mutated, error) + + +def test_native_ankh_generation_uses_an_explicit_decoder_prompt() -> None: + adapter = _AnkhGenerationAdapter() + request = { + "model_id": "ankh_base", + "family": "ankh", + "generation_policy": "required", + "reference_repo_id": "ElnaggarLab/ankh-base", + "reference_revision": "immutable-revision", + "seed": 42, + } + + contract = _generation_contract( + None, + object(), + request, + torch.device("cpu"), + adapter=adapter, + ) + + assert contract is not None + assert contract["decoder_prompt_contract"] == "explicit-task-prompt" + assert contract["decoder_input_ids"] == [[0, 2, 3]] + assert contract["decoder_attention_mask"] == [[1, 1, 1]] + assert contract["output_tokens"] == [[0, 2, 3, 9]] + assert len(contract["decoder_input_fingerprint"]) == 64 + assert adapter.loads == [ + { + "reference_repo_id": "ElnaggarLab/ankh-base", + "reference_revision": "immutable-revision", + "device": torch.device("cpu"), + "dtype": torch.float32, + } + ] + call = adapter.model.generation_calls[0] + assert torch.equal(call["input_ids"], torch.tensor([[4, 5, 1]])) + assert torch.equal(call["decoder_input_ids"], torch.tensor([[0, 2, 3]])) + assert call["do_sample"] is False + + +def test_native_generation_policy_rejects_missing_required_evidence() -> None: + with pytest.raises(RuntimeError, match="has no generation contract"): + _record_generation_contract( + {}, + None, + object(), + { + "model_id": "unsupported", + "family": "esm2", + "generation_policy": "required", + }, + torch.device("cpu"), + adapter=object(), + ) + + metadata: dict[str, object] = {} + _record_generation_contract( + metadata, + None, + object(), + { + "model_id": "esm2_8m", + "family": "esm2", + "generation_policy": "not_applicable", + }, + torch.device("cpu"), + adapter=object(), + ) + assert metadata == {} diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 0000000..74a185f --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,949 @@ +from __future__ import annotations + +import configparser +import json +import os +import re +import subprocess +import sys +import pytest +from pathlib import Path + +from fastplms.registry import FileDigest, RegistryError, load_model_registry + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_model_manifest_is_complete_and_typed() -> None: + registry = load_model_registry() + + assert registry.schema_version == 1 + assert {item.path for item in registry.legal_files} == { + "LICENSE", + "THIRD_PARTY_NOTICES.md", + } + assert len(registry) == 29 + assert set(registry.upstreams) == { + "ankh", + "biohub-esm", + "biohub-transformers", + "boltz", + "dplm", + "e1", + "fair-esm", + "openfold", + "protein-ttt", + } + assert {source_id: source.path for source_id, source in registry.upstreams.items()} == { + "ankh": "vendor/upstream/ankh", + "biohub-esm": "vendor/upstream/biohub-esm", + "biohub-transformers": "vendor/upstream/biohub-transformers", + "boltz": "vendor/upstream/boltz", + "dplm": "vendor/upstream/dplm", + "e1": "vendor/upstream/e1", + "fair-esm": "vendor/upstream/fair-esm", + "openfold": "vendor/upstream/openfold", + "protein-ttt": "vendor/upstream/protein-ttt", + } + assert registry["esm2_8m"].fast.repo_id == "Synthyra/ESM2-8M" + assert registry["esm2_8m"].artifact_source == "fast" + assert registry["esm2_8m"].artifact_checkpoint is registry["esm2_8m"].fast + assert registry["esm2_8m"].is_deep_reference + assert registry["esm2_650m"].size_category == "large" + assert registry["esmc_small"].family.reference_adapter.endswith(".esm_plusplus") + assert registry.families["esmfold2"].backbone_model == "esmc_6b" + assert registry[registry.families["esmfold2"].backbone_model].fast.repo_id == ( + "Synthyra/ESMplusplus_6B" + ) + assert registry["e1_150m"].family.conversion_provenance.startswith("Input:") + assert registry["dplm2_150m"].family.tokenizer_class == ( + "fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer" + ) + assert all( + family.tokenizer_class is None + for family_id, family in registry.families.items() + if family_id != "dplm2" + ) + assert { + family.id: family.bf16_execution + for family in registry.families.values() + if family.bf16_execution == "fp32_parameters_autocast" + } == { + "boltz2": "fp32_parameters_autocast", + "dplm": "fp32_parameters_autocast", + "dplm2": "fp32_parameters_autocast", + "esm2": "fp32_parameters_autocast", + "esm3": "fp32_parameters_autocast", + "esmfold": "fp32_parameters_autocast", + "esmfold2": "fp32_parameters_autocast", + } + assert all( + family.bf16_execution in {"static_parameters", "fp32_parameters_autocast"} + for family in registry.families.values() + ) + for model in registry.by_family("ankh"): + assert model.artifact_source == "official" + assert model.artifact_checkpoint is model.official + assert re.fullmatch(r"[0-9a-f]{64}", model.canonical_state_sha256 or "") + assert not model.family.requires_complete_weight_publication + assert all( + not model.family.requires_complete_weight_publication + for model in registry.values() + ) + for model in registry.by_family("dplm2"): + assert model.artifact_source == "official" + assert model.artifact_checkpoint is model.official + assert re.fullmatch(r"[0-9a-f]{64}", model.canonical_state_sha256 or "") + assert all( + model.canonical_state_sha256 is None + for model in registry.values() + if model.artifact_source == "fast" + ) + assert registry.families["dplm"].hub_license == "apache-2.0" + assert all(family.weights_publication_allowed for family in registry.families.values()) + dplm_source = registry.upstreams["dplm"] + assert {item.path for item in dplm_source.distribution_files} == { + "LICENSE", + "PROVENANCE.md", + } + for source in registry.upstreams.values(): + assert tuple(item.path for item in source.license_digests) == source.license_files + assert source.distribution_files + + +def test_weight_publication_permission_must_be_explicit(tmp_path: Path) -> None: + manifest_path = ROOT / "src" / "fastplms" / "models.toml" + manifest = manifest_path.read_text(encoding="utf-8") + assert manifest.count("weights_publication_allowed = true") == 10 + missing_decision = manifest.replace("weights_publication_allowed = true\n", "", 1) + candidate = tmp_path / "models.toml" + candidate.write_text(missing_decision, encoding="utf-8") + + with pytest.raises( + RegistryError, + match=r"weights_publication_allowed must be declared explicitly", + ): + load_model_registry(candidate) + + +def test_official_artifact_canonical_state_commitment_is_required(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + missing_commitment = re.sub( + r'^canonical_state_sha256 = "[0-9a-f]{64}"\n', + "", + manifest, + count=1, + flags=re.MULTILINE, + ) + candidate = tmp_path / "models.toml" + candidate.write_text(missing_commitment, encoding="utf-8") + + with pytest.raises(RegistryError, match="canonical_state_sha256 must be a SHA-256"): + load_model_registry(candidate) + + +def test_bf16_calibrations_are_scoped_to_models_and_backends() -> None: + import torch + + from tests.parity.test_model_parity import ( + BF16_CONTRACT, + ESM2_3B_SDPA_BF16_CONTRACT, + ESM2_OPTIMIZED_BF16_CONTRACT, + ESMC_ALTERNATE_BF16_CONTRACT, + FP32_CONTRACT, + _numeric_contract, + ) + + registry = load_model_registry() + spec = registry["esm2_3b"] + + assert _numeric_contract(spec, torch.bfloat16, None) is ESM2_3B_SDPA_BF16_CONTRACT + assert _numeric_contract(spec, torch.bfloat16, "sdpa") is ESM2_3B_SDPA_BF16_CONTRACT + assert _numeric_contract(spec, torch.bfloat16, "eager") is BF16_CONTRACT + assert _numeric_contract(spec, torch.bfloat16, "flex_attention") is ( + ESM2_OPTIMIZED_BF16_CONTRACT + ) + assert _numeric_contract(spec, torch.float32, None) is FP32_CONTRACT + + esmc = registry["esmc_6b"] + assert esmc.family.attention == ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + assert _numeric_contract(esmc, torch.bfloat16, "sdpa") is BF16_CONTRACT + assert _numeric_contract(esmc, torch.bfloat16, "flash_attention_2") is ( + ESMC_ALTERNATE_BF16_CONTRACT + ) + assert ESMC_ALTERNATE_BF16_CONTRACT.relative_l2_target == 0.029 + assert ESMC_ALTERNATE_BF16_CONTRACT.relative_q999_target == 0.049 + assert ESMC_ALTERNATE_BF16_CONTRACT.residue_cosine_target == 0.997 + assert ESMC_ALTERNATE_BF16_CONTRACT.jsd_target == 0.0004 + + +def test_generation_contracts_are_explicit_and_exact() -> None: + registry = load_model_registry() + required = { + "ankh_base", + "ankh_large", + "ankh2_large", + "ankh3_large", + "ankh3_xl", + "dplm_150m", + "dplm_650m", + "dplm_3b", + "dplm2_150m", + "dplm2_650m", + } + unavailable = {"dplm2_3b"} + + assert {model.id for model in registry.values() if model.generation_contract == "required"} == ( + required + ) + assert { + model.id + for model in registry.values() + if model.generation_contract == "official_unavailable" + } == unavailable + assert { + model.id for model in registry.values() if model.generation_contract == "not_applicable" + } == set(registry).difference(required, unavailable) + + +def test_esmfold2_ccd_runtime_asset_is_typed_and_immutable() -> None: + registry = load_model_registry() + assert set(registry.runtime_assets) == {"esmfold2_ccd"} + asset = registry.runtime_assets["esmfold2_ccd"] + assert asset.repository == "biohub/ESMFold2" + assert asset.revision == "1ebf0e3481a5184eb6171d40615c79e384b48796" + assert asset.path == "ccd.pkl" + assert asset.sha256 == "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5" + assert asset.size == 417306584 + assert asset.consumer_family == "esmfold2" + assert asset.trust_kind == "hash_pinned_pickle" + assert asset.license_expression == "MIT" + assert asset.offline_behavior == "requires_cached_verified_file" + + +def test_attention_kernel_revisions_are_typed_and_immutable() -> None: + registry = load_model_registry() + assert { + implementation: kernel.revision + for implementation, kernel in registry.attention_kernels.items() + } == { + "flash_attention_2": "db6b51744f0cd7061386442c09df890fc6d9f47e", + "flash_attention_3": "43f0bd269777115d94ff826e0d113ce9c1c9087b", + } + assert { + implementation: kernel.dtypes + for implementation, kernel in registry.attention_kernels.items() + } == { + "flash_attention_2": ("bfloat16",), + "flash_attention_3": ("bfloat16",), + } + assert registry.supported_attention_dtypes("esm2", "sdpa") == ( + "float32", + "bfloat16", + ) + assert registry.supported_attention_dtypes("esm2", "flash_attention_3") == ( + "bfloat16", + ) + + +def test_attention_kernel_lock_matches_manifest_and_h100_variants() -> None: + registry = load_model_registry() + entries = json.loads((ROOT / "kernels.lock").read_text(encoding="utf-8")) + locked = {entry["repo_id"]: entry for entry in entries} + assert set(locked) == { + "kernels-community/flash-attn2", + "kernels-community/flash-attn3", + } + for kernel in registry.attention_kernels.values(): + assert locked[kernel.repository]["sha"] == kernel.revision + assert locked[kernel.repository]["variants"] + + expected_h100 = { + "kernels-community/flash-attn2": ( + "torch213-cxx11-cu130-x86_64-linux", + "sha256-238cdad1945962331ad685a07119bb9e893ed976f11ecbf257e03d36682f95e4", + ), + "kernels-community/flash-attn3": ( + "torch-stable-abi29-cu130-x86_64-linux", + "sha256-8dc3c4645b8ed2c5ce27873f8c6deb4ecf60060b5f08f538389b3d79e8842a2f", + ), + } + for repository, (variant, digest) in expected_h100.items(): + variant_lock = locked[repository]["variants"][variant] + assert variant_lock == {"hash": digest, "hash_type": "git_lfs_concat"} + + assert { + kernel.repository: kernel.version + for kernel in registry.attention_kernels.values() + } == { + "kernels-community/flash-attn2": 2, + "kernels-community/flash-attn3": 1, + } + + +def test_manifest_rejects_mutable_attention_kernel_revision(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + 'revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"', + 'revision = "main"', + 1, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises( + RegistryError, + match=r"attention_kernels\[0\]\.revision must be an immutable", + ): + load_model_registry(path) + + +@pytest.mark.parametrize( + ("anchor", "unknown_field", "context"), + ( + ("schema_version = 1\n", 'unknown_root = "value"\n', "manifest"), + ("[[attention_kernels]]\n", 'unknown_kernel = "value"\n', "attention_kernels[0]"), + ("[[runtime_assets]]\n", 'unknown_asset = "value"\n', "runtime_assets[0]"), + ("[[upstreams]]\n", 'unknown_upstream = "value"\n', "upstreams[0]"), + ("[families.esm2]\n", 'unknown_family = "value"\n', "families.esm2"), + ("[[models]]\n", 'unknown_model = "value"\n', "models[0]"), + ), +) +def test_manifest_rejects_unknown_table_fields( + tmp_path: Path, + anchor: str, + unknown_field: str, + context: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace(anchor, anchor + unknown_field, 1) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=rf"{re.escape(context)} contains unknown fields"): + load_model_registry(path) + + +@pytest.mark.parametrize("path_value", ("official/ankh", "vendor/upstream/ankh/nested", "../ankh")) +def test_manifest_rejects_noncanonical_upstream_paths( + tmp_path: Path, + path_value: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + 'path = "vendor/upstream/ankh"', + f'path = "{path_value}"', + 1, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="normalized directory directly under"): + load_model_registry(path) + + +@pytest.mark.parametrize( + ("old", "new", "message"), + ( + ('extra = "core"', 'extra = "gpu"', r"families\.esm2\.extra must be one of"), + ( + 'test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]', + 'test_tiers = ["check", "compliance", "feature", "artifact", "nightly"]', + r"families\.esm2\.test_tiers contains unsupported tiers", + ), + ( + 'vram_tier = "sequence"', + 'vram_tier = "host-specific"', + r"families\.esm2\.vram_tier must be one of", + ), + ( + 'bf16_execution = "fp32_parameters_autocast"', + 'bf16_execution = "implicit"', + r"families\.esm2\.bf16_execution must be one of", + ), + ( + 'reference_container = "reference-esm2"', + 'reference_container = "../reference-esm2"', + r"families\.esm2\.reference_container must be a portable", + ), + ( + 'reference_adapter = "tests.parity.support.reference_adapters.esm2"', + 'reference_adapter = "fastplms.reference_adapters.esm2"', + r"families\.esm2\.reference_adapter must name one module", + ), + ( + 'documentation = "docs/models.md#esm2"', + 'documentation = "../models.md#esm2"', + r"families\.esm2\.documentation must reference a normalized Markdown file", + ), + ( + 'documentation = "docs/models.md#esm2"', + 'documentation = "docs/models.md#ESM2"', + r"families\.esm2\.documentation has an invalid heading fragment", + ), + ), +) +def test_manifest_rejects_invalid_family_enums_and_paths( + tmp_path: Path, + old: str, + new: str, + message: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace(old, new, 1) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=message): + load_model_registry(path) + + +def test_manifest_rejects_unknown_backbone_model(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + 'backbone_model = "esmc_6b"', + 'backbone_model = "missing_esmc"', + 1, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="references unknown backbone model 'missing_esmc'"): + load_model_registry(path) + + +@pytest.mark.parametrize( + ("old", "new", "message"), + ( + ( + 'repository = "biohub/ESMFold2"', + 'repository = "not-a-repository"', + r"runtime_assets\[0\]\.repository must be a Hugging Face repository ID", + ), + ( + 'revision = "1ebf0e3481a5184eb6171d40615c79e384b48796"', + 'revision = "main"', + r"runtime_assets\[0\]\.revision must be an immutable", + ), + ('path = "ccd.pkl"', 'path = "../ccd.pkl"', "Runtime asset path is not portable"), + ( + 'sha256 = "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5"', + 'sha256 = "unresolved"', + "Invalid runtime asset SHA-256", + ), + ("size = 417306584", "size = 0", r"runtime_assets\[0\]\.size must be a positive"), + ( + 'consumer_family = "esmfold2"', + 'consumer_family = "unknown"', + r"runtime_assets\[0\]\.consumer_family references unknown family", + ), + ( + 'trust_kind = "hash_pinned_pickle"', + 'trust_kind = "pickle"', + r"runtime_assets\[0\]\.trust_kind must be one of", + ), + ( + 'path = "ccd.pkl"', + 'path = "ccd.bin"', + r"runtime_assets\[0\]\.path must end in '.pkl'", + ), + ( + 'offline_behavior = "requires_cached_verified_file"', + 'offline_behavior = "download_if_missing"', + r"runtime_assets\[0\]\.offline_behavior is unsupported", + ), + ), +) +def test_manifest_rejects_invalid_runtime_asset_fields( + tmp_path: Path, + old: str, + new: str, + message: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace(old, new, 1) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=message): + load_model_registry(path) + + +def test_manifest_accepts_an_alternative_hash_pinned_runtime_asset(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + alternative = manifest.replace( + "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5", + "aff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5", + 1, + ) + path = tmp_path / "models.toml" + path.write_text(alternative, encoding="utf-8") + + registry = load_model_registry(path) + assert registry.runtime_assets["esmfold2_ccd"].sha256.startswith("aff44b") + + +@pytest.mark.parametrize( + ("old", "new", "message"), + ( + ( + 'generation_contract = "not_applicable"', + 'generation_contract = "best_effort"', + r"models\[0\]\.generation_contract must be one of", + ), + ( + 'generation_contract = "not_applicable"\n', + "", + r"models\[0\]\.generation_contract must be a non-empty string", + ), + ), +) +def test_manifest_rejects_invalid_generation_contracts( + tmp_path: Path, + old: str, + new: str, + message: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace(old, new, 1) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=message): + load_model_registry(path) + + +def test_hub_license_metadata_is_typed_and_complete() -> None: + registry = load_model_registry() + assert {family_id: family.hub_license for family_id, family in registry.families.items()} == { + "ankh": "cc-by-nc-sa-4.0", + "boltz2": "mit", + "dplm": "apache-2.0", + "dplm2": "apache-2.0", + "e1": "other", + "esm2": "mit", + "esm3": "mit", + "esm_plusplus": "mit", + "esmfold": "mit", + "esmfold2": "mit", + } + e1 = registry.families["e1"] + assert dict(e1.hub_license_metadata) == { + "license": "other", + "license_name": "profluent-e1-clickthrough-license-agreement", + "license_link": ( + "https://github.com/Profluent-AI/E1/blob/main/LICENSE" + ), + } + for family_id, family in registry.families.items(): + if family_id != "e1": + assert dict(family.hub_license_metadata) == {"license": family.hub_license} + + +def test_dplm_checkpoint_license_evidence_is_immutable_and_complete() -> None: + registry = load_model_registry() + revision = registry.upstreams["dplm"].revision + evidence = (ROOT / "LICENSES" / "dplm" / "PROVENANCE.md").read_text(encoding="utf-8") + + assert revision == "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d" + assert f"https://github.com/bytedance/dplm/blob/{revision}/LICENSE" in evidence + assert f"https://github.com/bytedance/dplm/blob/{revision}/README.md#overview" in evidence + for family_id in ("dplm", "dplm2"): + family = registry.families[family_id] + assert family.checkpoint_license == "Apache-2.0" + assert family.hub_license == "apache-2.0" + assert family.weights_publication_allowed + assert "LICENSES/dplm/PROVENANCE.md" in family.conversion_provenance + + +@pytest.mark.parametrize( + ("old", "new", "message"), + ( + ('hub_license = "mit"\n', "", "families.esm2.hub_license"), + ( + 'hub_license = "mit"\n', + 'hub_license = "proprietary"\n', + "supported Hugging Face license identifier", + ), + ( + 'hub_license = "mit"\n', + 'hub_license = "apache-2.0"\n', + "must be 'mit' for checkpoint terms", + ), + ( + 'hub_license_name = "Profluent-E1 Clickthrough License Agreement"\n', + "", + "must define hub_license_name and hub_license_link", + ), + ( + "https://github.com/Profluent-AI/E1/blob/", + "http://github.com/Profluent-AI/E1/blob/", + "hub_license_link must be an absolute HTTPS URL", + ), + ( + 'hub_license = "mit"\n', + 'hub_license = "mit"\n' + 'hub_license_name = "Unexpected custom terms"\n' + 'hub_license_link = "https://example.invalid/LICENSE"\n', + "may define hub_license_name and hub_license_link only", + ), + ( + 'hub_license = "mit"\n', + 'hub_license = "mit"\nhub_license_nam = "misspelled"\n', + "contains unsupported Hub license fields", + ), + ), +) +def test_manifest_rejects_invalid_hub_license_metadata( + tmp_path: Path, + old: str, + new: str, + message: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace(old, new, 1) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=message): + load_model_registry(path) + + +def test_esmfold2_support_is_exactly_the_approved_four() -> None: + registry = load_model_registry() + family = registry.by_family("esmfold2")[0].family + assert {model.official.repo_id for model in registry.by_family("esmfold2")} == { + "biohub/ESMFold2", + "biohub/ESMFold2-Fast", + "biohub/ESMFold2-Experimental-Cutoff2025", + "biohub/ESMFold2-Experimental-Fast-Cutoff2025", + } + assert set(family.precisions) == { + "auto", + "fp32", + "bf16", + "fp8", + } + assert family.experimental_precisions == ("fp8",) + assert family.stable_precisions == ("auto", "fp32", "bf16") + + +def test_experimental_precisions_must_be_declared_precisions(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + 'experimental_precisions = ["fp8"]', + 'experimental_precisions = ["fp8", "int4"]', + 1, + ) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="must be a subset of precisions"): + load_model_registry(path) + + +def test_boltz2_is_explicitly_provisional() -> None: + registry = load_model_registry() + spec = registry["boltz2"] + assert set(spec.family.test_tiers) == {"structure", "artifact", "benchmark"} + assert "check" not in spec.family.test_tiers + assert "compliance" not in spec.family.test_tiers + assert "provisional in FastPLMs 1.0" in spec.notes + assert "does not claim official inference equivalence" in spec.notes + + +def test_esm2_native_oracle_assets_are_hash_pinned() -> None: + registry = load_model_registry() + for model in registry.by_family("esm2"): + assert set(model.oracle_asset_map) == {"weights", "contact_regression"} + official_name = model.official.repo_id.split("/", maxsplit=1)[1] + assert model.oracle_asset_map["weights"].path == f"models/{official_name}.pt" + assert model.oracle_asset_map["contact_regression"].path == ( + f"regression/{official_name}-contact-regression.pt" + ) + for asset in model.oracle_assets: + assert asset.url == f"https://dl.fbaipublicfiles.com/fair-esm/{asset.path}" + assert len(asset.sha256) == 64 + assert asset.size > 0 + + esmfold = registry["esmfold"] + assert set(esmfold.oracle_asset_map) == {"weights"} + assert esmfold.oracle_asset_map["weights"].path == "models/esmfold_3B_v1.pt" + assert esmfold.oracle_asset_map["weights"].sha256 == ( + "e9a52579027e77d2d2e0a18218e755821f395730e86624cab9413dc117f5ca62" + ) + assert esmfold.oracle_asset_map["weights"].size == 2771653574 + + for model in registry.values(): + if model.family.id not in {"esm2", "esmfold"}: + assert not model.oracle_assets + + +def test_checkpoint_provenance_is_explicit_and_release_gated() -> None: + registry = load_model_registry() + unresolved_count = 0 + for model in registry.values(): + for checkpoint in (model.fast, model.official): + assert len(checkpoint.revision) == 40 + assert checkpoint.files + assert all(item.algorithm in {"git-sha1", "sha256"} for item in checkpoint.files) + assert set(checkpoint.file_map).isdisjoint(checkpoint.unresolved_files) + unresolved_count += len(checkpoint.unresolved_files) + if model.family.tokenizer_mode == "tokenizer": + assert any("tokenizer" in path or "vocab" in path for path in model.fast.file_map) + + assert unresolved_count == 0 + assert ( + registry["esm2_3b"].fast.file_map["model-00003-of-00003.safetensors"].digest + == "a6b3a55b9e3b2e1778de34c665c3dd17bdfdf6da9d6d5c97730c57168709ccae" + ) + registry.require_resolved("esm2_8m") + registry.require_resolved("esm2_35m") + registry.require_resolved() + + +def test_runtime_paths_cannot_include_official_sources() -> None: + registry = load_model_registry() + for family in registry.families.values(): + assert all(not path.startswith("vendor/") for path in family.runtime_paths) + assert "models/__init__.py" in family.runtime_paths + assert (ROOT / "src" / "fastplms" / "models" / "__init__.py").is_file() + + +def test_gitmodules_matches_manifest_paths_and_urls() -> None: + registry = load_model_registry() + parser = configparser.ConfigParser() + parser.read(ROOT / ".gitmodules", encoding="utf-8") + + configured = {} + for section in parser.sections(): + configured[parser[section]["path"]] = parser[section]["url"] + assert configured == {source.path: source.url for source in registry.upstreams.values()} + + +def test_manifest_rejects_an_alternative_pinned_esmfold2_repository(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + alternative = manifest.replace( + 'official_repo = "biohub/ESMFold2-Experimental-Cutoff2025"', + 'official_repo = "biohub/ESMFold2-Experimental"', + 1, + ) + path = tmp_path / "models.toml" + path.write_text(alternative, encoding="utf-8") + + with pytest.raises(RegistryError, match="exactly the four approved"): + load_model_registry(path) + + +def test_manifest_rejects_a_fifth_esmfold2_checkpoint(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + 'id = "esmfold"\nfamily = "esmfold"', + 'id = "esmfold_legacy_fifth"\nfamily = "esmfold2"\nmsa_conditioning = false', + 1, + ) + invalid = invalid.replace( + "tests/goldens/esmfold.json", + "tests/goldens/esmfold_legacy_fifth.json", + 1, + ).replace( + "tests/goldens/esmfold.safetensors", + "tests/goldens/esmfold_legacy_fifth.safetensors", + 1, + ) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="exactly the four approved"): + load_model_registry(path) + + +def test_manifest_rejects_missing_e1_legal_notice(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = re.sub( + r'^ "MODIFICATIONS\.md=sha256:[0-9a-f]{64}",\n', + "", + manifest, + count=1, + flags=re.MULTILINE, + ) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="missing E1 legal files"): + load_model_registry(path) + + +def test_manifest_rejects_missing_conversion_record(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = re.sub( + r"^conversion_provenance = .*\n", + "", + manifest, + count=1, + flags=re.MULTILINE, + ) + assert invalid != manifest + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="conversion_provenance"): + load_model_registry(path) + + +@pytest.mark.parametrize( + "runtime_paths", + [ + '["../secret.py"]', + '["C:/secret.py"]', + '["./__init__.py"]', + r"['models\unsafe.py']", + '["__init__.py", "__init__.py"]', + '["models/__pycache__"]', + '["NUL"]', + '["models/con.py"]', + '["models/bad:name.py"]', + '["models/trailing."]', + ], +) +def test_manifest_rejects_nonportable_runtime_paths( + tmp_path: Path, + runtime_paths: str, +) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = re.sub( + r"^runtime_paths = .*?$", + lambda _: f"runtime_paths = {runtime_paths}", + manifest, + count=1, + flags=re.MULTILINE, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match=r"runtime_paths|Unsafe runtime path"): + load_model_registry(path) + + +def test_manifest_rejects_non_boolean_complete_publication_policy(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + "requires_complete_weight_publication = false", + 'requires_complete_weight_publication = "yes"', + 1, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="must be a boolean"): + load_model_registry(path) + + +def test_manifest_rejects_unpinned_oracle_asset(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + invalid = manifest.replace( + "46f002a9870c9bdecd0ea887acb1f9a38a6b561e8f8bf8a6990b679b9d31b928", + "unresolved", + 1, + ) + path = tmp_path / "models.toml" + path.write_text(invalid, encoding="utf-8") + + with pytest.raises(RegistryError, match="Invalid oracle asset SHA-256"): + load_model_registry(path) + + +def test_manifest_parses_optional_hash_pinned_official_golden(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + declaration = ( + 'official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:' + + "a" * 64 + + '", tensors = "tests/goldens/esm2_8m.safetensors=sha256:' + + "b" * 64 + + '" }\n' + ) + modified = re.sub( + r'^official_golden = \{ metadata = "tests/goldens/esm2_8m\.json=.*\n', + declaration, + manifest, + count=1, + flags=re.MULTILINE, + ) + path = tmp_path / "models.toml" + path.write_text(modified, encoding="utf-8") + + golden = load_model_registry(path)["esm2_8m"].official_golden + assert golden is not None + assert golden.metadata.path == "tests/goldens/esm2_8m.json" + assert golden.metadata.digest == "a" * 64 + assert golden.tensors.path == "tests/goldens/esm2_8m.safetensors" + assert golden.tensors.digest == "b" * 64 + + +def test_manifest_rejects_an_unsafe_official_golden_path(tmp_path: Path) -> None: + manifest = (ROOT / "src" / "fastplms" / "models.toml").read_text(encoding="utf-8") + declaration = ( + 'official_golden = { metadata = "../esm2_8m.json=sha256:' + + "a" * 64 + + '", tensors = "tests/goldens/esm2_8m.safetensors=sha256:' + + "b" * 64 + + '" }\n' + ) + modified = re.sub( + r'^official_golden = \{ metadata = "tests/goldens/esm2_8m\.json=.*\n', + declaration, + manifest, + count=1, + flags=re.MULTILINE, + ) + path = tmp_path / "models.toml" + path.write_text(modified, encoding="utf-8") + + with pytest.raises(RegistryError, match="Checkpoint file path is not portable"): + load_model_registry(path) + + +@pytest.mark.parametrize( + "value", + [ + "../model.safetensors=sha256:" + "a" * 64, + "model.safetensors=md5:" + "a" * 32, + "model.safetensors=sha256:short", + ], +) +def test_file_digest_rejects_unsafe_or_unverifiable_values(value: str) -> None: + with pytest.raises(RegistryError): + FileDigest.parse(value) + + +def test_top_level_import_does_not_import_torch() -> None: + environment = os.environ.copy() + environment["PYTHONPATH"] = str(ROOT / "src") + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys, fastplms; assert fastplms.__version__ == '1.0.0'; " + "assert 'torch' not in sys.modules", + ], + check=False, + capture_output=True, + text=True, + env=environment, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/unit/test_remote_runner_contract.py b/tests/unit/test_remote_runner_contract.py new file mode 100644 index 0000000..2d1b3b9 --- /dev/null +++ b/tests/unit/test_remote_runner_contract.py @@ -0,0 +1,765 @@ +"""Portable remote-runner policy tests.""" + +from __future__ import annotations + +import hashlib +import json +import subprocess +import pytest +from pathlib import Path + +from tools.remote.run import ( + REMOTE_CLEANUP_SCRIPT, + SENSITIVE_SUFFIXES, + SUITES, + RemoteRunner, + RunnerConfig, + _artifact_tree_summary, + _host_hardware_preflight, + _is_sensitive, + _kernel_capability_preflight, + _reference_container_image_identity, + _require_clean_repository, + _require_matching_archive_digest, + _run_report, + remote_cleanup_command, +) + + +def test_required_remote_suites_are_available() -> None: + assert {"check", "compliance", "structure", "feature", "artifact", "benchmark"}.issubset(SUITES) + + +def test_sensitive_files_are_never_archived() -> None: + assert _is_sensitive(Path(".secrets.env")) + assert _is_sensitive(Path("credentials.json")) + for directory in (".agents", ".claude", ".codex"): + assert _is_sensitive(Path(directory) / "workspace-state.json") + for suffix in SENSITIVE_SUFFIXES: + assert _is_sensitive(Path(f"identity{suffix}")) + assert _is_sensitive(Path("vendor/upstream/example/.git/config")) + + +def test_uploaded_source_archive_digest_is_verified_against_local_bytes() -> None: + expected = hashlib.sha256(b"archive").hexdigest() + _require_matching_archive_digest(f"{expected} source.tar.gz\n", expected) + + with pytest.raises(RuntimeError, match="differs from local bytes"): + _require_matching_archive_digest(f"{'0' * 64} source.tar.gz\n", expected) + with pytest.raises(RuntimeError, match="differs from local bytes"): + _require_matching_archive_digest("", expected) + + +def test_remote_archive_inventory_is_tracked_only() -> None: + source = (Path(__file__).resolve().parents[2] / "tools" / "remote" / "run.py").read_text( + encoding="utf-8" + ) + git_files = source.split("def _git_files", maxsplit=1)[1].split( + "def _require_clean_repository", maxsplit=1 + )[0] + assert '"--cached"' in git_files + assert '"--others"' not in git_files + assert '"--exclude-standard"' not in git_files + + +def test_connection_details_are_runtime_only() -> None: + source = (Path(__file__).resolve().parents[2] / "tools" / "remote" / "run.py").read_text( + encoding="utf-8" + ) + assert 'parser.add_argument("--host", required=True' in source + assert 'parser.add_argument("--identity", required=True' in source + assert ".ssh/" not in source + assert source.count('"IdentitiesOnly=yes"') == 2 + assert "secrets.token_hex(8)" in source + + +def test_recursive_cleanup_is_verified_by_remote_realpath() -> None: + command = remote_cleanup_command( + "/home/ubuntu/fastplms-runs", + "/home/ubuntu/fastplms-runs/20260714T120000Z-1234abcd", + ) + assert command[:2] == ("sh", "-c") + assert command[2] == REMOTE_CLEANUP_SCRIPT + assert 'base=$(realpath -e -- "$1")' in REMOTE_CLEANUP_SCRIPT + assert 'workspace=$(realpath -e -- "$2")' in REMOTE_CLEANUP_SCRIPT + assert '"$base"/*' in REMOTE_CLEANUP_SCRIPT + assert 'test "$workspace" != "$base"' in REMOTE_CLEANUP_SCRIPT + assert 'rm -rf -- "$workspace"' in REMOTE_CLEANUP_SCRIPT + + +def test_remote_runner_retrieves_the_complete_artifact_tree() -> None: + source = (Path(__file__).resolve().parents[2] / "tools" / "remote" / "run.py").read_text( + encoding="utf-8" + ) + assert "remote_workspace}/artifacts/." in source + + +@pytest.mark.parametrize( + ("failure_phase", "failing_command"), + ( + ("create-remote-workspace", "mkdir"), + ("upload-source-archive", "scp"), + ("verify-source-archive", "sha256sum"), + ("extract-source-archive", "tar"), + ("remove-source-archive", "rm"), + ), +) +def test_remote_staging_failures_are_reported_and_cleaned( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + failure_phase: str, + failing_command: str, +) -> None: + repository = tmp_path / "repository" + repository.mkdir() + identity = tmp_path / "identity" + identity.write_text("test", encoding="utf-8") + artifacts = tmp_path / "artifacts" + revision = "a" * 40 + + monkeypatch.setattr("tools.remote.run._require_clean_repository", lambda _repository: None) + monkeypatch.setattr("tools.remote.run._git_head_revision", lambda _repository: revision) + + def write_archive(_repository: Path, destination: Path) -> dict[str, dict[str, object]]: + destination.write_bytes(b"archive") + return {} + + monkeypatch.setattr("tools.remote.run.create_source_archive", write_archive) + + runner = RemoteRunner( + RunnerConfig( + host="gpu-host", + identity=identity, + repository=repository, + suite="unit", + artifacts=artifacts, + ) + ) + monkeypatch.setattr( + runner, + "_capture_host_hardware", + lambda: _host_hardware_preflight( + "aarch64", + "NVIDIA GH200 480GB, GPU-test, 580.1, 97871\n", + ), + ) + monkeypatch.setattr(runner, "_remote_base", lambda: "/remote/fastplms-runs") + ssh_commands: list[tuple[str, ...]] = [] + + def run_ssh(command, *, capture=False, timeout_seconds=None): + del capture, timeout_seconds + value = tuple(command) + ssh_commands.append(value) + if failing_command in {"mkdir", "sha256sum", "tar", "rm"} and value[0] == failing_command: + raise subprocess.CalledProcessError(7, value) + stdout = "" + if value[0] == "sha256sum": + stdout = hashlib.sha256(b"archive").hexdigest() + " source.tar.gz\n" + return subprocess.CompletedProcess(value, 0, stdout=stdout) + + monkeypatch.setattr(runner, "_ssh", run_ssh) + + def run_local(command, *, check, **kwargs): + del kwargs + value = tuple(command) + if "-r" in value: + assert check is False + return subprocess.CompletedProcess(value, 1) + if failing_command == "scp": + assert check is True + raise subprocess.CalledProcessError(6, value) + return subprocess.CompletedProcess(value, 0) + + monkeypatch.setattr(subprocess, "run", run_local) + + with pytest.raises(subprocess.CalledProcessError): + runner.run() + + cleanup = remote_cleanup_command( + "/remote/fastplms-runs", + f"/remote/fastplms-runs/{runner.run_id}", + ) + assert cleanup in ssh_commands + report_path = artifacts / runner.run_id / "remote-run.json" + report = json.loads(report_path.read_text(encoding="utf-8")) + assert report["status"] == "failed" + assert report["failure"]["phase"] == failure_phase + assert report["remote_cleanup"] == "succeeded" + + +def test_remote_runner_rejects_a_dirty_exact_head( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + def dirty_status(*args, **kwargs): + del args, kwargs + return subprocess.CompletedProcess([], 0, stdout="?? scratch.py\n") + + monkeypatch.setattr(subprocess, "run", dirty_status) + + with pytest.raises(RuntimeError, match="clean Git worktree"): + _require_clean_repository(tmp_path) + + +def test_remote_run_report_is_machine_readable_and_secret_free() -> None: + suite = SUITES["unit"] + report = _run_report( + run_id="20260714T120000Z-1234abcd", + suite_name="unit", + suite=suite, + started_at="2026-07-14T12:00:00+00:00", + finished_at="2026-07-14T12:01:00+00:00", + source_archive_sha256="a" * 64, + git_revision="b" * 40, + submodule_revisions={"vendor/upstream/example": "c" * 40}, + execution_environment={ + "host_kernel": "Linux 6.8.0 x86_64", + "docker_server": {"Version": "28.0.0"}, + "gpus": ["NVIDIA H100, 580.1"], + "images": {"candidate": {"id": "sha256:" + "d" * 64}}, + }, + failure_phase=None, + failure=None, + artifact_retrieval_returncode=0, + cleanup_status="succeeded", + ) + assert report["status"] == "passed" + assert report["schema_version"] == 5 + assert report["git_revision"] == "b" * 40 + assert report["submodule_revisions"] == {"vendor/upstream/example": "c" * 40} + environment = report["execution_environment"] + assert isinstance(environment, dict) + assert environment["images"]["candidate"]["id"] == "sha256:" + "d" * 64 + assert report["suite_contract"] == { + "bake_targets": list(suite.bake_targets), + "pre_commands": [list(command) for command in suite.pre_commands], + "command": list(suite.command), + "required_paths": list(suite.required_paths), + "biohub_reference_targets": [], + "reference_targets": [], + "host_hardware_binding_required": True, + "attention_backends": [], + "kernel_downloads_allowed": False, + "same_host_candidate_reference_required": False, + "timeouts_seconds": { + "control": 300, + "transfer": 1_800, + "build": suite.build_timeout_seconds, + "pre_command": suite.pre_command_timeout_seconds, + "command": suite.command_timeout_seconds, + }, + } + assert report["phase_durations_seconds"] == {} + assert report["cache_telemetry"] == {} + assert report["artifact_inventory"] is None + assert report["host_hardware_preflight"] is None + assert report["kernel_capability_preflight"] is None + serialized = str(report).lower() + assert "identity" not in serialized + assert "private-key" not in serialized + assert "gpu-host" not in serialized + + +def test_retrieved_artifact_inventory_is_content_addressed_and_secret_free( + tmp_path: Path, +) -> None: + (tmp_path / "junit").mkdir() + (tmp_path / "junit" / "report.xml").write_text("", encoding="utf-8") + first = _artifact_tree_summary(tmp_path) + second = _artifact_tree_summary(tmp_path) + + assert first == second + assert first["status"] == "captured" + assert first["file_count"] == 1 + assert first["total_bytes"] == len("") + assert len(str(first["tree_sha256"])) == 64 + assert "report.xml" not in str(first) + + (tmp_path / "credentials.json").write_text("secret", encoding="utf-8") + with pytest.raises(RuntimeError, match="sensitive path"): + _artifact_tree_summary(tmp_path) + + +def test_remote_run_report_records_failure_without_exception_text() -> None: + failure = subprocess.CalledProcessError( + 7, + ["ssh", "-i", "/secret/private-key", "user@gpu-host"], + ) + report = _run_report( + run_id="20260714T120000Z-1234abcd", + suite_name="unit", + suite=SUITES["unit"], + started_at="2026-07-14T12:00:00+00:00", + finished_at="2026-07-14T12:01:00+00:00", + source_archive_sha256="a" * 64, + git_revision="b" * 40, + submodule_revisions={}, + execution_environment=None, + failure_phase="suite", + failure=failure, + artifact_retrieval_returncode=0, + cleanup_status="succeeded", + ) + assert report["failure"] == { + "phase": "suite", + "type": "CalledProcessError", + "returncode": 7, + } + serialized = str(report) + assert "/secret/private-key" not in serialized + assert "user@gpu-host" not in serialized + + +def test_compliance_runs_native_services_before_candidate_comparison() -> None: + suite = SUITES["compliance"] + commands = "\n".join(" ".join(command) for command in suite.pre_commands) + assert "tools.artifacts.build_all" in commands + assert "kernels download" not in commands + assert "tools.remote.prepare_references" in commands + for service in ( + "reference-esm2", + "reference-biohub-esm", + "reference-e1", + "reference-dplm", + "reference-ankh", + ): + assert service in suite.bake_targets + assert service in commands + assert "references" not in suite.bake_targets + for service in ("reference-esmfold", "reference-esmfold2"): + assert service in suite.bake_targets + assert service in commands + assert "reference-boltz2" not in suite.bake_targets + assert "tests.structure.support.boltz2_bundle" not in commands + command = " ".join(suite.command) + assert " fp8 " in f" {command} " + assert "tests/parity/test_native_results.py" in command + fp8_stack_test = ( + "tests/release/test_validation_stack.py::" + "test_fp8_validation_stack_uses_the_cuda13_transformer_engine_core" + ) + assert fp8_stack_test in command + assert f"--deselect={fp8_stack_test}" not in command + assert "tests/structure/test_boltz2_folding_compliance.py" not in command + assert "tests/structure/test_esmfold_folding_compliance.py" in command + assert "tests/structure/test_esmfold2_folding_compliance.py" in command + assert "tests/structure/test_esmfold2_fp8_compliance.py" in command + assert "biohub-biotraj-wheel" in suite.bake_targets + + +def test_host_hardware_preflight_binds_exact_gh200_arm64_identity() -> None: + preflight = _host_hardware_preflight( + "aarch64\n", + "NVIDIA GH200 480GB, GPU-1234, 580.1, 97871\n", + ) + + assert preflight["status"] == "passed" + assert preflight["uname_machine"] == "aarch64" + assert preflight["architecture"] == "arm64" + assert preflight["container_platform"] == "linux/arm64" + assert preflight["gpus"] == [ + { + "name": "NVIDIA GH200 480GB", + "uuid": "GPU-1234", + "driver_version": "580.1", + "memory_total_mib": 97871, + } + ] + assert len(str(preflight["identity_sha256"])) == 64 + + +def test_every_suite_accepts_the_bound_gh200_hardware_contract() -> None: + preflight = _host_hardware_preflight( + "aarch64", + "NVIDIA GH200 480GB, GPU-1234, 580.1, 97871\n", + ) + assert preflight["status"] == "passed" + assert all(suite.bake_targets for suite in SUITES.values()) + for suite in SUITES.values(): + if {"reference-biohub-esm", "reference-esmfold2"}.intersection( + suite.bake_targets + ): + assert "biohub-biotraj-wheel" in suite.bake_targets + for suite_name in ("compliance", "structure", "release"): + assert { + "reference-biohub-esm", + "reference-esmfold2", + }.intersection(SUITES[suite_name].bake_targets) + + +def test_gh200_kernel_policy_is_explicit_no_download_and_fail_closed() -> None: + hardware = _host_hardware_preflight( + "aarch64", + "NVIDIA GH200 480GB, GPU-1234, 580.1, 97871\n", + ) + policy = _kernel_capability_preflight( + hardware, + ("eager", "sdpa", "flex_attention"), + ) + + assert policy["status"] == "passed" + assert policy["network_downloads"] is False + assert policy["source_builds"] is False + assert policy["selected_backends"] == ["eager", "sdpa", "flex_attention"] + backends = policy["backends"] + assert backends["flash_attention_2"]["status"] == "prior_focused_evidence_only" + assert backends["flash_attention_2"]["selected"] is False + assert backends["flash_attention_3"]["status"] == "unavailable" + assert backends["flash_attention_3"]["selected"] is False + + rejected = _kernel_capability_preflight(hardware, ("sdpa", "flash_attention_3")) + assert rejected["status"] == "failed" + assert "flash_attention_3" in str(rejected["reason"]) + + +def test_reference_container_identity_is_stable_and_excludes_ephemeral_fields() -> None: + digest = "sha256:" + "a" * 64 + identity = _reference_container_image_identity( + { + "container_platform": "linux/arm64", + "docker_server": { + "Version": "28.0.0", + "ApiVersion": "1.48", + "Arch": "arm64", + "Os": "linux", + "Name": "ephemeral-hostname", + }, + "docker_buildx": "github.com/docker/buildx v0.25.0 deadbeef", + "images": { + "reference-biohub-esm": { + "tag": "local/ephemeral:tag", + "id": digest, + "content_digest": digest, + "created": "2026-07-22T00:00:00Z", + "os": "linux", + "architecture": "arm64", + "resolved_platform": "linux/arm64", + } + }, + } + ) + + assert identity == { + "schema_version": 1, + "resolved_platform": "linux/arm64", + "docker_server": { + "Version": "28.0.0", + "ApiVersion": "1.48", + "Os": "linux", + "Arch": "arm64", + }, + "docker_buildx": "github.com/docker/buildx v0.25.0 deadbeef", + "images": { + "reference-biohub-esm": { + "content_digest": digest, + "image_id": digest, + "os": "linux", + "architecture": "arm64", + "resolved_platform": "linux/arm64", + } + }, + } + assert "tag" not in str(identity).lower() + assert "created" not in str(identity).lower() + assert "ephemeral-hostname" not in str(identity) + + +def test_gh200_hardware_binding_happens_before_archive_or_build( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + repository = tmp_path / "repository" + repository.mkdir() + identity = tmp_path / "identity" + identity.write_text("test", encoding="utf-8") + artifacts = tmp_path / "artifacts" + revision = "a" * 40 + archive_called = False + hardware = _host_hardware_preflight( + "aarch64", + "NVIDIA GH200 480GB, GPU-test, 580.1, 97871\n", + ) + + monkeypatch.setattr("tools.remote.run._require_clean_repository", lambda _root: None) + monkeypatch.setattr("tools.remote.run._git_head_revision", lambda _root: revision) + + def stop_after_preflight(_repository: Path, _destination: Path): + nonlocal archive_called + archive_called = True + raise RuntimeError("stop after hardware preflight") + + monkeypatch.setattr("tools.remote.run.create_source_archive", stop_after_preflight) + runner = RemoteRunner( + RunnerConfig( + host="gpu-host", + identity=identity, + repository=repository, + suite="compliance", + artifacts=artifacts, + ) + ) + monkeypatch.setattr(runner, "_remote_base", lambda: "/remote/fastplms-runs") + monkeypatch.setattr(runner, "_capture_host_hardware", lambda: hardware) + + with pytest.raises(RuntimeError, match="stop after hardware preflight"): + runner.run() + + assert archive_called + report = json.loads( + (artifacts / runner.run_id / "remote-run.json").read_text(encoding="utf-8") + ) + assert report["status"] == "failed" + assert report["failure"]["phase"] == "create-source-archive" + assert report["host_hardware_preflight"] == hardware + + +def test_check_uses_candidate_goldens_without_artifacts_or_live_references() -> None: + suite = SUITES["check"] + assert suite.bake_targets == ("candidate-structure",) + pre_commands = "\n".join(" ".join(command) for command in suite.pre_commands) + command = " ".join(suite.command) + assert "tests/unit" in command + assert "tests/integration" in command + assert "tests/release" in command + assert "tools.artifacts.build_all" not in pre_commands + assert "tests/release/test_published_automodel.py" not in pre_commands + assert "kernels download" not in pre_commands + assert "tests/integration/test_official_goldens.py" in pre_commands + assert "tests/structure/test_structure_official_goldens.py" in pre_commands + assert "reference-" not in pre_commands + assert "tests.parity.support.native_reference" not in pre_commands + assert suite.attention_backends == ("eager", "sdpa", "flex_attention") + + +def test_remote_builds_targets_together_and_enforces_remote_timeouts() -> None: + source = (Path(__file__).resolve().parents[2] / "tools" / "remote" / "run.py").read_text( + encoding="utf-8" + ) + assert "*suite.bake_targets" in source + assert 'f"*.platform={host_hardware_preflight[' in source + assert 'f"*.platform={container_platform}"' in source + assert "value.get(\"Architecture\") != expected_architecture" in source + assert "Remote host hardware identity changed during the build" in source + assert '"--kill-after=30s"' in source + assert "timeout_seconds=suite.build_timeout_seconds" in source + assert "timeout_seconds=suite.pre_command_timeout_seconds" in source + assert "timeout_seconds=suite.command_timeout_seconds" in source + assert "timeout=_TRANSFER_TIMEOUT_SECONDS" in source + assert 'cache_telemetry["before_build"]' in source + assert 'cache_telemetry["after_run"]' in source + assert "kernels download" not in source + assert "persist-reference-container-identity" in source + assert "artifacts/reference/environment/container-images.json" in source + + +def test_gpu_golden_smoke_uses_checked_in_results_without_reference_images() -> None: + suite = SUITES["gpu-golden-smoke"] + assert suite.bake_targets == ("candidate-structure",) + assert suite.pre_commands == () + command = " ".join(suite.command) + assert "test_release_hopper_sm90_gpu_is_available_without_running_a_model" in command + assert "tests/integration/test_official_goldens.py" in command + assert "tests/structure/test_structure_official_goldens.py" in command + assert "gpu and not large" in command + assert "reference-" not in command + + +def test_nightly_uses_candidate_goldens_features_artifacts_fp8_and_throughput() -> None: + suite = SUITES["nightly"] + assert suite.bake_targets == ( + "candidate", + "candidate-structure", + "candidate-fp8", + "candidate-artifact", + ) + pre_commands = "\n".join(" ".join(command) for command in suite.pre_commands) + command = " ".join(suite.command) + assert "tools.artifacts.build_all" in pre_commands + assert "tests/release/test_published_automodel.py" in pre_commands + assert "test_esmfold2_fp8_compliance.py" in pre_commands + assert "artifacts/benchmarks/nightly-h100.json" in pre_commands + assert "--artifact-root dist/hub" in pre_commands + assert pre_commands.index("tools.artifacts.build_all") < pre_commands.index( + "artifacts/benchmarks/nightly-h100.json" + ) + assert "test_official_goldens.py" in pre_commands + assert "test_structure_official_goldens.py" in pre_commands + assert "test_flash_attention_backends.py" not in command + assert "kernels download" not in pre_commands + assert "--backends eager sdpa flex_attention" in pre_commands + assert "test_fine_tuning_example.py" in command + assert "reference-" not in pre_commands + + +def test_benchmark_requires_a_tracked_baseline_and_capture_is_descriptive() -> None: + gated = SUITES["benchmark"] + capture = SUITES["benchmark-capture"] + gated_command = " ".join(gated.command) + capture_command = " ".join(capture.command) + assert gated.required_paths == ("benchmarks/baselines/h100.json",) + assert "--baseline benchmarks/baselines/h100.json" in gated_command + assert "--baseline" not in capture_command + assert "h100-baseline-candidate.json" in capture_command + for suite, command in ((gated, gated_command), (capture, capture_command)): + assert suite.bake_targets == ("candidate", "candidate-fp8") + pre_commands = "\n".join(" ".join(item) for item in suite.pre_commands) + assert "tools.artifacts.build_all" in pre_commands + assert "--benchmark-suite" in pre_commands + assert "kernels download" not in pre_commands + assert "--artifact-root dist/hub" in command + assert "--backends eager sdpa flex_attention" in command + assert "--junit-output artifacts/junit/" in command + assert suite.attention_backends == ("eager", "sdpa", "flex_attention") + assert suite.pre_command_timeout_seconds >= 14_400 + + +def test_live_release_and_benchmark_suites_enforce_hopper_sm90_hardware() -> None: + compliance = " ".join(SUITES["compliance"].command) + benchmark = " ".join(SUITES["benchmark"].command) + + assert "test_release_hopper_sm90_gpu_is_available_without_running_a_model" in compliance + assert "benchmarks/baselines/h100.json" in benchmark # Legacy compatibility filename. + source = (Path(__file__).resolve().parents[2] / "benchmarks" / "suite.py").read_text( + encoding="utf-8" + ) + assert "validate_hopper_sm90_environment(environment)" in source + assert '"validated_hopper_sm90_exact_device"' in source + + +def test_unit_suite_uses_the_structure_dependency_superset() -> None: + suite = SUITES["unit"] + assert suite.bake_targets == ("candidate-structure",) + assert " structure " in f" {' '.join(suite.command)} " + + +def test_artifact_suite_builds_every_artifact_before_offline_probe() -> None: + suite = SUITES["artifact"] + assert "candidate" in suite.bake_targets + assert "candidate-artifact" in suite.bake_targets + assert any("tools.artifacts.build_all" in command for command in suite.pre_commands) + assert not any("kernels download" in " ".join(command) for command in suite.pre_commands) + assert "not test_local_artifact_locked_flash_backend" in " ".join(suite.command) + + +def test_release_suite_aggregates_exact_head_artifact_reference_and_gpu_gates() -> None: + suite = SUITES["release"] + commands = [" ".join(command) for command in suite.pre_commands] + joined = "\n".join(commands) + for target in ( + "candidate", + "candidate-structure", + "candidate-fp8", + "candidate-artifact", + "biohub-biotraj-wheel", + "reference-esmfold", + "reference-esmfold2", + ): + assert target in suite.bake_targets + assert "kernels download" not in joined + assert joined.index("tools.artifacts.build_all") < joined.index( + "tests/release/test_published_automodel.py" + ) + assert "tests.parity.support.native_reference" in joined + assert "tests.structure.support.esmfold2_bundle" in joined + assert "--precision bf16" in joined + assert "--precision fp8" not in joined + assert "tests.structure.support.boltz2_bundle" not in joined + assert "tools.remote.python_matrix" in joined + assert "artifacts/benchmarks/release-h100.json" in joined + assert "--artifact-root dist/hub" in joined + assert "--backends eager sdpa flex_attention" in joined + assert "--junit-output artifacts/junit/release-benchmark.xml" in joined + assert joined.index("tools.artifacts.build_all") < joined.index( + "artifacts/benchmarks/release-h100.json" + ) + command = " ".join(suite.command) + assert " structure " in f" {command} " + assert "reference-boltz2" not in suite.bake_targets + assert "--ignore=tests/structure/test_structure_models.py" in command + assert "--ignore=tests/structure/test_esmfold2_fp8_compliance.py" in command + assert "--ignore=tests/integration/test_flash_attention_backends.py" in command + fp8_stack_test = ( + "tests/release/test_validation_stack.py::" + "test_fp8_validation_stack_uses_the_cuda13_transformer_engine_core" + ) + assert f"--deselect={fp8_stack_test}" in command + assert "test_boltz2_live_folding_matches_pinned_official" in command + assert "test_esmfold2_isolated_bf16_and_fp8_folding_compliance" not in command + for path in ("tests/unit", "tests/integration", "tests/release", "tests/structure"): + assert path in command + assert "tests/parity/test_native_results.py" in command + for path in ( + "tests/parity/test_esmfold2_common_parity.py", + "tests/parity/test_esmfold2_protein_data_parity.py", + "tests/parity/test_esmfold2_reimplemented_source_parity.py", + "tests/parity/test_esmfold2_residue_config_parity.py", + "tests/parity/test_esmfold2_source_slice3_parity.py", + "tests/parity/test_esmfold2_source_slice4_parity.py", + ): + assert path in command + assert "tests/parity/test_model_parity.py" not in command + assert "tests/parity/test_ankh_seq2seq_parity.py" not in command + assert "tests/parity/test_e1_source_independence_parity.py" not in command + assert "not artifact" in command + + +def test_feature_suite_does_not_install_or_select_fp8() -> None: + suite = SUITES["feature"] + assert suite.bake_targets == ("candidate-structure",) + command = " ".join(suite.command) + assert " structure " in f" {command} " + assert " fp8 " not in f" {command} " + assert "tests/integration/test_dplm_generation.py" in command + assert "tests/integration/test_esm3.py" in command + assert "tests/release/test_conversion_tools.py" in command + + +def test_integration_suite_uses_the_structure_dependency_image() -> None: + suite = SUITES["integration"] + assert suite.bake_targets == ("candidate-structure",) + command = " ".join(suite.command) + assert " structure " in f" {command} " + assert "tests/integration" in command + + +def test_structure_suite_produces_isolated_folding_bundles_before_gating() -> None: + suite = SUITES["structure"] + commands = [" ".join(command) for command in suite.pre_commands] + joined = "\n".join(commands) + assert "reference-boltz2" in suite.bake_targets + assert "reference-esmfold" in suite.bake_targets + assert "reference-esmfold2" in suite.bake_targets + assert "tests.structure.support.boltz2_bundle prepare" in joined + assert "reference-boltz2 python -m tests.structure.support.boltz2_bundle" in joined + assert "tests.structure.support.boltz2_bundle produce-reference" in joined + assert "tests.structure.support.boltz2_bundle produce-candidate" in joined + assert "tests.structure.support.esmfold_bundle prepare" in joined + assert "reference-esmfold python -m tests.structure.support.esmfold_bundle" in joined + assert "produce-reference --exchange-root /exchange" in joined + assert "produce-candidate --exchange-root /workspace/artifacts/reference" in joined + assert "tests.structure.support.esmfold2_bundle prepare" in joined + suite_command = " ".join(suite.command) + assert "tests/structure" in suite_command + assert "tests/parity/test_boltz_source_refactor.py" in suite_command + assert "--ignore=tests/structure/test_structure_models.py" in suite_command + assert "-m structure" in suite_command + assert "structure and gpu" not in suite_command + esmfold2_commands = [ + command for command in commands if "tests.structure.support.esmfold2_bundle" in command + ] + assert any( + "reference-esmfold2" in command and "produce-reference" in command and "--all" in command + for command in esmfold2_commands + ) + assert any( + "produce-candidate" in command and "--all" in command and "--precision bf16" in command + for command in esmfold2_commands + ) + assert any( + "produce-candidate" in command and "--all" in command and "--precision fp8" in command + for command in esmfold2_commands + ) + runner_source = (Path(__file__).resolve().parents[2] / "tools" / "remote" / "run.py").read_text( + encoding="utf-8" + ) + assert "_ESMFOLD2_MODEL_IDS" not in runner_source diff --git a/tests/unit/test_semantic_config.py b/tests/unit/test_semantic_config.py new file mode 100644 index 0000000..8a31587 --- /dev/null +++ b/tests/unit/test_semantic_config.py @@ -0,0 +1,41 @@ +import pytest +import torch.nn as nn +from types import SimpleNamespace + +from tests.parity.support.semantic_config import SEMANTIC_PATHS, semantic_config + + +class _Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace( + vocab_size=64, + hidden_size=32, + num_hidden_layers=4, + num_attention_heads=8, + classifier_dropout=0.1, + initializer_range=0.02, + tie_word_embeddings=True, + ) + + +def test_shared_semantic_config_includes_esmc_checkpoint_fields() -> None: + assert {"classifier_dropout", "initializer_range", "tie_word_embeddings"}.issubset( + SEMANTIC_PATHS + ) + assert semantic_config(_Model()) == { + "vocab_size": 64, + "d_model": 32, + "n_layers": 4, + "n_heads": 8, + "classifier_dropout": 0.1, + "initializer_range": 0.02, + "tie_word_embeddings": True, + } + + +def test_shared_semantic_config_fails_closed_on_missing_required_field() -> None: + model = _Model() + del model.config.num_attention_heads + with pytest.raises(RuntimeError, match="n_heads"): + semantic_config(model) diff --git a/tests/unit/test_structure_output_contracts.py b/tests/unit/test_structure_output_contracts.py new file mode 100644 index 0000000..c10ed52 --- /dev/null +++ b/tests/unit/test_structure_output_contracts.py @@ -0,0 +1,513 @@ +"""Fast, injected-core contracts for advertised structure AutoModels.""" + +from __future__ import annotations + +import warnings +import pytest +import torch +from collections.abc import Iterable +from pathlib import Path +from typing import Any +from torch import Tensor, nn +from transformers.models.esm.modeling_esmfold import EsmForProteinFolding + +from fastplms.models.boltz import modeling_boltz2 +from fastplms.models.boltz.modeling_boltz2 import ( + Boltz2Config, + Boltz2Model, + Boltz2ModelOutput, +) +from fastplms.models.esmfold.modeling_fast_esmfold import ( + FastEsmFoldConfig, + FastEsmForProteinFolding, + FastEsmForProteinFoldingOutput, +) +from fastplms.models.esmfold2.configuration_esmfold2 import ESMFold2Config +from fastplms.models.esmfold2.modeling_esmfold2 import ( + ESMFold2Model, + ESMFold2Output, +) +from fastplms.models.esmfold2.modeling_esmfold2_common import NUM_RES_TYPES +from fastplms.models.esmfold2.modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel, +) +from fastplms.models.esmfold2.protein_utils import prepare_protein_features +from fastplms.models.esmfold2.reproducibility import seed_context + + +def _assert_nested_close(left: Any, right: Any) -> None: + if torch.is_tensor(left) and torch.is_tensor(right): + torch.testing.assert_close(left, right) + return + if isinstance(left, tuple) and isinstance(right, tuple): + assert len(left) == len(right) + for left_item, right_item in zip(left, right, strict=True): + _assert_nested_close(left_item, right_item) + return + assert left == right + + +def _assert_tuple_matches_output( + tuple_output: tuple[Any, ...], + structured_output: Iterable[Any], +) -> None: + expected = tuple(structured_output) + assert len(tuple_output) == len(expected) + for actual, reference in zip(tuple_output, expected, strict=True): + _assert_nested_close(actual, reference) + + +class _TinyBoltzCore(nn.Module): + def __init__(self, width: int = 3) -> None: + super().__init__() + self.weight = nn.Parameter(torch.linspace(0.5, 1.0, width)) # (d=width,) + + def forward(self, feats: dict[str, Tensor], **_kwargs: Any) -> dict[str, Tensor]: + # feats["signal"]: (b, l, d) + signal = feats["signal"] * self.weight # (b, l, d) + pair = signal[:, :, None, :] + signal[:, None, :, :] # (b, l, l, d) + return { + "pdistogram": pair.unsqueeze(-1), # (b, l, l, d, 1) + "s": signal, # (b, l, d) + "z": pair, # (b, l, l, d) + "sample_atom_coords": signal[..., :1].expand(-1, -1, 3), # (b, l, xyz=3) + } + + +def test_boltz_public_forward_honors_output_controls_backward_and_reload( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr(modeling_boltz2, "Boltz2InferenceCore", _TinyBoltzCore) + config = Boltz2Config(core_kwargs={"width": 3}) + model = Boltz2Model(config) + features = { + "signal": torch.arange(6, dtype=torch.float32).reshape(1, 2, 3) # (b=1, l=2, d=3) + } + + structured = model( + feats=features, + output_hidden_states=True, + return_dict=True, + ) + tuple_output = model( + feats=features, + output_hidden_states=True, + return_dict=False, + ) + + assert isinstance(structured, Boltz2ModelOutput) + assert structured.last_hidden_state is structured.s + assert structured.hidden_states is not None + assert structured.hidden_states[0] is structured.s + assert structured.hidden_states[1] is structured.z + _assert_tuple_matches_output(tuple_output, structured.to_tuple()) + assert structured.last_hidden_state is not None + structured.last_hidden_state.square().mean().backward() + assert model.core.weight.grad is not None + assert torch.isfinite(model.core.weight.grad).all() + with pytest.raises(NotImplementedError, match="output_attentions=True"): + model(feats=features, output_attentions=True) + with pytest.raises(TypeError): + model(feats=features, silently_ignored=True) + + model.save_pretrained(tmp_path) + reloaded = Boltz2Model.from_pretrained(tmp_path, local_files_only=True) + reloaded_output = reloaded(feats=features, return_dict=True) + torch.testing.assert_close( + reloaded_output.last_hidden_state, + structured.last_hidden_state, + ) + + +def test_fast_esmfold_public_forward_honors_output_controls_and_backward( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def native_forward( + self: FastEsmForProteinFolding, + input_ids: Tensor, + **_kwargs: Any, + ) -> dict[str, Tensor]: + # input_ids: (b, l) + state = self.contract_weight.expand(input_ids.shape[0], input_ids.shape[1], 2) + return {"s_s": state, "plddt": state[..., :1]} # (b, l, d=2); (b, l, 1) + + monkeypatch.setattr(EsmForProteinFolding, "forward", native_forward) + model = FastEsmForProteinFolding.__new__(FastEsmForProteinFolding) + nn.Module.__init__(model) + model.register_parameter( + "contract_weight", + nn.Parameter(torch.ones(1, 1, 2)), # (1, 1, d=2) + ) + input_ids = torch.zeros((1, 2), dtype=torch.int64) # (b=1, l=2) + + structured = model( + input_ids, + output_hidden_states=True, + return_dict=True, + ) + tuple_output = model( + input_ids, + output_hidden_states=True, + return_dict=False, + ) + + assert isinstance(structured, FastEsmForProteinFoldingOutput) + assert structured.last_hidden_state is structured.s_s + assert structured.hidden_states is not None + assert structured.hidden_states[0] is structured.s_s + assert torch.equal(structured.plddt, torch.full((1, 2, 1), 100.0)) + _assert_tuple_matches_output(tuple_output, structured.to_tuple()) + assert structured.last_hidden_state is not None + structured.last_hidden_state.square().mean().backward() + assert model.contract_weight.grad is not None + attention_output = model(input_ids, output_attentions=True) + assert attention_output.attentions == () + with pytest.raises(TypeError): + model(input_ids, silently_ignored=True) + + +def _tiny_fast_esmfold_config( + *, + bypass_lm: bool, + attn_backend: str = "sdpa", +) -> FastEsmFoldConfig: + return FastEsmFoldConfig( + vocab_size=33, + hidden_size=8, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=16, + max_position_embeddings=16, + pad_token_id=1, + mask_token_id=32, + position_embedding_type="rotary", + is_folding_model=True, + attn_backend=attn_backend, + esmfold_config={ + "fp16_esm": False, + "bypass_lm": bypass_lm, + "lddt_head_hid_dim": 4, + "trunk": { + "num_blocks": 1, + "sequence_state_dim": 8, + "pairwise_state_dim": 4, + "sequence_head_width": 4, + "pairwise_head_width": 2, + "position_bins": 4, + "max_recycles": 1, + "chunk_size": None, + "structure_module": { + "sequence_dim": 8, + "pairwise_dim": 4, + "ipa_dim": 2, + "resnet_dim": 4, + "num_heads_ipa": 2, + "num_qk_points": 1, + "num_v_points": 1, + "dropout_rate": 0.0, + "num_blocks": 1, + "num_transition_layers": 1, + "num_resnet_blocks": 1, + "num_angles": 7, + }, + }, + }, + ) + + +def test_fast_esmfold_output_attentions_uses_masked_per_call_eager_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def folding_stub( + self: FastEsmForProteinFolding, + input_ids: Tensor, + attention_mask: Tensor | None = None, + **_kwargs: Any, + ) -> dict[str, Tensor]: + # input_ids: (b, l); attention_mask: (b, l) or None + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) # (b, l) + esmaa = self.af2_idx_to_esm_idx(input_ids, attention_mask) # (b, l) + representations = self.compute_language_model_representations( + esmaa + ) # (b, l, n_layers, d) + state = representations[:, :, -1, :] # (b, l, d) + return { + "s_s": state, # (b, l, d) + "plddt": torch.ones((*state.shape[:2], 1), device=state.device), # (b, l, 1) + } + + monkeypatch.setattr(EsmForProteinFolding, "forward", folding_stub) + model = FastEsmForProteinFolding( + _tiny_fast_esmfold_config(bypass_lm=False, attn_backend="sdpa") + ).eval() + input_ids = torch.tensor(((0, 1, 2), (3, 4, 0)), dtype=torch.int64) # (b=2, l=3) + attention_mask = torch.tensor( # (b=2, l=3) + ((1, 1, 1), (1, 1, 0)), dtype=torch.int64 + ) + configured_encoder_backend = model.esm.encoder.attention_backend + attention_module = model.esm.encoder.layer[0].attention.self + configured_layer_backend = attention_module.attn_backend + + with pytest.warns( + RuntimeWarning, + match=r"output_attentions=True.*requested 'sdpa'.*using 'eager'.*call only", + ) as captured: + output = model( + input_ids, + attention_mask=attention_mask, + output_attentions=True, + ) + + assert len(captured) == 1 + assert output.attentions is not None and len(output.attentions) == 1 + attention = output.attentions[0] # (b=2, h=2, l=3, l=3) + assert attention.shape == (2, 2, 3, 3) + torch.testing.assert_close( + attention[1, :, :, 2], + torch.zeros_like(attention[1, :, :, 2]), + rtol=0.0, + atol=0.0, + ) + assert model.config.attn_backend == "sdpa" + assert model.esm.encoder.attention_backend == configured_encoder_backend + assert attention_module.attn_backend == configured_layer_backend + + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + subsequent = model( + input_ids, + attention_mask=attention_mask, + output_attentions=False, + ) + assert subsequent.attentions is None + assert attention_module.attn_backend == configured_layer_backend + + +def test_fast_esmfold_tiny_model_saves_and_reloads_exact_state(tmp_path: Path) -> None: + config = _tiny_fast_esmfold_config(bypass_lm=True) + source = FastEsmForProteinFolding(config).eval() + source.save_pretrained(tmp_path, safe_serialization=True) + + restored = FastEsmForProteinFolding.from_pretrained( + tmp_path, + local_files_only=True, + ).eval() + assert set(restored.state_dict()) == set(source.state_dict()) + for name, tensor in source.state_dict().items(): + torch.testing.assert_close( + restored.state_dict()[name], + tensor, + rtol=0.0, + atol=0.0, + ) + + +def _tiny_esmfold2_config(model_type: str) -> ESMFold2Config: + atom_token_width = 8 + input_feature_width = atom_token_width // 2 + 2 * NUM_RES_TYPES + 1 + return ESMFold2Config( + type=model_type, + d_single=8, + d_pair=8, + num_loops=0, + num_diffusion_samples=1, + lm_d_model=8, + lm_num_layers=1, + inputs={ + "d_inputs": input_feature_width, + "atom_encoder": { + "d_atom": 8, + "d_token": atom_token_width, + "n_blocks": 0, + "n_heads": 2, + "swa_window_size": 32, + "expansion_ratio": 2, + "n_spatial_rope_pairs_per_axis": 1, + "n_uid_rope_pairs": 1, + }, + }, + folding_trunk={"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + structure_head={ + "diffusion_module": { + "c_atom": 8, + "c_token": 8, + "c_z": 8, + "c_s_inputs": input_feature_width, + "fourier_dim": 8, + "atom_num_blocks": 0, + "atom_num_heads": 2, + "token_num_blocks": 0, + "token_num_heads": 2, + "transition_multiplier": 2, + }, + "distogram_bins": 8, + "inference_num_steps": 1, + }, + confidence_head={ + "enabled": False, + "folding_trunk": {"n_layers": 0, "n_heads": 2, "dropout": 0.0}, + "num_plddt_bins": 4, + "num_pde_bins": 4, + "num_pae_bins": 4, + "distogram_bins": 8, + }, + msa_encoder={ + "enabled": True, + "d_msa": 8, + "d_hidden": 4, + "n_layers": 0, + "n_heads_msa": 2, + "msa_head_width": 4, + }, + msa_conditioning=True, + lm_encoder={"enabled": False, "n_layers": 0}, + parcae={"enabled": True, "min_steps": 1, "max_steps": 1, "coda_n_layers": 0}, + ) + + +class _TinyStructureHead(nn.Module): + def __init__(self) -> None: + super().__init__() + self.observed: dict[str, Any] = {} + + def sample(self, **kwargs: Any) -> dict[str, Tensor]: + self.observed = { + name: kwargs[name] + for name in ( + "noise_scale", + "step_scale", + "max_inference_sigma", + "denoising_early_exit_rmsd", + ) + } + coords = kwargs["ref_pos"].float() # (b, a, xyz=3) + multiplicity = int(kwargs["num_diffusion_samples"]) + return { + "sample_atom_coords": coords.repeat_interleave( + multiplicity, dim=0 + ) # (b * multiplicity, a, xyz=3) + } + + +class _TinyConfidenceHead(nn.Module): + def forward( + self, + z: Tensor, + x_pred: Tensor, + num_diffusion_samples: int, + **_kwargs: Any, + ) -> dict[str, Tensor]: + # z: (b, l, l, d_pair); x_pred: (b * num_diffusion_samples, a, xyz=3) + _batch_size, sequence_length = z.shape[:2] + score = z.float().mean(dim=(1, 2, 3)).repeat_interleave( + num_diffusion_samples + ) # (b * num_diffusion_samples,) + return { + "plddt": score[:, None].expand(-1, sequence_length), # (b * samples, l) + "complex_plddt": score, # (b * samples,) + "ptm": score, # (b * samples,) + "iptm": score, # (b * samples,) + "pae": z.new_zeros( + (x_pred.shape[0], sequence_length, sequence_length) + ), # (b * samples, l, l) + } + + +@pytest.mark.parametrize( + ("model_class", "model_type"), + ( + (ESMFold2Model, "release"), + (ESMFold2ExperimentalModel, "experimental"), + ), +) +def test_esmfold2_public_forward_honors_output_controls_and_sampler_overrides( + model_class: type[ESMFold2Model] | type[ESMFold2ExperimentalModel], + model_type: str, +) -> None: + model = model_class(_tiny_esmfold2_config(model_type)).eval() + structure_head = _TinyStructureHead() + model.structure_head = structure_head + if model_type == "release": + model.confidence_head = _TinyConfidenceHead() + features = prepare_protein_features("AC") # batched tensors for b=1, l=2 residues + common_kwargs = { + "num_loops": 0, + "num_sampling_steps": 1, + "num_diffusion_samples": 1, + "noise_scale": 0.25, + "step_scale": 1.5, + "max_inference_sigma": 32.0, + "early_exit": True, + } + if model_type == "experimental": + common_kwargs.update({"calculate_confidence": False, "seed": 7}) + else: + common_kwargs.update( + {"msa_column_mask_rate": 0.0, "msa_subsample_at_inference": False} + ) + + with seed_context(31): + structured = model( + **features, + **common_kwargs, + output_hidden_states=True, + return_dict=True, + ) + with seed_context(31): + tuple_output = model( + **features, + **common_kwargs, + output_hidden_states=True, + return_dict=False, + ) + + assert isinstance(structured, ESMFold2Output) + assert structured.last_hidden_state is not None + assert structured.hidden_states is not None + assert structured.hidden_states[-1] is structured.last_hidden_state + assert structured.sample_atom_coords is not None + _assert_tuple_matches_output(tuple_output, structured.to_tuple()) + assert structure_head.observed == { + "noise_scale": 0.25, + "step_scale": 1.5, + "max_inference_sigma": 32.0, + "denoising_early_exit_rmsd": 0.10, + } + with pytest.raises(NotImplementedError, match="output_attentions=True"): + model(**features, output_attentions=True) + with pytest.raises(TypeError): + model(**features, silently_ignored=True) + + model.zero_grad(set_to_none=True) + if model_type == "experimental": + res_type_soft = torch.nn.functional.one_hot( # (b=1, l=2, c=NUM_RES_TYPES) + features["res_type"].long(), + num_classes=NUM_RES_TYPES, + ).float() + res_type_soft.requires_grad_(True) + differentiable = model( + **features, + **common_kwargs, + res_type_soft=res_type_soft, + return_dict=True, + ) + else: + differentiable = model( + **features, + **common_kwargs, + return_dict=True, + ) + assert differentiable.distogram_logits is not None + differentiable.distogram_logits.square().mean().backward() + if model_type == "experimental": + assert res_type_soft.grad is not None + assert torch.isfinite(res_type_soft.grad).all() + gradients = [ + parameter.grad + for parameter in model.parameters() + if parameter.requires_grad and parameter.grad is not None + ] + assert gradients + assert all(torch.isfinite(gradient).all() for gradient in gradients) diff --git a/tests/unit/test_structure_state_contract.py b/tests/unit/test_structure_state_contract.py new file mode 100644 index 0000000..e165303 --- /dev/null +++ b/tests/unit/test_structure_state_contract.py @@ -0,0 +1,107 @@ +"""Unit contracts for compact exact structure-checkpoint metadata.""" + +from __future__ import annotations + +import copy +import pytest +import torch +import torch.nn as nn +from collections.abc import Callable + +from tests.structure.support.state_contract import ( + exact_state_contract, + semantic_config_contract, + tensor_sha256, + validate_exact_state_contract, + validate_semantic_config_contract, +) + + +class SharedParameterModel(nn.Module): + def __init__(self) -> None: + super().__init__() + parameter = nn.Parameter(torch.tensor([[1.0, 2.0]], dtype=torch.float32)) # (1, 2) + self.left = parameter + self.right = parameter + self.register_buffer("counter", torch.tensor(3, dtype=torch.int64)) # () + + +def test_exact_state_contract_covers_scalars_hashes_and_aliases() -> None: + model = SharedParameterModel() + + contract = exact_state_contract(model) + + assert contract["tensors"]["counter"] == { + "dtype": "int64", + "shape": [], + "sha256": tensor_sha256(model.counter), + } + assert contract["aliases"] == [["left", "right"]] + validate_exact_state_contract(contract) + + +def test_exact_state_contract_applies_names_and_exclusions() -> None: + model = SharedParameterModel() + + contract = exact_state_contract( + model, + name_transform=lambda name: (f"canonical.{name}",), + excluded_prefixes=("counter",), + ) + + assert set(contract["tensors"]) == {"canonical.left", "canonical.right"} + assert contract["aliases"] == [["canonical.left", "canonical.right"]] + + +@pytest.mark.parametrize( + "mutation", + [ + lambda contract: contract.__setitem__("sha256", "0" * 64), + lambda contract: contract["tensors"]["counter"].__setitem__("shape", [-1]), + lambda contract: contract["tensors"]["counter"].__setitem__("sha256", "not-a-hash"), + lambda contract: contract["aliases"].append(["left", "missing"]), + ], +) +def test_exact_state_contract_rejects_corruption( + mutation: Callable[[dict[str, object]], None], +) -> None: + contract = exact_state_contract(SharedParameterModel()) + mutation(contract) + + with pytest.raises(ValueError, match="Structure state"): + validate_exact_state_contract(contract) + + +def test_semantic_config_contract_removes_packaging_fields_recursively() -> None: + contract = semantic_config_contract( + { + "hidden_size": 8, + "fastplms_model_id": "toy", + "nested": { + "architectures": ["PackagingOnly"], + "depth": 2, + "fastplms_checkpoint_hash": "a" * 64, + }, + "dtype": torch.bfloat16, + } + ) + + assert contract["fields"] == { + "hidden_size": 8, + "nested": {"depth": 2}, + } + validate_semantic_config_contract(contract) + + +def test_semantic_config_contract_rejects_corruption() -> None: + contract = semantic_config_contract({"hidden_size": 8}) + changed = copy.deepcopy(contract) + changed["fields"]["hidden_size"] = 16 + + with pytest.raises(ValueError, match="digest mismatch"): + validate_semantic_config_contract(changed) + + changed = copy.deepcopy(contract) + changed["fields"]["auto_map"] = {"AutoModel": "remote.Model"} + with pytest.raises(ValueError, match="packaging fields"): + validate_semantic_config_contract(changed) diff --git a/tests/unit/test_tokenizer_contract.py b/tests/unit/test_tokenizer_contract.py new file mode 100644 index 0000000..036fd42 --- /dev/null +++ b/tests/unit/test_tokenizer_contract.py @@ -0,0 +1,755 @@ +"""Tokenizer contract tests for all FastPLMs sequence checkpoints.""" + +from __future__ import annotations + +import json +import shutil +import pytest +import torch +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch +from transformers import AutoModelForMaskedLM, AutoTokenizer, EsmTokenizer, PretrainedConfig + +from fastplms.models.ankh.modeling_ankh import ( + FastAnkhConfig, + _load_ankh_tokenizer, +) +from fastplms.models.dplm.modeling_dplm import ( + DPLMConfig, + DPLMForMaskedLM, + DPLMPreTrainedModel, +) +from fastplms.models.dplm2.modeling_dplm2 import ( + DPLM2Config, + DPLM2ForMaskedLM, + DPLM2PreTrainedModel, + _normalize_dplm2_input_ids, +) +from fastplms.models.dplm2.tokenization_dplm2 import DPLM2Tokenizer +from fastplms.models.e1.modeling_e1 import E1BatchPreparer, E1Config, E1ForMaskedLM, get_tokenizer +from fastplms.models.esm2.modeling_fastesm import ( + FastEsmConfig, + FastEsmForMaskedLM, + FastEsmPreTrainedModel, + FastEsmTokenizer, +) +from fastplms.models.esm3.modeling_esm3 import ( + SEQUENCE_VOCAB as ESM3_SEQUENCE_VOCAB, +) +from fastplms.models.esm3.modeling_esm3 import ( + EsmSequenceTokenizer as ESM3SequenceTokenizer, +) +from fastplms.models.esm_plusplus.modeling_esm_plusplus import EsmSequenceTokenizer +from tests.conftest import CANONICAL_AAS, FULL_MODEL_REGISTRY, mark_by_size + + +TOKENIZER_REFERENCE_KEYS = [ + key + for key, value in FULL_MODEL_REGISTRY.items() + if value["uses_tokenizer"] and value["model_type"] not in {"DPLM2", "ESM3"} +] +ESM3_MODEL_KEYS = [ + key for key, value in FULL_MODEL_REGISTRY.items() if value["model_type"] == "ESM3" +] +DPLM2_MODEL_KEYS = [ + key for key, value in FULL_MODEL_REGISTRY.items() if value["model_type"] == "DPLM2" +] + + +def test_dplm2_model_tokenizer_uses_checkpoint_provenance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + expected_tokenizer = object() + requests: list[tuple[str, dict[str, str]]] = [] + + def load_tokenizer(source: str, **kwargs: str) -> object: + requests.append((source, kwargs)) + return expected_tokenizer + + monkeypatch.setattr(DPLM2Tokenizer, "from_pretrained", staticmethod(load_tokenizer)) + owner = SimpleNamespace( + _fastplms_tokenizer=None, + config=SimpleNamespace( + _name_or_path="local-dplm2-artifact", + _commit_hash="b" * 40, + ), + ) + + actual = DPLM2PreTrainedModel.tokenizer.fget(owner) + + assert actual is expected_tokenizer + assert owner.__dict__["_fastplms_tokenizer"] is expected_tokenizer + assert requests == [("local-dplm2-artifact", {"revision": "b" * 40})] + + +def test_dplm2_model_tokenizer_rejects_missing_checkpoint_provenance() -> None: + owner = SimpleNamespace( + _fastplms_tokenizer=None, + config=SimpleNamespace(_name_or_path="", _commit_hash=None), + ) + + with pytest.raises(RuntimeError, match="loaded with from_pretrained"): + DPLM2PreTrainedModel.tokenizer.fget(owner) + + +def test_dplm2_sequence_adapter_adds_amino_acid_boundaries_without_generic_specials() -> None: + observed: dict[str, object] = {} + + class RecordingTokenizer: + aa_cls_token = "" + aa_eos_token = "" + + def __call__(self, sequences: list[str], **kwargs: object) -> dict[str, object]: + observed.update(sequences=sequences, kwargs=kwargs) + return {"input_ids": sequences} + + tokenizer = RecordingTokenizer() + owner = SimpleNamespace(tokenizer=tokenizer) + + encoded = DPLM2PreTrainedModel._tokenize_sequence_batch( + owner, + ["AC", "M"], + tokenizer=tokenizer, + return_tensors="pt", + padding=True, + ) + + assert encoded == {"input_ids": ["AC", "M"]} + assert observed == { + "sequences": ["AC", "M"], + "kwargs": { + "add_special_tokens": False, + "return_tensors": "pt", + "padding": True, + }, + } + + +@pytest.mark.parametrize( + ("model_class", "model_name"), + ( + (DPLMPreTrainedModel, "DPLM"), + (FastEsmPreTrainedModel, "ESM2"), + ), +) +def test_esm_tokenizer_loaders_use_checkpoint_provenance( + monkeypatch: pytest.MonkeyPatch, + model_class: type, + model_name: str, +) -> None: + expected_tokenizer = object() + requests: list[tuple[str, dict[str, str]]] = [] + + def load_tokenizer(source: str, **kwargs: str) -> object: + requests.append((source, kwargs)) + return expected_tokenizer + + monkeypatch.setattr(EsmTokenizer, "from_pretrained", staticmethod(load_tokenizer)) + owner = SimpleNamespace( + _fastplms_tokenizer=None, + config=SimpleNamespace( + _name_or_path=f"local-{model_name.lower()}-artifact", + _commit_hash="c" * 40, + ), + ) + + actual = model_class.tokenizer.fget(owner) + + assert actual is expected_tokenizer + assert requests == [(f"local-{model_name.lower()}-artifact", {"revision": "c" * 40})] + + +def test_esm2_tokenizer_normalizes_cls_as_bos( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tokenizer = SimpleNamespace(bos_token_id=None, cls_token="") + monkeypatch.setattr( + EsmTokenizer, + "from_pretrained", + staticmethod(lambda *_args, **_kwargs: tokenizer), + ) + owner = SimpleNamespace( + _fastplms_tokenizer=None, + config=SimpleNamespace(_name_or_path="local-esm2-artifact", _commit_hash=None), + ) + + actual = FastEsmPreTrainedModel.tokenizer.fget(owner) + + assert actual.bos_token == "" + + +def test_esm2_tokenizer_rejects_residues_outside_official_alphabet() -> None: + tokenizer = object.__new__(FastEsmTokenizer) + tokenizer._token_to_id = {"A": 5} + + assert tokenizer._convert_token_to_id("A") == 5 + with pytest.raises(KeyError, match="J"): + tokenizer._convert_token_to_id("J") + + +@pytest.mark.parametrize( + ("model_class", "model_name"), + ( + (DPLMPreTrainedModel, "DPLM"), + (FastEsmPreTrainedModel, "ESM2"), + ), +) +def test_esm_tokenizer_loaders_reject_missing_checkpoint_provenance( + model_class: type, + model_name: str, +) -> None: + owner = SimpleNamespace( + _fastplms_tokenizer=None, + config=SimpleNamespace(_name_or_path="", _commit_hash=None), + ) + + with pytest.raises(RuntimeError, match=rf"{model_name} tokenizer loading requires"): + model_class.tokenizer.fget(owner) + + +def _tiny_esm_family_config( + config_class: type[PretrainedConfig], + vocab_size: int, +) -> PretrainedConfig: + kwargs = { + "vocab_size": vocab_size, + "hidden_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "intermediate_size": 16, + "hidden_dropout_prob": 0.0, + "attention_probs_dropout_prob": 0.0, + "max_position_embeddings": 16, + "pad_token_id": 1, + "mask_token_id": min(32, vocab_size - 1), + "attn_backend": "sdpa", + } + if config_class is FastEsmConfig: + kwargs["position_embedding_type"] = "absolute" + kwargs["attn_backend"] = "eager" + return config_class(**kwargs) + + +@pytest.mark.parametrize( + ("model_class", "config_class", "tokenizer_class", "vocab_size"), + ( + (DPLMForMaskedLM, DPLMConfig, EsmTokenizer, 33), + (DPLM2ForMaskedLM, DPLM2Config, DPLM2Tokenizer, 64), + (FastEsmForMaskedLM, FastEsmConfig, FastEsmTokenizer, 33), + ), +) +def test_tokenizer_context_survives_loading_info_and_concurrent_lazy_loads( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + model_class: type, + config_class: type, + tokenizer_class: type, + vocab_size: int, +) -> None: + roots = [tmp_path / "first", tmp_path / "second"] + for root in roots: + model = model_class(_tiny_esm_family_config(config_class, vocab_size)).eval() + model.save_pretrained(root / "nested") + + loaded_models = [] + loading_infos = [] + for index, root in enumerate(roots): + loaded, loading_info = model_class.from_pretrained( + root, + subfolder="nested", + revision=f"requested-{index}", + cache_dir=tmp_path / f"cache-{index}", + local_files_only=True, + force_download=False, + trust_remote_code=False, + token=f"secret-{index}", + output_loading_info=True, + ) + loaded.config._commit_hash = str(index) * 40 + loaded_models.append(loaded) + loading_infos.append(loading_info) + + assert all(not info["missing_keys"] for info in loading_infos) + assert all(not info["unexpected_keys"] for info in loading_infos) + assert all( + "secret-" not in json.dumps(model.config.to_dict(), default=str) + for model in loaded_models + ) + + requests: list[tuple[object, dict[str, object]]] = [] + rendezvous = Barrier(2) + + def load_tokenizer(source: object, **kwargs: object) -> object: + requests.append((source, kwargs)) + rendezvous.wait(timeout=5) + return SimpleNamespace(bos_token_id=0, cls_token="") + + monkeypatch.setattr( + tokenizer_class, + "from_pretrained", + staticmethod(load_tokenizer), + ) + with ThreadPoolExecutor(max_workers=2) as executor: + tokenizers = list(executor.map(lambda model: model.tokenizer, loaded_models)) + + assert len(tokenizers) == 2 + observed = {str(source): kwargs for source, kwargs in requests} + for index, root in enumerate(roots): + assert observed[str(root)] == { + "cache_dir": tmp_path / f"cache-{index}", + "force_download": False, + "local_files_only": True, + "revision": str(index) * 40, + "subfolder": "nested", + "token": f"secret-{index}", + "trust_remote_code": False, + } + + +def test_ankh_tokenizer_loader_uses_checkpoint_provenance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Tokenizer: + backend_tokenizer = SimpleNamespace(pre_tokenizer=None) + + expected_tokenizer = Tokenizer() + requests: list[tuple[str, dict[str, str]]] = [] + + def load_tokenizer(source: str, **kwargs: str) -> object: + requests.append((source, kwargs)) + return expected_tokenizer + + monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(load_tokenizer)) + config = _tiny_ankh_config("local-ankh-artifact", "d" * 40) + + assert _load_ankh_tokenizer(config) is expected_tokenizer + assert requests == [("local-ankh-artifact", {"revision": "d" * 40})] + assert expected_tokenizer.backend_tokenizer.pre_tokenizer is not None + + +CANONICAL_SEQUENCES = [ + "M" + CANONICAL_AAS, + "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSH", + "MXXBZUOACDEFGHIKLMNPQRSTVWY", +] + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[2] + + +def _e1_tokenizer_json() -> Path: + return _repo_root() / "src" / "fastplms" / "models" / "e1" / "tokenizer.json" + + +def _tiny_ankh_config( + name_or_path: str = "", + revision: str | None = None, +) -> FastAnkhConfig: + config = FastAnkhConfig( + vocab_size=4, + d_model=8, + d_kv=4, + d_ff=16, + num_heads=2, + num_layers=1, + ) + config._name_or_path = name_or_path + config._commit_hash = revision + return config + + +def _fast_tokenizer(config: dict[str, Any]) -> Any: + if config["model_type"] == "ANKH": + # ANKH artifacts are built from the manifest's pinned official source, + # including its byte-exact tokenizer assets. + return _load_ankh_tokenizer( + _tiny_ankh_config(config["official_path"], config["official_revision"]) + ) + if config["model_type"] == "ESMC": + return EsmSequenceTokenizer() + if config["model_type"] in ("ESM2", "DPLM"): + return EsmTokenizer.from_pretrained( + config["fast_path"], + revision=config["fast_revision"], + ) + return AutoTokenizer.from_pretrained( + config["fast_path"], + revision=config["fast_revision"], + trust_remote_code=True, + ) + + +def _reference_tokenizer(config: dict[str, Any]) -> Any: + if config["model_type"] == "ANKH": + return _load_ankh_tokenizer( + _tiny_ankh_config(config["official_path"], config["official_revision"]) + ) + if config["model_type"] == "ESMC": + return EsmSequenceTokenizer() + if config["model_type"] in ("ESM2", "DPLM"): + return EsmTokenizer.from_pretrained( + config["official_path"], + revision=config["official_revision"], + ) + return AutoTokenizer.from_pretrained( + config["official_path"], + revision=config["official_revision"], + trust_remote_code=True, + ) + + +def _token_ids(tokenizer: Any, sequence: str) -> torch.Tensor: + encoded = tokenizer( + sequence, + return_tensors="pt", + ) + return encoded["input_ids"] # (b=1, l) + + +def _special_token_ids(tokenizer: Any) -> dict[str, int | None]: + return { + "pad_token_id": tokenizer.pad_token_id, + "cls_token_id": tokenizer.cls_token_id, + "eos_token_id": tokenizer.eos_token_id, + "mask_token_id": tokenizer.mask_token_id, + "unk_token_id": tokenizer.unk_token_id, + } + + +@pytest.mark.parametrize( + "model_key", + mark_by_size(TOKENIZER_REFERENCE_KEYS, FULL_MODEL_REGISTRY), +) +def test_sequence_tokenizer_matches_reference(model_key: str) -> None: + config = FULL_MODEL_REGISTRY[model_key] + fast_tok = _fast_tokenizer(config) + reference_tok = _reference_tokenizer(config) + + fast_vocab = fast_tok.get_vocab() + reference_vocab = reference_tok.get_vocab() + assert len(fast_vocab) == len(reference_vocab), ( + f"{model_key}: vocab size mismatch fast={len(fast_vocab)} reference={len(reference_vocab)}" + ) + + missing_in_fast = [token for token in reference_vocab if token not in fast_vocab] + assert not missing_in_fast, ( + f"{model_key}: tokens missing from fast tokenizer: {missing_in_fast[:5]}" + ) + + id_mismatches = [ + (token, reference_vocab[token], fast_vocab[token]) + for token in reference_vocab + if reference_vocab[token] != fast_vocab[token] + ] + assert not id_mismatches, f"{model_key}: token id mismatches: {id_mismatches[:5]}" + + assert _special_token_ids(fast_tok) == _special_token_ids(reference_tok), ( + f"{model_key}: special token ids differ" + ) + + for sequence in CANONICAL_SEQUENCES: + fast_ids = _token_ids(fast_tok, sequence) + reference_ids = _token_ids(reference_tok, sequence) + assert torch.equal(fast_ids, reference_ids), ( + f"{model_key}: encoded ids differ for {sequence[:16]} " + f"fast={fast_ids[0, :8].tolist()} " + f"reference={reference_ids[0, :8].tolist()}" + ) + + +@pytest.mark.parametrize( + "model_key", + mark_by_size(ESM3_MODEL_KEYS, FULL_MODEL_REGISTRY), +) +def test_esm3_sequence_tokenizer_contract(model_key: str) -> None: + tokenizer = ESM3SequenceTokenizer() + expected_vocab = {token: token_id for token_id, token in enumerate(ESM3_SEQUENCE_VOCAB)} + + assert tokenizer.get_vocab() == expected_vocab, f"{model_key}: ESM3 sequence vocabulary changed" + assert _special_token_ids(tokenizer) == { + "pad_token_id": 1, + "cls_token_id": 0, + "eos_token_id": 2, + "mask_token_id": 32, + "unk_token_id": 3, + } + + for sequence in CANONICAL_SEQUENCES: + encoded = _token_ids(tokenizer, sequence) + expected_ids = [0] + [expected_vocab[token] for token in sequence] + [2] + assert encoded[0].tolist() == expected_ids, ( + f"{model_key}: encoded ids differ for {sequence[:16]}" + ) + + +def test_ankh_tokenizer_loader_rejects_missing_checkpoint_provenance() -> None: + with pytest.raises(RuntimeError, match="ANKH tokenizer loading requires"): + _load_ankh_tokenizer(_tiny_ankh_config()) + + +@pytest.mark.parametrize( + "model_key", + mark_by_size(DPLM2_MODEL_KEYS, FULL_MODEL_REGISTRY), +) +def test_dplm2_tokenizer_special_ids_normalize_in_range(model_key: str) -> None: + config = FULL_MODEL_REGISTRY[model_key] + fast_config = DPLM2Config.from_pretrained( + config["fast_path"], + revision=config["fast_revision"], + ) + tokenizer = DPLM2Tokenizer.from_pretrained( + config["fast_path"], + revision=config["fast_revision"], + ) + + generic_special_ids = torch.tensor( + [ + [ + fast_config.vocab_size, + fast_config.vocab_size + 1, + fast_config.vocab_size + 2, + fast_config.vocab_size + 3, + -100, + ] + ] + ) + expected = torch.tensor([[2, 3, 0, 32, -100]]) + normalized_special_ids = _normalize_dplm2_input_ids( + generic_special_ids, + vocab_size=fast_config.vocab_size, + ) + assert torch.equal(normalized_special_ids, expected) + + aa_sequences = [ + f"{tokenizer.aa_cls_token}{sequence}{tokenizer.aa_eos_token}" + for sequence in CANONICAL_SEQUENCES + ] + encoded = tokenizer( + aa_sequences, + add_special_tokens=False, + return_tensors="pt", + padding=True, + ) + normalized_input_ids = _normalize_dplm2_input_ids( + encoded["input_ids"], + vocab_size=fast_config.vocab_size, + ) + valid_ids = normalized_input_ids[normalized_input_ids.ge(0)] + assert bool(valid_ids.lt(fast_config.vocab_size).all()) + + +def test_dplm2_tokenizer_preserves_multimodal_contract(tmp_path: Path) -> None: + model_key = DPLM2_MODEL_KEYS[0] + config = FULL_MODEL_REGISTRY[model_key] + tokenizer = DPLM2Tokenizer.from_pretrained( + config["official_path"], + revision=config["official_revision"], + ) + + expected_tokens = { + "": 0, + "": 1, + "": 2, + "": 3, + "": 32, + "": 33, + "": 34, + "": 35, + "0000": 36, + "8191": 8227, + "": 8229, + } + vocabulary = tokenizer.get_vocab() + assert {token: vocabulary[token] for token in expected_tokens} == expected_tokens + assert tokenizer.vocab_size == 8229 + assert len(tokenizer) == 8229 + assert _special_token_ids(tokenizer) == { + "pad_token_id": 1, + "cls_token_id": None, + "eos_token_id": None, + "mask_token_id": None, + "unk_token_id": None, + } + assert tokenizer.all_special_ids == [0, 2, 3, 32, 33, 34, 35, 8229, 1] + + aa_track = tokenizer( + "AC", + add_special_tokens=False, + )["input_ids"] + structure_track = tokenizer( + "00000042", + add_special_tokens=False, + )["input_ids"] + assert aa_track == [0, 5, 23, 2] + assert structure_track == [33, 36, 78, 34] + with pytest.raises(ValueError): + tokenizer("AC") + + tokenizer.save_pretrained(tmp_path) + saved_config = json.loads((tmp_path / "tokenizer_config.json").read_text(encoding="utf-8")) + for name in DPLM2Tokenizer.SPECIAL_TOKENS_ATTRIBUTES: + assert saved_config[name] == getattr(tokenizer, name) + reloaded = DPLM2Tokenizer.from_pretrained(tmp_path, local_files_only=True) + assert reloaded.get_vocab() == vocabulary + assert reloaded.all_special_ids == tokenizer.all_special_ids + assert reloaded("AC", add_special_tokens=False) == tokenizer( + "AC", + add_special_tokens=False, + ) + + +def test_e1_config_uses_static_token_constants() -> None: + with patch( + "fastplms.models.e1.modeling_e1.get_tokenizer", + side_effect=AssertionError("E1Config should not load tokenizer.json"), + ): + config = E1Config() + + assert config.vocab_size == 34 + assert config.pad_token_id == 0 + assert config.bos_token_id == 1 + assert config.eos_token_id == 2 + + +def test_e1_sequence_preparer_uses_model_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observed: dict[str, object] = {} + + class RecordingPreparer: + def __init__(self, *, data_prep_config, **kwargs: object) -> None: + observed["config"] = data_prep_config + observed["kwargs"] = kwargs + + monkeypatch.setattr( + "fastplms.models.e1.modeling_e1.E1BatchPreparer", + RecordingPreparer, + ) + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + max_num_sequences=7, + max_num_positions_within_seq=31, + max_num_positions_global=64, + ) + model = object.__new__(E1ForMaskedLM) + torch.nn.Module.__init__(model) + model.config = config + model.__dict__["_fastplms_tokenizer_kwargs"] = {} + model.__dict__["_fastplms_prep_tokens"] = None + + assert model.prep_tokens is model.prep_tokens + prep_config = observed["config"] + assert prep_config.max_num_sequences == 7 + assert prep_config.max_num_positions_within_seq == 31 + + +def test_e1_get_tokenizer_prefers_local_model_dir(tmp_path: Path) -> None: + shutil.copyfile(_e1_tokenizer_json(), tmp_path / "tokenizer.json") + + with patch( + "huggingface_hub.hf_hub_download", + side_effect=AssertionError("local tokenizer load should not call Hub download"), + ) as hf_hub_download: + tokenizer = get_tokenizer(tmp_path, local_files_only=True) + + assert not hf_hub_download.called + assert tokenizer.token_to_id("") == 0 + assert tokenizer.get_vocab_size() == 34 + + +def test_e1_get_tokenizer_local_files_only_missing_local_source_raises( + tmp_path: Path, +) -> None: + with ( + patch("fastplms.models.e1.preparation.os.path.isfile", return_value=False), + patch( + "huggingface_hub.hf_hub_download", + side_effect=AssertionError("missing local tokenizer should not call Hub download"), + ) as hf_hub_download, + pytest.raises(FileNotFoundError), + ): + get_tokenizer(tmp_path, local_files_only=True) + + assert not hf_hub_download.called + + +def test_e1_automodel_local_files_only_uses_local_tokenizer(tmp_path: Path) -> None: + config = E1Config( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + max_num_sequences=4, + max_num_positions_within_seq=64, + max_num_positions_global=128, + ) + config.auto_map = { + "AutoConfig": "modeling_e1.E1Config", + "AutoModelForMaskedLM": "modeling_e1.E1ForMaskedLM", + } + model = E1ForMaskedLM(config) + model.save_pretrained(tmp_path) + shutil.copyfile(_e1_tokenizer_json(), tmp_path / "tokenizer.json") + e1_source = _repo_root() / "src" / "fastplms" / "models" / "e1" + for source_name in ( + "attention.py", + "cache.py", + "modeling_e1.py", + "preparation.py", + "retrieval.py", + ): + shutil.copyfile(e1_source / source_name, tmp_path / source_name) + + with patch( + "huggingface_hub.hf_hub_download", + side_effect=AssertionError("local AutoModel load should not call Hub download"), + ) as hf_hub_download: + loaded = AutoModelForMaskedLM.from_pretrained( + tmp_path, + trust_remote_code=True, + local_files_only=True, + ) + + assert not hf_hub_download.called + assert loaded.prep_tokens.tokenizer.token_to_id("") == 0 + + +def test_e1_sequence_mode_tokenizer_contract() -> None: + tokenizer = get_tokenizer() + preparer = E1BatchPreparer(tokenizer=tokenizer) + sequences = [ + "M" + CANONICAL_AAS, + "M" + CANONICAL_AAS[::-1], + ] + + assert tokenizer.token_to_id("") == 0 + for token in ("", "", "1", "2", "?", "X"): + token_id = tokenizer.token_to_id(token) + assert token_id is not None, f"E1 token missing from tokenizer: {token}" + + batch = preparer.get_batch_kwargs( + sequences, + device=torch.device("cpu"), + ) + input_ids = batch["input_ids"] + sequence_ids = batch["sequence_ids"] + within_seq_position_ids = batch["within_seq_position_ids"] + global_position_ids = batch["global_position_ids"] + + assert input_ids.shape == sequence_ids.shape + assert input_ids.shape == within_seq_position_ids.shape + assert input_ids.shape == global_position_ids.shape + assert input_ids.shape[0] == len(sequences) + assert bool((sequence_ids == -1).eq(input_ids == tokenizer.token_to_id("")).all()) + assert bool((within_seq_position_ids[sequence_ids != -1] >= 0).all()) + assert bool((global_position_ids[sequence_ids != -1] >= 0).all()) diff --git a/tools/artifacts/__init__.py b/tools/artifacts/__init__.py new file mode 100644 index 0000000..72e9742 --- /dev/null +++ b/tools/artifacts/__init__.py @@ -0,0 +1,28 @@ +"""Deterministic local Hugging Face artifact construction.""" + +from .build import ( + ArtifactError, + build_artifact, + build_local_artifact, + canonicalize_checkpoint_weights, + hash_file, + render_model_card, + validate_artifact, + validate_repository_legal_inventory, + validate_weight_artifact, + verify_checkpoint, +) + + +__all__ = [ + "ArtifactError", + "build_artifact", + "build_local_artifact", + "canonicalize_checkpoint_weights", + "hash_file", + "render_model_card", + "validate_artifact", + "validate_repository_legal_inventory", + "validate_weight_artifact", + "verify_checkpoint", +] diff --git a/tools/artifacts/build.py b/tools/artifacts/build.py new file mode 100644 index 0000000..e5d7937 --- /dev/null +++ b/tools/artifacts/build.py @@ -0,0 +1,3790 @@ +"""Build deterministic, offline-loadable local model artifacts. + +This tool only reads an already downloaded checkpoint snapshot. It never logs +in, downloads weights, creates a Hub repository, or uploads files. +""" + +from __future__ import annotations + +import argparse +import ast +import base64 +import hashlib +import io +import json +import re +import shutil +import subprocess +import sys +import tarfile +import tempfile +from collections.abc import Iterable, Mapping +from pathlib import Path, PurePosixPath +from typing import Any, cast +from zipfile import ZIP_DEFLATED, BadZipFile, ZipFile, ZipInfo + +from fastplms import __version__ +from fastplms.registry import ( + CheckpointSource, + FileDigest, + ModelRegistry, + ModelSpec, + RegistryError, + _portable_relative_path, + get_model_registry, +) +from tools.artifacts.license_metadata import ( + parse_hub_license_metadata, + validate_hub_license_metadata, +) +from tools.conversion import StateTransformError, apply_state_transform +from tools.source_provenance import ( + ARCHIVE_PROVENANCE_NAME, + SourceProvenanceError, + validate_archived_root, + validate_archived_submodule, +) + + +_MAX_SHARD_BYTES = 5 * 1024**3 +_IGNORED_PARTS = frozenset({".cache", ".git", "__pycache__"}) +_WEIGHT_SUFFIXES = frozenset({".bin", ".ckpt", ".pt", ".pth", ".safetensors"}) +_WEIGHT_INDEX = "model.safetensors.index.json" +_SHARD_NAME_RE = re.compile(r"^model-(\d{5})-of-(\d{5})\.safetensors$") +_BF16_EXECUTION_POLICIES = frozenset({"static_parameters", "fp32_parameters_autocast"}) +_PROVENANCE_SCHEMA_VERSION = 4 +_ARTIFACT_GENERATOR_VERSION = 4 +_CANONICAL_STATE_SCHEMA_VERSION = 1 +_CANONICAL_STATE_DOMAIN = b"fastplms-canonical-state-v1\0" +_CANONICAL_TENSOR_DOMAIN = b"fastplms-canonical-tensor-v1\0" +_CONVERSION_ATTESTATION_SCHEMA_VERSION = 1 +_RUNTIME_ATTESTATION_SCHEMA_VERSION = 2 +_RUNTIME_ATTESTATION_NAME = "runtime-attestation.json" +_MODEL_CARD_RUNTIME_REVISION_PLACEHOLDER = "" +_MODEL_CARD_RUNTIME_PROVENANCE = ( + "- Runtime revision: recorded separately in the built artifact and published commit" +) +_MODEL_CARD_DIGEST_PROVENANCE = ( + "- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json`" +) +_ARTIFACT_REQUIREMENT_INPUTS = ( + "requirements/core.in", + "requirements/features/flash.in", + "requirements/features/structure.in", +) +_RELEASE_TOOL_SCOPE_PATHS = ( + *_ARTIFACT_REQUIREMENT_INPUTS, + "src/fastplms/__init__.py", + "src/fastplms/models.toml", + "src/fastplms/registry.py", + "tools/artifacts/__init__.py", + "tools/artifacts/build.py", + "tools/artifacts/build_all.py", + "tools/artifacts/generate_docs.py", + "tools/artifacts/license_metadata.py", + "tools/artifacts/offline_probe.py", + "tools/artifacts/publish.py", + "tools/artifacts/resolve_fair_esm_assets.py", + "tools/artifacts/resolve_manifest_hashes.py", + "tools/conversion/__init__.py", + "tools/conversion/extract_esmfold2_geometry.py", + "tools/conversion/state_transforms.py", + "tools/conversion/state_validation.py", + "tools/remote/biohub_reference_environment.py", + "tools/source_provenance.py", +) +_RELEASE_TOOL_SCOPE_ROOTS = ( + *_ARTIFACT_REQUIREMENT_INPUTS, + "src/fastplms/__init__.py", + "src/fastplms/models.toml", + "src/fastplms/registry.py", + "tools/artifacts", + "tools/conversion", + "tools/remote/biohub_reference_environment.py", + "tools/source_provenance.py", +) +_RELEASE_TOOL_DIGEST_DOMAIN = b"fastplms-release-tools-v1\0" +_GENERATED_RUNTIME_UPDATE_PATHS = frozenset( + { + "README.md", + "config.json", + "fastplms_bundle.py", + "modeling_fastplms.py", + "requirements.txt", + "THIRD_PARTY_NOTICES.md", + _RUNTIME_ATTESTATION_NAME, + } +) +_RUNTIME_SOURCE_SUFFIXES = frozenset({".json", ".lock", ".py", ".toml"}) +_MAX_RUNTIME_SOURCE_BYTES = 8 * 1024**2 +_MAX_RUNTIME_ARCHIVE_BYTES = 128 * 1024**2 +_MAX_RUNTIME_ARCHIVE_EXPANDED_BYTES = 256 * 1024**2 +_MAX_RUNTIME_ARCHIVE_MEMBERS = 4096 +_SENSITIVE_SOURCE_NAMES = frozenset( + { + ".env", + ".netrc", + "credentials", + "credentials.json", + "id_ed25519", + "id_rsa", + "secrets.json", + "token", + "token.txt", + } +) +_SENSITIVE_SOURCE_SUFFIXES = frozenset({".key", ".p12", ".pfx", ".pem"}) +_SENSITIVE_SOURCE_STEMS = frozenset( + {"credential", "credentials", "secret", "secrets", "token"} +) +_TOKENIZER_FILE_NAMES = frozenset( + { + "added_tokens.json", + "merges.txt", + "sentencepiece.bpe.model", + "special_tokens_map.json", + "spiece.model", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "vocab.txt", + } +) + + +class ArtifactError(RuntimeError): + """Raised when an artifact cannot be built or validated safely.""" + + +def _update_length_prefixed(digest: Any, value: bytes) -> None: + digest.update(len(value).to_bytes(8, "big")) + digest.update(value) + + +def _canonical_tensor_leaf(name: str, tensor: Any) -> bytes: + """Hash one logical tensor independent of safetensors sharding.""" + + try: + import torch + except ImportError as error: + raise ArtifactError("Canonical state hashing requires torch.") from error + if sys.byteorder != "little": + raise ArtifactError("Canonical state hashing requires a little-endian host.") + if not isinstance(name, str) or not name or not torch.is_tensor(tensor): + raise ArtifactError(f"Canonical state contains an invalid tensor entry: {name!r}.") + if tensor.layout != torch.strided: + raise ArtifactError(f"Canonical state tensor {name!r} is not strided.") + canonical = tensor.detach().to(device="cpu").contiguous() + raw = canonical.reshape(-1).view(torch.uint8).numpy().tobytes() + leaf = hashlib.sha256() + leaf.update(_CANONICAL_TENSOR_DOMAIN) + _update_length_prefixed(leaf, name.encode("utf-8")) + _update_length_prefixed(leaf, str(canonical.dtype).removeprefix("torch.").encode("ascii")) + _update_length_prefixed( + leaf, + json.dumps(list(canonical.shape), separators=(",", ":")).encode("ascii"), + ) + _update_length_prefixed(leaf, raw) + return leaf.digest() + + +def _canonical_state_sha256(state: Mapping[str, Any]) -> str: + """Return a deterministic digest of tensor names, metadata, and values.""" + + if not state: + raise ArtifactError("Canonical state cannot be empty.") + leaves = {name: _canonical_tensor_leaf(name, state[name]) for name in sorted(state)} + digest = hashlib.sha256() + digest.update(_CANONICAL_STATE_DOMAIN) + digest.update(_CANONICAL_STATE_SCHEMA_VERSION.to_bytes(4, "big")) + digest.update(len(leaves).to_bytes(8, "big")) + for name in sorted(leaves): + _update_length_prefixed(digest, name.encode("utf-8")) + _update_length_prefixed(digest, leaves[name]) + return digest.hexdigest() + + +def hash_file(path: Path, algorithm: str = "sha256") -> str: + """Return a normal SHA-256 or Git-blob SHA-1 digest for one file.""" + + if algorithm == "sha256": + digest = hashlib.sha256() + elif algorithm == "git-sha1": + digest = hashlib.sha1(usedforsecurity=False) + digest.update(f"blob {path.stat().st_size}\0".encode("ascii")) + else: + raise ArtifactError(f"Unsupported digest algorithm: {algorithm!r}") + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_legal_bytes(path: Path) -> bytes: + """Return UTF-8 legal text with Git-canonical LF line endings.""" + + try: + raw = path.read_bytes() + text = raw.decode("utf-8") + except (OSError, UnicodeDecodeError) as error: + raise ArtifactError(f"Legal text must be readable UTF-8: {path}") from error + if "\x00" in text: + raise ArtifactError(f"Legal text contains a NUL byte: {path}") + return text.replace("\r\n", "\n").replace("\r", "\n").encode("utf-8") + + +def _hash_canonical_legal_file(path: Path, algorithm: str) -> str: + if algorithm != "sha256": + raise ArtifactError(f"Legal texts require SHA-256, received {algorithm!r}.") + return hashlib.sha256(_canonical_legal_bytes(path)).hexdigest() + + +def verify_checkpoint(snapshot: Path, source: CheckpointSource) -> None: + """Verify every manifest-pinned file in a local checkpoint snapshot.""" + + snapshot = snapshot.resolve() + if not snapshot.is_dir(): + raise ArtifactError(f"Checkpoint snapshot does not exist: {snapshot}") + failures: list[str] = [] + for expected in source.files: + path = snapshot.joinpath(*PurePosixPath(expected.path).parts) + if not path.is_file(): + failures.append(f"missing {expected.path}") + continue + actual = hash_file(path, expected.algorithm) + if actual != expected.digest: + failures.append( + f"{expected.path}: expected {expected.encoded}, " + f"received {expected.algorithm}:{actual}" + ) + if failures: + detail = "\n - ".join(failures) + raise ArtifactError(f"Checkpoint verification failed for {source.repo_id}:\n - {detail}") + + +def _is_weight_file(path: str) -> bool: + return PurePosixPath(path).suffix.lower() in _WEIGHT_SUFFIXES + + +def _is_runtime_update_path(path: str) -> bool: + relative = PurePosixPath(path) + return ( + path in _GENERATED_RUNTIME_UPDATE_PATHS + or (relative.parts and relative.parts[0] in {"fastplms", "LICENSES"}) + or (len(relative.parts) == 1 and relative.name in _TOKENIZER_FILE_NAMES) + ) + + +def _copy_checkpoint_assets( + snapshot: Path, + destination: Path, + source: CheckpointSource, +) -> None: + """Copy only pinned, non-weight checkpoint assets into an artifact.""" + + for expected in source.files: + if _is_weight_file(expected.path): + continue + relative = PurePosixPath(expected.path) + _copy_verified_checkpoint_file( + snapshot.joinpath(*relative.parts), + destination.joinpath(*relative.parts), + expected, + ) + + +def _copy_official_tokenizer_assets( + snapshot: Path, + destination: Path, + source: CheckpointSource, +) -> None: + """Copy byte-exact tokenizer files from the pinned official snapshot.""" + + selected = [ + item for item in source.files if PurePosixPath(item.path).name in _TOKENIZER_FILE_NAMES + ] + if not selected: + raise ArtifactError(f"Official checkpoint {source.repo_id} declares no tokenizer files.") + for expected in selected: + relative = PurePosixPath(expected.path) + source_path = snapshot.joinpath(*relative.parts) + _copy_verified_checkpoint_file( + source_path, + destination.joinpath(*relative.parts), + expected, + ) + + +def _validated_weight_snapshot( + snapshot: Path, + source: CheckpointSource, + destination: Path, +) -> Path: + """Copy pinned weight bytes into a private builder-owned snapshot.""" + + selected = tuple(item for item in source.files if _is_weight_file(item.path)) + if not selected: + raise ArtifactError(f"Checkpoint {source.repo_id} declares no weight files.") + for expected in selected: + relative = PurePosixPath(expected.path) + _copy_verified_checkpoint_file( + snapshot.joinpath(*relative.parts), + destination.joinpath(*relative.parts), + expected, + ) + return destination + + +def _load_checkpoint_state(snapshot: Path, source: CheckpointSource) -> dict[str, Any]: + """Load a builder-owned, hash-verified state without unrestricted pickle.""" + + try: + import torch + from safetensors.torch import load_file + except ImportError as error: + raise ArtifactError( + "Checkpoint canonicalization requires the core torch and safetensors dependencies." + ) from error + + weight_files = sorted( + (item for item in source.files if _is_weight_file(item.path)), + key=lambda item: item.path, + ) + if not weight_files: + raise ArtifactError(f"Checkpoint {source.repo_id} does not declare any weight files.") + + state: dict[str, Any] = {} + for expected in weight_files: + path = snapshot.joinpath(*PurePosixPath(expected.path).parts) + suffix = path.suffix.lower() + try: + if suffix == ".safetensors": + loaded: object = load_file(path, device="cpu") + elif suffix == ".bin": + loaded = torch.load(path, map_location="cpu", weights_only=True) + else: + raise ArtifactError( + f"Artifact conversion does not accept {suffix!r} weight files. " + "Convert the trusted source to safetensors or a hash-pinned .bin first." + ) + except ArtifactError: + raise + except Exception as error: + raise ArtifactError(f"Unable to load verified weight file: {path}") from error + + if isinstance(loaded, Mapping) and set(loaded) == {"state_dict"}: + loaded = loaded["state_dict"] + if not isinstance(loaded, Mapping): + raise ArtifactError(f"Weight file does not contain a state dictionary: {path}") + + loaded_keys = list(loaded) + if any(not isinstance(name, str) or not name for name in loaded_keys): + raise ArtifactError(f"Weight file contains an invalid parameter key: {path}") + for name in sorted(loaded_keys): + tensor = loaded[name] + if name in state: + raise ArtifactError(f"Duplicate parameter {name!r} across checkpoint shards.") + if not torch.is_tensor(tensor): + raise ArtifactError(f"State entry {name!r} in {path.name} is not a tensor.") + if tensor.layout != torch.strided: + raise ArtifactError( + f"State entry {name!r} uses unsupported layout {tensor.layout}." + ) + # Break shared storage deterministically because safetensors stores + # each state key independently. Parameter alias contracts are tested + # after model loading rather than encoded as shared file storage. + state[name] = tensor.detach().to(device="cpu").contiguous().clone() + del loaded + if not state: + raise ArtifactError(f"Checkpoint {source.repo_id} contains an empty state dictionary.") + return state + + +def canonicalize_checkpoint_weights( + snapshot: Path, + source: CheckpointSource, + destination: Path, + *, + state_transform: str = "identity", + source_is_canonical: bool = False, + max_shard_bytes: int = _MAX_SHARD_BYTES, +) -> dict[str, Any]: + """Transform weights, then write deterministic safetensors shards and an index.""" + + if max_shard_bytes <= 256: + raise ArtifactError("max_shard_bytes must exceed 256 bytes.") + try: + from safetensors.torch import save_file + except ImportError as error: + raise ArtifactError("Writing artifacts requires safetensors.") from error + + snapshot = snapshot.resolve() + destination = destination.resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=".fastplms-validated-checkpoint-", + dir=destination.parent, + ) as directory: + owned_snapshot = _validated_weight_snapshot(snapshot, source, Path(directory)) + state = _load_checkpoint_state(owned_snapshot, source) + try: + state = apply_state_transform( + state_transform, + state, + expected_keys=state if source_is_canonical else None, + ) + except StateTransformError as error: + raise ArtifactError( + f"Unable to apply declared state transform {state_transform!r} " + f"for {source.repo_id}: {error}" + ) from error + state_sha256 = _canonical_state_sha256(state) + margin = min(16 * 1024**2, max(128, max_shard_bytes // 100)) + payload_limit = max_shard_bytes - margin + groups: list[list[str]] = [] + current: list[str] = [] + current_bytes = 0 + total_size = 0 + for name in sorted(state): + tensor = state[name] + tensor_bytes = tensor.numel() * tensor.element_size() + total_size += tensor_bytes + if tensor_bytes > payload_limit: + raise ArtifactError( + f"Parameter {name!r} requires {tensor_bytes} bytes and cannot fit in a " + f"{max_shard_bytes}-byte safetensors shard." + ) + if current and current_bytes + tensor_bytes > payload_limit: + groups.append(current) + current = [] + current_bytes = 0 + current.append(name) + current_bytes += tensor_bytes + if current: + groups.append(current) + if not groups or len(groups) > 99_999: + raise ArtifactError(f"Invalid canonical shard count: {len(groups)}") + + destination.mkdir(parents=True, exist_ok=True) + weight_map: dict[str, str] = {} + shard_hashes: dict[str, str] = {} + shard_count = len(groups) + for index, names in enumerate(groups, start=1): + shard_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors" + shard_path = destination / shard_name + shard_state = {name: state[name] for name in names} + try: + save_file(shard_state, shard_path, metadata={"format": "pt"}) + except Exception as error: + raise ArtifactError(f"Unable to write canonical shard: {shard_path}") from error + if shard_path.stat().st_size > max_shard_bytes: + raise ArtifactError(f"Generated shard {shard_name} exceeds {max_shard_bytes} bytes.") + for name in names: + weight_map[name] = shard_name + shard_hashes[shard_name] = f"sha256:{hash_file(shard_path)}" + + index_path = destination / _WEIGHT_INDEX + _write_json( + index_path, + { + "metadata": {"total_size": total_size}, + "weight_map": weight_map, + }, + ) + validate_weight_artifact( + destination, + max_shard_bytes=max_shard_bytes, + expected_state_sha256=state_sha256, + ) + return { + "format": "safetensors", + "index": _WEIGHT_INDEX, + "index_digest": f"sha256:{hash_file(index_path)}", + "max_shard_bytes": max_shard_bytes, + "shards": shard_hashes, + "source_schema": "canonical" if source_is_canonical else "official", + "state_transform": state_transform, + "state_digest": { + "schema_version": _CANONICAL_STATE_SCHEMA_VERSION, + "algorithm": "sha256", + "sha256": state_sha256, + }, + "tensor_count": len(weight_map), + "total_size": total_size, + } + + +def validate_weight_artifact( + path: Path, + *, + max_shard_bytes: int = _MAX_SHARD_BYTES, + expected_state_sha256: str | None = None, +) -> dict[str, Any]: + """Validate an explicit safetensors shard set and its index.""" + + path = path.resolve() + index_path = path / _WEIGHT_INDEX + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError(f"Unable to read weight index: {index_path}") from error + if not isinstance(index, dict) or set(index) != {"metadata", "weight_map"}: + raise ArtifactError("Weight index must contain exactly metadata and weight_map.") + metadata = index["metadata"] + weight_map = index["weight_map"] + if ( + not isinstance(metadata, dict) + or not isinstance(metadata.get("total_size"), int) + or metadata["total_size"] < 0 + or not isinstance(weight_map, dict) + or not weight_map + or any(not isinstance(key, str) or not key for key in weight_map) + or any(not isinstance(value, str) for value in weight_map.values()) + ): + raise ArtifactError("Weight index metadata or weight_map is invalid.") + + shard_names = sorted(set(weight_map.values())) + parsed_names: list[tuple[int, int, str]] = [] + for shard_name in shard_names: + match = _SHARD_NAME_RE.fullmatch(shard_name) + if match is None: + raise ArtifactError(f"Invalid explicit shard name: {shard_name!r}") + parsed_names.append((int(match.group(1)), int(match.group(2)), shard_name)) + expected_count = len(shard_names) + if {index for index, _, _ in parsed_names} != set(range(1, expected_count + 1)) or { + count for _, count, _ in parsed_names + } != {expected_count}: + raise ArtifactError("Weight index does not describe a complete shard sequence.") + + actual_shards = {item.name for item in path.glob("*.safetensors") if item.is_file()} + if actual_shards != set(shard_names): + raise ArtifactError( + "Weight index and artifact shard files differ: " + f"index={shard_names}, files={sorted(actual_shards)}" + ) + legacy_weights = sorted( + item.name + for item in path.iterdir() + if item.is_file() and item.suffix.lower() in {".bin", ".ckpt", ".pt", ".pth"} + ) + if legacy_weights: + raise ArtifactError(f"Artifact contains non-safetensors weights: {legacy_weights}") + + try: + from safetensors import safe_open + except ImportError as error: + raise ArtifactError("Validating artifact weights requires safetensors.") from error + observed_keys: set[str] = set() + observed_total_size = 0 + tensor_leaves: dict[str, bytes] = {} + for shard_name in shard_names: + shard_path = path / shard_name + if not shard_path.is_file() or shard_path.stat().st_size > max_shard_bytes: + raise ArtifactError( + f"Shard {shard_name} is missing or exceeds {max_shard_bytes} bytes." + ) + expected_keys = {key for key, value in weight_map.items() if value == shard_name} + try: + with safe_open(shard_path, framework="pt", device="cpu") as handle: + actual_keys = set(handle.keys()) + if actual_keys != expected_keys: + raise ArtifactError(f"Shard {shard_name} keys differ from the weight index.") + for key in sorted(actual_keys): + tensor = handle.get_tensor(key) + observed_total_size += tensor.numel() * tensor.element_size() + tensor_leaves[key] = _canonical_tensor_leaf(key, tensor) + del tensor + except ArtifactError: + raise + except Exception as error: + raise ArtifactError(f"Unable to validate safetensors shard: {shard_name}") from error + if observed_keys.intersection(actual_keys): + raise ArtifactError(f"Duplicate tensor keys found in shard {shard_name}.") + observed_keys.update(actual_keys) + if observed_keys != set(weight_map): + raise ArtifactError("Weight index does not cover every stored tensor.") + if observed_total_size != metadata["total_size"]: + raise ArtifactError( + "Weight index total_size differs from the stored tensors: " + f"expected {metadata['total_size']}, received {observed_total_size}." + ) + state_digest = hashlib.sha256() + state_digest.update(_CANONICAL_STATE_DOMAIN) + state_digest.update(_CANONICAL_STATE_SCHEMA_VERSION.to_bytes(4, "big")) + state_digest.update(len(tensor_leaves).to_bytes(8, "big")) + for name in sorted(tensor_leaves): + _update_length_prefixed(state_digest, name.encode("utf-8")) + _update_length_prefixed(state_digest, tensor_leaves[name]) + actual_state_sha256 = state_digest.hexdigest() + if expected_state_sha256 is not None and actual_state_sha256 != expected_state_sha256: + raise ArtifactError( + "Canonical state digest differs from the trusted conversion commitment: " + f"expected sha256:{expected_state_sha256}, " + f"received sha256:{actual_state_sha256}." + ) + return index + + +def _verify_expected_file(path: Path, expected: FileDigest, label: str) -> None: + if not path.is_file(): + raise ArtifactError(f"Missing required {label}: {path}") + actual = hash_file(path, expected.algorithm) + if actual != expected.digest: + raise ArtifactError( + f"Required {label} differs from the manifest: {path}; " + f"expected {expected.encoded}, received {expected.algorithm}:{actual}" + ) + + +def _verify_expected_legal_file(path: Path, expected: FileDigest, label: str) -> None: + if not path.is_file(): + raise ArtifactError(f"Missing required {label}: {path}") + actual = _hash_canonical_legal_file(path, expected.algorithm) + if actual != expected.digest: + raise ArtifactError( + f"Required {label} differs from the manifest after canonical LF normalization: " + f"{path}; expected {expected.encoded}, received {expected.algorithm}:{actual}" + ) + + +def validate_repository_legal_inventory( + source_root: Path, + registry: ModelRegistry, + spec: ModelSpec | None = None, +) -> None: + """Verify legal texts and conversion records before release packaging.""" + + source_root = source_root.resolve() + for expected in registry.legal_files: + path = source_root.joinpath(*PurePosixPath(expected.path).parts) + _verify_expected_legal_file(path, expected, "repository legal file") + + source_ids = spec.family.upstreams if spec is not None else tuple(registry.upstreams) + for source_id in source_ids: + source = registry.upstreams[source_id] + for expected in source.license_digests: + path = source_root / source.path + path = path.joinpath(*PurePosixPath(expected.path).parts) + _verify_expected_legal_file(path, expected, f"{source_id} canonical legal file") + for expected in source.distribution_files: + path = source_root / "LICENSES" / source_id + path = path.joinpath(*PurePosixPath(expected.path).parts) + _verify_expected_legal_file(path, expected, f"{source_id} distribution legal file") + + families = (spec.family,) if spec is not None else tuple(registry.families.values()) + for family in families: + required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:") + if ( + not family.state_transform + or family.state_transform not in family.conversion_provenance + or any(section not in family.conversion_provenance for section in required_sections) + ): + raise ArtifactError(f"Model family {family.id!r} is missing conversion provenance.") + + +def _iter_files(root: Path) -> Iterable[Path]: + for path in sorted(root.rglob("*"), key=lambda item: item.as_posix()): + relative = path.relative_to(root) + if any(part in _IGNORED_PARTS for part in relative.parts): + continue + if path.is_symlink(): + raise ArtifactError(f"Symlinks are not allowed in artifacts: {path}") + if path.is_file() and path.suffix not in {".pyc", ".pyo"}: + yield path + + +def _iter_runtime_source_files(root: Path) -> Iterable[Path]: + """Yield only bounded runtime sources from one artifact scope.""" + + if root.is_symlink(): + raise ArtifactError(f"Symlinks are not allowed in runtime sources: {root}") + candidates = (root,) if root.is_file() else root.rglob("*") + for path in sorted(candidates, key=lambda item: item.as_posix()): + if path == root and root.is_dir(): + continue + relative = path.relative_to(root) if path != root else Path(path.name) + if any(part in _IGNORED_PARTS for part in relative.parts): + continue + if path.is_symlink(): + raise ArtifactError(f"Symlinks are not allowed in runtime sources: {path}") + if not path.is_file(): + continue + suffix = path.suffix.lower() + lowered_parts = tuple(part.lower() for part in path.parts) + if ( + any(part in _SENSITIVE_SOURCE_NAMES for part in lowered_parts) + or any( + PurePosixPath(part).stem in _SENSITIVE_SOURCE_STEMS + for part in lowered_parts + ) + or suffix in _SENSITIVE_SOURCE_SUFFIXES + ): + raise ArtifactError(f"Runtime source contains a sensitive path: {path}") + if suffix not in _RUNTIME_SOURCE_SUFFIXES: + raise ArtifactError( + f"Runtime source has an unapproved extension {suffix!r}: {path}" + ) + size = path.stat().st_size + if size > _MAX_RUNTIME_SOURCE_BYTES: + raise ArtifactError( + f"Runtime source exceeds {_MAX_RUNTIME_SOURCE_BYTES} bytes: {path}" + ) + yield path + + +def _copy_file(source: Path, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, destination) + + +def _copy_verified_checkpoint_file( + source: Path, + destination: Path, + expected: FileDigest, +) -> None: + """Preserve one mutable source file, then validate the owned bytes.""" + + if source.is_symlink(): + # Hugging Face snapshots intentionally use links into their blob cache. + # The copied destination is a regular builder-owned file whose complete + # content is checked below, so source-link identity is not trusted. + if not source.is_file(): + raise ArtifactError(f"Pinned checkpoint link is unavailable: {source}") + elif not source.is_file(): + raise ArtifactError(f"Pinned checkpoint file is missing: {source}") + try: + _copy_file(source, destination) + actual = hash_file(destination, expected.algorithm) + except OSError as error: + raise ArtifactError(f"Unable to preserve pinned checkpoint file: {source}") from error + if actual != expected.digest: + destination.unlink(missing_ok=True) + raise ArtifactError( + "Preserved checkpoint bytes differ from the pinned source after copy: " + f"{expected.path}; expected {expected.encoded}, " + f"received {expected.algorithm}:{actual}." + ) + + +def _copy_canonical_legal_file(source: Path, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(_canonical_legal_bytes(source)) + + +def _copy_tree(source: Path, destination: Path) -> None: + if not source.exists(): + raise ArtifactError(f"Required artifact source path does not exist: {source}") + copied = False + for path in _iter_runtime_source_files(source): + target = destination if source.is_file() else destination / path.relative_to(source) + _copy_file(path, target) + copied = True + if not copied: + raise ArtifactError(f"Runtime source scope contains no distributable files: {source}") + + +def _git_runtime_revision( + source_root: Path, + scopes: Iterable[Path], + entries: Iterable[tuple[Path, PurePosixPath]], +) -> str | None: + """Return HEAD after proving every selected repository source is tracked and clean.""" + + git_metadata = source_root / ".git" + if not (git_metadata.exists() or git_metadata.is_symlink()): + return None + selected_scopes = tuple(scopes) + relative_scopes: list[str] = [] + for scope in selected_scopes: + try: + relative = scope.resolve().relative_to(source_root) + except ValueError as error: + raise ArtifactError(f"Release source escapes the repository: {scope}") from error + relative_scopes.append(relative.as_posix()) + command_prefix = [ + "git", + "-c", + f"safe.directory={source_root.as_posix()}", + ] + try: + status = subprocess.run( + [ + *command_prefix, + "status", + "--porcelain=v1", + "--untracked-files=all", + "--", + *relative_scopes, + ], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if status: + raise ArtifactError( + "Artifact runtime inputs must be tracked and clean; scoped Git status: " + + status.replace("\n", "; ") + ) + tracked = subprocess.run( + [*command_prefix, "ls-files", "-z", "--", *relative_scopes], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + tracked_names = { + raw.decode("utf-8") for raw in tracked.split(b"\0") if raw + } + selected_names = { + path.resolve().relative_to(source_root).as_posix() for path, _ in entries + } + missing = sorted(selected_names.difference(tracked_names)) + if missing: + raise ArtifactError( + f"Artifact runtime inputs contain untracked files: {missing[:10]}" + ) + revision = subprocess.run( + [*command_prefix, "rev-parse", "HEAD"], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError, UnicodeDecodeError) as error: + raise ArtifactError("Unable to validate tracked artifact runtime sources.") from error + if re.fullmatch(r"[0-9a-f]{40}", revision) is None: + raise ArtifactError(f"Git returned an invalid runtime revision: {revision!r}") + return revision + + +def _release_tool_scope_digest(payloads: Mapping[str, bytes]) -> str: + """Hash path-bound release-tool bytes without trusting filesystem metadata.""" + + digest = hashlib.sha256() + digest.update(_RELEASE_TOOL_DIGEST_DOMAIN) + for relative_name, payload in sorted(payloads.items()): + encoded_name = relative_name.encode("utf-8") + digest.update(len(encoded_name).to_bytes(4, "big")) + digest.update(encoded_name) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(hashlib.sha256(payload).digest()) + return digest.hexdigest() + + +def _git_archived_regular_files( + source_root: Path, + revision: str, + relative_names: Iterable[str], + *, + label: str, +) -> dict[str, bytes]: + """Read an exact regular-file allowlist from one immutable Git tree.""" + + names = tuple(sorted(relative_names)) + command_prefix = ["git", "-c", f"safe.directory={source_root.as_posix()}"] + try: + archive = subprocess.run( + [*command_prefix, "archive", "--format=tar", revision, "--", *names], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + payloads: dict[str, bytes] = {} + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as handle: + for member in handle.getmembers(): + if member.isdir(): + continue + name = PurePosixPath(member.name).as_posix() + if not member.isfile() or name in payloads: + raise ArtifactError( + f"Tracked {label} is not a unique regular file: {member.name}" + ) + extracted = handle.extractfile(member) + if extracted is None: + raise ArtifactError(f"Unable to read tracked {label}: {member.name}") + payloads[name] = extracted.read(_MAX_RUNTIME_SOURCE_BYTES + 1) + except ArtifactError: + raise + except (OSError, subprocess.CalledProcessError, tarfile.TarError) as error: + raise ArtifactError(f"Unable to archive immutable {label} bytes.") from error + if set(payloads) != set(names): + raise ArtifactError(f"Tracked {label} archive differs from its exact allowlist.") + oversized = [ + name for name, payload in payloads.items() if len(payload) > _MAX_RUNTIME_SOURCE_BYTES + ] + if oversized: + raise ArtifactError(f"Tracked {label} files exceed the size limit: {oversized[:10]}") + return payloads + + +def _validated_release_tool_snapshot( + source_root: Path, + *, + _allow_untracked_for_tests: bool = False, +) -> tuple[str, str, dict[str, bytes]]: + """Return immutable bytes and identity for every artifact release tool. + + A Git worktree must have an exact, clean tracked inventory. A portable + Git-free runner must carry a validated root archive attestation. The test + escape hatch is private and deliberately produces a content-only identity. + """ + + source_root = source_root.resolve() + payloads: dict[str, bytes] + if _allow_untracked_for_tests: + payloads = { + relative_name: source_root.joinpath( + *PurePosixPath(relative_name).parts + ).read_bytes() + for relative_name in _ARTIFACT_REQUIREMENT_INPUTS + } + payloads["test-only-release-tool-scope"] = b"untracked test fixture" + tool_digest = _release_tool_scope_digest(payloads) + return f"release-tools-sha256:{tool_digest}", tool_digest, payloads + git_metadata = source_root / ".git" + if git_metadata.exists() or git_metadata.is_symlink(): + command_prefix = ["git", "-c", f"safe.directory={source_root.as_posix()}"] + try: + revision = subprocess.run( + [*command_prefix, "rev-parse", "HEAD"], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + status = subprocess.run( + [ + *command_prefix, + "status", + "--porcelain=v1", + "--untracked-files=all", + "--", + *_RELEASE_TOOL_SCOPE_ROOTS, + ], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if status: + raise ArtifactError( + "Artifact release tools must be tracked and clean; scoped Git status: " + + status.replace("\n", "; ") + ) + tracked = subprocess.run( + [ + *command_prefix, + "ls-files", + "-z", + "--", + *_RELEASE_TOOL_SCOPE_ROOTS, + ], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + except ArtifactError: + raise + except (OSError, subprocess.CalledProcessError, UnicodeDecodeError) as error: + raise ArtifactError("Unable to validate tracked artifact release tools.") from error + if re.fullmatch(r"[0-9a-f]{40}", revision) is None: + raise ArtifactError(f"Git returned an invalid release-tool revision: {revision!r}") + try: + tracked_names = { + raw.decode("utf-8") for raw in tracked.split(b"\0") if raw + } + except UnicodeDecodeError as error: + raise ArtifactError("Git returned a non-UTF-8 release-tool path.") from error + if tracked_names != set(_RELEASE_TOOL_SCOPE_PATHS): + missing = sorted(set(_RELEASE_TOOL_SCOPE_PATHS).difference(tracked_names)) + extra = sorted(tracked_names.difference(_RELEASE_TOOL_SCOPE_PATHS)) + raise ArtifactError( + "Tracked artifact release-tool inventory differs from the exact allowlist; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + payloads = _git_archived_regular_files( + source_root, + revision, + _RELEASE_TOOL_SCOPE_PATHS, + label="artifact release tool", + ) + tool_digest = _release_tool_scope_digest(payloads) + return revision, tool_digest, payloads + + archive_marker = source_root / ARCHIVE_PROVENANCE_NAME + if archive_marker.exists() or archive_marker.is_symlink(): + try: + _diagnostic_head, inventory = validate_archived_root(source_root) + except SourceProvenanceError as error: + raise ArtifactError( + f"Git-free release-tool attestation is invalid: {error}" + ) from error + scoped_names = { + relative_name + for relative_name in inventory + if any( + relative_name == scope + or relative_name.startswith(scope.rstrip("/") + "/") + for scope in _RELEASE_TOOL_SCOPE_ROOTS + ) + } + if scoped_names != set(_RELEASE_TOOL_SCOPE_PATHS): + missing = sorted(set(_RELEASE_TOOL_SCOPE_PATHS).difference(scoped_names)) + extra = sorted(scoped_names.difference(_RELEASE_TOOL_SCOPE_PATHS)) + raise ArtifactError( + "Attested artifact release-tool inventory differs from the exact allowlist; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + payloads = {} + for relative_name in _RELEASE_TOOL_SCOPE_PATHS: + path = source_root.joinpath(*PurePosixPath(relative_name).parts) + record = inventory[relative_name] + if path.is_symlink() or not path.is_file() or record.get("mode") not in { + "100644", + "100755", + }: + raise ArtifactError( + f"Attested artifact release tool is not a regular file: {path}" + ) + try: + payload = path.read_bytes() + except OSError as error: + raise ArtifactError(f"Unable to read artifact release tool: {path}") from error + if ( + record.get("size") != len(payload) + or record.get("sha256") != hashlib.sha256(payload).hexdigest() + ): + raise ArtifactError( + f"Attested artifact release tool mutated during snapshot: {relative_name}" + ) + payloads[relative_name] = payload + tool_digest = _release_tool_scope_digest(payloads) + return f"release-tools-sha256:{tool_digest}", tool_digest, payloads + + raise ArtifactError( + "Artifact release tools require either a clean verifiable Git worktree " + "or a content-attested tracked remote source archive." + ) + + +def _requirement_lines(payload: bytes, source_name: str) -> tuple[str, ...]: + """Parse one direct dependency declaration used by Hub artifacts.""" + + try: + text = payload.decode("utf-8") + except UnicodeDecodeError as error: + raise ArtifactError(f"Artifact dependency input is not UTF-8: {source_name}") from error + + requirements: list[str] = [] + for line_number, raw_line in enumerate(text.splitlines(), start=1): + requirement = raw_line.strip() + if not requirement or requirement.startswith("#"): + continue + if requirement.startswith(("-", "--")): + raise ArtifactError( + f"Artifact dependency input must contain direct requirements only: " + f"{source_name}:{line_number}" + ) + requirements.append(requirement) + if not requirements: + raise ArtifactError(f"Artifact dependency input is empty: {source_name}") + return tuple(requirements) + + +def _artifact_requirement_paths(spec: ModelSpec) -> tuple[str, ...]: + paths = ["requirements/core.in"] + if spec.family.extra == "structure": + paths.append("requirements/features/structure.in") + if any(name.startswith("flash_attention_") for name in spec.family.attention): + paths.append("requirements/features/flash.in") + return tuple(paths) + + +def _render_artifact_requirements( + spec: ModelSpec, + release_tool_payloads: Mapping[str, bytes], +) -> str: + """Render the direct dependencies shipped beside one Hub model.""" + + requirements: list[str] = [] + for source_name in _artifact_requirement_paths(spec): + try: + payload = release_tool_payloads[source_name] + except KeyError as error: + raise ArtifactError( + f"Release-tool snapshot is missing artifact dependencies: {source_name}" + ) from error + for requirement in _requirement_lines(payload, source_name): + if requirement not in requirements: + requirements.append(requirement) + + return "\n".join( + ( + f"# Direct runtime dependencies for {spec.fast.repo_id}.", + "# FastPLMs source is embedded in this model repository.", + *requirements, + "", + ) + ) + + +def _runtime_source_entries( + source_root: Path, + spec: ModelSpec, +) -> tuple[tuple[Path, PurePosixPath], ...]: + """Map every approved runtime input to its packaged relative path.""" + + package_source = source_root / "src" / "fastplms" + entries: dict[str, Path] = {} + for relative_name in spec.family.runtime_paths: + target_root = PurePosixPath(relative_name) + source = package_source.joinpath(*target_root.parts) + selected = tuple(_iter_runtime_source_files(source)) + if not selected: + raise ArtifactError( + f"Runtime source scope contains no distributable files: {source}" + ) + for path in selected: + target = ( + target_root + if source.is_file() + else target_root / path.relative_to(source).as_posix() + ) + target_name = target.as_posix() + if target_name in entries: + raise ArtifactError( + f"Runtime source scopes overlap at packaged path {target_name!r}." + ) + entries[target_name] = path + if any(name.startswith("flash_attention_") for name in spec.family.attention): + kernel_source = source_root / "kernels.lock" + if tuple(_iter_runtime_source_files(kernel_source)) != (kernel_source,): + raise ArtifactError(f"Runtime kernel lock is missing or invalid: {kernel_source}") + entries["kernels.lock"] = kernel_source + return tuple( + (entries[name], PurePosixPath(name)) + for name in sorted(entries) + ) + + +def _snapshot_runtime_sources( + source_root: Path, + entries: Iterable[tuple[Path, PurePosixPath]], + revision: str | None, +) -> dict[str, bytes]: + """Retain runtime bytes, using immutable Git blobs whenever available.""" + + selected = tuple(entries) + source_names = { + source.resolve().relative_to(source_root).as_posix(): target.as_posix() + for source, target in selected + } + if len(source_names) != len(selected): + raise ArtifactError("Runtime source inputs contain duplicate tracked paths.") + if revision is None: + try: + payloads = { + target.as_posix(): source.read_bytes() + for source, target in selected + } + except OSError as error: + raise ArtifactError("Unable to snapshot runtime source inputs.") from error + else: + try: + archive = subprocess.run( + [ + "git", + "-c", + f"safe.directory={source_root.as_posix()}", + "archive", + "--format=tar", + revision, + "--", + *source_names, + ], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + archived: dict[str, bytes] = {} + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as handle: + for member in handle.getmembers(): + if member.isdir(): + continue + if not member.isfile(): + raise ArtifactError( + f"Tracked runtime source is not a regular file: {member.name}" + ) + extracted = handle.extractfile(member) + if extracted is None: + raise ArtifactError( + f"Unable to read tracked runtime blob: {member.name}" + ) + archived[PurePosixPath(member.name).as_posix()] = extracted.read() + except ArtifactError: + raise + except (OSError, subprocess.CalledProcessError, tarfile.TarError) as error: + raise ArtifactError( + f"Unable to materialize tracked runtime blobs at {revision}." + ) from error + if set(archived) != set(source_names): + raise ArtifactError( + "Tracked runtime archive differs from the validated source allowlist." + ) + payloads = { + target_name: archived[source_name] + for source_name, target_name in source_names.items() + } + oversized = [ + name for name, payload in payloads.items() if len(payload) > _MAX_RUNTIME_SOURCE_BYTES + ] + if oversized: + raise ArtifactError(f"Tracked runtime sources exceed the size limit: {oversized[:10]}") + return payloads + + +def _write_runtime_snapshot(destination: Path, payloads: Mapping[str, bytes]) -> None: + for relative_name, payload in sorted(payloads.items()): + target = destination.joinpath(*PurePosixPath(relative_name).parts) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(payload) + + +def _runtime_payload_tree_sha256(payloads: Mapping[str, bytes]) -> str: + inventory = { + name: f"sha256:{hashlib.sha256(payload).hexdigest()}" + for name, payload in payloads.items() + } + return hashlib.sha256( + json.dumps(inventory, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + + +def _archived_runtime_payloads( + source_root: Path, + spec: ModelSpec, + entries: Iterable[tuple[Path, PurePosixPath]], +) -> dict[str, bytes]: + """Validate and snapshot runtime bytes from a content-attested Git-free archive.""" + + selected = tuple(entries) + try: + _diagnostic_head, inventory = validate_archived_root(source_root) + except SourceProvenanceError as error: + raise ArtifactError(f"Git-free runtime source attestation is invalid: {error}") from error + + package_source = source_root / "src" / "fastplms" + runtime_scopes = [ + package_source.joinpath(*PurePosixPath(relative_name).parts) + for relative_name in spec.family.runtime_paths + ] + if any(name.startswith("flash_attention_") for name in spec.family.attention): + runtime_scopes.append(source_root / "kernels.lock") + + scope_names: list[str] = [] + actual_scope_files: set[str] = set() + for scope in runtime_scopes: + try: + relative_scope = scope.relative_to(source_root).as_posix() + except ValueError as error: + raise ArtifactError( + f"Archived runtime scope escapes the source root: {scope}" + ) from error + scope_names.append(relative_scope) + if scope.is_symlink(): + raise ArtifactError(f"Archived runtime scope is a symlink: {scope}") + if not scope.exists(): + raise ArtifactError(f"Archived runtime scope is missing: {scope}") + candidates = (scope,) if scope.is_file() else scope.rglob("*") + for path in candidates: + if path.is_symlink(): + raise ArtifactError(f"Archived runtime source is a symlink: {path}") + if path.is_file(): + actual_scope_files.add(path.relative_to(source_root).as_posix()) + + expected_scope_files = { + relative_name + for relative_name in inventory + if any( + relative_name == scope_name + or relative_name.startswith(scope_name.rstrip("/") + "/") + for scope_name in scope_names + ) + } + if actual_scope_files != expected_scope_files: + missing = sorted(expected_scope_files.difference(actual_scope_files)) + extra = sorted(actual_scope_files.difference(expected_scope_files)) + raise ArtifactError( + "Git-free runtime source inventory differs from the tracked archive; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + + source_names: dict[str, str] = {} + for source, target in selected: + try: + source_name = source.relative_to(source_root).as_posix() + except ValueError as error: + raise ArtifactError( + f"Archived runtime source escapes the source root: {source}" + ) from error + if source_name in source_names: + raise ArtifactError(f"Archived runtime source is selected twice: {source_name}") + source_names[source_name] = target.as_posix() + if set(source_names) != actual_scope_files: + missing = sorted(actual_scope_files.difference(source_names)) + extra = sorted(set(source_names).difference(actual_scope_files)) + raise ArtifactError( + "Approved runtime allowlist differs from the archived tracked scope; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + + payloads = _snapshot_runtime_sources(source_root, selected, None) + for source_name, target_name in source_names.items(): + record = inventory.get(source_name) + payload = payloads[target_name] + if ( + not isinstance(record, Mapping) + or record.get("mode") not in {"100644", "100755"} + or record.get("size") != len(payload) + or record.get("sha256") != hashlib.sha256(payload).hexdigest() + ): + raise ArtifactError( + f"Archived runtime source mutated during snapshot: {source_name}" + ) + return payloads + + +def _validated_runtime_snapshot( + source_root: Path, + registry: ModelRegistry, + spec: ModelSpec, + *, + _allow_untracked_for_tests: bool = False, +) -> tuple[str, dict[str, bytes], str]: + """Return clean, tracked runtime bytes and their immutable identities.""" + + package_source = source_root / "src" / "fastplms" + runtime_scopes = [ + package_source.joinpath(*PurePosixPath(relative_name).parts) + for relative_name in spec.family.runtime_paths + ] + if any(name.startswith("flash_attention_") for name in spec.family.attention): + runtime_scopes.append(source_root / "kernels.lock") + entries = _runtime_source_entries(source_root, spec) + git_revision = _git_runtime_revision(source_root, runtime_scopes, entries) + if git_revision is None: + archive_marker = source_root / ARCHIVE_PROVENANCE_NAME + if archive_marker.exists() or archive_marker.is_symlink(): + payloads = _archived_runtime_payloads(source_root, spec, entries) + elif _allow_untracked_for_tests: + payloads = _snapshot_runtime_sources(source_root, entries, None) + else: + raise ArtifactError( + "Artifact runtime sources require either a clean verifiable Git worktree " + "or a content-attested tracked remote source archive." + ) + else: + payloads = _snapshot_runtime_sources( + source_root, + entries, + git_revision, + ) + _validate_attention_kernel_lock(payloads.get("kernels.lock"), registry, spec) + source_tree_sha256 = _runtime_payload_tree_sha256(payloads) + runtime_revision = git_revision or f"source-tree-sha256:{source_tree_sha256}" + return runtime_revision, payloads, source_tree_sha256 + + +def _tree_sha256(root: Path) -> str: + """Hash a portable path/digest inventory rather than filesystem metadata.""" + + inventory = { + path.relative_to(root).as_posix(): f"sha256:{hash_file(path)}" + for path in _iter_files(root) + } + return hashlib.sha256( + json.dumps(inventory, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + + +def _validate_attention_kernel_lock( + payload: bytes | None, + registry: ModelRegistry, + spec: ModelSpec, +) -> None: + """Validate the snapshotted kernel lock for advertised FlashAttention backends.""" + + implementations = tuple( + name for name in spec.family.attention if name.startswith("flash_attention_") + ) + if not implementations: + return + if payload is None: + raise ArtifactError("The runtime snapshot is missing kernels.lock.") + try: + entries = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ArtifactError("Unable to read the snapshotted kernel lock.") from error + if not isinstance(entries, list) or any(not isinstance(entry, dict) for entry in entries): + raise ArtifactError("kernels.lock must contain a list of JSON objects.") + locked: dict[str, str] = {} + for entry in entries: + repo_id = entry.get("repo_id") + revision = entry.get("sha") + if ( + not isinstance(repo_id, str) + or not isinstance(revision, str) + or re.fullmatch(r"[0-9a-f]{40}", revision) is None + or repo_id in locked + ): + raise ArtifactError("kernels.lock contains an invalid or duplicate identity.") + locked[repo_id] = revision + for implementation in implementations: + kernel = registry.attention_kernels[implementation] + if locked.get(kernel.repository) != kernel.revision: + raise ArtifactError( + f"kernels.lock does not match {implementation!r}: expected " + f"{kernel.repository}@{kernel.revision}." + ) + + +def _copy_attention_kernel_lock( + source_root: Path, + package_target: Path, + registry: ModelRegistry, + spec: ModelSpec, +) -> None: + """Compatibility helper for focused tests; builds use the retained snapshot.""" + + implementations = tuple( + name for name in spec.family.attention if name.startswith("flash_attention_") + ) + if not implementations: + return + source = source_root / "kernels.lock" + try: + payload = source.read_bytes() + except OSError as error: + raise ArtifactError(f"Unable to read the repository kernel lock: {source}") from error + _validate_attention_kernel_lock(payload, registry, spec) + destination = package_target / "kernels.lock" + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(payload) + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n", + encoding="utf-8", + newline="\n", + ) + + +def _artifact_auto_map(spec: ModelSpec) -> dict[str, str]: + return { + auto_class: f"modeling_fastplms.{class_path.rsplit('.', maxsplit=1)[1]}" + for auto_class, class_path in spec.auto_map.items() + } + + +def _apply_artifact_config_contract(spec: ModelSpec, config: dict[str, Any]) -> None: + """Materialize runtime-only config invariants without rewriting source provenance.""" + + if spec.family.id == "dplm2": + config.update( + { + "is_decoder": False, + "add_cross_attention": False, + "use_cache": False, + } + ) + if spec.family.id == "esmfold2": + if spec.msa_conditioning is None: + raise ArtifactError( + f"ESMFold2 checkpoint {spec.id!r} has no MSA-conditioning contract." + ) + msa_encoder = config.get("msa_encoder") + if not isinstance(msa_encoder, dict): + raise ArtifactError("ESMFold2 config.msa_encoder must be an object.") + config_msa_conditioning = msa_encoder.get("enabled") + if not isinstance(config_msa_conditioning, bool): + raise ArtifactError("ESMFold2 config.msa_encoder.enabled must be a boolean.") + if config_msa_conditioning != spec.msa_conditioning: + raise ArtifactError( + f"ESMFold2 config.msa_encoder.enabled={config_msa_conditioning!r} " + f"differs from models.toml msa_conditioning={spec.msa_conditioning!r}." + ) + declared = config.get("msa_conditioning") + if "msa_conditioning" in config and ( + not isinstance(declared, bool) or declared != spec.msa_conditioning + ): + raise ArtifactError( + "ESMFold2 config.msa_conditioning differs from the models.toml contract." + ) + config["msa_conditioning"] = spec.msa_conditioning + + +def _configure_custom_tokenizer(path: Path, spec: ModelSpec) -> None: + """Point a manifest-declared custom tokenizer at the flat artifact bridge.""" + + class_path = spec.family.tokenizer_class + if class_path is None: + return + config_path = path / "tokenizer_config.json" + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError( + f"Custom tokenizer {class_path!r} requires a valid tokenizer_config.json." + ) from error + if not isinstance(config, dict): + raise ArtifactError("tokenizer_config.json must contain a JSON object.") + raw_auto_map = config.get("auto_map") + if raw_auto_map is None or isinstance(raw_auto_map, (list, tuple)): + auto_map: dict[str, Any] = {} + elif isinstance(raw_auto_map, dict): + auto_map = dict(raw_auto_map) + else: + raise ArtifactError("tokenizer_config.json auto_map must be an object or legacy list.") + class_name = class_path.rsplit(".", maxsplit=1)[1] + auto_map["AutoTokenizer"] = [f"modeling_fastplms.{class_name}", None] + config["auto_map"] = auto_map + _write_json(config_path, config) + + +def _runtime_archive_files(package_root: Path) -> dict[str, bytes]: + files: dict[str, bytes] = {} + expanded_size = 0 + for path in _iter_files(package_root): + relative_name = path.relative_to(package_root).as_posix() + try: + relative = _portable_relative_path(relative_name, "Runtime archive path") + except RegistryError as error: + raise ArtifactError(f"Runtime archive path is invalid: {relative_name!r}") from error + if relative.suffix.lower() not in _RUNTIME_SOURCE_SUFFIXES: + raise ArtifactError( + f"Runtime archive path has an unapproved extension: {relative_name!r}" + ) + try: + size = path.stat().st_size + except OSError as error: + raise ArtifactError(f"Unable to inspect runtime archive source: {path}") from error + if size > _MAX_RUNTIME_SOURCE_BYTES: + raise ArtifactError(f"Runtime archive source exceeds its size limit: {path}") + if len(files) >= _MAX_RUNTIME_ARCHIVE_MEMBERS: + raise ArtifactError("The artifact runtime archive contains too many source files.") + archive_name = (PurePosixPath("fastplms") / relative).as_posix() + try: + with path.open("rb") as handle: + payload = handle.read(_MAX_RUNTIME_SOURCE_BYTES + 1) + except OSError as error: + raise ArtifactError(f"Unable to read runtime archive source: {path}") from error + if len(payload) > _MAX_RUNTIME_SOURCE_BYTES or len(payload) != size: + raise ArtifactError(f"Runtime archive source changed or exceeded its limit: {path}") + expanded_size += len(payload) + if expanded_size > _MAX_RUNTIME_ARCHIVE_EXPANDED_BYTES: + raise ArtifactError("The expanded artifact runtime archive exceeds its size limit.") + files[archive_name] = payload + if not files: + raise ArtifactError("The artifact runtime source archive would be empty.") + return files + + +def _build_runtime_archive(package_root: Path) -> bytes: + """Return a deterministic archive of unchanged package runtime sources.""" + + files = _runtime_archive_files(package_root) + if len(files) > _MAX_RUNTIME_ARCHIVE_MEMBERS: + raise ArtifactError("The artifact runtime archive contains too many source files.") + if sum(map(len, files.values())) > _MAX_RUNTIME_ARCHIVE_EXPANDED_BYTES: + raise ArtifactError("The expanded artifact runtime archive exceeds its size limit.") + buffer = io.BytesIO() + with ZipFile(buffer, mode="w", compression=ZIP_DEFLATED, compresslevel=9) as archive: + for archive_path, contents in sorted(files.items()): + info = ZipInfo(archive_path, date_time=(1980, 1, 1, 0, 0, 0)) + info.compress_type = ZIP_DEFLATED + info.external_attr = 0o100644 << 16 + archive.writestr(info, contents, compress_type=ZIP_DEFLATED, compresslevel=9) + payload = buffer.getvalue() + if len(payload) > _MAX_RUNTIME_ARCHIVE_BYTES: + raise ArtifactError("The compressed artifact runtime archive exceeds its size limit.") + return payload + + +def _render_runtime_bundle(package_root: Path) -> tuple[str, bytes]: + archive = _build_runtime_archive(package_root) + archive_hash = hashlib.sha256(archive).hexdigest() + encoded = base64.b85encode(archive).decode("ascii") + chunks = (encoded[index : index + 100] for index in range(0, len(encoded), 100)) + lines = [ + '"""Generated deterministic archive of unchanged FastPLMs runtime sources."""', + "", + f'RUNTIME_HASH = "{archive_hash}"', + "RUNTIME_DATA = (", + *(f" {chunk!r}" for chunk in chunks), + ")", + "", + ] + return archive_hash, "\n".join(lines).encode("utf-8") + + +def _write_runtime_bundle(path: Path, package_root: Path) -> str: + """Write the flat source bundle consumed by Transformers remote code.""" + + archive_hash, payload = _render_runtime_bundle(package_root) + path.write_bytes(payload) + return archive_hash + + +def _decode_runtime_bundle(path: Path) -> tuple[str, bytes]: + """Decode a data-only generated bundle without executing artifact code.""" + + try: + source = path.read_text(encoding="utf-8") + module = ast.parse(source, filename=str(path)) + except (OSError, UnicodeDecodeError, SyntaxError) as error: + raise ArtifactError(f"Runtime bundle is missing or invalid: {path}") from error + if len(module.body) != 3 or not ( + isinstance(module.body[0], ast.Expr) + and isinstance(module.body[0].value, ast.Constant) + and isinstance(module.body[0].value.value, str) + ): + raise ArtifactError("Runtime bundle must contain only its docstring and data assignments.") + assignments: dict[str, Any] = {} + for statement in module.body[1:]: + if not ( + isinstance(statement, ast.Assign) + and len(statement.targets) == 1 + and isinstance(statement.targets[0], ast.Name) + and statement.targets[0].id in {"RUNTIME_HASH", "RUNTIME_DATA"} + ): + raise ArtifactError("Runtime bundle contains executable or unknown statements.") + name = statement.targets[0].id + if name in assignments: + raise ArtifactError(f"Runtime bundle repeats {name}.") + try: + assignments[name] = ast.literal_eval(statement.value) + except (ValueError, TypeError) as error: + raise ArtifactError(f"Runtime bundle {name} is not literal data.") from error + runtime_hash = assignments.get("RUNTIME_HASH") + encoded = assignments.get("RUNTIME_DATA") + if ( + not isinstance(runtime_hash, str) + or re.fullmatch(r"[0-9a-f]{64}", runtime_hash) is None + or not isinstance(encoded, str) + or not encoded + ): + raise ArtifactError("Runtime bundle data identities are missing or invalid.") + try: + archive = base64.b85decode(encoded.encode("ascii")) + except (UnicodeEncodeError, ValueError) as error: + raise ArtifactError("Runtime bundle contains invalid base85 archive data.") from error + if hashlib.sha256(archive).hexdigest() != runtime_hash: + raise ArtifactError("Runtime bundle archive bytes differ from RUNTIME_HASH.") + if len(archive) > _MAX_RUNTIME_ARCHIVE_BYTES: + raise ArtifactError("Runtime bundle archive exceeds its compressed size limit.") + return runtime_hash, archive + + +def _validate_runtime_bundle( + path: Path, + package_root: Path, + expected_hash: str, +) -> None: + """Bind the executable bundle archive to the validated package-source bytes.""" + + runtime_hash, archive = _decode_runtime_bundle(path) + if runtime_hash != expected_hash: + raise ArtifactError("Runtime bundle identity differs from provenance.") + expected = _runtime_archive_files(package_root) + try: + with ZipFile(io.BytesIO(archive)) as bundle: + members = bundle.infolist() + if not members or len(members) > _MAX_RUNTIME_ARCHIVE_MEMBERS: + raise ArtifactError("Runtime bundle archive has an invalid member count.") + member_names = [member.filename for member in members] + if len(member_names) != len(set(member_names)): + raise ArtifactError("Runtime bundle archive contains duplicate paths.") + if set(member_names) != set(expected): + raise ArtifactError( + "Runtime bundle archive inventory differs from packaged runtime sources." + ) + for member in members: + mode = member.external_attr >> 16 + expected_payload = expected[member.filename] + if ( + member.is_dir() + or member.flag_bits & 0x1 + or member.compress_type != ZIP_DEFLATED + or mode != 0o100644 + or member.file_size != len(expected_payload) + or member.file_size > _MAX_RUNTIME_SOURCE_BYTES + ): + raise ArtifactError( + f"Runtime bundle archive member is not canonical: {member.filename!r}." + ) + if bundle.read(member) != expected_payload: + raise ArtifactError( + f"Runtime bundle archive differs at {member.filename!r}." + ) + except ArtifactError: + raise + except (BadZipFile, KeyError, RuntimeError, OSError) as error: + raise ArtifactError("Runtime bundle archive is invalid.") from error + + +def _render_bootstrap(spec: ModelSpec, runtime_hash: str) -> str: + """Render the flat Transformers bridge to the bundled unchanged sources.""" + + grouped: dict[str, list[str]] = {} + class_paths = list(spec.auto_map.values()) + if spec.family.tokenizer_class is not None: + class_paths.append(spec.family.tokenizer_class) + for class_path in class_paths: + module_name, class_name = class_path.rsplit(".", maxsplit=1) + grouped.setdefault(module_name, []).append(class_name) + lines = [ + '"""Generated bridge to the embedded FastPLMs runtime sources."""', + "", + "import base64", + "import hashlib", + "import importlib", + "import importlib.util", + "import sys", + "import tempfile", + "from io import BytesIO", + "from pathlib import Path", + "from zipfile import ZIP_DEFLATED, ZipFile", + "", + "from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH", + "", + f'if RUNTIME_HASH != "{runtime_hash}":', + ' raise RuntimeError("FastPLMs runtime identity differs from the bridge.")', + "", + "_RUNTIME_TEMPORARIES = []", + "", + "def _archive_runtime_hashes(payload):", + " result = {}", + " with ZipFile(BytesIO(payload)) as archive:", + " for member in archive.infolist():", + " name = member.filename", + " parts = Path(name).parts", + " if (", + " member.is_dir()", + ' or "\\\\" in name', + " or not parts", + ' or parts[0] != "fastplms"', + " or len(parts) < 2", + ' or any(part in {"", ".", ".."} for part in parts)', + ' or Path(name).suffix in {".pyc", ".pyo"}', + " or member.flag_bits & 0x1", + " or member.compress_type != ZIP_DEFLATED", + " or member.external_attr >> 16 != 0o100644", + " ):", + ' raise RuntimeError("Embedded FastPLMs archive has an unsafe path.")', + " relative = Path(*parts[1:]).as_posix()", + " if relative in result:", + ' raise RuntimeError("Embedded FastPLMs archive repeats a path.")', + " result[relative] = hashlib.sha256(archive.read(member)).hexdigest()", + " return result", + "", + "def _ensure_runtime():", + ' payload = base64.b85decode("".join(RUNTIME_DATA))', + " if hashlib.sha256(payload).hexdigest() != RUNTIME_HASH:", + ' raise RuntimeError("Embedded FastPLMs runtime hash mismatch.")', + " expected = _archive_runtime_hashes(payload)", + ' temporary = tempfile.TemporaryDirectory(prefix="fastplms-artifact-runtime-")', + " try:", + " runtime_root = Path(temporary.name)", + " with ZipFile(BytesIO(payload)) as archive:", + " for member in archive.infolist():", + " target = runtime_root.joinpath(*Path(member.filename).parts)", + " target.parent.mkdir(parents=True, exist_ok=True)", + ' with target.open("xb") as handle:', + " handle.write(archive.read(member))", + ' package_root = runtime_root / "fastplms"', + " if _runtime_file_hashes(package_root) != expected:", + " raise RuntimeError(", + ' "Private FastPLMs runtime differs from the embedded archive."', + " )", + " except BaseException:", + " temporary.cleanup()", + " raise", + " _RUNTIME_TEMPORARIES.append(temporary)", + " return package_root", + "", + "def _runtime_file_hashes(package_root):", + " result = {}", + ' for path in sorted(package_root.rglob("*")):', + " relative = path.relative_to(package_root)", + " if path.is_symlink():", + ' raise RuntimeError("Private FastPLMs runtime contains a symlink.")', + " if path.is_dir():", + " continue", + ' if path.suffix in {".pyc", ".pyo"}:', + ' raise RuntimeError("Private FastPLMs runtime contains bytecode.")', + " if not path.is_file():", + ' raise RuntimeError("Private FastPLMs runtime contains a non-file entry.")', + " result[relative.as_posix()] = hashlib.sha256(path.read_bytes()).hexdigest()", + " return result", + "", + "def _extend_loaded_package_paths(package_root):", + " for name, module in list(sys.modules.items()):", + ' if name != "fastplms" and not name.startswith("fastplms."):', + " continue", + ' paths = getattr(module, "__path__", None)', + " if paths is None:", + " continue", + ' relative = name.split(".")[1:]', + " candidate = package_root.joinpath(*relative)", + " candidate_text = str(candidate)", + " if candidate.is_dir() and candidate_text not in paths:", + " paths.append(candidate_text)", + "", + "def _merge_runtime(package, package_root):", + " incoming = _runtime_file_hashes(package_root)", + ' known = getattr(package, "__fastplms_artifact_runtime_files__", None)', + " if not isinstance(known, dict):", + " raise RuntimeError(", + ' "A non-artifact fastplms module is already loaded. Load the Hub artifact "', + ' "in a separate Python process."', + " )", + " conflicts = sorted(", + " relative", + " for relative, digest in incoming.items()", + " if relative in known and known[relative] != digest", + " )", + " if conflicts:", + " raise RuntimeError(", + ' "FastPLMs artifacts contain incompatible runtime sources at "', + ' + ", ".join(repr(path) for path in conflicts[:5])', + ' + ". Load incompatible releases in separate Python processes."', + " )", + " known = dict(known)", + " known.update(incoming)", + ' package.__fastplms_artifact_runtime_files__ = known', + ' roots = list(getattr(package, "__fastplms_artifact_runtime_roots__", ()))', + " if str(package_root) not in roots:", + " roots.append(str(package_root))", + ' package.__fastplms_artifact_runtime_roots__ = tuple(roots)', + " temporaries = list(", + ' getattr(package, "__fastplms_artifact_runtime_temporaries__", ())', + " )", + " for temporary in _RUNTIME_TEMPORARIES:", + " if temporary not in temporaries:", + " temporaries.append(temporary)", + ' package.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries)', + ' hashes = set(getattr(package, "__fastplms_artifact_runtime_hashes__", ()))', + " hashes.add(RUNTIME_HASH)", + ' package.__fastplms_artifact_runtime_hashes__ = frozenset(hashes)', + " _extend_loaded_package_paths(package_root)", + " return package", + "", + "def _import_without_bytecode(module_name):", + " previous = sys.dont_write_bytecode", + " sys.dont_write_bytecode = True", + " try:", + " return importlib.import_module(module_name)", + " finally:", + " sys.dont_write_bytecode = previous", + "", + "def _install_runtime():", + ' package = sys.modules.get("fastplms")', + ' hashes = getattr(package, "__fastplms_artifact_runtime_hashes__", ())', + " if RUNTIME_HASH in hashes:", + " return package", + " package_root = _ensure_runtime()", + " if package is not None:", + " return _merge_runtime(package, package_root)", + " spec = importlib.util.spec_from_file_location(", + ' "fastplms",', + ' package_root / "__init__.py",', + " submodule_search_locations=[str(package_root)],", + " )", + " if spec is None or spec.loader is None:", + ' raise ImportError("Unable to load the embedded FastPLMs runtime.")', + " package = importlib.util.module_from_spec(spec)", + " package.__fastplms_artifact_runtime_hash__ = RUNTIME_HASH", + " package.__fastplms_artifact_runtime_hashes__ = frozenset({RUNTIME_HASH})", + " package.__fastplms_artifact_runtime_files__ = _runtime_file_hashes(package_root)", + " package.__fastplms_artifact_runtime_roots__ = (str(package_root),)", + " package.__fastplms_artifact_runtime_temporaries__ = tuple(", + " _RUNTIME_TEMPORARIES", + " )", + ' sys.modules["fastplms"] = package', + " previous = sys.dont_write_bytecode", + " sys.dont_write_bytecode = True", + " try:", + " try:", + " spec.loader.exec_module(package)", + " except BaseException:", + ' sys.modules.pop("fastplms", None)', + " raise", + " finally:", + " sys.dont_write_bytecode = previous", + " return package", + "", + "_install_runtime()", + ] + for module_name in sorted(grouped): + variable = f"_module_{len(lines)}" + lines.append(f'{variable} = _import_without_bytecode("{module_name}")') + for class_name in sorted(set(grouped[module_name])): + lines.extend( + ( + f"{class_name} = {variable}.{class_name}", + f"{class_name}.__module__ = __name__", + ) + ) + lines.append("") + return "\n".join(lines) + + +def _write_bootstrap(path: Path, spec: ModelSpec, runtime_hash: str) -> None: + """Write a flat Transformers bridge to the bundled unchanged sources.""" + + path.write_text( + _render_bootstrap(spec, runtime_hash), + encoding="utf-8", + newline="\n", + ) + + +def _validate_bootstrap(path: Path, spec: ModelSpec, runtime_hash: str) -> None: + try: + actual = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError) as error: + raise ArtifactError(f"Artifact bootstrap is missing or invalid: {path}") from error + if actual != _render_bootstrap(spec, runtime_hash): + raise ArtifactError("Artifact bootstrap differs from the current deterministic generator.") + + +def render_model_card(spec: ModelSpec) -> str: + """Render the canonical generated card used by documentation and artifacts.""" + + from tools.artifacts.generate_docs import render_model_card as render_canonical_card + + return render_canonical_card(spec, allow_generic_family=True) + + +def _validated_model_card_template( + source_root: Path, + spec: ModelSpec, + *, + release_tool_revision: str, + _allow_untracked_for_tests: bool = False, +) -> str: + """Read one immutable tracked card template, or render with immutable tools.""" + + source_root = source_root.resolve() + card_relative = f"model_cards/{spec.id}.md" + card_source = source_root.joinpath(*PurePosixPath(card_relative).parts) + if _allow_untracked_for_tests: + if card_source.is_file() and not card_source.is_symlink(): + return card_source.read_text(encoding="utf-8") + return render_model_card(spec) + git_metadata = source_root / ".git" + if git_metadata.exists() or git_metadata.is_symlink(): + if re.fullmatch(r"[0-9a-f]{40}", release_tool_revision) is None: + raise ArtifactError("Git model-card templates require a Git release-tool revision.") + command_prefix = ["git", "-c", f"safe.directory={source_root.as_posix()}"] + try: + status = subprocess.run( + [ + *command_prefix, + "status", + "--porcelain=v1", + "--untracked-files=all", + "--", + card_relative, + ], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if status: + raise ArtifactError( + "Artifact model-card template must be tracked and clean; scoped Git " + f"status: {status.replace(chr(10), '; ')}" + ) + tracked = subprocess.run( + [*command_prefix, "ls-files", "-z", "--", card_relative], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + except ArtifactError: + raise + except (OSError, subprocess.CalledProcessError) as error: + raise ArtifactError("Unable to validate the tracked model-card template.") from error + tracked_names = { + raw.decode("utf-8") for raw in tracked.split(b"\0") if raw + } + if tracked_names: + if tracked_names != {card_relative}: + raise ArtifactError("Git returned an ambiguous model-card template.") + payload = _git_archived_regular_files( + source_root, + release_tool_revision, + (card_relative,), + label="model-card template", + )[card_relative] + try: + return payload.decode("utf-8") + except UnicodeDecodeError as error: + raise ArtifactError("Tracked model-card template is not UTF-8.") from error + if card_source.exists() or card_source.is_symlink(): + raise ArtifactError(f"Artifact model-card template is untracked: {card_source}") + return render_model_card(spec) + + archive_marker = source_root / ARCHIVE_PROVENANCE_NAME + if archive_marker.exists() or archive_marker.is_symlink(): + try: + _diagnostic_head, inventory = validate_archived_root(source_root) + except SourceProvenanceError as error: + raise ArtifactError( + f"Git-free model-card template attestation is invalid: {error}" + ) from error + record = inventory.get(card_relative) + if record is None: + if card_source.exists() or card_source.is_symlink(): + raise ArtifactError( + f"Artifact model-card template is absent from archive provenance: {card_source}" + ) + return render_model_card(spec) + if ( + card_source.is_symlink() + or not card_source.is_file() + or record.get("mode") not in {"100644", "100755"} + ): + raise ArtifactError( + f"Attested model-card template is not a regular file: {card_source}" + ) + try: + payload = card_source.read_bytes() + except OSError as error: + raise ArtifactError( + f"Unable to read attested model-card template: {card_source}" + ) from error + if ( + record.get("size") != len(payload) + or record.get("sha256") != hashlib.sha256(payload).hexdigest() + ): + raise ArtifactError("Attested model-card template mutated during snapshot.") + try: + return payload.decode("utf-8") + except UnicodeDecodeError as error: + raise ArtifactError("Attested model-card template is not UTF-8.") from error + + raise ArtifactError( + "Artifact model-card templates require either a clean verifiable Git worktree " + "or a content-attested tracked remote source archive." + ) + + +def _materialize_model_card( + template: str, + *, + runtime_revision: str, + source_tree_sha256: str, + runtime_bundle_sha256: str, +) -> str: + """Bind an immutable card template to one exact packaged runtime identity.""" + + if re.fullmatch( + r"(?:[0-9a-f]{40}|source-tree-sha256:[0-9a-f]{64})", + runtime_revision, + ) is None: + raise ArtifactError(f"Invalid model-card runtime revision: {runtime_revision!r}") + for label, digest in ( + ("source-tree", source_tree_sha256), + ("runtime-bundle", runtime_bundle_sha256), + ): + if re.fullmatch(r"[0-9a-f]{64}", digest) is None: + raise ArtifactError(f"Invalid model-card {label} SHA-256: {digest!r}") + if _MODEL_CARD_RUNTIME_REVISION_PLACEHOLDER in template: + raise ArtifactError("Model-card template retains a runtime-revision placeholder.") + if template.count(_MODEL_CARD_RUNTIME_PROVENANCE) != 1: + raise ArtifactError("Model-card template lacks the canonical runtime provenance line.") + if template.count(_MODEL_CARD_DIGEST_PROVENANCE) != 1: + raise ArtifactError("Model-card template lacks the canonical runtime-digest line.") + materialized = template + if _MODEL_CARD_RUNTIME_REVISION_PLACEHOLDER in materialized: + raise ArtifactError("Materialized model card retains a runtime placeholder.") + return materialized.rstrip() + "\n" + + +def _validate_vendor_revisions(source_root: Path, registry: ModelRegistry, spec: ModelSpec) -> None: + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + checkout = source_root / source.path + if not checkout.is_dir(): + raise ArtifactError( + f"Official source {source_id!r} is not initialized. Run " + "'git submodule update --init --recursive'." + ) + checkout_git_metadata = checkout / ".git" + checkout_has_git = checkout_git_metadata.exists() or checkout_git_metadata.is_symlink() + source_git_metadata = source_root / ".git" + source_has_git = source_git_metadata.exists() or source_git_metadata.is_symlink() + if not checkout_has_git: + if source_has_git: + raise ArtifactError( + f"Official source {source_id!r} is not initialized. Run " + "'git submodule update --init --recursive'." + ) + try: + validate_archived_submodule( + source_root, + relative_path=source.path, + expected_revision=source.revision, + ) + except SourceProvenanceError as error: + raise ArtifactError( + f"Official source {source_id!r} has invalid archive provenance: {error}" + ) from error + continue + result = subprocess.run( + [ + "git", + "-c", + f"safe.directory={checkout.as_posix()}", + "-C", + str(checkout), + "rev-parse", + "HEAD", + ], + check=False, + capture_output=True, + text=True, + ) + revision = result.stdout.strip() + if result.returncode != 0 or revision != source.revision: + raise ArtifactError( + f"Official source {source_id!r} must be at {source.revision}; " + f"received {revision or result.stderr.strip()!r}." + ) + status = subprocess.run( + [ + "git", + "-c", + f"safe.directory={checkout.as_posix()}", + "-C", + str(checkout), + "status", + "--porcelain=v1", + "--untracked-files=all", + ], + check=False, + capture_output=True, + text=True, + ) + dirty_paths = status.stdout.strip() + if status.returncode != 0 or dirty_paths: + detail = dirty_paths or status.stderr.strip() or "unable to inspect worktree" + raise ArtifactError( + f"Official source {source_id!r} must have a clean worktree at " + f"{source.revision}; received {detail!r}." + ) + + +def _copy_licenses( + temporary: Path, + source_root: Path, + registry: ModelRegistry, + spec: ModelSpec, +) -> None: + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + for expected in source.distribution_files: + relative_path = PurePosixPath(expected.path) + license_source = (source_root / "LICENSES" / source_id).joinpath(*relative_path.parts) + license_target = (temporary / "LICENSES" / source_id).joinpath(*relative_path.parts) + _copy_canonical_legal_file(license_source, license_target) + _copy_canonical_legal_file( + source_root / "LICENSE", + temporary / "LICENSES" / "FastPLMs-Apache-2.0.txt", + ) + _copy_canonical_legal_file( + source_root / "THIRD_PARTY_NOTICES.md", + temporary / "THIRD_PARTY_NOTICES.md", + ) + + +def _checkpoint_provenance(source: CheckpointSource) -> dict[str, Any]: + return { + "repo_id": source.repo_id, + "revision": source.revision, + "files": {item.path: item.encoded for item in source.files}, + "unresolved_files": list(source.unresolved_files), + } + + +def _upstream_provenance( + registry: ModelRegistry, + spec: ModelSpec, +) -> list[dict[str, Any]]: + return [ + { + "id": source.id, + "license": source.license_expression, + "canonical_license_files": { + item.path: item.encoded for item in source.license_digests + }, + "distribution_files": { + item.path: item.encoded for item in source.distribution_files + }, + "path": source.path, + "revision": source.revision, + "url": source.url, + } + for source_id in spec.family.upstreams + for source in (registry.upstreams[source_id],) + ] + + +def _runtime_asset_provenance( + registry: ModelRegistry, + spec: ModelSpec, +) -> list[dict[str, Any]]: + return [ + { + "id": asset.id, + "repository": asset.repository, + "revision": asset.revision, + "path": asset.path, + "sha256": asset.sha256, + "size": asset.size, + "license": asset.license_expression, + "consumer_family": asset.consumer_family, + "trust_kind": asset.trust_kind, + "offline_behavior": asset.offline_behavior, + "cache_identity": hashlib.sha256( + ( + f"{asset.repository}@{asset.revision}:{asset.path}:" + f"{asset.sha256}:{asset.size}" + ).encode() + ).hexdigest(), + } + for asset in registry.runtime_assets.values() + if asset.consumer_family == spec.family.id + ] + + +def _expected_registry_provenance( + registry: ModelRegistry, + spec: ModelSpec, +) -> dict[str, Any]: + selected_checkpoint = spec.artifact_checkpoint + tokenizer_checkpoint = _tokenizer_checkpoint(registry, spec) + return { + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "generator": { + "name": "tools.artifacts.build", + "version": _ARTIFACT_GENERATOR_VERSION, + }, + "fastplms_version": __version__, + "model_id": spec.id, + "architecture": spec.family.architecture, + "auto_map": dict(spec.auto_map), + "tokenizer_class": spec.family.tokenizer_class, + "tokenizer_auto_map": ( + [ + "modeling_fastplms." + + spec.family.tokenizer_class.rsplit(".", maxsplit=1)[1], + None, + ] + if spec.family.tokenizer_class is not None + else None + ), + "bf16_execution": spec.family.bf16_execution, + **( + {"msa_conditioning": spec.msa_conditioning} + if spec.family.id == "esmfold2" + else {} + ), + "checkpoint_license": spec.family.checkpoint_license, + "weights_license_status": ( + "resolved" if spec.family.weights_publication_allowed else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + "hub_license_metadata": dict(spec.family.hub_license_metadata), + "legal_files": {item.path: item.encoded for item in registry.legal_files}, + "artifact_source": spec.artifact_source, + "artifact_checkpoint": _checkpoint_provenance(selected_checkpoint), + "weights_revision": selected_checkpoint.revision, + "fast_checkpoint": _checkpoint_provenance(spec.fast), + "official_checkpoint": _checkpoint_provenance(spec.official), + "tokenizer_checkpoint": ( + { + "repo_id": tokenizer_checkpoint.repo_id, + "revision": tokenizer_checkpoint.revision, + "files": { + item.path: item.encoded + for item in tokenizer_checkpoint.files + if PurePosixPath(item.path).name in _TOKENIZER_FILE_NAMES + }, + "unresolved_files": [ + path + for path in tokenizer_checkpoint.unresolved_files + if PurePosixPath(path).name in _TOKENIZER_FILE_NAMES + ], + } + if spec.family.tokenizer_mode == "tokenizer" + else None + ), + "oracle_assets": [ + { + "role": asset.role, + "path": asset.path, + "url": asset.url, + "sha256": asset.sha256, + "size": asset.size, + } + for asset in spec.oracle_assets + ], + "runtime_assets": _runtime_asset_provenance(registry, spec), + "state_transform": spec.family.state_transform, + "conversion": { + "id": spec.family.state_transform, + "record": spec.family.conversion_provenance, + }, + "conversion_equality_attestation": _conversion_equality_attestation(spec), + "upstreams": _upstream_provenance(registry, spec), + } + + +def _validate_registry_provenance( + provenance: Mapping[str, Any], + registry: ModelRegistry, + spec: ModelSpec, +) -> None: + try: + registered = registry[spec.id] + except KeyError as error: + raise ArtifactError(f"Artifact model {spec.id!r} is absent from the registry.") from error + if registered != spec: + raise ArtifactError(f"Artifact model {spec.id!r} differs from the supplied registry.") + expected = _expected_registry_provenance(registry, spec) + mismatches = [ + name for name, value in expected.items() if provenance.get(name) != value + ] + if mismatches: + raise ArtifactError( + "Artifact provenance differs from the current registry for fields: " + + ", ".join(sorted(mismatches)) + ) + + +def _provenance( + registry: ModelRegistry, + spec: ModelSpec, + canonical_weights: Mapping[str, Any], + *, + runtime_revision: str, + source_tree_sha256: str, + runtime_bundle_sha256: str, + release_tool_revision: str, + release_tool_sha256: str, +) -> dict[str, Any]: + selected_checkpoint = spec.artifact_checkpoint + return { + **_expected_registry_provenance(registry, spec), + "runtime_revision": runtime_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "attestations": { + "complete_artifact": { + "scope": "weights+runtime", + "weights_revision": selected_checkpoint.revision, + "runtime_revision": runtime_revision, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" + if spec.family.weights_publication_allowed + else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + }, + "runtime_update": { + "path": _RUNTIME_ATTESTATION_NAME, + "scope": "runtime-only", + "weights_repo_id": spec.fast.repo_id, + "weights_revision": spec.fast.revision, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" + if spec.family.weights_publication_allowed + else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + }, + }, + "canonical_weights": dict(canonical_weights), + } + + +def _tokenizer_checkpoint( + registry: ModelRegistry, + spec: ModelSpec, +) -> CheckpointSource: + """Resolve the manifest-declared official tokenizer identity for a model.""" + + if spec.tokenizer_source_id is None: + return spec.official + try: + return registry[spec.tokenizer_source_id].official + except KeyError as error: + raise ArtifactError( + f"Unknown tokenizer source {spec.tokenizer_source_id!r} for {spec.id}." + ) from error + + +def _checkpoint_identity_hash_fields( + repo_id: str, + revision: str, + files: Mapping[str, str], +) -> str: + payload = { + "repo_id": repo_id, + "revision": revision, + "files": [{"path": path, "digest": digest} for path, digest in sorted(files.items())], + } + return hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + + +def _checkpoint_identity_hash(source: CheckpointSource) -> str: + """Return a deterministic digest of one immutable checkpoint identity.""" + + return _checkpoint_identity_hash_fields( + source.repo_id, + source.revision, + {item.path: item.encoded for item in source.files}, + ) + + +def _conversion_equality_attestation(spec: ModelSpec) -> dict[str, Any] | None: + """Return the registry-owned conversion commitment for official-source state.""" + + if spec.artifact_source != "official": + return None + expected_state_sha256 = spec.canonical_state_sha256 + if ( + not isinstance(expected_state_sha256, str) + or re.fullmatch(r"[0-9a-f]{64}", expected_state_sha256) is None + ): + raise ArtifactError( + f"Official-source model {spec.id!r} lacks a canonical-state commitment." + ) + source = spec.artifact_checkpoint + payload: dict[str, Any] = { + "schema_version": _CONVERSION_ATTESTATION_SCHEMA_VERSION, + "model_id": spec.id, + "source_checkpoint": { + "repo_id": source.repo_id, + "revision": source.revision, + "identity_sha256": _checkpoint_identity_hash(source), + }, + "state_transform": spec.family.state_transform, + "conversion_record_sha256": hashlib.sha256( + spec.family.conversion_provenance.encode("utf-8") + ).hexdigest(), + "canonical_state": { + "schema_version": _CANONICAL_STATE_SCHEMA_VERSION, + "algorithm": "sha256", + "sha256": expected_state_sha256, + }, + } + payload["attestation_sha256"] = hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + return payload + + +def _content_manifest(root: Path) -> dict[str, str]: + result: dict[str, str] = {} + for path in _iter_files(root): + relative = path.relative_to(root).as_posix() + if relative != "artifact-manifest.json": + result[relative] = f"sha256:{hash_file(path)}" + return result + + +def _runtime_attestation( + root: Path, + spec: ModelSpec, + *, + weights_revision: str, + runtime_revision: str, + source_tree_sha256: str, + runtime_bundle_sha256: str, + release_tool_revision: str, + release_tool_sha256: str, +) -> dict[str, Any]: + """Create a deliberately runtime-scoped attestation for files-only updates.""" + + canonical_weights = _load_json_object_for_build(root / "provenance.json").get( + "canonical_weights" + ) + if not isinstance(canonical_weights, Mapping): + raise ArtifactError("Artifact provenance is missing canonical weight metadata.") + index = canonical_weights.get("index") + shards = canonical_weights.get("shards") + if not isinstance(index, str) or not isinstance(shards, Mapping): + raise ArtifactError("Artifact provenance has invalid canonical weight metadata.") + excluded = { + "artifact-manifest.json", + "provenance.json", + _RUNTIME_ATTESTATION_NAME, + index, + *(str(name) for name in shards), + } + files = { + relative_name: encoded + for relative_name, encoded in _content_manifest(root).items() + if ( + relative_name not in excluded + and not _is_weight_file(relative_name) + and _is_runtime_update_path(relative_name) + ) + } + if not files or not any(name.startswith("fastplms/") for name in files): + raise ArtifactError("Runtime attestation would contain no packaged runtime sources.") + return { + "schema_version": _RUNTIME_ATTESTATION_SCHEMA_VERSION, + "scope": "runtime-only", + "model_id": spec.id, + "weights": { + "repo_id": spec.fast.repo_id, + "revision": weights_revision, + }, + "runtime_revision": runtime_revision, + "source_tree_sha256": source_tree_sha256, + "runtime_bundle_sha256": runtime_bundle_sha256, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": ( + "resolved" if spec.family.weights_publication_allowed else "unresolved" + ), + "redistributable": spec.family.weights_publication_allowed, + "files": files, + } + + +def _load_json_object_for_build(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError(f"Unable to read JSON object: {path}") from error + if not isinstance(value, dict): + raise ArtifactError(f"JSON document must contain an object: {path}") + return value + + +def _resolve_artifact_manifest_path(root: Path, relative_name: str) -> Path: + """Resolve one portable manifest path while keeping it inside the artifact.""" + + try: + relative = _portable_relative_path(relative_name, "Artifact manifest path") + except RegistryError as error: + raise ArtifactError( + f"invalid artifact manifest path: {relative_name!r}" + ) from error + resolved = root.joinpath(*relative.parts).resolve() + try: + resolved.relative_to(root) + except ValueError as error: + raise ArtifactError( + f"artifact manifest path escapes the artifact root: {relative_name!r}" + ) from error + return resolved + + +def _artifact_build_slots( + output_root: Path, + repository_name: str, +) -> tuple[Path, Path, Path]: + """Return the only three sibling paths an artifact build may mutate.""" + + try: + relative = _portable_relative_path(repository_name, "Artifact repository name") + except RegistryError as error: + raise ArtifactError(f"Invalid artifact repository name: {repository_name!r}") from error + if ( + len(relative.parts) != 1 + or relative.name != repository_name + or repository_name.startswith(".") + ): + raise ArtifactError(f"Invalid artifact repository name: {repository_name!r}") + slots = ( + output_root / repository_name, + output_root / f".{repository_name}.tmp", + output_root / f".{repository_name}.backup", + ) + for slot in slots: + if slot.parent != output_root: + raise ArtifactError(f"Artifact build path escapes the output root: {slot}") + return slots + + +def _is_artifact_path_link(path: Path) -> bool: + is_junction = getattr(path, "is_junction", None) + return path.is_symlink() or bool(callable(is_junction) and is_junction()) + + +def _assert_artifact_directory_slot(path: Path, output_root: Path, label: str) -> None: + """Reject escaped, linked, or non-directory artifact transaction slots.""" + + if path.parent != output_root: + raise ArtifactError(f"{label} escapes the resolved output root: {path}") + if _is_artifact_path_link(path): + raise ArtifactError(f"{label} must not be a symlink or junction: {path}") + if path.exists() and not path.is_dir(): + raise ArtifactError(f"{label} must be a directory: {path}") + + +def _remove_artifact_directory(path: Path, output_root: Path, label: str) -> None: + """Remove one checked sibling transaction directory without following links.""" + + _assert_artifact_directory_slot(path, output_root, label) + if not path.exists(): + return + try: + shutil.rmtree(path) + except OSError as error: + raise ArtifactError(f"Unable to remove {label}: {path}") from error + + +def _atomic_artifact_rename( + source: Path, + destination: Path, + output_root: Path, +) -> None: + """Atomically rename one artifact slot to another on the same filesystem.""" + + _assert_artifact_directory_slot(source, output_root, "Artifact rename source") + _assert_artifact_directory_slot(destination, output_root, "Artifact rename destination") + if not source.is_dir(): + raise ArtifactError(f"Artifact rename source is missing: {source}") + if destination.exists(): + raise ArtifactError(f"Artifact rename destination already exists: {destination}") + try: + source.rename(destination) + except OSError as error: + raise ArtifactError( + f"Unable to atomically rename artifact slot {source} to {destination}." + ) from error + + +def _validate_artifact_slot( + path: Path, + output_root: Path, + spec: ModelSpec, + registry: ModelRegistry, + label: str, +) -> None: + _assert_artifact_directory_slot(path, output_root, label) + if not path.is_dir(): + raise ArtifactError(f"{label} is missing: {path}") + try: + validate_artifact(path, spec=spec, registry=registry) + except ArtifactError as error: + raise ArtifactError(f"{label} is not a valid recoverable artifact: {path}") from error + + +def _recover_artifact_transaction( + destination: Path, + temporary: Path, + backup: Path, + output_root: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> None: + """Recover or clean a transaction left by an interrupted artifact build.""" + + for path, label in ( + (destination, "Artifact destination"), + (temporary, "Artifact temporary directory"), + (backup, "Artifact rollback backup"), + ): + _assert_artifact_directory_slot(path, output_root, label) + + destination_validated = False + if backup.exists(): + if destination.exists(): + _validate_artifact_slot( + destination, + output_root, + spec, + registry, + "Artifact destination beside a stale rollback backup", + ) + destination_validated = True + _remove_artifact_directory( + backup, + output_root, + "stale artifact rollback backup", + ) + else: + _validate_artifact_slot( + backup, + output_root, + spec, + registry, + "Interrupted artifact rollback backup", + ) + _atomic_artifact_rename(backup, destination, output_root) + destination_validated = True + + if temporary.exists(): + if not destination.exists(): + raise ArtifactError( + "An incomplete artifact temporary directory exists without a valid destination " + f"or rollback backup: {temporary}" + ) + if not destination_validated: + _validate_artifact_slot( + destination, + output_root, + spec, + registry, + "Artifact destination beside a stale temporary directory", + ) + _remove_artifact_directory( + temporary, + output_root, + "stale artifact temporary directory", + ) + + +def _restore_artifact_backup( + destination: Path, + temporary: Path, + backup: Path, + output_root: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> None: + """Restore the prior artifact while retaining every valid copy until validation.""" + + if not backup.exists(): + raise ArtifactError(f"Artifact rollback backup is missing: {backup}") + if destination.exists(): + if temporary.exists(): + raise ArtifactError( + "Cannot quarantine a failed replacement because the temporary slot exists." + ) + _atomic_artifact_rename(destination, temporary, output_root) + _atomic_artifact_rename(backup, destination, output_root) + _validate_artifact_slot( + destination, + output_root, + spec, + registry, + "Restored artifact rollback backup", + ) + if temporary.exists(): + _remove_artifact_directory( + temporary, + output_root, + "failed artifact replacement", + ) + + +def _commit_artifact_transaction( + destination: Path, + temporary: Path, + backup: Path, + output_root: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> None: + """Install a validated temporary artifact with rollback-safe directory swaps.""" + + for path, label in ( + (destination, "Artifact destination"), + (temporary, "Validated artifact temporary directory"), + (backup, "Artifact rollback backup"), + ): + _assert_artifact_directory_slot(path, output_root, label) + if not temporary.is_dir(): + raise ArtifactError(f"Validated artifact temporary directory is missing: {temporary}") + if backup.exists(): + raise ArtifactError(f"Artifact rollback backup was not recovered: {backup}") + + if not destination.exists(): + _atomic_artifact_rename(temporary, destination, output_root) + return + + try: + _atomic_artifact_rename(destination, backup, output_root) + except BaseException as error: + try: + if backup.exists() and not destination.exists(): + _restore_artifact_backup( + destination, + temporary, + backup, + output_root, + spec, + registry, + ) + elif destination.exists(): + _validate_artifact_slot( + destination, + output_root, + spec, + registry, + "Artifact destination after a failed backup rename", + ) + _remove_artifact_directory( + temporary, + output_root, + "aborted artifact replacement", + ) + except BaseException as rollback_error: + raise ArtifactError( + "Unable to move the prior artifact into its rollback backup, and automatic " + "recovery did not complete. The transaction slots were retained." + ) from rollback_error + raise ArtifactError( + "Unable to move the prior artifact into its rollback backup; the prior artifact " + "was restored." + ) from error + + try: + _atomic_artifact_rename(temporary, destination, output_root) + except BaseException as error: + try: + _restore_artifact_backup( + destination, + temporary, + backup, + output_root, + spec, + registry, + ) + except BaseException as rollback_error: + raise ArtifactError( + "Artifact replacement failed and automatic rollback did not complete. " + "The transaction slots were retained for deterministic recovery." + ) from rollback_error + raise ArtifactError( + "Artifact replacement failed; the prior validated artifact was restored." + ) from error + + try: + _remove_artifact_directory( + backup, + output_root, + "completed artifact rollback backup", + ) + except ArtifactError as error: + raise ArtifactError( + "The new artifact was installed, but its rollback backup could not be removed. " + "The next invocation will validate the destination before cleanup." + ) from error + + +def build_artifact( + spec: ModelSpec, + registry: ModelRegistry, + checkpoint_dir: Path, + output_root: Path, + source_root: Path, + *, + tokenizer_dir: Path | None = None, + replace: bool = False, + _allow_untracked_runtime_for_tests: bool = False, +) -> Path: + """Build one local artifact from an already verified checkpoint snapshot.""" + + try: + registered_spec = registry[spec.id] + except KeyError as error: + raise ArtifactError(f"Model {spec.id!r} is absent from the supplied registry.") from error + if registered_spec != spec: + raise ArtifactError( + f"Model {spec.id!r} differs from the current supplied registry contract." + ) + try: + registry.require_resolved(spec.id) + except (KeyError, RegistryError) as error: + raise ArtifactError(str(error)) from error + checkpoint_dir = checkpoint_dir.resolve() + source_root = source_root.resolve() + validate_repository_legal_inventory(source_root, registry, spec) + selected_checkpoint = spec.artifact_checkpoint + tokenizer_checkpoint = _tokenizer_checkpoint(registry, spec) + resolved_tokenizer_dir: Path | None = None + if spec.family.tokenizer_mode == "tokenizer": + if tokenizer_dir is not None: + resolved_tokenizer_dir = tokenizer_dir.resolve() + elif selected_checkpoint == tokenizer_checkpoint: + resolved_tokenizer_dir = checkpoint_dir + else: + raise ArtifactError( + f"{spec.id} packages a FastPLMs checkpoint and requires the pinned official " + "tokenizer snapshot via tokenizer_dir/--tokenizer-dir." + ) + if _is_artifact_path_link(output_root): + raise ArtifactError( + f"Artifact output root must not be a symlink or junction: {output_root}" + ) + output_root = output_root.resolve() + output_root.mkdir(parents=True, exist_ok=True) + repository_parts = spec.fast.repo_id.split("/", maxsplit=1) + if len(repository_parts) != 2: + raise ArtifactError(f"Invalid artifact repository id: {spec.fast.repo_id!r}") + repository_name = repository_parts[1] + destination, temporary, backup = _artifact_build_slots(output_root, repository_name) + _recover_artifact_transaction( + destination, + temporary, + backup, + output_root, + spec, + registry, + ) + if destination.exists() and not replace: + raise ArtifactError(f"Artifact already exists: {destination}") + temporary.mkdir(parents=True) + try: + _copy_checkpoint_assets(checkpoint_dir, temporary, selected_checkpoint) + if resolved_tokenizer_dir is not None: + _copy_official_tokenizer_assets( + resolved_tokenizer_dir, + temporary, + tokenizer_checkpoint, + ) + _configure_custom_tokenizer(temporary, spec) + package_target = temporary / "fastplms" + runtime_revision, runtime_payloads, expected_source_tree_sha256 = ( + _validated_runtime_snapshot( + source_root, + registry, + spec, + _allow_untracked_for_tests=_allow_untracked_runtime_for_tests, + ) + ) + release_tool_revision, release_tool_sha256, release_tool_payloads = ( + _validated_release_tool_snapshot( + source_root, + _allow_untracked_for_tests=_allow_untracked_runtime_for_tests, + ) + ) + canonical_weights = canonicalize_checkpoint_weights( + checkpoint_dir, + selected_checkpoint, + temporary, + state_transform=spec.family.state_transform, + source_is_canonical=spec.artifact_source == "fast", + ) + conversion_attestation = _conversion_equality_attestation(spec) + if conversion_attestation is not None and canonical_weights.get("state_digest") != ( + conversion_attestation["canonical_state"] + ): + raise ArtifactError( + f"Canonical state for {spec.id} differs from the registry-owned " + "conversion equality commitment." + ) + _write_runtime_snapshot(package_target, runtime_payloads) + source_tree_sha256 = _tree_sha256(package_target) + if source_tree_sha256 != expected_source_tree_sha256: + raise ArtifactError( + "Packaged runtime sources differ from the validated tracked snapshot." + ) + + config_path = temporary / "config.json" + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError(f"Unable to read checkpoint config: {config_path}") from error + if not isinstance(config, dict): + raise ArtifactError("Checkpoint config.json must contain a JSON object.") + _apply_artifact_config_contract(spec, config) + config["auto_map"] = _artifact_auto_map(spec) + config["fastplms_model_id"] = spec.id + config["fastplms_checkpoint_repo_id"] = selected_checkpoint.repo_id + config["fastplms_checkpoint_revision"] = selected_checkpoint.revision + config["fastplms_checkpoint_hash"] = _checkpoint_identity_hash(selected_checkpoint) + config["fastplms_weights_revision"] = selected_checkpoint.revision + config["fastplms_runtime_revision"] = runtime_revision + config["fastplms_source_tree_sha256"] = source_tree_sha256 + config["fastplms_release_tool_revision"] = release_tool_revision + config["fastplms_release_tool_sha256"] = release_tool_sha256 + _write_json(config_path, config) + runtime_hash = _write_runtime_bundle( + temporary / "fastplms_bundle.py", + package_target, + ) + config["fastplms_runtime_bundle_sha256"] = runtime_hash + _write_json(config_path, config) + _write_bootstrap(temporary / "modeling_fastplms.py", spec, runtime_hash) + + card_template = _validated_model_card_template( + source_root, + spec, + release_tool_revision=release_tool_revision, + _allow_untracked_for_tests=_allow_untracked_runtime_for_tests, + ) + try: + card_license = parse_hub_license_metadata(card_template) + except ValueError as error: + raise ArtifactError(f"Model-card template metadata is invalid: {error}") from error + if card_license != dict(spec.family.hub_license_metadata): + raise ArtifactError( + "Model-card template license metadata differs from models.toml." + ) + card_text = _materialize_model_card( + card_template, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_hash, + ) + (temporary / "requirements.txt").write_text( + _render_artifact_requirements(spec, release_tool_payloads), + encoding="utf-8", + newline="\n", + ) + (temporary / "README.md").write_text( + card_text, + encoding="utf-8", + newline="\n", + ) + _copy_licenses(temporary, source_root, registry, spec) + _write_json( + temporary / "provenance.json", + _provenance( + registry, + spec, + canonical_weights, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_hash, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + ), + ) + _write_json( + temporary / _RUNTIME_ATTESTATION_NAME, + _runtime_attestation( + temporary, + spec, + weights_revision=spec.fast.revision, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_hash, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + ), + ) + _write_json(temporary / "artifact-manifest.json", _content_manifest(temporary)) + validate_artifact(temporary, spec=spec, registry=registry) + except BaseException: + if temporary.exists(): + _remove_artifact_directory( + temporary, + output_root, + "failed artifact temporary directory", + ) + raise + _commit_artifact_transaction( + destination, + temporary, + backup, + output_root, + spec, + registry, + ) + return destination + + +def build_local_artifact( + model_id: str, + checkpoint_dir: Path, + output_root: Path = Path("dist/hub"), + source_root: Path | None = None, + *, + tokenizer_dir: Path | None = None, + replace: bool = False, +) -> Path: + """Validate provenance and build one manifest-selected local artifact.""" + + registry = get_model_registry() + try: + spec = registry[model_id] + except KeyError as error: + raise ArtifactError(f"Unknown model ID: {model_id!r}") from error + root = source_root or Path(__file__).resolve().parents[2] + _validate_vendor_revisions(root.resolve(), registry, spec) + return build_artifact( + spec=spec, + registry=registry, + checkpoint_dir=checkpoint_dir, + output_root=output_root, + source_root=root, + tokenizer_dir=tokenizer_dir, + replace=replace, + ) + + +def validate_artifact( + path: Path, + *, + spec: ModelSpec | None = None, + registry: ModelRegistry | None = None, +) -> None: + """Verify an artifact and optionally bind it to a current registry model.""" + + if registry is not None and spec is None: + raise ArtifactError("Artifact registry validation requires a selected model spec.") + path = path.resolve() + manifest_path = path / "artifact-manifest.json" + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError(f"Unable to read artifact manifest: {manifest_path}") from error + if not isinstance(manifest, dict) or not manifest: + raise ArtifactError("artifact-manifest.json must contain a non-empty object.") + failures: list[str] = [] + required_paths = { + "README.md", + "config.json", + "provenance.json", + "requirements.txt", + _RUNTIME_ATTESTATION_NAME, + "THIRD_PARTY_NOTICES.md", + "LICENSES/FastPLMs-Apache-2.0.txt", + _WEIGHT_INDEX, + } + missing_required = sorted(required_paths.difference(manifest)) + if missing_required: + failures.append(f"missing required artifact entries: {', '.join(missing_required)}") + weight_validation_attempted = False + for relative_name, encoded_digest in sorted(manifest.items()): + if not isinstance(relative_name, str) or not isinstance(encoded_digest, str): + failures.append("manifest keys and values must be strings") + continue + try: + algorithm, expected = encoded_digest.split(":", maxsplit=1) + except ValueError: + failures.append(f"invalid digest entry for {relative_name}") + continue + try: + artifact_file = _resolve_artifact_manifest_path(path, relative_name) + except ArtifactError as error: + failures.append(str(error)) + continue + if not artifact_file.is_file(): + failures.append(f"missing {relative_name}") + continue + actual = hash_file(artifact_file, algorithm) + if actual != expected: + failures.append(f"digest mismatch for {relative_name}") + unlisted = sorted(set(_content_manifest(path)).difference(manifest)) + if unlisted: + failures.append(f"unlisted files: {', '.join(unlisted)}") + try: + config = json.loads((path / "config.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + config = None + failures.append("config.json is missing or invalid") + if not isinstance(config, dict): + config = None + failures.append("config.json must contain an object") + provenance_path = path / "provenance.json" + try: + card_text = (path / "README.md").read_text(encoding="utf-8") + card_license = parse_hub_license_metadata(card_text) + except (OSError, ValueError) as error: + card_text = None + card_license = None + failures.append(f"README.md has invalid Hub license metadata: {error}") + try: + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + provenance = None + failures.append("provenance.json is missing or invalid") + if isinstance(provenance, dict): + if spec is not None and registry is not None: + try: + _validate_registry_provenance(provenance, registry, spec) + except ArtifactError as error: + failures.append(str(error)) + if provenance.get("schema_version") != _PROVENANCE_SCHEMA_VERSION: + failures.append("provenance schema version is missing or unsupported") + if provenance.get("generator") != { + "name": "tools.artifacts.build", + "version": _ARTIFACT_GENERATOR_VERSION, + }: + failures.append("artifact generator identity is missing or unsupported") + if spec is not None and provenance.get("model_id") != spec.id: + failures.append("artifact model identity differs from the current registry") + weights_license_status = provenance.get("weights_license_status") + redistributable = provenance.get("redistributable") + if ( + weights_license_status not in {"resolved", "unresolved"} + or not isinstance(redistributable, bool) + or redistributable != (weights_license_status == "resolved") + ): + failures.append("weight-license publication status is missing or invalid") + if spec is not None and redistributable != spec.family.weights_publication_allowed: + failures.append("weight-license publication status differs from the registry") + raw_hub_license = provenance.get("hub_license_metadata") + try: + if not isinstance(raw_hub_license, Mapping): + raise TypeError("Hub license metadata must be a mapping") + provenance_license = validate_hub_license_metadata(raw_hub_license) + except (TypeError, ValueError) as error: + provenance_license = None + failures.append(f"Hub license provenance is missing or invalid: {error}") + if card_license is not None and provenance_license != card_license: + failures.append("README Hub license metadata differs from provenance") + bf16_execution = provenance.get("bf16_execution") + if bf16_execution not in _BF16_EXECUTION_POLICIES: + failures.append("BF16 execution provenance is missing or invalid") + elif card_text is not None and f"`{bf16_execution}`" not in card_text: + failures.append("README BF16 execution policy differs from provenance") + artifact_source = provenance.get("artifact_source") + artifact_checkpoint = provenance.get("artifact_checkpoint") + if artifact_source not in {"fast", "official"}: + failures.append("artifact_source is missing or invalid") + if ( + not isinstance(artifact_checkpoint, dict) + or not isinstance(artifact_checkpoint.get("repo_id"), str) + or not artifact_checkpoint.get("repo_id") + or not isinstance(artifact_checkpoint.get("revision"), str) + or re.fullmatch( + r"[0-9a-f]{40}", artifact_checkpoint.get("revision", "") + ) + is None + ): + failures.append("selected artifact checkpoint provenance is missing") + elif config is not None: + checkpoint_files = artifact_checkpoint.get("files") + if not isinstance(checkpoint_files, dict) or any( + not isinstance(name, str) or not isinstance(digest, str) + for name, digest in checkpoint_files.items() + ): + failures.append("selected artifact checkpoint file identities are invalid") + else: + expected_identity = _checkpoint_identity_hash_fields( + artifact_checkpoint["repo_id"], + artifact_checkpoint["revision"], + checkpoint_files, + ) + expected_config_identity = { + "fastplms_model_id": provenance.get("model_id"), + "fastplms_checkpoint_repo_id": artifact_checkpoint["repo_id"], + "fastplms_checkpoint_revision": artifact_checkpoint["revision"], + "fastplms_checkpoint_hash": expected_identity, + "fastplms_weights_revision": artifact_checkpoint["revision"], + "fastplms_runtime_revision": provenance.get("runtime_revision"), + "fastplms_source_tree_sha256": provenance.get("source_tree_sha256"), + "fastplms_runtime_bundle_sha256": provenance.get( + "runtime_bundle_sha256" + ), + "fastplms_release_tool_revision": provenance.get( + "release_tool_revision" + ), + "fastplms_release_tool_sha256": provenance.get( + "release_tool_sha256" + ), + } + if spec is not None and spec.family.id == "esmfold2": + expected_config_identity["msa_conditioning"] = spec.msa_conditioning + msa_encoder = config.get("msa_encoder") + if not isinstance(msa_encoder, dict) or ( + msa_encoder.get("enabled") != spec.msa_conditioning + ): + failures.append( + "config MSA encoder policy differs from the current registry" + ) + if any( + config.get(name) != expected + for name, expected in expected_config_identity.items() + ): + failures.append("config packaging identity differs from checkpoint provenance") + if provenance.get("weights_revision") != ( + artifact_checkpoint.get("revision") + if isinstance(artifact_checkpoint, dict) + else None + ): + failures.append("weights revision differs from checkpoint provenance") + runtime_revision = provenance.get("runtime_revision") + source_tree_sha256 = provenance.get("source_tree_sha256") + runtime_bundle_sha256 = provenance.get("runtime_bundle_sha256") + release_tool_revision = provenance.get("release_tool_revision") + release_tool_sha256 = provenance.get("release_tool_sha256") + if not isinstance(runtime_revision, str) or re.fullmatch( + r"(?:[0-9a-f]{40}|source-tree-sha256:[0-9a-f]{64})", + runtime_revision, + ) is None: + failures.append("runtime revision provenance is missing") + if not isinstance(source_tree_sha256, str) or re.fullmatch( + r"[0-9a-f]{64}", source_tree_sha256 + ) is None: + failures.append("runtime source-tree digest is missing or invalid") + else: + try: + actual_source_tree_sha256 = _tree_sha256(path / "fastplms") + except ArtifactError: + actual_source_tree_sha256 = None + if actual_source_tree_sha256 != source_tree_sha256: + failures.append("runtime source-tree digest differs from packaged sources") + if ( + isinstance(runtime_revision, str) + and runtime_revision.startswith("source-tree-sha256:") + and runtime_revision != f"source-tree-sha256:{source_tree_sha256}" + ): + failures.append("content-addressed runtime revision differs from source tree") + if not isinstance(runtime_bundle_sha256, str) or re.fullmatch( + r"[0-9a-f]{64}", runtime_bundle_sha256 + ) is None: + failures.append("runtime bundle digest is missing or invalid") + else: + try: + _validate_runtime_bundle( + path / "fastplms_bundle.py", + path / "fastplms", + runtime_bundle_sha256, + ) + if spec is not None: + _validate_bootstrap( + path / "modeling_fastplms.py", + spec, + runtime_bundle_sha256, + ) + except ArtifactError as error: + failures.append(str(error)) + if not isinstance(release_tool_revision, str) or re.fullmatch( + r"(?:[0-9a-f]{40}|release-tools-sha256:[0-9a-f]{64})", + release_tool_revision, + ) is None: + failures.append("release-tool revision provenance is missing or invalid") + if not isinstance(release_tool_sha256, str) or re.fullmatch( + r"[0-9a-f]{64}", + release_tool_sha256, + ) is None: + failures.append("release-tool digest provenance is missing or invalid") + elif ( + isinstance(release_tool_revision, str) + and release_tool_revision.startswith("release-tools-sha256:") + and release_tool_revision != f"release-tools-sha256:{release_tool_sha256}" + ): + failures.append("content-addressed release-tool revision differs from its digest") + if card_text is not None: + if _MODEL_CARD_RUNTIME_REVISION_PLACEHOLDER in card_text: + failures.append("README retains an unresolved runtime-revision placeholder") + expected_card_lines = ( + _MODEL_CARD_RUNTIME_PROVENANCE, + _MODEL_CARD_DIGEST_PROVENANCE, + ) + if any(line not in card_text for line in expected_card_lines): + failures.append("README runtime identity differs from provenance") + attestations = provenance.get("attestations") + complete_attestation = ( + attestations.get("complete_artifact") if isinstance(attestations, dict) else None + ) + runtime_attestation_record = ( + attestations.get("runtime_update") if isinstance(attestations, dict) else None + ) + fast_checkpoint = provenance.get("fast_checkpoint") + expected_attestations = { + "complete_artifact": { + "scope": "weights+runtime", + "weights_revision": provenance.get("weights_revision"), + "runtime_revision": runtime_revision, + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": weights_license_status, + "redistributable": redistributable, + }, + "runtime_update": { + "path": _RUNTIME_ATTESTATION_NAME, + "scope": "runtime-only", + "weights_repo_id": ( + fast_checkpoint.get("repo_id") + if isinstance(fast_checkpoint, dict) + else None + ), + "weights_revision": ( + fast_checkpoint.get("revision") + if isinstance(fast_checkpoint, dict) + else None + ), + "release_tool_revision": release_tool_revision, + "release_tool_sha256": release_tool_sha256, + "weights_license_status": weights_license_status, + "redistributable": redistributable, + }, + } + if ( + not isinstance(complete_attestation, dict) + or not isinstance(runtime_attestation_record, dict) + or attestations != expected_attestations + ): + failures.append("scoped artifact attestations are missing or invalid") + if spec is not None: + expected_checkpoint = ( + spec.fast if provenance.get("artifact_source") == "fast" else spec.official + ) + expected_record = { + "repo_id": expected_checkpoint.repo_id, + "revision": expected_checkpoint.revision, + "files": {item.path: item.encoded for item in expected_checkpoint.files}, + "unresolved_files": list(expected_checkpoint.unresolved_files), + } + if artifact_checkpoint != expected_record: + failures.append("artifact checkpoint differs from the current registry") + canonical_weights = provenance.get("canonical_weights") + canonical_state = ( + canonical_weights.get("state_digest") + if isinstance(canonical_weights, dict) + else None + ) + if ( + not isinstance(canonical_weights, dict) + or canonical_weights.get("format") != "safetensors" + or canonical_weights.get("index") != _WEIGHT_INDEX + or canonical_weights.get("source_schema") not in {"canonical", "official"} + or not isinstance(canonical_weights.get("state_transform"), str) + or not canonical_weights.get("state_transform") + or not isinstance(canonical_weights.get("shards"), dict) + or not canonical_weights.get("shards") + or not isinstance(canonical_state, dict) + or canonical_state.get("schema_version") != _CANONICAL_STATE_SCHEMA_VERSION + or canonical_state.get("algorithm") != "sha256" + or not isinstance(canonical_state.get("sha256"), str) + or re.fullmatch(r"[0-9a-f]{64}", canonical_state.get("sha256", "")) is None + ): + failures.append("canonical weight provenance is missing") + else: + expected_index_digest = canonical_weights.get("index_digest") + expected_shards = cast(dict[str, str], canonical_weights["shards"]) + if manifest.get(_WEIGHT_INDEX) != expected_index_digest: + failures.append("canonical weight index digest differs from artifact manifest") + if any(manifest.get(name) != digest for name, digest in expected_shards.items()): + failures.append("canonical shard digests differ from artifact manifest") + weight_validation_attempted = True + try: + validate_weight_artifact( + path, + expected_state_sha256=canonical_state["sha256"], + ) + except ArtifactError as error: + failures.append(str(error)) + conversion_attestation = provenance.get("conversion_equality_attestation") + if artifact_source == "official": + if spec is None or registry is None: + failures.append( + "official-source artifact validation requires a current registry commitment" + ) + else: + try: + expected_attestation = _conversion_equality_attestation(spec) + except ArtifactError as error: + failures.append(str(error)) + else: + if expected_attestation is None: + failures.append( + "official-source artifact has no registry conversion commitment" + ) + elif conversion_attestation != expected_attestation: + failures.append( + "conversion equality attestation differs from the current registry" + ) + elif isinstance(canonical_state, dict) and ( + canonical_state != expected_attestation["canonical_state"] + ): + failures.append( + "canonical state differs from the registry conversion commitment" + ) + elif conversion_attestation is not None: + failures.append("canonical-source artifact has an unexpected conversion attestation") + selected_record = provenance.get(f"{artifact_source}_checkpoint") + if isinstance(artifact_checkpoint, dict) and artifact_checkpoint != selected_record: + failures.append("selected artifact checkpoint differs from source record") + conversion = provenance.get("conversion") + if ( + not isinstance(conversion, dict) + or not isinstance(conversion.get("id"), str) + or not conversion.get("id") + or not isinstance(conversion.get("record"), str) + or not conversion.get("record") + ): + failures.append("conversion provenance is missing") + elif isinstance(canonical_weights, dict) and ( + canonical_weights.get("state_transform") != conversion.get("id") + ): + failures.append("canonical weights did not use the declared conversion") + oracle_assets = provenance.get("oracle_assets") + if not isinstance(oracle_assets, list): + failures.append("oracle asset provenance is missing") + else: + required_oracle_fields = {"role", "path", "url", "sha256", "size"} + for asset in oracle_assets: + if not isinstance(asset, dict) or set(asset) != required_oracle_fields: + failures.append("invalid oracle asset provenance entry") + continue + if ( + not isinstance(asset["role"], str) + or not isinstance(asset["path"], str) + or not isinstance(asset["url"], str) + or not isinstance(asset["sha256"], str) + or not isinstance(asset["size"], int) + ): + failures.append("invalid oracle asset provenance value") + runtime_assets = provenance.get("runtime_assets") + if not isinstance(runtime_assets, list): + failures.append("runtime asset provenance is missing") + else: + required_runtime_asset_fields = { + "id", + "repository", + "revision", + "path", + "sha256", + "size", + "license", + "consumer_family", + "trust_kind", + "offline_behavior", + "cache_identity", + } + for asset in runtime_assets: + if not isinstance(asset, dict) or set(asset) != required_runtime_asset_fields: + failures.append("invalid runtime asset provenance entry") + continue + cache_material = ( + f"{asset['repository']}@{asset['revision']}:{asset['path']}:" + f"{asset['sha256']}:{asset['size']}" + ).encode() + if ( + not isinstance(asset["id"], str) + or not isinstance(asset["repository"], str) + or re.fullmatch(r"[0-9a-f]{40}", str(asset["revision"])) is None + or not isinstance(asset["path"], str) + or re.fullmatch(r"[0-9a-f]{64}", str(asset["sha256"])) is None + or isinstance(asset["size"], bool) + or not isinstance(asset["size"], int) + or asset["size"] <= 0 + or not isinstance(asset["license"], str) + or not asset["license"] + or asset["trust_kind"] != "hash_pinned_pickle" + or asset["offline_behavior"] != "requires_cached_verified_file" + or asset["cache_identity"] != hashlib.sha256(cache_material).hexdigest() + ): + failures.append("invalid runtime asset provenance value") + tokenizer_checkpoint = provenance.get("tokenizer_checkpoint") + tokenizer_auto_map = provenance.get("tokenizer_auto_map") + if tokenizer_auto_map is not None: + try: + tokenizer_config = json.loads( + (path / "tokenizer_config.json").read_text(encoding="utf-8") + ) + except (OSError, json.JSONDecodeError): + tokenizer_config = None + configured_auto_map = ( + tokenizer_config.get("auto_map", {}).get("AutoTokenizer") + if isinstance(tokenizer_config, dict) + and isinstance(tokenizer_config.get("auto_map"), dict) + else None + ) + if configured_auto_map != tokenizer_auto_map: + failures.append("custom tokenizer AutoTokenizer mapping differs from provenance") + if tokenizer_checkpoint is not None: + tokenizer_files = ( + tokenizer_checkpoint.get("files") + if isinstance(tokenizer_checkpoint, dict) + else None + ) + if ( + not isinstance(tokenizer_checkpoint, dict) + or not isinstance(tokenizer_checkpoint.get("repo_id"), str) + or not isinstance(tokenizer_checkpoint.get("revision"), str) + or not isinstance(tokenizer_files, dict) + or not tokenizer_files + ): + failures.append("official tokenizer provenance is missing or invalid") + else: + for relative_name, encoded_digest in tokenizer_files.items(): + if PurePosixPath(relative_name).name not in _TOKENIZER_FILE_NAMES: + failures.append( + f"tokenizer provenance contains a non-tokenizer file: {relative_name}" + ) + continue + if ( + PurePosixPath(relative_name).name == "tokenizer_config.json" + and tokenizer_auto_map is not None + ): + # This file is intentionally rewritten to point at the + # artifact-local bridge. Its final digest is enforced by + # artifact-manifest.json and the mapping above. + continue + try: + algorithm, expected = encoded_digest.split(":", maxsplit=1) + actual = hash_file(path / relative_name, algorithm) + except (AttributeError, ArtifactError, OSError, ValueError): + failures.append(f"invalid tokenizer provenance for {relative_name}") + continue + if actual != expected: + failures.append(f"tokenizer digest mismatch for {relative_name}") + upstreams = provenance.get("upstreams") + if not isinstance(upstreams, list) or not upstreams: + failures.append("upstream legal provenance is missing") + else: + for upstream in upstreams: + if not isinstance(upstream, dict): + failures.append("invalid upstream provenance entry") + continue + source_id = upstream.get("id") + distributed = upstream.get("distribution_files") + if not isinstance(source_id, str) or not isinstance(distributed, dict): + failures.append("invalid upstream distribution record") + continue + for relative_name, encoded_digest in distributed.items(): + artifact_name = f"LICENSES/{source_id}/{relative_name}" + if artifact_name not in manifest: + failures.append(f"missing required legal artifact {artifact_name}") + continue + if not isinstance(encoded_digest, str): + failures.append(f"invalid legal digest for {artifact_name}") + continue + try: + algorithm, expected = encoded_digest.split(":", maxsplit=1) + actual = hash_file(path / artifact_name, algorithm) + except (ArtifactError, OSError, ValueError): + failures.append(f"invalid legal digest for {artifact_name}") + continue + if actual != expected: + failures.append(f"legal digest mismatch for {artifact_name}") + try: + runtime_attestation = _load_json_object_for_build( + path / _RUNTIME_ATTESTATION_NAME + ) + except ArtifactError as error: + runtime_attestation = None + failures.append(str(error)) + if isinstance(runtime_attestation, dict): + canonical = provenance.get("canonical_weights") + raw_shards = canonical.get("shards") if isinstance(canonical, dict) else None + weight_paths = { + canonical.get("index") if isinstance(canonical, dict) else None, + *(raw_shards if isinstance(raw_shards, dict) else ()), + } + excluded = { + None, + "artifact-manifest.json", + "provenance.json", + _RUNTIME_ATTESTATION_NAME, + *weight_paths, + } + expected_runtime_files = { + name: digest + for name, digest in manifest.items() + if ( + name not in excluded + and not _is_weight_file(name) + and _is_runtime_update_path(name) + ) + } + expected_attestation_fields = { + "schema_version": _RUNTIME_ATTESTATION_SCHEMA_VERSION, + "scope": "runtime-only", + "model_id": provenance.get("model_id"), + "weights": { + "repo_id": ( + provenance.get("fast_checkpoint", {}).get("repo_id") + if isinstance(provenance.get("fast_checkpoint"), dict) + else None + ), + "revision": ( + provenance.get("fast_checkpoint", {}).get("revision") + if isinstance(provenance.get("fast_checkpoint"), dict) + else None + ), + }, + "runtime_revision": provenance.get("runtime_revision"), + "source_tree_sha256": provenance.get("source_tree_sha256"), + "runtime_bundle_sha256": provenance.get("runtime_bundle_sha256"), + "release_tool_revision": provenance.get("release_tool_revision"), + "release_tool_sha256": provenance.get("release_tool_sha256"), + "weights_license_status": provenance.get("weights_license_status"), + "redistributable": provenance.get("redistributable"), + "files": expected_runtime_files, + } + if runtime_attestation != expected_attestation_fields: + failures.append( + "runtime-only attestation differs from artifact contents or provenance" + ) + if not weight_validation_attempted: + try: + validate_weight_artifact(path) + except ArtifactError as error: + failures.append(str(error)) + if failures: + raise ArtifactError("Artifact validation failed: " + "; ".join(failures)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model_id", help="Stable model ID from src/fastplms/models.toml") + parser.add_argument("checkpoint_dir", type=Path, help="Pinned local Hub snapshot") + parser.add_argument("--output-root", type=Path, default=Path("dist/hub")) + parser.add_argument("--source-root", type=Path, default=None) + parser.add_argument( + "--tokenizer-dir", + type=Path, + default=None, + help="Pinned official tokenizer snapshot (required for FastPLMs tokenizer checkpoints)", + ) + parser.add_argument("--replace", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + registry = get_model_registry() + try: + spec = registry[args.model_id] + except KeyError as error: + raise ArtifactError(f"Unknown model ID: {args.model_id}") from error + destination = build_local_artifact( + model_id=args.model_id, + checkpoint_dir=args.checkpoint_dir, + output_root=args.output_root, + source_root=args.source_root, + tokenizer_dir=args.tokenizer_dir, + replace=args.replace, + ) + validate_artifact(destination, spec=spec, registry=registry) + print(destination) + + +if __name__ == "__main__": + main() + + +__all__ = [ + "ArtifactError", + "build_artifact", + "build_local_artifact", + "canonicalize_checkpoint_weights", + "hash_file", + "main", + "render_model_card", + "validate_artifact", + "validate_repository_legal_inventory", + "validate_weight_artifact", + "verify_checkpoint", +] diff --git a/tools/artifacts/build_all.py b/tools/artifacts/build_all.py new file mode 100644 index 0000000..f8ff673 --- /dev/null +++ b/tools/artifacts/build_all.py @@ -0,0 +1,125 @@ +"""Materialize every manifest artifact from immutable Hub snapshots. + +The low-level builder remains network-free. This orchestration command resolves +only the repository IDs and revisions pinned by ``models.toml``, then hands the +verified local snapshots to that builder. It never uploads or deletes a Hub +repository. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Iterable +from pathlib import Path +from huggingface_hub import snapshot_download + +from benchmarks.suite import benchmark_artifact_model_ids +from fastplms.registry import get_model_registry +from tools.artifacts.build import ( + _TOKENIZER_FILE_NAMES, + _tokenizer_checkpoint, + build_local_artifact, + validate_artifact, +) + + +def build_all_artifacts( + *, + output_root: Path, + source_root: Path, + model_ids: Iterable[str] | None = None, + benchmark_suite: bool = False, + replace: bool = False, +) -> tuple[Path, ...]: + """Build selected artifacts after resolving their immutable snapshots.""" + + registry = get_model_registry() + if model_ids is not None and benchmark_suite: + raise ValueError("model_ids and benchmark_suite are mutually exclusive") + if model_ids is not None: + selected = tuple(model_ids) + elif benchmark_suite: + selected = benchmark_artifact_model_ids() + else: + selected = tuple(registry) + unknown = sorted(set(selected).difference(registry)) + if unknown: + raise ValueError(f"Unknown model IDs: {unknown}") + + destinations: list[Path] = [] + for model_id in selected: + spec = registry[model_id] + checkpoint = spec.artifact_checkpoint + snapshot = Path( + snapshot_download( + repo_id=checkpoint.repo_id, + revision=checkpoint.revision, + allow_patterns=[item.path for item in checkpoint.files], + ) + ) + tokenizer_snapshot: Path | None = None + if spec.family.tokenizer_mode == "tokenizer": + tokenizer_checkpoint = _tokenizer_checkpoint(registry, spec) + tokenizer_files = [ + item.path + for item in tokenizer_checkpoint.files + if Path(item.path).name in _TOKENIZER_FILE_NAMES + ] + if not tokenizer_files: + raise RuntimeError(f"{model_id}: official tokenizer files are not declared") + if checkpoint == tokenizer_checkpoint: + tokenizer_snapshot = snapshot + else: + tokenizer_snapshot = Path( + snapshot_download( + repo_id=tokenizer_checkpoint.repo_id, + revision=tokenizer_checkpoint.revision, + allow_patterns=tokenizer_files, + ) + ) + destination = build_local_artifact( + model_id=model_id, + checkpoint_dir=snapshot, + output_root=output_root, + source_root=source_root, + tokenizer_dir=tokenizer_snapshot, + replace=replace, + ) + # Keep the completed artifact bound to the same current registry entry + # used for construction. This is required for official-source + # artifacts, whose canonical-state commitment cannot be authenticated + # from their self-authored provenance alone. + validate_artifact(destination, spec=spec, registry=registry) + destinations.append(destination) + return tuple(destinations) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model_ids", nargs="*") + parser.add_argument("--output-root", type=Path, default=Path("dist/hub")) + parser.add_argument("--source-root", type=Path, default=Path.cwd()) + parser.add_argument( + "--benchmark-suite", + action="store_true", + help="Build only the fixed benchmark matrix artifacts and nested backbones.", + ) + parser.add_argument("--replace", action="store_true") + return parser.parse_args() + + +def main() -> None: + arguments = _parse_args() + paths = build_all_artifacts( + output_root=arguments.output_root, + source_root=arguments.source_root, + model_ids=arguments.model_ids or None, + benchmark_suite=arguments.benchmark_suite, + replace=arguments.replace, + ) + for path in paths: + print(path) + + +if __name__ == "__main__": + main() diff --git a/tools/artifacts/generate_docs.py b/tools/artifacts/generate_docs.py new file mode 100644 index 0000000..c2073ec --- /dev/null +++ b/tools/artifacts/generate_docs.py @@ -0,0 +1,4118 @@ +"""Generate model support data and model cards from the typed manifest.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import random +import statistics +import tempfile +import textwrap +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path + +from fastplms.registry import ModelFamily, ModelRegistry, ModelSpec, get_model_registry +from tools.artifacts.license_metadata import ( + render_checkpoint_terms, + render_hub_license_yaml, +) +from tools.remote.biohub_reference_environment import ( + BiohubReferenceEnvironmentError, + validate_biohub_reference_environment_evidence, +) + + +GENERATED_MARKER = "" +BINDER_IMAGE_URL = ( + "https://raw.githubusercontent.com/Synthyra/FastPLMs/main/" + "docs/assets/egfr_fastplms_binder_design.png" +) +ESMC_RELEASE_DOCUMENTATION = """\ +Detailed backend measurements, release guardrails, and the GH200 package +compatibility exception are maintained in the +[attention backend guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/attention_backends.md) +and +[release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md). +""" + +FAMILY_DOCUMENTATION = { + "esm2": ("../models.md#esm2", "../../examples/embedding_and_retrieval.py"), + "esm_plusplus": ( + "../models.md#esm-and-esmc", + "../../examples/attention_switching.py", + ), + "esm3": ("../models.md#esm3", "../../examples/generation.py"), + "e1": ("../models.md#e1", "../../examples/e1_rag.py"), + "dplm": ("../models.md#dplm", "../../examples/generation.py"), + "dplm2": ("../models.md#dplm2", "../../examples/generation.py"), + "ankh": ("../models.md#ankh", "../../examples/ankh_embeddings.py"), + "boltz2": ("../models.md#boltz2", "../../examples/structure_preparation.py"), + "esmfold": ("../models.md#esmfold", "../../examples/structure_preparation.py"), + "esmfold2": ("../esmfold2.md", "../../examples/structure_preparation.py"), +} + +AUTO_CLASS_STATUS = { + "AutoConfig": "FastPLMs extension", + "AutoModel": "pretrained", + "AutoModelForMaskedLM": "pretrained", + "AutoModelForProteinFolding": "pretrained", + "AutoModelForSeq2SeqLM": "pretrained", + "AutoModelForSequenceClassification": "base weights + untrained task head", + "AutoModelForTokenClassification": "base weights + untrained task head", +} + +EMBEDDING_FAMILIES = frozenset( + { + "ankh", + "dplm", + "dplm2", + "e1", + "esm2", + "esm3", + "esm_plusplus", + "esmfold2", + } +) + +SEQUENCE_TTT_AUTO_CLASSES = { + "ankh": "AutoModelForMaskedLM", + "dplm": "AutoModelForMaskedLM", + "dplm2": "AutoModelForMaskedLM", + "e1": "AutoModelForMaskedLM", + "esm2": "AutoModelForMaskedLM", + "esm3": "AutoModel", + "esm_plusplus": "AutoModelForMaskedLM", +} + + +@dataclass(frozen=True) +class EvidenceSelector: + """One executable validation selector and the exact scope it supports.""" + + tier: str + targets: tuple[str, ...] + scope: str + + +@dataclass(frozen=True) +class CapabilityEvidenceRow: + """A documented capability backed by explicit evidence selectors.""" + + capability: str + guide: str + example: str + evidence: tuple[str, ...] + + +class EsmcReportError(ValueError): + """Raised when frozen ESMC release evidence is incomplete or invalid.""" + + +@dataclass(frozen=True, slots=True) +class EsmcRuntimeIdentity: + """Artifact-derived runtime identity required by every ESMC report.""" + + runtime_revision: str + source_tree_sha256: str + runtime_bundle_sha256: str + + +@dataclass(frozen=True, slots=True) +class EsmcReportSet: + """One complete, cross-device-consistent 30-record ESMC evidence set.""" + + reports: tuple[dict[str, object], ...] + runtime_identity: EsmcRuntimeIdentity + candidate_environment: dict[str, object] + reference_environment: dict[str, object] + + def select(self, model_id: str) -> tuple[dict[str, object], ...]: + """Return the ten backend/panel reports for one checkpoint.""" + + return tuple(report for report in self.reports if report["model_id"] == model_id) + + def get(self, model_id: str, backend: str, panel: str) -> dict[str, object]: + """Return one uniquely keyed report from the complete evidence set.""" + + matches = tuple( + report + for report in self.reports + if report["model_id"] == model_id + and report["configured_backend"] == backend + and isinstance(report["panel"], Mapping) + and report["panel"]["kind"] == panel + ) + if len(matches) != 1: + raise EsmcReportError( + f"ESMC evidence key {(model_id, backend, panel)!r} resolved to " + f"{len(matches)} reports" + ) + return matches[0] + + +ESMC_DIAGNOSTIC_SCHEMA_VERSION = 3 +ESMC_MODEL_IDS = ("esmc_small", "esmc_large", "esmc_6b") +ESMC_PANEL_KINDS = ("generated_kernel_boundary", "real_biological_holdout") +ESMC_REFERENCE_SOURCE_NAMES = ("biohub-esm", "biohub-transformers") +ESMC_BACKENDS = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", +) +ESMC_MEASURED_BACKENDS = ("eager", "sdpa", "flex_attention") +ESMC_UNAVAILABLE_BACKENDS = ("flash_attention_2", "flash_attention_3") +ESMC_REPORT_COUNT = len(ESMC_MODEL_IDS) * len(ESMC_BACKENDS) * len(ESMC_PANEL_KINDS) +ESMC_REPORT_MAX_BYTES = 16 * 1024 * 1024 +ESMC_RELEASE_GATE_MODES = { + "sdpa": "exact", + "eager": "strict_numeric", + "flex_attention": "diagnostic_with_catastrophe_gate", +} +ESMC_CATASTROPHE_UPPER = { + "relative_l2": 0.25, + "relative_q999": 0.50, +} +ESMC_CATASTROPHE_LOWER = { + "residue_cosine_p01": 0.90, + "pooled_cosine_min": 0.95, +} +ESMC_TOP_LEVEL_FIELDS = { + "schema_version", + "model_id", + "candidate", + "reference", + "record_status", + "unavailability", + "configured_backend", + "effective_backend", + "dtype", + "panel", + "environment", + "kernel", + "panel_tensor_metrics", + "panel_logits_metrics", + "cases", + "published_band_violations", + "catastrophic_gate", + "release_gate", + "report_sha256", +} + + +CAPABILITY_EVIDENCE_SELECTORS: dict[str, EvidenceSelector] = { + "cpu:autoclass-runtime": EvidenceSelector( + tier="cpu_contract", + targets=( + "tests/cpu/test_autoclass_evidence_matrix.py::" + "test_autoclass_runtime_evidence_matrix_exactly_matches_all_37_entries", + "tests/cpu/test_autoclass_evidence_matrix.py::" + "test_autoclass_runtime_evidence_targets_are_collected_cpu_tests", + ), + scope="Every family-level AutoClass entry and its explicit tiny runtime contracts.", + ), + "artifact:checkpoint-autoclasses": EvidenceSelector( + tier="artifact", + targets=( + "tests/release/test_published_automodel.py::" + "test_local_artifact_offline_autoclass_parity", + ), + scope="Every advertised AutoClass for every built checkpoint, grouped by checkpoint.", + ), + "compliance:sequence-primary-head": EvidenceSelector( + tier="compliance", + targets=( + "tests/parity/test_native_results.py::test_native_exact_checkpoint_contract", + "tests/parity/test_native_results.py::test_native_every_checkpoint_bf16_inference", + ), + scope=( + "The official-parity head only: AutoModel for ANKH or a family without MaskedLM; " + "otherwise AutoModelForMaskedLM." + ), + ), + "compliance:ankh-seq2seq": EvidenceSelector( + tier="compliance", + targets=( + "tests/parity/test_native_results.py::" + "test_native_ankh_explicit_decoder_prompt_generation", + ), + scope="ANKH AutoModelForSeq2SeqLM explicit-prompt generation only.", + ), + "compliance:structure-automodel": EvidenceSelector( + tier="compliance", + targets=( + "tests/structure/test_esmfold_folding_compliance.py", + "tests/structure/test_esmfold2_folding_compliance.py", + ), + scope="ESMFold and ESMFold2 AutoModel folding paths only.", + ), + "benchmark:claim-eligible-primary-head": EvidenceSelector( + tier="benchmark", + targets=("benchmarks/suite.py::benchmark_cases[claim_eligible=True]",), + scope=( + "The benchmark-selected head for representative sequence checkpoints and " + "ESMFold2 projection cases; startup and embedding cases are excluded." + ), + ), + "cpu:attention-contracts": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_attention_contracts.py",), + scope=( + "Portable dispatch, masks, fallback, fake FA2/FA3, ESMC Flex/FA3, and " + "eager/SDPA gradient contracts." + ), + ), + "nightly:sequence-backends": EvidenceSelector( + tier="nightly", + targets=("tests/integration/test_backend_consistency.py",), + scope=( + "Current GH200 eager, SDPA, and Flex forward/backward paths. Flash kernels " + "are not downloaded, built, or executed in the current locked environment." + ), + ), + "historical:fa2-focused": EvidenceSelector( + tier="historical", + targets=("tools/remote/run.py::_kernel_capability_preflight",), + scope=( + "Policy records prior real FlashAttention 2 focused execution, but the immutable " + "execution report is not bundled in this repository and no current GH200 " + "numerical claim is inferred from it." + ), + ), + "compliance:flash-unavailable-gh200": EvidenceSelector( + tier="compliance", + targets=( + "tests/parity/test_native_results.py::" + "test_esmc_bf16_calibration_and_biological_holdout", + ), + scope=( + "Complete report-bound FA2/FA3 unavailability records and fail-closed " + "dispatch on the frozen release environment." + ), + ), + "compliance:deep-backends": EvidenceSelector( + tier="compliance", + targets=("tests/parity/test_native_results.py::test_native_representatives_all_backends",), + scope="Every advertised backend on the pinned deep sequence representative per family.", + ), + "benchmark:claim-eligible-backends": EvidenceSelector( + tier="benchmark", + targets=("benchmarks/suite.py::benchmark_cases[claim_eligible=True]",), + scope="Backends emitted by claim-eligible sequence and ESMFold2 benchmark cases.", + ), + "cpu:embedding-contracts": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_embedding_contracts.py",), + scope="Ordered inputs, biological masking, pooling, streaming, and persistence.", + ), + "cpu:e1-embeddings": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_e1_contracts.py",), + scope="E1 raw-sequence and MSA embedding persistence.", + ), + "feature:e1-rag": EvidenceSelector( + tier="feature", + targets=("tests/integration/test_e1_rag.py",), + scope="E1 retrieval, MSA preparation, cache, scoring, and embedding flows.", + ), + "cpu:ankh-contracts": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_ankh_contracts.py",), + scope="ANKH encoder and explicit-decoder embeddings, layers, masks, and T5 views.", + ), + "cpu:generation-contracts": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_generation_contracts.py",), + scope="Tiny deterministic DPLM, DPLM2, and ESM3 generation contracts.", + ), + "feature:generation": EvidenceSelector( + tier="feature", + targets=( + "tests/integration/test_dplm_generation.py", + "tests/integration/test_esm3.py", + ), + scope="DPLM, DPLM2, and ESM3 generation behavior in the feature suite.", + ), + "cpu:peft": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_peft_contracts.py",), + scope="Real initializer, collators, one optimizer step, and adapter/classifier reload.", + ), + "nightly:peft": EvidenceSelector( + tier="nightly", + targets=("tests/unit/test_fine_tuning_example.py",), + scope="Fine-tuning example contracts in the nightly feature job.", + ), + "cpu:ttt": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_ttt_contracts.py",), + scope="Seeded TTT initialization, update, reset, save, reload, and family isolation.", + ), + "feature:ttt": EvidenceSelector( + tier="feature", + targets=("tests/integration/test_ttt.py",), + scope="TTT integration behavior in the feature suite.", + ), + "cpu:structure-contracts": EvidenceSelector( + tier="cpu_contract", + targets=("tests/cpu/test_structure_contracts.py",), + scope="Tiny injected structure cores, public outputs, save/reload, and binder batching.", + ), + "structure:public-contracts": EvidenceSelector( + tier="structure", + targets=("tests/structure/test_structure_public_helpers.py",), + scope="Seeded Boltz helper, linker masking, real features, losses, and binder gradients.", + ), + "structure:full-suite": EvidenceSelector( + tier="structure", + targets=("tests/structure",), + scope="The declared GPU structure suite for folding and preparation behavior.", + ), + "feature:binder": EvidenceSelector( + tier="feature", + targets=("tests/integration/test_binder_design.py",), + scope="Seeded binder workflow, atom padding, critic ranking, and traceability.", + ), + "cpu:artifact-example": EvidenceSelector( + tier="cpu_contract", + targets=( + "tests/cpu/test_documentation_contracts.py::" + "test_artifact_loading_example_executes_local_only_autoconfig", + ), + scope="The offline local-artifact example with AutoConfig.", + ), + "cpu:task-head-example": EvidenceSelector( + tier="cpu_contract", + targets=( + "tests/cpu/test_documentation_contracts.py::" + "test_task_head_example_executes_all_advertised_heads_offline", + ), + scope=( + "Offline ESM2 masked-LM scoring, contacts, sequence classification, " + "and token classification through the documented example." + ), + ), +} + + +EMBEDDING_CAPABILITY_ROWS = ( + CapabilityEvidenceRow( + "Sequence list or streaming FASTA", + "[embedding API](../embedding_api.md)", + "[embedding and retrieval](../../examples/embedding_and_retrieval.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Ordered mapping or one-shot generator", + "[embedding API](../embedding_api.md)", + "[runnable API contracts](../../tests/cpu/test_embedding_contracts.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Biological-residue `max_length`, bounded token windows, and stable order", + "[embedding API](../embedding_api.md#bounded-streaming-and-length-policy)", + "[runnable API contracts](../../tests/cpu/test_embedding_contracts.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Mean and standard-deviation pooling", + "[embedding API](../embedding_api.md#pooling)", + "[embedding and retrieval](../../examples/embedding_and_retrieval.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Max/norm/median/variance/CLS/PARTI pooling", + "[embedding API](../embedding_api.md#pooling)", + "[runnable pooler contract](../../tests/cpu/test_embedding_contracts.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Full-residue and all-selected-layer output", + "[embedding API](../embedding_api.md#full-residue-embeddings)", + "[ANKH layers](../../examples/ankh_embeddings.py)", + ("cpu:embedding-contracts", "cpu:ankh-contracts"), + ), + CapabilityEvidenceRow( + "Transactional sharded safetensors and exact resume", + "[embedding API](../embedding_api.md#safetensors-storage)", + "[embedding and retrieval](../../examples/embedding_and_retrieval.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Read-only SQLite and ordered duplicate-preserving filters", + "[embedding API](../embedding_api.md#sqlite-streaming-retrieval-and-resume)", + "[embedding and retrieval](../../examples/embedding_and_retrieval.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "Legacy SQLite conversion without pickle deserialization", + "[embedding API](../embedding_api.md#sqlite-streaming-retrieval-and-resume)", + "[runnable converter contract](../../tests/cpu/test_embedding_contracts.py)", + ("cpu:embedding-contracts",), + ), + CapabilityEvidenceRow( + "E1 raw-sequence and MSA-aware ordered embeddings", + "[E1 guide](../models.md#e1)", + "[E1 RAG](../../examples/e1_rag.py)", + ("cpu:e1-embeddings", "feature:e1-rag"), + ), + CapabilityEvidenceRow( + "ANKH encoder/explicit-decoder hidden-state selection", + "[ANKH guide](../models.md#ankh)", + "[ANKH layers](../../examples/ankh_embeddings.py)", + ("cpu:ankh-contracts",), + ), +) + + +GENERATION_CAPABILITY_ROWS = ( + CapabilityEvidenceRow( + "ESM2 pretrained masked-LM scoring and contact prediction", + "[ESM2](../models.md#esm2)", + "[task heads](../../examples/task_heads.py)", + ("cpu:task-head-example", "cpu:autoclass-runtime"), + ), + CapabilityEvidenceRow( + "ESM2 sequence/token classification with explicitly untrained task heads", + "[ESM2](../models.md#esm2)", + "[task heads](../../examples/task_heads.py)", + ("cpu:task-head-example", "cpu:autoclass-runtime"), + ), + CapabilityEvidenceRow( + "DPLM amino-acid diffusion generation", + "[DPLM](../models.md#dplm)", + "[generation](../../examples/generation.py)", + ("cpu:generation-contracts", "feature:generation"), + ), + CapabilityEvidenceRow( + "DPLM2 modality-aware sequence/structure co-generation", + "[DPLM2](../models.md#dplm2)", + "[generation](../../examples/generation.py)", + ("cpu:generation-contracts", "feature:generation"), + ), + CapabilityEvidenceRow( + "ESM3 multimodal-conditioned generation", + "[ESM3](../models.md#esm3)", + "[generation](../../examples/generation.py)", + ("cpu:generation-contracts", "feature:generation"), + ), + CapabilityEvidenceRow( + "ANKH task-prompted sequence-to-sequence generation", + "[ANKH](../models.md#ankh)", + "[ANKH embeddings and generation](../../examples/ankh_embeddings.py)", + ("cpu:ankh-contracts", "compliance:ankh-seq2seq"), + ), + CapabilityEvidenceRow( + "Trainer/PEFT LoRA with immutable inputs and verified save/reload", + "[fine-tuning](../finetuning.md)", + "[fine-tuning](../../examples/fine_tuning.py)", + ("cpu:peft", "nightly:peft"), + ), + CapabilityEvidenceRow( + "Seeded TTT adapter initialize/update/reset/save/reload", + "[TTT](../ttt.md)", + "[TTT](../../examples/ttt.py)", + ("cpu:ttt", "feature:ttt"), + ), +) + + +STRUCTURE_CAPABILITY_ROWS = ( + CapabilityEvidenceRow( + "ESMFold single-chain folding and multimer-linker confidence masking", + "[models](../models.md#esmfold)", + "[structure preparation](../../examples/structure_preparation.py)", + ( + "cpu:structure-contracts", + "structure:public-contracts", + "structure:full-suite", + "compliance:structure-automodel", + ), + ), + CapabilityEvidenceRow( + "Seed-scoped Boltz2 protein helper and BF16 execution policy", + "[Boltz2](../models.md#boltz2)", + "[structure preparation](../../examples/structure_preparation.py)", + ( + "cpu:structure-contracts", + "structure:public-contracts", + "structure:full-suite", + ), + ), + CapabilityEvidenceRow( + "Atom-dense binder optimization and critic reporting", + "[binder design](../binder_design.md)", + "[binder design](../../examples/binder_design_fastplms.py)", + ( + "cpu:structure-contracts", + "structure:public-contracts", + "feature:binder", + ), + ), + CapabilityEvidenceRow( + "Offline local artifact AutoClass loading", + "[artifacts](../artifacts.md)", + "[artifact loading](../../examples/artifact_loading.py)", + ("cpu:artifact-example", "artifact:checkpoint-autoclasses"), + ), +) + + +def _esmfold2_structure_capability_rows( + registry: ModelRegistry, +) -> tuple[CapabilityEvidenceRow, ...]: + rows: list[CapabilityEvidenceRow] = [] + for spec in registry.by_family("esmfold2"): + if spec.msa_conditioning is None: + raise ValueError(f"{spec.id}: ESMFold2 MSA conditioning is undeclared") + if spec.msa_conditioning: + capability = ( + f"`{spec.id}` 48-block full ESMFold2: single-sequence or optional " + "MSA-conditioned protein inputs, typed complexes, ligands, nucleic acids, " + "modifications, bonds, and distograms; pocket requests fail closed" + ) + else: + capability = ( + f"`{spec.id}` 24-block Fast ESMFold2: inference-optimized " + "single-sequence conditioning with typed multichain and multimolecule " + "inputs; every protein must have `msa=None` and MSA inputs fail closed" + ) + rows.append( + CapabilityEvidenceRow( + capability, + "[ESMFold2](../esmfold2.md)", + "[structure preparation](../../examples/structure_preparation.py)", + ( + "cpu:structure-contracts", + "structure:full-suite", + "compliance:structure-automodel", + ), + ) + ) + return tuple(rows) + + +CURATED_EXAMPLE_CPU_CASES: dict[str, tuple[str, ...]] = { + "embedding_and_retrieval.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_embedding_and_retrieval_example_executes_with_ordered_sqlite", + ), + "attention_switching.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_attention_switching_main_executes_optimized_and_masked_fallback", + ), + "ankh_embeddings.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_ankh_embedding_example_executes_encoder_and_decoder_layers", + ), + "generation.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_dplm_branch_offline", + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_dplm2_branch_offline", + "tests/cpu/test_documentation_contracts.py::" + "test_generation_example_executes_seeded_esm3_trace", + ), + "e1_rag.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_e1_rag_example_executes_local_msa_and_shared_persistence", + ), + "ttt.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_ttt_example_executes_seeded_adapt_save_and_reset", + ), + "structure_preparation.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_structure_preparation_example_executes_each_public_branch", + ), + "artifact_loading.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_artifact_loading_example_executes_local_only_autoconfig", + ), + "task_heads.py": ( + "tests/cpu/test_documentation_contracts.py::" + "test_task_head_example_executes_all_advertised_heads_offline", + ), + "fine_tuning.py": ( + "tests/cpu/test_peft_contracts.py::" + "test_fine_tuning_main_wires_both_tasks_without_external_io", + "tests/cpu/test_peft_contracts.py::" + "test_shipped_collators_create_tokenizer_aware_sequence_and_pair_batches", + "tests/cpu/test_peft_contracts.py::" + "test_shipped_initializer_drives_one_peft_step_and_atomic_final_reload", + ), + "binder_design_fastplms.py": ( + "tests/cpu/test_structure_contracts.py::" + "test_public_binder_workflow_pads_heterogeneous_prepared_atoms_without_truncation", + "tests/cpu/test_structure_contracts.py::" + "test_binder_example_main_wires_explicit_offline_cli_arguments", + "tests/cpu/test_structure_contracts.py::" + "test_binder_structure_loss_is_finite_and_differentiable", + ), +} + + +def _code(values: Iterable[str]) -> str: + return ", ".join(f"`{value}`" for value in values) + + +def _table_row(*cells: str) -> str: + return "| " + " | ".join(cells) + " |" + + +def _append_rows(lines: list[str], rows: Iterable[tuple[str, ...]]) -> None: + lines.extend(_table_row(*row) for row in rows) + + +def _render_evidence_keys(keys: Iterable[str]) -> str: + values = tuple(keys) + missing = sorted(set(values).difference(CAPABILITY_EVIDENCE_SELECTORS)) + if missing: + raise ValueError("Unknown capability evidence selectors: " + ", ".join(missing)) + if not values: + raise ValueError("Every advertised capability requires at least one evidence selector.") + return _code(values) + + +def _append_capability_rows( + lines: list[str], + rows: Iterable[CapabilityEvidenceRow], +) -> None: + for row in rows: + lines.append( + _table_row( + row.capability, + row.guide, + row.example, + _render_evidence_keys(row.evidence), + ) + ) + + +def _primary_sequence_auto_class(family: ModelFamily) -> str: + advertised = set(family.auto_map) + if family.id == "ankh" or "AutoModelForMaskedLM" not in advertised: + selected = "AutoModel" + else: + selected = "AutoModelForMaskedLM" + if selected not in advertised: + raise ValueError(f"{family.id} does not advertise required primary class {selected}.") + return selected + + +def benchmark_autoclass_evidence_pairs( + registry: ModelRegistry, +) -> frozenset[tuple[str, str]]: + """Return only family/AutoClass pairs emitted by claim-eligible benchmarks.""" + + pairs: set[tuple[str, str]] = set() + for spec in registry.values(): + family = spec.family + if "benchmark" not in family.test_tiers: + continue + if not (spec.is_deep_reference or family.id == "esmfold2"): + continue + if family.tokenizer_mode == "structure" and family.id != "esmfold2": + continue + pairs.add((family.id, _primary_sequence_auto_class(family))) + return frozenset(pairs) + + +def benchmark_backend_evidence(registry: ModelRegistry) -> frozenset[str]: + """Return only backends emitted by claim-eligible benchmark cases.""" + + from benchmarks.suite import benchmark_cases + + backends = frozenset( + str(case.backend) + for case in benchmark_cases( + family=None, + quick=False, + local_files_only=True, + ) + if case.claim_eligible + ) + advertised = {backend for family in registry.families.values() for backend in family.attention} + unexpected = sorted(backends.difference(advertised)) + if unexpected: + raise ValueError( + "Claim-eligible benchmark cases advertise unknown backends: " + ", ".join(unexpected) + ) + return backends + + +def autoclass_evidence_keys( + registry: ModelRegistry, + family_id: str, + auto_class: str, +) -> tuple[str, ...]: + """Map one advertised AutoClass to its actually executable evidence.""" + + family = registry.families[family_id] + if auto_class not in family.auto_map: + raise ValueError(f"{family_id} does not advertise {auto_class}.") + + evidence = ["cpu:autoclass-runtime", "artifact:checkpoint-autoclasses"] + if family.id == "esm2" and auto_class in { + "AutoModelForMaskedLM", + "AutoModelForSequenceClassification", + "AutoModelForTokenClassification", + }: + evidence.append("cpu:task-head-example") + if family.tokenizer_mode != "structure": + if auto_class == _primary_sequence_auto_class(family): + evidence.append("compliance:sequence-primary-head") + if family.id == "ankh" and auto_class == "AutoModelForSeq2SeqLM": + evidence.append("compliance:ankh-seq2seq") + elif family.id in {"esmfold", "esmfold2"} and auto_class == "AutoModel": + evidence.append("compliance:structure-automodel") + + if (family_id, auto_class) in benchmark_autoclass_evidence_pairs(registry): + evidence.append("benchmark:claim-eligible-primary-head") + return tuple(evidence) + + +def _autoclass_workflow_example(family: ModelFamily, auto_class: str) -> str: + if family.id == "esm2" and auto_class in { + "AutoModelForMaskedLM", + "AutoModelForSequenceClassification", + "AutoModelForTokenClassification", + }: + return "../../examples/task_heads.py" + return FAMILY_DOCUMENTATION[family.id][1] + + +def attention_backend_evidence_keys( + registry: ModelRegistry, + backend: str, +) -> tuple[str, ...]: + """Map an advertised backend to scoped CPU, GPU, parity, and benchmark evidence.""" + + advertising_families = tuple( + family for family in registry.families.values() if backend in family.attention + ) + if not advertising_families: + raise ValueError(f"No family advertises attention backend {backend!r}.") + + evidence = ["cpu:attention-contracts"] + sequence_families = tuple( + family for family in advertising_families if family.tokenizer_mode != "structure" + ) + if sequence_families and backend in ESMC_MEASURED_BACKENDS: + evidence.append("nightly:sequence-backends") + if backend == "flash_attention_2": + evidence.append("historical:fa2-focused") + evidence.append("compliance:flash-unavailable-gh200") + elif backend == "flash_attention_3": + evidence.append("compliance:flash-unavailable-gh200") + elif any( + spec.is_deep_reference + and spec.family.tokenizer_mode != "structure" + and backend in spec.family.attention + for spec in registry.values() + ): + evidence.append("compliance:deep-backends") + if backend in ESMC_MEASURED_BACKENDS and backend in benchmark_backend_evidence(registry): + evidence.append("benchmark:claim-eligible-backends") + return tuple(evidence) + + +def _render_evidence_selector_catalog() -> list[str]: + lines = [ + "## Executable evidence selectors", + "", + "Only the selectors below are claimed. Their scopes are intentionally narrower", + "than a whole family or validation tier. A tier appearing on another row does not", + "automatically apply to the capability in this row.", + "", + _table_row("Selector", "Tier/job", "Executable target", "Scope"), + _table_row("---", "---", "---", "---"), + ] + for key, selector in CAPABILITY_EVIDENCE_SELECTORS.items(): + targets = "
".join(f"`{target}`" for target in selector.targets) + lines.append(_table_row(f"`{key}`", f"`{selector.tier}`", targets, selector.scope)) + lines.append("") + return lines + + +def _render_curated_example_cpu_evidence() -> list[str]: + lines = [ + "## Curated offline example execution", + "", + "Every curated example is routed to the exact collected CPU test nodes below.", + "These tests run under the required offline `cpu_contract` gate.", + "", + _table_row("Example", "Tier", "Exact executable CPU node"), + _table_row("---", "---", "---"), + ] + for example_name, nodeids in CURATED_EXAMPLE_CPU_CASES.items(): + for nodeid in nodeids: + lines.append( + _table_row( + f"[`{example_name}`](../../examples/{example_name})", + "`cpu_contract`", + f"`{nodeid}`", + ) + ) + lines.append("") + return lines + + +def _precision_contract(family: ModelFamily) -> str: + experimental = set(family.experimental_precisions) + return ", ".join( + f"`{value}` (experimental)" if value in experimental else f"`{value}`" + for value in family.precisions + ) + + +def _hub_license_label(family: ModelFamily) -> str: + label = f"`{family.hub_license}`" + if family.hub_license == "other": + label += f" ({render_checkpoint_terms(family)})" + return label + + +def _tokenizer_class_label(family: ModelFamily) -> str: + if family.tokenizer_class is None: + return "`n/a`" + return f"`{family.tokenizer_class}`" + + +def _auto_class_status(family: ModelFamily, auto_class: str) -> str: + """Describe whether an advertised entry point has trained checkpoint state.""" + + if family.id == "ankh" and auto_class == "AutoModelForMaskedLM": + return "FastPLMs extension" + try: + return AUTO_CLASS_STATUS[auto_class] + except KeyError as error: + raise ValueError(f"No model-card weight status is defined for {auto_class!r}.") from error + + +def _platform_requirements(family: ModelFamily) -> str: + requirements = ["Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required."] + if family.tokenizer_mode == "structure": + requirements.append( + "The artifact requirements include the direct structure dependencies. " + "The published execution contract requires a CUDA device. The current " + "validated release target is the exact NVIDIA GH200 on Linux aarch64; " + "Linux x86-64, CPU-only, Windows, and macOS structure runs are not " + "current release evidence." + ) + elif any(name.startswith("flash_attention_") for name in family.attention): + requirements.append( + "The artifact requirements include the direct FlashAttention loader " + "dependency. FlashAttention also requires compatible CUDA hardware and " + "BF16 execution." + ) + else: + requirements.append( + "The declared CPU gate covers tiny offline contracts; published " + "checkpoint throughput and parity require the documented device tier." + ) + return " ".join(requirements) + + +def _installation_section(spec: ModelSpec) -> str: + return f"""\ +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \\ + "https://huggingface.co/{spec.fast.repo_id}/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository and loaded +by Transformers through `trust_remote_code=True`. + +{_platform_requirements(spec.family)} The Hub quick start below requires network +access on first download. For an air-gapped run, first build the manifest-pinned +local artifact and use the offline form shown in the example. + +""" + + +def _esmc_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise EsmcReportError(f"ESMC JSON contains duplicate key {key!r}") + result[key] = value + return result + + +def _esmc_reject_json_constant(value: str) -> object: + raise EsmcReportError(f"ESMC JSON contains non-finite constant {value!r}") + + +def _esmc_decode_json(encoded: str, *, context: str) -> dict[str, object]: + try: + payload = json.loads( + encoded, + object_pairs_hook=_esmc_json_object, + parse_constant=_esmc_reject_json_constant, + ) + except (json.JSONDecodeError, UnicodeError) as error: + raise EsmcReportError(f"{context} is not strict UTF-8 JSON: {error}") from error + if not isinstance(payload, dict): + raise EsmcReportError(f"{context} must contain one JSON object") + return payload + + +def _esmc_read_json(path: Path) -> dict[str, object]: + try: + size = path.stat().st_size + except OSError as error: + raise EsmcReportError(f"Unable to stat ESMC evidence file {path}: {error}") from error + if size <= 0 or size > ESMC_REPORT_MAX_BYTES: + raise EsmcReportError( + f"ESMC evidence file {path.name!r} has invalid size {size}; " + f"maximum is {ESMC_REPORT_MAX_BYTES} bytes" + ) + try: + encoded = path.read_text(encoding="utf-8") + except (OSError, UnicodeError) as error: + raise EsmcReportError(f"Unable to read ESMC evidence file {path}: {error}") from error + return _esmc_decode_json(encoded, context=f"ESMC evidence file {path.name!r}") + + +def _esmc_require_mapping( + value: object, + fields: set[str], + *, + context: str, +) -> dict[str, object]: + if not isinstance(value, dict) or set(value) != fields: + raise EsmcReportError(f"{context} fields differ from schema v3") + return value + + +def _esmc_require_object(value: object, *, context: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise EsmcReportError(f"{context} must be a JSON object") + return value + + +def _esmc_require_list(value: object, *, context: str) -> list[object]: + if not isinstance(value, list): + raise EsmcReportError(f"{context} must be a JSON array") + return value + + +def _esmc_require_text(value: object, *, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise EsmcReportError(f"{context} must be a nonempty string") + return value + + +def _esmc_require_sha256(value: object, *, context: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or value != value.lower() + or any(character not in "0123456789abcdef" for character in value) + ): + raise EsmcReportError(f"{context} must be a canonical lowercase SHA-256 digest") + return value + + +def _esmc_require_finite(value: object, *, context: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise EsmcReportError(f"{context} must be a finite number") + numeric = float(value) + if not math.isfinite(numeric): + raise EsmcReportError(f"{context} must be a finite number") + return numeric + + +def _esmc_require_gpu_capability( + value: object, + *, + context: str, +) -> tuple[int, int]: + if not isinstance(value, list) or len(value) != 2: + raise EsmcReportError(f"{context} must contain exactly two integers") + major, minor = value + if ( + isinstance(major, bool) + or not isinstance(major, int) + or major < 0 + or isinstance(minor, bool) + or not isinstance(minor, int) + or minor < 0 + ): + raise EsmcReportError(f"{context} must contain exactly two non-negative integers") + return major, minor + + +def _esmc_report_sha256(payload: Mapping[str, object]) -> str: + digest_payload = dict(payload) + digest_payload.pop("report_sha256", None) + encoded = json.dumps( + digest_payload, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _esmc_public_case(case: Mapping[str, object]) -> dict[str, object]: + return { + "case_id": case["case_id"], + "sequence_length": case["sequence_length"], + "sequence_sha256": case["sequence_sha256"], + "source": case.get("source"), + "source_sha256": case.get("source_sha256"), + } + + +def _esmc_panel_identity(kind: str, cases: list[dict[str, object]]) -> dict[str, object]: + definition = { + "schema_version": 1, + "kind": kind, + "seed": 42, + "cases": cases, + } + definition_sha256 = hashlib.sha256( + json.dumps(definition, separators=(",", ":"), sort_keys=True).encode("utf-8") + ).hexdigest() + return { + "schema_version": 1, + "kind": kind, + "seed": 42, + "definition_sha256": definition_sha256, + "cases": [_esmc_public_case(case) for case in cases], + } + + +def _expected_esmc_panels(source_root: Path) -> dict[str, dict[str, object]]: + alphabet = "ACDEFGHIKLMNPQRSTVWY" + generator = random.Random(42) + generated_cases: list[dict[str, object]] = [] + for length in (13, 15, 16, 17, 29, 31, 32, 33, 61, 127, 128, 129): + sequence = "M" + "".join(generator.choices(alphabet, k=length - 1)) + generated_cases.append( + { + "case_id": f"generated-boundary-{length}", + "sequence": sequence, + "sequence_length": length, + "sequence_sha256": hashlib.sha256(sequence.encode("ascii")).hexdigest(), + } + ) + + fixture_path = source_root / "tests" / "parity" / "fixtures" / "esmc_biological_holdout.json" + try: + fixture_text = fixture_path.read_text(encoding="utf-8") + except (OSError, UnicodeError) as error: + raise EsmcReportError(f"Unable to read immutable ESMC panel fixture: {error}") from error + fixture = _esmc_decode_json(fixture_text, context="ESMC biological holdout fixture") + if set(fixture) != {"schema_version", "cases"} or fixture["schema_version"] != 1: + raise EsmcReportError("ESMC biological holdout fixture fields differ from schema v1") + raw_cases = fixture["cases"] + if not isinstance(raw_cases, list) or not raw_cases: + raise EsmcReportError("ESMC biological holdout fixture has no ordered cases") + biological_cases: list[dict[str, object]] = [] + for index, raw_case in enumerate(raw_cases): + case = _esmc_require_mapping( + raw_case, + {"case_id", "sequence", "sequence_sha256", "source", "source_sha256"}, + context=f"ESMC biological holdout case {index}", + ) + sequence = _esmc_require_text( + case["sequence"], context=f"ESMC biological holdout case {index} sequence" + ) + if not sequence.isupper() or not set(sequence).issubset(set(alphabet)): + raise EsmcReportError( + f"ESMC biological holdout case {index} is not canonical uppercase protein" + ) + sequence_sha256 = _esmc_require_sha256( + case["sequence_sha256"], context=f"ESMC biological holdout case {index} sequence" + ) + if sequence_sha256 != hashlib.sha256(sequence.encode("ascii")).hexdigest(): + raise EsmcReportError(f"ESMC biological holdout case {index} sequence digest drifted") + _esmc_require_sha256( + case["source_sha256"], context=f"ESMC biological holdout case {index} source" + ) + biological_cases.append({**case, "sequence_length": len(sequence)}) + + return { + "generated_kernel_boundary": _esmc_panel_identity( + "generated_kernel_boundary", generated_cases + ), + "real_biological_holdout": _esmc_panel_identity( + "real_biological_holdout", biological_cases + ), + } + + +def _expected_biohub_source_contracts( + source_root: Path, +) -> dict[str, dict[str, object]]: + expected_fields = { + "import_name", + "import_root", + "package_version", + "schema_version", + "source_revision", + "tree_sha256", + } + file_names = { + "biohub-esm": "biohub-esm-source.json", + "biohub-transformers": "biohub-transformers-source.json", + } + contracts: dict[str, dict[str, object]] = {} + for source_name in ESMC_REFERENCE_SOURCE_NAMES: + path = source_root / "docker" / "constraints" / file_names[source_name] + try: + encoded = path.read_text(encoding="utf-8") + except (OSError, UnicodeError) as error: + raise EsmcReportError( + f"Unable to read pinned {source_name} source contract: {error}" + ) from error + contract = _esmc_decode_json(encoded, context=f"{source_name} source contract") + if set(contract) != expected_fields or contract["schema_version"] != 1: + raise EsmcReportError(f"{source_name} source contract differs from schema v1") + _esmc_require_text(contract["import_name"], context=f"{source_name} import name") + _esmc_require_text(contract["import_root"], context=f"{source_name} import root") + _esmc_require_text(contract["package_version"], context=f"{source_name} package version") + revision = contract["source_revision"] + if ( + not isinstance(revision, str) + or len(revision) != 40 + or revision != revision.lower() + or any(character not in "0123456789abcdef" for character in revision) + ): + raise EsmcReportError( + f"{source_name} source revision is not canonical lowercase 40-hex" + ) + _esmc_require_sha256(contract["tree_sha256"], context=f"{source_name} tree") + contracts[source_name] = contract + return contracts + + +def _validate_esmc_reference_sources( + value: object, + expected_contracts: Mapping[str, Mapping[str, object]], +) -> dict[str, object]: + sources = _esmc_require_mapping( + value, + set(ESMC_REFERENCE_SOURCE_NAMES), + context="ESMC reference sources", + ) + source_fields = { + "schema_version", + "source_revision", + "tree_sha256", + "attestation_sha256", + "file_count", + "import_name", + "import_root", + "import_file", + "package_version", + } + for source_name in ESMC_REFERENCE_SOURCE_NAMES: + source = _esmc_require_mapping( + sources[source_name], + source_fields, + context=f"ESMC reference source {source_name}", + ) + expected_contract = expected_contracts[source_name] + if source["schema_version"] != 1: + raise EsmcReportError(f"ESMC reference source {source_name} schema is unsupported") + for name in ( + "source_revision", + "tree_sha256", + "import_name", + "import_root", + "package_version", + ): + if source[name] != expected_contract[name]: + raise EsmcReportError( + f"ESMC reference source {source_name} {name} differs from the pin" + ) + _esmc_require_sha256( + source["attestation_sha256"], + context=f"ESMC reference source {source_name} attestation", + ) + file_count = source["file_count"] + if isinstance(file_count, bool) or not isinstance(file_count, int) or file_count <= 0: + raise EsmcReportError(f"ESMC reference source {source_name} file count is invalid") + expected_import_file = f"{expected_contract['import_root']}/__init__.py" + if source["import_file"] != expected_import_file: + raise EsmcReportError( + f"ESMC reference source {source_name} import file differs from its root" + ) + return sources + + +def _esmc_runtime_identity_from_source( + source_root: Path, + registry: ModelRegistry, +) -> EsmcRuntimeIdentity: + try: + from tools.artifacts.build import ( + _render_runtime_bundle, + _validated_runtime_snapshot, + _write_runtime_snapshot, + ) + + identities: set[tuple[str, str, str]] = set() + with tempfile.TemporaryDirectory(prefix="fastplms-esmc-runtime-") as directory: + temporary_root = Path(directory) + for spec in (registry[model_id] for model_id in ESMC_MODEL_IDS): + runtime_revision, payloads, source_tree_sha256 = _validated_runtime_snapshot( + source_root, + registry, + spec, + ) + package_root = temporary_root / spec.id / "fastplms" + _write_runtime_snapshot(package_root, payloads) + runtime_bundle_sha256, _ = _render_runtime_bundle(package_root) + identities.add((runtime_revision, source_tree_sha256, runtime_bundle_sha256)) + except Exception as error: + raise EsmcReportError( + "Unable to derive the clean tracked ESMC runtime identity required for release " + f"evidence: {error}" + ) from error + if len(identities) != 1: + raise EsmcReportError( + "ESMC checkpoints do not resolve to one shared runtime/source/bundle identity" + ) + runtime_revision, source_tree_sha256, runtime_bundle_sha256 = identities.pop() + return EsmcRuntimeIdentity( + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ) + + +def _validate_esmc_runtime_identity(identity: EsmcRuntimeIdentity) -> None: + source_digest = _esmc_require_sha256( + identity.source_tree_sha256, context="ESMC expected source tree" + ) + _esmc_require_sha256(identity.runtime_bundle_sha256, context="ESMC expected runtime bundle") + revision = identity.runtime_revision + is_git_revision = ( + isinstance(revision, str) + and len(revision) == 40 + and revision == revision.lower() + and all(character in "0123456789abcdef" for character in revision) + ) + if not is_git_revision and revision != f"source-tree-sha256:{source_digest}": + raise EsmcReportError( + "ESMC runtime revision must be a clean 40-hex Git revision or the exact " + "source-tree-sha256 fallback" + ) + + +def _esmc_runtime_platform_identity( + reference_environment: Mapping[str, object], +) -> tuple[str, str]: + runtime = _esmc_require_object( + reference_environment.get("runtime"), + context="ESMC locked reference runtime", + ) + operating_system = _esmc_require_text( + runtime.get("operating_system"), context="ESMC locked operating system" + ) + architecture = _esmc_require_text( + runtime.get("architecture"), context="ESMC locked architecture" + ) + gpu = _esmc_require_object(runtime.get("gpu"), context="ESMC locked GPU identity") + gpu_name = _esmc_require_text(gpu.get("name"), context="ESMC locked GPU name") + capability = gpu.get("capability") + if ( + not isinstance(capability, list) + or len(capability) != 2 + or any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in capability + ) + ): + raise EsmcReportError("ESMC locked GPU capability is malformed") + return ( + f"{operating_system.lower()}/{architecture.lower()}", + f"{gpu_name}/SM{capability[0]}{capability[1]}", + ) + + +def _esmc_unavailability_identity( + backend: str, + reference_environment: Mapping[str, object], +) -> dict[str, str]: + platform_identity, accelerator_identity = _esmc_runtime_platform_identity(reference_environment) + if backend == "flash_attention_2": + historical_evidence = "separate_historical_focused_evidence_only" + reason = ( + f"The locked {platform_identity} {accelerator_identity} release environment " + "has no validated FlashAttention 2 " + "kernel. Prior focused execution evidence is historical and is not part of " + "the current ESMC release distribution." + ) + elif backend == "flash_attention_3": + historical_evidence = "none" + reason = ( + "The manifest-pinned FlashAttention 3 kernel has no validated artifact for " + f"the locked {platform_identity} {accelerator_identity} release environment." + ) + else: + raise EsmcReportError(f"ESMC backend {backend!r} is not a structured unavailable backend") + return { + "code": "locked_platform_kernel_unavailable", + "platform": platform_identity, + "accelerator": accelerator_identity, + "dispatch_contract": "fail_closed_without_dispatch", + "historical_evidence": historical_evidence, + "reason": reason, + } + + +def _validate_esmc_candidate_environment(value: object) -> dict[str, object]: + environment = _esmc_require_mapping( + value, + { + "python", + "torch", + "transformers", + "cuda_runtime", + "cuda_driver", + "gpu", + "packages", + }, + context="ESMC candidate environment", + ) + for name in ("python", "torch", "transformers", "cuda_runtime", "cuda_driver"): + _esmc_require_text(environment[name], context=f"ESMC candidate environment {name}") + packages = _esmc_require_mapping( + environment["packages"], + { + "fastplms", + "huggingface-hub", + "kernels", + "tokenizers", + "transformer-engine", + "transformer-engine-torch", + }, + context="ESMC candidate package inventory", + ) + for name, version in packages.items(): + if version is not None: + _esmc_require_text(version, context=f"ESMC candidate package {name}") + for required in ("fastplms", "huggingface-hub", "kernels", "tokenizers"): + if packages[required] is None: + raise EsmcReportError(f"ESMC candidate package {required!r} is unavailable") + gpu = _esmc_require_mapping( + environment["gpu"], + {"name", "capability", "total_memory_bytes"}, + context="ESMC candidate GPU identity", + ) + _esmc_require_text(gpu["name"], context="ESMC candidate GPU name") + capability = gpu["capability"] + if ( + not isinstance(capability, list) + or len(capability) != 2 + or any( + isinstance(item, bool) or not isinstance(item, int) or item < 0 for item in capability + ) + ): + raise EsmcReportError("ESMC candidate GPU capability is invalid") + memory = gpu["total_memory_bytes"] + if isinstance(memory, bool) or not isinstance(memory, int) or memory <= 0: + raise EsmcReportError("ESMC candidate GPU memory identity is invalid") + return environment + + +def _validate_esmc_reference_environment(value: object) -> dict[str, object]: + if not isinstance(value, dict): + raise EsmcReportError("ESMC reference environment is missing") + required = { + "cuda_device", + "cuda_device_capability", + "cuda_total_memory", + "cuda_runtime", + "packages", + "python", + "torch", + } + if not required.issubset(value): + raise EsmcReportError("ESMC reference environment fields are incomplete") + for name in ("cuda_device", "cuda_runtime", "python", "torch"): + _esmc_require_text(value[name], context=f"ESMC reference environment {name}") + _esmc_require_text(value["cuda_device"], context="ESMC reference environment CUDA device") + capability = value["cuda_device_capability"] + if ( + not isinstance(capability, list) + or len(capability) != 2 + or any( + isinstance(item, bool) or not isinstance(item, int) or item < 0 for item in capability + ) + ): + raise EsmcReportError("ESMC reference CUDA capability is invalid") + memory = value["cuda_total_memory"] + if isinstance(memory, bool) or not isinstance(memory, int) or memory <= 0: + raise EsmcReportError("ESMC reference GPU memory identity is invalid") + packages_text = _esmc_require_text( + value["packages"], context="ESMC reference package inventory" + ) + packages = _esmc_decode_json(packages_text, context="ESMC reference package inventory") + if not isinstance(packages, dict): + raise EsmcReportError("ESMC reference package inventory is not an object") + return value + + +def _validate_locked_esmc_reference_environment( + value: object, + *, + source_root: Path, +) -> dict[str, object]: + try: + return validate_biohub_reference_environment_evidence( + value, + repository_root=source_root, + contract_path=source_root / "docker/constraints/biohub-reference-lock.json", + ) + except BiohubReferenceEnvironmentError as error: + raise EsmcReportError(f"ESMC locked reference environment is invalid: {error}") from error + + +def _validate_esmc_environment_binding( + candidate_environment: Mapping[str, object], + dynamic_reference_environment: Mapping[str, object], + locked_reference_environment: Mapping[str, object], +) -> None: + candidate_gpu = _esmc_require_object( + candidate_environment.get("gpu"), context="ESMC candidate GPU binding" + ) + locked_runtime = _esmc_require_object( + locked_reference_environment.get("runtime"), + context="ESMC locked reference runtime binding", + ) + locked_gpu = _esmc_require_object( + locked_runtime.get("gpu"), context="ESMC locked reference GPU binding" + ) + candidate_identity = { + "python": candidate_environment.get("python"), + "torch": candidate_environment.get("torch"), + "cuda_runtime": candidate_environment.get("cuda_runtime"), + "cuda_driver": candidate_environment.get("cuda_driver"), + "gpu": dict(candidate_gpu), + } + dynamic_reference_identity = { + "python": dynamic_reference_environment.get("python"), + "torch": dynamic_reference_environment.get("torch"), + "cuda_runtime": dynamic_reference_environment.get("cuda_runtime"), + "cuda_driver": candidate_environment.get("cuda_driver"), + "gpu": { + "name": dynamic_reference_environment.get("cuda_device"), + "capability": dynamic_reference_environment.get("cuda_device_capability"), + "total_memory_bytes": dynamic_reference_environment.get("cuda_total_memory"), + }, + } + locked_identity = { + "python": locked_runtime.get("python_version"), + "torch": locked_runtime.get("torch"), + "cuda_runtime": locked_runtime.get("cuda_runtime"), + "cuda_driver": locked_runtime.get("cuda_driver"), + "gpu": dict(locked_gpu), + } + if candidate_identity != dynamic_reference_identity: + raise EsmcReportError("ESMC candidate and native reference environments differ") + if candidate_identity != locked_identity: + raise EsmcReportError( + "ESMC candidate environment differs from the locked reference runtime" + ) + + +def _validate_esmc_kernel( + value: object, + backend: str, + environment: Mapping[str, object], + registry: ModelRegistry, +) -> None: + kernel_spec = registry.attention_kernels.get(backend) + if kernel_spec is None: + expected = { + "implementation": backend, + "provider": "torch", + "torch_version": environment["torch"], + } + else: + packages = _esmc_require_object( + environment["packages"], context="ESMC candidate package inventory" + ) + expected = { + "implementation": backend, + "provider": "huggingface_kernels", + "repository": kernel_spec.repository, + "revision": kernel_spec.revision, + "version": kernel_spec.version, + "expected_variant": kernel_spec.expected_variant, + "supported_dtypes": list(kernel_spec.dtypes), + "kernels_package_version": packages["kernels"], + } + if value != expected: + raise EsmcReportError( + f"ESMC {backend} kernel identity differs from the manifest/runtime contract" + ) + + +def _validate_esmc_logits_metrics(value: object, *, context: str) -> None: + metrics = _esmc_require_mapping( + value, + {"confident_top1_agreement", "mean_jsd"}, + context=f"{context} logits metrics", + ) + agreement = _esmc_require_finite( + metrics["confident_top1_agreement"], context=f"{context} top-1 agreement" + ) + mean_jsd = _esmc_require_finite(metrics["mean_jsd"], context=f"{context} mean JSD") + if not 0.80 <= agreement <= 1.000001: + raise EsmcReportError(f"{context} top-1 agreement fails the catastrophe gate") + if not -1e-7 <= mean_jsd <= 0.05: + raise EsmcReportError(f"{context} mean JSD fails the catastrophe gate") + + +def _validate_esmc_tensor_metrics( + value: object, + *, + context: str, + expected_metric_context: str, +) -> tuple[tuple[str, int | None], ...]: + if not isinstance(value, list) or not value: + raise EsmcReportError(f"{context} tensor metrics are missing") + layout: list[tuple[str, int | None]] = [] + hidden_layers: list[int] = [] + output_counts = {"last_hidden_state": 0, "logits": 0} + fields = { + "context", + "output", + "layer_index", + "relative_l2", + "relative_q999", + "residue_cosine_p01", + "pooled_cosine_min", + } + for index, raw_metric in enumerate(value): + metric = _esmc_require_mapping( + raw_metric, fields, context=f"{context} tensor metric {index}" + ) + if metric["context"] != expected_metric_context: + raise EsmcReportError(f"{context} tensor metric context is stale or misaligned") + output = metric["output"] + layer_index = metric["layer_index"] + if output == "hidden_state": + if isinstance(layer_index, bool) or not isinstance(layer_index, int) or layer_index < 0: + raise EsmcReportError(f"{context} hidden-state layer index is invalid") + hidden_layers.append(layer_index) + elif output in output_counts: + if layer_index is not None: + raise EsmcReportError(f"{context} {output} layer index must be null") + output_counts[output] += 1 + else: + raise EsmcReportError(f"{context} tensor output {output!r} is unsupported") + for metric_name, upper in ESMC_CATASTROPHE_UPPER.items(): + numeric = _esmc_require_finite(metric[metric_name], context=f"{context} {metric_name}") + if not 0 <= numeric <= upper: + raise EsmcReportError(f"{context} {metric_name} fails the catastrophe gate") + for metric_name, lower in ESMC_CATASTROPHE_LOWER.items(): + numeric = _esmc_require_finite(metric[metric_name], context=f"{context} {metric_name}") + if not lower <= numeric <= 1.000001: + raise EsmcReportError(f"{context} {metric_name} fails the catastrophe gate") + if not isinstance(output, str) or not (layer_index is None or isinstance(layer_index, int)): + raise EsmcReportError(f"{context} tensor metric layout is invalid") + layout.append((output, layer_index)) + if hidden_layers != list(range(len(hidden_layers))): + raise EsmcReportError(f"{context} hidden-state layers are incomplete or unordered") + if output_counts != {"last_hidden_state": 1, "logits": 1}: + raise EsmcReportError( + f"{context} must contain exactly one last-hidden-state and one logits metric" + ) + return tuple(layout) + + +def _validate_esmc_report( + payload: dict[str, object], + *, + spec: ModelSpec, + backend: str, + panel: str, + expected_panel: Mapping[str, object], + expected_reference_sources: Mapping[str, Mapping[str, object]], + runtime_identity: EsmcRuntimeIdentity, + registry: ModelRegistry, + source_root: Path, +) -> None: + if set(payload) != ESMC_TOP_LEVEL_FIELDS or ( + payload.get("schema_version") != ESMC_DIAGNOSTIC_SCHEMA_VERSION + ): + raise EsmcReportError("ESMC diagnostic fields differ from schema v3") + report_sha256 = _esmc_require_sha256( + payload["report_sha256"], context="ESMC report self-digest" + ) + if report_sha256 != _esmc_report_sha256(payload): + raise EsmcReportError("ESMC report self-digest does not match its canonical payload") + if payload["model_id"] != spec.id or payload["dtype"] != "bfloat16": + raise EsmcReportError("ESMC report model or dtype identity is stale") + if payload["configured_backend"] != backend: + raise EsmcReportError("ESMC configured backend identity is invalid") + record_status = payload["record_status"] + if record_status not in {"measured", "unavailable"}: + raise EsmcReportError("ESMC record status is invalid") + + candidate = _esmc_require_mapping( + payload["candidate"], + { + "repo_id", + "manifest_revision", + "resolved_commit", + "checkpoint_repo_id", + "checkpoint_revision", + "weights_revision", + "runtime_revision", + "source_tree_sha256", + "runtime_bundle_sha256", + }, + context="ESMC candidate identity", + ) + expected_candidate = { + "repo_id": spec.fast.repo_id, + "manifest_revision": spec.fast.revision, + "checkpoint_repo_id": spec.artifact_checkpoint.repo_id, + "checkpoint_revision": spec.artifact_checkpoint.revision, + "weights_revision": spec.artifact_checkpoint.revision, + "runtime_revision": runtime_identity.runtime_revision, + "source_tree_sha256": runtime_identity.source_tree_sha256, + "runtime_bundle_sha256": runtime_identity.runtime_bundle_sha256, + } + for name, expected in expected_candidate.items(): + if candidate[name] != expected: + raise EsmcReportError(f"ESMC candidate {name} differs from frozen release identity") + if candidate["resolved_commit"] not in {None, spec.fast.revision}: + raise EsmcReportError("ESMC candidate resolved Hub commit is stale") + + reference = _esmc_require_mapping( + payload["reference"], + { + "repo_id", + "revision", + "state_transform", + "environment", + "reference_environment", + "reference_sources", + }, + context="ESMC reference identity", + ) + if ( + reference["repo_id"] != spec.official.repo_id + or reference["revision"] != spec.official.revision + or reference["state_transform"] != spec.family.state_transform + ): + raise EsmcReportError("ESMC reference identity differs from the pinned manifest") + _validate_esmc_reference_sources( + reference["reference_sources"], + expected_reference_sources, + ) + dynamic_reference_environment = _validate_esmc_reference_environment(reference["environment"]) + locked_reference_environment = _validate_locked_esmc_reference_environment( + reference["reference_environment"], source_root=source_root + ) + candidate_environment = _validate_esmc_candidate_environment(payload["environment"]) + _validate_esmc_environment_binding( + candidate_environment, + dynamic_reference_environment, + locked_reference_environment, + ) + _validate_esmc_kernel(payload["kernel"], backend, candidate_environment, registry) + + report_panel = payload["panel"] + if report_panel != expected_panel: + raise EsmcReportError(f"ESMC panel {panel!r} differs from its immutable definition") + report_panel = _esmc_require_object(report_panel, context="ESMC panel identity") + panel_cases = report_panel["cases"] + cases = payload["cases"] + if ( + not isinstance(panel_cases, list) + or not isinstance(cases, list) + or (len(cases) != len(panel_cases)) + ): + raise EsmcReportError("ESMC panel and per-case metrics are not aligned") + identity_fields = { + "case_id", + "sequence_length", + "sequence_sha256", + "source", + "source_sha256", + } + violations = payload["published_band_violations"] + if not isinstance(violations, list) or any( + not isinstance(item, str) or not item.strip() for item in violations + ): + raise EsmcReportError("ESMC published-band violations must be a string list") + release_gate = _esmc_require_mapping( + payload["release_gate"], {"mode", "status"}, context="ESMC release gate" + ) + if record_status == "unavailable": + if backend not in ESMC_UNAVAILABLE_BACKENDS: + raise EsmcReportError("Only locked Flash backends may be unavailable") + if payload["effective_backend"] is not None: + raise EsmcReportError("Unavailable ESMC records must not claim effective dispatch") + if payload["unavailability"] != _esmc_unavailability_identity( + backend, locked_reference_environment + ): + raise EsmcReportError("ESMC structured unavailability identity is invalid") + if payload["catastrophic_gate"] != "not_run": + raise EsmcReportError("Unavailable ESMC catastrophe gate must be not run") + if release_gate != {"mode": "availability", "status": "unavailable"}: + raise EsmcReportError("Unavailable ESMC release-gate identity is invalid") + if ( + payload["panel_tensor_metrics"] is not None + or payload["panel_logits_metrics"] is not None + or violations + ): + raise EsmcReportError("Unavailable ESMC records must not contain measurements") + if cases != panel_cases: + raise EsmcReportError("Unavailable ESMC cases must be immutable panel identities only") + return + + if backend not in ESMC_MEASURED_BACKENDS: + raise EsmcReportError("Current frozen measurements are limited to eager, SDPA, and Flex") + if payload["effective_backend"] != backend: + raise EsmcReportError("ESMC effective backend identity is invalid or fell back") + if payload["unavailability"] is not None: + raise EsmcReportError("Measured ESMC records carry unavailability metadata") + if payload["catastrophic_gate"] != "passed": + raise EsmcReportError("Measured ESMC report catastrophe gate did not pass") + if release_gate != {"mode": ESMC_RELEASE_GATE_MODES[backend], "status": "passed"}: + raise EsmcReportError("Measured ESMC release-gate identity is invalid") + + metric_context = f"{spec.id}:bf16:{backend}:{panel}" + panel_layout = _validate_esmc_tensor_metrics( + payload["panel_tensor_metrics"], + context="ESMC panel", + expected_metric_context=metric_context, + ) + _validate_esmc_logits_metrics(payload["panel_logits_metrics"], context="ESMC panel") + for index, (panel_case, raw_case) in enumerate(zip(panel_cases, cases, strict=True)): + case = _esmc_require_mapping( + raw_case, + identity_fields | {"tensor_metrics", "logits_metrics"}, + context=f"ESMC case {index}", + ) + if not isinstance(panel_case, Mapping) or any( + case[name] != panel_case[name] for name in identity_fields + ): + raise EsmcReportError(f"ESMC case {index} identity is misaligned with its panel") + _esmc_require_sha256(case["sequence_sha256"], context=f"ESMC case {index} sequence") + if case["source_sha256"] is not None: + _esmc_require_sha256(case["source_sha256"], context=f"ESMC case {index} source") + case_id = _esmc_require_text(case["case_id"], context=f"ESMC case {index} ID") + case_layout = _validate_esmc_tensor_metrics( + case["tensor_metrics"], + context=f"ESMC case {case_id}", + expected_metric_context=f"{metric_context}:case={case_id}", + ) + if case_layout != panel_layout: + raise EsmcReportError(f"ESMC case {case_id} metric layout differs from its panel") + _validate_esmc_logits_metrics(case["logits_metrics"], context=f"ESMC case {case_id}") + if backend in {"sdpa", "eager"} and violations: + raise EsmcReportError(f"ESMC strict backend {backend} has published-band violations") + + +def load_esmc_report_set( + report_root: Path, + registry: ModelRegistry, + *, + source_root: Path | None = None, + expected_runtime_identity: EsmcRuntimeIdentity | None = None, +) -> EsmcReportSet: + """Load exactly 30 immutable schema-v3 records and fail closed on any drift.""" + + source_root = (source_root or Path(__file__).resolve().parents[2]).resolve() + expected_specs = tuple(spec.id for spec in registry.by_family("esm_plusplus")) + if expected_specs != ESMC_MODEL_IDS: + raise EsmcReportError( + f"ESMC manifest inventory {expected_specs!r} differs from {ESMC_MODEL_IDS!r}" + ) + family = registry.families["esm_plusplus"] + supported_backends = tuple( + backend + for backend in family.attention + if "bfloat16" in registry.supported_attention_dtypes(family.id, backend) + ) + if supported_backends != ESMC_BACKENDS: + raise EsmcReportError( + f"ESMC BF16 backend inventory {supported_backends!r} differs from {ESMC_BACKENDS!r}" + ) + runtime_identity = expected_runtime_identity or _esmc_runtime_identity_from_source( + source_root, registry + ) + _validate_esmc_runtime_identity(runtime_identity) + panels = _expected_esmc_panels(source_root) + expected_reference_sources = _expected_biohub_source_contracts(source_root) + expected_names = { + f"{model_id}-{backend}-{panel}.json" + for model_id in ESMC_MODEL_IDS + for backend in ESMC_BACKENDS + for panel in ESMC_PANEL_KINDS + } + if report_root.is_symlink(): + raise EsmcReportError(f"ESMC report root must not be a symlink: {report_root}") + report_root = report_root.resolve() + if not report_root.exists() or not report_root.is_dir(): + raise EsmcReportError(f"ESMC report root is not a real directory: {report_root}") + entries = tuple(report_root.iterdir()) + if any(entry.is_symlink() or not entry.is_file() for entry in entries): + raise EsmcReportError("ESMC report root contains a symlink or non-file entry") + observed_names = {entry.name for entry in entries} + if len(observed_names) != len(entries): + raise EsmcReportError("ESMC report root contains duplicate path identities") + if observed_names != expected_names: + missing = sorted(expected_names.difference(observed_names)) + unexpected = sorted(observed_names.difference(expected_names)) + raise EsmcReportError( + "ESMC release evidence must contain exactly 30 records; " + f"missing={missing}, unexpected={unexpected}" + ) + + reports: list[dict[str, object]] = [] + for model_id in ESMC_MODEL_IDS: + spec = registry[model_id] + for backend in ESMC_BACKENDS: + for panel in ESMC_PANEL_KINDS: + path = report_root / f"{model_id}-{backend}-{panel}.json" + payload = _esmc_read_json(path) + _validate_esmc_report( + payload, + spec=spec, + backend=backend, + panel=panel, + expected_panel=panels[panel], + expected_reference_sources=expected_reference_sources, + runtime_identity=runtime_identity, + registry=registry, + source_root=source_root, + ) + reports.append(payload) + if len(reports) != ESMC_REPORT_COUNT: + raise EsmcReportError( + f"ESMC release evidence contains {len(reports)} validated records, expected 30" + ) + measured_count = sum(report["record_status"] == "measured" for report in reports) + unavailable_count = sum(report["record_status"] == "unavailable" for report in reports) + if measured_count != 18 or unavailable_count != 12: + raise EsmcReportError( + "ESMC release evidence must contain exactly 18 measured and 12 structured " + "unavailable records" + ) + + candidate_environments = { + json.dumps(report["environment"], sort_keys=True) for report in reports + } + reference_environments = { + json.dumps(report["reference"]["environment"], sort_keys=True) + for report in reports + if isinstance(report["reference"], Mapping) + } + locked_reference_environments = { + json.dumps(report["reference"]["reference_environment"], sort_keys=True) + for report in reports + if isinstance(report["reference"], Mapping) + } + reference_sources = { + json.dumps(report["reference"]["reference_sources"], sort_keys=True) + for report in reports + if isinstance(report["reference"], Mapping) + } + if ( + len(candidate_environments) != 1 + or len(reference_environments) != 1 + or len(locked_reference_environments) != 1 + or len(reference_sources) != 1 + ): + raise EsmcReportError( + "ESMC release evidence crosses candidate/reference devices, software " + "environments, or source attestations" + ) + candidate_environment = reports[0]["environment"] + reference = reports[0]["reference"] + if not isinstance(candidate_environment, dict): + raise EsmcReportError("Validated ESMC candidate environment is not an object") + if not isinstance(reference, dict): + raise EsmcReportError("Validated ESMC reference identity is not an object") + reference_environment = reference["reference_environment"] + if not isinstance(reference_environment, dict): + raise EsmcReportError("Validated ESMC reference environment is not an object") + return EsmcReportSet( + reports=tuple(reports), + runtime_identity=runtime_identity, + candidate_environment=candidate_environment, + reference_environment=reference_environment, + ) + + +def _esmc_kernel_label(kernel: object) -> str: + kernel = _esmc_require_object(kernel, context="Rendered ESMC kernel identity") + if kernel["provider"] == "torch": + return f"Torch {kernel['torch_version']}" + return ( + f"{kernel['repository']} " + f"v{kernel['version']} ({kernel['expected_variant']})" + ) + + +def _esmc_reference_source_table(value: object) -> list[str]: + sources = _esmc_require_mapping( + value, + set(ESMC_REFERENCE_SOURCE_NAMES), + context="Rendered ESMC reference sources", + ) + lines = [ + "Every report carries both official reference source attestations:", + "", + _table_row( + "Source", + "Schema", + "Package", + "Revision", + "Import file", + "Tree SHA-256", + "Attestation SHA-256", + "Files", + ), + _table_row("---", "---", "---", "---", "---", "---", "---", "---"), + ] + for source_name in ESMC_REFERENCE_SOURCE_NAMES: + source = _esmc_require_object( + sources[source_name], + context=f"Rendered ESMC reference source {source_name}", + ) + lines.append( + _table_row( + f"`{source_name}`", + f"`{source['schema_version']}`", + f"`{source['import_name']} {source['package_version']}`", + f"`{source['source_revision']}`", + f"`{source['import_file']}` under `{source['import_root']}`", + f"`{source['tree_sha256']}`", + f"`{source['attestation_sha256']}`", + f"`{source['file_count']}`", + ) + ) + return lines + + +def _esmc_number(value: object) -> str: + numeric = _esmc_require_finite(value, context="Rendered ESMC metric") + if numeric == 0: + return "0" + return f"{numeric:.6g}" + + +def _esmc_range(values: Iterable[object]) -> str: + numbers = [ + _esmc_require_finite(value, context="Rendered ESMC range metric") + for value in values + ] + return f"{_esmc_number(min(numbers))} to {_esmc_number(max(numbers))}" + + +def _esmc_distribution(values: Iterable[object]) -> str: + numbers = [ + _esmc_require_finite(value, context="Rendered ESMC distribution metric") + for value in values + ] + return ( + f"{_esmc_number(min(numbers))} / " + f"{_esmc_number(statistics.median(numbers))} / {_esmc_number(max(numbers))}" + ) + + +def _esmc_pip_check_disclosure( + evidence: EsmcReportSet | None, + *, + heading: str, +) -> str: + if evidence is None: + status = "The frozen oracle lock permits" + exception: Mapping[str, object] = { + "accepted_diagnostic": ( + "nvidia-cusparselt-cu13 0.8.1 is not supported on this platform" + ), + "distribution": "nvidia-cusparselt-cu13", + "version": "0.8.1", + "wheel_filename": ("nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl"), + "wheel_sha256": ("4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f"), + "filename_platform_tag": "py3-none-manylinux2014_aarch64", + "wheel_metadata_platform_tag": "py3-none-manylinux2014_sbsa", + "target_hardware": "NVIDIA GH200 480GB", + "target_operating_system": "linux", + "target_architecture": "aarch64", + "resolution": "validated-vendor-metadata-exception-no-wheel-rewrite", + } + else: + status = "The validated oracle environment recorded" + pip_check = _esmc_require_mapping( + evidence.reference_environment.get("pip_check"), + { + "status", + "returncode", + "diagnostics", + "accepted_platform_exceptions", + }, + context="Rendered ESMC pip-check evidence", + ) + diagnostics = _esmc_require_list( + pip_check["diagnostics"], context="Rendered ESMC pip-check diagnostics" + ) + exceptions = _esmc_require_list( + pip_check["accepted_platform_exceptions"], + context="Rendered ESMC pip-check platform exceptions", + ) + if ( + pip_check["status"] != "accepted-platform-exception" + or pip_check["returncode"] != 1 + or len(diagnostics) != 1 + or len(exceptions) != 1 + ): + raise EsmcReportError("Rendered ESMC pip-check exception identity is invalid") + exception = _esmc_require_object( + exceptions[0], context="Rendered ESMC pip-check platform exception" + ) + if diagnostics[0] != exception.get("accepted_diagnostic"): + raise EsmcReportError("Rendered ESMC pip-check diagnostic is not attested") + return f"""\ +{heading} Locked oracle package compatibility exception + +{status} exactly one nonzero `pip check` diagnostic: +`{exception["accepted_diagnostic"]}`. It applies only to +`{exception["distribution"]}=={exception["version"]}` on +`{exception["target_hardware"]}` / `{exception["target_operating_system"]}` / +`{exception["target_architecture"]}`. The vendor filename tag is +`{exception["filename_platform_tag"]}`, while the wheel metadata declares +`{exception["wheel_metadata_platform_tag"]}`. The exact wheel is +`{exception["wheel_filename"]}` with SHA-256 `{exception["wheel_sha256"]}`. +FastPLMs accepts this vendor metadata mismatch only after the lock, installed +inventory, wheel bytes, metadata tag, and target identity all match. The wheel +is not rewritten (`{exception["resolution"]}`). Any additional diagnostic or +identity drift fails closed. +""" + + +def _esmc_diagnostic_table( + backends: Iterable[tuple[str, str]], + *, + model_id: str, + evidence: EsmcReportSet | None, +) -> str: + backend_rows = tuple(backends) + if evidence is None: + lines = [ + _table_row("Backend", "Support", "Measurement status"), + _table_row("---", "---", "---"), + _table_row( + "`sdpa`", + "Recommended fidelity path", + "Pending release measurement", + ), + ] + for backend, support in backend_rows: + if backend in ESMC_UNAVAILABLE_BACKENDS: + status = "Unavailable on current GH200/aarch64 lock" + else: + status = "Pending release measurement" + lines.append( + _table_row( + f"`{backend}`", + support, + status, + ) + ) + return "\n".join(lines) + + display_backends = ("sdpa", *(backend for backend, _ in backend_rows)) + model_reports = tuple( + evidence.get(model_id, backend, panel) + for backend in display_backends + for panel in ESMC_PANEL_KINDS + ) + if len(model_reports) != len(display_backends) * len(ESMC_PANEL_KINDS): + raise EsmcReportError(f"ESMC evidence for {model_id!r} is incomplete") + measured_reports = tuple( + report for report in model_reports if report["record_status"] == "measured" + ) + unavailable_reports = tuple( + report for report in model_reports if report["record_status"] == "unavailable" + ) + expected_measured = len(set(display_backends).intersection(ESMC_MEASURED_BACKENDS)) * len( + ESMC_PANEL_KINDS + ) + expected_unavailable = len(set(display_backends).intersection(ESMC_UNAVAILABLE_BACKENDS)) * len( + ESMC_PANEL_KINDS + ) + if ( + len(measured_reports) != expected_measured + or len(unavailable_reports) != expected_unavailable + ): + raise EsmcReportError( + f"ESMC evidence for {model_id!r} must contain {expected_measured} measurements " + f"and {expected_unavailable} structured unavailable records" + ) + gpu = _esmc_require_object( + evidence.candidate_environment["gpu"], + context="Rendered ESMC candidate GPU identity", + ) + capability = _esmc_require_gpu_capability( + gpu["capability"], + context="Rendered ESMC candidate GPU capability", + ) + reference = _esmc_require_object( + model_reports[0]["reference"], context="Rendered ESMC reference identity" + ) + lines = [ + "The following values come from the complete validated schema-v3 release set.", + f"All reports used `{gpu['name']}` (SM{capability[0]}{capability[1]}, " + f"{gpu['total_memory_bytes']} bytes), BF16, runtime " + f"`{evidence.runtime_identity.runtime_revision}`, source tree " + f"`{evidence.runtime_identity.source_tree_sha256}`, and runtime bundle " + f"`{evidence.runtime_identity.runtime_bundle_sha256}`. Results are evidence for this", + "exact accelerator identity and are not cross-device equivalence claims.", + "", + ] + lines.extend(_esmc_reference_source_table(reference["reference_sources"])) + lines.extend( + [ + "", + "### Measurement identity", + "", + _table_row( + "Panel", + "Configured/effective", + "dtype", + "Kernel", + "Release gate", + "Catastrophe gate", + "Band warnings", + "Report SHA-256", + ), + _table_row("---", "---", "---", "---", "---", "---", "---", "---"), + ] + ) + for report in model_reports: + panel = _esmc_require_object(report["panel"], context="Rendered ESMC panel identity") + release_gate = _esmc_require_object( + report["release_gate"], context="Rendered ESMC release gate" + ) + violations = _esmc_require_list( + report["published_band_violations"], + context="Rendered ESMC band warnings", + ) + lines.append( + _table_row( + f"`{panel['kind']}` (`{str(panel['definition_sha256'])[:12]}`)", + ( + f"`{report['configured_backend']}` / `{report['effective_backend']}`" + if report["effective_backend"] is not None + else f"`{report['configured_backend']}` / not dispatched" + ), + f"`{report['dtype']}`", + _esmc_kernel_label(report["kernel"]), + f"`{release_gate['mode']}` / `{release_gate['status']}`", + f"`{report['catastrophic_gate']}`", + str(len(violations)), + f"`{report['report_sha256']}`", + ) + ) + if unavailable_reports: + lines.extend( + ( + "", + "### Locked-platform unavailable backends", + "", + "These records are availability evidence, not numerical measurements. The", + "backend remains supported, but dispatch fails closed when its locked kernel", + "is unavailable on the exact report-bound release environment named below.", + "", + _table_row( + "Backend", + "Panel", + "Platform", + "Dispatch contract", + "Historical evidence", + "Reason", + "Report SHA-256", + ), + _table_row("---", "---", "---", "---", "---", "---", "---"), + ) + ) + for report in unavailable_reports: + panel = _esmc_require_object( + report["panel"], context="Rendered unavailable ESMC panel identity" + ) + unavailable = _esmc_require_object( + report["unavailability"], context="Rendered ESMC unavailability identity" + ) + lines.append( + _table_row( + f"`{report['configured_backend']}`", + f"`{panel['kind']}`", + f"`{unavailable['platform']}` / `{unavailable['accelerator']}`", + f"`{unavailable['dispatch_contract']}`", + f"`{unavailable['historical_evidence']}`", + str(unavailable["reason"]), + f"`{report['report_sha256']}`", + ) + ) + lines.extend( + ( + "", + "### Panel aggregates", + "", + "Tensor cells are the minimum-to-maximum range across every hidden-state layer,", + "last hidden state, and logits entry in `panel_tensor_metrics`. Top-1 and JSD are", + "the panel-level `panel_logits_metrics` aggregates. These are measured values,", + "not release thresholds.", + "", + _table_row( + "Backend", + "Panel", + "Relative L2", + "Q99.9", + "Residue cosine P01", + "Pooled cosine min", + "Top-1", + "JSD", + ), + _table_row("---", "---", "---", "---", "---", "---", "---", "---"), + ) + ) + for report in measured_reports: + panel = _esmc_require_object(report["panel"], context="Rendered ESMC panel identity") + raw_metrics = _esmc_require_list( + report["panel_tensor_metrics"], context="Rendered ESMC panel metrics" + ) + metrics = [ + _esmc_require_object(metric, context="Rendered ESMC panel tensor metric") + for metric in raw_metrics + ] + logits_metrics = _esmc_require_object( + report["panel_logits_metrics"], context="Rendered ESMC logits metrics" + ) + lines.append( + _table_row( + f"`{report['configured_backend']}`", + f"`{panel['kind']}`", + _esmc_range(metric["relative_l2"] for metric in metrics), + _esmc_range(metric["relative_q999"] for metric in metrics), + _esmc_range(metric["residue_cosine_p01"] for metric in metrics), + _esmc_range(metric["pooled_cosine_min"] for metric in metrics), + _esmc_number(logits_metrics["confident_top1_agreement"]), + _esmc_number(logits_metrics["mean_jsd"]), + ) + ) + lines.extend( + ( + "", + "### Per-case distributions", + "", + "Tensor cells are minimum / median / maximum across every case, output, and", + "hidden-state layer in `cases[].tensor_metrics`. Top-1 and JSD use the same", + "minimum / median / maximum summary over `cases[].logits_metrics`.", + "", + _table_row( + "Backend", + "Panel", + "Relative L2", + "Q99.9", + "Residue cosine P01", + "Pooled cosine min", + "Top-1", + "JSD", + ), + _table_row("---", "---", "---", "---", "---", "---", "---", "---"), + ) + ) + for report in measured_reports: + panel = _esmc_require_object(report["panel"], context="Rendered ESMC panel identity") + raw_cases = _esmc_require_list(report["cases"], context="Rendered ESMC cases") + cases = [_esmc_require_object(case, context="Rendered ESMC case") for case in raw_cases] + tensor_metrics: list[Mapping[str, object]] = [] + case_logits: list[Mapping[str, object]] = [] + for case in cases: + raw_case_metrics = _esmc_require_list( + case["tensor_metrics"], context="Rendered ESMC case tensor metrics" + ) + tensor_metrics.extend( + _esmc_require_object(metric, context="Rendered ESMC case tensor metric") + for metric in raw_case_metrics + ) + case_logits.append( + _esmc_require_object( + case["logits_metrics"], context="Rendered ESMC case logits metrics" + ) + ) + lines.append( + _table_row( + f"`{report['configured_backend']}`", + f"`{panel['kind']}`", + _esmc_distribution(metric["relative_l2"] for metric in tensor_metrics), + _esmc_distribution(metric["relative_q999"] for metric in tensor_metrics), + _esmc_distribution(metric["residue_cosine_p01"] for metric in tensor_metrics), + _esmc_distribution(metric["pooled_cosine_min"] for metric in tensor_metrics), + _esmc_distribution(metric["confident_top1_agreement"] for metric in case_logits), + _esmc_distribution(metric["mean_jsd"] for metric in case_logits), + ) + ) + return "\n".join(lines) + + +def _render_esmc_capability_evidence(evidence: EsmcReportSet | None) -> list[str]: + lines = [ + "## Frozen ESMC release evidence", + "", + ] + if evidence is None: + lines.extend( + ( + "**Status: pending.** Default documentation generation never discovers or", + "trusts reports implicitly. Release rendering requires an explicitly selected,", + "complete schema-v3 set of exactly 30 records on one exact GH200/aarch64", + "target: 18 measured eager, SDPA, and Flex records plus 12 structured", + "FlashAttention 2/3 unavailable records across three checkpoints and two", + "immutable sequence panels.", + "The set must also carry the final candidate/reference image identities,", + "dependency lock, installed inventory, and official source attestations.", + "A partial, stale, malformed, self-digest-invalid, or cross-device set fails", + "closed and cannot replace this status.", + "", + ) + ) + lines.extend(_esmc_pip_check_disclosure(evidence, heading="###").rstrip().splitlines()) + lines.append("") + return lines + + gpu = _esmc_require_object( + evidence.candidate_environment["gpu"], + context="Rendered ESMC candidate GPU identity", + ) + capability = _esmc_require_gpu_capability( + gpu["capability"], + context="Rendered ESMC candidate GPU capability", + ) + reference = _esmc_require_object( + evidence.reports[0]["reference"], context="Rendered ESMC reference identity" + ) + lines.extend( + ( + f"**Status: validated complete set ({len(evidence.reports)}/30 records).**", + "The set contains 18 measured eager, SDPA, and Flex records and 12", + "structured FlashAttention 2/3 locked-platform unavailable records.", + "", + f"Exact device: `{gpu['name']}`; capability: `SM{capability[0]}" + f"{capability[1]}`; memory: `{gpu['total_memory_bytes']}` bytes; " + "dtype: `bfloat16`.", + f"Runtime revision: `{evidence.runtime_identity.runtime_revision}`; source-tree " + f"SHA-256: `{evidence.runtime_identity.source_tree_sha256}`; runtime-bundle " + f"SHA-256: `{evidence.runtime_identity.runtime_bundle_sha256}`.", + "", + ) + ) + lines.extend(_esmc_reference_source_table(reference["reference_sources"])) + lines.extend(("",)) + lines.extend(_esmc_pip_check_disclosure(evidence, heading="###").rstrip().splitlines()) + lines.extend( + ( + "", + "Results are not transferred to another accelerator identity. Each model card", + "defines and publishes the corresponding per-case minimum/median/maximum", + "distributions.", + "", + _table_row( + "Checkpoint", + "Backend", + "Panel", + "Relative L2 range", + "Q99.9 range", + "Residue cosine range", + "Pooled cosine range", + "Top-1", + "JSD", + "Band warnings", + ), + _table_row("---", "---", "---", "---", "---", "---", "---", "---", "---", "---"), + ) + ) + measured_reports = tuple( + report for report in evidence.reports if report["record_status"] == "measured" + ) + unavailable_reports = tuple( + report for report in evidence.reports if report["record_status"] == "unavailable" + ) + if len(measured_reports) != 18 or len(unavailable_reports) != 12: + raise EsmcReportError( + "Rendered ESMC set must contain 18 measurements and 12 unavailable records" + ) + for report in measured_reports: + panel = _esmc_require_object(report["panel"], context="Rendered ESMC panel identity") + raw_metrics = _esmc_require_list( + report["panel_tensor_metrics"], context="Rendered ESMC panel metrics" + ) + metrics = [ + _esmc_require_object(metric, context="Rendered ESMC panel tensor metric") + for metric in raw_metrics + ] + logits = _esmc_require_object( + report["panel_logits_metrics"], context="Rendered ESMC logits metrics" + ) + violations = _esmc_require_list( + report["published_band_violations"], + context="Rendered ESMC band warnings", + ) + lines.append( + _table_row( + f"`{report['model_id']}`", + f"`{report['configured_backend']}`", + f"`{panel['kind']}` (`{str(panel['definition_sha256'])[:12]}`)", + _esmc_range(metric["relative_l2"] for metric in metrics), + _esmc_range(metric["relative_q999"] for metric in metrics), + _esmc_range(metric["residue_cosine_p01"] for metric in metrics), + _esmc_range(metric["pooled_cosine_min"] for metric in metrics), + _esmc_number(logits["confident_top1_agreement"]), + _esmc_number(logits["mean_jsd"]), + str(len(violations)), + ) + ) + lines.extend( + ( + "", + "### Current locked-platform Flash availability", + "", + _table_row( + "Checkpoint", + "Backend", + "Panel", + "Status", + "Dispatch contract", + "Historical evidence", + "Reason", + ), + _table_row("---", "---", "---", "---", "---", "---", "---"), + ) + ) + for report in unavailable_reports: + panel = _esmc_require_object( + report["panel"], context="Rendered unavailable ESMC panel identity" + ) + unavailable = _esmc_require_object( + report["unavailability"], context="Rendered ESMC unavailability identity" + ) + lines.append( + _table_row( + f"`{report['model_id']}`", + f"`{report['configured_backend']}`", + f"`{panel['kind']}`", + (f"`unavailable` on `{unavailable['platform']}` / `{unavailable['accelerator']}`"), + f"`{unavailable['dispatch_contract']}`", + f"`{unavailable['historical_evidence']}`", + str(unavailable["reason"]), + ) + ) + lines.append("") + return lines + + +def render_support(registry: ModelRegistry) -> str: + """Render the complete support matrix without importing model code.""" + + lines = [ + GENERATED_MARKER, + "", + "# Model support", + "", + "This file is generated from `src/fastplms/models.toml`. A listed capability is", + "selectable. Strict-parity exceptions are documented in the checkpoint cards.", + "", + "## Family interfaces", + "", + "| Family | Architecture | Checkpoints | Public input | AutoClasses | Tokenizer class |", + "| --- | --- | ---: | --- | --- | --- |", + ] + for family in registry.families.values(): + count = len(registry.by_family(family.id)) + lines.append( + "| " + + " | ".join( + ( + f"`{family.id}`", + family.architecture, + str(count), + family.public_input.replace("|", "\\|"), + _code(sorted(family.auto_map)), + _tokenizer_class_label(family), + ) + ) + + " |" + ) + + lines.extend( + ( + "", + "## AutoClass weight status", + "", + "`pretrained` means the advertised head is present in the checkpoint. " + "`base weights + untrained task head` means the task head must be " + "trained before use. `FastPLMs extension` is an integration or head " + "that is not an official pretrained ANKH capability.", + "", + "| Family | AutoClass | Weight status |", + "| --- | --- | --- |", + ) + ) + for family in registry.families.values(): + for auto_class in sorted(family.auto_map): + lines.append( + f"| `{family.id}` | `{auto_class}` | `{_auto_class_status(family, auto_class)}` |" + ) + + lines.extend( + ( + "", + "## Family execution", + "", + "| Family | Attention | Precision | BF16 execution | Extra | Reference |", + "| --- | --- | --- | --- | --- | --- |", + ) + ) + for family in registry.families.values(): + lines.append( + "| " + + " | ".join( + ( + f"`{family.id}`", + _code(family.attention), + _precision_contract(family), + f"`{family.bf16_execution}`", + f"`{family.extra}`", + f"`{family.reference_container}`", + ) + ) + + " |" + ) + + lines.extend( + ( + "", + "## Family release contracts", + "", + "| Family | Checkpoint terms | Hub license | Weight publication | Tiers |", + "| --- | --- | --- | --- | --- |", + ) + ) + for family in registry.families.values(): + lines.append( + "| " + + " | ".join( + ( + f"`{family.id}`", + family.checkpoint_license.replace("|", "\\|"), + _hub_license_label(family), + ( + "blocked" + if not family.weights_publication_allowed + else "complete checkpoint required" + if family.requires_complete_weight_publication + else "manifest policy" + ), + _code(family.test_tiers), + ) + ) + + " |" + ) + + lines.extend( + ( + "", + "## Runtime assets", + "", + _table_row( + "ID", + "Family", + "Repository", + "Path", + "SHA-256", + "Size", + "License", + "Trust boundary", + "Offline behavior", + ), + _table_row("---", "---", "---", "---", "---", "---:", "---", "---", "---"), + ) + ) + for asset in registry.runtime_assets.values(): + lines.append( + "| " + + " | ".join( + ( + f"`{asset.id}`", + f"`{asset.consumer_family}`", + f"`{asset.repository}`", + f"`{asset.path}`", + f"`{asset.sha256}`", + str(asset.size), + f"`{asset.license_expression}`", + f"`{asset.trust_kind}`", + f"`{asset.offline_behavior}`", + ) + ) + + " |" + ) + + lines.extend( + ( + "", + "## Checkpoints", + "", + "| ID | Family | Size | FastPLMs checkpoint | Official checkpoint | " + "Artifact source | State transform | Generation contract | MSA conditioning | " + "Unresolved files |", + "| --- | --- | --- | --- | --- | --- | --- | --- | --- | ---: |", + ) + ) + for spec in registry.values(): + fast_url = f"https://huggingface.co/{spec.fast.repo_id}" + official_url = f"https://huggingface.co/{spec.official.repo_id}" + unresolved = len(spec.fast.unresolved_files) + len(spec.official.unresolved_files) + if spec.family.id == "esmfold2": + if spec.msa_conditioning is None: + raise ValueError(f"{spec.id}: ESMFold2 MSA conditioning is undeclared") + msa_conditioning = ( + "`optional` (full checkpoint)" + if spec.msa_conditioning + else "`none` (Fast; MSA inputs rejected)" + ) + else: + msa_conditioning = "not applicable" + lines.append( + "| " + + " | ".join( + ( + f"`{spec.id}`", + f"`{spec.family.id}`", + f"`{spec.size_category}`", + f"[{spec.fast.repo_id}]({fast_url})", + f"[{spec.official.repo_id}]({official_url})", + f"`{spec.artifact_source}`", + f"`{spec.family.state_transform}`", + f"`{spec.generation_contract}`", + msa_conditioning, + str(unresolved), + ) + ) + + " |" + ) + lines.extend( + ( + "", + "A nonzero unresolved-file count blocks release. It is not permission to", + "omit that file from checkpoint, tokenizer, artifact, or compliance checks.", + "", + ) + ) + return "\n".join(lines) + + +def render_capability_evidence( + registry: ModelRegistry, + *, + esmc_evidence: EsmcReportSet | None = None, +) -> str: + """Render the release evidence required for every advertised capability.""" + + missing_families = sorted(set(registry.families).difference(FAMILY_DOCUMENTATION)) + if missing_families: + raise ValueError( + "Capability evidence has no documentation mapping for: " + ", ".join(missing_families) + ) + + lines = [ + GENERATED_MARKER, + "", + "# Capability-to-evidence manifest", + "", + "This manifest maps every advertised FastPLMs 1.0 capability to its user", + "documentation, runnable example, and required validation tier. It is a", + "coverage contract, not a statement that an unreported run passed. The exact", + "checkpoint list and family declarations come from `src/fastplms/models.toml`.", + "", + "The Example column links a curated CLI when that interface exposes the whole", + "capability. Programmatic-only forms instead link their runnable CPU contract so", + "the manifest does not imply broader CLI coverage than the example provides.", + "", + ] + lines.extend(_render_esmc_capability_evidence(esmc_evidence)) + lines.extend(_render_evidence_selector_catalog()) + lines.extend(_render_curated_example_cpu_evidence()) + lines.extend( + ( + "## Families and AutoClasses", + "", + _table_row( + "Family", + "Tokenizer mode", + "AutoClass", + "Weight status", + "Guide", + "Family workflow and runnable entry-point contract", + "Required evidence", + ), + _table_row("---", "---", "---", "---", "---", "---", "---"), + ) + ) + for family in registry.families.values(): + guide, _ = FAMILY_DOCUMENTATION[family.id] + for auto_class in sorted(family.auto_map): + example = _autoclass_workflow_example(family, auto_class) + evidence = _render_evidence_keys( + autoclass_evidence_keys(registry, family.id, auto_class) + ) + lines.append( + "| " + + " | ".join( + ( + f"`{family.id}`", + f"`{family.tokenizer_mode}`", + f"`{auto_class}`", + f"`{_auto_class_status(family, auto_class)}`", + f"[guide]({guide})", + ( + f"[family workflow]({example}); " + "[runnable AutoClass contract](../../tests/cpu/" + "test_autoclass_evidence_matrix.py)" + ), + evidence, + ) + ) + + " |" + ) + + lines.extend( + ( + "", + "## Attention backends", + "", + "| Backend | Advertising families | Guide | Example | Required evidence |", + "| --- | --- | --- | --- | --- |", + ) + ) + advertised_backends = sorted( + {backend for family in registry.families.values() for backend in family.attention} + ) + for backend in advertised_backends: + families = sorted( + family.id for family in registry.families.values() if backend in family.attention + ) + evidence = _render_evidence_keys(attention_backend_evidence_keys(registry, backend)) + lines.append( + f"| `{backend}` | {_code(families)} | " + "[guide](../attention_backends.md) | " + "[example](../../examples/attention_switching.py) | " + f"{evidence} |" + ) + + lines.extend( + ( + "", + "## Input, embedding, and storage contracts", + "", + _table_row("Capability", "Guide", "Example", "Required evidence"), + _table_row("---", "---", "---", "---"), + ) + ) + _append_capability_rows(lines, EMBEDDING_CAPABILITY_ROWS) + + lines.extend( + ( + "", + "## Generation and adaptation contracts", + "", + _table_row("Capability", "Guide", "Example", "Required evidence"), + _table_row("---", "---", "---", "---"), + ) + ) + _append_capability_rows(lines, GENERATION_CAPABILITY_ROWS) + + lines.extend( + ( + "", + "## Structure contracts", + "", + _table_row("Capability", "Guide", "Example", "Required evidence"), + _table_row("---", "---", "---", "---"), + ) + ) + _append_capability_rows(lines, STRUCTURE_CAPABILITY_ROWS) + _append_capability_rows(lines, _esmfold2_structure_capability_rows(registry)) + lines.extend( + ( + "", + "Release evidence must name the exact head, checkpoint and runtime revisions,", + "tokenizer identity, backend, dtype, hardware, sequence or structure panel,", + "seed, environment, and input hash. Missing evidence remains visibly pending;", + "it must not be replaced by a synthetic benchmark number or an inferred claim.", + "", + ) + ) + return "\n".join(lines) + + +def _preferred_auto_class(spec: ModelSpec) -> str: + preference = ( + "AutoModel", + "AutoModelForMaskedLM", + "AutoModelForSeq2SeqLM", + "AutoModelForProteinFolding", + ) + for name in preference: + if name in spec.auto_map: + return name + return sorted(spec.auto_map)[0] + + +def _feature_statuses(spec: ModelSpec) -> tuple[tuple[str, str], ...]: + """Return concise, checkpoint-specific public capability statuses.""" + + family_id = spec.family.id + sequence_head = "AutoModelForSequenceClassification" in spec.auto_map + token_head = "AutoModelForTokenClassification" in spec.auto_map + + if family_id == "esmfold2": + embedding = "Special: ESMC state mixture to 256-wide residue embeddings" + elif family_id == "ankh": + embedding = "Special: encoder or explicitly prepared decoder states" + elif family_id == "e1": + embedding = "Special: tokenizer-free raw-sequence preparation" + elif family_id in EMBEDDING_FAMILIES: + embedding = "Supported: shared ordered embedding API" + else: + embedding = "Unavailable for this structure-only checkpoint" + + if family_id in SEQUENCE_TTT_AUTO_CLASSES: + ttt = "Supported: low-rank masked-residue adaptation" + elif family_id == "esmfold2" and "experimental" not in spec.id: + ttt = "Special: opt-in folding TTT on the ESMC backbone" + elif family_id == "esmfold": + ttt = "Unavailable: the checkpoint has no trained MLM head" + elif family_id == "esmfold2": + ttt = "Unavailable for this experimental checkpoint" + else: + ttt = "Unavailable for this inference-only checkpoint" + + if spec.family.id == "esm_plusplus": + attention = "Special: SDPA fidelity path; alternate backends have explicit bands" + else: + attention = f"Supported: {_code(spec.family.attention)}" + + compliance = ( + "Declared: exact release evidence is required" + if "compliance" in spec.family.test_tiers + else "Unavailable: this provisional family has no compliance tier" + ) + + return ( + ( + "Sequence classification", + "Supported: base weights with an untrained task head" + if sequence_head + else "Unavailable: no advertised AutoClass", + ), + ( + "Token classification", + "Supported: base weights with an untrained task head" + if token_head + else "Unavailable: no advertised AutoClass", + ), + ( + "PEFT fine-tuning", + "Supported pattern: preserve the separately trained `classifier`" + if sequence_head + else "Supported pattern: attach LoRA to the pretrained model", + ), + ("Embeddings", embedding), + ("Test-time training", ttt), + ("Attention variants", attention), + ("Compliance", compliance), + ) + + +def _capability_summary(spec: ModelSpec) -> str: + lines = [ + "## Capabilities", + "", + "| Feature | Status |", + "| --- | --- |", + ] + lines.extend(f"| {feature} | {status} |" for feature, status in _feature_statuses(spec)) + lines.extend( + ( + "", + "A supported interface is not a pretrained downstream predictor. Classification", + "heads start untrained, and declared compliance metadata is not a claim that an", + "arbitrary local build passed its release gate.", + "", + ) + ) + return "\n".join(lines) + + +def _task_head_usage(spec: ModelSpec) -> str: + if not { + "AutoModelForSequenceClassification", + "AutoModelForTokenClassification", + }.issubset(spec.auto_map): + return "" + + model_id = spec.fast.repo_id + if spec.family.id == "e1": + preparation = """\ +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = sequence_model.prep_tokens.get_batch_kwargs( + sequences, + device=sequence_model.device, +) +biological = batch["sequence_ids"].ne(-1) +""" + else: + preparation = f"""\ +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = tokenizer(sequences, padding=True, return_tensors="pt") +biological = batch["attention_mask"].bool() +for special_id in tokenizer.all_special_ids: + biological &= batch["input_ids"].ne(special_id) +""" + + tokenizer_import = ( + "" if spec.family.id == "e1" else "from transformers import AutoTokenizer\n" + ) + return f"""\ +## Downstream classification + +Both downstream AutoClasses reuse the checkpoint backbone and initialize a new, +untrained `classifier`. Sequence labels have shape `(b,)`; residue labels have +shape `(b, l)` and use `-100` outside biological positions: + +```python +import torch +{tokenizer_import}\ +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "{model_id}" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +{preparation} +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +""" + + +def _peft_usage(spec: ModelSpec) -> str: + has_classifier = "AutoModelForSequenceClassification" in spec.auto_map + if has_classifier: + model_name = "sequence_model" + task_import = ", TaskType" + task_type = " task_type=TaskType.SEQ_CLS,\n" + modules_to_save = ' modules_to_save=["classifier"],\n' + persistence = ( + "This checkpoint advertises a classification head, so the separately " + "trained `classifier` is saved with the adapter." + ) + else: + model_name = "model" + task_import = "" + task_type = "" + modules_to_save = "" + persistence = ( + "This checkpoint has no advertised classifier. Supply the task-specific " + "objective and preserve any new head through `modules_to_save`." + ) + return f"""\ +## PEFT fine-tuning + +Install the direct training dependencies, then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig{task_import}, get_peft_model + +peft_model = get_peft_model( + {model_name}, + LoraConfig( +{task_type}\ + r=8, + lora_alpha=16, + target_modules="all-linear", +{modules_to_save}\ + ), +) +``` + +{textwrap.fill(persistence, width=79)} +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +""" + + +def _sequence_ttt_usage(spec: ModelSpec) -> str: + auto_class = SEQUENCE_TTT_AUTO_CLASSES.get(spec.family.id) + if auto_class is None: + return "" + return f"""\ +## Test-time training + +TTT samples masked views of one protein and updates only injected low-rank +adapters. Base checkpoint weights remain frozen: + +```python +from transformers import {auto_class} + +ttt_model = {auto_class}.from_pretrained( + "{spec.fast.repo_id}", + trust_remote_code=True, +) +metrics = ttt_model.ttt( + seq="MSTNPKPQRKTKRNT", + ttt_config={{"steps": 3, "batch_size": 1, "seed": 7}}, +) +ttt_model.save_pretrained("adapted", safe_serialization=True) +ttt_model.ttt_reset() +print(metrics) +``` + +Persisted adapters retain their deterministic reset state. TTT adds latency +and memory, can worsen an output, and does not establish biological function. + +""" + + +def _attention_usage(spec: ModelSpec) -> str: + recommended = "sdpa" if "sdpa" in spec.family.attention else spec.family.attention[0] + declared = textwrap.fill( + f"Declared variants are {_code(spec.family.attention)}. An unavailable " + "requested backend raises instead of silently switching implementations.", + width=79, + break_long_words=False, + break_on_hyphens=False, + ) + if "compliance" in spec.family.test_tiers: + compliance = ( + "This family declares the `compliance` tier. Release evidence binds the exact " + "checkpoint, backend, dtype, hardware, inputs, and reference revision." + ) + else: + compliance = ( + "This family does not declare the `compliance` tier. Boltz2 remains " + "provisional and its structure checks must not be broadened into parity claims." + ) + return f"""\ +## Attention and compliance + +The quick start selects `{recommended}` explicitly. {declared} +`output_attentions=True` may use the documented, one-call eager fallback solely +to materialize attention tensors; the configured backend remains unchanged. + +{textwrap.fill(compliance, width=79)} + +""" + + +def _sequence_forward_usage(spec: ModelSpec) -> str: + if spec.family.id not in {"esm2", "esm_plusplus", "dplm", "ankh"}: + return "" + if spec.family.id == "ankh": + return f"""\ +## Tokenization and forward inference + +`{spec.fast.repo_id}` contains the complete encoder-decoder checkpoint. +`AutoModel` loads the encoder view without allocating the decoder, while +`AutoModelForSeq2SeqLM` loads the encoder, decoder, cross-attention, and +language-model head. + +Use the tokenizer owned by the loaded model so tokenizer files, revision, +offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned. +Pass raw protein strings without inserted residue spaces: + +```python +import torch + +tokenizer = model.tokenizer +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +""" + return f"""\ +## Tokenization and forward inference + +Load the tokenizer from the same artifact as the model. Padding is represented +explicitly by the attention mask: + +```python +import torch +from transformers import AutoTokenizer + +model_id = "{spec.fast.repo_id}" +tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=True, +) +batch = tokenizer( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + padding=True, + return_tensors="pt", +) + +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +``` + +""" + + +def _embedding_usage(spec: ModelSpec) -> str: + if spec.family.id not in EMBEDDING_FAMILIES or spec.family.id == "esmfold2": + return "" + if spec.family.id == "ankh": + return f"""\ +## Dataset embeddings + +Dataset embeddings default to the encoder final state. Select a native encoder +layer directly: + +```python +encoder_result = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="encoder", + hidden_state_index=-1, + full_embeddings=True, +) +print(encoder_result[0].tensor.shape) # (l, d) +``` + +Decoder representations require `AutoModelForSeq2SeqLM` and exactly one +aligned decoder input. ANKH does not invent a shifted target: + +```python +from transformers import AutoModelForSeq2SeqLM + +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + "{spec.fast.repo_id}", + trust_remote_code=True, +).eval() +decoder_result = seq2seq.embed_dataset( + ["MSTNPKPQRKTKRNT"], + hidden_state_source="decoder", + hidden_state_index=-1, + decoder_inputs=["M"], + full_embeddings=True, +) +print(decoder_result[0].tensor.shape) # (decoder_length, d) +``` + +Pooling excludes boundary, padding, sentinel, and other non-biological +positions. Persisted results record the selected stack, layer, inputs, masks, +and alignment policy. + +""" + return """\ +## Dataset embeddings + +The shared embedding mixin preserves input order and biological-position +masking. It accepts sequences, identified records, mappings, or a FASTA path: + +```python +pooled = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean", "std"), +) +residues = model.embed_dataset( + ["MSTNPKPQRKTKRNT"], + full_embeddings=True, +) +print(pooled[0].tensor.shape) # (2 * d,) +print(residues[0].tensor.shape) # (l, d) +``` + +Set `output` and `format="safetensors"` or `"sqlite"` for transactional, +bounded-memory persistence. Resume verifies input order, model state, tokenizer +policy, backend, dtype, and pooling configuration before appending. + +""" + + +def _family_usage_notes( + spec: ModelSpec, + *, + allow_generic: bool = False, + esmc_evidence: EsmcReportSet | None = None, +) -> str: + family_id = spec.family.id + model_id = spec.fast.repo_id + if family_id == "esm2": + return f"""\ +## Masked language modeling and contacts + +Use the masked-language-model AutoClass when logits are required: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "{model_id}" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +masked_model = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + logits = masked_model(**batch).logits + contacts = masked_model.predict_contacts( + batch["input_ids"], + batch["attention_mask"], + ) + +print(logits.shape, contacts.shape) +``` + +Contact prediction materializes attention maps and should not be enabled in a +high-throughput embedding path unless those maps are required. + +Plain `AutoModel` omits the optional ESM pooler because this masked-language- +model checkpoint contains no trained pooler weights. Pass +`add_pooling_layer=True` only when intentionally initializing and training that +head. + +""" + if family_id == "esm_plusplus": + esmc_table = _esmc_diagnostic_table( + ( + ("eager", "Supported"), + ("flash_attention_2", "Supported"), + ("flex_attention", "Supported, numerically divergent"), + ("flash_attention_3", "Supported, numerically divergent"), + ), + model_id=spec.id, + evidence=esmc_evidence, + ) + return f"""\ +## ESMC behavior + +This artifact exposes the Biohub ESMC sequence encoder and masked-language-model +head through Transformers. It is also the language-model family used by +ESMFold2. SDPA is the default and the recommended choice for highest numerical +fidelity. Flex Attention and FlashAttention 3 are supported, non-experimental +backends, but their BF16 arithmetic may be numerically divergent from SDPA. +Those deviations produce diagnostic warnings rather than strict parity +failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic +biological disagreement remain hard gates. + +The current GH200/aarch64 release environment validates eager, SDPA, and Flex. +Flash requests fail closed because compatible locked kernels are unavailable +on this platform. + +When `sequence_id` is supplied, it is authoritative for ESMC attention grouping +and padding, and `attention_mask` is ignored. Values greater than or equal to +zero are valid sequence-group IDs; `-1` denotes padding. Omit `sequence_id` to +use `attention_mask` as the padding contract. + +{esmc_table} + +{ESMC_RELEASE_DOCUMENTATION} + +""" + if family_id == "esm3": + return """\ +## Sequence inference and masked-sequence generation + +ESM3 owns its sequence preparation. This example exercises the sequence track; +the public input contract also supports structure and function tracks through +the multimodal helpers: + +```python +import torch + +batch = model.tokenize_sequences( + ["MKTAYIAKQ", "GGGG"], + device=model.device, +) +with torch.inference_mode(): + output = model(**batch) + +print(output.last_hidden_state.shape) +print(output.logits.shape) +print(output.structure_logits.shape) +print(output.function_logits.shape) +``` + +When `return_dict=False`, ESM3 follows the standard base-model tuple prefix: +`last_hidden_state`, then requested `hidden_states` and `attentions`. Multimodal +logits and extensions follow that prefix. Prefer named fields for individual +tracks. + +Generate masked sequence positions with an explicit seed: + +```python +from fastplms.models.esm3.modeling_esm3 import FastESM3GenerationConfig + +config = FastESM3GenerationConfig( + num_steps=8, + temperature=1.0, + seed=7, +) +generated = model.generate("MK____A", config) +print(generated) +``` + +Underscores mark positions to generate. Model outputs are predictions over +tracks, not experimental measurements of structure or function. + +""" + if family_id == "e1": + return """\ +## Tokenizer-free E1 input + +E1 has no tokenizer. The model retains native raw-sequence preparation, +boundary tokens, sequence positions, and retrieval-augmented context behavior. +The ordinary representation path accepts sequences directly: + +```python +result = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + pooling=("mean",), +) +print(result[0].tensor.shape) +``` + +Lower-level masked-language-model calls must use the E1 batch preparer rather +than an `AutoTokenizer`. E1 launch messages and distributed legal files retain +the attribution required by the upstream agreement. + +""" + if family_id == "dplm": + license_url = "https://github.com/bytedance/dplm/blob/main/LICENSE" + readme_url = "https://github.com/bytedance/dplm/blob/main/README.md#overview" + return f"""\ +## Diffusion sequence generation + +DPLM defines the requested length from biological positions in a tokenized +input, masks those positions, and iteratively retains confident predictions: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "{model_id}" +tokenizer = AutoTokenizer.from_pretrained(model_id) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +input_ids = tokenizer("A" * 64, return_tensors="pt")["input_ids"].cuda() + +with torch.inference_mode(): + generated_ids = generator.generate(input_ids, max_iter=100) + +sequence = tokenizer.decode( + generated_ids[0], + skip_special_tokens=True, +).replace(" ", "") +print(sequence) +``` + +Omitting `max_iter` uses the official 500-step schedule. A shorter schedule +changes the sampling process rather than providing an equivalent faster mode. + +Plain `AutoModel` omits the optional ESM pooler because this diffusion +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +DPLM1 and DPLM2 checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE]({license_url}) is Apache-2.0 and the +[README]({readme_url}) +explicitly scopes the repository release to the pretrained DPLM1 and DPLM2 +weights. FastPLMs artifacts record `weights_license_status="resolved"` and +`redistributable=true`; complete publication is permitted only after all +artifact, legal, parity, and atomic-publication preflights pass. + +""" + if family_id == "dplm2": + license_url = "https://github.com/bytedance/dplm/blob/main/LICENSE" + readme_url = "https://github.com/bytedance/dplm/blob/main/README.md#overview" + return f"""\ +## Amino-acid and structure co-generation + +DPLM2 uses separate structure and amino-acid tracks with modality-specific +boundary and mask tokens: + +```python +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +model_id = "{model_id}" +tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) +generator = AutoModelForMaskedLM.from_pretrained( + model_id, + trust_remote_code=True, +).cuda().eval() +vocab = tokenizer.get_vocab() +l = 64 +structure = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +amino_acids = [ + vocab[""], + *([vocab[""]] * l), + vocab[""], +] +input_ids = torch.tensor([structure + amino_acids], device="cuda") + +with torch.inference_mode(): + generated = generator.generate(input_ids, max_iter=100)["output_tokens"] +print(generated.shape) +``` + +Generic `cls_token`, `eos_token`, `mask_token`, and `unk_token` aliases are +intentionally unset. Callers constructing multimodal tensors must choose the +amino-acid or structure token explicitly. Raw amino-acid sequences remain +supported by `model.embed_dataset(...)`. + +Plain `AutoModel` omits the optional ESM pooler because this co-generation +checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True` +only when intentionally initializing and training that head. + +The checkpoint weights are Apache-2.0. The maintained ByteDance +[LICENSE]({license_url}) and [README]({readme_url}) document the license +basis for the pretrained DPLM1 and DPLM2 weights. Complete publication remains +subject to all artifact, legal, parity, and atomic-publication preflights. + +""" + if family_id == "ankh": + return f"""\ +## Encoder and sequence-to-sequence use + +`{spec.fast.repo_id}` contains the complete ANKH encoder-decoder checkpoint. +Use `AutoModel` for encoder embeddings and `AutoModelForSeq2SeqLM` for +task-specific decoding: + +```python +import torch +from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer + +repo_id = "{spec.fast.repo_id}" +tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) +encoder = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval() +seq2seq = AutoModelForSeq2SeqLM.from_pretrained( + repo_id, + trust_remote_code=True, +).eval() +batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt") + +with torch.inference_mode(): + encoder_hidden = encoder(**batch).last_hidden_state + generated_ids = seq2seq.generate(**batch, max_new_tokens=16) +print(encoder_hidden.shape) +print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) +``` + +ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the +official heads from FastPLMs extensions. The complete checkpoint is larger than +the former encoder-only mirror while preserving encoder-output parity. + +""" + if family_id == "boltz2": + return """\ +## Protein structure prediction + +The high-level helper prepares a protein-only input, runs the declared Boltz2 +inference core, and returns coordinates and confidence fields: + +```python +import torch + +model = model.cuda().eval() +output = model.predict_structure( + amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG", + recycling_steps=3, + num_sampling_steps=50, + diffusion_samples=1, + seed=7, +) +model.save_as_cif(output, "prediction.cif") + +print(output.sample_atom_coords.shape) +print(output.plddt, output.ptm, output.iptm) +``` + +The validation boundary below describes the currently supported inference +subset and its provisional status. The helper scopes and restores Python, +NumPy, CPU Torch, and CUDA RNG state. Parameters and prepared features remain +FP32; supported CUDA inference executes inside BF16 autocast. + +""" + if family_id == "esmfold": + return """\ +## Protein structure prediction + +ESMFold accepts a raw sequence and returns structure tensors and confidence: + +```python +import torch + +model = model.cuda().eval() +with torch.inference_mode(): + output = model.infer( + "MKTLLILAVVAAALA", + num_recycles=4, + ) + +print(output["mean_plddt"]) + +summary = model.fold_protein( + "MKTLLILAVVAAALA", + return_pdb_string=True, +) +with open("prediction.pdb", "w", encoding="utf-8") as handle: + handle.write(summary["pdb_string"]) +print(summary["plddt"], summary["ptm"]) +``` + +FastPLMs does not expose ProteinTTT for ESMFold. The pinned folding checkpoint +does not contain a trained masked-language-model head for that objective, so +`ttt()` and TTT folding requests raise explicitly. + +""" + if family_id == "esmfold2": + if spec.msa_conditioning is None: + raise ValueError(f"{spec.id}: ESMFold2 MSA conditioning is undeclared") + ttt_note = "" + binder_note = "" + esmc_table = _esmc_diagnostic_table( + ( + ("eager", "Supported"), + ("flex_attention", "Supported, numerically divergent"), + ), + model_id="esmc_6b", + evidence=esmc_evidence, + ) + if spec.msa_conditioning: + msa_contract = """\ +## Alignment-conditioning contract + +This is a full 48-block ESMFold2 checkpoint. It supports both +single-sequence inference and optional MSA-conditioned inference. Typed +multichain and multimolecule inputs may attach an MSA to each applicable +protein chain. + +""" + typed_input_contract = """\ +The typed interface also supports RNA, protein MSAs, modifications, covalent +bonds, and distogram conditioning.""" + else: + msa_contract = """\ +## Alignment-conditioning contract + +This 24-block Fast checkpoint is inference-optimized for single-sequence +conditioning and was trained without MSA conditioning. It is not +MSA-conditioned and rejects `ProteinInput.msa` and low-level MSA-derived +features. Typed multichain and multimolecule inputs remain supported when +every protein chain uses `msa=None`. Use the corresponding full ESMFold2 +checkpoint for MSA-conditioned inference. This follows the official Biohub +architecture description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf). + +""" + typed_input_contract = """\ +The typed interface also supports RNA, modifications, covalent bonds, and +distogram conditioning. Protein MSA inputs are not supported by this Fast +checkpoint; every protein chain must use `msa=None`.""" + if "experimental" not in spec.id: + ttt_note = """\ +## Optional folding TTT + +The standard and Fast checkpoints expose opt-in folding TTT on their ESMC +backbone: + +```python +adapted = model.fold_protein_ttt( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=50, + seed=7, + ttt_config={"steps": 3, "batch_size": 1, "seed": 7}, +) +print(adapted.ttt_metrics) +``` + +Entering a gradient-enabled path reloads canonical BF16 ESMC weights. TTT adds +latency and memory, can worsen a prediction, and does not calibrate confidence +or establish biological validity. Folding TTT is result-scoped: its transient +ESMC adapter modules are excluded from checkpoint state, so it is not a generic +`save_pretrained` adapter-persistence path. + +""" + else: + ttt_note = """\ +## Test-time training + +This experimental checkpoint does not expose folding TTT. Use the corresponding +standard or Fast checkpoint when opt-in ESMC-backbone adaptation is required. + +""" + binder_note = f"""\ +## Binder-design research example + +The FastPLMs binder-design workflow uses the experimental Fast Cutoff2025 +checkpoint for differentiable inversion, both experimental Cutoff2025 +checkpoints as critics, and ESM++ as the sequence prior: + +![FastPLMs EGFR minibinder design]({BINDER_IMAGE_URL}) + +```bash +python examples/binder_design_fastplms.py \\ + --target-name pd-l1 \\ + --binder-name minibinder \\ + --batch-size 4 \\ + --steps 150 \\ + --output-dir artifacts/binder-design +``` + +The workflow ranks candidates by mean iPTM across the approved critics after +the minibinder isoelectric-point filter. These are model-based prioritization +signals, not experimental evidence of affinity or specificity. See the +[complete workflow](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md). + +""" + return f"""\ +{msa_contract} +## Protein folding + +The single-protein helper returns typed structure and confidence outputs: + +```python +result = model.fold_protein( + "MSTNPKPQRKTKRNT", + num_loops=1, + num_sampling_steps=200, + num_diffusion_samples=1, + seed=7, +) +pdb_text = model.result_to_pdb(result) +cif_text = model.result_to_cif(result) +print(result.ptm, result.plddt.mean().item()) +``` + +No target structure is required. For complexes, construct the input from the +types exposed by the loaded artifact: + +```python +types = model.input_types +complex_input = types.StructurePredictionInput( + sequences=[ + types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"), + types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"), + types.DNAInput(id="C", sequence="ATGC"), + types.LigandInput(id="L", smiles="O"), + ] +) +complex_result = model.fold( + complex_input, + num_loops=1, + num_sampling_steps=200, + seed=7, +) +print(complex_result.ptm, complex_result.plddt.mean().item()) +``` + +{typed_input_contract} The public schema recognizes +`PocketConditioning`, but the pinned official runtime discards it and hard-codes +a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning +instead of silently ignoring it. Prepared `ref_pos` values are component +reference geometries created during featurization, not target coordinates. +Predicted coordinates and confidence scores are outputs and do not establish +biochemical activity. + +## Learned representation and ESMC precision + +ESMFold2 applies its learned state mixture and projection as +`H: (b, l, 81, 2560) -> Z: (b, l, 256)`. Retrieve `Z` through the public +embedding API: + +```python +representations = model.embed_dataset( + ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"], + batch_size=2, + full_embeddings=True, +) +print(representations[0].tensor.shape) # (sequence_length, 256) +``` + +`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue +tensor per single-chain input. It rejects complexes, ligands, MSAs, +chain-separated inputs, `cls`, and `parti` in the embedding path. + +Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading. +`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only, +and strict: + +```python +model.reload_esmc(precision="fp8", device="cuda:0") +print(model.esmc_precision_status) +``` + +FP8 raises when the validated CUDA and Transformer Engine path is unavailable. +Canonical BF16 weights are retained, and transient quantization state is never +serialized. + +The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex +Attention is supported and non-experimental but can be numerically divergent; +ESMFold2 does not advertise FlashAttention for the folding interface. + +{esmc_table} + +{ESMC_RELEASE_DOCUMENTATION} + +## Hash-pinned CCD runtime asset + +Structure preparation requires `ccd.pkl` from +`biohub/ESMFold2`. The manifest pins +its 417,306,584-byte size and SHA-256 +`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5` +under MIT terms. This is a trusted-deserialization boundary: FastPLMs only +allows the exact manifest repository/revision snapshot link to resolve within +that repository's contained blob directory; user-supplied asset and `cache_dir` +symlinks are rejected. The loader creates a private temporary snapshot, verifies +its size and SHA-256, and unpickles only that loader-owned snapshot, closing +path-replacement and in-place source-write races. Offline execution requires the +exact cache object and never downloads a replacement. + +{ttt_note}{binder_note}""" + if allow_generic: + return "" + raise ValueError(f"Unsupported model-card family: {family_id!r}") + + +def render_model_card( + spec: ModelSpec, + *, + allow_generic_family: bool = False, + esmc_evidence: EsmcReportSet | None = None, +) -> str: + """Render one checkpoint card whose claims are limited to manifest evidence.""" + + auto_class = _preferred_auto_class(spec) + unresolved = len(spec.fast.unresolved_files) + len(spec.official.unresolved_files) + license_yaml = render_hub_license_yaml(spec.family) + checkpoint_terms = render_checkpoint_terms(spec.family) + canonical_state_provenance = "" + tokenizer_provenance = "" + notes = "" + capability_summary = _capability_summary(spec) + attention_usage = _attention_usage(spec) + sequence_forward = _sequence_forward_usage(spec) + embedding_usage = _embedding_usage(spec) + task_head_usage = _task_head_usage(spec) + peft_usage = _peft_usage(spec) + sequence_ttt_usage = _sequence_ttt_usage(spec) + family_usage = _family_usage_notes( + spec, + allow_generic=allow_generic_family, + esmc_evidence=esmc_evidence, + ) + local_artifact = spec.fast.repo_id.rsplit("/", maxsplit=1)[-1] + public_input_intro = textwrap.fill( + "Accepted inputs are " + f"{spec.family.public_input[0].lower() + spec.family.public_input[1:]}.", + width=79, + ) + auto_class_intro = textwrap.fill( + f"Supported Transformers entry points are {_code(sorted(spec.auto_map))}.", + width=79, + ) + recommended_attention = ( + "sdpa" if "sdpa" in spec.family.attention else spec.family.attention[0] + ) + if spec.family.tokenizer_class is not None: + tokenizer_provenance = f"- Tokenizer class: `{spec.family.tokenizer_class}`\n" + if spec.canonical_state_sha256 is not None: + canonical_state_provenance = ( + "- Canonical transformed state SHA-256: " + f"`{spec.canonical_state_sha256}`\n" + "- Conversion equality attestation: recorded in `provenance.json`\n" + ) + if spec.notes and spec.family.id != "esm_plusplus": + wrapped_notes = textwrap.fill( + spec.notes, + width=79, + break_long_words=False, + break_on_hyphens=False, + ) + notes = f"""\ +## Notes and limitations + +{wrapped_notes} + +""" + auto_status = ", ".join( + f"`{name}` = `{_auto_class_status(spec.family, name)}`" for name in sorted(spec.auto_map) + ) + weights_allowed = str(spec.family.weights_publication_allowed).lower() + weights_license_status = "resolved" if spec.family.weights_publication_allowed else "unresolved" + complete_weights = str(spec.family.requires_complete_weight_publication).lower() + return f"""--- +library_name: transformers +{license_yaml} +tags: + - protein-language-model + - fastplms +--- + +{GENERATED_MARKER} + +# {spec.fast.repo_id} + +This checkpoint packages the FastPLMs `{spec.family.architecture}` implementation. + +{public_input_intro} +{auto_class_intro} + +{capability_summary} +{_installation_section(spec)}## Quick start + +```python +from transformers import {auto_class} + +model_id = "{spec.fast.repo_id}" +model = {auto_class}.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="{recommended_attention}", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/{local_artifact}` path and pass `local_files_only=True`. + +{attention_usage}{sequence_forward}{embedding_usage}{task_head_usage}{peft_usage}\ +{sequence_ttt_usage}{family_usage}{notes}## Runtime contract + +- Public input: {spec.family.public_input} +- Advertised AutoClasses: {_code(sorted(spec.auto_map))} +- AutoClass weight status: {auto_status} +- Attention implementations: {_code(spec.family.attention)} +- Precision policies: {_precision_contract(spec.family)} +- BF16 execution: `{spec.family.bf16_execution}` +- Generation contract: `{spec.generation_contract}` +- Artifact dependency set: `{"core + structure" if spec.family.extra == "structure" else "core"}` +- Weight publication allowed: `{weights_allowed}` +- Weight license status: `{weights_license_status}` +- Redistributable: `{weights_allowed}` +- Complete weight publication required: `{complete_weights}` + +## Release record + +- FastPLMs weights: `{spec.fast.repo_id}` +- Runtime revision: recorded separately in the built artifact and published commit +- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json` +{canonical_state_provenance}\ +- Official checkpoint: `{spec.official.repo_id}` +- Artifact source: `{spec.artifact_source}` +- State transform: `{spec.family.state_transform}` +{tokenizer_provenance}- Pinned upstreams: {_code(spec.family.upstreams)} +- Release tiers: {_code(spec.family.test_tiers)} +- Unresolved required file identities: `{unresolved}` + +`provenance.json` records exact file identities, conversion, source revisions, +legal texts, schema, and attestations. A nonzero unresolved count blocks release. + +## Validation boundary + +Declared tiers compare applicable configuration, tokenizer behavior, state, +and representative inference with the pinned reference. Metadata alone does +not claim a build passed, a backend is faster, or an output is biologically +valid. + +## License + +Checkpoint terms: {checkpoint_terms}. The Hub model-card identifier is +`{spec.family.hub_license}`. Applicable source licenses, notices, attribution, +and conversion records are distributed with the local artifact. Review them +before use. +""" + + +def expected_outputs( + root: Path, + registry: ModelRegistry, + *, + esmc_evidence: EsmcReportSet | None = None, +) -> dict[Path, str]: + """Return every generated path and its deterministic UTF-8 content.""" + + output = { + root / "docs" / "generated" / "support.md": render_support(registry), + root / "docs" / "generated" / "capability_evidence.md": render_capability_evidence( + registry, + esmc_evidence=esmc_evidence, + ), + } + for spec in registry.values(): + output[root / "model_cards" / f"{spec.id}.md"] = render_model_card( + spec, + esmc_evidence=esmc_evidence, + ) + return output + + +def synchronize( + root: Path, + *, + check: bool, + esmc_report_root: Path | None = None, + require_esmc_release_evidence: bool = False, +) -> list[str]: + """Write generated files or return descriptions of stale files.""" + + registry = get_model_registry() + esmc_evidence = None + if esmc_report_root is not None or require_esmc_release_evidence: + selected_root = esmc_report_root + if selected_root is None: + selected_root = Path( + os.environ.get( + "FASTPLMS_DIAGNOSTIC_REPORTS", + "artifacts/diagnostics/esmc", + ) + ) + if not selected_root.is_absolute(): + selected_root = root / selected_root + esmc_evidence = load_esmc_report_set( + selected_root, + registry, + source_root=root, + ) + outputs = expected_outputs(root, registry, esmc_evidence=esmc_evidence) + failures: list[str] = [] + for path, content in outputs.items(): + rendered = content.rstrip() + "\n" + current = path.read_text(encoding="utf-8") if path.is_file() else None + if current == rendered: + continue + if check: + failures.append(f"stale or missing generated file: {path.relative_to(root)}") + else: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(rendered, encoding="utf-8", newline="\n") + + expected_cards = {path.resolve() for path in outputs if path.parent.name == "model_cards"} + for path in sorted((root / "model_cards").glob("*.md")): + if path.name == "README.md" or path.resolve() in expected_cards: + continue + try: + generated = GENERATED_MARKER in path.read_text(encoding="utf-8") + except OSError: + generated = False + if generated and check: + failures.append(f"stale generated model card: {path.relative_to(root)}") + elif generated: + path.unlink() + return failures + + +def main(argv: Iterable[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--check", action="store_true") + parser.add_argument( + "--source-root", + type=Path, + default=Path(__file__).resolve().parents[2], + ) + parser.add_argument( + "--esmc-report-root", + type=Path, + help=( + "strictly validate and render one explicit complete 30-record schema-v3 " + "ESMC release-evidence set" + ), + ) + parser.add_argument( + "--require-esmc-release-evidence", + action="store_true", + help=( + "require release evidence from --esmc-report-root, " + "FASTPLMS_DIAGNOSTIC_REPORTS, or artifacts/diagnostics/esmc" + ), + ) + arguments = parser.parse_args(argv) + try: + failures = synchronize( + arguments.source_root.resolve(), + check=arguments.check, + esmc_report_root=arguments.esmc_report_root, + require_esmc_release_evidence=arguments.require_esmc_release_evidence, + ) + except EsmcReportError as error: + print(f"invalid ESMC release evidence: {error}") + return 1 + if failures: + for failure in failures: + print(failure) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/artifacts/license_metadata.py b/tools/artifacts/license_metadata.py new file mode 100644 index 0000000..c9db301 --- /dev/null +++ b/tools/artifacts/license_metadata.py @@ -0,0 +1,93 @@ +"""Render validated Hugging Face model-card license metadata.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Mapping +from typing import cast +from urllib.parse import urlparse +from huggingface_hub import ModelCard + +from fastplms.registry import HUB_LICENSE_IDENTIFIERS, ModelFamily + + +_HUB_LICENSE_NAME_RE = re.compile(r"^[a-z0-9][a-z0-9.-]*$") + + +def render_hub_license_yaml(family: ModelFamily) -> str: + """Return deterministic YAML fields from the typed family contract.""" + + return "\n".join( + f"{key}: {json.dumps(value, ensure_ascii=False)}" + for key, value in family.hub_license_metadata.items() + ) + + +def render_checkpoint_terms(family: ModelFamily) -> str: + """Return precise Markdown for the checkpoint's governing terms.""" + + if family.hub_license == "other": + if family.hub_license_name is None or family.hub_license_link is None: + raise ValueError("Custom Hub licenses require a name and link") + return f"[{family.hub_license_name}]({family.hub_license_link})" + return cast(str, family.checkpoint_license) + + +def validate_hub_license_metadata(metadata: Mapping[str, object]) -> dict[str, str]: + """Validate and normalize one Hub license metadata mapping.""" + + allowed_fields = {"license", "license_name", "license_link"} + if not metadata or not set(metadata).issubset(allowed_fields): + raise ValueError("Hub license metadata has missing or unknown fields") + identifier = metadata.get("license") + name = metadata.get("license_name") + link = metadata.get("license_link") + if not isinstance(identifier, str) or identifier not in HUB_LICENSE_IDENTIFIERS: + raise ValueError("Model card has an unsupported Hugging Face license identifier") + normalized = {"license": identifier} + if identifier != "other": + if name is not None or link is not None: + raise ValueError("Standard Hub licenses may not define custom license fields") + return normalized + if not isinstance(name, str) or not name.strip(): + raise ValueError("Custom Hub license is missing license_name") + if _HUB_LICENSE_NAME_RE.fullmatch(name) is None: + raise ValueError("Custom Hub license_name must be a lowercase Hub slug") + if not isinstance(link, str) or not link.strip(): + raise ValueError("Custom Hub license is missing license_link") + parsed_link = urlparse(link) + if ( + parsed_link.scheme != "https" + or not parsed_link.netloc + or not parsed_link.path + or parsed_link.username is not None + or parsed_link.password is not None + ): + raise ValueError("Custom Hub license_link must be an absolute HTTPS URL") + normalized["license_name"] = name + normalized["license_link"] = link + return normalized + + +def parse_hub_license_metadata(card_text: str) -> dict[str, str]: + """Parse and validate a card's Hugging Face license fields.""" + + try: + data = ModelCard(card_text).data + except Exception as error: + raise ValueError(f"Invalid model-card metadata: {error}") from error + metadata: dict[str, object] = {} + for key in ("license", "license_name", "license_link"): + value = getattr(data, key, None) + if value is not None: + metadata[key] = value + return validate_hub_license_metadata(metadata) + + +__all__ = [ + "parse_hub_license_metadata", + "render_checkpoint_terms", + "render_hub_license_yaml", + "validate_hub_license_metadata", +] diff --git a/tools/artifacts/offline_probe.py b/tools/artifacts/offline_probe.py new file mode 100644 index 0000000..7f06ed9 --- /dev/null +++ b/tools/artifacts/offline_probe.py @@ -0,0 +1,1363 @@ +"""Probe one local Hub artifact without importing FastPLMs from outside it.""" + +from __future__ import annotations + +import argparse +import builtins +import contextlib +import dataclasses +import gc +import hashlib +import importlib +import importlib.abc +import importlib.util +import io +import json +import os +import subprocess +import sys +import tempfile +from collections.abc import Callable, Iterable, Mapping +from pathlib import Path +from typing import Any, cast + + +_CPU_CONTRACT_MARKER = ".fastplms-cpu-contract.json" +_CPU_FORBIDDEN_READ_ROOTS = tuple( + path.resolve() + for path in ( + Path(__file__).resolve().parents[2] / "vendor" / "upstream", + Path(__file__).resolve().parents[2] / ".git" / "modules", + Path(__file__).resolve().parents[2] / "official", + ) +) + + +@dataclasses.dataclass(frozen=True) +class ProbeCase: + """One advertised AutoClass contract within a checkpoint probe.""" + + auto_class: str + class_path: str + expected_missing_key_prefixes: tuple[str, ...] = () + expected_unexpected_key_prefixes: tuple[str, ...] = () + + +class _BlockExternalFastPLMs(importlib.abc.MetaPathFinder): + """Prevent a probe from satisfying artifact imports from external source.""" + + def find_spec( + self, + fullname: str, + path: object = None, + target: object = None, + ) -> None: + del path, target + if fullname == "fastplms" and fullname not in sys.modules: + raise ModuleNotFoundError( + "Artifact remote code attempted to import FastPLMs from outside its " + "embedded runtime." + ) + return None + + +def _runtime_site_packages() -> tuple[Path, ...]: + """Return dependency roots without processing editable-install ``.pth`` files.""" + + paths: set[Path] = set() + for entry in sys.path: + if not entry: + continue + path = Path(entry).resolve() + if path.name in {"site-packages", "dist-packages"} and path.is_dir(): + paths.add(path) + return tuple(sorted(paths, key=lambda path: path.as_posix())) + + +def _add_runtime_site_packages(paths: Iterable[Path]) -> None: + """Expose installed dependencies passed explicitly to a ``python -I -S`` probe.""" + + for path in paths: + resolved = path.resolve() + if not resolved.is_dir(): + raise RuntimeError(f"Runtime site-packages path does not exist: {resolved}") + value = str(resolved) + if value not in sys.path: + sys.path.append(value) + + +def _require_artifact_isolation() -> None: + """Reject loaded FastPLMs state and guard against external source imports.""" + + if not sys.flags.isolated: + raise RuntimeError("Artifact mode must run under python -I") + if "fastplms" in sys.modules: + raise RuntimeError("FastPLMs must not be imported before artifact loading") + if not any(isinstance(finder, _BlockExternalFastPLMs) for finder in sys.meta_path): + sys.meta_path.insert(0, _BlockExternalFastPLMs()) + + +def _tensor_digest(tensor: Any) -> str: + value = tensor.detach().cpu().contiguous() + digest = hashlib.sha256() + digest.update(str(value.dtype).encode()) + digest.update(json.dumps(list(value.shape)).encode()) + digest.update(value.reshape(-1).view(__import__("torch").uint8).numpy().tobytes()) + return digest.hexdigest() + + +def _matches_key_prefix(name: str, prefixes: Iterable[str]) -> bool: + return any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes) + + +def _state_digest( + model: Any, + *, + excluded_prefixes: Iterable[str] = (), +) -> str: + digest = hashlib.sha256() + for name, tensor in sorted(model.state_dict().items()): + if _matches_key_prefix(name, excluded_prefixes): + continue + digest.update(name.encode()) + digest.update(_tensor_digest(tensor).encode()) + return digest.hexdigest() + + +def _normalize(value: Any, *, depth: int = 0, seen: set[int] | None = None) -> Any: + import torch + + if seen is None: + seen = set() + if torch.is_tensor(value): + return { + "dtype": str(value.dtype), + "shape": list(value.shape), + "sha256": _tensor_digest(value), + } + if value is None or isinstance(value, (bool, int, float, str)): + return value + if depth >= 6 or id(value) in seen: + return f"<{type(value).__module__}.{type(value).__qualname__}>" + seen.add(id(value)) + if isinstance(value, Mapping): + return { + str(key): _normalize(item, depth=depth + 1, seen=seen) + for key, item in sorted(value.items(), key=lambda pair: str(pair[0])) + } + if isinstance(value, (list, tuple)): + return [_normalize(item, depth=depth + 1, seen=seen) for item in value] + if dataclasses.is_dataclass(value): + return { + field.name: _normalize(getattr(value, field.name), depth=depth + 1, seen=seen) + for field in dataclasses.fields(value) + } + to_tuple = getattr(value, "to_tuple", None) + if callable(to_tuple): + return _normalize(to_tuple(), depth=depth + 1, seen=seen) + values = getattr(value, "__dict__", None) + if isinstance(values, dict): + return { + str(key): _normalize(item, depth=depth + 1, seen=seen) + for key, item in sorted(values.items()) + if not str(key).startswith("_") + } + return repr(value) + + +def _output_digest(output: Any) -> str: + encoded = json.dumps( + _normalize(output), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _semantic_config(config: Any) -> dict[str, Any]: + """Remove save-location metadata before configuration comparison.""" + + values = config.to_dict() + for name in ( + "_commit_hash", + "_name_or_path", + "auto_map", + "fastplms_checkpoint_hash", + "fastplms_checkpoint_repo_id", + "fastplms_checkpoint_revision", + "fastplms_model_id", + "fastplms_runtime_bundle_sha256", + "fastplms_runtime_revision", + "fastplms_source_tree_sha256", + "fastplms_weights_revision", + "transformers_version", + ): + values.pop(name, None) + return values + + +def _save_model_for_probe(model: Any, save_path: Path, implementation: str) -> None: + """Exercise the unmodified Transformers save path for every implementation.""" + + del implementation + model.save_pretrained(save_path, safe_serialization=True) + + +def _load_class(implementation: str, auto_class: str, class_path: str) -> type: + if implementation == "artifact": + import transformers + + return getattr(transformers, auto_class) + module_name, class_name = class_path.rsplit(".", maxsplit=1) + source_class = getattr(importlib.import_module(module_name), class_name) + register = getattr(source_class, "register_for_auto_class", None) + if not callable(register): + raise RuntimeError(f"{class_path} cannot register for {auto_class}") + register(auto_class) + if getattr(source_class, "_auto_class", None) != auto_class: + raise RuntimeError(f"{class_path} did not register for {auto_class}") + config_class = getattr(source_class, "config_class", None) + if auto_class != "AutoConfig" and config_class is not None: + config_register = getattr(config_class, "register_for_auto_class", None) + if not callable(config_register): + raise RuntimeError(f"{class_path} config cannot register for AutoConfig") + config_register("AutoConfig") + if getattr(config_class, "_auto_class", None) != "AutoConfig": + raise RuntimeError(f"{class_path} config did not register for AutoConfig") + return source_class + + +def _assert_complete_saved_auto_map( + save_path: Path, + *, + expected_auto_classes: set[str], +) -> None: + """Require every advertised AutoClass to survive normal serialization.""" + + try: + config = json.loads((save_path / "config.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("Saved AutoClass config is missing or malformed") from error + auto_map = config.get("auto_map") if isinstance(config, dict) else None + if not isinstance(auto_map, dict) or set(auto_map) != expected_auto_classes: + raise RuntimeError("Saved config does not preserve the complete advertised auto_map") + invalid = { + name: target + for name, target in auto_map.items() + if not isinstance(target, str) or not target or target.endswith(".None") + } + if invalid: + raise RuntimeError( + "Saved config contains null or invalid AutoClass targets: " + + json.dumps(invalid, sort_keys=True) + ) + + +def _load_kwargs( + family: str, + bf16_execution: str, + torch: Any, + attn_implementation: str | None = None, +) -> dict[str, Any]: + dtype = torch.float32 if bf16_execution == "fp32_parameters_autocast" else torch.bfloat16 + kwargs: dict[str, Any] = { + "local_files_only": True, + "dtype": dtype, + "device_map": torch.device("cuda"), + } + if family == "esmfold2": + kwargs["load_esmc"] = False + if attn_implementation is not None: + kwargs["attn_implementation"] = attn_implementation + return kwargs + + +def _tokenizer(artifact: Path, config: Any) -> Any: + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained( + artifact, + config=config, + local_files_only=True, + trust_remote_code=True, + ) + + +def _exercise( + model: Any, + artifact: Path, + family: str, + bf16_execution: str, + torch: Any, +) -> Any: + sequence = "MSTNPKPQ" + numeric_context = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if bf16_execution == "fp32_parameters_autocast" + else contextlib.nullcontext() + ) + if family == "boltz2": + with torch.inference_mode(), numeric_context: + return model.predict_structure( + sequence, + recycling_steps=1, + num_sampling_steps=2, + diffusion_samples=1, + ) + if family == "esmfold": + with torch.inference_mode(), numeric_context: + return model.fold_protein(sequence, return_pdb_string=False) + if family == "esmfold2": + # H follows Biohub's embedding-plus-80-block ordering and has shape + # (b, l, 81, 2560). This exercises the advertised learned projection + # without loading the separately pinned 6B ESMC checkpoint. + hidden_states = torch.arange( + 2 * 81 * 2560, + device="cuda", + dtype=torch.bfloat16, + ).reshape(1, 2, 81, 2560) + residue_mask = torch.tensor([[True, False]], device="cuda") + with torch.inference_mode(), numeric_context: + return model.project_esmc_hidden_states(hidden_states, residue_mask) + + prep_tokens = getattr(getattr(model, "model", None), "prep_tokens", None) + if family == "dplm2": + tokenizer = _tokenizer(artifact, model.config) + aa_sequence = f"{tokenizer.aa_cls_token}{sequence}{tokenizer.aa_eos_token}" + encoded = tokenizer( + [aa_sequence], + add_special_tokens=False, + return_tensors="pt", + padding=True, + ) + inputs = { + name: value.to("cuda") for name, value in encoded.items() if torch.is_tensor(value) + } + elif prep_tokens is not None: + batch = prep_tokens.get_batch_kwargs([sequence], device=torch.device("cuda")) + inputs = dict(batch) + # The raw-sequence preparer returns masked-LM training labels. Artifact + # inference validates each advertised AutoClass without imposing those + # token-level labels on unrelated sequence-classification heads. + inputs.pop("labels", None) + inputs["attention_mask"] = batch["sequence_ids"].ne(-1).long() + else: + encoded = _tokenizer(artifact, model.config)( + [sequence], + return_tensors="pt", + padding=True, + ) + inputs = { + name: value.to("cuda") for name, value in encoded.items() if torch.is_tensor(value) + } + if family == "esm_plusplus": + inputs["sequence_id"] = inputs["attention_mask"].bool() + # AutoModel intentionally exposes the encoder-only ANKH view while + # retaining the official T5 ``is_encoder_decoder`` configuration value. + # Decoder inputs belong only to models that actually allocate a decoder, + # such as AutoModelForSeq2SeqLM. + if getattr(model.config, "is_encoder_decoder", False) and hasattr(model, "decoder"): + inputs["decoder_input_ids"] = inputs["input_ids"] + inputs["decoder_attention_mask"] = inputs["attention_mask"] + with torch.inference_mode(), numeric_context: + return model(**inputs) + + +def _load_saved_artifact( + *, + artifact: Path, + family: str, + bf16_execution: str, + auto_class: str, + implementation: str, + attn_implementation: str | None = None, +) -> dict[str, Any]: + """Load one saved directory exactly once through its advertised AutoClass.""" + + if implementation == "artifact": + _require_artifact_isolation() + auto_type = _load_class("artifact", auto_class, "") + if auto_class == "AutoConfig": + config = auto_type.from_pretrained( + artifact, + local_files_only=True, + trust_remote_code=True, + ) + semantic = _semantic_config(config) + return { + "config": hashlib.sha256( + json.dumps(semantic, sort_keys=True, default=str).encode() + ).hexdigest() + } + + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("Offline artifact validation requires a CUDA GPU") + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + model = _load_model_exact( + auto_type, + artifact, + trust_remote_code=True, + **_load_kwargs(family, bf16_execution, torch, attn_implementation), + ).eval() + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + output = _exercise(model, artifact, family, bf16_execution, torch) + return { + "state": _state_digest(model), + "output": _output_digest(output), + } + + +def _run_isolated_reload( + *, + artifact: Path, + family: str, + bf16_execution: str, + auto_class: str, + class_path: str, + implementation: str, + source_root: Path | None, + attn_implementation: str | None = None, +) -> dict[str, Any]: + """Reload a saved directory through its AutoClass in a fresh process.""" + + with tempfile.TemporaryDirectory(prefix="fastplms-isolated-reload-") as directory: + isolation_root = Path(directory) + output = isolation_root / "reload.json" + command = [ + sys.executable, + "-I", + "-S", + str(Path(__file__).resolve()), + "--artifact", + str(artifact.resolve()), + "--family", + family, + "--bf16-execution", + bf16_execution, + "--auto-class", + auto_class, + "--class-path", + class_path, + "--implementation", + implementation, + "--output", + str(output), + "--reload-only", + ] + if source_root is not None: + command.extend(("--source-root", str(source_root.resolve()))) + for path in _runtime_site_packages(): + command.extend(("--runtime-site-package", str(path))) + if attn_implementation is not None: + command.extend(("--attn-implementation", attn_implementation)) + + environment = os.environ.copy() + environment.pop("PYTHONHOME", None) + environment.pop("PYTHONPATH", None) + environment["HF_HOME"] = ( + os.environ.get("HF_HOME", str(isolation_root / "hf-home")) + if attn_implementation is not None + else str(isolation_root / "hf-home") + ) + environment["HF_MODULES_CACHE"] = str(isolation_root / "modules") + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + environment["PYTHONNOUSERSITE"] = "1" + completed = subprocess.run( + command, + cwd=isolation_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + details = (completed.stdout + completed.stderr).strip() + raise RuntimeError( + "Isolated saved-artifact reload failed" + (f":\n{details}" if details else ".") + ) + try: + result = json.loads(output.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("Isolated saved-artifact reload produced invalid output") from error + if not isinstance(result, dict): + raise RuntimeError("Isolated saved-artifact reload output must be an object") + return result + + +def _load_model_exact( + auto_type: Any, + artifact: Path, + *, + expected_missing_key_prefixes: Iterable[str] = (), + expected_unexpected_key_prefixes: Iterable[str] = (), + **kwargs: Any, +) -> Any: + """Load a model while rejecting every undeclared weight-loading outcome.""" + + loaded = auto_type.from_pretrained( + artifact, + output_loading_info=True, + **kwargs, + ) + if not isinstance(loaded, tuple) or len(loaded) != 2: + raise RuntimeError("Transformers did not return model loading diagnostics") + model, loading_info = loaded + if not isinstance(loading_info, dict): + raise RuntimeError("Transformers returned invalid model loading diagnostics") + diagnostics: dict[str, list[Any]] = {} + for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"): + values = loading_info.get(name, []) + if not isinstance(values, (list, tuple, set, frozenset)): + raise RuntimeError(f"Transformers returned invalid {name} loading diagnostics") + diagnostics[name] = sorted(values, key=repr) + + unexpected_missing = [ + key + for key in diagnostics["missing_keys"] + if not isinstance(key, str) or not _matches_key_prefix(key, expected_missing_key_prefixes) + ] + unexpected_checkpoint_keys = [ + key + for key in diagnostics["unexpected_keys"] + if not isinstance(key, str) + or not _matches_key_prefix(key, expected_unexpected_key_prefixes) + ] + failures = { + name: values + for name, values in ( + ("missing_keys", unexpected_missing), + ("unexpected_keys", unexpected_checkpoint_keys), + ("mismatched_keys", diagnostics["mismatched_keys"]), + ("error_msgs", diagnostics["error_msgs"]), + ) + if values + } + if failures: + raise RuntimeError( + "Validated AutoModel weight loading failed: " + + json.dumps(failures, sort_keys=True, default=str) + ) + return model + + +def _prepare_probe_environment( + implementation: str, + source_root: Path | None, + *, + reload_only: bool, +) -> None: + if implementation == "artifact": + if source_root is not None: + raise ValueError("Artifact mode must not receive a repository source root") + _require_artifact_isolation() + return + if source_root is None: + raise ValueError("Source mode requires --source-root") + source_path = str(source_root.resolve()) + if source_path not in sys.path: + sys.path.insert(0, source_path) + + +def probe( + *, + artifact: Path, + family: str, + bf16_execution: str, + auto_class: str, + class_path: str, + implementation: str, + source_root: Path | None, + reload_only: bool = False, + attn_implementation: str | None = None, + expected_missing_key_prefixes: Iterable[str] = (), + expected_unexpected_key_prefixes: Iterable[str] = (), + _environment_prepared: bool = False, +) -> dict[str, Any]: + """Load, infer, save, and reload one advertised class.""" + + if not _environment_prepared: + _prepare_probe_environment( + implementation, + source_root, + reload_only=reload_only, + ) + + if reload_only: + return _load_saved_artifact( + artifact=artifact, + family=family, + bf16_execution=bf16_execution, + auto_class=auto_class, + implementation=implementation, + attn_implementation=attn_implementation, + ) + + import torch + + auto_type = _load_class(implementation, auto_class, class_path) + trust_remote_code = implementation == "artifact" + if auto_class == "AutoConfig": + config = auto_type.from_pretrained( + artifact, + local_files_only=True, + trust_remote_code=trust_remote_code, + ) + expected_auto_classes = set(config.auto_map) + first = _semantic_config(config) + result = { + "config": hashlib.sha256( + json.dumps(first, sort_keys=True, default=str).encode() + ).hexdigest() + } + with tempfile.TemporaryDirectory(prefix="fastplms-config-reload-") as directory: + config.save_pretrained(directory) + save_path = Path(directory) + _assert_complete_saved_auto_map( + save_path, + expected_auto_classes=expected_auto_classes, + ) + reloaded_result = _run_isolated_reload( + artifact=save_path, + family=family, + bf16_execution=bf16_execution, + auto_class=auto_class, + class_path=class_path, + implementation=implementation, + source_root=source_root, + attn_implementation=attn_implementation, + ) + if reloaded_result != result: + raise AssertionError("Configuration changed across isolated AutoClass save/reload") + return result + + if not torch.cuda.is_available(): + raise RuntimeError("Offline artifact validation requires a CUDA GPU") + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + + load_kwargs = { + "trust_remote_code": trust_remote_code, + **_load_kwargs(family, bf16_execution, torch, attn_implementation), + } + model = _load_model_exact( + auto_type, + artifact, + expected_missing_key_prefixes=expected_missing_key_prefixes, + expected_unexpected_key_prefixes=expected_unexpected_key_prefixes, + **load_kwargs, + ).eval() + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + output = _exercise(model, artifact, family, bf16_execution, torch) + result = { + "state": _state_digest(model), + "pretrained_state": _state_digest( + model, + excluded_prefixes=expected_missing_key_prefixes, + ), + "output": _output_digest(output), + } + + with tempfile.TemporaryDirectory(prefix="fastplms-model-reload-") as directory: + save_path = Path(directory) + expected_auto_classes = set(model.config.auto_map) + _save_model_for_probe(model, save_path, implementation) + _assert_complete_saved_auto_map( + save_path, + expected_auto_classes=expected_auto_classes, + ) + try: + tokenizer = _tokenizer(artifact, model.config) + except (OSError, ValueError): + tokenizer = None + if tokenizer is not None: + tokenizer.save_pretrained(save_path) + del output, model + torch.cuda.empty_cache() + torch.manual_seed(314159) + torch.cuda.manual_seed_all(314159) + independently_loaded = _load_model_exact( + auto_type, + artifact, + expected_missing_key_prefixes=expected_missing_key_prefixes, + expected_unexpected_key_prefixes=expected_unexpected_key_prefixes, + **load_kwargs, + ).eval() + independently_loaded_state = _state_digest( + independently_loaded, + excluded_prefixes=expected_missing_key_prefixes, + ) + del independently_loaded + torch.cuda.empty_cache() + if independently_loaded_state != result["pretrained_state"]: + raise RuntimeError( + "Pretrained AutoModel weights depend on the initialization seed; " + "one or more shared/base weights were not loaded from the checkpoint" + ) + reloaded_result = _run_isolated_reload( + artifact=save_path, + family=family, + bf16_execution=bf16_execution, + auto_class=auto_class, + class_path=class_path, + implementation=implementation, + source_root=source_root, + attn_implementation=attn_implementation, + ) + expected_reload_result = { + "state": result["state"], + "output": result["output"], + } + if reloaded_result != expected_reload_result: + raise AssertionError("Model changed across isolated AutoClass save/reload") + return result + + +def _load_probe_cases(path: Path) -> tuple[ProbeCase, ...]: + """Load a fail-closed batch description produced by the release test.""" + + try: + raw_cases = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"Invalid AutoClass probe case file: {path}") from error + if not isinstance(raw_cases, list) or not raw_cases: + raise RuntimeError("AutoClass probe case file must contain a non-empty list") + + required = { + "auto_class", + "class_path", + "expected_missing_key_prefixes", + "expected_unexpected_key_prefixes", + } + cases: list[ProbeCase] = [] + names: set[str] = set() + for index, raw_case in enumerate(raw_cases): + if not isinstance(raw_case, dict) or set(raw_case) != required: + raise RuntimeError( + f"AutoClass probe case {index} must contain exactly {sorted(required)}" + ) + auto_class = raw_case["auto_class"] + class_path = raw_case["class_path"] + missing = raw_case["expected_missing_key_prefixes"] + unexpected = raw_case["expected_unexpected_key_prefixes"] + if not isinstance(auto_class, str) or not auto_class: + raise RuntimeError(f"AutoClass probe case {index} has an invalid auto_class") + if auto_class in names: + raise RuntimeError(f"Duplicate AutoClass probe case: {auto_class}") + if not isinstance(class_path, str) or not class_path: + raise RuntimeError(f"AutoClass probe case {index} has an invalid class_path") + if ( + not isinstance(missing, list) + or not all(isinstance(prefix, str) and prefix for prefix in missing) + or not isinstance(unexpected, list) + or not all(isinstance(prefix, str) and prefix for prefix in unexpected) + ): + raise RuntimeError(f"AutoClass probe case {index} has invalid key allowances") + names.add(auto_class) + cases.append( + ProbeCase( + auto_class=auto_class, + class_path=class_path, + expected_missing_key_prefixes=tuple(missing), + expected_unexpected_key_prefixes=tuple(unexpected), + ) + ) + return tuple(cases) + + +def _release_case_memory() -> None: + """Release model objects and allocator cache between checkpoint views.""" + + gc.collect() + torch = sys.modules.get("torch") + cuda = getattr(torch, "cuda", None) + if cuda is not None and cuda.is_available(): + cuda.empty_cache() + + +def _assert_nested_cpu_output(actual: Any, expected: Any, torch: Any) -> None: + """Compare tuple and ModelOutput views without flattening nested tensors.""" + + if torch.is_tensor(expected): + if not torch.is_tensor(actual): + raise AssertionError( + f"Expected a tensor in tuple output, received {type(actual).__name__}." + ) + torch.testing.assert_close(actual, expected) + return + if isinstance(expected, (tuple, list)): + if not isinstance(actual, type(expected)) or len(actual) != len(expected): + raise AssertionError("Nested tuple/list output differs from ModelOutput.to_tuple().") + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_nested_cpu_output(actual_item, expected_item, torch) + return + if actual != expected: + raise AssertionError( + f"Tuple output differs from ModelOutput.to_tuple(): {actual!r} != {expected!r}." + ) + + +def _cpu_sequence_inputs( + family: str, + auto_class: str, + config: Any, + torch: Any, +) -> tuple[dict[str, Any], bool]: + """Return tiny tensor-only inputs and whether the advertised head owns a loss.""" + + if family == "e1": + input_ids = torch.tensor([[1, 5, 6, 2]], dtype=torch.long) + inputs: dict[str, Any] = { + "input_ids": input_ids, + "within_seq_position_ids": torch.arange(4).unsqueeze(0), + "global_position_ids": torch.arange(4).unsqueeze(0), + "sequence_ids": torch.zeros((1, 4), dtype=torch.long), + } + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + elif family == "ankh": + input_ids = torch.tensor([[2, 3, 1, 0], [4, 5, 1, 0]], dtype=torch.long) + attention_mask = input_ids.ne(0) + inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + else: + input_ids = torch.tensor( + [[0, 3, 4, 2, 1], [0, 6, 2, 1, 1]], + dtype=torch.long, + ) + attention_mask = input_ids.ne(1) + inputs = {"input_ids": input_ids, "attention_mask": attention_mask} + if family == "esm3": + inputs = { + "sequence_tokens": input_ids, + "sequence_id": attention_mask, + } + if family == "esm_plusplus": + inputs["sequence_id"] = attention_mask + + has_loss = auto_class != "AutoModel" + if auto_class == "AutoModelForSequenceClassification": + inputs["labels"] = torch.arange(input_ids.shape[0], dtype=torch.long).remainder( + int(config.num_labels) + ) + elif auto_class == "AutoModelForTokenClassification": + labels = input_ids.remainder(int(config.num_labels)) + inputs["labels"] = labels.masked_fill(~attention_mask, -100) + elif auto_class == "AutoModelForSeq2SeqLM": + decoder_input_ids = torch.tensor( + [[0, 5, 1, 0], [0, 6, 1, 0]], + dtype=torch.long, + ) + decoder_attention_mask = decoder_input_ids.ne(0) + inputs.update( + { + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + "labels": decoder_input_ids.masked_fill(~decoder_attention_mask, -100), + "use_cache": False, + } + ) + elif auto_class == "AutoModelForMaskedLM": + inputs["labels"] = input_ids.masked_fill(~attention_mask, -100) + else: + has_loss = False + return inputs, has_loss + + +def _cpu_primary_tensor(output: Any, torch: Any) -> Any: + for name in ("last_hidden_state", "logits", "embeddings", "sequence_logits"): + value = getattr(output, name, None) + if torch.is_tensor(value): + return value + for value in output.to_tuple(): + if torch.is_tensor(value): + return value + raise AssertionError("Advertised AutoClass output contains no tensor for backward().") + + +def _cpu_state_digest(model: Any) -> str: + return _state_digest(model) + + +def _cpu_resize_and_setter_contract(model: Any) -> int: + input_embeddings = model.get_input_embeddings() + if input_embeddings is None or not hasattr(input_embeddings, "num_embeddings"): + raise AssertionError("Advertised sequence AutoClass has no token input embeddings.") + model.set_input_embeddings(input_embeddings) + if model.get_input_embeddings() is not input_embeddings: + raise AssertionError("set_input_embeddings() did not install the supplied module.") + + output_embeddings = model.get_output_embeddings() + if output_embeddings is not None: + model.set_output_embeddings(output_embeddings) + if model.get_output_embeddings() is not output_embeddings: + raise AssertionError("set_output_embeddings() did not install the supplied module.") + + resized_vocab = int(input_embeddings.num_embeddings) + 1 + resized = model.resize_token_embeddings(resized_vocab) + if int(resized.num_embeddings) != resized_vocab: + raise AssertionError("resize_token_embeddings() returned the wrong vocabulary size.") + if int(model.get_input_embeddings().num_embeddings) != resized_vocab: + raise AssertionError("Input embeddings did not retain the resized vocabulary.") + resized_output = model.get_output_embeddings() + output_size = getattr(resized_output, "out_features", resized_vocab) + if resized_output is not None and int(output_size) != resized_vocab: + raise AssertionError("Output embeddings did not retain the resized vocabulary.") + return resized_vocab + + +def _probe_tiny_cpu_model( + *, + artifact: Path, + family: str, + auto_class: str, + config: Any, +) -> dict[str, Any]: + """Exercise one tiny remote AutoClass without checkpoint or accelerator access.""" + + import torch + + if family == "boltz2": + from torch import nn + + class _TinyBoltzCore(nn.Module): + def __init__(self, width: int = 3) -> None: + super().__init__() + self.weight = nn.Parameter(torch.linspace(0.5, 1.0, width)) + + module = sys.modules.get("fastplms.models.boltz.modeling_boltz2") + if module is None: + raise RuntimeError("The remote Boltz2 runtime module was not loaded.") + dynamic_module: Any = module + dynamic_module.Boltz2InferenceCore = _TinyBoltzCore + + auto_type: Any = _load_class("artifact", auto_class, "") + model = auto_type.from_config(config, trust_remote_code=True).eval() + if not model.is_remote_code() or model._auto_class != auto_class: + raise AssertionError(f"{auto_class} did not retain its remote AutoClass registration.") + expected_class = config.auto_map[auto_class].rsplit(".", maxsplit=1)[1] + if type(model).__name__ != expected_class: + raise AssertionError( + f"{auto_class} dispatched {type(model).__name__}, expected {expected_class}." + ) + + if family in {"boltz2", "esmfold", "esmfold2"}: + # Full public structure forwards use injected native-compatible cores in + # tests/unit/test_structure_output_contracts.py. This isolated artifact + # check owns remote Auto dispatch plus exact state persistence so those + # behavior tests do not need a second copy of the runtime bundler. + before = _cpu_state_digest(model) + with tempfile.TemporaryDirectory(prefix="fastplms-cpu-structure-reload-") as directory: + save_path = Path(directory) + _save_model_for_probe(model, save_path, "artifact") + reload_kwargs: dict[str, Any] = { + "local_files_only": True, + "trust_remote_code": True, + } + if family == "esmfold2": + reload_kwargs["load_esmc"] = False + reloaded = _load_model_exact(auto_type, save_path, **reload_kwargs).eval() + if _cpu_state_digest(reloaded) != before: + raise AssertionError("Structure AutoClass state changed across save/reload.") + return { + "class": type(model).__name__, + "state": before, + "structure_forward_delegated": True, + } + + inputs, has_loss = _cpu_sequence_inputs(family, auto_class, config, torch) + output_flags = { + "output_attentions": True, + "output_hidden_states": True, + } + structured = model(**inputs, **output_flags, return_dict=True) + tuple_output = model(**inputs, **output_flags, return_dict=False) + if not hasattr(structured, "to_tuple") or not isinstance(tuple_output, tuple): + raise AssertionError("Advertised AutoClass does not honor return_dict.") + _assert_nested_cpu_output(tuple_output, structured.to_tuple(), torch) + + loss = getattr(structured, "loss", None) + if has_loss: + if loss is None or not torch.isfinite(loss): + raise AssertionError("Advertised task AutoClass did not return a finite loss.") + objective = loss + else: + objective = _cpu_primary_tensor(structured, torch).float().square().mean() + objective.backward() + gradients = [ + parameter.grad + for parameter in model.parameters() + if parameter.requires_grad and parameter.grad is not None + ] + if not gradients or not all(torch.isfinite(gradient).all() for gradient in gradients): + raise AssertionError("Advertised AutoClass backward pass produced invalid gradients.") + + resized_vocab = _cpu_resize_and_setter_contract(model) + reloaded_inputs = {name: value for name, value in inputs.items() if name != "labels"} + with torch.inference_mode(): + expected_output = model(**reloaded_inputs, return_dict=True) + expected_digest = _output_digest(expected_output) + expected_state = _cpu_state_digest(model) + resaved = False + with tempfile.TemporaryDirectory(prefix="fastplms-cpu-autoclass-reload-") as directory: + save_path = Path(directory) / "first" + _save_model_for_probe(model, save_path, "artifact") + reloaded = _load_model_exact( + auto_type, + save_path, + local_files_only=True, + trust_remote_code=True, + ).eval() + if _cpu_state_digest(reloaded) != expected_state: + raise AssertionError("AutoClass state changed across save/reload.") + with torch.inference_mode(): + observed_output = reloaded(**reloaded_inputs, return_dict=True) + if _output_digest(observed_output) != expected_digest: + raise AssertionError("AutoClass output changed across save/reload.") + if family in {"ankh", "dplm", "dplm2", "esm2", "esm3", "esm_plusplus"} and ( + auto_class == "AutoModel" + ): + resave_path = Path(directory) / "second" + _save_model_for_probe(reloaded, resave_path, "artifact") + resaved_model = _load_model_exact( + auto_type, + resave_path, + local_files_only=True, + trust_remote_code=True, + ).eval() + if _cpu_state_digest(resaved_model) != expected_state: + raise AssertionError("Remote AutoClass state changed across save-resave.") + with torch.inference_mode(): + resaved_output = resaved_model(**reloaded_inputs, return_dict=True) + if _output_digest(resaved_output) != expected_digest: + raise AssertionError("Remote AutoClass output changed across save-resave.") + resaved = True + + return { + "class": type(model).__name__, + "loss": has_loss, + "resized_vocab": resized_vocab, + "state": expected_state, + "tuple_fields": len(tuple_output), + "resaved": resaved, + } + + +def probe_tiny_cpu_many( + *, + artifact: Path, + family: str, + cases: Iterable[ProbeCase], +) -> dict[str, dict[str, Any]]: + """Run every tiny family AutoClass in one isolated, offline CPU process.""" + + marker_path = artifact / _CPU_CONTRACT_MARKER + try: + marker = json.loads(marker_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError( + "--tiny-cpu-contract requires an explicit non-release CPU test artifact marker." + ) from error + if marker != { + "release_artifact": False, + "schema_version": 1, + "scope": "tests/cpu", + }: + raise RuntimeError("Invalid non-release CPU test artifact marker.") + forbidden_release_files = [ + name + for name in ("artifact-manifest.json", "provenance.json", "runtime-attestation.json") + if (artifact / name).exists() + ] + if forbidden_release_files: + raise RuntimeError( + "Tiny CPU AutoClass probes refuse release-shaped artifacts: " + + ", ".join(forbidden_release_files) + ) + + _prepare_probe_environment("artifact", None, reload_only=False) + case_tuple = tuple(cases) + config_cases = [case for case in case_tuple if case.auto_class == "AutoConfig"] + if len(config_cases) != 1: + raise ValueError("A tiny CPU AutoClass batch requires exactly one AutoConfig case.") + + from transformers import AutoConfig + + config = AutoConfig.from_pretrained( + artifact, + local_files_only=True, + trust_remote_code=True, + ) + if getattr(config, "fastplms_cpu_contract_only", None) is not True: + raise RuntimeError("Tiny CPU AutoClass config lacks its non-release identity.") + expected_config_class = config.auto_map["AutoConfig"].rsplit(".", maxsplit=1)[1] + if type(config).__name__ != expected_config_class: + raise AssertionError( + f"AutoConfig dispatched {type(config).__name__}, expected {expected_config_class}." + ) + with tempfile.TemporaryDirectory(prefix="fastplms-cpu-config-reload-") as directory: + config.save_pretrained(directory) + reloaded_config = AutoConfig.from_pretrained( + directory, + local_files_only=True, + trust_remote_code=True, + ) + if _semantic_config(reloaded_config) != _semantic_config(config): + raise AssertionError("Tiny remote AutoConfig changed across save/reload.") + + results: dict[str, dict[str, Any]] = { + "AutoConfig": { + "class": type(config).__name__, + "config": hashlib.sha256( + json.dumps(_semantic_config(config), sort_keys=True, default=str).encode() + ).hexdigest(), + } + } + for case in case_tuple: + if case.auto_class == "AutoConfig": + continue + if case.auto_class in results: + raise ValueError(f"Duplicate AutoClass probe case: {case.auto_class}") + try: + # Some model constructors normalize or annotate the config. Give + # every advertised view a fresh remote configuration instance. + model_config = AutoConfig.from_pretrained( + artifact, + local_files_only=True, + trust_remote_code=True, + ) + results[case.auto_class] = _probe_tiny_cpu_model( + artifact=artifact, + family=family, + auto_class=case.auto_class, + config=model_config, + ) + finally: + _release_case_memory() + return results + + +def _install_cpu_probe_hermetic_guards() -> None: + """Deny network access and reference-source reads in a tiny isolated probe.""" + + import socket + + def blocked(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("Network access is forbidden in the tiny CPU AutoClass probe.") + + socket_module: Any = socket + socket_type: Any = socket.socket + socket_module.create_connection = blocked + socket_module.getaddrinfo = blocked + socket_type.connect = blocked + socket_type.connect_ex = blocked + socket_type.sendto = blocked + if hasattr(socket.socket, "sendmsg"): + socket_type.sendmsg = blocked + + import huggingface_hub + import huggingface_hub._snapshot_download + import huggingface_hub.file_download + + hub_module: Any = huggingface_hub + file_download_module: Any = huggingface_hub.file_download + snapshot_download_module: Any = huggingface_hub._snapshot_download + hub_module.hf_hub_download = blocked + hub_module.snapshot_download = blocked + file_download_module.hf_hub_download = blocked + file_download_module.http_get = blocked + snapshot_download_module.snapshot_download = blocked + + def assert_portable_path(file: object) -> None: + if isinstance(file, str): + path_value = file + elif isinstance(file, os.PathLike): + path_value = os.fspath(file) + if not isinstance(path_value, str): + return + else: + return + try: + resolved = Path(path_value).resolve() + except (OSError, TypeError, ValueError): + return + if any(resolved == root or root in resolved.parents for root in _CPU_FORBIDDEN_READ_ROOTS): + raise RuntimeError( + f"Tiny CPU AutoClass probes may not access submodule/reference path: {resolved}" + ) + + original_builtin_open = cast(Callable[..., Any], builtins.open) + original_io_open = cast(Callable[..., Any], io.open) + original_os_open = cast(Callable[..., int], os.open) + + def guarded_builtin_open(file: object, *args: Any, **kwargs: Any) -> Any: + assert_portable_path(file) + return original_builtin_open(file, *args, **kwargs) + + def guarded_io_open(file: object, *args: Any, **kwargs: Any) -> Any: + assert_portable_path(file) + return original_io_open(file, *args, **kwargs) + + def guarded_os_open(file: object, *args: Any, **kwargs: Any) -> int: + assert_portable_path(file) + return original_os_open(file, *args, **kwargs) + + builtins.open = guarded_builtin_open + io.open = guarded_io_open + os.open = guarded_os_open # type: ignore[assignment] + + +def probe_many( + *, + artifact: Path, + family: str, + bf16_execution: str, + cases: Iterable[ProbeCase], + implementation: str, + source_root: Path | None, + attn_implementation: str | None = None, +) -> dict[str, dict[str, Any]]: + """Exercise all checkpoint AutoClasses sequentially in one process.""" + + _prepare_probe_environment(implementation, source_root, reload_only=False) + results: dict[str, dict[str, Any]] = {} + for case in cases: + if case.auto_class in results: + raise ValueError(f"Duplicate AutoClass probe case: {case.auto_class}") + try: + results[case.auto_class] = probe( + artifact=artifact, + family=family, + bf16_execution=bf16_execution, + auto_class=case.auto_class, + class_path=case.class_path, + implementation=implementation, + source_root=source_root, + attn_implementation=attn_implementation, + expected_missing_key_prefixes=case.expected_missing_key_prefixes, + expected_unexpected_key_prefixes=case.expected_unexpected_key_prefixes, + _environment_prepared=True, + ) + finally: + _release_case_memory() + if not results: + raise ValueError("At least one AutoClass probe case is required") + return results + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--family", required=True) + parser.add_argument( + "--bf16-execution", + required=True, + choices=("static_parameters", "fp32_parameters_autocast"), + ) + parser.add_argument("--auto-class") + parser.add_argument("--class-path") + parser.add_argument("--cases-file", type=Path) + parser.add_argument("--implementation", choices=("artifact", "source"), required=True) + parser.add_argument( + "--attn-implementation", + choices=("flash_attention_2", "flash_attention_3"), + ) + parser.add_argument("--source-root", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--reload-only", action="store_true", help=argparse.SUPPRESS) + parser.add_argument( + "--tiny-cpu-contract", + action="store_true", + help="Exercise a config-only tiny remote-code artifact on CPU", + ) + parser.add_argument( + "--expected-missing-key-prefix", + action="append", + default=[], + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--expected-unexpected-key-prefix", + action="append", + default=[], + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--runtime-site-package", + action="append", + default=[], + type=Path, + help=argparse.SUPPRESS, + ) + return parser + + +def main(argv: Iterable[str] | None = None) -> int: + parser = build_parser() + arguments = parser.parse_args(argv) + _add_runtime_site_packages(arguments.runtime_site_package) + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + if arguments.cases_file is not None: + if arguments.auto_class is not None or arguments.class_path is not None: + parser.error("--cases-file cannot be combined with --auto-class/--class-path") + if arguments.reload_only: + parser.error("--cases-file cannot be combined with --reload-only") + cases = _load_probe_cases(arguments.cases_file) + if arguments.tiny_cpu_contract: + if arguments.implementation != "artifact": + parser.error("--tiny-cpu-contract requires --implementation artifact") + if arguments.source_root is not None or arguments.attn_implementation is not None: + parser.error( + "--tiny-cpu-contract cannot receive repository source or Flash backend options" + ) + _install_cpu_probe_hermetic_guards() + result = probe_tiny_cpu_many( + artifact=arguments.artifact.resolve(), + family=arguments.family, + cases=cases, + ) + else: + result = probe_many( + artifact=arguments.artifact.resolve(), + family=arguments.family, + bf16_execution=arguments.bf16_execution, + cases=cases, + implementation=arguments.implementation, + source_root=arguments.source_root, + attn_implementation=arguments.attn_implementation, + ) + else: + if arguments.tiny_cpu_contract: + parser.error("--tiny-cpu-contract requires --cases-file") + if arguments.auto_class is None or arguments.class_path is None: + parser.error("--auto-class and --class-path are required without --cases-file") + result = probe( + artifact=arguments.artifact.resolve(), + family=arguments.family, + bf16_execution=arguments.bf16_execution, + auto_class=arguments.auto_class, + class_path=arguments.class_path, + implementation=arguments.implementation, + source_root=arguments.source_root, + reload_only=arguments.reload_only, + attn_implementation=arguments.attn_implementation, + expected_missing_key_prefixes=arguments.expected_missing_key_prefix, + expected_unexpected_key_prefixes=arguments.expected_unexpected_key_prefix, + ) + arguments.output.write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/artifacts/publish.py b/tools/artifacts/publish.py new file mode 100644 index 0000000..d35adf2 --- /dev/null +++ b/tools/artifacts/publish.py @@ -0,0 +1,1927 @@ +"""Publish validated artifacts through parent-guarded Hub commits. + +Runtime-only updates preserve the exact bytes validated during preflight. A +complete update includes weights and both attestations in one atomic commit and +may remove only obsolete registry-pinned paths. Neither mode creates +repositories, removes unpinned remote files, or changes settings. +""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import io +import json +import os +import re +import subprocess +import sys +import tarfile +import tempfile +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any, BinaryIO +from huggingface_hub import CommitOperationAdd, CommitOperationDelete, HfApi + +from fastplms.registry import FileDigest, ModelRegistry, ModelSpec, get_model_registry +from tools.artifacts.build import ( + _RUNTIME_ATTESTATION_NAME, + _RUNTIME_ATTESTATION_SCHEMA_VERSION, + _WEIGHT_INDEX, + _WEIGHT_SUFFIXES, + ArtifactError, + _artifact_auto_map, + _checkpoint_identity_hash, + _is_runtime_update_path, + _materialize_model_card, + _render_artifact_requirements, + _resolve_artifact_manifest_path, + _tree_sha256, + _validate_bootstrap, + _validate_registry_provenance, + _validate_runtime_bundle, + _validated_release_tool_snapshot, + _validated_runtime_snapshot, + hash_file, + render_model_card, + validate_artifact, +) + + +_COMPLETE_ATTESTATION_FILES = frozenset({"artifact-manifest.json", "provenance.json"}) +_REQUIRED_FILES_ONLY_PATHS = frozenset( + { + "README.md", + "config.json", + "fastplms_bundle.py", + "modeling_fastplms.py", + "requirements.txt", + "THIRD_PARTY_NOTICES.md", + "LICENSES/FastPLMs-Apache-2.0.txt", + _RUNTIME_ATTESTATION_NAME, + } +) +_GENERATED_RUNTIME_PATHS = frozenset( + { + "README.md", + "config.json", + "fastplms_bundle.py", + "modeling_fastplms.py", + "requirements.txt", + "THIRD_PARTY_NOTICES.md", + _RUNTIME_ATTESTATION_NAME, + } +) +_RUNTIME_SUFFIXES = frozenset({".json", ".lock", ".py", ".toml"}) +_SENSITIVE_NAMES = frozenset( + { + ".env", + ".netrc", + "credentials", + "credentials.json", + "id_ed25519", + "id_rsa", + "secrets.json", + "token", + "token.txt", + } +) +_SENSITIVE_SUFFIXES = frozenset({".key", ".p12", ".pfx", ".pem"}) +_SENSITIVE_STEMS = frozenset( + {"credential", "credentials", "secret", "secrets", "token"} +) +_MAX_RUNTIME_FILE_BYTES = 128 * 1024**2 +_MAX_RETAINED_COMPLETE_BYTES = 128 * 1024**2 +_MAX_DECLARED_ASSET_BYTES = 2 * 1024**3 +_MAX_RELEASE_TEXT_BYTES = 8 * 1024**2 +_REQUIRED_COMPLETE_AUTOMODEL_VIEWS = ("AutoModel", "AutoModelForSeq2SeqLM") + + +@dataclass(frozen=True) +class FilesOnlyPublishPlan: + """One preflighted add-only commit for a manifest-declared model.""" + + model_id: str + repo_id: str + revision: str + parent_commit: str + artifact_path: Path + files: tuple[str, ...] + payloads: tuple[tuple[str, bytes], ...] + runtime_revision: str + source_tree_sha256: str + runtime_bundle_sha256: str + release_tool_revision: str + release_tool_sha256: str + release_revision: str + release_source_sha256: str + + +@dataclass(frozen=True) +class CompletePublishPlan: + """One complete, atomic weights-plus-runtime commit.""" + + model_id: str + repo_id: str + revision: str + parent_commit: str + artifact_path: Path + files: tuple[str, ...] + digests: tuple[tuple[str, str], ...] + deletes: tuple[str, ...] = () + replacement_weight_paths: tuple[str, ...] = () + runtime_revision: str | None = None + source_tree_sha256: str | None = None + runtime_bundle_sha256: str | None = None + release_tool_revision: str | None = None + release_tool_sha256: str | None = None + release_revision: str | None = None + release_source_sha256: str | None = None + validation_manifest_sha256: str | None = None + validated_auto_classes: tuple[str, ...] = () + + +@dataclass(frozen=True) +class FilesOnlyPublishResult: + """Identity of one completed Hub commit.""" + + model_id: str + repo_id: str + commit_oid: str + commit_url: str + + +def _is_weight_path(path: str) -> bool: + relative = PurePosixPath(path) + name = relative.name.lower() + return ( + relative.suffix.lower() in _WEIGHT_SUFFIXES + or name == _WEIGHT_INDEX + or name.endswith(".safetensors.index.json") + or name.endswith(".bin.index.json") + ) + + +def _load_json_object(path: Path, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError(f"Unable to read {label}: {path}") from error + if not isinstance(value, dict): + raise ArtifactError(f"{label} must contain a JSON object: {path}") + return value + + +def _encoded_digest(path: Path, encoded: str) -> str: + try: + algorithm, expected = encoded.split(":", maxsplit=1) + except ValueError as error: + raise ArtifactError(f"Invalid artifact digest for {path}: {encoded!r}") from error + actual = hash_file(path, algorithm) + if actual != expected: + raise ArtifactError( + f"Artifact file digest differs for {path}: expected {expected}, received {actual}." + ) + return actual + + +def _read_validated_bytes( + path: Path, + encoded: str, + *, + max_bytes: int | None = None, +) -> bytes: + """Read once and validate the exact immutable payload retained by a plan.""" + + try: + algorithm, expected = encoded.split(":", maxsplit=1) + except ValueError as error: + raise ArtifactError(f"Invalid artifact digest for {path}: {encoded!r}") from error + with path.open("rb") as handle: + payload = handle.read(max_bytes + 1 if max_bytes is not None else -1) + if max_bytes is not None and len(payload) > max_bytes: + raise ArtifactError( + f"Artifact payload exceeds its {max_bytes}-byte retained publication limit: {path}" + ) + if algorithm == "sha256": + actual = hashlib.sha256(payload).hexdigest() + elif algorithm == "git-sha1": + digest = hashlib.sha1(usedforsecurity=False) + digest.update(f"blob {len(payload)}\0".encode("ascii")) + digest.update(payload) + actual = digest.hexdigest() + else: + raise ArtifactError(f"Unsupported artifact digest algorithm: {algorithm!r}") + if actual != expected: + raise ArtifactError( + f"Artifact file digest differs for {path}: expected {expected}, received {actual}." + ) + return payload + + +def _open_validated_payload( + handle: BinaryIO, + path: Path, + encoded: str, +) -> tuple[int, int, int, int]: + """Hash one retained open file and return its stable filesystem identity.""" + + try: + algorithm, expected = encoded.split(":", maxsplit=1) + except ValueError as error: + raise ArtifactError(f"Invalid artifact digest for {path}: {encoded!r}") from error + before = os.fstat(handle.fileno()) + if algorithm == "sha256": + digest = hashlib.sha256() + elif algorithm == "git-sha1": + digest = hashlib.sha1(usedforsecurity=False) + digest.update(f"blob {before.st_size}\0".encode("ascii")) + else: + raise ArtifactError(f"Unsupported artifact digest algorithm: {algorithm!r}") + handle.seek(0) + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + after = os.fstat(handle.fileno()) + identity = (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) + if identity != (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns): + raise ArtifactError(f"Artifact file changed while it was hashed: {path}") + actual = digest.hexdigest() + if actual != expected: + raise ArtifactError( + f"Artifact file digest differs for {path}: expected {expected}, received {actual}." + ) + handle.seek(0) + return identity + + +def _snapshot_validated_payload( + source: Path, + destination: Path, + encoded: str, +) -> None: + """Copy one large payload into publisher-owned storage while hashing it.""" + + try: + algorithm, expected = encoded.split(":", maxsplit=1) + except ValueError as error: + raise ArtifactError(f"Invalid artifact digest for {source}: {encoded!r}") from error + try: + with source.open("rb") as reader, destination.open("xb") as writer: + before = os.fstat(reader.fileno()) + if algorithm == "sha256": + digest = hashlib.sha256() + elif algorithm == "git-sha1": + digest = hashlib.sha1(usedforsecurity=False) + digest.update(f"blob {before.st_size}\0".encode("ascii")) + else: + raise ArtifactError( + f"Unsupported artifact digest algorithm: {algorithm!r}" + ) + for chunk in iter(lambda: reader.read(1024 * 1024), b""): + digest.update(chunk) + writer.write(chunk) + writer.flush() + os.fsync(writer.fileno()) + after = os.fstat(reader.fileno()) + except ArtifactError: + raise + except OSError as error: + raise ArtifactError(f"Unable to snapshot complete artifact payload: {source}") from error + if ( + before.st_dev, + before.st_ino, + before.st_size, + before.st_mtime_ns, + ) != ( + after.st_dev, + after.st_ino, + after.st_size, + after.st_mtime_ns, + ): + raise ArtifactError(f"Artifact file changed while it was snapshotted: {source}") + actual = digest.hexdigest() + if actual != expected: + raise ArtifactError( + f"Artifact file digest differs for {source}: expected {expected}, received {actual}." + ) + + +def _canonical_weight_paths(provenance: Mapping[str, Any]) -> frozenset[str]: + canonical = provenance.get("canonical_weights") + if not isinstance(canonical, Mapping): + raise ArtifactError("Artifact provenance is missing canonical_weights.") + index = canonical.get("index") + shards = canonical.get("shards") + if not isinstance(index, str) or not isinstance(shards, Mapping): + raise ArtifactError("Artifact provenance has invalid canonical weight metadata.") + if any(not isinstance(path, str) for path in shards): + raise ArtifactError("Artifact provenance contains an invalid canonical shard path.") + return frozenset({index, *shards}) + + +def _assert_current_registry_spec(spec: ModelSpec) -> ModelRegistry: + registry = get_model_registry() + try: + current = registry[spec.id] + except KeyError as error: + raise ArtifactError(f"Model {spec.id!r} is absent from the current registry.") from error + if current != spec: + raise ArtifactError( + f"Model {spec.id!r} differs from the current registry; rebuild the plan." + ) + return registry + + +def _sensitive_path(relative_name: str) -> bool: + relative = PurePosixPath(relative_name) + lowered = tuple(part.lower() for part in relative.parts) + return ( + any(part in _SENSITIVE_NAMES for part in lowered) + or any(PurePosixPath(part).stem in _SENSITIVE_STEMS for part in lowered) + or relative.suffix.lower() in _SENSITIVE_SUFFIXES + or any(part in {".git", "__pycache__"} for part in lowered) + ) + + +def _declared_non_weight_assets( + spec: ModelSpec, + registry: ModelRegistry, +) -> frozenset[str]: + paths = { + item.path for item in spec.artifact_checkpoint.files if not _is_weight_path(item.path) + } + tokenizer_source = ( + registry[spec.tokenizer_source_id].official + if spec.tokenizer_source_id is not None + else spec.official + ) + paths.update(item.path for item in tokenizer_source.files if not _is_weight_path(item.path)) + return frozenset(paths) + + +def _declared_legal_paths( + spec: ModelSpec, + registry: ModelRegistry, +) -> frozenset[str]: + paths = {"LICENSES/FastPLMs-Apache-2.0.txt"} + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + paths.update( + f"LICENSES/{source_id}/{item.path}" for item in source.distribution_files + ) + return frozenset(paths) + + +def _canonical_release_bytes(raw: bytes, path: Path) -> bytes: + try: + text = raw.decode("utf-8") + except UnicodeDecodeError as error: + raise ArtifactError(f"Release text must be valid UTF-8: {path}") from error + if "\x00" in text: + raise ArtifactError(f"Release text contains a NUL byte: {path}") + return text.replace("\r\n", "\n").replace("\r", "\n").encode("utf-8") + + +def _validated_release_text_snapshot( + spec: ModelSpec, + registry: ModelRegistry, + *, + runtime_revision: str, + source_tree_sha256: str, + runtime_bundle_sha256: str, +) -> tuple[str, str, dict[str, bytes]]: + """Return release texts from one clean, tracked immutable Git revision.""" + + source_root = Path(__file__).resolve().parents[2] + _validated_release_tool_snapshot(source_root) + git_metadata = source_root / ".git" + if not (git_metadata.exists() or git_metadata.is_symlink()): + raise ArtifactError("Publication release texts require a verifiable Git worktree.") + command_prefix = [ + "git", + "-c", + f"safe.directory={source_root.as_posix()}", + ] + card_source = source_root / "model_cards" / f"{spec.id}.md" + card_relative = card_source.relative_to(source_root).as_posix() + try: + revision = subprocess.run( + [*command_prefix, "rev-parse", "HEAD"], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + card_tree = subprocess.run( + [*command_prefix, "ls-tree", "-z", revision, "--", card_relative], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as error: + raise ArtifactError("Unable to inspect the tracked publication model card.") from error + if re.fullmatch(r"[0-9a-f]{40}", revision) is None: + raise ArtifactError(f"Git returned an invalid release revision: {revision!r}") + card_entries = [entry for entry in card_tree.split(b"\0") if entry] + if len(card_entries) > 1: + raise ArtifactError("Git returned ambiguous publication model-card entries.") + generated_card = not card_entries + if card_entries: + try: + metadata, archived_name = card_entries[0].decode("utf-8").split("\t", maxsplit=1) + mode, kind, _ = metadata.split(maxsplit=2) + except (UnicodeDecodeError, ValueError) as error: + raise ArtifactError("Git returned invalid publication model-card metadata.") from error + if mode != "100644" or kind != "blob" or archived_name != card_relative: + raise ArtifactError("The tracked publication model card is not a regular file.") + elif card_source.exists() or card_source.is_symlink(): + raise ArtifactError( + f"Publication model card is untracked at the validated revision: {card_source}" + ) + sources: dict[str, Path] = { + "LICENSES/FastPLMs-Apache-2.0.txt": source_root / "LICENSE", + "THIRD_PARTY_NOTICES.md": source_root / "THIRD_PARTY_NOTICES.md", + } + if generated_card: + sources[".source/tools/artifacts/build.py"] = source_root.joinpath( + "tools", "artifacts", "build.py" + ) + sources[".source/tools/artifacts/generate_docs.py"] = source_root.joinpath( + "tools", "artifacts", "generate_docs.py" + ) + else: + sources["README.md"] = card_source + for source_id in spec.family.upstreams: + source = registry.upstreams[source_id] + for item in source.distribution_files: + sources[f"LICENSES/{source_id}/{item.path}"] = source_root.joinpath( + "LICENSES", + source_id, + *PurePosixPath(item.path).parts, + ) + relative_sources: dict[str, str] = {} + for artifact_name, path in sources.items(): + if path.is_symlink() or not path.is_file(): + raise ArtifactError( + f"Publication release source must be a regular non-symlink file: {path}" + ) + try: + size = path.stat().st_size + except OSError as error: + raise ArtifactError(f"Unable to inspect publication release source: {path}") from error + if size > _MAX_RELEASE_TEXT_BYTES: + raise ArtifactError( + f"Publication release source exceeds {_MAX_RELEASE_TEXT_BYTES} bytes: {path}" + ) + try: + relative_sources[path.relative_to(source_root).as_posix()] = artifact_name + except ValueError as error: + raise ArtifactError(f"Release source escapes the repository: {path}") from error + if len(relative_sources) != len(sources): + raise ArtifactError("Release source paths contain an ambiguous duplicate.") + source_names = tuple(sorted(relative_sources)) + try: + status = subprocess.run( + [ + *command_prefix, + "status", + "--porcelain=v1", + "--untracked-files=all", + "--", + *source_names, + ], + cwd=source_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if status: + raise ArtifactError( + "Publication release texts must be tracked and clean; scoped Git status: " + + status.replace("\n", "; ") + ) + tracked = subprocess.run( + [*command_prefix, "ls-files", "-z", "--", *source_names], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + tracked_names = {raw.decode("utf-8") for raw in tracked.split(b"\0") if raw} + if tracked_names != set(source_names): + raise ArtifactError("Publication release texts contain untracked or missing files.") + archive = subprocess.run( + [ + *command_prefix, + "archive", + "--format=tar", + revision, + "--", + *source_names, + ], + cwd=source_root, + check=True, + capture_output=True, + ).stdout + archived: dict[str, bytes] = {} + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as handle: + for member in handle.getmembers(): + if member.isdir(): + continue + if not member.isfile() or member.name in archived: + raise ArtifactError( + f"Tracked release source is not a unique regular file: {member.name}" + ) + extracted = handle.extractfile(member) + if extracted is None: + raise ArtifactError(f"Unable to read tracked release source: {member.name}") + payload = extracted.read(_MAX_RELEASE_TEXT_BYTES + 1) + if len(payload) > _MAX_RELEASE_TEXT_BYTES: + raise ArtifactError( + f"Tracked release source exceeds {_MAX_RELEASE_TEXT_BYTES} bytes: " + f"{member.name}" + ) + archived[PurePosixPath(member.name).as_posix()] = payload + except ArtifactError: + raise + except ( + OSError, + subprocess.CalledProcessError, + UnicodeDecodeError, + tarfile.TarError, + ) as error: + raise ArtifactError("Unable to validate tracked publication release texts.") from error + if set(archived) != set(source_names): + raise ArtifactError("Tracked release archive differs from the validated source allowlist.") + payloads: dict[str, bytes] = {} + identity_payloads: dict[str, bytes] = {} + for source_name, artifact_name in relative_sources.items(): + raw = archived[source_name] + identity_payloads[artifact_name] = raw + if not artifact_name.startswith(".source/"): + payloads[artifact_name] = ( + raw + if artifact_name == "README.md" + else _canonical_release_bytes(raw, source_root / source_name) + ) + if generated_card: + payloads["README.md"] = render_model_card(spec).encode("utf-8") + try: + card_template = payloads["README.md"].decode("utf-8") + except UnicodeDecodeError as error: + raise ArtifactError("Publication model-card template is not valid UTF-8.") from error + payloads["README.md"] = _materialize_model_card( + card_template, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ).encode("utf-8") + identity_payloads["README.md"] = payloads["README.md"] + identity = { + name: hashlib.sha256(payload).hexdigest() + for name, payload in identity_payloads.items() + } + digest = hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + return revision, digest, payloads + + +def _assert_current_release_texts( + artifact_path: Path, + spec: ModelSpec, + registry: ModelRegistry, +) -> tuple[str, str]: + """Bind cards and legal texts to the current validated source tree.""" + + provenance = _load_json_object(artifact_path / "provenance.json", "provenance.json") + runtime_revision = provenance.get("runtime_revision") + source_tree_sha256 = provenance.get("source_tree_sha256") + runtime_bundle_sha256 = provenance.get("runtime_bundle_sha256") + if ( + not isinstance(runtime_revision, str) + or not isinstance(source_tree_sha256, str) + or not isinstance(runtime_bundle_sha256, str) + ): + raise ArtifactError("Artifact provenance lacks a complete model-card runtime identity.") + release_revision, release_digest, expected = _validated_release_text_snapshot( + spec, + registry, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ) + try: + actual_card = (artifact_path / "README.md").read_bytes() + except OSError as error: + raise ArtifactError("Artifact model card is missing.") from error + if b"" in actual_card: + raise ArtifactError( + "Artifact model card retains an unresolved runtime-revision placeholder." + ) + if actual_card != expected["README.md"]: + raise ArtifactError("Artifact model card differs from the current source tree.") + for relative_name, expected_payload in expected.items(): + if relative_name == "README.md": + continue + path = artifact_path.joinpath(*PurePosixPath(relative_name).parts) + try: + actual = path.read_bytes() + except OSError as error: + raise ArtifactError( + f"Artifact is missing current legal text {relative_name!r}." + ) from error + if actual != expected_payload: + raise ArtifactError( + f"Artifact legal text differs from the current source: {relative_name!r}." + ) + return release_revision, release_digest + + +def _validate_publishable_non_weight_path( + relative_name: str, + path: Path, + *, + declared_assets: frozenset[str], + declared_legal_paths: frozenset[str], +) -> None: + if _sensitive_path(relative_name): + raise ArtifactError(f"Artifact contains a sensitive publication path: {relative_name!r}") + relative = PurePosixPath(relative_name) + size = path.stat().st_size + allowed = False + size_limit = _MAX_RUNTIME_FILE_BYTES + if relative_name in _GENERATED_RUNTIME_PATHS: + allowed = True + elif relative.parts and relative.parts[0] == "fastplms": + allowed = relative.suffix.lower() in _RUNTIME_SUFFIXES + elif relative_name in declared_legal_paths: + allowed = True + elif relative_name in declared_assets: + allowed = True + size_limit = _MAX_DECLARED_ASSET_BYTES + if not allowed: + raise ArtifactError( + f"Artifact path is outside the publication allowlist: {relative_name!r}" + ) + if size > size_limit: + raise ArtifactError( + f"Artifact path exceeds its {size_limit}-byte publication limit: {relative_name!r}" + ) + + +def _artifact_inventory(root: Path) -> frozenset[str]: + result: set[str] = set() + for path in sorted(root.rglob("*"), key=lambda item: item.as_posix()): + if path.is_symlink(): + raise ArtifactError(f"Symlinks are forbidden in publication artifacts: {path}") + if not path.is_file(): + continue + relative_name = path.relative_to(root).as_posix() + if _sensitive_path(relative_name): + raise ArtifactError( + f"Artifact contains a sensitive publication path: {relative_name!r}" + ) + result.add(relative_name) + return frozenset(result) + + +def _expected_runtime_assets(spec: ModelSpec, registry: ModelRegistry) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for asset in registry.runtime_assets.values(): + if asset.consumer_family != spec.family.id: + continue + material = ( + f"{asset.repository}@{asset.revision}:{asset.path}:{asset.sha256}:{asset.size}" + ).encode() + records.append( + { + "id": asset.id, + "repository": asset.repository, + "revision": asset.revision, + "path": asset.path, + "sha256": asset.sha256, + "size": asset.size, + "license": asset.license_expression, + "consumer_family": asset.consumer_family, + "trust_kind": asset.trust_kind, + "offline_behavior": asset.offline_behavior, + "cache_identity": hashlib.sha256(material).hexdigest(), + } + ) + return records + + +def _assert_current_runtime_source( + spec: ModelSpec, + registry: ModelRegistry, + provenance: Mapping[str, Any], +) -> tuple[str, str]: + """Bind one artifact to the current clean tracked runtime source scope.""" + + source_root = Path(__file__).resolve().parents[2] + runtime_revision, _, source_tree_sha256 = _validated_runtime_snapshot( + source_root, + registry, + spec, + ) + if provenance.get("runtime_revision") != runtime_revision: + raise ArtifactError( + "Artifact runtime revision differs from the current clean source revision; " + "rebuild it." + ) + if provenance.get("source_tree_sha256") != source_tree_sha256: + raise ArtifactError( + "Artifact runtime source-tree digest differs from the current tracked sources; " + "rebuild it." + ) + return runtime_revision, source_tree_sha256 + + +def _assert_current_release_tool_source( + provenance: Mapping[str, Any], +) -> tuple[str, str, dict[str, bytes]]: + """Bind one artifact to the current immutable release-tool scope.""" + + source_root = Path(__file__).resolve().parents[2] + tool_revision, tool_sha256, payloads = _validated_release_tool_snapshot(source_root) + if provenance.get("release_tool_revision") != tool_revision: + raise ArtifactError( + "Artifact release-tool revision differs from the current clean tool scope; " + "rebuild it." + ) + if provenance.get("release_tool_sha256") != tool_sha256: + raise ArtifactError( + "Artifact release-tool digest differs from the current clean tool scope; " + "rebuild it." + ) + return tool_revision, tool_sha256, payloads + + +def _assert_artifact_requirements( + artifact_path: Path, + spec: ModelSpec, + release_tool_payloads: Mapping[str, bytes], +) -> None: + expected = _render_artifact_requirements(spec, release_tool_payloads).encode("utf-8") + path = artifact_path / "requirements.txt" + try: + actual = path.read_bytes() + except OSError as error: + raise ArtifactError(f"Artifact requirements are missing: {path}") from error + if len(actual) > _MAX_RELEASE_TEXT_BYTES: + raise ArtifactError("Artifact requirements exceed the release-text size limit.") + if actual != expected: + raise ArtifactError( + "Artifact requirements differ from the current direct dependency contract." + ) + + +def _revalidate_plan_runtime_source( + *, + model_id: str, + repo_id: str, + runtime_revision: str | None, + source_tree_sha256: str | None, + runtime_bundle_sha256: str | None, + release_tool_revision: str | None, + release_tool_sha256: str | None, + release_revision: str | None, + release_source_sha256: str | None, +) -> None: + """Reject a plan if scoped runtime or release bytes changed after preflight.""" + + if ( + runtime_revision is None + and source_tree_sha256 is None + and runtime_bundle_sha256 is None + and release_tool_revision is None + and release_tool_sha256 is None + and release_revision is None + and release_source_sha256 is None + ): + return + if ( + runtime_revision is None + or source_tree_sha256 is None + or runtime_bundle_sha256 is None + or release_tool_revision is None + or release_tool_sha256 is None + or release_revision is None + or release_source_sha256 is None + ): + raise ArtifactError(f"Publication plan for {repo_id} has partial source identity.") + registry = get_model_registry() + spec = registry.get(model_id) + if spec is None or spec.fast.repo_id != repo_id: + raise ArtifactError(f"Publication plan identity is stale for {repo_id}.") + current_revision, _, current_source_tree_sha256 = _validated_runtime_snapshot( + Path(__file__).resolve().parents[2], + registry, + spec, + ) + if ( + current_revision != runtime_revision + or current_source_tree_sha256 != source_tree_sha256 + ): + raise ArtifactError( + f"Scoped runtime sources changed after publication preflight for {repo_id}." + ) + current_tool_revision, current_tool_sha256, _tool_payloads = ( + _validated_release_tool_snapshot(Path(__file__).resolve().parents[2]) + ) + if ( + current_tool_revision != release_tool_revision + or current_tool_sha256 != release_tool_sha256 + ): + raise ArtifactError( + f"Release tools changed after publication preflight for {repo_id}." + ) + current_release_revision, current_release_source_sha256, _ = ( + _validated_release_text_snapshot( + spec, + registry, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + ) + ) + if ( + current_release_revision != release_revision + or current_release_source_sha256 != release_source_sha256 + ): + raise ArtifactError( + f"Scoped release texts changed after publication preflight for {repo_id}." + ) + + +def _run_required_complete_autoclass_probe( + spec: ModelSpec, + artifact_path: Path, +) -> tuple[str, ...]: + """Validate ANKH encoder and seq2seq views from the same local artifact.""" + + required = tuple( + name for name in _REQUIRED_COMPLETE_AUTOMODEL_VIEWS if name in spec.auto_map + ) + if required != _REQUIRED_COMPLETE_AUTOMODEL_VIEWS: + raise ArtifactError( + f"{spec.id} does not advertise both required complete-publication AutoClasses." + ) + cases = [ + { + "auto_class": auto_class, + "class_path": spec.auto_map[auto_class], + "expected_missing_key_prefixes": [], + "expected_unexpected_key_prefixes": [], + } + for auto_class in required + ] + with tempfile.TemporaryDirectory( + prefix=".fastplms-complete-probe-", + dir=artifact_path.parent, + ) as directory: + probe_root = Path(directory) + cases_path = probe_root / "cases.json" + output_path = probe_root / "result.json" + cases_path.write_text( + json.dumps(cases, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + command = [ + sys.executable, + "-I", + "-S", + str(Path(__file__).with_name("offline_probe.py")), + "--artifact", + str(artifact_path), + "--family", + spec.family.id, + "--bf16-execution", + spec.family.bf16_execution, + "--cases-file", + str(cases_path), + "--implementation", + "artifact", + "--output", + str(output_path), + ] + site_packages = sorted( + { + Path(entry).resolve() + for entry in sys.path + if entry + and Path(entry).name in {"site-packages", "dist-packages"} + and Path(entry).is_dir() + }, + key=lambda path: path.as_posix(), + ) + for path in site_packages: + command.extend(("--runtime-site-package", str(path))) + environment = os.environ.copy() + environment.pop("PYTHONHOME", None) + environment.pop("PYTHONPATH", None) + environment["HF_HOME"] = str(probe_root / "hf-home") + environment["HF_MODULES_CACHE"] = str(probe_root / "modules") + environment["HF_HUB_OFFLINE"] = "1" + environment["TRANSFORMERS_OFFLINE"] = "1" + environment["PYTHONNOUSERSITE"] = "1" + try: + completed = subprocess.run( + command, + cwd=probe_root, + env=environment, + capture_output=True, + text=True, + check=False, + timeout=60 * 60, + ) + except (OSError, subprocess.TimeoutExpired) as error: + raise ArtifactError( + f"Required complete-publication AutoClass probe failed for {spec.id}." + ) from error + if completed.returncode != 0: + details = (completed.stdout + completed.stderr).strip()[-4000:] + raise ArtifactError( + f"Required complete-publication AutoClass probe failed for {spec.id}:\n" + f"{details}" + ) + try: + results = json.loads(output_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactError( + f"Complete-publication AutoClass probe returned invalid output for {spec.id}." + ) from error + if not isinstance(results, dict) or set(results) != set(required) or any( + not isinstance(results[name], dict) for name in required + ): + raise ArtifactError( + f"Complete-publication AutoClass probe omitted a required view for {spec.id}." + ) + return required + + +def _validate_files_only_artifact( + artifact_path: Path, + spec: ModelSpec, +) -> tuple[ + tuple[str, ...], + tuple[tuple[str, bytes], ...], + str, + str, + str, + str, + str, + str, + str, +]: + """Validate only the files eligible for an add-only Hub commit. + + Weight files are identified from both provenance and filename rules, but + they are deliberately neither opened nor hashed. + """ + + registry = _assert_current_registry_spec(spec) + if spec.family.requires_complete_weight_publication: + raise ArtifactError( + f"{spec.id} requires one complete weights-plus-runtime publication; " + "files-only publication is forbidden." + ) + artifact_path = artifact_path.resolve() + if not artifact_path.is_dir(): + raise ArtifactError(f"Built artifact does not exist: {artifact_path}") + expected_name = spec.fast.repo_id.split("/", maxsplit=1)[1] + if artifact_path.name != expected_name: + raise ArtifactError( + f"Artifact directory {artifact_path.name!r} does not match " + f"manifest repository {expected_name!r}." + ) + + manifest = _load_json_object( + artifact_path / "artifact-manifest.json", + "artifact-manifest.json", + ) + provenance = _load_json_object( + artifact_path / "provenance.json", + "provenance.json", + ) + config = _load_json_object(artifact_path / "config.json", "config.json") + if provenance.get("model_id") != spec.id: + raise ArtifactError( + f"Artifact provenance model_id does not match selected model {spec.id!r}." + ) + if config.get("fastplms_model_id") != spec.id: + raise ArtifactError( + f"Artifact config fastplms_model_id does not match selected model {spec.id!r}." + ) + _validate_registry_provenance(provenance, registry, spec) + selected_checkpoint = spec.artifact_checkpoint + expected_config_identity = { + "auto_map": _artifact_auto_map(spec), + "fastplms_model_id": spec.id, + "fastplms_checkpoint_repo_id": selected_checkpoint.repo_id, + "fastplms_checkpoint_revision": selected_checkpoint.revision, + "fastplms_checkpoint_hash": _checkpoint_identity_hash(selected_checkpoint), + "fastplms_weights_revision": selected_checkpoint.revision, + "fastplms_runtime_revision": provenance.get("runtime_revision"), + "fastplms_source_tree_sha256": provenance.get("source_tree_sha256"), + "fastplms_runtime_bundle_sha256": provenance.get("runtime_bundle_sha256"), + "fastplms_release_tool_revision": provenance.get("release_tool_revision"), + "fastplms_release_tool_sha256": provenance.get("release_tool_sha256"), + } + if any(config.get(name) != value for name, value in expected_config_identity.items()): + raise ArtifactError( + "Artifact config packaging identity differs from the current registry or source." + ) + release_revision, release_source_sha256 = _assert_current_release_texts( + artifact_path, + spec, + registry, + ) + + expected_fast_checkpoint = { + "repo_id": spec.fast.repo_id, + "revision": spec.fast.revision, + "files": {item.path: item.encoded for item in spec.fast.files}, + "unresolved_files": list(spec.fast.unresolved_files), + } + if provenance.get("fast_checkpoint") != expected_fast_checkpoint: + raise ArtifactError( + f"Artifact Fast checkpoint provenance differs from the current registry for {spec.id}." + ) + if provenance.get("runtime_assets") != _expected_runtime_assets(spec, registry): + raise ArtifactError( + f"Artifact runtime-asset provenance differs from the current registry for {spec.id}." + ) + + inventory = _artifact_inventory(artifact_path) + expected_inventory = frozenset({"artifact-manifest.json", *manifest}) + if inventory != expected_inventory: + missing = sorted(expected_inventory.difference(inventory)) + extra = sorted(inventory.difference(expected_inventory)) + raise ArtifactError( + "Artifact file inventory differs from artifact-manifest.json; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + + canonical_weights = _canonical_weight_paths(provenance) + declared_assets = _declared_non_weight_assets(spec, registry) + declared_legal_paths = _declared_legal_paths(spec, registry) + selected: list[str] = [] + selected_digests: dict[str, str] = {} + payloads: list[tuple[str, bytes]] = [] + for relative_name, encoded_digest in sorted(manifest.items()): + if not isinstance(relative_name, str) or not isinstance(encoded_digest, str): + raise ArtifactError("artifact-manifest.json keys and values must be strings.") + is_weight = relative_name in canonical_weights or _is_weight_path(relative_name) + path = _resolve_artifact_manifest_path(artifact_path, relative_name) + if path.is_symlink() or not path.is_file(): + raise ArtifactError(f"Files-only artifact path is missing or unsafe: {path}") + if not is_weight and relative_name not in _COMPLETE_ATTESTATION_FILES: + _validate_publishable_non_weight_path( + relative_name, + path, + declared_assets=declared_assets, + declared_legal_paths=declared_legal_paths, + ) + if ( + is_weight + or relative_name in _COMPLETE_ATTESTATION_FILES + or not _is_runtime_update_path(relative_name) + ): + continue + payload = _read_validated_bytes( + path, + encoded_digest, + max_bytes=_MAX_RUNTIME_FILE_BYTES, + ) + selected.append(relative_name) + selected_digests[relative_name] = encoded_digest + payloads.append((relative_name, payload)) + + missing = sorted(_REQUIRED_FILES_ONLY_PATHS.difference(selected)) + if missing: + raise ArtifactError(f"Files-only artifact is missing required paths: {missing}") + if not any(path.startswith("fastplms/") for path in selected): + raise ArtifactError("Files-only artifact contains no packaged FastPLMs runtime sources.") + if any(_is_weight_path(path) for path in selected): + raise ArtifactError("Files-only upload plan unexpectedly contains a weight path.") + + runtime_attestation = _load_json_object( + artifact_path / _RUNTIME_ATTESTATION_NAME, + _RUNTIME_ATTESTATION_NAME, + ) + attested_files = runtime_attestation.get("files") + expected_attested_files = { + name: digest + for name, digest in selected_digests.items() + if name != _RUNTIME_ATTESTATION_NAME + } + expected_runtime_identity = { + "schema_version": _RUNTIME_ATTESTATION_SCHEMA_VERSION, + "scope": "runtime-only", + "model_id": spec.id, + "weights": {"repo_id": spec.fast.repo_id, "revision": spec.fast.revision}, + "runtime_revision": provenance.get("runtime_revision"), + "source_tree_sha256": provenance.get("source_tree_sha256"), + "runtime_bundle_sha256": provenance.get("runtime_bundle_sha256"), + "release_tool_revision": provenance.get("release_tool_revision"), + "release_tool_sha256": provenance.get("release_tool_sha256"), + "weights_license_status": provenance.get("weights_license_status"), + "redistributable": provenance.get("redistributable"), + "files": expected_attested_files, + } + if ( + runtime_attestation != expected_runtime_identity + or attested_files != expected_attested_files + ): + raise ArtifactError( + "Runtime-only attestation differs from the selected files or current weight identity." + ) + if _tree_sha256(artifact_path / "fastplms") != runtime_attestation.get( + "source_tree_sha256" + ): + raise ArtifactError("Runtime source-tree digest differs from packaged sources.") + runtime_bundle_sha256 = runtime_attestation.get("runtime_bundle_sha256") + if not isinstance(runtime_bundle_sha256, str): + raise ArtifactError("Runtime bundle digest is missing from the runtime attestation.") + _validate_runtime_bundle( + artifact_path / "fastplms_bundle.py", + artifact_path / "fastplms", + runtime_bundle_sha256, + ) + _validate_bootstrap( + artifact_path / "modeling_fastplms.py", + spec, + runtime_bundle_sha256, + ) + runtime_revision, source_tree_sha256 = _assert_current_runtime_source( + spec, + registry, + provenance, + ) + ( + release_tool_revision, + release_tool_sha256, + release_tool_payloads, + ) = _assert_current_release_tool_source(provenance) + _assert_artifact_requirements(artifact_path, spec, release_tool_payloads) + return ( + tuple(selected), + tuple(payloads), + runtime_revision, + source_tree_sha256, + runtime_bundle_sha256, + release_tool_revision, + release_tool_sha256, + release_revision, + release_source_sha256, + ) + + +def _attribute(value: object, name: str) -> Any: + if isinstance(value, Mapping): + return value.get(name) + return getattr(value, name, None) + + +def _remote_file_digest(sibling: object, expected: FileDigest) -> str | None: + if expected.algorithm == "sha256": + lfs = _attribute(sibling, "lfs") + digest = _attribute(lfs, "sha256") if lfs is not None else None + return digest if isinstance(digest, str) else None + if expected.algorithm == "git-sha1": + digest = _attribute(sibling, "blob_id") + return digest if isinstance(digest, str) else None + return None + + +def _verify_remote_weights(spec: ModelSpec, model_info: object) -> None: + siblings = _attribute(model_info, "siblings") + if not isinstance(siblings, Iterable): + raise ArtifactError(f"Hub metadata for {spec.fast.repo_id} contains no file listing.") + remote_files = { + path: sibling + for sibling in siblings + if isinstance(path := _attribute(sibling, "rfilename"), str) + } + expected_weights = tuple(item for item in spec.fast.files if _is_weight_path(item.path)) + if not expected_weights: + raise ArtifactError(f"{spec.id} declares no Fast checkpoint weight files.") + for expected in expected_weights: + sibling = remote_files.get(expected.path) + if sibling is None: + raise ArtifactError( + f"Hub repository {spec.fast.repo_id} is missing pinned weight {expected.path}." + ) + actual = _remote_file_digest(sibling, expected) + if actual is None: + raise ArtifactError( + f"Hub did not return {expected.algorithm} metadata for " + f"{spec.fast.repo_id}/{expected.path}." + ) + if actual != expected.digest: + raise ArtifactError( + f"Hub weight identity differs for {spec.fast.repo_id}/{expected.path}: " + f"expected {expected.digest}, received {actual}." + ) + + +def _obsolete_registry_pinned_paths( + spec: ModelSpec, + model_info: object, + new_inventory: Iterable[str], +) -> tuple[str, ...]: + """Verify current weights and select only superseded declared paths.""" + + siblings = _attribute(model_info, "siblings") + if not isinstance(siblings, Iterable): + raise ArtifactError(f"Hub metadata for {spec.fast.repo_id} contains no file listing.") + remote_files = { + path: sibling + for sibling in siblings + if isinstance(path := _attribute(sibling, "rfilename"), str) + } + pinned = { + item.path: item for item in spec.fast.files if _is_weight_path(item.path) + } + replacement_weights = { + relative_name for relative_name in new_inventory if _is_weight_path(relative_name) + } + remote_weights = { + relative_name for relative_name in remote_files if _is_weight_path(relative_name) + } + ambiguous = sorted(remote_weights.difference(pinned, replacement_weights)) + if ambiguous: + raise ArtifactError( + f"Hub repository {spec.fast.repo_id} contains unpinned competing weight files: " + f"{ambiguous}. Resolve their identity before complete replacement." + ) + for relative_name, expected in pinned.items(): + sibling = remote_files.get(relative_name) + if sibling is None: + raise ArtifactError( + f"Hub repository {spec.fast.repo_id} is missing pinned file {relative_name}." + ) + actual = _remote_file_digest(sibling, expected) + if actual is None: + raise ArtifactError( + f"Hub did not return {expected.algorithm} metadata for " + f"{spec.fast.repo_id}/{relative_name}." + ) + if actual != expected.digest: + raise ArtifactError( + f"Hub pinned-file identity differs for {spec.fast.repo_id}/{relative_name}: " + f"expected {expected.digest}, received {actual}." + ) + return tuple(sorted(set(pinned).difference(replacement_weights))) + + +def prepare_files_only_plan( + spec: ModelSpec, + *, + artifact_root: Path, + revision: str, + api: HfApi, +) -> FilesOnlyPublishPlan: + """Validate local non-weight files and the pinned remote weight identity.""" + + repository_name = spec.fast.repo_id.split("/", maxsplit=1)[1] + artifact_path = artifact_root.resolve() / repository_name + ( + files, + payloads, + runtime_revision, + source_tree_sha256, + runtime_bundle_sha256, + release_tool_revision, + release_tool_sha256, + release_revision, + release_source_sha256, + ) = _validate_files_only_artifact(artifact_path, spec) + info = api.model_info( + spec.fast.repo_id, + revision=revision, + files_metadata=True, + ) + parent_commit = _attribute(info, "sha") + if not isinstance(parent_commit, str) or not parent_commit: + raise ArtifactError(f"Hub repository {spec.fast.repo_id} has no commit identity.") + _verify_remote_weights(spec, info) + return FilesOnlyPublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision=revision, + parent_commit=parent_commit, + artifact_path=artifact_path, + files=files, + payloads=payloads, + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + release_revision=release_revision, + release_source_sha256=release_source_sha256, + ) + + +def prepare_files_only_plans( + specs: Iterable[ModelSpec], + *, + artifact_root: Path, + revision: str, + api: HfApi, +) -> tuple[FilesOnlyPublishPlan, ...]: + """Preflight every repository before the first remote mutation.""" + + selected = tuple(specs) + blocked = [ + spec.id for spec in selected if spec.family.requires_complete_weight_publication + ] + if blocked: + raise ArtifactError( + "Files-only publication is forbidden for models requiring complete weights: " + + ", ".join(blocked) + ) + return tuple( + prepare_files_only_plan( + spec, + artifact_root=artifact_root, + revision=revision, + api=api, + ) + for spec in selected + ) + + +def prepare_complete_plan( + spec: ModelSpec, + *, + artifact_root: Path, + revision: str, + api: HfApi, +) -> CompletePublishPlan: + """Preflight one complete artifact for a single atomic Hub commit.""" + + registry = _assert_current_registry_spec(spec) + if not spec.family.weights_publication_allowed: + raise ArtifactError( + f"{spec.id} checkpoint publication is blocked by unresolved weight-license terms." + ) + repository_name = spec.fast.repo_id.split("/", maxsplit=1)[1] + artifact_path = artifact_root.resolve() / repository_name + validate_artifact(artifact_path, spec=spec, registry=registry) + release_revision, release_source_sha256 = _assert_current_release_texts( + artifact_path, + spec, + registry, + ) + manifest = _load_json_object( + artifact_path / "artifact-manifest.json", + "artifact-manifest.json", + ) + inventory = _artifact_inventory(artifact_path) + if inventory != frozenset({"artifact-manifest.json", *manifest}): + raise ArtifactError( + "Complete artifact inventory differs from artifact-manifest.json." + ) + provenance = _load_json_object(artifact_path / "provenance.json", "provenance.json") + if ( + provenance.get("weights_license_status") != "resolved" + or provenance.get("redistributable") is not True + ): + raise ArtifactError( + f"Complete publication is forbidden for non-redistributable artifact {spec.id}." + ) + for label, checkpoint in (("fast", spec.fast), ("official", spec.official)): + expected_checkpoint = { + "repo_id": checkpoint.repo_id, + "revision": checkpoint.revision, + "files": {item.path: item.encoded for item in checkpoint.files}, + "unresolved_files": list(checkpoint.unresolved_files), + } + if provenance.get(f"{label}_checkpoint") != expected_checkpoint: + raise ArtifactError( + f"Complete artifact {label} checkpoint differs from the current registry." + ) + if provenance.get("runtime_assets") != _expected_runtime_assets(spec, registry): + raise ArtifactError( + "Complete artifact runtime-asset provenance differs from the current registry." + ) + runtime_revision, source_tree_sha256 = _assert_current_runtime_source( + spec, + registry, + provenance, + ) + ( + release_tool_revision, + release_tool_sha256, + release_tool_payloads, + ) = _assert_current_release_tool_source(provenance) + _assert_artifact_requirements(artifact_path, spec, release_tool_payloads) + runtime_bundle_sha256 = provenance.get("runtime_bundle_sha256") + if not isinstance(runtime_bundle_sha256, str): + raise ArtifactError("Complete artifact lacks a runtime-bundle identity.") + canonical_weights = _canonical_weight_paths(provenance) + declared_assets = _declared_non_weight_assets(spec, registry) + declared_legal_paths = _declared_legal_paths(spec, registry) + digests: list[tuple[str, str]] = [] + for relative_name in sorted(inventory): + path = _resolve_artifact_manifest_path(artifact_path, relative_name) + if relative_name not in {"artifact-manifest.json", "provenance.json"} and ( + relative_name not in canonical_weights and not _is_weight_path(relative_name) + ): + _validate_publishable_non_weight_path( + relative_name, + path, + declared_assets=declared_assets, + declared_legal_paths=declared_legal_paths, + ) + if relative_name == "artifact-manifest.json": + digest = f"sha256:{hash_file(path)}" + else: + encoded = manifest.get(relative_name) + if not isinstance(encoded, str): + raise ArtifactError( + f"Complete artifact manifest is missing {relative_name!r}." + ) + _encoded_digest(path, encoded) + digest = encoded + digests.append((relative_name, digest)) + required = { + "README.md", + "config.json", + "provenance.json", + "artifact-manifest.json", + _RUNTIME_ATTESTATION_NAME, + *canonical_weights, + } + missing = sorted(required.difference(inventory)) + if missing: + raise ArtifactError(f"Complete artifact is missing required atomic paths: {missing}") + validated_auto_classes: tuple[str, ...] = () + validation_manifest_sha256: str | None = None + if spec.family.requires_complete_weight_publication: + validated_auto_classes = _run_required_complete_autoclass_probe( + spec, + artifact_path, + ) + validation_manifest_sha256 = hash_file( + artifact_path / "artifact-manifest.json" + ) + info = api.model_info(spec.fast.repo_id, revision=revision, files_metadata=True) + parent_commit = _attribute(info, "sha") + if not isinstance(parent_commit, str) or not parent_commit: + raise ArtifactError(f"Hub repository {spec.fast.repo_id} has no commit identity.") + files = tuple(name for name, _ in digests) + deletes = _obsolete_registry_pinned_paths(spec, info, files) + return CompletePublishPlan( + model_id=spec.id, + repo_id=spec.fast.repo_id, + revision=revision, + parent_commit=parent_commit, + artifact_path=artifact_path, + files=files, + digests=tuple(digests), + deletes=deletes, + replacement_weight_paths=tuple(sorted(canonical_weights)), + runtime_revision=runtime_revision, + source_tree_sha256=source_tree_sha256, + runtime_bundle_sha256=runtime_bundle_sha256, + release_tool_revision=release_tool_revision, + release_tool_sha256=release_tool_sha256, + release_revision=release_revision, + release_source_sha256=release_source_sha256, + validation_manifest_sha256=validation_manifest_sha256, + validated_auto_classes=validated_auto_classes, + ) + + +def prepare_complete_plans( + specs: Iterable[ModelSpec], + *, + artifact_root: Path, + revision: str, + api: HfApi, +) -> tuple[CompletePublishPlan, ...]: + """Preflight all complete commits before the first remote mutation.""" + + selected = tuple(specs) + blocked = [spec.id for spec in selected if not spec.family.weights_publication_allowed] + if blocked: + raise ArtifactError( + "Complete checkpoint publication is blocked by unresolved weight licenses: " + + ", ".join(blocked) + ) + return tuple( + prepare_complete_plan( + spec, + artifact_root=artifact_root, + revision=revision, + api=api, + ) + for spec in selected + ) + + +def publish_files_only( + plans: Iterable[FilesOnlyPublishPlan], + *, + api: HfApi, + commit_message: str, + dry_run: bool = False, +) -> tuple[FilesOnlyPublishResult, ...]: + """Execute preflighted add-only commits, or print their dry-run plans.""" + + results: list[FilesOnlyPublishResult] = [] + for plan in plans: + _revalidate_plan_runtime_source( + model_id=plan.model_id, + repo_id=plan.repo_id, + runtime_revision=plan.runtime_revision, + source_tree_sha256=plan.source_tree_sha256, + runtime_bundle_sha256=plan.runtime_bundle_sha256, + release_tool_revision=plan.release_tool_revision, + release_tool_sha256=plan.release_tool_sha256, + release_revision=plan.release_revision, + release_source_sha256=plan.release_source_sha256, + ) + print( + f"{'[dry-run] ' if dry_run else ''}{plan.repo_id}: " + f"{len(plan.files)} non-weight files at {plan.parent_commit}" + ) + for relative_name in plan.files: + print(f" {relative_name}") + if dry_run: + continue + payload_by_name = dict(plan.payloads) + if set(payload_by_name) != set(plan.files): + raise ArtifactError( + f"Preflighted payload inventory differs for {plan.repo_id}; rebuild the plan." + ) + operations = [ + CommitOperationAdd( + path_in_repo=relative_name, + path_or_fileobj=io.BytesIO(payload_by_name[relative_name]), + ) + for relative_name in plan.files + ] + commit = api.create_commit( + repo_id=plan.repo_id, + repo_type="model", + revision=plan.revision, + parent_commit=plan.parent_commit, + operations=operations, + commit_message=commit_message, + commit_description=( + "Add-only FastPLMs files-only publication. " + "Checkpoint weights and complete-artifact attestations are unchanged." + ), + ) + oid = _attribute(commit, "oid") + url = _attribute(commit, "commit_url") + if not isinstance(oid, str) or not isinstance(url, str): + raise ArtifactError(f"Hub returned incomplete commit metadata for {plan.repo_id}.") + results.append( + FilesOnlyPublishResult( + model_id=plan.model_id, + repo_id=plan.repo_id, + commit_oid=oid, + commit_url=url, + ) + ) + return tuple(results) + + +def _revalidate_complete_plan( + plan: CompletePublishPlan, + api: HfApi, +) -> ModelSpec: + """Rebuild one complete preflight immediately before any remote mutation.""" + + registry = get_model_registry() + selected_spec = registry.get(plan.model_id) + if selected_spec is None: + raise ArtifactError( + f"Complete publication plan model {plan.model_id!r} is absent from the " + "current registry." + ) + if plan.repo_id != selected_spec.fast.repo_id: + raise ArtifactError( + f"Complete plan repository differs from the registry for {plan.model_id}." + ) + current_plan = prepare_complete_plan( + selected_spec, + artifact_root=plan.artifact_path.resolve().parent, + revision=plan.revision, + api=api, + ) + if current_plan != plan: + raise ArtifactError( + f"Complete plan for {plan.repo_id} differs from a current full preflight; " + "rebuild the plan." + ) + return selected_spec + + +def publish_complete( + plans: Iterable[CompletePublishPlan], + *, + api: HfApi, + commit_message: str, + dry_run: bool = False, +) -> tuple[FilesOnlyPublishResult, ...]: + """Publish each complete artifact in one parent-guarded atomic commit.""" + + results: list[FilesOnlyPublishResult] = [] + for plan in plans: + selected_spec = _revalidate_complete_plan(plan, api) + if selected_spec is not None: + if not selected_spec.family.weights_publication_allowed: + raise ArtifactError( + f"{plan.model_id} checkpoint publication is blocked by unresolved " + "weight-license terms." + ) + replacement_weights = set(plan.replacement_weight_paths) + if ( + replacement_weights != { + path for path in plan.files if _is_weight_path(path) + } + or _WEIGHT_INDEX not in replacement_weights + or not any( + path != _WEIGHT_INDEX and path.endswith(".safetensors") + for path in replacement_weights + ) + ): + raise ArtifactError( + f"Complete plan for {plan.repo_id} lacks a canonical replacement " + "weight set." + ) + if ( + selected_spec is not None + and selected_spec.family.requires_complete_weight_publication + ): + manifest_path = plan.artifact_path / "artifact-manifest.json" + if ( + plan.repo_id != selected_spec.fast.repo_id + or plan.validated_auto_classes != _REQUIRED_COMPLETE_AUTOMODEL_VIEWS + or plan.validation_manifest_sha256 is None + or not manifest_path.is_file() + or hash_file(manifest_path) != plan.validation_manifest_sha256 + ): + raise ArtifactError( + f"Complete plan for {plan.repo_id} lacks a current required AutoClass probe." + ) + if len(set(plan.deletes)) != len(plan.deletes): + raise ArtifactError(f"Complete delete paths are repeated for {plan.repo_id}.") + if set(plan.deletes).intersection(plan.files): + raise ArtifactError( + f"Complete add and delete paths overlap for {plan.repo_id}." + ) + if plan.deletes: + if selected_spec is None: + raise ArtifactError( + f"Complete plan for {plan.repo_id} cannot authorize deletes without " + "a current registry model." + ) + current_info = api.model_info( + plan.repo_id, + revision=plan.revision, + files_metadata=True, + ) + if _attribute(current_info, "sha") != plan.parent_commit: + raise ArtifactError( + f"Remote parent changed after complete preflight for {plan.repo_id}." + ) + permitted_deletes = _obsolete_registry_pinned_paths( + selected_spec, + current_info, + plan.files, + ) + if plan.deletes != permitted_deletes: + raise ArtifactError( + f"Complete plan for {plan.repo_id} contains an unproven obsolete " + "weight delete." + ) + print( + f"{'[dry-run] ' if dry_run else ''}{plan.repo_id}: " + f"{len(plan.files)} complete files, {len(plan.deletes)} guarded deletes " + f"at {plan.parent_commit}" + ) + for relative_name in plan.deletes: + print(f" delete {relative_name}") + if dry_run: + continue + expected = dict(plan.digests) + if set(expected) != set(plan.files): + raise ArtifactError( + f"Preflighted complete inventory differs for {plan.repo_id}; rebuild the plan." + ) + # Bounded files are retained as exact bytes. Large shards are copied to + # a publisher-owned temporary snapshot, then rehashed and held open + # through the synchronous Hub commit. + with contextlib.ExitStack() as open_payloads: + snapshot_root = Path( + open_payloads.enter_context( + tempfile.TemporaryDirectory( + prefix=".fastplms-complete-publish-", + dir=plan.artifact_path.parent, + ) + ) + ) + operations: list[CommitOperationAdd | CommitOperationDelete] = [] + retained_handles: list[ + tuple[BinaryIO, tuple[int, int, int, int], Path] + ] = [] + for index, relative_name in enumerate(plan.files): + path = _resolve_artifact_manifest_path(plan.artifact_path, relative_name) + if path.is_symlink() or not path.is_file(): + raise ArtifactError( + f"Preflighted complete artifact path changed: {relative_name!r}" + ) + if path.stat().st_size <= _MAX_RETAINED_COMPLETE_BYTES: + payload: BinaryIO = io.BytesIO( + _read_validated_bytes( + path, + expected[relative_name], + max_bytes=_MAX_RETAINED_COMPLETE_BYTES, + ) + ) + else: + snapshot = snapshot_root / f"{index:05d}.payload" + _snapshot_validated_payload( + path, + snapshot, + expected[relative_name], + ) + payload = open_payloads.enter_context(snapshot.open("rb")) + identity = _open_validated_payload( + payload, + snapshot, + expected[relative_name], + ) + retained_handles.append((payload, identity, snapshot)) + operations.append( + CommitOperationAdd( + path_in_repo=relative_name, + path_or_fileobj=payload, + ) + ) + operations.extend( + CommitOperationDelete(path_in_repo=relative_name) + for relative_name in plan.deletes + ) + for handle, expected_identity, path in retained_handles: + current = os.fstat(handle.fileno()) + if expected_identity != ( + current.st_dev, + current.st_ino, + current.st_size, + current.st_mtime_ns, + ): + raise ArtifactError( + f"Preflighted complete artifact changed before upload: {path}" + ) + handle.seek(0) + commit = api.create_commit( + repo_id=plan.repo_id, + repo_type="model", + revision=plan.revision, + parent_commit=plan.parent_commit, + operations=operations, + commit_message=commit_message, + commit_description=( + "Atomic FastPLMs complete publication. Checkpoint weights, tokenizer " + "assets, runtime sources, model card, legal texts, and scoped " + "attestations are updated together. Deletes are restricted to obsolete " + "current-registry-pinned paths." + ), + ) + oid = _attribute(commit, "oid") + url = _attribute(commit, "commit_url") + if not isinstance(oid, str) or not isinstance(url, str): + raise ArtifactError(f"Hub returned incomplete commit metadata for {plan.repo_id}.") + results.append( + FilesOnlyPublishResult( + model_id=plan.model_id, + repo_id=plan.repo_id, + commit_oid=oid, + commit_url=url, + ) + ) + return tuple(results) + + +def _selected_specs( + registry: ModelRegistry, + model_ids: Iterable[str], + *, + all_models: bool, +) -> tuple[ModelSpec, ...]: + selected = tuple(model_ids) + if all_models and selected: + raise ArtifactError("Pass model IDs or --all, not both.") + if all_models or not selected: + selected = tuple(registry) + unknown = sorted(set(selected).difference(registry)) + if unknown: + raise ArtifactError(f"Unknown model IDs: {unknown}") + if len(set(selected)) != len(selected): + raise ArtifactError("Model IDs must not be repeated.") + return tuple(registry[model_id] for model_id in selected) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model_ids", nargs="*", help="Stable IDs from src/fastplms/models.toml") + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--files-only", + action="store_true", + help="Required safety mode: upload no checkpoint weights and delete no remote files", + ) + mode.add_argument( + "--complete", + action="store_true", + help="Atomically upload a complete validated artifact, including checkpoint weights", + ) + parser.add_argument( + "--all", + action="store_true", + help="Explicit alias for the default behavior of publishing every manifest model", + ) + parser.add_argument("--artifact-root", type=Path, default=Path("dist/hub")) + parser.add_argument("--revision", default="main") + parser.add_argument( + "--commit-message", + default="Update FastPLMs runtime files", + ) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> None: + arguments = _parse_args() + if not arguments.files_only and not arguments.complete: + raise SystemExit( + "Refusing to publish without an explicit --files-only or --complete mode." + ) + if arguments.complete and (arguments.all or not arguments.model_ids): + raise SystemExit("Complete publication requires explicit model IDs and forbids --all.") + registry = get_model_registry() + try: + specs = _selected_specs( + registry, + arguments.model_ids, + all_models=arguments.all, + ) + api = HfApi() + if arguments.files_only: + files_only_plans = prepare_files_only_plans( + specs, + artifact_root=arguments.artifact_root, + revision=arguments.revision, + api=api, + ) + results = publish_files_only( + files_only_plans, + api=api, + commit_message=arguments.commit_message, + dry_run=arguments.dry_run, + ) + else: + complete_plans = prepare_complete_plans( + specs, + artifact_root=arguments.artifact_root, + revision=arguments.revision, + api=api, + ) + results = publish_complete( + complete_plans, + api=api, + commit_message=arguments.commit_message, + dry_run=arguments.dry_run, + ) + except ArtifactError as error: + raise SystemExit(str(error)) from error + for result in results: + print(f"{result.repo_id}: {result.commit_oid} {result.commit_url}") + + +if __name__ == "__main__": + main() + + +__all__ = [ + "CompletePublishPlan", + "FilesOnlyPublishPlan", + "FilesOnlyPublishResult", + "prepare_complete_plan", + "prepare_complete_plans", + "prepare_files_only_plan", + "prepare_files_only_plans", + "publish_complete", + "publish_files_only", +] diff --git a/tools/artifacts/resolve_fair_esm_assets.py b/tools/artifacts/resolve_fair_esm_assets.py new file mode 100644 index 0000000..5e0e69d --- /dev/null +++ b/tools/artifacts/resolve_fair_esm_assets.py @@ -0,0 +1,148 @@ +"""Download and hash the native fair-esm parity-oracle assets. + +This resolver reads only Meta's public fair-esm asset host. Files are written +transactionally to a caller-selected cache and the resulting URL, relative +path, byte size, and SHA-256 identity are emitted as JSON for manifest review. +It never uploads content or authenticates to a remote service. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from pathlib import Path, PurePosixPath +from urllib.parse import urlparse + + +_HOST = "dl.fbaipublicfiles.com" +_ROOT = f"https://{_HOST}/fair-esm" +_MODEL_NAMES = ( + "esm2_t6_8M_UR50D", + "esm2_t12_35M_UR50D", + "esm2_t30_150M_UR50D", + "esm2_t33_650M_UR50D", + "esm2_t36_3B_UR50D", +) + + +@dataclass(frozen=True) +class ResolvedAsset: + model: str + role: str + path: str + url: str + sha256: str + size: int + + +def _candidates() -> tuple[tuple[str, str, str, str], ...]: + assets: list[tuple[str, str, str, str]] = [] + for model in _MODEL_NAMES: + assets.extend( + ( + (model, "weights", f"models/{model}.pt", f"{_ROOT}/models/{model}.pt"), + ( + model, + "contact_regression", + f"regression/{model}-contact-regression.pt", + f"{_ROOT}/regression/{model}-contact-regression.pt", + ), + ) + ) + assets.append( + ( + "esmfold_3B_v1", + "weights", + "models/esmfold_3B_v1.pt", + f"{_ROOT}/models/esmfold_3B_v1.pt", + ) + ) + return tuple(assets) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(8 * 1024**2): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_url(url: str) -> None: + parsed = urlparse(url) + if parsed.scheme != "https" or parsed.hostname != _HOST: + raise RuntimeError(f"Refusing non-fair-esm asset URL: {url}") + + +def _download_one( + candidate: tuple[str, str, str, str], + cache: Path, +) -> ResolvedAsset: + model, role, relative_name, url = candidate + _validate_url(url) + relative = PurePosixPath(relative_name) + target = cache.joinpath(*relative.parts) + target.parent.mkdir(parents=True, exist_ok=True) + partial = target.with_name(f"{target.name}.part") + if not target.is_file(): + offset = partial.stat().st_size if partial.is_file() else 0 + request = urllib.request.Request( + url, + headers={"Range": f"bytes={offset}-"} if offset else {}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + _validate_url(response.geturl()) + status = getattr(response, "status", 200) + append = offset > 0 and status == 206 + if offset > 0 and not append: + offset = 0 + mode = "ab" if append else "wb" + with partial.open(mode) as handle: + while chunk := response.read(8 * 1024**2): + handle.write(chunk) + handle.flush() + os.fsync(handle.fileno()) + partial.replace(target) + return ResolvedAsset( + model=model, + role=role, + path=relative.as_posix(), + url=url, + sha256=_sha256(target), + size=target.stat().st_size, + ) + + +def resolve(cache: Path, jobs: int) -> list[ResolvedAsset]: + """Resolve all supported ESM2 native oracle assets.""" + + cache = cache.resolve() + cache.mkdir(parents=True, exist_ok=True) + with ThreadPoolExecutor(max_workers=jobs) as pool: + assets = list(pool.map(lambda item: _download_one(item, cache), _candidates())) + return sorted(assets, key=lambda item: (item.model, item.role)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cache", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--jobs", type=int, default=3) + args = parser.parse_args() + if args.jobs < 1: + parser.error("--jobs must be at least one") + document = [asdict(item) for item in resolve(args.cache, args.jobs)] + encoded = json.dumps(document, indent=2, sort_keys=True) + "\n" + if args.output is None: + print(encoded, end="") + else: + args.output.write_text(encoded, encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/tools/artifacts/resolve_manifest_hashes.py b/tools/artifacts/resolve_manifest_hashes.py new file mode 100644 index 0000000..3aeb48e --- /dev/null +++ b/tools/artifacts/resolve_manifest_hashes.py @@ -0,0 +1,105 @@ +"""Resolve manifest file identities at immutable Hugging Face revisions. + +This tool is intentionally read-only. It prints a JSON mapping that can be +reviewed before ``models.toml`` is changed. Small Git-managed files are +downloaded and hashed as Git blobs. LFS files use the SHA-256 identity exposed +by Hub metadata, so multi-gigabyte weight shards are never downloaded merely +to resolve provenance. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import tomllib +from pathlib import Path +from typing import Any +from huggingface_hub import HfApi, hf_hub_download + + +def _git_blob_sha1(payload: bytes) -> str: + header = f"blob {len(payload)}\0".encode() + return hashlib.sha1(header + payload, usedforsecurity=False).hexdigest() + + +def _lfs_sha256(sibling: Any) -> str | None: + lfs = sibling.lfs + if lfs is None: + return None + value = lfs.get("sha256") if isinstance(lfs, dict) else getattr(lfs, "sha256", None) + return str(value) if value else None + + +def resolve_manifest(manifest: Path) -> dict[str, dict[str, str]]: + """Return identities for every unresolved manifest path.""" + with manifest.open("rb") as handle: + document = tomllib.load(handle) + + api = HfApi() + resolved: dict[str, dict[str, str]] = {} + failures: list[str] = [] + for model in document["models"]: + model_id = model["id"] + model_result: dict[str, str] = {} + for label in ("fast", "official"): + paths = model.get(f"{label}_unresolved_files", ()) + if not paths: + continue + repo_id = model[f"{label}_repo"] + revision = model[f"{label}_revision"] + try: + info = api.model_info( + repo_id=repo_id, + revision=revision, + files_metadata=True, + ) + except Exception as error: # pragma: no cover - network diagnostic + failures.append(f"{repo_id}@{revision}: {type(error).__name__}: {error}") + continue + siblings = {sibling.rfilename: sibling for sibling in info.siblings} + for path in paths: + key = f"{label}:{path}" + sibling = siblings.get(path) + if sibling is None: + failures.append(f"{repo_id}@{revision}:{path}: missing from Hub metadata") + continue + sha256 = _lfs_sha256(sibling) + if sha256 is not None: + model_result[key] = f"sha256:{sha256}" + continue + try: + downloaded = hf_hub_download( + repo_id=repo_id, + filename=path, + revision=revision, + ) + payload = Path(downloaded).read_bytes() + except Exception as error: # pragma: no cover - network diagnostic + failures.append(f"{repo_id}@{revision}:{path}: {type(error).__name__}: {error}") + continue + model_result[key] = f"git-sha1:{_git_blob_sha1(payload)}" + if model_result: + resolved[model_id] = model_result + + if failures: + detail = "\n".join(failures) + raise RuntimeError(f"Manifest identities could not be resolved:\n{detail}") + return resolved + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("manifest", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + result = resolve_manifest(args.manifest) + encoded = json.dumps(result, indent=2, sort_keys=True) + "\n" + if args.output is None: + print(encoded, end="") + else: + args.output.write_text(encoded, encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/tools/conversion/__init__.py b/tools/conversion/__init__.py new file mode 100644 index 0000000..c276ed3 --- /dev/null +++ b/tools/conversion/__init__.py @@ -0,0 +1,18 @@ +"""Local-only deterministic checkpoint state conversion utilities. + +Build complete Hub-format artifacts with :mod:`tools.artifacts.build`. This +package contains pure tensor transforms and exact validators only. +""" + +from tools.conversion.state_transforms import ( + StateTransformError, + apply_state_transform, + available_state_transforms, +) + + +__all__ = [ + "StateTransformError", + "apply_state_transform", + "available_state_transforms", +] diff --git a/tools/conversion/extract_esmfold2_geometry.py b/tools/conversion/extract_esmfold2_geometry.py new file mode 100644 index 0000000..a0b3bd0 --- /dev/null +++ b/tools/conversion/extract_esmfold2_geometry.py @@ -0,0 +1,87 @@ +"""Extract the pinned ESMFold2 protein geometry table as declarative JSON.""" + +from __future__ import annotations + +import argparse +import ast +import json +import math +from pathlib import Path + + +def _validate_geometry(residues: object) -> dict[str, dict[str, tuple[float, float, float]]]: + if not isinstance(residues, dict) or not residues: + raise ValueError("PROTEIN_REF_POS must be a non-empty dictionary.") + + validated: dict[str, dict[str, tuple[float, float, float]]] = {} + for residue, atoms in residues.items(): + if not isinstance(residue, str) or not residue: + raise ValueError("PROTEIN_REF_POS residue names must be non-empty strings.") + if not isinstance(atoms, dict) or not atoms: + raise ValueError(f"PROTEIN_REF_POS[{residue!r}] must be a non-empty dictionary.") + validated_atoms: dict[str, tuple[float, float, float]] = {} + for atom, coordinates in atoms.items(): + if not isinstance(atom, str) or not atom: + raise ValueError(f"PROTEIN_REF_POS[{residue!r}] has an invalid atom name.") + if not isinstance(coordinates, (tuple, list)) or len(coordinates) != 3: + raise ValueError( + f"PROTEIN_REF_POS[{residue!r}][{atom!r}] must contain three coordinates." + ) + values: list[float] = [] + for coordinate in coordinates: + if isinstance(coordinate, bool) or not isinstance(coordinate, (int, float)): + raise ValueError( + f"PROTEIN_REF_POS[{residue!r}][{atom!r}] contains a non-numeric coordinate." + ) + value = float(coordinate) + if not math.isfinite(value): + raise ValueError( + f"PROTEIN_REF_POS[{residue!r}][{atom!r}] contains a non-finite coordinate." + ) + values.append(value) + validated_atoms[atom] = (values[0], values[1], values[2]) + validated[residue] = validated_atoms + return validated + + +def extract_geometry(source: Path) -> dict[str, object]: + """Read ``PROTEIN_REF_POS`` without importing or executing upstream code.""" + + module = ast.parse(source.read_text(encoding="utf-8"), filename=str(source)) + for statement in module.body: + if ( + isinstance(statement, ast.AnnAssign) + and isinstance(statement.target, ast.Name) + and statement.target.id == "PROTEIN_REF_POS" + and statement.value is not None + ): + try: + residues = ast.literal_eval(statement.value) + except (SyntaxError, TypeError, ValueError) as error: + raise ValueError("PROTEIN_REF_POS must be a Python literal.") from error + return { + "schema": "fastplms.esmfold2.reference_geometry.v1", + "provenance": { + "manifest_family": "esmfold2", + "contract": "biohub_esmfold2_input_v1", + }, + "dtype": "float32", + "residues": _validate_geometry(residues), + } + raise ValueError(f"PROTEIN_REF_POS was not found in {source}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("source", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + payload = extract_geometry(args.source) + args.output.write_text( + json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n", + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/tools/conversion/state_transforms.py b/tools/conversion/state_transforms.py new file mode 100644 index 0000000..220bda8 --- /dev/null +++ b/tools/conversion/state_transforms.py @@ -0,0 +1,404 @@ +"""Pure, deterministic checkpoint state transforms declared by ``models.toml``. + +The functions in this module operate only on in-memory tensor mappings. They +cannot download a checkpoint, authenticate to a service, or mutate a Hub +repository. Artifact assembly and safetensors sharding remain centralized in +``tools.artifacts.build``. +""" + +from __future__ import annotations + +import re +import torch +from collections.abc import Callable, Iterable, Mapping + + +StateDict = dict[str, torch.Tensor] +Transform = Callable[[Mapping[str, torch.Tensor], frozenset[str] | None], StateDict] + +# Checkpoint entries are heterogeneous tensors. ``(...)`` denotes the +# arbitrary rank retained by every value unless a transform changes only dtype. + + +class StateTransformError(RuntimeError): + """Raised when a declared state transform cannot be applied exactly.""" + + +def _clone_tensor(key: str, value: object) -> torch.Tensor: + if not torch.is_tensor(value): + raise StateTransformError(f"State entry {key!r} is not a tensor.") + return value.detach().cpu().clone() # (...) + + +def _validate_expected(state: StateDict, expected_keys: frozenset[str] | None) -> None: + if expected_keys is None: + return + actual = frozenset(state) + missing = sorted(expected_keys - actual) + unexpected = sorted(actual - expected_keys) + if missing or unexpected: + raise StateTransformError( + "Transformed state keys do not match the expected model state. " + f"Missing: {missing[:20]}; unexpected: {unexpected[:20]}." + ) + + +def _map_state( + state: Mapping[str, torch.Tensor], + key_mapper: Callable[[str], str | None], + expected_keys: frozenset[str] | None, +) -> StateDict: + # state[key]: (...) + if not state: + raise StateTransformError("A checkpoint state dictionary cannot be empty.") + transformed: StateDict = {} + for key in sorted(state): + if not isinstance(key, str) or not key: + raise StateTransformError(f"Invalid state key: {key!r}.") + mapped = key_mapper(key) + if mapped is None: + continue + if mapped in transformed: + raise StateTransformError(f"State-key collision while mapping {key!r} to {mapped!r}.") + transformed[mapped] = _clone_tensor(key, state[key]) # (...) + _validate_expected(transformed, expected_keys) + return transformed + + +def _identity( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + return _map_state(state, lambda key: key, expected_keys) + + +def _cast_floating( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, + dtype: torch.dtype, +) -> StateDict: + # state[key]: (...) + transformed = _identity(state, expected_keys) # transformed[key]: (...) + return { + key: value.to(dtype=dtype) if value.is_floating_point() else value # (...) + for key, value in transformed.items() + } + + +def _drop_unused_rotary_position_table( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + return _map_state( + state, + lambda key: None if key == "esm.embeddings.position_embeddings.weight" else key, + expected_keys, + ) + + +def _esm2( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + """Map pinned fair-esm ESM2 names to the canonical FastPLMs schema.""" + + # state[key]: (...) + if not state: + raise StateTransformError("A checkpoint state dictionary cannot be empty.") + keys = frozenset(state) + official_schema = "embed_tokens.weight" in keys or any( + key.startswith(("layers.", "emb_layer_norm_after.", "contact_head.")) for key in keys + ) + canonical_schema = any(key.startswith("esm.") for key in keys) or any( + key.startswith("lm_head.decoder.") for key in keys + ) + if official_schema and canonical_schema: + raise StateTransformError("ESM2 checkpoint mixes official and canonical parameter schemas.") + if canonical_schema: + return _identity(state, expected_keys) + + transformed: StateDict = {} + + def store(source_key: str, target_key: str) -> None: + if target_key in transformed: + raise StateTransformError( + f"State-key collision while mapping {source_key!r} to {target_key!r}." + ) + transformed[target_key] = _clone_tensor(source_key, state[source_key]) # (...) + + projection_names = {"q_proj": "query", "k_proj": "key", "v_proj": "value"} + for key in sorted(state): + target: str | None = None + if key == "embed_tokens.weight": + target = "esm.embeddings.word_embeddings.weight" + elif key.startswith("layers."): + match = re.fullmatch(r"layers\.(\d+)\.(.+)", key) + if match is None: + raise StateTransformError(f"Unrecognized official ESM2 layer key: {key!r}.") + layer, suffix = match.groups() + prefix = f"esm.encoder.layer.{layer}." + if suffix == "self_attn.rot_emb.inv_freq": + target = f"{prefix}attention.self.rotary_embeddings.inv_freq" + for source_name, target_name in projection_names.items(): + if suffix.startswith(f"self_attn.{source_name}."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.self.{target_name}.{parameter}" + break + if suffix.startswith("self_attn.out_proj."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.output.dense.{parameter}" + elif suffix.startswith("self_attn_layer_norm."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}attention.LayerNorm.{parameter}" + elif suffix.startswith("fc1."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}intermediate.dense.{parameter}" + elif suffix.startswith("fc2."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}output.dense.{parameter}" + elif suffix.startswith("final_layer_norm."): + parameter = suffix.rsplit(".", 1)[-1] + target = f"{prefix}LayerNorm.{parameter}" + elif key.startswith("emb_layer_norm_after."): + target = f"esm.encoder.{key}" + elif key.startswith("contact_head."): + target = f"esm.{key}" + elif key == "lm_head.weight": + target = "lm_head.decoder.weight" + elif key == "lm_head.bias": + store(key, "lm_head.bias") + continue + elif key.startswith("lm_head."): + target = key + + if target is None: + raise StateTransformError(f"Unrecognized official ESM2 state key: {key!r}.") + store(key, target) + + _validate_expected(transformed, expected_keys) + return transformed + + +def _esmc_key(key: str) -> str | None: + if key.endswith("._extra_state"): + return None + if key.startswith("esmc."): + key = key[len("esmc.") :] + if key.startswith("lm_head."): + key = f"sequence_head.{key[len('lm_head.') :]}" + replacements = ( + (".attn.layernorm_qkv.layer_norm_bias", ".attn.layernorm_qkv.0.bias"), + (".attn.layernorm_qkv.layer_norm_weight", ".attn.layernorm_qkv.0.weight"), + (".attn.layernorm_qkv.weight", ".attn.layernorm_qkv.1.weight"), + (".ffn.layer_norm_bias", ".ffn.0.bias"), + (".ffn.layer_norm_weight", ".ffn.0.weight"), + (".ffn.fc1_weight", ".ffn.1.weight"), + (".ffn.fc2_weight", ".ffn.3.weight"), + ) + for old, new in replacements: + key = key.replace(old, new) + return key + + +def _esmc( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + return _map_state(state, _esmc_key, expected_keys) + + +def _esm3( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + # state[key]: (...) + transformed = _map_state( # transformed[key]: (...) + state, + lambda key: key if key.startswith("esm3.") else f"esm3.{key}", + expected_keys, + ) + return { + key: value.to(dtype=torch.float32) if value.is_floating_point() else value # (...) + for key, value in transformed.items() + } + + +def _e1( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + return _cast_floating(state, expected_keys, torch.bfloat16) + + +def _ankh_t5( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + """Preserve ANKH's complete official encoder-decoder T5 state exactly.""" + + keys = frozenset(state) + required_exact = { + "shared.weight", + "encoder.embed_tokens.weight", + "decoder.embed_tokens.weight", + "lm_head.weight", + } + missing = sorted(required_exact - keys) + has_encoder_block = any(key.startswith("encoder.block.") for key in keys) + has_decoder_block = any(key.startswith("decoder.block.") for key in keys) + has_cross_attention = any(".EncDecAttention." in key for key in keys) + if missing or not has_encoder_block or not has_decoder_block or not has_cross_attention: + raise StateTransformError( + "ANKH publication requires the complete official T5 state, including shared, " + "encoder, decoder, cross-attention, and language-model-head parameters. " + f"Missing required keys: {missing}; encoder blocks: {has_encoder_block}; " + f"decoder blocks: {has_decoder_block}; cross-attention: {has_cross_attention}." + ) + return _identity(state, expected_keys) + + +def _boltz2( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + # state[key]: (...) + if expected_keys is None: + raise StateTransformError( + "boltz2_inference_core_v1 requires the expected FastPLMs core keys." + ) + transformed: StateDict = {} + unsupported: list[str] = [] + for source_key in sorted(state): + if source_key.startswith("ema."): + continue + key = source_key + if key.startswith("model."): + key = key[len("model.") :] + if key.startswith("module."): + key = key[len("module.") :] + canonical = key if key.startswith("core.") else f"core.{key}" + if canonical not in expected_keys: + bare = canonical[len("core.") :] + if bare.startswith(("template_module.", "bfactor_module.")): + continue + unsupported.append(source_key) + continue + if canonical in transformed: + raise StateTransformError( + f"State-key collision while mapping {source_key!r} to {canonical!r}." + ) + transformed[canonical] = _clone_tensor(source_key, state[source_key]) # (...) + if unsupported: + raise StateTransformError( + f"Boltz2 checkpoint contains undeclared non-inference parameters: {unsupported[:20]}." + ) + _validate_expected(transformed, expected_keys) + return transformed + + +_ESMFOLD_DERIVED_BUFFERS = frozenset( + { + "positional_encoding._float_tensor", + "trunk.structure_module.atom_mask", + "trunk.structure_module.default_frames", + "trunk.structure_module.group_idx", + "trunk.structure_module.lit_positions", + } +) + + +def _esmfold( + state: Mapping[str, torch.Tensor], + expected_keys: frozenset[str] | None, +) -> StateDict: + """Map native Meta ESMFold and prior canonical mirrors to package state.""" + + # state[key]: (...) + if not state: + raise StateTransformError("ESMFold checkpoint state cannot be empty.") + canonical = any(key.startswith("esm.encoder.") for key in state) + if canonical: + if expected_keys is not None: + expected_keys = frozenset( + key + for key in expected_keys + if key not in _ESMFOLD_DERIVED_BUFFERS + and not key.startswith(("mlm_head.", "esm.contact_head.")) + ) + transformed = _map_state( # transformed[key]: (...) + state, + lambda key: ( + None + if key in _ESMFOLD_DERIVED_BUFFERS + or key.startswith(("mlm_head.", "esm.contact_head.")) + else key + ), + None, + ) + _validate_expected(transformed, expected_keys) + return transformed + + folding: StateDict = {} + native_esm: StateDict = {} + for key in sorted(state): + if key in _ESMFOLD_DERIVED_BUFFERS: + continue + if key.startswith("esm."): + inner = key.removeprefix("esm.") + if inner.startswith(("lm_head.", "contact_head.")): + continue + native_esm[inner] = _clone_tensor(key, state[key]) # (...) + continue + folding[key] = _clone_tensor(key, state[key]) # (...) + mapped_esm = _esm2(native_esm, None) # mapped_esm[key]: (...) + overlap = sorted(set(folding).intersection(mapped_esm)) + if overlap: + raise StateTransformError(f"ESMFold state-key collision: {overlap[:20]}.") + transformed = {**folding, **mapped_esm} # transformed[key]: (...) + _validate_expected(transformed, expected_keys) + return transformed + + +_TRANSFORMS: dict[str, Transform] = { + "identity": _identity, + "esm2_hf_to_fastplms_v1": _esm2, + "esmc_to_fastplms_v1": _esmc, + "esm3_to_fastplms_v1": _esm3, + "e1_to_fastplms_v1": _e1, + "dplm_to_fastplms_v1": _drop_unused_rotary_position_table, + "dplm2_to_fastplms_v1": _drop_unused_rotary_position_table, + "ankh_t5_to_fastplms_v1": _ankh_t5, + "boltz2_inference_core_v1": _boltz2, + "esmfold_meta_to_fastplms_v1": _esmfold, +} + + +def available_state_transforms() -> tuple[str, ...]: + """Return stable transform identifiers accepted by the local converter.""" + + return tuple(sorted(_TRANSFORMS)) + + +def apply_state_transform( + transform_id: str, + state: Mapping[str, torch.Tensor], + *, + expected_keys: Iterable[str] | None = None, +) -> StateDict: + """Apply one manifest-declared transform without mutating ``state``.""" + + # state[key]: (...) + try: + transform = _TRANSFORMS[transform_id] + except KeyError as error: + raise StateTransformError(f"Unknown state transform: {transform_id!r}.") from error + expected = frozenset(expected_keys) if expected_keys is not None else None + return transform(state, expected) + + +__all__ = [ + "StateDict", + "StateTransformError", + "apply_state_transform", + "available_state_transforms", +] diff --git a/tools/conversion/state_validation.py b/tools/conversion/state_validation.py new file mode 100644 index 0000000..50df31f --- /dev/null +++ b/tools/conversion/state_validation.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import torch +from collections.abc import Mapping + + +def assert_model_parameters_fp32(model: torch.nn.Module, model_name: str) -> None: + non_fp32: list[dict[str, str]] = [] + parameter_count = 0 + for name, parameter in model.named_parameters(): + # parameter: (...) + parameter_count += 1 + if parameter.dtype != torch.float32: + non_fp32.append({"name": name, "dtype": str(parameter.dtype)}) + + if parameter_count == 0: + raise AssertionError(f"{model_name} has no parameters.") + if non_fp32: + raise AssertionError( + f"{model_name} parameters must all be torch.float32. " + f"non_fp32_count={len(non_fp32)} sample={non_fp32[:5]}" + ) + + +def assert_state_dict_floating_tensors_fp32( + state_dict: Mapping[str, torch.Tensor], + state_dict_name: str, +) -> None: + non_fp32: list[dict[str, str]] = [] + for tensor_name in sorted(state_dict.keys()): + tensor = state_dict[tensor_name] # (...) + if not torch.is_tensor(tensor): + raise AssertionError( + f"{state_dict_name} state_dict entry must be a tensor. " + f"name={tensor_name} type={type(tensor)}" + ) + if tensor.is_floating_point() and tensor.dtype != torch.float32: + non_fp32.append({"name": tensor_name, "dtype": str(tensor.dtype)}) + + if non_fp32: + raise AssertionError( + f"{state_dict_name} floating tensors must be torch.float32. " + f"non_fp32_count={len(non_fp32)} sample={non_fp32[:5]}" + ) + + +def assert_state_dict_equal( + reference_state_dict: Mapping[str, torch.Tensor], + candidate_state_dict: Mapping[str, torch.Tensor], + context: str, + max_report: int = 10, +) -> None: + reference_keys = set(reference_state_dict) + candidate_keys = set(candidate_state_dict) + missing = sorted(reference_keys - candidate_keys) + unexpected = sorted(candidate_keys - reference_keys) + errors: list[str] = [] + if missing: + errors.append(f"missing keys: {missing[:max_report]}") + if unexpected: + errors.append(f"unexpected keys: {unexpected[:max_report]}") + for name in sorted(reference_keys & candidate_keys): + reference = reference_state_dict[name] # (...) + candidate = candidate_state_dict[name] # (...) + if not torch.is_tensor(reference) or not torch.is_tensor(candidate): + errors.append(f"{name}: both entries must be tensors") + continue + if reference.shape != candidate.shape: + errors.append(f"{name}: shape {tuple(reference.shape)} != {tuple(candidate.shape)}") + continue + if reference.dtype != candidate.dtype: + errors.append(f"{name}: dtype {reference.dtype} != {candidate.dtype}") + continue + if not torch.equal(reference, candidate): + errors.append(f"{name}: tensor values differ") + if errors: + raise AssertionError( + f"{context} state_dict parity failed: {' | '.join(errors[:max_report])}" + ) + + +def assert_models_fp32_and_equal( + reference_model: torch.nn.Module, + candidate_model: torch.nn.Module, + context: str, + max_report: int = 5, +) -> None: + assert_model_parameters_fp32(model=reference_model, model_name=f"{context} reference model") + assert_model_parameters_fp32(model=candidate_model, model_name=f"{context} candidate model") + assert_state_dict_equal( + reference_state_dict=reference_model.state_dict(), + candidate_state_dict=candidate_model.state_dict(), + context=context, + max_report=max_report, + ) diff --git a/tools/debug/README.md b/tools/debug/README.md new file mode 100644 index 0000000..af1c5b7 --- /dev/null +++ b/tools/debug/README.md @@ -0,0 +1,13 @@ +# Maintained diagnostics + +This directory contains small, reusable repository checks that do not belong in the +runtime package. Diagnostic scripts must be deterministic, accept their inputs at +runtime, and remain free of credentials, Hub mutations, cache patching, and +machine-specific paths. + +`check_notation.py` enforces the documentation and comment notation contract. It is +also exercised by the release test suite. + +One-off parity investigations belong in an untracked work directory. Promote an +investigation into this directory only when it becomes a supported diagnostic with +tests and documented inputs. diff --git a/tools/debug/analyze_boltz_opm_projection.py b/tools/debug/analyze_boltz_opm_projection.py new file mode 100644 index 0000000..b3da0be --- /dev/null +++ b/tools/debug/analyze_boltz_opm_projection.py @@ -0,0 +1,178 @@ +"""Compare bounded runtime variants for Boltz2's first OPM output projection.""" + +from __future__ import annotations + +import argparse +import json +import torch +import torch.nn.functional as F +from collections.abc import Callable, Sequence +from pathlib import Path +from safetensors.torch import load_file + + +_PREFIX = "msa_module__layers__0__outer_product_mean__proj_o" + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("candidate", type=Path) + parser.add_argument("reference", type=Path) + parser.add_argument("--output", type=Path) + parser.add_argument( + "--device", + choices=("auto", "cpu", "cuda"), + default="auto", + help="Execution device for bounded projection variants.", + ) + return parser + + +def _comparison(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, object]: + # actual, expected: (...); both tensors have the same shape. + difference = actual.float() - expected.float() # (...) + scale = torch.linalg.vector_norm(expected.float()).clamp_min( # () + torch.finfo(torch.float32).tiny + ) + unequal = torch.ne(actual, expected) # (...) + return { + "exact": bool(torch.equal(actual, expected)), + "unequal_values": int(unequal.sum().item()), + "max_absolute_error": float(difference.abs().max().item()), + "relative_l2": float((torch.linalg.vector_norm(difference) / scale).item()), + } + + +def _chunked_linear( + X: torch.Tensor, + W: torch.Tensor, + bias: torch.Tensor, + chunk_size: int, +) -> torch.Tensor: + # X: (..., d_in); W: (d_out, d_in); bias: (d_out) + output = torch.zeros((*X.shape[:-1], W.shape[0]), device=X.device) # (..., d_out) + for start in range(0, X.shape[-1], chunk_size): + stop = min(start + chunk_size, X.shape[-1]) + # X[..., start:stop]: (..., d_chunk); W[:, start:stop].T: (d_chunk, d_out) + output.add_(X[..., start:stop] @ W[:, start:stop].T) # (..., d_out) + return output + bias # (..., d_out) + + +def _autocast_linear( + X: torch.Tensor, + W: torch.Tensor, + bias: torch.Tensor, + *, + allow_reduced_bf16_reduction: bool, +) -> torch.Tensor: + # X: (..., d_in); W: (d_out, d_in); bias: (d_out) + if not X.is_cuda: + raise RuntimeError("The BF16 reduction-policy probe requires CUDA.") + previous = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + try: + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = ( + allow_reduced_bf16_reduction + ) + with torch.autocast("cuda", dtype=torch.bfloat16): + return F.linear(X, W, bias) # (..., d_out) + finally: + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = previous + + +def main(argv: Sequence[str] | None = None) -> int: + """Report isolated projection errors against the native BF16 oracle.""" + + arguments = _parser().parse_args(argv) + candidate = load_file(arguments.candidate, device="cpu") # values: (...) + reference = load_file(arguments.reference, device="cpu") # values: (...) + input_key = f"{_PREFIX}__call_000__args__0" + output_key = f"{_PREFIX}__call_000__output" + weight_key = f"{_PREFIX}__parameter__weight" + bias_key = f"{_PREFIX}__parameter__bias" + contract_equal = { + "input": torch.equal(candidate[input_key], reference[input_key]), + "weight": torch.equal(candidate[weight_key], reference[weight_key]), + "bias": torch.equal(candidate[bias_key], reference[bias_key]), + } + for name, equal in contract_equal.items(): + if not equal: + raise RuntimeError(f"Projection trace {name} differs.") + + if arguments.device == "auto": + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(arguments.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable.") + # d_in and d_out are the projection input and output widths. + X = candidate[input_key].to(device) # (..., d_in) + W = candidate[weight_key].to(device) # (d_out, d_in) + bias = candidate[bias_key].to(device) # (d_out) + expected = reference[output_key] # (..., d_out) + output_dtype = expected.dtype + X_bf16 = X.to(torch.bfloat16).float() # (..., d_in) + W_bf16 = W.to(torch.bfloat16).float() # (d_out, d_in) + bias_bf16 = bias.to(torch.bfloat16).float() # (d_out) + + # Every variant returns (..., d_out). + variants: dict[str, Callable[[], torch.Tensor]] = { + "recorded_autocast_bf16": lambda: candidate[output_key].to(device), + "fp32_linear": lambda: F.linear(X.float(), W.float(), bias.float()).to(output_dtype), + "bf16_operands_fp32_linear": lambda: F.linear(X_bf16, W_bf16, bias_bf16).to(output_dtype), + "bf16_operands_fp32_chunk32": lambda: _chunked_linear(X_bf16, W_bf16, bias_bf16, 32).to( + output_dtype + ), + "bf16_operands_fp32_chunk64": lambda: _chunked_linear(X_bf16, W_bf16, bias_bf16, 64).to( + output_dtype + ), + } + if device.type == "cuda": + variants.update( + { + "autocast_bf16_reduced_reduction_on": lambda: _autocast_linear( + X, + W, + bias, + allow_reduced_bf16_reduction=True, + ), + "autocast_bf16_reduced_reduction_off": lambda: _autocast_linear( + X, + W, + bias, + allow_reduced_bf16_reduction=False, + ), + } + ) + previous_tf32 = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + results = { + name: _comparison(compute().cpu(), expected) for name, compute in variants.items() + } + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_tf32 + + payload = { + "candidate_torch": torch.__version__, + "device": str(device), + "localization": { + "operation": "msa_module.layers.0.outer_product_mean.proj_o", + "first_differing_kernel": "autocast_bf16_linear_output", + "input_equal": contract_equal["input"], + "weight_equal": contract_equal["weight"], + "bias_equal": contract_equal["bias"], + "recorded_output_equal": torch.equal(candidate[output_key], reference[output_key]), + }, + "output_dtype": str(output_dtype), + "variants": results, + } + serialized = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if arguments.output is not None: + arguments.output.parent.mkdir(parents=True, exist_ok=True) + arguments.output.write_text(serialized, encoding="utf-8") + print(serialized, end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/check_notation.py b/tools/debug/check_notation.py new file mode 100644 index 0000000..130e90f --- /dev/null +++ b/tools/debug/check_notation.py @@ -0,0 +1,256 @@ +"""Reject noncanonical tensor-shape notation in documentation and comments.""" + +from __future__ import annotations + +import argparse +import ast +import io +import re +import tokenize +from collections.abc import Iterable, Iterator +from dataclasses import dataclass +from pathlib import Path + + +TEXT_SUFFIXES = frozenset( + {".hcl", ".in", ".md", ".rst", ".toml", ".txt", ".yaml", ".yml"} +) +SKIP_PARTS = frozenset( + { + ".git", + ".pytest_cache", + ".ruff_cache", + ".venv", + "LICENSES", + "artifacts", + "dist", + "vendor", + "__pycache__", + } +) +SQUARE_CANDIDATE = re.compile( + r"\[\s*(?:[A-Za-z][A-Za-z0-9_]*|\d+|\.\.\.|\*)" + r"(?:\s*,\s*(?:[A-Za-z][A-Za-z0-9_]*|\d+|\.\.\.|\*))*\s*\]" +) +UPPER_PAREN_CANDIDATE = re.compile( + r"\(\s*(?:[A-Za-z][A-Za-z0-9_]*|\d+)(?:\s*,\s*(?:[A-Za-z][A-Za-z0-9_]*|\d+)){1,}\s*\)" +) +COMMON_DIMENSIONS = frozenset( + { + "b", + "l", + "d", + "n", + "h", + "q", + "k", + "v", + "s", + "t", + "c", + "m", + "r", + "p", + "batch", + "batch_size", + "length", + "seq_len", + "sequence_length", + "hidden_size", + "d_model", + "d_head", + "n_layers", + "n_heads", + "heads", + "layers", + "atoms", + "channels", + "items", + "residues", + "samples", + "sequences", + "tokens", + } +) + + +@dataclass(frozen=True, slots=True) +class Violation: + path: Path + line: int + column: int + message: str + excerpt: str + + def render(self, root: Path) -> str: + relative = self.path.resolve().relative_to(root.resolve()) + return ( + f"{relative.as_posix()}:{self.line}:{self.column}: {self.message}: " + f"{self.excerpt.strip()}" + ) + + +def _tokens(candidate: str) -> tuple[str, ...]: + return tuple(token.strip() for token in candidate[1:-1].split(",")) + + +def _is_dimension_name(token: str) -> bool: + lowered = token.lower() + return ( + lowered in COMMON_DIMENSIONS + or lowered.endswith(("_dim", "_dims", "_len", "_length", "_size")) + or lowered.startswith( + ( + "b_", + "l_", + "d_", + "n_", + "h_", + "q_", + "k_", + "v_", + "s_", + "t_", + "c_", + "m_", + "r_", + "p_", + "batch_", + "seq_", + "sequence_", + "hidden_", + "head_", + "layer_", + ) + ) + ) + + +def violations_in_text( + text: str, + *, + path: Path, + first_line: int = 1, +) -> Iterator[Violation]: + """Yield notation violations in one documentation string or comment.""" + + for offset, line in enumerate(text.splitlines() or (text,)): + line_number = first_line + offset + for match in SQUARE_CANDIDATE.finditer(line): + dimensions = _tokens(match.group()) + if any(_is_dimension_name(token) for token in dimensions): + yield Violation( + path, + line_number, + match.start() + 1, + "shape signatures must use parentheses", + match.group(), + ) + for match in UPPER_PAREN_CANDIDATE.finditer(line): + dimensions = _tokens(match.group()) + uppercase_dimensions = [ + token + for token in dimensions + if token.isidentifier() and token.upper() == token and _is_dimension_name(token) + ] + if uppercase_dimensions: + yield Violation( + path, + line_number, + match.start() + 1, + "shape dimensions must use lowercase symbols", + match.group(), + ) + + +def _python_documentation(path: Path) -> Iterator[tuple[str, int]]: + """Yield comments and real docstrings with source line numbers.""" + + source = path.read_text(encoding="utf-8") + try: + tokens = tokenize.generate_tokens(io.StringIO(source).readline) + for token in tokens: + if token.type == tokenize.COMMENT: + yield token.string[1:], token.start[0] + except (IndentationError, tokenize.TokenError) as error: + raise ValueError(f"Cannot tokenize {path}: {error}") from error + try: + tree = ast.parse(source, filename=str(path)) + except SyntaxError as error: + raise ValueError(f"Cannot parse {path}: {error}") from error + documentable = (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + for node in ast.walk(tree): + if not isinstance(node, documentable) or not node.body: + continue + expression = node.body[0] + if ( + isinstance(expression, ast.Expr) + and isinstance(expression.value, ast.Constant) + and isinstance(expression.value.value, str) + ): + yield expression.value.value, expression.lineno + + +def iter_repository_files(root: Path) -> Iterator[Path]: + """Yield tracked documentation and Python paths inside the source boundary.""" + + candidates = ( + root / "README.md", + root / "AGENTS.md", + root / "THIRD_PARTY_NOTICES.md", + root / "LICENSES" / "README.md", + root / "vendor" / "README.md", + root / "requirements", + root / "docs", + root / "model_cards", + root / "docker", + root / "src", + root / "tests", + root / "benchmarks", + root / "tools", + root / "examples", + ) + for candidate in candidates: + if candidate.is_file(): + yield candidate + continue + if not candidate.is_dir(): + continue + for path in sorted(candidate.rglob("*")): + if not path.is_file() or any(part in SKIP_PARTS for part in path.parts): + continue + if path.suffix in TEXT_SUFFIXES or path.suffix == ".py" or path.name == "Dockerfile": + yield path + + +def scan_repository(root: Path) -> list[Violation]: + """Return every shape-notation violation in repository prose.""" + + result: list[Violation] = [] + for path in iter_repository_files(root): + if path.suffix == ".py": + regions = _python_documentation(path) + else: + regions = ((path.read_text(encoding="utf-8"), 1),) + for text, first_line in regions: + result.extend(violations_in_text(text, path=path, first_line=first_line)) + return result + + +def main(argv: Iterable[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source-root", + type=Path, + default=Path(__file__).resolve().parents[2], + ) + arguments = parser.parse_args(argv) + root = arguments.source_root.resolve() + violations = scan_repository(root) + for violation in violations: + print(violation.render(root)) + return 1 if violations else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/compare_structure_bundles.py b/tools/debug/compare_structure_bundles.py new file mode 100644 index 0000000..5aa1b9c --- /dev/null +++ b/tools/debug/compare_structure_bundles.py @@ -0,0 +1,105 @@ +"""Report tensor-level differences between two structure compliance bundles.""" + +from __future__ import annotations + +import argparse +import torch +from collections.abc import Sequence +from pathlib import Path +from safetensors.torch import load_file + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("actual", type=Path) + parser.add_argument("expected", type=Path) + parser.add_argument( + "--prefix", + default="", + help="Only compare tensor keys beginning with this prefix.", + ) + parser.add_argument( + "--contains", + action="append", + default=[], + help="Only compare keys containing every supplied fragment.", + ) + parser.add_argument( + "--max-differences", + type=int, + default=1, + help="Maximum number of differing values to display per tensor.", + ) + parser.add_argument( + "--largest-first", + action="store_true", + help="Display the largest absolute differences instead of index order.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Compare safetensors bundles without loading package or upstream code.""" + + arguments = _parser().parse_args(argv) + actual = load_file(arguments.actual, device="cpu") # values: (...) + expected = load_file(arguments.expected, device="cpu") # values: (...) + keys = sorted( + key + for key in set(actual) | set(expected) + if key.startswith(arguments.prefix) + and all(fragment in key for fragment in arguments.contains) + ) + print( + "key\tactual_dtype\texpected_dtype\tactual_shape\texpected_shape\texact\t" + "unequal_values\tmax_absolute_error\trelative_l2\tfirst_difference" + ) + for key in keys: + if key not in actual or key not in expected: + present = "actual" if key in actual else "expected" + print(f"{key}\tpresent_only_in_{present}") + continue + # r is the rank of this bundle entry; shapes vary by tensor key. + X = actual[key] # (...) + X_ref = expected[key] # (...) + shape_matches = X.shape == X_ref.shape + dtype_matches = X.dtype == X_ref.dtype + exact = shape_matches and dtype_matches and torch.equal(X, X_ref) + first_difference = "" + if shape_matches and X.numel() > 0: + unequal = torch.ne(X, X_ref) # (...) + unequal_count = int(unequal.sum().item()) + if unequal_count: + indices = unequal.nonzero(as_tuple=False) # (n_diff, r) + if arguments.largest_first: + errors = (X.float() - X_ref.float()).abs()[unequal] # (n_diff,) + order = torch.argsort(errors, descending=True) # (n_diff,) + indices = indices[order] # (n_diff, r) + differences = [] + for raw_index in indices[: arguments.max_differences]: + # raw_index: (r,) + index = tuple(raw_index.tolist()) + differences.append(f"{index}: {X[index].item()} != {X_ref[index].item()}") + first_difference = "; ".join(differences) + max_error = (X.float() - X_ref.float()).abs().max().item() + error = f"{max_error:.9g}" + difference_norm = torch.linalg.vector_norm(X.float() - X_ref.float()) # () + reference_norm = torch.linalg.vector_norm(X_ref.float()).clamp_min( + torch.finfo(torch.float32).tiny + ) # () + relative_l2 = f"{(difference_norm / reference_norm).item():.9g}" + else: + unequal_count = "n/a" + error = "n/a" + relative_l2 = "n/a" + print( + f"{key}\t{X.dtype}\t{X_ref.dtype}\t{tuple(X.shape)}\t" + f"{tuple(X_ref.shape)}\t" + f"{exact}\t{unequal_count}\t{error}\t{relative_l2}\t" + f"{first_difference}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/export_boltz_conformers.py b/tools/debug/export_boltz_conformers.py new file mode 100644 index 0000000..6883b8c --- /dev/null +++ b/tools/debug/export_boltz_conformers.py @@ -0,0 +1,77 @@ +"""Export pinned Boltz molecule metadata as dependency-free Python data.""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from pathlib import Path + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("molecule_dir", type=Path) + parser.add_argument( + "--name", + action="append", + dest="names", + help="Residue name to export; repeat for more than one residue.", + ) + parser.add_argument("--output", type=Path, help="Write JSON to this path.") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Print atom metadata and conformer coordinates from official pickles.""" + + from boltz.data import const + from boltz.data.mol import load_molecules + + arguments = _parser().parse_args(argv) + names = ( + sorted(arguments.names) + if arguments.names + else sorted(path.stem for path in arguments.molecule_dir.glob("*.pkl")) + ) + molecules = load_molecules(arguments.molecule_dir, names) + unknown_chirality = const.chirality_type_ids[const.unk_chirality_type] + payload = {} + for name in names: + molecule = molecules[name] + atoms = [] + for atom in molecule.GetAtoms(): + atoms.append( + { + "name": atom.GetProp("name"), + "charge": atom.GetFormalCharge(), + "chirality": const.chirality_type_ids.get( + atom.GetChiralTag().name, + unknown_chirality, + ), + } + ) + conformers = [] + for conformer in molecule.GetConformers(): + conformers.append( + { + atom["name"]: [ + conformer.GetAtomPosition(index).x, + conformer.GetAtomPosition(index).y, + conformer.GetAtomPosition(index).z, + ] + for index, atom in enumerate(atoms) + } + ) + payload[name] = {"atoms": atoms, "conformers": conformers} + serialized = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if arguments.output is None: + print(serialized, end="") + else: + arguments.output.parent.mkdir(parents=True, exist_ok=True) + arguments.output.write_text(serialized, encoding="utf-8") + print(arguments.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/generate_boltz_conformer_patch.py b/tools/debug/generate_boltz_conformer_patch.py new file mode 100644 index 0000000..f4af426 --- /dev/null +++ b/tools/debug/generate_boltz_conformer_patch.py @@ -0,0 +1,134 @@ +"""Generate an apply-patch update from exported official Boltz conformers.""" + +from __future__ import annotations + +import argparse +import ast +import json +import pprint +import textwrap +from collections.abc import Sequence +from pathlib import Path + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("target", type=Path) + parser.add_argument("conformers", type=Path) + parser.add_argument("--name", action="append", dest="names") + parser.add_argument( + "--repair-chirality", + action="store_true", + help="Regenerate the canonical chiral-atom table from the export.", + ) + parser.add_argument( + "--patch-path-label", + help="Path written into the patch when it differs from the readable target.", + ) + return parser + + +def _existing_table(source: str) -> dict[str, dict[str, list[float]]]: + tree = ast.parse(source) + for node in tree.body: + if ( + isinstance(node, ast.AnnAssign) + and isinstance(node.target, ast.Name) + and node.target.id == "_RDKIT_CONFORMERS" + ): + value = ast.literal_eval(node.value) + if isinstance(value, dict): + return value + raise RuntimeError("Target source omits _RDKIT_CONFORMERS.") + + +def _assignment_lines(source: str, variable: str) -> list[str]: + tree = ast.parse(source) + lines = source.splitlines() + for node in tree.body: + if isinstance(node, ast.AnnAssign): + target = node.target + elif isinstance(node, ast.Assign): + target = node.targets[0] if len(node.targets) == 1 else None + else: + target = None + if isinstance(target, ast.Name) and target.id == variable: + assert node.end_lineno is not None + return lines[node.lineno - 1 : node.end_lineno] + raise RuntimeError(f"Target source omits {variable}.") + + +def _mapping_entry_lines(source: str, variable: str, entry: str) -> list[str]: + """Return every source line occupied by one top-level mapping entry.""" + + tree = ast.parse(source) + lines = source.splitlines() + for node in tree.body: + if isinstance(node, ast.AnnAssign): + target = node.target + value = node.value + elif isinstance(node, ast.Assign): + target = node.targets[0] if len(node.targets) == 1 else None + value = node.value + else: + continue + if not (isinstance(target, ast.Name) and target.id == variable): + continue + if not isinstance(value, ast.Dict): + raise TypeError(f"{variable} must be a dictionary literal.") + for key_node, value_node in zip(value.keys, value.values, strict=True): + if ( + isinstance(key_node, ast.Constant) + and key_node.value == entry + and value_node.end_lineno is not None + ): + return lines[key_node.lineno - 1 : value_node.end_lineno] + raise KeyError(f"{variable} omits {entry!r}.") + raise RuntimeError(f"Target source omits {variable}.") + + +def main(argv: Sequence[str] | None = None) -> int: + """Print a patch retaining full-precision coordinates for selected names.""" + + arguments = _parser().parse_args(argv) + source = arguments.target.read_text(encoding="utf-8") + existing = _existing_table(source) + exported = json.loads(arguments.conformers.read_text(encoding="utf-8")) + if not arguments.names and not arguments.repair_chirality: + raise ValueError("Select at least one conformer name or --repair-chirality.") + print("*** Begin Patch") + print(f"*** Update File: {arguments.patch_path_label or arguments.target}") + if arguments.repair_chirality: + old_lines = _assignment_lines(source, "_CHIRAL_ATOMS") + chiral_atoms = { + name: frozenset(atom["name"] for atom in value["atoms"] if atom["chirality"] != 0) + for name, value in exported.items() + if any(atom["chirality"] != 0 for atom in value["atoms"]) + } + rendered = pprint.pformat(chiral_atoms, width=88, sort_dicts=False) + rendered_lines = rendered.splitlines() + new_lines = [f"_CHIRAL_ATOMS = {rendered_lines[0]}", *rendered_lines[1:]] + print("@@") + for line in old_lines: + print(f"-{line}") + for line in new_lines: + print(f"+{line}") + for name in arguments.names or (): + old_lines = _mapping_entry_lines(source, "_RDKIT_CONFORMERS", name) + raw = exported[name]["conformers"][0] + entry = {atom: raw[atom] for atom in existing[name]} + rendered = pprint.pformat(entry, width=88, sort_dicts=False) + lines = textwrap.indent(rendered, " ").splitlines() + new_lines = [f' "{name}": {lines[0].lstrip()}', *lines[1:]] + new_lines[-1] += "," + print("@@") + for line in old_lines: + print(f"-{line}") + for line in new_lines: + print(f"+{line}") + print("*** End Patch") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/probe_cuda_loader.py b/tools/debug/probe_cuda_loader.py new file mode 100644 index 0000000..8fd16df --- /dev/null +++ b/tools/debug/probe_cuda_loader.py @@ -0,0 +1,49 @@ +"""Report CUDA library resolution and one required cuBLASLt symbol.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + + +REQUIRED_SYMBOL = "cublasLtGroupedMatrixLayoutInit_internal" +TE_LIBRARY = Path( + "/opt/venv/lib/python3.12/site-packages/transformer_engine/wheel_lib/libtransformer_engine.so" +) +CUDA_CUBLAS_LT = Path("/usr/local/cuda/lib64/libcublasLt.so.13") +PYTHON_CUBLAS_LT = Path("/opt/venv/lib/python3.12/site-packages/nvidia/cu13/lib/libcublasLt.so.13") + + +def _output(*command: str) -> str: + return subprocess.run( + command, + check=True, + capture_output=True, + text=True, + ).stdout + + +def _symbol_matches(path: Path) -> list[str]: + output = _output("objdump", "-T", str(path)) + return [line.strip() for line in output.splitlines() if REQUIRED_SYMBOL in line] + + +def main() -> None: + paths = (CUDA_CUBLAS_LT, PYTHON_CUBLAS_LT) + result = { + "required_symbol": REQUIRED_SYMBOL, + "transformer_engine_ldd": _output("ldd", str(TE_LIBRARY)).splitlines(), + "cublas_lt": { + str(path): { + "exists": path.exists(), + "required_symbol_matches": _symbol_matches(path) if path.exists() else [], + } + for path in paths + }, + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/debug/probe_flash_attention_forward.py b/tools/debug/probe_flash_attention_forward.py new file mode 100644 index 0000000..d608871 --- /dev/null +++ b/tools/debug/probe_flash_attention_forward.py @@ -0,0 +1,223 @@ +"""Validate precompiled FlashAttention kernels on shared and model paths.""" + +from __future__ import annotations + +import json +import torch +from types import SimpleNamespace +from typing import Any +from torch.nn import functional as F + +from fastplms.attention import FASTPLMS_ATTENTION_FUNCTIONS +from fastplms.attention._core import kernels_flash_attention_func +from fastplms.models.esm2.modeling_fastesm import FastEsmConfig, FastEsmModel +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, +) + + +BACKENDS = ("flash_attention_2", "flash_attention_3") +MODEL_BACKENDS = { + "esm2": BACKENDS, + "esm_plusplus": BACKENDS, +} + + +def _metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + # actual, expected: (..., d) + actual_float = actual.float() # (..., d) + expected_float = expected.float() # (..., d) + relative_l2 = torch.linalg.vector_norm( + actual_float - expected_float + ) / torch.linalg.vector_norm(expected_float).clamp_min(1e-12) # () + cosine = F.cosine_similarity( + actual_float.reshape(-1, actual.shape[-1]), # (n, d) + expected_float.reshape(-1, expected.shape[-1]), # (n, d) + dim=-1, + ) # (n,) + result = { + "relative_l2": relative_l2.item(), + "minimum_cosine": cosine.min().item(), + } + if result["relative_l2"] > 1e-2 or result["minimum_cosine"] < 0.999: + raise RuntimeError(f"FlashAttention parity failed: {result}") + return result + + +def _sdpa_reference( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, +) -> torch.Tensor: + # query_states, key_states, value_states: (b, l, h, d_h) + return F.scaled_dot_product_attention( + query_states.transpose(1, 2), # (b, h, l, d_h) + key_states.transpose(1, 2), # (b, h, l, d_h) + value_states.transpose(1, 2), # (b, h, l, d_h) + ).transpose(1, 2) # (b, l, h, d_h) + + +def _shared_results() -> dict[str, Any]: + torch.manual_seed(17) + query_states = torch.randn( # (b=2, l=17, h=4, d_h=16) + 2, 17, 4, 16, device="cuda", dtype=torch.bfloat16 + ) + key_states = torch.randn_like(query_states) # (b, l, h, d_h) + value_states = torch.randn_like(query_states) # (b, l, h, d_h) + attention_mask = torch.tensor( + [[1] * 17, [1] * 9 + [0] * 8], + device="cuda", + dtype=torch.bool, + ) # (b, l) + dense_reference = _sdpa_reference( # (b, l, h, d_h) + query_states, key_states, value_states + ) + mixed_reference = torch.zeros_like(dense_reference) # (b, l, h, d_h) + for batch_index, length in enumerate((17, 9)): + mixed_reference[batch_index, :length] = _sdpa_reference( # (l_i, h, d_h) + query_states[batch_index : batch_index + 1, :length], # (1, l_i, h, d_h) + key_states[batch_index : batch_index + 1, :length], # (1, l_i, h, d_h) + value_states[batch_index : batch_index + 1, :length], # (1, l_i, h, d_h) + )[0] # (l_i, h, d_h) + + results: dict[str, Any] = {} + module = SimpleNamespace(training=False, is_causal=False) + for backend in BACKENDS: + dense = kernels_flash_attention_func( # (b, l, h, d_h) + query_states, + key_states, + value_states, + implementation=backend, + ) + mixed = kernels_flash_attention_func( # (b, l, h, d_h) + query_states, + key_states, + value_states, + attention_mask_2d=attention_mask, + implementation=backend, + ) + interface_output, interface_weights = FASTPLMS_ATTENTION_FUNCTIONS[backend]( + module, + query_states.transpose(1, 2), # (b, h, l, d_h) + key_states.transpose(1, 2), # (b, h, l, d_h) + value_states.transpose(1, 2), # (b, h, l, d_h) + attention_mask, # (b, l) + ) + # interface_output: (b, l, h, d_h); interface_weights: None + if interface_weights is not None or not torch.equal(interface_output, mixed): + raise RuntimeError(f"{backend} AttentionInterface dispatch disagrees with core.") + results[backend] = { + "dense": _metrics(dense, dense_reference), + "mixed_padding": _metrics( + mixed[attention_mask], # (n_valid, h, d_h) + mixed_reference[attention_mask], # (n_valid, h, d_h) + ), + "padding_is_zero": bool( + torch.equal( + mixed[~attention_mask], + torch.zeros_like(mixed[~attention_mask]), + ) + ), + } + if not results[backend]["padding_is_zero"]: + raise RuntimeError(f"{backend} returned nonzero values at padded positions.") + return results + + +def _model_specs() -> tuple[tuple[str, type[torch.nn.Module], object], ...]: + esm_kwargs = { + "vocab_size": 33, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "intermediate_size": 128, + "hidden_dropout_prob": 0.0, + "attention_probs_dropout_prob": 0.0, + "max_position_embeddings": 64, + "pad_token_id": 1, + "mask_token_id": 32, + "position_embedding_type": "rotary", + "attn_backend": "eager", + } + return ( + ("esm2", FastEsmModel, FastEsmConfig(**esm_kwargs, token_dropout=False)), + ( + "esm_plusplus", + ESMplusplusModel, + ESMplusplusConfig( + vocab_size=33, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + dropout=0.0, + attn_backend="eager", + pad_token_id=1, + ), + ), + ) + + +def _last_hidden_state(output: object) -> torch.Tensor: + value = getattr(output, "last_hidden_state", None) + if not torch.is_tensor(value): + raise TypeError("Model output omitted last_hidden_state.") + return value # (b, l, d) + + +def _model_results() -> dict[str, Any]: + input_ids = torch.tensor( + [ + [0, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 2, 1, 1, 1, 1], + [0, 20, 19, 18, 17, 16, 15, 14, 2, 1, 1, 1, 1, 1, 1, 1, 1], + ], + device="cuda", + ) # (b=2, l=17) + attention_mask = input_ids.ne(1) # (b, l) + results: dict[str, Any] = {} + for family, model_class, config in _model_specs(): + torch.manual_seed(23) + model = model_class(config).eval().to(device="cuda", dtype=torch.bfloat16) + outputs: dict[str, torch.Tensor] = {} + backends = MODEL_BACKENDS[family] + with torch.inference_mode(): + for backend in ("eager", "sdpa", *backends): + model.attn_backend = backend + outputs[backend] = _last_hidden_state( # (b, l, d) + model(input_ids=input_ids, attention_mask=attention_mask) + ).detach() + valid = attention_mask # (b, l) + family_results: dict[str, Any] = {} + for backend in backends: + family_results[backend] = { + "vs_eager": _metrics( + outputs[backend][valid], # (n_valid, d) + outputs["eager"][valid], # (n_valid, d) + ), + "vs_sdpa": _metrics( + outputs[backend][valid], # (n_valid, d) + outputs["sdpa"][valid], # (n_valid, d) + ), + "finite": bool(torch.isfinite(outputs[backend]).all()), + } + if not family_results[backend]["finite"]: + raise RuntimeError(f"{family} {backend} produced non-finite values.") + results[family] = family_results + del model + return results + + +def main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for FlashAttention validation.") + result = { + "device": torch.cuda.get_device_name(0), + "shared": _shared_results(), + "models": _model_results(), + "torch": torch.__version__, + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/debug/probe_flash_checkpoint_forward.py b/tools/debug/probe_flash_checkpoint_forward.py new file mode 100644 index 0000000..98a1ab7 --- /dev/null +++ b/tools/debug/probe_flash_checkpoint_forward.py @@ -0,0 +1,143 @@ +"""Run warm-cache FlashAttention parity on representative checkpoints.""" + +from __future__ import annotations + +import argparse +import contextlib +import gc +import json +import torch +from collections.abc import Sequence +from typing import Any +from torch.nn import functional as F +from transformers import EsmTokenizer + +from fastplms.models.dplm.modeling_dplm import DPLMModel +from fastplms.models.esm2.modeling_fastesm import FastEsmForMaskedLM +from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusForMaskedLM, + EsmSequenceTokenizer, +) +from fastplms.registry import get_model_registry + + +BACKENDS = ("flash_attention_2", "flash_attention_3") +CHECKPOINTS = { + "esm2_8m": FastEsmForMaskedLM, + "esmc_small": ESMplusplusForMaskedLM, + "dplm_150m": DPLMModel, +} +BACKENDS_BY_CHECKPOINT = { + "esm2_8m": BACKENDS, + "esmc_small": BACKENDS, + "dplm_150m": ("flash_attention_3",), +} + + +def _metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + # actual, expected: (n_valid, d) + actual_float = actual.float() # (n_valid, d) + expected_float = expected.float() # (n_valid, d) + relative_l2 = ( + torch.linalg.vector_norm(actual_float - expected_float) + / torch.linalg.vector_norm(expected_float).clamp_min(1e-12) + ).item() # scalar + cosine_values = F.cosine_similarity( # (n_valid,) + actual_float, + expected_float, + dim=-1, + ) + cosine = cosine_values.min().item() + return {"relative_l2": relative_l2, "minimum_cosine": cosine} + + +def _hidden_state(output: object) -> torch.Tensor: + value = getattr(output, "last_hidden_state", None) + if not torch.is_tensor(value): + raise TypeError("Checkpoint output omitted last_hidden_state.") + return value # (b, l, d) + + +def _run_checkpoint( + model_id: str, + model_class: type[torch.nn.Module], +) -> dict[str, Any]: + spec = get_model_registry()[model_id] + if spec.family.id == "esm_plusplus": + tokenizer = EsmSequenceTokenizer() + else: + tokenizer = EsmTokenizer.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + ) + batch = tokenizer( + ["MSTNPKPQRKTKRNTNR", "ACDEFGHIK"], + return_tensors="pt", + padding=True, + ) + # Each tokenized batch value has shape (b=2, l). + batch = {name: value.to("cuda") for name, value in batch.items()} # values: (b, l) + use_bf16_autocast = spec.family.bf16_execution == "fp32_parameters_autocast" + load_dtype = torch.float32 if use_bf16_autocast else torch.bfloat16 + model = ( + model_class.from_pretrained( + spec.fast.repo_id, + revision=spec.fast.revision, + dtype=load_dtype, + ) + .eval() + .to("cuda") + ) + outputs: dict[str, torch.Tensor] = {} + with torch.inference_mode(): + for backend in ("sdpa", *BACKENDS_BY_CHECKPOINT[model_id]): + model.set_attn_implementation(backend) + numeric_context = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if use_bf16_autocast + else contextlib.nullcontext() + ) + with numeric_context: + outputs[backend] = _hidden_state(model(**batch)).detach() # (b, l, d) + valid = batch["attention_mask"].bool() # (b, l) + result = { + backend: { + **_metrics( + outputs[backend][valid], # (n_valid, d) + outputs["sdpa"][valid], # (n_valid, d) + ), + "finite": bool(torch.isfinite(outputs[backend]).all()), + } + for backend in BACKENDS_BY_CHECKPOINT[model_id] + } + if not all(value["finite"] for value in result.values()): + raise RuntimeError(f"{model_id} produced non-finite checkpoint output.") + del model + gc.collect() + torch.cuda.empty_cache() + return { + "checkpoint": spec.fast.repo_id, + "revision": spec.fast.revision, + "backends": result, + } + + +def main(argv: Sequence[str] | None = None) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for checkpoint FlashAttention validation.") + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-id", action="append", choices=tuple(CHECKPOINTS)) + arguments = parser.parse_args(argv) + model_ids = tuple(arguments.model_id or CHECKPOINTS) + result = { + "device": torch.cuda.get_device_name(0), + "models": { + model_id: _run_checkpoint(model_id, CHECKPOINTS[model_id]) for model_id in model_ids + }, + "torch": torch.__version__, + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/debug/probe_flash_kernels.py b/tools/debug/probe_flash_kernels.py new file mode 100644 index 0000000..dff1b92 --- /dev/null +++ b/tools/debug/probe_flash_kernels.py @@ -0,0 +1,79 @@ +"""Probe the two precompiled Hugging Face FlashAttention repositories.""" + +from __future__ import annotations + +import inspect +import json +from importlib.metadata import version +from pathlib import Path +from kernels import get_kernel_variants, has_kernel + +from fastplms.attention._kernel_lock import load_locked_kernel +from fastplms.registry import get_model_registry + + +API_NAMES = ( + "fwd", + "varlen_fwd", + "flash_attn_func", + "flash_attn_varlen_func", +) + + +def _variant_summary(decisions: list[object]) -> dict[str, object]: + """Keep the diagnostic concise while preserving every accepted variant.""" + return { + "accepted": [ + repr(getattr(decision, "variant", decision)) + for decision in decisions + if type(decision).__name__ == "VariantAccepted" + ], + "rejected_count": sum( + type(decision).__name__ == "VariantRejected" for decision in decisions + ), + } + + +def main() -> None: + result = { + "kernels_version": version("kernels"), + "api_signatures": { + "get_kernel_variants": str(inspect.signature(get_kernel_variants)), + "has_kernel": str(inspect.signature(has_kernel)), + }, + "repositories": {}, + } + for spec in get_model_registry().attention_kernels.values(): + repository = spec.repository + available = has_kernel( + repository, + revision=spec.revision, + ) + compatible_variants = get_kernel_variants( + repository, + revision=spec.revision, + ) + try: + kernel = load_locked_kernel(repository, spec.revision) + except Exception as error: + result["repositories"][repository] = { + "error": f"{type(error).__name__}: {error}", + "revision": spec.revision, + "version": spec.version, + "has_kernel": available, + "variants": _variant_summary(compatible_variants), + } + continue + result["repositories"][repository] = { + "revision": spec.revision, + "version": spec.version, + "has_kernel": available, + "variants": _variant_summary(compatible_variants), + "module_file": str(Path(kernel.__file__).resolve()), + "api": {name: callable(getattr(kernel, name, None)) for name in API_NAMES}, + } + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/debug/probe_transformer_engine.py b/tools/debug/probe_transformer_engine.py new file mode 100644 index 0000000..de7af3f --- /dev/null +++ b/tools/debug/probe_transformer_engine.py @@ -0,0 +1,52 @@ +"""Fail-closed Transformer Engine import and FP8 capability probe.""" + +from __future__ import annotations + +import json +import platform +import torch +from importlib.metadata import PackageNotFoundError, version + + +def _package_version(name: str) -> str | None: + try: + return version(name) + except PackageNotFoundError: + return None + + +def main() -> int: + report: dict[str, object] = { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "transformer_engine": _package_version("transformer-engine"), + "transformer_engine_cu12": _package_version("transformer-engine-cu12"), + "transformer_engine_cu13": _package_version("transformer-engine-cu13"), + "transformer_engine_torch": _package_version("transformer-engine-torch"), + } + try: + import transformer_engine.pytorch as te + + try: + result = te.is_fp8_available(return_reason=True) + except TypeError: + result = te.is_fp8_available() + if isinstance(result, tuple): + available = bool(result[0]) + reason = str(result[1]) if len(result) > 1 else "" + else: + available = bool(result) + reason = "" + report.update(fp8_available=available, reason=reason) + except (ImportError, OSError, RuntimeError) as error: + report.update( + fp8_available=False, + reason=f"{type(error).__name__}: {error}", + ) + print(json.dumps(report, sort_keys=True)) + return 0 if report["fp8_available"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/debug/probe_transformer_engine_fp8.py b/tools/debug/probe_transformer_engine_fp8.py new file mode 100644 index 0000000..c4e7c06 --- /dev/null +++ b/tools/debug/probe_transformer_engine_fp8.py @@ -0,0 +1,51 @@ +"""Run one minimal Transformer Engine FP8 linear layer on CUDA.""" + +from __future__ import annotations + +import json +import torch +from importlib.metadata import version + + +def main() -> None: + import transformer_engine.pytorch as te + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the Transformer Engine FP8 probe.") + torch.manual_seed(17) + device = torch.device("cuda") + X = torch.randn(32, 16, device=device, dtype=torch.bfloat16) # (n=32, d=16) + linear = te.Linear( + 16, + 32, + bias=False, + params_dtype=torch.bfloat16, + device=device, + ).eval() + autocast = getattr(te, "autocast", None) + if autocast is None: + autocast = te.fp8_autocast + with torch.inference_mode(), autocast(enabled=True): + Z = linear(X) # (n=32, d_out=32) + torch.cuda.synchronize() + if Z.shape != (32, 32) or not torch.isfinite(Z).all(): + raise RuntimeError("Transformer Engine FP8 linear output is invalid.") + print( + json.dumps( + { + "cuda": torch.version.cuda, + "device": torch.cuda.get_device_name(device), + "input_dtype": str(X.dtype), + "output_dtype": str(Z.dtype), + "output_shape": list(Z.shape), + "torch": torch.__version__, + "transformer_engine": version("transformer-engine"), + "transformer_engine_torch": version("transformer-engine-torch"), + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/debug/trace_boltz2_modules.py b/tools/debug/trace_boltz2_modules.py new file mode 100644 index 0000000..5f25107 --- /dev/null +++ b/tools/debug/trace_boltz2_modules.py @@ -0,0 +1,225 @@ +"""Export bounded Boltz2 activation traces for official/local diagnostics.""" + +from __future__ import annotations + +import argparse +import gc +import json +import torch +from collections import defaultdict +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any, Literal +from safetensors.torch import load_file, save_file + +from tests.structure.support import boltz2_bundle + + +_MSA_LAYER_PATHS = tuple( + f"msa_module.layers.{layer_index}.{module_name}" + for layer_index in range(8) + for module_name in ( + "pair_weighted_averaging", + "msa_transition", + "outer_product_mean", + "outer_product_mean.norm", + "outer_product_mean.proj_a", + "outer_product_mean.proj_b", + "outer_product_mean.proj_o", + "pairformer_layer", + ) +) + +_MODULE_PATHS = ( + "input_embedder", + "s_init", + "z_init_1", + "z_init_2", + "rel_pos", + "token_bonds", + "contact_conditioning", + "s_norm", + "z_norm", + "s_recycle", + "z_recycle", + "msa_module.msa_proj", + "msa_module.s_proj", + "msa_module", + "pairformer_module", + "distogram_module", + "diffusion_conditioning", + "structure_module.score_model", + "confidence_module.s_inputs_norm", + "confidence_module.s_norm", + "confidence_module.z_norm", + "confidence_module.rel_pos", + "confidence_module.token_bonds", + "confidence_module.contact_conditioning", + "confidence_module.s_to_z", + "confidence_module.s_to_z_transpose", + "confidence_module.dist_bin_pairwise_embed", + "confidence_module.pairformer_stack", + "confidence_module.confidence_heads", + *_MSA_LAYER_PATHS, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--producer", choices=("reference", "candidate"), required=True) + parser.add_argument("--exchange-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--feature-bundle", + type=Path, + help="Override prepared inputs with feature tensors from this bundle.", + ) + return parser + + +def _tensor_leaves(value: Any, path: str = "value") -> dict[str, torch.Tensor]: + if torch.is_tensor(value): + return {path: value.detach().cpu().contiguous().clone()} # value: (...) + if isinstance(value, Mapping): + leaves: dict[str, torch.Tensor] = {} + for key in sorted(value, key=str): + leaves.update(_tensor_leaves(value[key], f"{path}.{key}")) + return leaves + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + leaves = {} + for index, item in enumerate(value): + leaves.update(_tensor_leaves(item, f"{path}.{index}")) + return leaves + return {} + + +def _register_trace_hooks( + core: torch.nn.Module, + traces: dict[str, torch.Tensor], +) -> list[torch.utils.hooks.RemovableHandle]: + call_counts: defaultdict[str, int] = defaultdict(int) + handles: list[torch.utils.hooks.RemovableHandle] = [] + + for module_path in _MODULE_PATHS: + try: + module = core.get_submodule(module_path) + except AttributeError: + continue + + for parameter_name, parameter in module.named_parameters(recurse=False): + # parameter: (...) + key = f"{module_path}__parameter__{parameter_name}".replace(".", "__") + if key not in traces: + traces[key] = parameter.detach().cpu().contiguous().clone() # (...) + for buffer_name, buffer in module.named_buffers(recurse=False): + # buffer: (...) + key = f"{module_path}__buffer__{buffer_name}".replace(".", "__") + if key not in traces: + traces[key] = buffer.detach().cpu().contiguous().clone() # (...) + + def hook( + _module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + path: str = module_path, + ) -> None: + call_index = call_counts[path] + call_counts[path] += 1 + prefix = f"{path}__call_{call_index:03d}" + values = { + **_tensor_leaves(args, "args"), + **_tensor_leaves(kwargs, "kwargs"), + **_tensor_leaves(output, "output"), + } # values[path]: (...) + for value_path, X in values.items(): + # X: (...) + key = f"{prefix}__{value_path}".replace(".", "__") + if key in traces: + raise RuntimeError(f"Duplicate trace tensor key: {key}") + traces[key] = X + + handles.append(module.register_forward_hook(hook, with_kwargs=True)) + return handles + + +def _load( + producer: Literal["reference", "candidate"], + request: Mapping[str, Any], +) -> tuple[torch.nn.Module, dict[str, torch.Tensor]]: + if producer == "reference": + archive = boltz2_bundle._download_official_file( + request, + boltz2_bundle._molecule_archive, + ) + checkpoint = boltz2_bundle._download_official_file(request, "boltz2_conf.ckpt") + molecule_dir = boltz2_bundle._extract_molecules(archive, str(request["sequence"])) + features = boltz2_bundle._prepare_reference_features( # values: (...) + request, molecule_dir + ) + model = boltz2_bundle._load_reference_model(request, checkpoint) + else: + features = boltz2_bundle._prepare_candidate_features(request) # values: (...) + model = boltz2_bundle._load_candidate_model(request) + return model, features + + +def main(argv: Sequence[str] | None = None) -> int: + """Run one deterministic forward and export the selected module trace.""" + + arguments = _parser().parse_args(argv) + producer: Literal["reference", "candidate"] = arguments.producer + request_path = ( + arguments.exchange_root + / "structure" + / "requests" + / boltz2_bundle.reference_container + / f"{boltz2_bundle.model_id}.json" + ) + request = boltz2_bundle.load_request(request_path) + model, features = _load(producer, request) # features values: (...) + if arguments.feature_bundle is not None: + stored = load_file(arguments.feature_bundle, device="cpu") # values: (...) + features = { + name.removeprefix("feature__"): X # (...) + for name, X in stored.items() + if name.startswith("feature__") + } # values: (...) + if set(features) != set(boltz2_bundle._feature_names): + raise RuntimeError("Feature override does not match the Boltz2 contract.") + core = model.core if hasattr(model, "core") else model + traces: dict[str, torch.Tensor] = {} + handles = _register_trace_hooks(core, traces) + try: + bundle = boltz2_bundle._run_model(model, features, request) # values: (...) + finally: + for handle in handles: + handle.remove() + del model + gc.collect() + torch.cuda.empty_cache() + + for name, X in bundle.items(): + # X: (...) + if name.startswith(("noise__", "output__")): + traces[f"bundle__{name}"] = X.detach().cpu().contiguous().clone() # (...) + arguments.output.parent.mkdir(parents=True, exist_ok=True) + save_file(dict(sorted(traces.items())), arguments.output) + metadata = { + "producer": producer, + "request_sha256": request["request_sha256"], + "environment": boltz2_bundle._environment_metadata(), + "tensor_count": len(traces), + "tensor_sha256": boltz2_bundle.tensor_set_sha256(traces), + } + arguments.output.with_suffix(".json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(arguments.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/goldens/__init__.py b/tools/goldens/__init__.py new file mode 100644 index 0000000..cee021e --- /dev/null +++ b/tools/goldens/__init__.py @@ -0,0 +1,36 @@ +"""Deterministic producer and read-only validator contracts for official goldens.""" + +from .bundle import ( + GoldenBundleRecord, + GoldenError, + require_declared_goldens, + validate_golden_bundle, + write_golden_bundle, +) +from .from_native import ( + GoldenMatrixEntry, + NativeGoldenRecord, + check_tier_specs, + convert_native_result, + detect_native_result_kind, + golden_generation_matrix, + missing_check_golden_ids, + require_complete_check_goldens, +) + + +__all__ = [ + "GoldenBundleRecord", + "GoldenError", + "GoldenMatrixEntry", + "NativeGoldenRecord", + "check_tier_specs", + "convert_native_result", + "detect_native_result_kind", + "golden_generation_matrix", + "missing_check_golden_ids", + "require_complete_check_goldens", + "require_declared_goldens", + "validate_golden_bundle", + "write_golden_bundle", +] diff --git a/tools/goldens/__main__.py b/tools/goldens/__main__.py new file mode 100644 index 0000000..86cd0ab --- /dev/null +++ b/tools/goldens/__main__.py @@ -0,0 +1,6 @@ +"""Command-line entry point for native official-golden conversion.""" + +from .from_native import main + + +raise SystemExit(main()) diff --git a/tools/goldens/bundle.py b/tools/goldens/bundle.py new file mode 100644 index 0000000..9278656 --- /dev/null +++ b/tools/goldens/bundle.py @@ -0,0 +1,464 @@ +"""Safetensors official-golden bundles with strict provenance metadata. + +This module does not load models or download checkpoints. The producer accepts +already generated tensors and records the immutable manifest provenance. The +validator is read-only and verifies every declared digest and tensor. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from fastplms.registry import ModelRegistry, ModelSpec, OfficialGolden + + +_SCHEMA_VERSION = 1 +_HEX = frozenset("0123456789abcdef") + + +class GoldenError(RuntimeError): + """Raised when an official golden cannot be produced or verified exactly.""" + + +@dataclass(frozen=True, slots=True) +class GoldenBundleRecord: + """Content identities returned by golden production or validation.""" + + metadata_sha256: str + tensors_sha256: str + tensor_hashes: Mapping[str, str] + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_json(value: object) -> bytes: + return json.dumps( + value, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def _is_sha256(value: object) -> bool: + return ( + isinstance(value, str) + and len(value) == 64 + and set(value).issubset(_HEX) + ) + + +def _tensor_hash(tensor: Any) -> str: + import torch + + if not torch.is_tensor(tensor) or tensor.layout != torch.strided: + raise GoldenError("Official golden entries must be strided Torch tensors.") + # Golden tensors may have any rank; their exact shapes are metadata. + T = tensor.detach().to(device="cpu").contiguous() # (...) + shape = list(T.shape) + dtype = str(T.dtype).removeprefix("torch.") + raw = T.reshape(-1).view(torch.uint8).numpy().tobytes() + digest = hashlib.sha256() + digest.update(_canonical_json({"dtype": dtype, "shape": shape})) + digest.update(b"\0") + digest.update(raw) + return digest.hexdigest() + + +def _normalize_tensors(tensors: Mapping[str, Any]) -> dict[str, Any]: + import torch + + if not tensors: + raise GoldenError("An official golden must contain at least one tensor.") + normalized: dict[str, Any] = {} + for name in sorted(tensors): + if not isinstance(name, str) or not name: + raise GoldenError(f"Invalid official golden tensor name: {name!r}.") + tensor = tensors[name] # (...) + if not torch.is_tensor(tensor) or tensor.layout != torch.strided: + raise GoldenError(f"Official golden entry {name!r} is not a strided tensor.") + normalized[name] = tensor.detach().to(device="cpu").contiguous().clone() # (...) + return normalized + + +def _validate_output_paths(metadata_path: Path, tensors_path: Path) -> None: + if metadata_path.suffix != ".json" or tensors_path.suffix != ".safetensors": + raise GoldenError("Official goldens require one .json and one .safetensors file.") + if metadata_path.parent.resolve() != tensors_path.parent.resolve(): + raise GoldenError("Official golden metadata and tensors must share one directory.") + if metadata_path.stem != tensors_path.stem: + raise GoldenError("Official golden metadata and tensor basenames must match.") + + +def _environment_record(environment: Mapping[str, str]) -> dict[str, object]: + if not environment or any( + not isinstance(key, str) + or not key + or not isinstance(value, str) + or not value + for key, value in environment.items() + ): + raise GoldenError("Environment details must be a non-empty string mapping.") + details = {key: environment[key] for key in sorted(environment)} + return { + "details": details, + "fingerprint": hashlib.sha256(_canonical_json(details)).hexdigest(), + } + + +def _source_file_record(source_files: Mapping[str, str] | None) -> dict[str, str]: + if source_files is None: + return {} + result: dict[str, str] = {} + if not isinstance(source_files, Mapping): + raise GoldenError("Official-golden source files must be a mapping.") + for name, digest in source_files.items(): + if not isinstance(name, str): + raise GoldenError(f"Invalid official-golden source file record: {name!r}.") + path = Path(name) + if ( + not name + or path.is_absolute() + or ".." in path.parts + or path.as_posix() != name + or not _is_sha256(digest) + ): + raise GoldenError(f"Invalid official-golden source file record: {name!r}.") + result[name] = digest + return dict(sorted(result.items())) + + +def _limitation_records( + limitations: Sequence[Mapping[str, str]] | None, +) -> list[dict[str, str]]: + """Validate portable, fail-closed official capability limitations.""" + + if limitations is None: + return [] + required = { + "capability", + "status", + "public_method", + "exception_type", + "reason", + } + result: list[dict[str, str]] = [] + for limitation in limitations: + if not isinstance(limitation, Mapping) or set(limitation) != required: + raise GoldenError("Official-golden capability limitation schema is invalid.") + record = {key: limitation[key] for key in sorted(required)} + if ( + record["capability"] != "generation" + or record["status"] != "official_unavailable" + or any(not isinstance(value, str) or not value for value in record.values()) + ): + raise GoldenError("Official-golden capability limitation is invalid.") + result.append(record) + capabilities = [record["capability"] for record in result] + if len(capabilities) != len(set(capabilities)): + raise GoldenError("Official-golden capability limitations contain duplicates.") + return sorted(result, key=lambda record: record["capability"]) + + +def _source_records(spec: ModelSpec, registry: ModelRegistry) -> list[dict[str, str]]: + return [ + { + "id": source_id, + "revision": registry.upstreams[source_id].revision, + "url": registry.upstreams[source_id].url, + } + for source_id in spec.family.upstreams + ] + + +def _checkpoint_record(spec: ModelSpec) -> dict[str, object]: + return { + "repo_id": spec.official.repo_id, + "revision": spec.official.revision, + "files": { + item.path: item.encoded + for item in sorted(spec.official.files, key=lambda value: value.path) + }, + } + + +def write_golden_bundle( + spec: ModelSpec, + registry: ModelRegistry, + tensors: Mapping[str, Any], + *, + metadata_path: Path, + tensors_path: Path, + generation_command: Sequence[str], + environment: Mapping[str, str], + input_fingerprint: str, + source_files: Mapping[str, str] | None = None, + limitations: Sequence[Mapping[str, str]] | None = None, + replace: bool = False, +) -> GoldenBundleRecord: + """Persist supplied official outputs without performing model generation.""" + + from safetensors.torch import save_file + + metadata_path = metadata_path.resolve() + tensors_path = tensors_path.resolve() + _validate_output_paths(metadata_path, tensors_path) + if isinstance(generation_command, (str, bytes)) or not generation_command or any( + not isinstance(argument, str) or not argument for argument in generation_command + ): + raise GoldenError("generation_command must be a non-empty string argument sequence.") + if not _is_sha256(input_fingerprint): + raise GoldenError("input_fingerprint must be a lowercase SHA-256 digest.") + if not replace and (metadata_path.exists() or tensors_path.exists()): + raise GoldenError("Official golden output already exists; pass replace=True explicitly.") + + normalized = _normalize_tensors(tensors) # values: (...) + tensor_metadata = { + name: { + "dtype": str(T.dtype).removeprefix("torch."), + "shape": list(T.shape), + "sha256": _tensor_hash(T), + } + for name, T in normalized.items() # T: (...) + } + metadata_path.parent.mkdir(parents=True, exist_ok=True) + temporary_tensors = tensors_path.with_name(f".{tensors_path.name}.{os.getpid()}.tmp") + temporary_metadata = metadata_path.with_name(f".{metadata_path.name}.{os.getpid()}.tmp") + try: + save_file( + normalized, + temporary_tensors, + # The JSON sidecar owns the schema. A single stable safetensors + # metadata entry avoids implementation-dependent map ordering. + metadata={"format": "pt"}, + ) + tensors_sha256 = _sha256_file(temporary_tensors) + metadata = { + "schema_version": _SCHEMA_VERSION, + "model_id": spec.id, + "sources": _source_records(spec, registry), + "checkpoint": _checkpoint_record(spec), + "environment": _environment_record(environment), + "generation_command": list(generation_command), + "input_fingerprint": input_fingerprint, + "source_files": _source_file_record(source_files), + "tensor_file": { + "path": tensors_path.name, + "sha256": tensors_sha256, + }, + "tensors": tensor_metadata, + } + normalized_limitations = _limitation_records(limitations) + if normalized_limitations: + metadata["limitations"] = normalized_limitations + temporary_metadata.write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + metadata_sha256 = _sha256_file(temporary_metadata) + os.replace(temporary_tensors, tensors_path) + os.replace(temporary_metadata, metadata_path) + except BaseException: + temporary_tensors.unlink(missing_ok=True) + temporary_metadata.unlink(missing_ok=True) + raise + return GoldenBundleRecord( + metadata_sha256=metadata_sha256, + tensors_sha256=tensors_sha256, + tensor_hashes={name: value["sha256"] for name, value in tensor_metadata.items()}, + ) + + +def _read_metadata(path: Path) -> dict[str, object]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise GoldenError(f"Unable to read official golden metadata: {path}.") from error + if not isinstance(value, dict): + raise GoldenError("Official golden metadata must contain a JSON object.") + expected = { + "schema_version", + "model_id", + "sources", + "checkpoint", + "environment", + "generation_command", + "input_fingerprint", + "source_files", + "tensor_file", + "tensors", + } + observed = frozenset(value) + if ( + observed not in {frozenset(expected), frozenset((*expected, "limitations"))} + or value.get("schema_version") != _SCHEMA_VERSION + ): + raise GoldenError("Official golden metadata schema is invalid.") + return value + + +def validate_golden_bundle( + spec: ModelSpec, + registry: ModelRegistry, + *, + metadata_path: Path, + tensors_path: Path, + declaration: OfficialGolden | None = None, +) -> GoldenBundleRecord: + """Read and exactly validate one golden bundle without changing any file.""" + + from safetensors.torch import load_file + + metadata_path = metadata_path.resolve() + tensors_path = tensors_path.resolve() + _validate_output_paths(metadata_path, tensors_path) + if not metadata_path.is_file() or not tensors_path.is_file(): + raise GoldenError(f"Missing required official golden for check tier: {spec.id}.") + metadata_sha256 = _sha256_file(metadata_path) + tensors_sha256 = _sha256_file(tensors_path) + if declaration is not None: + if metadata_sha256 != declaration.metadata.digest: + raise GoldenError(f"Official golden metadata digest mismatch for {spec.id}.") + if tensors_sha256 != declaration.tensors.digest: + raise GoldenError(f"Official golden tensor-file digest mismatch for {spec.id}.") + + metadata = _read_metadata(metadata_path) + if metadata["model_id"] != spec.id: + raise GoldenError(f"Official golden model identity mismatch for {spec.id}.") + if metadata["sources"] != _source_records(spec, registry): + raise GoldenError(f"Official golden source revisions mismatch for {spec.id}.") + if metadata["checkpoint"] != _checkpoint_record(spec): + raise GoldenError(f"Official golden checkpoint provenance mismatch for {spec.id}.") + if "limitations" in metadata: + try: + limitations = _limitation_records(metadata["limitations"]) + except GoldenError as error: + raise GoldenError( + f"Official golden capability limitations are invalid for {spec.id}." + ) from error + if metadata["limitations"] != limitations: + raise GoldenError( + f"Official golden capability limitations are invalid for {spec.id}." + ) + + environment = metadata["environment"] + if not isinstance(environment, dict) or set(environment) != {"details", "fingerprint"}: + raise GoldenError(f"Official golden environment record is invalid for {spec.id}.") + details = environment["details"] + if not isinstance(details, dict) or not details or any( + not isinstance(key, str) + or not key + or not isinstance(value, str) + or not value + for key, value in details.items() + ): + raise GoldenError(f"Official golden environment details are invalid for {spec.id}.") + expected_environment = hashlib.sha256(_canonical_json(details)).hexdigest() + if environment["fingerprint"] != expected_environment: + raise GoldenError(f"Official golden environment fingerprint mismatch for {spec.id}.") + + command = metadata["generation_command"] + if not isinstance(command, list) or not command or any( + not isinstance(argument, str) or not argument for argument in command + ): + raise GoldenError(f"Official golden generation command is invalid for {spec.id}.") + if not _is_sha256(metadata["input_fingerprint"]): + raise GoldenError(f"Official golden input fingerprint is invalid for {spec.id}.") + source_files = metadata["source_files"] + if not isinstance(source_files, dict): + raise GoldenError(f"Official golden source-file records are invalid for {spec.id}.") + try: + normalized_source_files = _source_file_record(source_files) + except GoldenError as error: + raise GoldenError( + f"Official golden source-file records are invalid for {spec.id}." + ) from error + if normalized_source_files != source_files: + raise GoldenError(f"Official golden source-file records are invalid for {spec.id}.") + tensor_file = metadata["tensor_file"] + if ( + not isinstance(tensor_file, dict) + or set(tensor_file) != {"path", "sha256"} + or tensor_file["path"] != tensors_path.name + or tensor_file["sha256"] != tensors_sha256 + ): + raise GoldenError(f"Official golden tensor-file record mismatch for {spec.id}.") + + tensor_records = metadata["tensors"] + if not isinstance(tensor_records, dict) or not tensor_records: + raise GoldenError(f"Official golden tensor records are invalid for {spec.id}.") + try: + tensors = load_file(tensors_path, device="cpu") # values: (...) + except Exception as error: + raise GoldenError(f"Unable to load official golden tensors for {spec.id}.") from error + if set(tensors) != set(tensor_records): + raise GoldenError(f"Official golden tensor names mismatch for {spec.id}.") + tensor_hashes: dict[str, str] = {} + for name, T in tensors.items(): + # T: (...) + record = tensor_records[name] + if ( + not isinstance(record, dict) + or set(record) != {"dtype", "shape", "sha256"} + or record["dtype"] != str(T.dtype).removeprefix("torch.") + or record["shape"] != list(T.shape) + or not _is_sha256(record["sha256"]) + ): + raise GoldenError(f"Official golden tensor metadata mismatch for {spec.id}:{name}.") + actual_hash = _tensor_hash(T) + if record["sha256"] != actual_hash: + raise GoldenError(f"Official golden tensor hash mismatch for {spec.id}:{name}.") + tensor_hashes[name] = actual_hash + return GoldenBundleRecord( + metadata_sha256=metadata_sha256, + tensors_sha256=tensors_sha256, + tensor_hashes=tensor_hashes, + ) + + +def _declared_path(root: Path, relative: str) -> Path: + root = root.resolve() + path = (root / relative).resolve() + if not path.is_relative_to(root): + raise GoldenError(f"Official golden path escapes the repository: {relative!r}.") + return path + + +def require_declared_goldens( + root: Path, + registry: ModelRegistry, + *, + tier: str = "check", +) -> tuple[GoldenBundleRecord, ...]: + """Validate manifest-declared goldens only when running the check tier.""" + + if tier != "check": + return () + records: list[GoldenBundleRecord] = [] + for spec in registry.values(): + declaration = spec.official_golden + if declaration is None: + continue + records.append( + validate_golden_bundle( + spec, + registry, + metadata_path=_declared_path(root, declaration.metadata.path), + tensors_path=_declared_path(root, declaration.tensors.path), + declaration=declaration, + ) + ) + return tuple(records) diff --git a/tools/goldens/from_native.py b/tools/goldens/from_native.py new file mode 100644 index 0000000..2ce924d --- /dev/null +++ b/tools/goldens/from_native.py @@ -0,0 +1,704 @@ +"""Convert isolated official-reference results into compact golden bundles. + +The converter never loads a model or imports an upstream package. It accepts +only the normalized, hash-checkable interchange formats written by the native +reference services, verifies their identity against ``models.toml``, and keeps +the minimum tensors needed for a routine candidate regression. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import torch +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal +from safetensors.torch import load_file + +from fastplms.registry import ModelRegistry, ModelSpec, get_model_registry +from tools.goldens.bundle import GoldenBundleRecord, GoldenError, write_golden_bundle + + +NativeResultKind = Literal["sequence", "structure"] +_SEQUENCE_REQUIRED_TENSORS = frozenset( + {"residue_mask", "output__last_hidden_state"} +) +_MAX_GOLDEN_TENSOR_BYTES = 64 * 1024 * 1024 +_DPLM2_3B_GENERATION_LIMITATION = { + "status": "official_unavailable", + "public_method": "EsmForDPLM.generate", + "exception_type": "TypeError", + "reason": ( + "The checkpoint-selected EsmForDPLM sampler uses tokenizer.cls_token_id " + "as bos_id, but the pinned DPLM2 tokenizer defines no cls_token_id." + ), +} +_EXPECTED_GENERATION_LIMITATIONS = { + "dplm2_3b": _DPLM2_3B_GENERATION_LIMITATION, +} + + +@dataclass(frozen=True, slots=True) +class NativeGoldenRecord: + """Paths and immutable digests for one converted official result.""" + + model_id: str + metadata_path: Path + tensors_path: Path + bundle: GoldenBundleRecord + + def manifest_declaration(self, repository_root: Path) -> str: + """Return the exact TOML declaration for files under ``tests/goldens``.""" + + root = repository_root.resolve() + try: + metadata_path = self.metadata_path.resolve().relative_to(root).as_posix() + tensors_path = self.tensors_path.resolve().relative_to(root).as_posix() + except ValueError as error: + raise GoldenError("Golden output is outside the repository.") from error + expected_metadata = f"tests/goldens/{self.model_id}.json" + expected_tensors = f"tests/goldens/{self.model_id}.safetensors" + if (metadata_path, tensors_path) != (expected_metadata, expected_tensors): + raise GoldenError( + "Only validated outputs under tests/goldens can be declared in models.toml." + ) + metadata = f"{metadata_path}=sha256:{self.bundle.metadata_sha256}" + tensors = f"{tensors_path}=sha256:{self.bundle.tensors_sha256}" + return ( + "official_golden = { " + f'metadata = "{metadata}", tensors = "{tensors}"' + " }" + ) + + +@dataclass(frozen=True, slots=True) +class GoldenMatrixEntry: + """Manifest-derived paths and readiness state for one check-tier golden.""" + + model_id: str + family: str + kind: NativeResultKind + reference_container: str + request_path: Path + native_result_path: Path + native_ready: bool + metadata_path: Path + tensors_path: Path + converted_ready: bool + declared: bool + + def as_dict(self) -> dict[str, object]: + """Return a stable JSON representation for remote orchestration.""" + + return { + "converted_ready": self.converted_ready, + "declared": self.declared, + "family": self.family, + "kind": self.kind, + "metadata_path": self.metadata_path.as_posix(), + "model_id": self.model_id, + "native_ready": self.native_ready, + "native_result_path": self.native_result_path.as_posix(), + "reference_container": self.reference_container, + "request_path": self.request_path.as_posix(), + "tensors_path": self.tensors_path.as_posix(), + } + + +def check_tier_specs(registry: ModelRegistry) -> tuple[ModelSpec, ...]: + """Return every checkpoint whose family declares the check tier.""" + + return tuple(spec for spec in registry.values() if "check" in spec.family.test_tiers) + + +def _structure_result_paths(native_root: Path, spec: ModelSpec) -> tuple[Path, ...]: + """Return the canonical structure bundle, preferring BF16-compute leaves.""" + + root = native_root / "structure" / "results" / "reference" / spec.id + paths: list[Path] = [] + if (root / "metadata.json").is_file() and (root / "bundle.safetensors").is_file(): + paths.append(root) + if root.is_dir(): + paths.extend( + path.parent + for path in sorted(root.rglob("bundle.safetensors")) + if path.parent != root and (path.parent / "metadata.json").is_file() + ) + preferred = [path for path in paths if path.name in {"bf16", "bfloat16"}] + if len(preferred) == 1: + return tuple(preferred) + return tuple(paths) + + +def _matrix_native_result_path(native_root: Path, spec: ModelSpec) -> Path: + if spec.family.tokenizer_mode != "structure": + return native_root / "results" / spec.id + available = _structure_result_paths(native_root, spec) + if len(available) == 1: + return available[0] + return native_root / "structure" / "results" / "reference" / spec.id + + +def golden_generation_matrix( + registry: ModelRegistry, + native_root: Path, + output_root: Path, +) -> tuple[GoldenMatrixEntry, ...]: + """Describe the generation path for every check-tier manifest entry.""" + + entries: list[GoldenMatrixEntry] = [] + for spec in check_tier_specs(registry): + kind: NativeResultKind = ( + "structure" if spec.family.tokenizer_mode == "structure" else "sequence" + ) + if kind == "sequence": + request_path = ( + native_root + / "requests" + / spec.family.reference_container + / f"{spec.id}.json" + ) + native_result_path = _matrix_native_result_path(native_root, spec) + native_tensors_name = "bf16.safetensors" + else: + request_path = ( + native_root + / "structure" + / "requests" + / spec.family.reference_container + / f"{spec.id}.json" + ) + native_result_path = _matrix_native_result_path(native_root, spec) + native_tensors_name = "bundle.safetensors" + metadata_path = output_root / f"{spec.id}.json" + tensors_path = output_root / f"{spec.id}.safetensors" + entries.append( + GoldenMatrixEntry( + model_id=spec.id, + family=spec.family.id, + kind=kind, + reference_container=spec.family.reference_container, + request_path=request_path, + native_result_path=native_result_path, + native_ready=(native_result_path / "metadata.json").is_file() + and (native_result_path / native_tensors_name).is_file(), + metadata_path=metadata_path, + tensors_path=tensors_path, + converted_ready=metadata_path.is_file() and tensors_path.is_file(), + declared=spec.official_golden is not None, + ) + ) + return tuple(entries) + + +def missing_check_golden_ids(registry: ModelRegistry) -> tuple[str, ...]: + """Report undeclared check-tier goldens without treating partial work as complete.""" + + return tuple( + spec.id for spec in check_tier_specs(registry) if spec.official_golden is None + ) + + +def require_complete_check_goldens(registry: ModelRegistry) -> None: + """Fail with the exact undeclared list instead of accepting partial coverage.""" + + missing = missing_check_golden_ids(registry) + if missing: + raise GoldenError( + "Check-tier official goldens are incomplete: " + ", ".join(missing) + "." + ) + + +def _canonical_json(value: object) -> bytes: + return json.dumps( + value, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _tensor_bytes(T: torch.Tensor) -> bytes: + # T: (...) + return T.detach().to(device="cpu").contiguous().view(torch.uint8).numpy().tobytes() + + +def _raw_tensor_sha256(T: torch.Tensor) -> str: + # T: (...) + return hashlib.sha256(_tensor_bytes(T)).hexdigest() + + +def _tensor_set_fingerprint(tensors: Mapping[str, torch.Tensor]) -> str: + if not tensors: + raise GoldenError("A native golden input fingerprint requires tensors.") + digest = hashlib.sha256() + for name in sorted(tensors): + T = tensors[name].detach().to(device="cpu").contiguous() # (...) + digest.update( + _canonical_json( + {"dtype": str(T.dtype), "name": name, "shape": list(T.shape)} + ) + ) + digest.update(b"\0") + digest.update(_tensor_bytes(T)) + digest.update(b"\0") + return digest.hexdigest() + + +def _ensure_compact(tensors: Mapping[str, torch.Tensor], *, model_id: str) -> None: + # tensors[name]: (...) + size = sum(T.numel() * T.element_size() for T in tensors.values()) + if size > _MAX_GOLDEN_TENSOR_BYTES: + raise GoldenError( + f"{model_id}: compact golden tensors require {size} bytes, exceeding the " + f"{_MAX_GOLDEN_TENSOR_BYTES}-byte limit. Add an explicit deterministic " + "projection instead of committing a large fixture." + ) + + +def _read_metadata(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise GoldenError(f"Unable to read native result metadata: {path}.") from error + if not isinstance(value, dict): + raise GoldenError(f"Native result metadata must be a JSON object: {path}.") + return value + + +def _environment(raw: object, *, model_id: str) -> dict[str, str]: + if not isinstance(raw, Mapping) or not raw: + raise GoldenError( + f"{model_id}: native result has no environment record; regenerate it in the " + "pinned reference container." + ) + result: dict[str, str] = {} + for key, value in raw.items(): + if not isinstance(key, str) or not key: + raise GoldenError(f"{model_id}: native environment contains an invalid key.") + if isinstance(value, str): + normalized = value + elif value is None: + normalized = "null" + elif isinstance(value, (bool, int, float)): + normalized = json.dumps(value, separators=(",", ":")) + elif isinstance(value, (Mapping, list, tuple)): + normalized = _canonical_json(value).decode("ascii") + else: + raise GoldenError( + f"{model_id}: native environment value {key!r} is not serializable." + ) + if not normalized: + raise GoldenError(f"{model_id}: native environment value {key!r} is empty.") + result[key] = normalized + return result + + +def _expected_files(spec: ModelSpec) -> list[dict[str, str]]: + return [ + {"algorithm": item.algorithm, "digest": item.digest, "path": item.path} + for item in spec.official.files + ] + + +def _validate_identity(metadata: Mapping[str, Any], spec: ModelSpec) -> None: + if metadata.get("schema_version") != 1: + raise GoldenError(f"{spec.id}: unsupported native result schema.") + if metadata.get("model_id") != spec.id: + raise GoldenError(f"{spec.id}: native result model identity mismatch.") + + official = metadata.get("official") + if isinstance(official, Mapping): + repo_id = official.get("repo_id") + revision = official.get("revision") + files = official.get("files") + else: + repo_id = metadata.get("reference_repo_id") + revision = metadata.get("reference_revision") + files = metadata.get("reference_files") + if repo_id != spec.official.repo_id or revision != spec.official.revision: + raise GoldenError(f"{spec.id}: native official checkpoint identity mismatch.") + if files != _expected_files(spec): + raise GoldenError( + f"{spec.id}: native official checkpoint file identities mismatch; regenerate " + "the native result from the current manifest request." + ) + + +def _load_sequence_result( + result_dir: Path, + spec: ModelSpec, +) -> tuple[ + dict[str, torch.Tensor], + dict[str, str], + str, + tuple[dict[str, str], ...], +]: + metadata_path = result_dir / "metadata.json" + tensors_path = result_dir / "bf16.safetensors" + metadata = _read_metadata(metadata_path) + _validate_identity(metadata, spec) + if metadata.get("state_transform") != spec.family.state_transform: + raise GoldenError(f"{spec.id}: native state transform mismatch.") + expected_limitation = _EXPECTED_GENERATION_LIMITATIONS.get(spec.id) + observed_limitation = metadata.get("generation_limitation") + generation = metadata.get("generation") + if expected_limitation is None: + if observed_limitation is not None: + raise GoldenError( + f"{spec.id}: undeclared official generation limitation." + ) + limitations: tuple[dict[str, str], ...] = () + else: + if generation is not None: + raise GoldenError( + f"{spec.id}: native result cannot claim generation parity and an " + "official generation limitation." + ) + if observed_limitation != expected_limitation: + raise GoldenError( + f"{spec.id}: native official generation limitation mismatch." + ) + limitations = ( + {"capability": "generation", **expected_limitation}, + ) + if ( + spec.family.id in {"dplm", "dplm2"} + and expected_limitation is None + and not isinstance(generation, Mapping) + ): + raise GoldenError(f"{spec.id}: native result omits required generation parity.") + if not tensors_path.is_file(): + raise GoldenError(f"{spec.id}: native BF16 result is missing: {tensors_path}.") + try: + tensors = load_file(tensors_path, device="cpu") # values: (...) + except Exception as error: + raise GoldenError(f"{spec.id}: unable to load native BF16 tensors.") from error + precision_keys = metadata.get("precision_tensor_keys") + if not isinstance(precision_keys, Mapping) or precision_keys.get("bf16") != sorted(tensors): + raise GoldenError(f"{spec.id}: native BF16 tensor-key contract mismatch.") + + missing = sorted(_SEQUENCE_REQUIRED_TENSORS.difference(tensors)) + input_names = sorted(name for name in tensors if name.startswith("input__")) + if missing or not input_names: + raise GoldenError( + f"{spec.id}: native BF16 result omits required golden tensors: " + f"{missing or ['input__*']}." + ) + selected_names = [ + *input_names, + "residue_mask", + "output__last_hidden_state", + ] + if "output__logits" in tensors: + selected_names.append("output__logits") + selected = {name: tensors[name] for name in selected_names} # values: (...) + input_tensors = { + name: selected[name] # (...) + for name in (*input_names, "residue_mask") + } # values: (...) + return ( + selected, + _environment(metadata.get("environment"), model_id=spec.id), + _tensor_set_fingerprint(input_tensors), + limitations, + ) + + +def _load_structure_result( + result_dir: Path, + spec: ModelSpec, +) -> tuple[dict[str, torch.Tensor], dict[str, str], str]: + metadata_path = result_dir / "metadata.json" + tensors_path = result_dir / "bundle.safetensors" + metadata = _read_metadata(metadata_path) + _validate_identity(metadata, spec) + if metadata.get("producer") != "reference": + raise GoldenError(f"{spec.id}: only an official reference bundle can become a golden.") + request_sha256 = metadata.get("request_sha256") + if ( + not isinstance(request_sha256, str) + or len(request_sha256) != 64 + or any(character not in "0123456789abcdef" for character in request_sha256) + ): + raise GoldenError(f"{spec.id}: structure request fingerprint is invalid.") + if not tensors_path.is_file(): + raise GoldenError(f"{spec.id}: native structure bundle is missing: {tensors_path}.") + try: + tensors = load_file(tensors_path, device="cpu") # values: (...) + except Exception as error: + raise GoldenError(f"{spec.id}: unable to load native structure tensors.") from error + if metadata.get("tensor_keys") != sorted(tensors): + raise GoldenError(f"{spec.id}: native structure tensor-key contract mismatch.") + observed_hashes = { + name: _raw_tensor_sha256(T) # T: (...) + for name, T in sorted(tensors.items()) + } + if metadata.get("tensor_hashes") != observed_hashes: + raise GoldenError(f"{spec.id}: native structure tensor hash mismatch.") + if not any(name.startswith("output__") for name in tensors): + raise GoldenError(f"{spec.id}: native structure result contains no outputs.") + return ( + dict(tensors), + _environment(metadata.get("environment"), model_id=spec.id), + request_sha256, + ) + + +def detect_native_result_kind(result_dir: Path) -> NativeResultKind: + """Identify one normalized result directory from its immutable files.""" + + has_sequence = (result_dir / "bf16.safetensors").is_file() + has_structure = (result_dir / "bundle.safetensors").is_file() + if has_sequence == has_structure: + raise GoldenError( + f"Native result must contain exactly one BF16 or structure tensor file: {result_dir}." + ) + return "sequence" if has_sequence else "structure" + + +def convert_native_result( + spec: ModelSpec, + registry: ModelRegistry, + result_dir: Path, + output_root: Path, + *, + generation_command: Sequence[str], + replace: bool = False, +) -> NativeGoldenRecord: + """Validate and convert one isolated official output without model loading.""" + + result_dir = result_dir.resolve() + kind = detect_native_result_kind(result_dir) + if kind == "sequence": + if spec.family.tokenizer_mode == "structure": + raise GoldenError(f"{spec.id}: sequence native result used for a structure model.") + ( + tensors, + environment, + input_fingerprint, + limitations, + ) = _load_sequence_result(result_dir, spec) + else: + if spec.family.tokenizer_mode != "structure": + raise GoldenError(f"{spec.id}: structure native result used for a sequence model.") + tensors, environment, input_fingerprint = _load_structure_result( # values: (...) + result_dir, spec + ) + limitations = () + _ensure_compact(tensors, model_id=spec.id) + + native_tensor_name = ( + "native/bf16.safetensors" if kind == "sequence" else "native/bundle.safetensors" + ) + source_files = { + "native/metadata.json": _sha256_file(result_dir / "metadata.json"), + native_tensor_name: _sha256_file( + result_dir / ("bf16.safetensors" if kind == "sequence" else "bundle.safetensors") + ), + } + + output_root = output_root.resolve() + metadata_path = output_root / f"{spec.id}.json" + tensors_path = output_root / f"{spec.id}.safetensors" + bundle = write_golden_bundle( + spec, + registry, + tensors, + metadata_path=metadata_path, + tensors_path=tensors_path, + generation_command=generation_command, + environment=environment, + input_fingerprint=input_fingerprint, + source_files=source_files, + limitations=limitations, + replace=replace, + ) + return NativeGoldenRecord( + model_id=spec.id, + metadata_path=metadata_path, + tensors_path=tensors_path, + bundle=bundle, + ) + + +def _find_native_result(native_root: Path, spec: ModelSpec) -> Path: + sequence_path = native_root / "results" / spec.id + if spec.family.tokenizer_mode == "structure": + existing = list(_structure_result_paths(native_root, spec)) + detail = str( + native_root / "structure" / "results" / "reference" / spec.id + ) + else: + existing = [sequence_path] if sequence_path.is_dir() else [] + detail = str(sequence_path) + if len(existing) != 1: + raise GoldenError( + f"{spec.id}: expected exactly one normalized native result under: {detail}." + ) + return existing[0] + + +def _canonical_generation_command(spec: ModelSpec) -> tuple[str, ...]: + return ( + "python", + "-m", + "tools.goldens", + "--native-root", + "artifacts/reference", + "--output-root", + "tests/goldens", + "--model", + spec.id, + ) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--native-root", type=Path, default=Path("artifacts/reference")) + parser.add_argument("--output-root", type=Path, default=Path("tests/goldens")) + parser.add_argument("--model", action="append", dest="model_ids") + parser.add_argument( + "--native-result", + type=Path, + help="Explicit result directory; valid only with exactly one --model.", + ) + parser.add_argument("--replace", action="store_true") + parser.add_argument( + "--status-only", + action="store_true", + help="Inspect manifest declaration completeness without converting results.", + ) + parser.add_argument( + "--report-missing", + action="store_true", + help="Print all undeclared check-tier model IDs after conversion.", + ) + parser.add_argument( + "--report-matrix", + action="store_true", + help="Print the manifest-wide generation and readiness matrix as JSON.", + ) + parser.add_argument( + "--require-complete", + action="store_true", + help="Fail after conversion unless every check-tier checkpoint is declared.", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Convert selected results and print declarations without editing the manifest.""" + + args = _parser().parse_args(argv) + registry = get_model_registry() + if args.status_only: + if args.native_result is not None or args.model_ids or args.replace: + raise GoldenError( + "--status-only cannot be combined with conversion inputs or --replace." + ) + if args.report_missing: + print( + json.dumps( + {"undeclared_check_goldens": missing_check_golden_ids(registry)} + ) + ) + if args.report_matrix: + print( + json.dumps( + { + "check_golden_matrix": [ + entry.as_dict() + for entry in golden_generation_matrix( + registry, + args.native_root, + args.output_root, + ) + ] + }, + sort_keys=True, + ) + ) + if args.require_complete: + require_complete_check_goldens(registry) + return 0 + model_ids = args.model_ids or [spec.id for spec in check_tier_specs(registry)] + unknown = sorted(set(model_ids).difference(registry)) + if unknown: + raise GoldenError(f"Unknown model IDs: {unknown}.") + if args.native_result is not None and len(model_ids) != 1: + raise GoldenError("--native-result requires exactly one --model.") + + records: list[NativeGoldenRecord] = [] + for model_id in model_ids: + spec = registry[model_id] + result_dir = ( + args.native_result + if args.native_result is not None + else _find_native_result(args.native_root, spec) + ) + records.append( + convert_native_result( + spec, + registry, + result_dir, + args.output_root, + generation_command=_canonical_generation_command(spec), + replace=args.replace, + ) + ) + + repository_root = Path.cwd() + declarable = args.output_root.resolve() == (repository_root / "tests/goldens").resolve() + for record in records: + if declarable: + print(f"[{record.model_id}] {record.manifest_declaration(repository_root)}") + else: + print( + json.dumps( + { + "metadata_path": str(record.metadata_path), + "metadata_sha256": record.bundle.metadata_sha256, + "model_id": record.model_id, + "tensors_path": str(record.tensors_path), + "tensors_sha256": record.bundle.tensors_sha256, + }, + sort_keys=True, + ) + ) + if args.report_missing: + print(json.dumps({"undeclared_check_goldens": missing_check_golden_ids(registry)})) + if args.report_matrix: + print( + json.dumps( + { + "check_golden_matrix": [ + entry.as_dict() + for entry in golden_generation_matrix( + registry, + args.native_root, + args.output_root, + ) + ] + }, + sort_keys=True, + ) + ) + if args.require_complete: + require_complete_check_goldens(registry) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/goldens/generate_e1_sampling.py b/tools/goldens/generate_e1_sampling.py new file mode 100644 index 0000000..171d744 --- /dev/null +++ b/tools/goldens/generate_e1_sampling.py @@ -0,0 +1,121 @@ +"""Generate deterministic E1 MSA-sampling goldens from the pinned oracle. + +This producer is intended for the native E1 reference environment. Candidate +tests consume only its JSON output and never import the upstream package. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1024 * 1024): + digest.update(block) + return digest.hexdigest() + + +def _write_fixture(path: Path) -> None: + path.write_text( + ">query\nACDEFGHI\n>near\nACDEYGH-\n>gapped\nAC-EFGHI\n>mid\nTCD-FGHI\n>far\nTTTTTTTT\n", + encoding="utf-8", + ) + + +def _git_revision(upstream_root: Path) -> str: + result = subprocess.run( + ["git", "-C", str(upstream_root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def _generate(upstream_root: Path) -> dict[str, Any]: + sys.path.insert(0, str(upstream_root / "src")) + from E1.msa_sampling import ( # type: ignore[import-not-found] + ContextSpecification, + sample_context, + sample_multiple_contexts, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + msa_path = Path(temp_dir) / "parity.a3m" + _write_fixture(msa_path) + common = { + "msa_path": str(msa_path), + "max_num_samples": 3, + "max_token_length": 32, + "max_query_similarity": 0.99, + "min_query_similarity": 0.0, + "neighbor_similarity_lower_bound": 0.8, + "device": "cpu", + } + single = {str(seed): sample_context(seed=seed, **common) for seed in (0, 3, 11)} + multiple = sample_multiple_contexts( + msa_path=str(msa_path), + context_specifications=[ + ContextSpecification( + max_num_samples=3, + max_token_length=16, + max_query_similarity=0.99, + min_query_similarity=0.0, + neighbor_similarity_lower_bound=0.8, + ), + ContextSpecification( + max_num_samples=4, + max_token_length=32, + max_query_similarity=1.0, + min_query_similarity=0.2, + neighbor_similarity_lower_bound=0.8, + ), + ], + seed=7, + device="cpu", + ) + return {"single_context": single, "multiple_context": multiple} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--upstream-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + upstream_root = args.upstream_root.resolve() + source_path = upstream_root / "src" / "E1" / "msa_sampling.py" + payload = { + "provenance": { + "upstream_revision": _git_revision(upstream_root), + "source_path": "src/E1/msa_sampling.py", + "source_sha256": _sha256_file(source_path), + "generation_command": [ + "python", + "tools/goldens/generate_e1_sampling.py", + "--upstream-root", + str(args.upstream_root), + "--output", + str(args.output), + ], + }, + "goldens": _generate(upstream_root), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/tools/remote/README.md b/tools/remote/README.md new file mode 100644 index 0000000..be119c8 --- /dev/null +++ b/tools/remote/README.md @@ -0,0 +1,158 @@ +# Remote Hopper/SM90 runner + +The runner archives only Git-tracked files plus tracked files from initialized, +pinned upstream submodules, copies them into a unique remote directory, builds +the requested Docker targets, executes the suite, and retrieves `artifacts/`. +Untracked and ignored files, common credential names, private-key extensions, +caches, and `.git` metadata are excluded. + +Connection details are required at runtime and are never written into tracked +configuration: + +```bash +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/ssh-key \ + --accept-new-host-key \ + --suite check +``` + +Available suites include `check`, `gpu-golden-smoke`, `nightly`, `compliance`, +`structure`, `artifact`, `benchmark`, `benchmark-capture`, `release`, and +`python-matrix`, plus the focused `unit`, `integration`, and `feature` suites. +Remote source is removed after artifact retrieval unless `--keep-remote` is +passed. Persistent model and compiler caches are Docker volumes defined by +`docker/compose.yaml`, not part of the synchronized tree. + +Every invocation writes `remote-run.json` beside the retrieved outputs. The +report records the source-archive digest, exact suite command graph, pre-build +host-hardware binding, phase durations and timeouts, normalized Docker +cache telemetry, built image IDs, structured no-download kernel availability, +artifact-retrieval result, cleanup result, and a digest over the retrieved +artifact tree. Reports are written atomically into a unique run directory. They +never record the SSH destination, identity path, raw subprocess output, command +exception text, or secret values. + +The routine `check` suite builds only `candidate-structure`. It runs portable +units, imports, local integration and release checks, and compares candidates +with the checked-in sequence and structure goldens. It does not build local Hub +artifacts, download attention kernels, or build/import a live official +reference implementation. Artifact construction remains in `artifact`, +`nightly`, and `release`. +The repository does not use GitHub Actions. Run CPU, source, reference, and GPU +validation explicitly on the workstation before merge or release. + +`gpu-golden-smoke` is the conditional Hopper/SM90 tier. The current release run +uses the exact containerized Linux aarch64 environment on the configured GH200, +builds only the structure candidate superset, and compares sequence plus +structure candidates with the checked-in, hash-validated goldens. H100 and H200 +remain supported Hopper-class execution devices, but their results do not +substitute for this GH200 release evidence. +It never builds or imports an official reference implementation. Large cases +remain reserved for `nightly`. + +`nightly` builds candidate, structure, FP8, and artifact images together. It +exercises the complete checkpoint golden panels, the eager/SDPA/Flex GH200 +matrix, generation, TTT, PEFT, binder/structure flows, offline artifact loading, +FP8 reloads, and a descriptive family throughput report. It neither builds live +official references nor downloads/builds Flash kernels. FA2 remains separate +prior focused evidence; FA3 is explicitly unavailable in the current arm64 +lock. + +The `compliance` suite is the live-reference release-candidate tier. Its build, +reference, and test phases have explicit cancellation timeouts. It compares +every release-gated sequence checkpoint with its pinned official implementation +and runs the complete ESMFold and four-variant ESMFold2 folding gates, including +ESMFold2 FP8 validation. Boltz2 remains provisional and runs only in the focused +`structure`, `artifact`, and `benchmark` tiers. + +The Biohub oracle has a platform-specific, fully pinned and hash-attested GH200 +lock, including any source-built BioTraj wheel. Before source archiving or +Buildx, every remote suite records `uname -m`, normalized OCI architecture, GPU +name, UUID, driver, and total memory. Bake receives that exact native platform; +the runner rejects an image whose resolved platform or digest does not match the +preflight and also rejects hardware drift during the build. A GH200 therefore +runs `linux/arm64` images directly rather than emulated `linux/amd64` images. +This is exact GH200 evidence, not a claim that Docker erases ABI differences or +that an unvalidated architecture is interchangeable. + +Immediately after image inspection and before any reference command, the runner +writes `artifacts/reference/environment/container-images.json` (mounted as +`/exchange/environment/container-images.json`). Schema version 1 contains the +resolved platform, stable Docker server and Buildx identities, and each Bake +target's content digest, image ID, OS, and architecture. It excludes tags, +creation timestamps, hostnames, and other ephemeral fields. Biohub suites also +bind the `biohub-biotraj-wheel` builder image identity. + +The focused `structure` suite first produces isolated Meta ESMFold, Boltz2, and +Biohub ESMFold2 reference bundles, then produces candidate bundles from the same +immutable requests before running the metric gates. Reference containers contain +only their pinned upstream sources, the normalization protocol, and required +license notices. + +The `feature` suite uses the BF16 structure candidate. It does not install the +FP8 dependency profile because test-time training and all gradient-enabled paths +are required to remain BF16. + +`benchmark` is intentionally gated: it requires the tracked immutable +`benchmarks/baselines/h100.json` and fails before remote work if that baseline is +absent. The filename is a legacy automation identifier; the current release +baseline must record the exact GH200 model, Linux aarch64 architecture, and +environment, and regression comparison requires an exact match. No baseline is +synthesized by the runner. +`benchmark-capture` produces an ungated, descriptive candidate report +containing separate cold compilation, first-forward, warmup, and steady-state +measurements. Review that report before adding a baseline in a separate change. +Full release benchmark and ESMC evidence must retain the same preflight hardware +identity as the candidate and official-reference measurements. The GH200 lock +makes that same-host contract available natively on `linux/arm64`; evidence from +another architecture or GPU UUID is not substituted or combined. + +Every benchmark-producing invocation is self-contained. The focused benchmark +and capture suites use `tools.artifacts.build_all --benchmark-suite`; the +nightly throughput phase and aggregate `release` suite build their complete +artifact validation set. Each does so inside its own remote +workspace before invoking `benchmarks.suite --artifact-root dist/hub`. The +benchmark therefore loads registry-validated local artifacts and never assumes +that `dist/hub` from another GitHub matrix job or remote invocation is shared. +Official-source artifacts such as ANKH and DPLM2 are revalidated against the +current model registry after construction. The embedded nightly and aggregate +release reports remain descriptive; only the focused `benchmark` suite applies +the checked-in regression baseline. + +Remote archives never carry `.git` metadata. Before upload, the runner records +the clean tracked root inventory with portable modes, sizes, symlink targets, +and content digests, then verifies the uploaded archive SHA-256 before +extraction. Artifact construction independently validates that inventory and +rejects missing, extra, linked, sensitive, oversized, or mutated runtime-scope +files. Because an extracted manifest cannot authenticate a Git commit object, +Git-free builds use `source-tree-sha256:` as `runtime_revision`; the +outer remote report separately records the clean source HEAD and archive +SHA-256. Clean Git worktrees continue to use the exact Git revision directly. + +Run validation tiers directly through `python -m tools.remote` against the +GH200 Linux aarch64 workstation. Bind release evidence to the exact candidate +revision and keep only one accelerator-heavy suite active at a time. + +## Python source-support matrix + +Python 3.12 remains the canonical GPU validation environment. Before release, +the explicit remote `python-matrix` suite runs the non-canonical 3.11, 3.13, +and 3.14 members concurrently. It creates a separate CPU-only environment for +each interpreter and installs `requirements/profiles/runtime.in` with the +validation constraints: + +```bash +python -m tools.remote \ + --host user@gpu-host \ + --identity /path/to/ssh-key \ + --suite python-matrix +``` + +Each environment imports from the explicit repository source root. Its +offline, CPU-only smoke imports every advertised runtime class, compiles the +runtime source, parses the model registry, constructs a small ESM2 encoder, and +runs one finite forward. The suite writes `artifacts/python-matrix.json` and +`artifacts/junit/python-matrix.xml`. Raw installer output is represented only +by byte counts and SHA-256 digests, never copied into reports. A missing +interpreter or dependency wheel is a source-support failure, not a skip. diff --git a/tools/remote/__init__.py b/tools/remote/__init__.py new file mode 100644 index 0000000..08a1d91 --- /dev/null +++ b/tools/remote/__init__.py @@ -0,0 +1,6 @@ +"""Host-agnostic remote test orchestration.""" + +from .run import RemoteRunner, RunnerConfig + + +__all__ = ["RemoteRunner", "RunnerConfig"] diff --git a/tools/remote/__main__.py b/tools/remote/__main__.py new file mode 100644 index 0000000..d9c0a04 --- /dev/null +++ b/tools/remote/__main__.py @@ -0,0 +1,7 @@ +"""Run the remote test orchestrator.""" + +from .run import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/biohub_reference_environment.py b/tools/remote/biohub_reference_environment.py new file mode 100644 index 0000000..8be7341 --- /dev/null +++ b/tools/remote/biohub_reference_environment.py @@ -0,0 +1,400 @@ +"""Capture and validate the complete native Biohub reference environment.""" + +from __future__ import annotations + +import hashlib +import json +import platform +import re +import subprocess +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from tools.remote.biohub_reference_lock import ( + BiohubReferenceLockError, + expected_installed_inventory, + load_biohub_reference_lock_contract, + normalize_distribution_version, + verify_biohub_reference_lock_contract, + verify_current_installed_inventory, + verify_current_pip_check, +) + + +REFERENCE_ENVIRONMENT_SCHEMA_VERSION = 2 +BIOHUB_BUILD_TARGET = "biohub-biotraj-wheel" +BIOHUB_REFERENCE_TARGETS = frozenset({"reference-biohub-esm", "reference-esmfold2"}) +_IMAGE_ID = re.compile(r"sha256:[0-9a-f]{64}") +_SHA256 = re.compile(r"[0-9a-f]{64}") + + +class BiohubReferenceEnvironmentError(RuntimeError): + """The native runtime or persisted image identity differs from its lock.""" + + +def _sha256(path: Path) -> str: + try: + return hashlib.sha256(path.read_bytes()).hexdigest() + except OSError as error: + raise BiohubReferenceEnvironmentError(f"Unable to hash {path}.") from error + + +def _canonical_json_bytes(value: object) -> bytes: + return (json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n").encode() + + +def _digest_object(value: object) -> str: + return hashlib.sha256(_canonical_json_bytes(value)).hexdigest() + + +def _load_canonical_json(path: Path) -> tuple[dict[str, object], str]: + try: + serialized = path.read_bytes() + value: Any = json.loads(serialized.decode("utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: + raise BiohubReferenceEnvironmentError( + f"Unable to read reference container identity: {path}" + ) from error + if not isinstance(value, dict) or not all(isinstance(key, str) for key in value): + raise BiohubReferenceEnvironmentError("Reference container identity must be a JSON object.") + normalized = {str(key): item for key, item in value.items()} + if serialized != _canonical_json_bytes(normalized): + raise BiohubReferenceEnvironmentError("Reference container identity is not canonical JSON.") + return normalized, hashlib.sha256(serialized).hexdigest() + + +def _validated_image_identity(value: object, *, target: str) -> dict[str, str]: + fields = {"content_digest", "image_id", "os", "architecture", "resolved_platform"} + if not isinstance(value, Mapping) or set(value) != fields: + raise BiohubReferenceEnvironmentError( + f"Container image identity fields differ for {target!r}." + ) + digest = value["content_digest"] + if not isinstance(digest, str) or _IMAGE_ID.fullmatch(digest) is None: + raise BiohubReferenceEnvironmentError(f"Container image digest is invalid for {target!r}.") + if ( + value["image_id"] != digest + or value["os"] != "linux" + or value["architecture"] != "arm64" + or value["resolved_platform"] != "linux/arm64" + ): + raise BiohubReferenceEnvironmentError(f"Container image platform differs for {target!r}.") + return {field: str(value[field]) for field in sorted(fields)} + + +def _validated_container_identity(value: object) -> dict[str, object]: + fields = {"schema_version", "resolved_platform", "docker_server", "docker_buildx", "images"} + if not isinstance(value, Mapping) or set(value) != fields: + raise BiohubReferenceEnvironmentError("Reference container identity fields differ.") + if value["schema_version"] != 1 or value["resolved_platform"] != "linux/arm64": + raise BiohubReferenceEnvironmentError("Reference container platform is not linux/arm64.") + buildx = value["docker_buildx"] + if not isinstance(buildx, str) or not buildx.strip(): + raise BiohubReferenceEnvironmentError("Docker Buildx identity is missing.") + server = value["docker_server"] + allowed_server_fields = { + "Version", + "ApiVersion", + "MinAPIVersion", + "GitCommit", + "Os", + "Arch", + "KernelVersion", + } + required_server_fields = {"Version", "ApiVersion", "Os", "Arch"} + if ( + not isinstance(server, Mapping) + or not required_server_fields.issubset(server) + or not set(server).issubset(allowed_server_fields) + or server.get("Os") != "linux" + or server.get("Arch") not in {"arm64", "aarch64"} + or any(not isinstance(item, (str, int, float, bool)) for item in server.values()) + ): + raise BiohubReferenceEnvironmentError("Docker server identity is invalid.") + images = value["images"] + if not isinstance(images, Mapping) or not images: + raise BiohubReferenceEnvironmentError("Reference container image map is missing.") + normalized_images: dict[str, dict[str, str]] = {} + for raw_target, identity in images.items(): + if not isinstance(raw_target, str) or not raw_target: + raise BiohubReferenceEnvironmentError("Reference container target name is invalid.") + normalized_images[raw_target] = _validated_image_identity(identity, target=raw_target) + return { + "schema_version": 1, + "resolved_platform": "linux/arm64", + "docker_server": {str(key): server[key] for key in sorted(server)}, + "docker_buildx": buildx.strip(), + "images": {key: normalized_images[key] for key in sorted(normalized_images)}, + } + + +def _driver_version() -> str: + try: + completed = subprocess.run( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError) as error: + raise BiohubReferenceEnvironmentError( + "Biohub reference evidence requires the NVIDIA driver version." + ) from error + versions = {line.strip() for line in completed.stdout.splitlines() if line.strip()} + if len(versions) != 1: + raise BiohubReferenceEnvironmentError("NVIDIA driver version is ambiguous.") + return versions.pop() + + +def _runtime_identity(installed_inventory: Mapping[str, str]) -> dict[str, object]: + import torch + + if not torch.cuda.is_available(): + raise BiohubReferenceEnvironmentError("Biohub reference evidence requires CUDA.") + system = platform.system().lower() + machine = platform.machine() + implementation = platform.python_implementation() + python_version = platform.python_version() + torch_version = normalize_distribution_version("torch", torch.__version__) + cuda_runtime = str(torch.version.cuda or "") + properties = torch.cuda.get_device_properties(0) + gpu_name = properties.name + capability = list(torch.cuda.get_device_capability(0)) + if ( + system != "linux" + or machine != "aarch64" + or implementation != "CPython" + or not python_version.startswith("3.12.") + or torch_version != installed_inventory.get("torch") + or not cuda_runtime.startswith("13.0") + or gpu_name != "NVIDIA GH200 480GB" + or capability != [9, 0] + ): + raise BiohubReferenceEnvironmentError("Active runtime differs from the GH200 lock target.") + uname = platform.uname() + return { + "operating_system": system, + "architecture": machine, + "python_implementation": implementation, + "python_version": python_version, + "torch": torch_version, + "cuda_runtime": cuda_runtime, + "cuda_driver": _driver_version(), + "gpu": { + "name": gpu_name, + "capability": capability, + "total_memory_bytes": int(properties.total_memory), + }, + "uname": { + "system": uname.system, + "release": uname.release, + "version": uname.version, + "machine": uname.machine, + }, + } + + +def capture_biohub_reference_environment( + repository_root: Path, + contract_path: Path, + container_identity_path: Path, + *, + reference_target: str, +) -> dict[str, object]: + """Capture the locked dependency, image, hardware, and runtime identity.""" + + if reference_target not in BIOHUB_REFERENCE_TARGETS: + raise BiohubReferenceEnvironmentError( + f"Unsupported Biohub reference target: {reference_target!r}." + ) + try: + contract = load_biohub_reference_lock_contract(contract_path) + locks = verify_biohub_reference_lock_contract(repository_root, contract_path) + inventory = verify_current_installed_inventory( + repository_root, + contract_path, + profile="final", + ) + pip_check = verify_current_pip_check(repository_root, contract_path) + except BiohubReferenceLockError as error: + raise BiohubReferenceEnvironmentError(str(error)) from error + container_identity, manifest_sha256 = _load_canonical_json(container_identity_path) + normalized_container = _validated_container_identity(container_identity) + images = normalized_container["images"] + if not isinstance(images, Mapping): + raise BiohubReferenceEnvironmentError( + "Reference container identity images must be a mapping." + ) + missing_targets = {BIOHUB_BUILD_TARGET, reference_target}.difference(images) + if missing_targets: + raise BiohubReferenceEnvironmentError( + f"Reference container identity omits targets: {sorted(missing_targets)}." + ) + payload: dict[str, object] = { + "schema_version": REFERENCE_ENVIRONMENT_SCHEMA_VERSION, + "contract": contract.contract, + "contract_sha256": _sha256(contract_path), + "target": asdict(contract.target), + "build_container": asdict(contract.container), + "locks": locks, + "biotraj": asdict(contract.biotraj), + "installed_inventory": inventory, + "installed_inventory_sha256": _digest_object(inventory), + "pip_check": pip_check, + "reference_container_target": reference_target, + "container_identity": normalized_container, + "container_identity_sha256": manifest_sha256, + "runtime": _runtime_identity(inventory), + } + return validate_biohub_reference_environment_evidence( + payload, + repository_root=repository_root, + contract_path=contract_path, + ) + + +def validate_biohub_reference_environment_evidence( + value: object, + *, + repository_root: Path, + contract_path: Path, +) -> dict[str, object]: + """Validate portable Biohub environment evidence against checked-in locks.""" + + fields = { + "schema_version", + "contract", + "contract_sha256", + "target", + "build_container", + "locks", + "biotraj", + "installed_inventory", + "installed_inventory_sha256", + "pip_check", + "reference_container_target", + "container_identity", + "container_identity_sha256", + "runtime", + } + if not isinstance(value, Mapping) or set(value) != fields: + raise BiohubReferenceEnvironmentError("Biohub reference environment fields differ.") + if value["schema_version"] != REFERENCE_ENVIRONMENT_SCHEMA_VERSION: + raise BiohubReferenceEnvironmentError("Unsupported Biohub environment schema version.") + try: + contract = load_biohub_reference_lock_contract(contract_path) + expected_locks = verify_biohub_reference_lock_contract(repository_root, contract_path) + expected_inventory = expected_installed_inventory( + repository_root, + contract_path, + profile="final", + ) + except BiohubReferenceLockError as error: + raise BiohubReferenceEnvironmentError(str(error)) from error + expected_static = { + "contract": contract.contract, + "contract_sha256": _sha256(contract_path), + "target": asdict(contract.target), + "build_container": asdict(contract.container), + "locks": expected_locks, + "biotraj": asdict(contract.biotraj), + "installed_inventory": expected_inventory, + "installed_inventory_sha256": _digest_object(expected_inventory), + "pip_check": { + "status": "accepted-platform-exception", + "returncode": 1, + "diagnostics": [ + exception.accepted_diagnostic + for exception in contract.pip_check_platform_exceptions + ], + "accepted_platform_exceptions": [ + asdict(exception) for exception in contract.pip_check_platform_exceptions + ], + }, + } + for field, expected in expected_static.items(): + if value[field] != expected: + raise BiohubReferenceEnvironmentError( + f"Biohub reference environment {field} differs from the lock." + ) + target = value["reference_container_target"] + if not isinstance(target, str) or target not in BIOHUB_REFERENCE_TARGETS: + raise BiohubReferenceEnvironmentError("Biohub reference container target is invalid.") + container = _validated_container_identity(value["container_identity"]) + if value["container_identity"] != container: + raise BiohubReferenceEnvironmentError("Container identity is not canonically ordered.") + container_digest = value["container_identity_sha256"] + if ( + not isinstance(container_digest, str) + or _SHA256.fullmatch(container_digest) is None + or container_digest != _digest_object(container) + ): + raise BiohubReferenceEnvironmentError("Container identity digest differs.") + images = container["images"] + if not isinstance(images, Mapping): + raise BiohubReferenceEnvironmentError( + "Reference container identity images must be a mapping." + ) + if BIOHUB_BUILD_TARGET not in images or target not in images: + raise BiohubReferenceEnvironmentError("Required Biohub container image is absent.") + runtime = value["runtime"] + runtime_fields = { + "operating_system", + "architecture", + "python_implementation", + "python_version", + "torch", + "cuda_runtime", + "cuda_driver", + "gpu", + "uname", + } + if not isinstance(runtime, Mapping) or set(runtime) != runtime_fields: + raise BiohubReferenceEnvironmentError("Biohub runtime identity fields differ.") + if ( + runtime["operating_system"] != "linux" + or runtime["architecture"] != "aarch64" + or runtime["python_implementation"] != "CPython" + or not isinstance(runtime["python_version"], str) + or not runtime["python_version"].startswith("3.12.") + or runtime["torch"] != expected_inventory["torch"] + or not isinstance(runtime["cuda_runtime"], str) + or not runtime["cuda_runtime"].startswith("13.0") + or not isinstance(runtime["cuda_driver"], str) + or not runtime["cuda_driver"].strip() + ): + raise BiohubReferenceEnvironmentError("Biohub runtime differs from the target policy.") + gpu = runtime["gpu"] + if ( + not isinstance(gpu, Mapping) + or set(gpu) != {"name", "capability", "total_memory_bytes"} + or gpu["name"] != "NVIDIA GH200 480GB" + or gpu["capability"] != [9, 0] + or isinstance(gpu["total_memory_bytes"], bool) + or not isinstance(gpu["total_memory_bytes"], int) + or gpu["total_memory_bytes"] <= 0 + ): + raise BiohubReferenceEnvironmentError("Biohub GPU identity differs from GH200/SM90.") + uname = runtime["uname"] + if ( + not isinstance(uname, Mapping) + or set(uname) != {"system", "release", "version", "machine"} + or uname["system"] != "Linux" + or uname["machine"] != "aarch64" + or any(not isinstance(item, str) or not item for item in uname.values()) + ): + raise BiohubReferenceEnvironmentError("Biohub uname identity is invalid.") + return {field: value[field] for field in sorted(fields)} + + +__all__ = [ + "BIOHUB_BUILD_TARGET", + "BIOHUB_REFERENCE_TARGETS", + "REFERENCE_ENVIRONMENT_SCHEMA_VERSION", + "BiohubReferenceEnvironmentError", + "capture_biohub_reference_environment", + "validate_biohub_reference_environment_evidence", +] diff --git a/tools/remote/biohub_reference_lock.py b/tools/remote/biohub_reference_lock.py new file mode 100644 index 0000000..c8f3611 --- /dev/null +++ b/tools/remote/biohub_reference_lock.py @@ -0,0 +1,1075 @@ +"""Validate and attest the native GH200 Biohub reference dependency lock.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import os +import platform +import re +import subprocess +import sys +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path, PurePosixPath +from typing import Any +from urllib.parse import urlsplit + + +_SCHEMA_VERSION = 2 +_EVIDENCE_SCHEMA_VERSION = 1 +_DIGEST = re.compile(r"[0-9a-f]{64}") +_IMAGE_ID = re.compile(r"sha256:[0-9a-f]{64}") +_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") +_PINNED = re.compile(r"(?P[A-Za-z0-9][A-Za-z0-9._-]*)==(?P[^\s;]+)") +_HASH_OPTION = re.compile(r"(?:^|\s)--hash=sha256:([0-9a-f]{64})(?=\s|$)") +_RUNTIME_HEADER = ( + "# This file was autogenerated by uv via the following command:", + "# python -m tools.remote.biohub_reference_lock generate", +) +_BUILD_HEADER = ( + "# This file was autogenerated by uv via the following command:", + "# python -m tools.remote.biohub_reference_lock generate-build-lock", +) +_RUNTIME_DIRECTIVES = ( + "--index-url https://pypi.org/simple", + "--extra-index-url https://download.pytorch.org/whl/cu130", + "--no-binary biotraj", + "--only-binary :all:", +) +_MATERIALIZED_DIRECTIVES = ( + "--index-url https://pypi.org/simple", + "--extra-index-url https://download.pytorch.org/whl/cu130", + "--only-binary :all:", +) +_BUILD_DIRECTIVES = ( + "--index-url https://pypi.org/simple", + "--only-binary :all:", +) +_BUILD_INVENTORY = { + "cython": "3.2.4", + "numpy": "2.4.4", + "packaging": "26.2", + "pip": "26.1.1", + "setuptools": "83.0.0", + "setuptools-scm": "9.2.2", + "wheel": "0.47.0", +} +_RUNTIME_SENTINELS = { + "numpy": "1.26.4", + "pytest": "9.0.2", + "torch": "2.13.0+cu130", + "wheel": "0.47.0", + "zstd": "1.5.6.1", +} + + +class BiohubReferenceLockError(RuntimeError): + """The checked-in Biohub reference lock or runtime differs from its contract.""" + + +@dataclass(frozen=True) +class TargetPolicy: + """The one supported execution target for this reference lock.""" + + hardware: str + operating_system: str + architecture: str + container_platform: str + python_implementation: str + python_version: str + cuda_version: str + torch_backend: str + + +@dataclass(frozen=True) +class ContainerPolicy: + """Immutable build recipe and base-image identities.""" + + dockerfile_path: str + dockerfile_sha256: str + dockerfile_frontend: str + cuda_image: str + cuda_image_digest: str + python_image: str + python_image_digest: str + build_python_version: str + + +@dataclass(frozen=True) +class LockFilePolicy: + """One canonical input and its generated hash lock.""" + + input_path: str + input_sha256: str + lock_path: str + lock_sha256: str + package_count: int + + +@dataclass(frozen=True) +class BioTrajPolicy: + """Pinned BioTraj source and reproducibly built native wheel.""" + + version: str + sdist_url: str + sdist_sha256: str + wheel_filename: str + wheel_sha256: str + wheel_size: int + + +@dataclass(frozen=True) +class PipCheckPlatformException: + """One exact upstream wheel-tag defect accepted by the locked runtime.""" + + distribution: str + version: str + wheel_filename: str + wheel_sha256: str + filename_platform_tag: str + wheel_metadata_platform_tag: str + target_hardware: str + target_operating_system: str + target_architecture: str + accepted_diagnostic: str + resolution: str + + +@dataclass(frozen=True) +class BiohubReferenceLockContract: + """Strict schema for the native Biohub dependency evidence boundary.""" + + schema_version: int + contract: str + target: TargetPolicy + container: ContainerPolicy + runtime: LockFilePolicy + build: LockFilePolicy + biotraj: BioTrajPolicy + bootstrap_inventory: Mapping[str, str] + final_inventory_overlays: Mapping[str, str] + pip_check_platform_exceptions: tuple[PipCheckPlatformException, ...] + + +@dataclass(frozen=True) +class ParsedLock: + """Normalized inventory and hashes parsed from one requirements lock.""" + + inventory: Mapping[str, str] + hashes: Mapping[str, tuple[str, ...]] + directives: tuple[str, ...] + + +def canonical_distribution_name(value: str) -> str: + """Return the PEP 503 spelling used for inventory comparisons.""" + + return re.sub(r"[-_.]+", "-", value).lower() + + +def normalize_distribution_version(name: str, value: str) -> str: + """Normalize metadata spelling while retaining Torch's CUDA local version.""" + + normalized = value.strip().lower().replace("_", ".").replace("-", ".") + if canonical_distribution_name(name) != "torch": + return normalized + match = re.fullmatch(r"(\d+(?:\.\d+){2})(?:\+([a-z0-9.]+))?", normalized) + if match is None: + raise BiohubReferenceLockError(f"Invalid Torch version in inventory: {value!r}") + local = match.group(2) + if local in {"cu13.0", "cu130"}: + local = "cu130" + return match.group(1) + (f"+{local}" if local else "") + + +def _load_json(path: Path) -> dict[str, object]: + try: + raw: Any = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: + raise BiohubReferenceLockError(f"Unable to read Biohub lock contract: {path}") from error + if not isinstance(raw, dict) or not all(isinstance(key, str) for key in raw): + raise BiohubReferenceLockError("Biohub lock contract must be a JSON object.") + return {str(key): value for key, value in raw.items()} + + +def _exact_fields(raw: Mapping[str, object], expected: set[str], *, context: str) -> None: + if set(raw) == expected: + return + raise BiohubReferenceLockError( + f"{context} fields differ: missing={sorted(expected.difference(raw))}, " + f"extra={sorted(set(raw).difference(expected))}." + ) + + +def _mapping(value: object, *, field: str) -> dict[str, object]: + if not isinstance(value, Mapping) or not all(isinstance(key, str) for key in value): + raise BiohubReferenceLockError(f"{field} must be a JSON object.") + return {str(key): item for key, item in value.items()} + + +def _string(value: object, *, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise BiohubReferenceLockError(f"{field} must be a non-empty string.") + return value + + +def _digest(value: object, *, field: str) -> str: + digest = _string(value, field=field) + if _DIGEST.fullmatch(digest) is None: + raise BiohubReferenceLockError(f"{field} must be 64 lowercase hexadecimal characters.") + return digest + + +def _positive_int(value: object, *, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise BiohubReferenceLockError(f"{field} must be a positive integer.") + return value + + +def _portable_path(value: object, *, field: str) -> str: + rendered = _string(value, field=field) + path = PurePosixPath(rendered) + if path.is_absolute() or ".." in path.parts or path.as_posix() != rendered: + raise BiohubReferenceLockError(f"{field} must be a portable relative path.") + return rendered + + +def _parse_target(value: object) -> TargetPolicy: + raw = _mapping(value, field="target") + fields = set(TargetPolicy.__dataclass_fields__) + _exact_fields(raw, fields, context="target") + target = TargetPolicy( + **{field: _string(raw[field], field=f"target.{field}") for field in fields} + ) + required = TargetPolicy( + hardware="NVIDIA GH200 480GB", + operating_system="linux", + architecture="aarch64", + container_platform="linux/arm64", + python_implementation="CPython", + python_version="3.12", + cuda_version="13.0", + torch_backend="cu130", + ) + if target != required: + raise BiohubReferenceLockError(f"Unsupported Biohub lock target: {target!r}") + return target + + +def _parse_container(value: object) -> ContainerPolicy: + raw = _mapping(value, field="container") + fields = set(ContainerPolicy.__dataclass_fields__) + _exact_fields(raw, fields, context="container") + policy = ContainerPolicy( + dockerfile_path=_portable_path(raw["dockerfile_path"], field="container.dockerfile_path"), + dockerfile_sha256=_digest(raw["dockerfile_sha256"], field="container.dockerfile_sha256"), + dockerfile_frontend=_string( + raw["dockerfile_frontend"], field="container.dockerfile_frontend" + ), + cuda_image=_string(raw["cuda_image"], field="container.cuda_image"), + cuda_image_digest=_digest(raw["cuda_image_digest"], field="container.cuda_image_digest"), + python_image=_string(raw["python_image"], field="container.python_image"), + python_image_digest=_digest( + raw["python_image_digest"], field="container.python_image_digest" + ), + build_python_version=_string( + raw["build_python_version"], field="container.build_python_version" + ), + ) + expected_frontend = ( + "docker/dockerfile:1.19@sha256:" + "b6afd42430b15f2d2a4c5a02b919e98a525b785b1aaff16747d2f623364e39b6" + ) + if policy.dockerfile_path != "docker/biohub-reference-lock.Dockerfile": + raise BiohubReferenceLockError("Unexpected Biohub lock Dockerfile path.") + if policy.dockerfile_frontend != expected_frontend: + raise BiohubReferenceLockError("Unexpected Biohub lock Dockerfile frontend.") + if not policy.cuda_image.endswith("@sha256:" + policy.cuda_image_digest): + raise BiohubReferenceLockError("CUDA image does not use its declared ARM64 child digest.") + if not policy.python_image.endswith("@sha256:" + policy.python_image_digest): + raise BiohubReferenceLockError("Python image does not use its declared ARM64 child digest.") + if policy.build_python_version != "3.12.11": + raise BiohubReferenceLockError("Unexpected BioTraj build Python patch version.") + return policy + + +def _parse_lock_policy(value: object, *, field: str) -> LockFilePolicy: + raw = _mapping(value, field=field) + fields = set(LockFilePolicy.__dataclass_fields__) + _exact_fields(raw, fields, context=field) + return LockFilePolicy( + input_path=_portable_path(raw["input_path"], field=f"{field}.input_path"), + input_sha256=_digest(raw["input_sha256"], field=f"{field}.input_sha256"), + lock_path=_portable_path(raw["lock_path"], field=f"{field}.lock_path"), + lock_sha256=_digest(raw["lock_sha256"], field=f"{field}.lock_sha256"), + package_count=_positive_int(raw["package_count"], field=f"{field}.package_count"), + ) + + +def _parse_biotraj(value: object) -> BioTrajPolicy: + raw = _mapping(value, field="biotraj") + fields = set(BioTrajPolicy.__dataclass_fields__) + _exact_fields(raw, fields, context="biotraj") + policy = BioTrajPolicy( + version=_string(raw["version"], field="biotraj.version"), + sdist_url=_string(raw["sdist_url"], field="biotraj.sdist_url"), + sdist_sha256=_digest(raw["sdist_sha256"], field="biotraj.sdist_sha256"), + wheel_filename=_string(raw["wheel_filename"], field="biotraj.wheel_filename"), + wheel_sha256=_digest(raw["wheel_sha256"], field="biotraj.wheel_sha256"), + wheel_size=_positive_int(raw["wheel_size"], field="biotraj.wheel_size"), + ) + expected_wheel = "biotraj-1.2.2-cp312-cp312-linux_aarch64.whl" + if policy.version != "1.2.2" or policy.wheel_filename != expected_wheel: + raise BiohubReferenceLockError("BioTraj version or native wheel tag differs.") + parsed = urlsplit(policy.sdist_url) + if parsed.scheme != "https" or parsed.hostname != "files.pythonhosted.org": + raise BiohubReferenceLockError("BioTraj sdist must use the pinned Pythonhost URL.") + if parsed.fragment != "sha256=" + policy.sdist_sha256: + raise BiohubReferenceLockError("BioTraj sdist URL fragment differs from its digest.") + return policy + + +def _parse_inventory(value: object, *, field: str) -> dict[str, str]: + raw = _mapping(value, field=field) + result: dict[str, str] = {} + for raw_name, raw_version in raw.items(): + if _NAME.fullmatch(raw_name) is None: + raise BiohubReferenceLockError(f"Invalid distribution name in {field}: {raw_name!r}") + name = canonical_distribution_name(raw_name) + if name != raw_name or name in result: + raise BiohubReferenceLockError(f"{field} names must be unique and canonical.") + result[name] = normalize_distribution_version( + name, _string(raw_version, field=f"{field}.{name}") + ) + return dict(sorted(result.items())) + + +def _parse_pip_check_platform_exceptions( + value: object, +) -> tuple[PipCheckPlatformException, ...]: + if not isinstance(value, list) or len(value) != 1: + raise BiohubReferenceLockError( + "pip_check_platform_exceptions must contain the one declared vendor defect." + ) + raw = _mapping(value[0], field="pip_check_platform_exceptions[0]") + fields = set(PipCheckPlatformException.__dataclass_fields__) + _exact_fields(raw, fields, context="pip_check_platform_exceptions[0]") + policy = PipCheckPlatformException( + distribution=_string( + raw["distribution"], field="pip_check_platform_exceptions[0].distribution" + ), + version=_string(raw["version"], field="pip_check_platform_exceptions[0].version"), + wheel_filename=_string( + raw["wheel_filename"], field="pip_check_platform_exceptions[0].wheel_filename" + ), + wheel_sha256=_digest( + raw["wheel_sha256"], field="pip_check_platform_exceptions[0].wheel_sha256" + ), + filename_platform_tag=_string( + raw["filename_platform_tag"], + field="pip_check_platform_exceptions[0].filename_platform_tag", + ), + wheel_metadata_platform_tag=_string( + raw["wheel_metadata_platform_tag"], + field="pip_check_platform_exceptions[0].wheel_metadata_platform_tag", + ), + target_hardware=_string( + raw["target_hardware"], + field="pip_check_platform_exceptions[0].target_hardware", + ), + target_operating_system=_string( + raw["target_operating_system"], + field="pip_check_platform_exceptions[0].target_operating_system", + ), + target_architecture=_string( + raw["target_architecture"], + field="pip_check_platform_exceptions[0].target_architecture", + ), + accepted_diagnostic=_string( + raw["accepted_diagnostic"], + field="pip_check_platform_exceptions[0].accepted_diagnostic", + ), + resolution=_string(raw["resolution"], field="pip_check_platform_exceptions[0].resolution"), + ) + required = PipCheckPlatformException( + distribution="nvidia-cusparselt-cu13", + version="0.8.1", + wheel_filename=("nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl"), + wheel_sha256=("4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f"), + filename_platform_tag="py3-none-manylinux2014_aarch64", + wheel_metadata_platform_tag="py3-none-manylinux2014_sbsa", + target_hardware="NVIDIA GH200 480GB", + target_operating_system="linux", + target_architecture="aarch64", + accepted_diagnostic=("nvidia-cusparselt-cu13 0.8.1 is not supported on this platform"), + resolution="validated-vendor-metadata-exception-no-wheel-rewrite", + ) + if policy != required: + raise BiohubReferenceLockError("Unexpected pip-check platform exception policy.") + return (policy,) + + +def load_biohub_reference_lock_contract(path: Path) -> BiohubReferenceLockContract: + """Load the checked-in strict GH200 lock contract.""" + + raw = _load_json(path) + expected = { + "schema_version", + "contract", + "target", + "container", + "runtime", + "build", + "biotraj", + "bootstrap_inventory", + "final_inventory_overlays", + "pip_check_platform_exceptions", + } + _exact_fields(raw, expected, context="contract") + if raw["schema_version"] != _SCHEMA_VERSION: + raise BiohubReferenceLockError("Unsupported Biohub lock contract schema version.") + if raw["contract"] != "fastplms.biohub-reference-lock": + raise BiohubReferenceLockError("Unexpected Biohub lock contract identifier.") + bootstrap = _parse_inventory(raw["bootstrap_inventory"], field="bootstrap_inventory") + final_overlays = _parse_inventory( + raw["final_inventory_overlays"], field="final_inventory_overlays" + ) + if bootstrap != {"pip": "26.1.1"}: + raise BiohubReferenceLockError("Unexpected Biohub bootstrap inventory.") + expected_final = {"esm": "3.3.0", "transformers": "4.57.6", "uv": "0.10.12"} + if final_overlays != expected_final: + raise BiohubReferenceLockError("Unexpected Biohub final inventory overlays.") + contract = BiohubReferenceLockContract( + schema_version=_SCHEMA_VERSION, + contract="fastplms.biohub-reference-lock", + target=_parse_target(raw["target"]), + container=_parse_container(raw["container"]), + runtime=_parse_lock_policy(raw["runtime"], field="runtime"), + build=_parse_lock_policy(raw["build"], field="build"), + biotraj=_parse_biotraj(raw["biotraj"]), + bootstrap_inventory=bootstrap, + final_inventory_overlays=final_overlays, + pip_check_platform_exceptions=_parse_pip_check_platform_exceptions( + raw["pip_check_platform_exceptions"] + ), + ) + if contract.runtime.package_count != 108 or contract.build.package_count != 7: + raise BiohubReferenceLockError( + "Biohub lock package counts must be 108 runtime and 7 build." + ) + return contract + + +def _sha256(path: Path) -> str: + try: + return hashlib.sha256(path.read_bytes()).hexdigest() + except OSError as error: + raise BiohubReferenceLockError(f"Unable to hash required lock input: {path}") from error + + +def _contract_path(root: Path, relative: str) -> Path: + root = root.resolve() + path = root.joinpath(*PurePosixPath(relative).parts).resolve() + try: + path.relative_to(root) + except ValueError as error: + raise BiohubReferenceLockError( + f"Contract path escapes repository root: {relative}" + ) from error + return path + + +def _read_lock(path: Path) -> tuple[list[str], bytes]: + try: + serialized = path.read_bytes() + text = serialized.decode("utf-8") + except (OSError, UnicodeDecodeError) as error: + raise BiohubReferenceLockError(f"Unable to read requirements lock: {path}") from error + if not serialized.endswith(b"\n") or b"\r" in serialized: + raise BiohubReferenceLockError("Requirements lock must use canonical LF text ending in LF.") + return text.splitlines(), serialized + + +def _logical_entries(lines: Sequence[str]) -> tuple[tuple[str, ...], tuple[str, ...]]: + directives: list[str] = [] + entries: list[str] = [] + pending = "" + for line in lines: + if not line.strip(): + if pending: + raise BiohubReferenceLockError("Blank line interrupts a lock continuation.") + continue + if line.startswith("#"): + raise BiohubReferenceLockError("Unexpected comment in generated requirements body.") + if line.startswith("--") and not pending: + directives.append(line) + continue + continuation = line.rstrip().endswith("\\") + piece = line.rstrip()[:-1].strip() if continuation else line.strip() + pending = f"{pending} {piece}".strip() + if not continuation: + entries.append(pending) + pending = "" + if pending: + raise BiohubReferenceLockError("Requirements lock ends during a continuation.") + return tuple(directives), tuple(entries) + + +def _parse_hashed_lock( + path: Path, + *, + header: tuple[str, str], + directives: tuple[str, ...], + biotraj_url: str | None, + biotraj_digest: str | None, + biotraj_version: str, +) -> ParsedLock: + lines, _ = _read_lock(path) + if tuple(lines[:2]) != header: + raise BiohubReferenceLockError("Generated lock header differs from the stable command.") + observed_directives, entries = _logical_entries(lines[2:]) + if observed_directives != directives: + raise BiohubReferenceLockError( + f"Generated lock directives differ: {observed_directives!r}." + ) + inventory: dict[str, str] = {} + hashes_by_name: dict[str, tuple[str, ...]] = {} + observed_order: list[str] = [] + for entry in entries: + hashes = tuple(_HASH_OPTION.findall(entry)) + requirement = _HASH_OPTION.sub("", entry).strip() + pinned = _PINNED.fullmatch(requirement) + if pinned is not None: + raw_name = pinned.group("name") + version = pinned.group("version") + elif biotraj_url is not None and requirement == f"biotraj @ {biotraj_url}": + raw_name = "biotraj" + version = biotraj_version + if hashes != (biotraj_digest,): + raise BiohubReferenceLockError( + "BioTraj lock must contain only its selected digest." + ) + else: + raise BiohubReferenceLockError( + f"Unpinned or unsupported lock requirement: {requirement}" + ) + name = canonical_distribution_name(raw_name) + if raw_name != name or name in inventory: + raise BiohubReferenceLockError(f"Duplicate or non-canonical lock name: {raw_name!r}") + if not hashes or tuple(sorted(set(hashes))) != hashes: + raise BiohubReferenceLockError(f"Hashes for {name!r} must be unique and sorted.") + inventory[name] = normalize_distribution_version(name, version) + hashes_by_name[name] = hashes + observed_order.append(name) + if observed_order != sorted(observed_order): + raise BiohubReferenceLockError("Generated lock inventory is not sorted by canonical name.") + return ParsedLock( + inventory=dict(sorted(inventory.items())), + hashes=dict(sorted(hashes_by_name.items())), + directives=observed_directives, + ) + + +def verify_biohub_reference_lock_contract( + repository_root: Path, + contract_path: Path, +) -> dict[str, object]: + """Rehash every input and validate both generated dependency graphs.""" + + contract = load_biohub_reference_lock_contract(contract_path) + file_policies = (contract.runtime, contract.build) + for policy in file_policies: + for relative, expected in ( + (policy.input_path, policy.input_sha256), + (policy.lock_path, policy.lock_sha256), + ): + actual = _sha256(_contract_path(repository_root, relative)) + if actual != expected: + raise BiohubReferenceLockError( + f"Biohub lock file digest differs for {relative}: " + f"expected {expected}, received {actual}." + ) + dockerfile = _contract_path(repository_root, contract.container.dockerfile_path) + if _sha256(dockerfile) != contract.container.dockerfile_sha256: + raise BiohubReferenceLockError("Biohub reference lock Dockerfile digest differs.") + + runtime_lock = _parse_hashed_lock( + _contract_path(repository_root, contract.runtime.lock_path), + header=_RUNTIME_HEADER, + directives=_RUNTIME_DIRECTIVES, + biotraj_url=contract.biotraj.sdist_url, + biotraj_digest=contract.biotraj.sdist_sha256, + biotraj_version=contract.biotraj.version, + ) + build_lock = _parse_hashed_lock( + _contract_path(repository_root, contract.build.lock_path), + header=_BUILD_HEADER, + directives=_BUILD_DIRECTIVES, + biotraj_url=None, + biotraj_digest=None, + biotraj_version=contract.biotraj.version, + ) + if len(runtime_lock.inventory) != contract.runtime.package_count: + raise BiohubReferenceLockError("Runtime dependency graph package count differs.") + if len(build_lock.inventory) != contract.build.package_count: + raise BiohubReferenceLockError("BioTraj build graph package count differs.") + if dict(build_lock.inventory) != _BUILD_INVENTORY: + raise BiohubReferenceLockError("BioTraj PEP 517 build toolchain differs.") + for name, expected in _RUNTIME_SENTINELS.items(): + if runtime_lock.inventory.get(name) != expected: + raise BiohubReferenceLockError( + f"Runtime lock sentinel {name!r} differs from {expected!r}." + ) + if runtime_lock.inventory.get("biotraj") != contract.biotraj.version: + raise BiohubReferenceLockError("Runtime BioTraj version differs from source contract.") + for exception in contract.pip_check_platform_exceptions: + if runtime_lock.inventory.get(exception.distribution) != exception.version: + raise BiohubReferenceLockError( + "Pip-check exception distribution differs from the runtime lock." + ) + if exception.wheel_sha256 not in runtime_lock.hashes.get(exception.distribution, ()): + raise BiohubReferenceLockError( + "Pip-check exception wheel digest is absent from the runtime lock." + ) + forbidden = {"esm", "pip", "transformers", "uv"}.intersection(runtime_lock.inventory) + if forbidden: + raise BiohubReferenceLockError( + f"Source/bootstrap overlays leaked into runtime dependency graph: {sorted(forbidden)}" + ) + return { + "schema_version": _SCHEMA_VERSION, + "target": asdict(contract.target), + "runtime_lock_sha256": contract.runtime.lock_sha256, + "runtime_package_count": len(runtime_lock.inventory), + "build_lock_sha256": contract.build.lock_sha256, + "build_package_count": len(build_lock.inventory), + "biotraj_sdist_sha256": contract.biotraj.sdist_sha256, + "biotraj_wheel_sha256": contract.biotraj.wheel_sha256, + "pip_check_platform_exceptions": [ + asdict(exception) for exception in contract.pip_check_platform_exceptions + ], + } + + +def expected_installed_inventory( + repository_root: Path, + contract_path: Path, + *, + profile: str, +) -> dict[str, str]: + """Return the exact expected distributions for one build/runtime phase.""" + + contract = load_biohub_reference_lock_contract(contract_path) + verify_biohub_reference_lock_contract(repository_root, contract_path) + if profile == "build": + return dict(_BUILD_INVENTORY) + runtime_lock = _parse_hashed_lock( + _contract_path(repository_root, contract.runtime.lock_path), + header=_RUNTIME_HEADER, + directives=_RUNTIME_DIRECTIVES, + biotraj_url=contract.biotraj.sdist_url, + biotraj_digest=contract.biotraj.sdist_sha256, + biotraj_version=contract.biotraj.version, + ) + expected = dict(runtime_lock.inventory) + if profile not in {"runtime", "final"}: + raise BiohubReferenceLockError(f"Unsupported installed-inventory profile: {profile!r}") + for name, version in contract.bootstrap_inventory.items(): + if name in expected: + raise BiohubReferenceLockError(f"Bootstrap inventory collides with lock: {name}") + expected[name] = version + if profile == "final": + for name, version in contract.final_inventory_overlays.items(): + if name in expected: + raise BiohubReferenceLockError( + f"Final inventory overlay collides with lock: {name}" + ) + expected[name] = version + return dict(sorted(expected.items())) + + +def installed_distribution_inventory() -> dict[str, str]: + """Collect the current interpreter's exact normalized distribution inventory.""" + + inventory: dict[str, str] = {} + for distribution in importlib.metadata.distributions(): + raw_name = distribution.metadata.get("Name") + if not isinstance(raw_name, str) or not raw_name: + raise BiohubReferenceLockError("Installed distribution has no Name metadata.") + name = canonical_distribution_name(raw_name) + version = normalize_distribution_version(name, distribution.version) + if name in inventory: + raise BiohubReferenceLockError(f"Duplicate installed distribution metadata: {name}") + inventory[name] = version + return dict(sorted(inventory.items())) + + +def assert_exact_installed_inventory( + expected: Mapping[str, str], + observed: Mapping[str, str], +) -> dict[str, str]: + """Fail closed on a missing, extra, or version-divergent distribution.""" + + normalized_observed: dict[str, str] = {} + for raw_name, raw_version in observed.items(): + name = canonical_distribution_name(raw_name) + if name in normalized_observed: + raise BiohubReferenceLockError(f"Observed inventory contains duplicate name: {name}") + normalized_observed[name] = normalize_distribution_version(name, raw_version) + normalized_expected = { + canonical_distribution_name(name): normalize_distribution_version(name, version) + for name, version in expected.items() + } + missing = sorted(set(normalized_expected).difference(normalized_observed)) + extra = sorted(set(normalized_observed).difference(normalized_expected)) + changed = { + name: { + "expected": normalized_expected[name], + "observed": normalized_observed[name], + } + for name in sorted(set(normalized_expected).intersection(normalized_observed)) + if normalized_expected[name] != normalized_observed[name] + } + if missing or extra or changed: + raise BiohubReferenceLockError( + "Installed distribution inventory differs: " + f"missing={missing}, extra={extra}, changed={changed}." + ) + return dict(sorted(normalized_observed.items())) + + +def verify_current_installed_inventory( + repository_root: Path, + contract_path: Path, + *, + profile: str, +) -> dict[str, str]: + """Verify the current Python environment against one exact lock profile.""" + + expected = expected_installed_inventory(repository_root, contract_path, profile=profile) + return assert_exact_installed_inventory(expected, installed_distribution_inventory()) + + +def verify_current_pip_check( + repository_root: Path, + contract_path: Path, +) -> dict[str, object]: + """Run pip check while accepting only the attested NVIDIA SBSA tag defect.""" + + contract = load_biohub_reference_lock_contract(contract_path) + verify_biohub_reference_lock_contract(repository_root, contract_path) + verify_current_installed_inventory(repository_root, contract_path, profile="final") + if ( + platform.system().lower() != contract.target.operating_system + or platform.machine() != contract.target.architecture + ): + raise BiohubReferenceLockError( + "Pip-check platform exception is valid only on locked Linux/aarch64." + ) + exception = contract.pip_check_platform_exceptions[0] + if ( + exception.target_hardware != contract.target.hardware + or exception.target_operating_system != contract.target.operating_system + or exception.target_architecture != contract.target.architecture + ): + raise BiohubReferenceLockError( + "Pip-check platform exception target differs from the lock target." + ) + try: + distribution = importlib.metadata.distribution(exception.distribution) + wheel_metadata = distribution.read_text("WHEEL") + except (importlib.metadata.PackageNotFoundError, OSError) as error: + raise BiohubReferenceLockError( + "Unable to inspect the accepted pip-check exception distribution." + ) from error + if wheel_metadata is None: + raise BiohubReferenceLockError( + "Accepted pip-check exception distribution has no WHEEL metadata." + ) + wheel_tags = tuple( + line.removeprefix("Tag: ").strip() + for line in wheel_metadata.splitlines() + if line.startswith("Tag: ") + ) + if wheel_tags != (exception.wheel_metadata_platform_tag,): + raise BiohubReferenceLockError( + "NVIDIA cuSPARSELt WHEEL tag differs from the accepted vendor defect." + ) + try: + completed = subprocess.run( + [sys.executable, "-m", "pip", "check"], + check=False, + capture_output=True, + text=True, + timeout=60, + ) + except (OSError, subprocess.TimeoutExpired) as error: + raise BiohubReferenceLockError("Unable to execute pip check.") from error + diagnostics = tuple( + line.strip() + for stream in (completed.stdout, completed.stderr) + for line in stream.splitlines() + if line.strip() + ) + if completed.returncode != 1 or diagnostics != (exception.accepted_diagnostic,): + raise BiohubReferenceLockError( + "Pip check differs from the one accepted NVIDIA wheel-tag diagnostic: " + f"returncode={completed.returncode}, diagnostics={diagnostics!r}." + ) + return { + "status": "accepted-platform-exception", + "returncode": completed.returncode, + "diagnostics": list(diagnostics), + "accepted_platform_exceptions": [asdict(exception)], + } + + +def _atomic_write(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_bytes(payload) + temporary.replace(path) + + +def materialize_biotraj_wheel_lock( + repository_root: Path, + contract_path: Path, + wheel: Path, + output: Path, + *, + wheel_uri: str, +) -> ParsedLock: + """Replace the attested BioTraj sdist with its attested native wheel.""" + + contract = load_biohub_reference_lock_contract(contract_path) + verify_biohub_reference_lock_contract(repository_root, contract_path) + if wheel.name != contract.biotraj.wheel_filename: + raise BiohubReferenceLockError("BioTraj wheel filename differs from native target tag.") + if wheel.stat().st_size != contract.biotraj.wheel_size: + raise BiohubReferenceLockError("BioTraj wheel size differs from deterministic build.") + if _sha256(wheel) != contract.biotraj.wheel_sha256: + raise BiohubReferenceLockError("BioTraj wheel digest differs from deterministic build.") + parsed_uri = urlsplit(wheel_uri) + if ( + parsed_uri.scheme != "file" + or parsed_uri.netloc + or parsed_uri.query + or parsed_uri.fragment + or PurePosixPath(parsed_uri.path).name != contract.biotraj.wheel_filename + ): + raise BiohubReferenceLockError("Materialized BioTraj wheel URI is not an exact file URI.") + + lock_path = _contract_path(repository_root, contract.runtime.lock_path) + lines, _ = _read_lock(lock_path) + rendered: list[str] = [] + replaced = False + index = 0 + while index < len(lines): + line = lines[index] + if line == "--no-binary biotraj": + index += 1 + continue + if line.startswith("biotraj @ "): + if replaced: + raise BiohubReferenceLockError("Runtime lock contains duplicate BioTraj entries.") + while line.rstrip().endswith("\\"): + index += 1 + if index >= len(lines): + raise BiohubReferenceLockError("BioTraj lock entry ends during continuation.") + line = lines[index] + rendered.extend( + ( + f"biotraj @ {wheel_uri} \\", + f" --hash=sha256:{contract.biotraj.wheel_sha256}", + ) + ) + replaced = True + index += 1 + continue + rendered.append(line) + index += 1 + if not replaced: + raise BiohubReferenceLockError("Runtime lock contains no BioTraj source entry.") + _atomic_write(output, ("\n".join(rendered) + "\n").encode()) + materialized = _parse_hashed_lock( + output, + header=_RUNTIME_HEADER, + directives=_MATERIALIZED_DIRECTIVES, + biotraj_url=wheel_uri, + biotraj_digest=contract.biotraj.wheel_sha256, + biotraj_version=contract.biotraj.version, + ) + if len(materialized.inventory) != contract.runtime.package_count: + raise BiohubReferenceLockError("Materialized runtime lock package count differs.") + return materialized + + +def _first_command_line(command: Sequence[str]) -> str: + try: + completed = subprocess.run( + command, + check=True, + capture_output=True, + text=True, + timeout=15, + ) + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: + raise BiohubReferenceLockError( + f"Unable to attest toolchain command: {command!r}" + ) from error + lines = completed.stdout.splitlines() + if not lines: + raise BiohubReferenceLockError(f"Toolchain command emitted no version: {command!r}") + return lines[0] + + +def _gpu_name() -> str: + try: + completed = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + timeout=15, + ) + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: + raise BiohubReferenceLockError("Unable to attest the GH200 build host.") from error + names = tuple(line.strip() for line in completed.stdout.splitlines() if line.strip()) + if len(names) != 1: + raise BiohubReferenceLockError(f"Expected one GH200 device, received {names!r}.") + return names[0] + + +def write_biohub_reference_build_evidence( + repository_root: Path, + contract_path: Path, + wheel: Path, + output: Path, + *, + build_image_id: str, +) -> dict[str, object]: + """Write native image, toolchain, wheel, and build-inventory evidence.""" + + contract = load_biohub_reference_lock_contract(contract_path) + lock_evidence = verify_biohub_reference_lock_contract(repository_root, contract_path) + if _IMAGE_ID.fullmatch(build_image_id) is None: + raise BiohubReferenceLockError("Build image identity must be a sha256 Docker image ID.") + if platform.system().lower() != contract.target.operating_system: + raise BiohubReferenceLockError("Build operating system differs from lock target.") + if platform.machine() != contract.target.architecture: + raise BiohubReferenceLockError("Build architecture differs from lock target.") + if platform.python_implementation() != contract.target.python_implementation: + raise BiohubReferenceLockError("Build Python implementation differs from lock target.") + python_version = platform.python_version() + if python_version != contract.container.build_python_version: + raise BiohubReferenceLockError("Build Python patch version differs from pinned image.") + hardware = _gpu_name() + if hardware != contract.target.hardware: + raise BiohubReferenceLockError( + f"Build hardware differs: expected {contract.target.hardware!r}, received {hardware!r}." + ) + if wheel.name != contract.biotraj.wheel_filename: + raise BiohubReferenceLockError("Evidence wheel filename differs from contract.") + if wheel.stat().st_size != contract.biotraj.wheel_size: + raise BiohubReferenceLockError("Evidence wheel size differs from contract.") + if _sha256(wheel) != contract.biotraj.wheel_sha256: + raise BiohubReferenceLockError("Evidence wheel digest differs from contract.") + inventory = verify_current_installed_inventory(repository_root, contract_path, profile="build") + payload: dict[str, object] = { + "schema_version": _EVIDENCE_SCHEMA_VERSION, + "contract_sha256": _sha256(contract_path), + "target": asdict(contract.target), + "container": { + **asdict(contract.container), + "build_image_id": build_image_id, + }, + "observed": { + "hardware": hardware, + "operating_system": platform.system().lower(), + "architecture": platform.machine(), + "python_implementation": platform.python_implementation(), + "python_version": python_version, + "gcc": _first_command_line(("gcc", "--version")), + "binutils": _first_command_line(("ld", "--version")), + }, + "locks": lock_evidence, + "biotraj": { + **asdict(contract.biotraj), + "observed_wheel_sha256": _sha256(wheel), + "observed_wheel_size": wheel.stat().st_size, + }, + "installed_build_inventory": inventory, + } + _atomic_write(output, (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode()) + return payload + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + verify = subparsers.add_parser("verify-contract") + verify.add_argument("--root", type=Path, required=True) + verify.add_argument("--contract", type=Path, required=True) + inventory = subparsers.add_parser("verify-inventory") + inventory.add_argument("--root", type=Path, required=True) + inventory.add_argument("--contract", type=Path, required=True) + inventory.add_argument("--profile", choices=("build", "runtime", "final"), required=True) + pip_check = subparsers.add_parser("verify-pip-check") + pip_check.add_argument("--root", type=Path, required=True) + pip_check.add_argument("--contract", type=Path, required=True) + materialize = subparsers.add_parser("materialize-wheel-lock") + materialize.add_argument("--root", type=Path, required=True) + materialize.add_argument("--contract", type=Path, required=True) + materialize.add_argument("--wheel", type=Path, required=True) + materialize.add_argument("--wheel-uri", required=True) + materialize.add_argument("--output", type=Path, required=True) + evidence = subparsers.add_parser("write-build-evidence") + evidence.add_argument("--root", type=Path, required=True) + evidence.add_argument("--contract", type=Path, required=True) + evidence.add_argument("--wheel", type=Path, required=True) + evidence.add_argument( + "--build-image-id", + default=os.environ.get("FASTPLMS_BUILD_IMAGE_ID"), + required="FASTPLMS_BUILD_IMAGE_ID" not in os.environ, + ) + evidence.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Validate the lock, materialize its wheel, or write native build evidence.""" + + arguments = _parser().parse_args(argv) + if arguments.command == "verify-contract": + result: object = verify_biohub_reference_lock_contract(arguments.root, arguments.contract) + elif arguments.command == "verify-inventory": + result = verify_current_installed_inventory( + arguments.root, arguments.contract, profile=arguments.profile + ) + elif arguments.command == "verify-pip-check": + result = verify_current_pip_check(arguments.root, arguments.contract) + elif arguments.command == "materialize-wheel-lock": + parsed = materialize_biotraj_wheel_lock( + arguments.root, + arguments.contract, + arguments.wheel, + arguments.output, + wheel_uri=arguments.wheel_uri, + ) + result = {"package_count": len(parsed.inventory), "output": str(arguments.output)} + else: + result = write_biohub_reference_build_evidence( + arguments.root, + arguments.contract, + arguments.wheel, + arguments.output, + build_image_id=arguments.build_image_id, + ) + print(json.dumps(result, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/biohub_reference_requirements.py b/tools/remote/biohub_reference_requirements.py new file mode 100644 index 0000000..b990147 --- /dev/null +++ b/tools/remote/biohub_reference_requirements.py @@ -0,0 +1,112 @@ +"""Extract Biohub ESM dependencies without resolving its mutable Transformers URL.""" + +from __future__ import annotations + +import argparse +import re +import tomllib +from collections.abc import Sequence +from pathlib import Path + + +_PINNED_TRANSFORMERS_REQUIREMENT = ( + "transformers @ git+https://github.com/Biohub/transformers.git@main" +) +_SAFE_PEP508_SUBSET = re.compile( + r"[A-Za-z0-9][A-Za-z0-9._-]*" + r"(?:\[[A-Za-z0-9._-]+(?:,[A-Za-z0-9._-]+)*\])?" + r"(?:" + r"(?:===|==|!=|~=|<=|>=|<|>)[A-Za-z0-9][A-Za-z0-9.*+!_-]*" + r"(?:,(?:===|==|!=|~=|<=|>=|<|>)[A-Za-z0-9][A-Za-z0-9.*+!_-]*)*" + r")?" +) + + +class BiohubReferenceRequirementsError(RuntimeError): + """The pinned Biohub ESM dependency contract is missing or has drifted.""" + + +def _canonical_name(requirement: str) -> str: + name = re.split(r"[<>=!~ @\[]", requirement, maxsplit=1)[0] + return re.sub(r"[-_.]+", "-", name).lower() + + +def extract_biohub_reference_requirements(pyproject: Path) -> tuple[str, ...]: + """Return every pinned Biohub ESM dependency except its mutable Transformers URL.""" + + try: + raw = tomllib.loads(pyproject.read_text(encoding="utf-8")) + except (OSError, tomllib.TOMLDecodeError) as error: + raise BiohubReferenceRequirementsError( + f"Unable to read the pinned Biohub ESM pyproject: {pyproject}" + ) from error + project = raw.get("project") + dependencies = project.get("dependencies") if isinstance(project, dict) else None + if not isinstance(dependencies, list) or not all( + isinstance(requirement, str) and requirement.strip() for requirement in dependencies + ): + raise BiohubReferenceRequirementsError( + "Pinned Biohub ESM pyproject must declare a non-empty string dependency list." + ) + + transformer_requirements = [ + requirement + for requirement in dependencies + if _canonical_name(requirement) == "transformers" + ] + if transformer_requirements != [_PINNED_TRANSFORMERS_REQUIREMENT]: + raise BiohubReferenceRequirementsError( + "Pinned Biohub ESM must declare exactly its known mutable Transformers main URL; " + f"received {transformer_requirements!r}." + ) + + filtered = tuple( + requirement + for requirement in dependencies + if requirement != _PINNED_TRANSFORMERS_REQUIREMENT + ) + unsafe_requirements = [ + requirement + for requirement in filtered + if _SAFE_PEP508_SUBSET.fullmatch(requirement) is None + ] + if unsafe_requirements: + raise BiohubReferenceRequirementsError( + "Biohub ESM contains a dependency outside the allowed PEP 508 subset: " + f"{unsafe_requirements!r}." + ) + return filtered + + +def write_biohub_reference_requirements(pyproject: Path, output: Path) -> tuple[str, ...]: + """Write the filtered dependency set atomically for a reference-image build.""" + + requirements = extract_biohub_reference_requirements(pyproject) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text( + "".join(f"{requirement}\n" for requirement in requirements), + encoding="utf-8", + ) + temporary.replace(output) + return requirements + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pyproject", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Render the fail-closed Biohub ESM non-Transformers requirement file.""" + + arguments = _parser().parse_args(argv) + requirements = write_biohub_reference_requirements(arguments.pyproject, arguments.output) + print(f"Wrote {len(requirements)} pinned Biohub ESM non-Transformers requirements.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/prepare_references.py b/tools/remote/prepare_references.py new file mode 100644 index 0000000..5d11457 --- /dev/null +++ b/tools/remote/prepare_references.py @@ -0,0 +1,117 @@ +"""Write immutable native-reference requests from the typed model manifest.""" + +from __future__ import annotations + +import argparse +import json +import random +from dataclasses import asdict +from pathlib import Path + +from fastplms.registry import get_model_registry +from tests.parity.support.esmc_calibration import ( + CANONICAL_AA_ALPHABET, + ESMC_CALIBRATION_SEED, + esmc_calibration_batches, +) +from tests.parity.support.reference_adapters.dplm2 import ( + DPLM2_3B_GENERATION_LIMITATION, +) + + +SCHEMA_VERSION = 1 +SEED = ESMC_CALIBRATION_SEED +CANONICAL_AAS = CANONICAL_AA_ALPHABET +MIXED_LENGTHS = (61, 29, 13) +EDGE_SEQUENCES = ( + "ACDEFGHIKLMNPQRSTVWY", + "AXBJOUZ", + "acdefghik", + "A C\nD\tE", + "", +) +_TOKENIZER_FILE_NAMES = frozenset( + { + "added_tokens.json", + "merges.txt", + "sentencepiece.bpe.model", + "special_tokens_map.json", + "spiece.model", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "vocab.txt", + } +) + + +def _sequence_batch() -> tuple[str, ...]: + generator = random.Random(SEED) + return tuple( + "M" + "".join(generator.choices(CANONICAL_AAS, k=length - 1)) for length in MIXED_LENGTHS + ) + + +def _esmc_calibration_batches() -> list[dict[str, object]]: + return [dict(batch) for batch in esmc_calibration_batches()] + + +def prepare_reference_requests(output_root: Path) -> tuple[Path, ...]: + """Write one self-contained request for every sequence checkpoint.""" + + registry = get_model_registry() + paths: list[Path] = [] + for spec in registry.values(): + if spec.family.tokenizer_mode == "structure": + continue + request = { + "schema_version": SCHEMA_VERSION, + "model_id": spec.id, + "family": spec.family.id, + "architecture": spec.family.architecture, + "adapter": spec.family.reference_adapter, + "reference_container": spec.family.reference_container, + "reference_repo_id": spec.official.repo_id, + "reference_revision": spec.official.revision, + "reference_files": [asdict(item) for item in spec.official.files], + "state_transform": spec.family.state_transform, + "tokenizer_mode": spec.family.tokenizer_mode, + "attention_implementations": list(spec.family.attention), + "deep_reference": spec.is_deep_reference, + "oracle_assets": [asdict(asset) for asset in spec.oracle_assets], + "tokenizer_files": [ + item.path + for item in spec.official.files + if Path(item.path).name in _TOKENIZER_FILE_NAMES + ], + "sequences": list(_sequence_batch()), + "edge_sequences": list(EDGE_SEQUENCES), + "generation_policy": spec.generation_contract, + "seed": SEED, + } + if spec.family.id == "esm_plusplus": + request["calibration_batches"] = _esmc_calibration_batches() + if spec.generation_contract == "official_unavailable": + request["official_generation_limitation"] = dict( + DPLM2_3B_GENERATION_LIMITATION + ) + path = output_root / "requests" / spec.family.reference_container / f"{spec.id}.json" + path.parent.mkdir(parents=True, exist_ok=True) + encoded = json.dumps(request, indent=2, sort_keys=True) + "\n" + if path.exists() and path.read_text(encoding="utf-8") != encoded: + raise RuntimeError(f"Refusing to replace a different reference request: {path}") + path.write_text(encoded, encoding="utf-8") + paths.append(path) + return tuple(paths) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-root", type=Path, default=Path("artifacts/reference")) + arguments = parser.parse_args() + for path in prepare_reference_requests(arguments.output_root): + print(path) + + +if __name__ == "__main__": + main() diff --git a/tools/remote/python_matrix.py b/tools/remote/python_matrix.py new file mode 100644 index 0000000..2a19e39 --- /dev/null +++ b/tools/remote/python_matrix.py @@ -0,0 +1,335 @@ +"""Run the non-canonical FastPLMs source support matrix with uv.""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import hashlib +import json +import os +import shutil +import subprocess +import tempfile +import time +import xml.etree.ElementTree as ET +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + + +CANONICAL_GPU_PYTHON = "3.12" +PYTHON_SUPPORT_VERSIONS = ("3.11", "3.13", "3.14") +OFFLINE_SMOKE_ENVIRONMENT = { + "CUDA_VISIBLE_DEVICES": "", + "HF_DATASETS_OFFLINE": "1", + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "PYTHONNOUSERSITE": "1", + "PYTHONPATH": "", + "UV_TORCH_BACKEND": "cpu", +} + + +class MatrixCommandError(RuntimeError): + """One uv or smoke command failed for a matrix member.""" + + def __init__(self, stage: str, completed: subprocess.CompletedProcess[str]) -> None: + super().__init__(f"{stage} failed with exit code {completed.returncode}") + self.stage = stage + self.completed = completed + + +def build_dependency_install_command( + uv: str, + python: Path, + project_root: Path, +) -> tuple[str, ...]: + """Build the CPU dependency installation command for one source smoke.""" + + return ( + uv, + "pip", + "install", + "--python", + str(python), + "--torch-backend=cpu", + "-r", + str(project_root / "requirements/profiles/runtime.in"), + "-c", + str(project_root / "requirements/constraints/validation.txt"), + ) + + +def build_smoke_environment(base: Mapping[str, str]) -> dict[str, str]: + """Return the no-network, no-GPU environment used by each smoke.""" + + environment = dict(base) + environment.update(OFFLINE_SMOKE_ENVIRONMENT) + return environment + + +def _run( + stage: str, + command: Sequence[str], + *, + cwd: Path, + environment: Mapping[str, str], +) -> subprocess.CompletedProcess[str]: + completed = subprocess.run( + command, + cwd=cwd, + env=environment, + check=False, + text=True, + capture_output=True, + ) + if completed.returncode: + raise MatrixCommandError(stage, completed) + return completed + + +def _output_fingerprint(text: str) -> dict[str, object]: + """Describe subprocess output without persisting URLs, tokens, or host paths.""" + + encoded = text.encode("utf-8", errors="replace") + return { + "bytes": len(encoded), + "lines": len(text.splitlines()), + "sha256": hashlib.sha256(encoded).hexdigest(), + } + + +def _run_member( + *, + uv: str, + project_root: Path, + temporary_root: Path, + target: str, +) -> dict[str, Any]: + started = time.perf_counter() + environment_root = temporary_root / f"python-{target.replace('.', '')}" + environment = dict(os.environ) + environment["UV_PYTHON_PREFERENCE"] = "only-managed" + stage = "python-install" + + try: + print(f"[{target}] installing the uv-managed interpreter", flush=True) + _run( + stage, + (uv, "python", "install", target), + cwd=project_root, + environment=environment, + ) + + stage = "environment-create" + print(f"[{target}] creating an isolated environment", flush=True) + _run( + stage, + (uv, "venv", "--python", target, str(environment_root)), + cwd=project_root, + environment=environment, + ) + python = environment_root / ("Scripts/python.exe" if os.name == "nt" else "bin/python") + if not python.is_file(): + raise RuntimeError(f"uv did not create the expected interpreter: {python}") + + stage = "dependency-install" + print(f"[{target}] installing the declared CPU runtime dependencies", flush=True) + _run( + stage, + build_dependency_install_command(uv, python, project_root), + cwd=temporary_root, + environment=environment, + ) + + stage = "offline-cpu-source-smoke" + print(f"[{target}] running the isolated repository-source smoke", flush=True) + completed = _run( + stage, + ( + str(python), + "-I", + str(project_root / "tools/remote/python_support_smoke.py"), + "--expected-python", + target, + "--source-root", + str(project_root / "src"), + ), + cwd=temporary_root, + environment=build_smoke_environment(environment), + ) + lines = [line for line in completed.stdout.splitlines() if line.strip()] + if not lines: + raise RuntimeError("The support smoke produced no JSON evidence.") + evidence = json.loads(lines[-1]) + if not isinstance(evidence, dict): + raise RuntimeError("The support smoke result is not a JSON object.") + + elapsed = time.perf_counter() - started + return { + "target": target, + "status": "passed", + "elapsed_seconds": round(elapsed, 3), + "evidence": evidence, + "stderr": _output_fingerprint(completed.stderr), + } + except MatrixCommandError as error: + elapsed = time.perf_counter() - started + return { + "target": target, + "status": "failed", + "stage": error.stage, + "elapsed_seconds": round(elapsed, 3), + "returncode": error.completed.returncode, + "stdout": _output_fingerprint(error.completed.stdout), + "stderr": _output_fingerprint(error.completed.stderr), + } + except Exception as error: + elapsed = time.perf_counter() - started + return { + "target": target, + "status": "failed", + "stage": stage, + "elapsed_seconds": round(elapsed, 3), + "error_type": type(error).__name__, + } + + +def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def _write_junit(path: Path, results: Sequence[Mapping[str, Any]], elapsed: float) -> None: + failures = sum(result["status"] != "passed" for result in results) + suite = ET.Element( + "testsuite", + { + "name": "fastplms-python-support", + "tests": str(len(results)), + "failures": str(failures), + "errors": "0", + "skipped": "0", + "time": f"{elapsed:.3f}", + }, + ) + properties = ET.SubElement(suite, "properties") + ET.SubElement( + properties, + "property", + {"name": "canonical_gpu_python", "value": CANONICAL_GPU_PYTHON}, + ) + for result in results: + case = ET.SubElement( + suite, + "testcase", + { + "classname": "tools.remote.python_matrix", + "name": f"python-{result['target']}", + "time": f"{float(result['elapsed_seconds']):.3f}", + }, + ) + if result["status"] != "passed": + failure = ET.SubElement( + case, + "failure", + { + "message": str(result.get("stage", "matrix-member")), + "type": "PythonSupportFailure", + }, + ) + failure.text = json.dumps(result, indent=2, sort_keys=True) + output = ET.SubElement(case, "system-out") + output.text = json.dumps(result, sort_keys=True) + + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + ET.ElementTree(suite).write(temporary, encoding="utf-8", xml_declaration=True) + temporary.replace(path) + + +def run_matrix( + *, + project_root: Path, + output: Path, + junit_output: Path, + versions: Sequence[str] = PYTHON_SUPPORT_VERSIONS, +) -> int: + """Run every support smoke and persist complete machine-readable results.""" + + project_root = project_root.resolve() + required_paths = ( + project_root / "src/fastplms", + project_root / "requirements/profiles/runtime.in", + project_root / "requirements/constraints/validation.txt", + ) + missing = [str(path) for path in required_paths if not path.exists()] + if missing: + raise FileNotFoundError(f"FastPLMs source workspace is incomplete: {missing}") + uv = shutil.which("uv") + if uv is None: + raise RuntimeError("uv is required for the Python support matrix.") + if not versions or len(versions) != len(set(versions)): + raise ValueError("Python support versions must be a non-empty unique sequence.") + + started = time.perf_counter() + with tempfile.TemporaryDirectory(prefix="fastplms-python-support-") as temporary: + temporary_root = Path(temporary) + with concurrent.futures.ThreadPoolExecutor( + max_workers=min(4, len(versions)), + thread_name_prefix="fastplms-python-support", + ) as executor: + futures = { + target: executor.submit( + _run_member, + uv=uv, + project_root=project_root, + temporary_root=temporary_root, + target=target, + ) + for target in versions + } + results = [futures[target].result() for target in versions] + + elapsed = time.perf_counter() - started + payload = { + "schema_version": 2, + "canonical_gpu_python": CANONICAL_GPU_PYTHON, + "support_matrix": list(versions), + "elapsed_seconds": round(elapsed, 3), + "results": results, + } + _atomic_json(output, payload) + _write_junit(junit_output, results, elapsed) + passed = sum(result["status"] == "passed" for result in results) + print(f"Python support matrix: {passed}/{len(results)} passed", flush=True) + print(output, flush=True) + print(junit_output, flush=True) + return 0 if passed == len(results) else 1 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--project-root", type=Path, default=Path.cwd()) + parser.add_argument("--output", type=Path, default=Path("artifacts/python-matrix.json")) + parser.add_argument( + "--junit-output", + type=Path, + default=Path("artifacts/junit/python-matrix.xml"), + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + return run_matrix( + project_root=arguments.project_root, + output=arguments.output, + junit_output=arguments.junit_output, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/python_support_smoke.py b/tools/remote/python_support_smoke.py new file mode 100644 index 0000000..9734522 --- /dev/null +++ b/tools/remote/python_support_smoke.py @@ -0,0 +1,179 @@ +"""Validate FastPLMs repository source without checkpoint, network, or GPU access.""" + +from __future__ import annotations + +import argparse +import importlib +import importlib.metadata +import importlib.util +import json +import os +import socket +import sys +from pathlib import Path + + +_OFFLINE_ENVIRONMENT = { + "CUDA_VISIBLE_DEVICES": "", + "HF_DATASETS_OFFLINE": "1", + "HF_HUB_OFFLINE": "1", + "PYTHONNOUSERSITE": "1", + "TRANSFORMERS_OFFLINE": "1", +} + + +def _network_blocked(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("Network access is forbidden in the repository-source smoke") + + +def _compile_sources(package_root: Path) -> int: + source_files = sorted(package_root.rglob("*.py")) + for path in source_files: + source = path.read_text(encoding="utf-8") + compile(source, str(path), "exec") + return len(source_files) + + +def run_smoke(expected_python: str, source_root: Path) -> dict[str, object]: + """Return evidence for an isolated, CPU-only repository-source environment.""" + + for name, expected_value in _OFFLINE_ENVIRONMENT.items(): + if os.environ.get(name) != expected_value: + raise AssertionError( + f"The repository-source smoke requires {name}={expected_value!r}." + ) + + socket.create_connection = _network_blocked # type: ignore[assignment] + socket.getaddrinfo = _network_blocked # type: ignore[assignment] + socket.socket.connect = _network_blocked # type: ignore[method-assign] + + expected = tuple(int(part) for part in expected_python.split(".")) + if sys.version_info[:2] != expected: + raise AssertionError( + f"Expected Python {expected_python}, found " + f"{sys.version_info.major}.{sys.version_info.minor}." + ) + if not (sys.version_info[:2] >= (3, 11) and sys.version_info[:2] < (3, 15)): + raise AssertionError("Interpreter is outside FastPLMs' supported Python range.") + + source_root = source_root.resolve() + expected_package_root = source_root / "fastplms" + if not (expected_package_root / "models.toml").is_file(): + raise AssertionError(f"FastPLMs source is incomplete: {expected_package_root}") + sys.path.insert(0, str(source_root)) + + import torch + + import fastplms + from fastplms.models.esm2.modeling_fastesm import FastEsmConfig, FastEsmModel + from fastplms.registry import get_model_registry + + package_root = Path(fastplms.__file__).resolve().parent + if package_root != expected_package_root: + raise AssertionError( + f"FastPLMs did not load from the requested source root: {package_root}" + ) + if fastplms.__version__ != "1.0.0": + raise AssertionError(f"Unexpected FastPLMs source version: {fastplms.__version__!r}") + if importlib.metadata.version("torch").split("+", maxsplit=1)[0] != "2.13.0": + raise AssertionError( + f"Expected Torch 2.13.0, found {importlib.metadata.version('torch')}." + ) + if importlib.metadata.version("transformers") != "5.13.0": + raise AssertionError( + "Expected Transformers 5.13.0, found " + f"{importlib.metadata.version('transformers')}." + ) + if torch.version.cuda is not None or torch.cuda.is_available(): + raise AssertionError("The repository-source smoke must use the CPU-only Torch build.") + if importlib.util.find_spec("flash_attn") is not None: + raise AssertionError("FlashAttention is present in the core source environment.") + + registry = get_model_registry() + if len(registry.families) != 10 or len(tuple(registry)) != 29: + raise AssertionError("The source registry must contain 10 families and 29 checkpoints.") + family_maps = {spec.family.id: spec.auto_map for spec in registry.values()} + advertised_entries = sum(len(auto_map) for auto_map in family_maps.values()) + if advertised_entries != 37: + raise AssertionError(f"Expected 37 advertised Auto entries, found {advertised_entries}.") + + imported_entries: list[str] = [] + for family_id, auto_map in sorted(family_maps.items()): + for auto_class, class_path in sorted(auto_map.items()): + module_name, separator, class_name = class_path.rpartition(".") + if not separator: + raise AssertionError(f"Invalid AutoMap path: {class_path!r}") + auto_class_type = getattr(importlib.import_module(module_name), class_name) + if not isinstance(auto_class_type, type): + raise AssertionError(f"AutoMap target is not a class: {class_path}") + imported_entries.append(f"{family_id}:{auto_class}") + + source_files = _compile_sources(package_root) + config = FastEsmConfig( + vocab_size=33, + mask_token_id=32, + pad_token_id=1, + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=32, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + max_position_embeddings=16, + position_embedding_type="rotary", + token_dropout=False, + attn_backend="eager", + ) + model = FastEsmModel(config, add_pooling_layer=False).eval() + input_ids = torch.tensor(((0, 5, 6, 2, 1),), dtype=torch.long, device="cpu") + attention_mask = input_ids.ne(1) + with torch.inference_mode(): + # H is the hidden-state tensor with shape (b, l, d). + hidden_states = model( + input_ids=input_ids, + attention_mask=attention_mask, + ).last_hidden_state + if tuple(hidden_states.shape) != (1, 5, 16): + raise AssertionError(f"Unexpected hidden-state shape: {tuple(hidden_states.shape)}") + if hidden_states.device.type != "cpu" or not torch.isfinite(hidden_states).all(): + raise AssertionError("The CPU construction smoke produced an invalid tensor.") + if torch.cuda.is_initialized(): # type: ignore[no-untyped-call] + raise AssertionError("The CPU support smoke initialized CUDA.") + + return { + "python": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}", + "fastplms": fastplms.__version__, + "torch": importlib.metadata.version("torch"), + "transformers": importlib.metadata.version("transformers"), + "package_root": str(package_root), + "model_families": len(registry.families), + "checkpoints": len(tuple(registry)), + "advertised_auto_entries": len(imported_entries), + "source_files": source_files, + "hidden_state_shape": list(hidden_states.shape), + "device": hidden_states.device.type, + "cpu_only_torch": torch.version.cuda is None, + "cuda_initialized": torch.cuda.is_initialized(), # type: ignore[no-untyped-call] + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--expected-python", required=True) + parser.add_argument("--source-root", type=Path, required=True) + return parser + + +def main() -> int: + arguments = build_parser().parse_args() + print( + json.dumps( + run_smoke(arguments.expected_python, arguments.source_root), + sort_keys=True, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/reference_source_attestation.py b/tools/remote/reference_source_attestation.py new file mode 100644 index 0000000..cf60fcc --- /dev/null +++ b/tools/remote/reference_source_attestation.py @@ -0,0 +1,495 @@ +"""Create and verify immutable source attestations for Git-free reference images.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import importlib.util +import json +import re +import sys +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path, PurePosixPath +from typing import Any + +from tools.source_provenance import actual_tree_paths, tracked_tree_digest + + +_SCHEMA_VERSION = 1 +REFERENCE_SOURCE_EVIDENCE_SCHEMA_VERSION = 1 +_HEX_REVISION_LENGTH = 40 +_HEX_DIGEST_LENGTH = 64 + + +class ReferenceSourceAttestationError(RuntimeError): + """A reference source tree or imported package differs from its pinned contract.""" + + +@dataclass(frozen=True) +class ReferenceSourceContract: + """Immutable identity and import contract for one copied reference source tree.""" + + schema_version: int + source_revision: str + tree_sha256: str + import_name: str + import_root: str + package_version: str + + +def _load_json_bytes(path: Path) -> tuple[dict[str, object], bytes]: + try: + serialized = path.read_bytes() + raw: Any = json.loads(serialized.decode("utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: + raise ReferenceSourceAttestationError(f"Unable to read source contract: {path}") from error + if not isinstance(raw, dict) or not all(isinstance(key, str) for key in raw): + raise ReferenceSourceAttestationError(f"Source contract is not a JSON object: {path}") + return {str(key): value for key, value in raw.items()}, serialized + + +def _load_json(path: Path) -> dict[str, object]: + return _load_json_bytes(path)[0] + + +def _is_lower_hex(value: object, length: int) -> bool: + return ( + isinstance(value, str) + and len(value) == length + and all(character in "0123456789abcdef" for character in value) + ) + + +def _portable_relative_path(value: object, *, field: str) -> str: + if not isinstance(value, str) or not value: + raise ReferenceSourceAttestationError(f"{field} must be a non-empty string.") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or path.as_posix() != value: + raise ReferenceSourceAttestationError(f"{field} is not a portable relative path: {value!r}") + return value + + +def _validate_contract(raw: Mapping[str, object]) -> ReferenceSourceContract: + required = { + "schema_version", + "source_revision", + "tree_sha256", + "import_name", + "import_root", + "package_version", + } + if set(raw) != required: + raise ReferenceSourceAttestationError( + "Reference source contract fields differ: " + f"missing={sorted(required.difference(raw))}, " + f"extra={sorted(set(raw).difference(required))}" + ) + if raw["schema_version"] != _SCHEMA_VERSION: + raise ReferenceSourceAttestationError("Unsupported reference source schema version.") + source_revision = raw["source_revision"] + tree_sha256 = raw["tree_sha256"] + if not isinstance(source_revision, str) or not _is_lower_hex( + source_revision, _HEX_REVISION_LENGTH + ): + raise ReferenceSourceAttestationError("Reference source revision must be 40 lowercase hex.") + if not isinstance(tree_sha256, str) or not _is_lower_hex( + tree_sha256, _HEX_DIGEST_LENGTH + ): + raise ReferenceSourceAttestationError("Reference tree digest must be 64 lowercase hex.") + import_name = raw["import_name"] + package_version = raw["package_version"] + if not isinstance(import_name, str) or re.fullmatch( + r"[A-Za-z_][A-Za-z0-9_]*", import_name + ) is None: + raise ReferenceSourceAttestationError( + "Reference import_name must name one top-level package." + ) + if not isinstance(package_version, str) or not package_version: + raise ReferenceSourceAttestationError("Reference package_version must be non-empty.") + return ReferenceSourceContract( + schema_version=_SCHEMA_VERSION, + source_revision=source_revision, + tree_sha256=tree_sha256, + import_name=import_name, + import_root=_portable_relative_path(raw["import_root"], field="import_root"), + package_version=package_version, + ) + + +def load_reference_source_contract(path: Path) -> ReferenceSourceContract: + """Load and validate one immutable source contract.""" + + return _validate_contract(_load_json(path)) + + +def _atomic_json(path: Path, payload: Mapping[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def create_reference_source_attestation( + source_root: Path, + contract_path: Path, + output: Path, +) -> dict[str, object]: + """Verify a Git-free copied tree and persist its tracked runtime proof.""" + + source_root = source_root.resolve() + contract = load_reference_source_contract(contract_path) + tracked_files = actual_tree_paths(source_root) + actual_digest = tracked_tree_digest(source_root, tracked_files) + if actual_digest != contract.tree_sha256: + raise ReferenceSourceAttestationError( + "Copied reference source tree differs from its pinned revision: " + f"expected {contract.tree_sha256}, received {actual_digest}." + ) + payload: dict[str, object] = { + **asdict(contract), + "file_count": len(tracked_files), + "tracked_files": list(tracked_files), + } + _atomic_json(output, payload) + return payload + + +def _load_attestation( + path: Path, +) -> tuple[ReferenceSourceContract, tuple[str, ...], str]: + raw, serialized = _load_json_bytes(path) + extra_fields = {"file_count", "tracked_files"} + contract_fields = set(ReferenceSourceContract.__dataclass_fields__) + if set(raw) != contract_fields | extra_fields: + raise ReferenceSourceAttestationError("Reference source attestation fields differ.") + contract = _validate_contract({field: raw[field] for field in contract_fields}) + tracked_files = raw["tracked_files"] + file_count = raw["file_count"] + if not isinstance(tracked_files, list) or not all( + isinstance(relative_name, str) and relative_name for relative_name in tracked_files + ): + raise ReferenceSourceAttestationError("Reference source attestation file list is invalid.") + normalized = tuple( + _portable_relative_path(relative_name, field="tracked_files entry") + for relative_name in tracked_files + ) + if normalized != tuple(sorted(normalized)) or len(normalized) != len(set(normalized)): + raise ReferenceSourceAttestationError( + "Reference source attestation file list is not unique and sorted." + ) + if ( + isinstance(file_count, bool) + or not isinstance(file_count, int) + or file_count != len(normalized) + ): + raise ReferenceSourceAttestationError("Reference source attestation file count differs.") + return contract, normalized, hashlib.sha256(serialized).hexdigest() + + +def validate_reference_source_evidence(value: object) -> dict[str, object]: + """Validate the stable, portable provenance record written to native results.""" + + expected_fields = { + "schema_version", + "source_revision", + "tree_sha256", + "attestation_sha256", + "file_count", + "import_name", + "import_root", + "import_file", + "package_version", + } + if not isinstance(value, Mapping) or set(value) != expected_fields: + raise ReferenceSourceAttestationError( + "Reference source evidence fields differ from schema v1." + ) + if value["schema_version"] != REFERENCE_SOURCE_EVIDENCE_SCHEMA_VERSION: + raise ReferenceSourceAttestationError( + "Unsupported reference source evidence schema version." + ) + if not _is_lower_hex(value["source_revision"], _HEX_REVISION_LENGTH): + raise ReferenceSourceAttestationError("Evidence source revision is invalid.") + for field in ("tree_sha256", "attestation_sha256"): + if not _is_lower_hex(value[field], _HEX_DIGEST_LENGTH): + raise ReferenceSourceAttestationError(f"Evidence {field} is invalid.") + file_count = value["file_count"] + if isinstance(file_count, bool) or not isinstance(file_count, int) or file_count <= 0: + raise ReferenceSourceAttestationError("Evidence file_count must be positive.") + import_name = value["import_name"] + if not isinstance(import_name, str) or re.fullmatch( + r"[A-Za-z_][A-Za-z0-9_]*", import_name + ) is None: + raise ReferenceSourceAttestationError("Evidence import_name is invalid.") + import_root = _portable_relative_path(value["import_root"], field="import_root") + import_file = _portable_relative_path(value["import_file"], field="import_file") + if not import_file.startswith(import_root + "/"): + raise ReferenceSourceAttestationError( + "Evidence import_file is outside the declared import_root." + ) + package_version = value["package_version"] + if not isinstance(package_version, str) or not package_version.strip(): + raise ReferenceSourceAttestationError("Evidence package_version is invalid.") + return {field: value[field] for field in sorted(expected_fields)} + + +def validate_reference_sources_evidence( + value: object, + *, + required_sources: Sequence[str], +) -> dict[str, dict[str, object]]: + """Validate a named, exact set of source attestations for one reference.""" + + required = tuple(required_sources) + if ( + not required + or len(required) != len(set(required)) + or any( + not isinstance(name, str) + or re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)*", name) is None + for name in required + ) + ): + raise ReferenceSourceAttestationError( + "Required reference source names must be unique lowercase slugs." + ) + if not isinstance(value, Mapping) or set(value) != set(required): + observed = sorted(str(name) for name in value) if isinstance(value, Mapping) else [] + raise ReferenceSourceAttestationError( + "Reference source evidence names differ: " + f"expected={sorted(required)}, observed={observed}." + ) + return { + name: validate_reference_source_evidence(value[name]) + for name in sorted(required) + } + + +def _assert_exact_tree_inventory( + source_root: Path, + tracked_files: tuple[str, ...], + *, + stage: str, +) -> None: + actual_files = actual_tree_paths(source_root) + if actual_files == tracked_files: + return + missing = sorted(set(tracked_files).difference(actual_files)) + extra = sorted(set(actual_files).difference(tracked_files)) + raise ReferenceSourceAttestationError( + f"Reference source inventory differs {stage}; " + f"missing={missing[:10]}, extra={extra[:10]}." + ) + + +def _assert_source_file( + raw_module_file: object, + import_root: Path, + *, + context: str, +) -> Path: + if not isinstance(raw_module_file, str): + raise ReferenceSourceAttestationError( + f"{context} reference package has no source file." + ) + module_file = Path(raw_module_file).resolve() + try: + module_file.relative_to(import_root) + except ValueError as error: + raise ReferenceSourceAttestationError( + f"{context} reference package resolves outside the pinned source: " + f"{module_file}" + ) from error + return module_file + + +def _module_source_file(module: object, import_root: Path, *, context: str) -> Path: + return _assert_source_file( + getattr(module, "__file__", None), + import_root, + context=context, + ) + + +def _prioritize_import_parent(import_root: Path) -> None: + import_parent = import_root.parent.resolve() + retained: list[str] = [] + for entry in sys.path: + try: + if Path(entry or ".").resolve() == import_parent: + continue + except OSError: + pass + retained.append(entry) + sys.path[:] = [str(import_parent), *retained] + + +def _validate_cached_package_modules( + import_name: str, + import_root: Path, +) -> object | None: + package_prefix = import_name + "." + cached_top_level: object | None = None + for module_name, module in tuple(sys.modules.items()): + if module_name != import_name and not module_name.startswith(package_prefix): + continue + if module is None: + raise ReferenceSourceAttestationError( + f"Cached reference module {module_name!r} has no module object." + ) + _module_source_file(module, import_root, context=f"Cached {module_name!r}") + if module_name == import_name: + cached_top_level = module + return cached_top_level + + +def verify_reference_source( + source_root: Path, + attestation_path: Path, + contract_path: Path, + *, + expected_revision: str, +) -> dict[str, object]: + """Rehash the pinned tree and prove that the imported package comes from it.""" + + if not _is_lower_hex(expected_revision, _HEX_REVISION_LENGTH): + raise ReferenceSourceAttestationError( + "Expected reference source revision must be 40 lowercase hex." + ) + source_root = source_root.resolve() + pinned_contract = load_reference_source_contract(contract_path) + contract, tracked_files, attestation_sha256 = _load_attestation(attestation_path) + if contract != pinned_contract: + raise ReferenceSourceAttestationError( + "Runtime source attestation differs from the checked-in contract." + ) + if contract.source_revision != expected_revision: + raise ReferenceSourceAttestationError( + "Reference source revision differs from the runtime expectation: " + f"expected {expected_revision}, received {contract.source_revision}." + ) + _assert_exact_tree_inventory(source_root, tracked_files, stage="before import") + actual_digest = tracked_tree_digest(source_root, tracked_files) + if actual_digest != contract.tree_sha256: + raise ReferenceSourceAttestationError( + "Reference source tree changed after image construction: " + f"expected {contract.tree_sha256}, received {actual_digest}." + ) + + import_root = source_root.joinpath(*PurePosixPath(contract.import_root).parts).resolve() + try: + import_root.relative_to(source_root) + except ValueError as error: + raise ReferenceSourceAttestationError( + "Reference import root escapes its attested source tree." + ) from error + _prioritize_import_parent(import_root) + + cached_module = _validate_cached_package_modules(contract.import_name, import_root) + if cached_module is not None: + _module_source_file(cached_module, import_root, context="Cached top-level") + else: + spec = importlib.util.find_spec(contract.import_name) + if spec is None or not isinstance(spec.origin, str): + raise ReferenceSourceAttestationError( + f"Pinned reference package {contract.import_name!r} is not importable." + ) + _assert_source_file( + spec.origin, + import_root, + context="Preflight", + ) + + previous_dont_write_bytecode = sys.dont_write_bytecode + try: + sys.dont_write_bytecode = True + module = importlib.import_module(contract.import_name) + finally: + sys.dont_write_bytecode = previous_dont_write_bytecode + module_file = _module_source_file(module, import_root, context="Imported") + if cached_module is not None and module is not cached_module: + raise ReferenceSourceAttestationError( + "Reference import cache identity changed during attestation." + ) + _validate_cached_package_modules(contract.import_name, import_root) + imported_version = getattr(module, "__version__", None) + if imported_version != contract.package_version: + raise ReferenceSourceAttestationError( + f"Imported {contract.import_name!r} version {imported_version!r}, " + f"expected {contract.package_version!r}." + ) + _assert_exact_tree_inventory(source_root, tracked_files, stage="after import") + post_import_digest = tracked_tree_digest(source_root, tracked_files) + if post_import_digest != actual_digest: + raise ReferenceSourceAttestationError( + "Reference source tree changed while importing the pinned package." + ) + import_file = module_file.relative_to(source_root).as_posix() + return validate_reference_source_evidence({ + "schema_version": REFERENCE_SOURCE_EVIDENCE_SCHEMA_VERSION, + "source_revision": contract.source_revision, + "tree_sha256": actual_digest, + "attestation_sha256": attestation_sha256, + "file_count": len(tracked_files), + "import_name": contract.import_name, + "import_root": contract.import_root, + "import_file": import_file, + "package_version": imported_version, + }) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + create = subparsers.add_parser("create") + create.add_argument("--source-root", type=Path, required=True) + create.add_argument("--contract", type=Path, required=True) + create.add_argument("--output", type=Path, required=True) + verify = subparsers.add_parser("verify") + verify.add_argument("--source-root", type=Path, required=True) + verify.add_argument("--attestation", type=Path, required=True) + verify.add_argument("--contract", type=Path, required=True) + verify.add_argument("--expected-revision", required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Create or verify a reference-source attestation.""" + + arguments = _parser().parse_args(argv) + if arguments.command == "create": + payload = create_reference_source_attestation( + arguments.source_root, + arguments.contract, + arguments.output, + ) + else: + payload = verify_reference_source( + arguments.source_root, + arguments.attestation, + arguments.contract, + expected_revision=arguments.expected_revision, + ) + summary = { + key: payload[key] + for key in ( + "schema_version", + "source_revision", + "tree_sha256", + "attestation_sha256", + "file_count", + "import_name", + "import_root", + "import_file", + "package_version", + ) + if key in payload + } + print(json.dumps(summary, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/run.py b/tools/remote/run.py new file mode 100644 index 0000000..8d17849 --- /dev/null +++ b/tools/remote/run.py @@ -0,0 +1,2041 @@ +"""Synchronize a clean workspace and run FastPLMs containers over SSH. + +The runner accepts the SSH host and identity only at invocation time. It does +not read credential files, copy ignored files, or persist workstation details +in the repository. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import io +import json +import re +import secrets +import shlex +import subprocess +import sys +import tarfile +import tempfile +import time +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + +from tools.source_provenance import ( + ARCHIVE_PROVENANCE_NAME, + archive_root_record, + render_archive_provenance, + tracked_tree_digest, +) + + +HOST_PATTERN = re.compile(r"^[A-Za-z0-9_.@:\-]+$") +RUN_PATTERN = re.compile(r"^[0-9]{8}T[0-9]{6}Z-[0-9a-f]{16}$") +GIT_REVISION_PATTERN = re.compile(r"^[0-9a-f]{40}$") +REMOTE_CLEANUP_SCRIPT = """set -eu +base=$(realpath -e -- "$1") +workspace=$(realpath -e -- "$2") +case "$workspace" in + "$base"/*) ;; + *) echo "refusing cleanup outside managed remote base" >&2; exit 64 ;; +esac +test "$workspace" != "$base" +rm -rf -- "$workspace" +""" +SENSITIVE_NAMES = { + ".agents", + ".claude", + ".codex", + ".env", + ".secrets.env", + "credentials", + "credentials.json", + "id_rsa", + "id_ed25519", +} +SENSITIVE_SUFFIXES = {".key", ".pem", ".p12", ".pfx"} +_ARTIFACT_TREE_DOMAIN = b"fastplms-remote-artifact-tree-v1\0" +_CONTROL_TIMEOUT_SECONDS = 300 +_TRANSFER_TIMEOUT_SECONDS = 1_800 +_BIOHUB_REFERENCE_TARGETS = frozenset({"reference-biohub-esm", "reference-esmfold2"}) +_BIOHUB_BUILD_TARGET = "biohub-biotraj-wheel" +_MACHINE_PATTERN = re.compile(r"^[A-Za-z0-9_.-]{1,64}$") +_GH200_RELEASE_BACKENDS = ("eager", "sdpa", "flex_attention") +_FLASH_ATTENTION_2_REVISION = "db6b51744f0cd7061386442c09df890fc6d9f47e" +_FLASH_ATTENTION_3_REVISION = "43f0bd269777115d94ff826e0d113ce9c1c9087b" +_REFERENCE_IMAGE_IDENTITY_PATH = ( + "artifacts/reference/environment/container-images.json" +) +_WRITE_JSON_SCRIPT = """import pathlib, sys +path = pathlib.Path(sys.argv[1]) +path.parent.mkdir(parents=True, exist_ok=True) +temporary = path.with_suffix(path.suffix + '.tmp') +temporary.write_text(sys.argv[2] + '\\n', encoding='utf-8') +temporary.replace(path) +""" + + +@dataclass(frozen=True) +class Suite: + """Images to build and the command executed in the remote workspace.""" + + bake_targets: tuple[str, ...] + command: tuple[str, ...] + pre_commands: tuple[tuple[str, ...], ...] = () + required_paths: tuple[str, ...] = () + build_timeout_seconds: int = 7_200 + pre_command_timeout_seconds: int = 7_200 + command_timeout_seconds: int = 7_200 + attention_backends: tuple[str, ...] = () + + +def _normalized_host_architecture(machine: str) -> str: + """Normalize trusted ``uname -m`` aliases without guessing unknown machines.""" + + value = machine.strip().lower() + if _MACHINE_PATTERN.fullmatch(value) is None: + return "unknown" + if value in {"amd64", "x86_64"}: + return "amd64" + if value in {"aarch64", "arm64"}: + return "arm64" + return value + + +def _host_hardware_preflight(machine: str, gpu_output: str) -> dict[str, object]: + """Return one exact, platform-neutral host architecture and GPU binding.""" + + uname_machine = machine.strip().lower() + architecture = _normalized_host_architecture(uname_machine) + if architecture == "unknown": + raise RuntimeError("Remote uname returned an invalid machine architecture") + gpus: list[dict[str, object]] = [] + seen_uuids: set[str] = set() + for raw_line in gpu_output.splitlines(): + if not raw_line.strip(): + continue + fields = [field.strip() for field in raw_line.split(",")] + if len(fields) != 4 or any(not field for field in fields): + raise RuntimeError("nvidia-smi returned an invalid GPU identity record") + name, uuid, driver_version, raw_memory = fields + if uuid in seen_uuids: + raise RuntimeError("nvidia-smi returned a duplicate GPU UUID") + try: + memory_total_mib = int(raw_memory) + except ValueError as error: + raise RuntimeError("nvidia-smi returned invalid total GPU memory") from error + if memory_total_mib <= 0: + raise RuntimeError("nvidia-smi returned non-positive total GPU memory") + seen_uuids.add(uuid) + gpus.append( + { + "name": name, + "uuid": uuid, + "driver_version": driver_version, + "memory_total_mib": memory_total_mib, + } + ) + if not gpus: + raise RuntimeError("Remote validation requires an identifiable NVIDIA GPU") + gpus.sort(key=lambda item: str(item["uuid"])) + if architecture not in {"amd64", "arm64"}: + raise RuntimeError( + f"Remote validation does not declare an OCI platform for {architecture!r}" + ) + identity = { + "uname_machine": uname_machine, + "architecture": architecture, + "container_platform": f"linux/{architecture}", + "gpus": gpus, + } + identity_sha256 = hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + return {"status": "passed", **identity, "identity_sha256": identity_sha256} + + +def _kernel_capability_preflight( + host_hardware: Mapping[str, object], + requested_backends: Sequence[str], +) -> dict[str, object]: + """Resolve the no-download attention matrix for one exact native platform.""" + + platform_name = host_hardware.get("container_platform") + architecture = host_hardware.get("architecture") + if platform_name != "linux/arm64" or architecture != "arm64": + return { + "schema_version": 1, + "status": "failed", + "policy": "gh200-native-no-flash-download-v1", + "platform": platform_name, + "selected_backends": list(requested_backends), + "network_downloads": False, + "source_builds": False, + "reason": "The current release kernel policy is bound to native GH200 linux/arm64.", + "backends": {}, + } + + requested = tuple(requested_backends) + if len(set(requested)) != len(requested): + return { + "schema_version": 1, + "status": "failed", + "policy": "gh200-native-no-flash-download-v1", + "platform": platform_name, + "selected_backends": list(requested), + "network_downloads": False, + "source_builds": False, + "reason": "The requested attention matrix contains duplicate backends.", + "backends": {}, + } + + backend_records: dict[str, dict[str, object]] = { + "eager": { + "status": "available", + "selected": "eager" in requested, + "provider": "torch", + "reason": "Framework eager attention is available without an external kernel.", + }, + "sdpa": { + "status": "available", + "selected": "sdpa" in requested, + "provider": "torch", + "reason": "Framework SDPA is available without an external kernel.", + }, + "flex_attention": { + "status": "available", + "selected": "flex_attention" in requested, + "provider": "torch", + "reason": "Framework Flex Attention is available without an external kernel.", + }, + "flash_attention_2": { + "status": "prior_focused_evidence_only", + "selected": False, + "provider": "kernels-community/flash-attn2", + "revision": _FLASH_ATTENTION_2_REVISION, + "reason": ( + "The GH200 release matrix reuses prior revision-pinned focused FA2 " + "evidence; it does not download, build, or execute FA2 in this run." + ), + }, + "flash_attention_3": { + "status": "unavailable", + "selected": False, + "provider": "kernels-community/flash-attn3", + "revision": _FLASH_ATTENTION_3_REVISION, + "reason": ( + "The manifest-pinned FA3 kernel has no validated linux/arm64 artifact " + "for the current GH200 release image." + ), + }, + } + unavailable = [ + backend + for backend in requested + if backend not in backend_records + or backend_records[backend]["status"] != "available" + ] + return { + "schema_version": 1, + "status": "failed" if unavailable else "passed", + "policy": "gh200-native-no-flash-download-v1", + "platform": platform_name, + "selected_backends": list(requested), + "excluded_backends": [ + backend for backend in backend_records if backend not in requested + ], + "network_downloads": False, + "source_builds": False, + "reason": ( + "Requested backends are unavailable under the native GH200 policy: " + + ", ".join(unavailable) + if unavailable + else None + ), + "backends": backend_records, + } + + +def _reference_container_image_identity( + execution_environment: Mapping[str, object], +) -> dict[str, object]: + """Return the stable image/runtime identity shared with reference containers.""" + + platform_name = execution_environment.get("container_platform") + if not isinstance(platform_name, str) or not platform_name.startswith("linux/"): + raise RuntimeError("Execution environment has no resolved Linux platform") + raw_images = execution_environment.get("images") + if not isinstance(raw_images, Mapping) or not raw_images: + raise RuntimeError("Execution environment has no built-image identity map") + images: dict[str, dict[str, str]] = {} + for raw_name, raw_identity in sorted(raw_images.items(), key=lambda item: str(item[0])): + if not isinstance(raw_name, str) or not isinstance(raw_identity, Mapping): + raise RuntimeError("Execution environment contains an invalid image identity") + content_digest = raw_identity.get("content_digest") + image_id = raw_identity.get("id") + os_name = raw_identity.get("os") + architecture = raw_identity.get("architecture") + resolved_platform = raw_identity.get("resolved_platform") + if ( + not isinstance(content_digest, str) + or re.fullmatch(r"sha256:[0-9a-f]{64}", content_digest) is None + or image_id != content_digest + or os_name != "linux" + or architecture != platform_name.split("/")[1] + or resolved_platform != platform_name + ): + raise RuntimeError(f"Built image {raw_name!r} has an invalid stable identity") + images[raw_name] = { + "content_digest": content_digest, + "image_id": content_digest, + "os": os_name, + "architecture": architecture, + "resolved_platform": platform_name, + } + + raw_server = execution_environment.get("docker_server") + if not isinstance(raw_server, Mapping): + raise RuntimeError("Execution environment has no Docker server identity") + server_fields = ( + "Version", + "ApiVersion", + "MinAPIVersion", + "GitCommit", + "Os", + "Arch", + "KernelVersion", + ) + docker_server = { + field: raw_server[field] + for field in server_fields + if isinstance(raw_server.get(field), (str, int, float, bool)) + } + required_server_fields = {"Version", "ApiVersion", "Os", "Arch"} + if not required_server_fields.issubset(docker_server): + raise RuntimeError("Docker server identity is missing required stable fields") + if docker_server["Os"] != "linux" or docker_server["Arch"] != platform_name.split("/")[1]: + raise RuntimeError("Docker server identity differs from the resolved native platform") + buildx = execution_environment.get("docker_buildx") + if not isinstance(buildx, str) or not buildx.strip(): + raise RuntimeError("Execution environment has no Docker Buildx identity") + return { + "schema_version": 1, + "resolved_platform": platform_name, + "docker_server": docker_server, + "docker_buildx": buildx.strip(), + "images": images, + } + + +def _run_report( + *, + run_id: str, + suite_name: str, + suite: Suite, + started_at: str, + finished_at: str, + source_archive_sha256: str, + git_revision: str, + submodule_revisions: Mapping[str, str], + execution_environment: Mapping[str, object] | None, + failure_phase: str | None, + failure: BaseException | None, + artifact_retrieval_returncode: int, + cleanup_status: str, + phase_durations_seconds: Mapping[str, float] | None = None, + cache_telemetry: Mapping[str, object] | None = None, + artifact_inventory: Mapping[str, object] | None = None, + host_hardware_preflight: Mapping[str, object] | None = None, + kernel_capability_preflight: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Build a secret-free machine report for one remote invocation.""" + + failure_record: dict[str, object] | None = None + if failure is not None: + failure_record = {"phase": failure_phase, "type": type(failure).__name__} + if isinstance(failure, subprocess.CalledProcessError): + failure_record["returncode"] = failure.returncode + elif artifact_retrieval_returncode != 0: + failure_record = { + "phase": "artifact-retrieval", + "type": "ArtifactRetrievalError", + "returncode": artifact_retrieval_returncode, + } + passed = failure_record is None and cleanup_status in {"succeeded", "retained"} + return { + "schema_version": 5, + "run_id": run_id, + "suite": suite_name, + "status": "passed" if passed else "failed", + "started_at_utc": started_at, + "finished_at_utc": finished_at, + "source_archive_sha256": source_archive_sha256, + "git_revision": git_revision, + "submodule_revisions": dict(sorted(submodule_revisions.items())), + "execution_environment": ( + dict(execution_environment) if execution_environment is not None else None + ), + "artifact_retrieval": { + "returncode": artifact_retrieval_returncode, + "status": ("succeeded" if artifact_retrieval_returncode == 0 else "failed"), + }, + "remote_cleanup": cleanup_status, + "phase_durations_seconds": { + key: round(value, 3) for key, value in sorted((phase_durations_seconds or {}).items()) + }, + "cache_telemetry": dict(cache_telemetry or {}), + "artifact_inventory": ( + dict(artifact_inventory) if artifact_inventory is not None else None + ), + "host_hardware_preflight": ( + dict(host_hardware_preflight) + if host_hardware_preflight is not None + else None + ), + "kernel_capability_preflight": ( + dict(kernel_capability_preflight) + if kernel_capability_preflight is not None + else None + ), + "failure": failure_record, + "suite_contract": { + "bake_targets": list(suite.bake_targets), + "pre_commands": [list(command) for command in suite.pre_commands], + "command": list(suite.command), + "required_paths": list(suite.required_paths), + "biohub_reference_targets": sorted( + _BIOHUB_REFERENCE_TARGETS.intersection(suite.bake_targets) + ), + "reference_targets": sorted( + target for target in suite.bake_targets if target.startswith("reference-") + ), + "host_hardware_binding_required": True, + "attention_backends": list(suite.attention_backends), + "kernel_downloads_allowed": False, + "same_host_candidate_reference_required": bool( + any(target.startswith("reference-") for target in suite.bake_targets) + ), + "timeouts_seconds": { + "control": _CONTROL_TIMEOUT_SECONDS, + "transfer": _TRANSFER_TIMEOUT_SECONDS, + "build": suite.build_timeout_seconds, + "pre_command": suite.pre_command_timeout_seconds, + "command": suite.command_timeout_seconds, + }, + }, + } + + +def _artifact_tree_summary(root: Path) -> dict[str, object]: + """Hash retrieved artifacts without recording potentially sensitive contents.""" + + root = root.resolve() + if not root.is_dir(): + raise RuntimeError(f"Retrieved artifact root does not exist: {root}") + digest = hashlib.sha256() + digest.update(_ARTIFACT_TREE_DOMAIN) + file_count = 0 + total_bytes = 0 + for path in sorted(root.rglob("*")): + relative_name = path.relative_to(root).as_posix() + relative = PurePosixPath(relative_name) + if path.is_symlink(): + raise RuntimeError("Retrieved artifacts may not contain symlinks") + if path.is_dir(): + continue + if not path.is_file(): + raise RuntimeError("Retrieved artifacts contain a non-regular entry") + if _is_sensitive(relative): + raise RuntimeError("Retrieved artifacts contain a sensitive path") + content = hashlib.sha256() + size = 0 + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + content.update(chunk) + size += len(chunk) + for value in ( + relative_name.encode("utf-8"), + size.to_bytes(8, "big"), + content.digest(), + ): + digest.update(len(value).to_bytes(8, "big")) + digest.update(value) + file_count += 1 + total_bytes += size + return { + "status": "captured", + "file_count": file_count, + "total_bytes": total_bytes, + "tree_sha256": digest.hexdigest(), + } + + +def _compose_run(service: str, *command: str) -> tuple[str, ...]: + return ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + service, + *command, + ) + + +_BUILD_ARTIFACTS = _compose_run( + "candidate", + "python", + "-m", + "tools.artifacts.build_all", + "--output-root", + "dist/hub", + "--source-root", + "/workspace", +) +_BUILD_BENCHMARK_ARTIFACTS = _compose_run( + "candidate", + "python", + "-m", + "tools.artifacts.build_all", + "--benchmark-suite", + "--output-root", + "dist/hub", + "--source-root", + "/workspace", +) +_PREPARE_REFERENCES = _compose_run( + "candidate", + "python", + "-m", + "tools.remote.prepare_references", + "--output-root", + "artifacts/reference", +) +_SEQUENCE_REFERENCE_CONTAINERS = ( + "reference-esm2", + "reference-biohub-esm", + "reference-e1", + "reference-dplm", + "reference-ankh", +) +_RUN_NATIVE_REFERENCES = tuple( + _compose_run( + container, + "python", + "-m", + "tests.parity.support.native_reference", + "--request-dir", + f"/exchange/requests/{container}", + "--output-dir", + "/exchange/results", + ) + for container in _SEQUENCE_REFERENCE_CONTAINERS +) +_RUN_CHECK_ARTIFACTS = _compose_run( + "artifact", + "python", + "-m", + "pytest", + "tests/release/test_published_automodel.py", + "tests/release/test_manifest_readiness.py", + "-m", + "artifact", + "-k", + "not test_local_artifact_locked_flash_backend", + "--junitxml=artifacts/junit/check-artifact.xml", +) +_RUN_CHECK_GOLDENS = _compose_run( + "structure", + "python", + "-m", + "pytest", + "tests/integration/test_official_goldens.py", + "tests/structure/test_structure_official_goldens.py", + "-m", + "gpu and not large", + "--junitxml=artifacts/junit/check-goldens.xml", +) +_RUN_PYTHON_MATRIX = _compose_run( + "candidate", + "python", + "-m", + "tools.remote.python_matrix", + "--output", + "artifacts/python-matrix.json", + "--junit-output", + "artifacts/junit/python-matrix.xml", +) +_RUN_NIGHTLY_FP8 = _compose_run( + "fp8", + "python", + "-m", + "pytest", + "tests/structure/test_esmfold2_fp8_compliance.py", + "--junitxml=artifacts/junit/nightly-fp8.xml", +) +_RUN_NIGHTLY_SEQUENCE_GOLDENS = _compose_run( + "structure", + "python", + "-m", + "pytest", + "tests/integration/test_official_goldens.py", + "--junitxml=artifacts/junit/nightly-sequence-goldens.xml", +) +_RUN_NIGHTLY_STRUCTURE_GOLDENS = _compose_run( + "structure", + "python", + "-m", + "pytest", + "tests/structure/test_structure_official_goldens.py", + "--junitxml=artifacts/junit/nightly-structure-goldens.xml", +) +_RUN_NIGHTLY_BENCHMARK = _compose_run( + "benchmark", + "--artifact-root", + "dist/hub", + "--backends", + *_GH200_RELEASE_BACKENDS, + "--output", + "artifacts/benchmarks/nightly-h100.json", + "--junit-output", + "artifacts/junit/nightly-benchmark.xml", +) +_RUN_RELEASE_BENCHMARK = _compose_run( + "benchmark", + "--artifact-root", + "dist/hub", + "--backends", + *_GH200_RELEASE_BACKENDS, + "--output", + "artifacts/benchmarks/release-h100.json", + "--junit-output", + "artifacts/junit/release-benchmark.xml", +) +_PREPARE_BOLTZ2_BUNDLE = _compose_run( + "structure", + "python", + "-m", + "tests.structure.support.boltz2_bundle", + "prepare", + "--exchange-root", + "/workspace/artifacts/reference", +) +_RUN_BOLTZ2_REFERENCE = _compose_run( + "reference-boltz2", + "python", + "-m", + "tests.structure.support.boltz2_bundle", + "produce-reference", + "--exchange-root", + "/exchange", +) +_RUN_BOLTZ2_CANDIDATE = _compose_run( + "structure", + "python", + "-m", + "tests.structure.support.boltz2_bundle", + "produce-candidate", + "--exchange-root", + "/workspace/artifacts/reference", +) +_PREPARE_ESMFOLD_BUNDLE = _compose_run( + "structure", + "python", + "-m", + "tests.structure.support.esmfold_bundle", + "prepare", + "--exchange-root", + "/workspace/artifacts/reference", +) +_RUN_ESMFOLD_REFERENCES = tuple( + _compose_run( + "reference-esmfold", + "python", + "-m", + "tests.structure.support.esmfold_bundle", + "produce-reference", + "--exchange-root", + "/exchange", + "--precision", + precision, + ) + for precision in ("fp32", "bf16") +) +_RUN_ESMFOLD_CANDIDATES = tuple( + _compose_run( + "structure", + "python", + "-m", + "tests.structure.support.esmfold_bundle", + "produce-candidate", + "--exchange-root", + "/workspace/artifacts/reference", + "--precision", + precision, + ) + for precision in ("fp32", "bf16") +) +_PREPARE_ESMFOLD2_BUNDLES = _compose_run( + "structure", + "python", + "-m", + "tests.structure.support.esmfold2_bundle", + "prepare", + "--exchange-root", + "/workspace/artifacts/reference", +) +_RUN_ESMFOLD2_REFERENCE = _compose_run( + "reference-esmfold2", + "python", + "-m", + "tests.structure.support.esmfold2_bundle", + "produce-reference", + "--exchange-root", + "/exchange", + "--all", +) +_RUN_ESMFOLD2_CANDIDATES = tuple( + _compose_run( + "fp8" if precision == "fp8" else "structure", + "python", + "-m", + "tests.structure.support.esmfold2_bundle", + "produce-candidate", + "--exchange-root", + "/workspace/artifacts/reference", + "--all", + "--precision", + precision, + ) + for precision in ("bf16", "fp8") +) +_RUN_STRUCTURE_REFERENCES = ( + _PREPARE_BOLTZ2_BUNDLE, + _RUN_BOLTZ2_REFERENCE, + _RUN_BOLTZ2_CANDIDATE, + _PREPARE_ESMFOLD_BUNDLE, + *_RUN_ESMFOLD_REFERENCES, + *_RUN_ESMFOLD_CANDIDATES, + _PREPARE_ESMFOLD2_BUNDLES, + _RUN_ESMFOLD2_REFERENCE, + *_RUN_ESMFOLD2_CANDIDATES, +) + +_RUN_RELEASE_STRUCTURE_REFERENCES = ( + _PREPARE_ESMFOLD_BUNDLE, + *_RUN_ESMFOLD_REFERENCES, + *_RUN_ESMFOLD_CANDIDATES, + _PREPARE_ESMFOLD2_BUNDLES, + _RUN_ESMFOLD2_REFERENCE, + _RUN_ESMFOLD2_CANDIDATES[0], +) + +# These source-parity modules are self-contained in the candidate environment. +# Direct model parity plus ANKH and E1 parity remain in the isolated native +# reference workflow because their official dependencies conflict with it. +_RELEASE_LOCAL_PARITY_TESTS = ( + "tests/parity/test_esmfold2_common_parity.py", + "tests/parity/test_esmfold2_protein_data_parity.py", + "tests/parity/test_esmfold2_reimplemented_source_parity.py", + "tests/parity/test_esmfold2_residue_config_parity.py", + "tests/parity/test_esmfold2_source_slice3_parity.py", + "tests/parity/test_esmfold2_source_slice4_parity.py", +) + + +SUITES = { + "check": Suite( + ("candidate-structure",), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/unit", + "tests/integration", + "tests/release", + "-m", + "not gpu and not slow and not structure and not artifact", + "--junitxml=artifacts/junit/check.xml", + ), + pre_commands=( + _RUN_CHECK_GOLDENS, + ), + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "gpu-golden-smoke": Suite( + ("candidate-structure",), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + ( + "tests/release/test_validation_stack.py::" + "test_release_hopper_sm90_gpu_is_available_without_running_a_model" + ), + "tests/integration/test_official_goldens.py", + "tests/structure/test_structure_official_goldens.py", + "-m", + "gpu and not large", + "--junitxml=artifacts/junit/gpu-golden-smoke.xml", + ), + ), + "unit": Suite( + ("candidate-structure",), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/unit", + "--junitxml=artifacts/junit/unit.xml", + ), + ), + "integration": Suite( + ("candidate-structure",), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/integration", + "--junitxml=artifacts/junit/integration.xml", + ), + ), + "compliance": Suite( + ( + "candidate", + "candidate-structure", + "candidate-fp8", + _BIOHUB_BUILD_TARGET, + "reference-esm2", + "reference-biohub-esm", + "reference-e1", + "reference-dplm", + "reference-ankh", + "reference-esmfold", + "reference-esmfold2", + ), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "fp8", + "python", + "-m", + "pytest", + "tests/parity/test_native_results.py", + ( + "tests/release/test_validation_stack.py::" + "test_release_hopper_sm90_gpu_is_available_without_running_a_model" + ), + ( + "tests/release/test_validation_stack.py::" + "test_fp8_validation_stack_uses_the_cuda13_transformer_engine_core" + ), + "tests/structure/test_esmfold_folding_compliance.py", + "tests/structure/test_esmfold2_folding_compliance.py", + "tests/structure/test_esmfold2_fp8_compliance.py", + "--junitxml=artifacts/junit/compliance.xml", + ), + pre_commands=( + _BUILD_ARTIFACTS, + _PREPARE_REFERENCES, + *_RUN_NATIVE_REFERENCES, + *_RUN_RELEASE_STRUCTURE_REFERENCES, + ), + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "structure": Suite( + ( + "candidate-structure", + "candidate-fp8", + _BIOHUB_BUILD_TARGET, + "reference-boltz2", + "reference-esmfold", + "reference-esmfold2", + ), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/structure", + "tests/parity/test_boltz_source_refactor.py", + "--ignore=tests/structure/test_structure_models.py", + "-m", + "structure", + "--junitxml=artifacts/junit/structure.xml", + ), + pre_commands=_RUN_STRUCTURE_REFERENCES, + ), + "feature": Suite( + ("candidate-structure",), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/integration/test_binder_design.py", + "tests/integration/test_dplm_generation.py", + "tests/integration/test_e1_rag.py", + "tests/integration/test_esm3.py", + "tests/integration/test_ttt.py", + "tests/release/test_conversion_tools.py", + "--junitxml=artifacts/junit/feature.xml", + ), + ), + "artifact": Suite( + ("candidate", "candidate-artifact"), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "artifact", + "python", + "-m", + "pytest", + "tests/release", + "-m", + "artifact", + "-k", + "not test_local_artifact_locked_flash_backend", + "--junitxml=artifacts/junit/artifact.xml", + ), + pre_commands=(_BUILD_ARTIFACTS,), + ), + "benchmark": Suite( + ("candidate", "candidate-fp8"), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "benchmark", + "--artifact-root", + "dist/hub", + "--backends", + *_GH200_RELEASE_BACKENDS, + "--output", + "artifacts/benchmarks/h100-current.json", + "--baseline", + "benchmarks/baselines/h100.json", + "--junit-output", + "artifacts/junit/benchmark.xml", + ), + pre_commands=(_BUILD_BENCHMARK_ARTIFACTS,), + required_paths=("benchmarks/baselines/h100.json",), + pre_command_timeout_seconds=14_400, + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "benchmark-capture": Suite( + ("candidate", "candidate-fp8"), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "benchmark", + "--artifact-root", + "dist/hub", + "--backends", + *_GH200_RELEASE_BACKENDS, + "--output", + "artifacts/benchmarks/h100-baseline-candidate.json", + "--junit-output", + "artifacts/junit/benchmark-capture.xml", + ), + pre_commands=(_BUILD_BENCHMARK_ARTIFACTS,), + pre_command_timeout_seconds=14_400, + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "nightly": Suite( + ( + "candidate", + "candidate-structure", + "candidate-fp8", + "candidate-artifact", + ), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/integration/test_backend_consistency.py", + "tests/integration/test_binder_design.py", + "tests/integration/test_dplm_generation.py", + "tests/integration/test_e1_rag.py", + "tests/integration/test_esm3.py", + "tests/integration/test_ttt.py", + "tests/unit/test_fine_tuning_example.py", + "--junitxml=artifacts/junit/nightly-features.xml", + ), + pre_commands=( + _BUILD_ARTIFACTS, + _RUN_CHECK_ARTIFACTS, + _RUN_NIGHTLY_SEQUENCE_GOLDENS, + _RUN_NIGHTLY_STRUCTURE_GOLDENS, + _RUN_NIGHTLY_FP8, + _RUN_NIGHTLY_BENCHMARK, + ), + pre_command_timeout_seconds=14_400, + command_timeout_seconds=21_600, + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "release": Suite( + ( + "candidate", + "candidate-structure", + "candidate-fp8", + "candidate-artifact", + _BIOHUB_BUILD_TARGET, + *_SEQUENCE_REFERENCE_CONTAINERS, + "reference-esmfold", + "reference-esmfold2", + ), + ( + "sudo", + "docker", + "compose", + "-f", + "docker/compose.yaml", + "run", + "--rm", + "structure", + "python", + "-m", + "pytest", + "tests/unit", + "tests/integration", + "tests/release", + "tests/parity/test_native_results.py", + *_RELEASE_LOCAL_PARITY_TESTS, + "tests/structure", + "tests/parity/test_boltz_source_refactor.py", + "--ignore=tests/structure/test_structure_models.py", + "--ignore=tests/structure/test_esmfold2_fp8_compliance.py", + "--ignore=tests/integration/test_flash_attention_backends.py", + ( + "--deselect=tests/release/test_validation_stack.py::" + "test_fp8_validation_stack_uses_the_cuda13_transformer_engine_core" + ), + ( + "--deselect=tests/structure/test_boltz2_folding_compliance.py::" + "test_boltz2_live_folding_matches_pinned_official" + ), + "-m", + "not artifact", + "--junitxml=artifacts/junit/release.xml", + ), + pre_commands=( + _BUILD_ARTIFACTS, + _RUN_CHECK_ARTIFACTS, + _PREPARE_REFERENCES, + *_RUN_NATIVE_REFERENCES, + *_RUN_RELEASE_STRUCTURE_REFERENCES, + _RUN_PYTHON_MATRIX, + _RUN_RELEASE_BENCHMARK, + ), + pre_command_timeout_seconds=21_600, + attention_backends=_GH200_RELEASE_BACKENDS, + ), + "python-matrix": Suite( + ("candidate",), + _RUN_PYTHON_MATRIX, + ), +} + + +@dataclass(frozen=True) +class RunnerConfig: + """Runtime-only remote connection and execution settings.""" + + host: str + identity: Path + repository: Path + suite: str = "check" + artifacts: Path = Path("artifacts/remote") + accept_new_host_key: bool = False + keep_remote: bool = False + remote_parent: str | None = None + + +def _run_id(_repository: Path) -> str: + timestamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%SZ") + return f"{timestamp}-{secrets.token_hex(8)}" + + +def _is_sensitive(path: PurePosixPath) -> bool: + lowered = tuple(part.lower() for part in path.parts) + return ( + any(part in SENSITIVE_NAMES for part in lowered) + or path.suffix.lower() in SENSITIVE_SUFFIXES + or ".git" in lowered + or "__pycache__" in lowered + ) + + +def _require_matching_archive_digest(output: str, expected_sha256: str) -> None: + fields = output.split() + if ( + re.fullmatch(r"[0-9a-f]{64}", expected_sha256) is None + or not fields + or fields[0] != expected_sha256 + ): + raise RuntimeError("Uploaded source archive SHA-256 differs from local bytes") + + +def _git_files(repository: Path) -> list[Path]: + command = [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "ls-files", + "-z", + "--cached", + ] + completed = subprocess.run(command, cwd=repository, check=True, capture_output=True) + return [Path(raw.decode()) for raw in completed.stdout.split(b"\0") if raw] + + +def _require_clean_repository(repository: Path) -> None: + completed = subprocess.run( + [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "status", + "--porcelain=v1", + "--untracked-files=all", + ], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + if completed.stdout.strip(): + raise RuntimeError( + "Remote runs require a clean Git worktree so the reported revision " + "identifies the exact source." + ) + + +def _require_clean_tracked_repository(repository: Path) -> None: + completed = subprocess.run( + [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "status", + "--porcelain=v1", + "--untracked-files=no", + "--ignore-submodules=all", + ], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + if completed.stdout.strip(): + raise RuntimeError( + "Source archives require clean tracked root files so the content " + "attestation identifies exact bytes." + ) + + +def _git_head_revision(repository: Path) -> str: + completed = subprocess.run( + [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "rev-parse", + "HEAD", + ], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + revision = completed.stdout.strip() + if GIT_REVISION_PATTERN.fullmatch(revision) is None: + raise RuntimeError(f"Git returned an invalid HEAD revision: {revision!r}") + return revision + + +def _is_tracked_file(repository: Path, relative_name: str) -> bool: + completed = subprocess.run( + [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "ls-files", + "--error-unmatch", + "--", + relative_name, + ], + cwd=repository, + check=False, + capture_output=True, + text=True, + ) + return completed.returncode == 0 and completed.stdout.strip() == relative_name + + +def _gitlink_revision(repository: Path, relative_root: Path) -> str: + completed = subprocess.run( + [ + "git", + "-c", + f"safe.directory={repository.resolve().as_posix()}", + "ls-files", + "--stage", + "-z", + "--", + relative_root.as_posix(), + ], + cwd=repository, + check=True, + capture_output=True, + ) + records = [record for record in completed.stdout.split(b"\0") if record] + if len(records) != 1: + raise RuntimeError(f"Expected one Git-link record for {relative_root.as_posix()!r}") + try: + metadata, encoded_path = records[0].split(b"\t", 1) + mode, revision, stage = metadata.decode("ascii").split() + recorded_path = encoded_path.decode() + except (UnicodeDecodeError, ValueError) as error: + raise RuntimeError( + f"Could not parse Git-link record for {relative_root.as_posix()!r}" + ) from error + if ( + mode != "160000" + or stage != "0" + or recorded_path != relative_root.as_posix() + or GIT_REVISION_PATTERN.fullmatch(revision) is None + ): + raise RuntimeError(f"Invalid Git-link record for {relative_root.as_posix()!r}") + return revision + + +def _submodule_files( + repository: Path, + submodule: Path, + relative_root: Path, +) -> tuple[list[tuple[Path, Path]], dict[str, object]]: + git_metadata = submodule / ".git" + if not (git_metadata.exists() or git_metadata.is_symlink()): + raise RuntimeError( + f"Submodule {relative_root.as_posix()!r} is not initialized. Run " + "'git submodule update --init --recursive'." + ) + safe_directory = f"safe.directory={submodule.resolve().as_posix()}" + head = subprocess.run( + ["git", "-c", safe_directory, "rev-parse", "HEAD"], + cwd=submodule, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + gitlink = _gitlink_revision(repository, relative_root) + if GIT_REVISION_PATTERN.fullmatch(head) is None or head != gitlink: + raise RuntimeError( + f"Submodule {relative_root.as_posix()!r} is at {head!r}, " + f"but its Git link records {gitlink!r}." + ) + status = subprocess.run( + ["git", "-c", safe_directory, "status", "--porcelain=v1", "--untracked-files=no"], + cwd=submodule, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if status: + raise RuntimeError( + f"Submodule {relative_root.as_posix()!r} has modified tracked files; " + "source archives require the exact pinned tree." + ) + completed = subprocess.run( + ["git", "-c", safe_directory, "ls-files", "-z"], + cwd=submodule, + check=True, + capture_output=True, + ) + output: list[tuple[Path, Path]] = [] + tracked_files: list[str] = [] + for raw in completed.stdout.split(b"\0"): + if not raw: + continue + child = Path(raw.decode()) + source = submodule / child + if not (source.is_file() or source.is_symlink()): + raise RuntimeError( + f"Tracked submodule path is unavailable or unsupported: " + f"{(relative_root / child).as_posix()}" + ) + child_name = child.as_posix() + if _is_sensitive(PurePosixPath(child_name)): + raise RuntimeError( + f"Submodule tracks forbidden source path: {(relative_root / child).as_posix()}" + ) + output.append((source, relative_root / child)) + tracked_files.append(child_name) + tracked_files.sort() + record: dict[str, object] = { + "file_count": len(tracked_files), + "gitlink_revision": gitlink, + "head_revision": head, + "tracked_files": tracked_files, + "tree_sha256": tracked_tree_digest(submodule, tracked_files), + } + return output, record + + +def create_source_archive( + repository: Path, + destination: Path, +) -> dict[str, dict[str, object]]: + """Archive tracked source plus initialized, pinned submodule tracked files.""" + + repository = repository.resolve() + _require_clean_tracked_repository(repository) + head_revision = _git_head_revision(repository) + files: list[tuple[Path, Path]] = [] + root_tracked_files: list[str] = [] + provenance: dict[str, dict[str, object]] = {} + for relative in _git_files(repository): + source = repository / relative + posix = PurePosixPath(relative.as_posix()) + if _is_sensitive(posix): + raise RuntimeError(f"Repository tracks forbidden source path: {posix.as_posix()!r}") + if posix.as_posix() == ARCHIVE_PROVENANCE_NAME: + raise RuntimeError("Repository may not track the generated source provenance marker") + if not source.exists() and not source.is_symlink(): + if posix.parts[:2] == ("vendor", "upstream"): + continue + raise RuntimeError(f"Tracked source path is unavailable: {posix.as_posix()!r}") + if source.is_file() or source.is_symlink(): + files.append((source, relative)) + root_tracked_files.append(posix.as_posix()) + elif posix.parts[:2] == ("vendor", "upstream"): + git_metadata = source / ".git" + if not (git_metadata.exists() or git_metadata.is_symlink()): + continue + submodule_files, record = _submodule_files(repository, source, relative) + files.extend(submodule_files) + provenance[posix.as_posix()] = record + else: + raise RuntimeError(f"Tracked source path has an unsupported type: {posix.as_posix()!r}") + + root_record = archive_root_record( + repository, + root_tracked_files, + head_revision=head_revision, + ) + + seen: set[str] = set() + with tarfile.open( + destination, + "w:gz", + format=tarfile.PAX_FORMAT, + dereference=False, + ) as archive: + for source, relative in sorted(files, key=lambda item: item[1].as_posix()): + archive_name = relative.as_posix() + if archive_name in seen or _is_sensitive(PurePosixPath(archive_name)): + continue + seen.add(archive_name) + archive.add(source, arcname=archive_name, recursive=False) + provenance_bytes = render_archive_provenance(provenance, root=root_record) + provenance_info = tarfile.TarInfo(ARCHIVE_PROVENANCE_NAME) + provenance_info.size = len(provenance_bytes) + provenance_info.mode = 0o644 + provenance_info.mtime = 0 + provenance_info.uid = 0 + provenance_info.gid = 0 + provenance_info.uname = "" + provenance_info.gname = "" + archive.addfile(provenance_info, io.BytesIO(provenance_bytes)) + _require_clean_tracked_repository(repository) + if _git_head_revision(repository) != head_revision: + destination.unlink(missing_ok=True) + raise RuntimeError("Repository revision changed while creating the source archive") + return provenance + + +def remote_cleanup_command(remote_base: str, remote_workspace: str) -> tuple[str, ...]: + """Build a fail-closed, remote-realpath-verified cleanup command.""" + + base = PurePosixPath(remote_base) + workspace = PurePosixPath(remote_workspace) + if not base.is_absolute() or not workspace.is_absolute(): + raise ValueError("Remote cleanup paths must be absolute") + if ".." in base.parts or ".." in workspace.parts: + raise ValueError("Remote cleanup paths may not contain '..'") + return ( + "sh", + "-c", + REMOTE_CLEANUP_SCRIPT, + "fastplms-cleanup", + str(base), + str(workspace), + ) + + +class RemoteRunner: + """Run one isolated Docker suite and retrieve its artifacts.""" + + def __init__(self, config: RunnerConfig) -> None: + if config.host.startswith("-") or not HOST_PATTERN.fullmatch(config.host): + raise ValueError("SSH host contains unsupported characters") + if config.suite not in SUITES: + raise ValueError(f"Unknown suite {config.suite!r}") + if not config.identity.is_file(): + raise FileNotFoundError(f"SSH identity does not exist: {config.identity}") + self.config = config + self.run_id = _run_id(config.repository) + if not RUN_PATTERN.fullmatch(self.run_id): + raise AssertionError("Generated invalid run ID") + + @property + def ssh_prefix(self) -> list[str]: + options = [ + "ssh", + "-i", + str(self.config.identity), + "-o", + "BatchMode=yes", + "-o", + "IdentitiesOnly=yes", + ] + if self.config.accept_new_host_key: + options.extend(["-o", "StrictHostKeyChecking=accept-new"]) + return options + + @property + def scp_prefix(self) -> list[str]: + options = [ + "scp", + "-i", + str(self.config.identity), + "-o", + "BatchMode=yes", + "-o", + "IdentitiesOnly=yes", + ] + if self.config.accept_new_host_key: + options.extend(["-o", "StrictHostKeyChecking=accept-new"]) + return options + + def _ssh( + self, + command: Sequence[str], + *, + capture: bool = False, + timeout_seconds: int | None = None, + ) -> subprocess.CompletedProcess[str]: + effective_timeout = ( + _CONTROL_TIMEOUT_SECONDS if timeout_seconds is None else timeout_seconds + 60 + ) + return subprocess.run( + [*self.ssh_prefix, self.config.host, shlex.join(command)], + check=True, + text=True, + capture_output=capture, + timeout=effective_timeout, + ) + + def _ssh_at( + self, + workspace: str, + command: Sequence[str], + *, + capture: bool = False, + timeout_seconds: int | None = None, + ) -> subprocess.CompletedProcess[str]: + """Run one cancellable command from ``workspace``. + + GNU ``timeout`` terminates the remote process group, while the slightly + longer local SSH timeout prevents a disconnected client from waiting + forever if the remote host becomes unresponsive. + """ + + script = f"cd {shlex.quote(workspace)} && exec {shlex.join(command)}" + remote_command: tuple[str, ...] + if timeout_seconds is None: + remote_command = ("sh", "-lc", script) + else: + remote_command = ( + "timeout", + "--signal=TERM", + "--kill-after=30s", + f"{timeout_seconds}s", + "sh", + "-lc", + script, + ) + return self._ssh( + remote_command, + capture=capture, + timeout_seconds=timeout_seconds, + ) + + def _capture_host_hardware(self) -> dict[str, object]: + """Capture the exact native architecture and NVIDIA devices before Docker.""" + + machine = self._ssh( + ("uname", "-m"), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ).stdout + gpu_output = self._ssh( + ( + "nvidia-smi", + "--query-gpu=name,uuid,driver_version,memory.total", + "--format=csv,noheader,nounits", + ), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ).stdout + return _host_hardware_preflight(machine, gpu_output) + + def _docker_cache_telemetry(self) -> dict[str, object]: + """Return stable Docker disk/cache counters without command output text.""" + + try: + completed = self._ssh( + ( + "sudo", + "docker", + "system", + "df", + "--format", + "{{json .}}", + ), + capture=True, + timeout_seconds=60, + ) + records: list[dict[str, object]] = [] + allowed_fields = {"Type", "TotalCount", "Active", "Size", "Reclaimable"} + for line in completed.stdout.splitlines(): + if not line.strip(): + continue + raw_record = json.loads(line) + if not isinstance(raw_record, dict): + raise ValueError("Docker cache telemetry record is not an object") + records.append( + { + str(key): value + for key, value in raw_record.items() + if key in allowed_fields and isinstance(value, (str, int, float)) + } + ) + except (OSError, ValueError, subprocess.SubprocessError): + return {"status": "unavailable"} + return {"status": "captured", "records": records} + + def _execution_environment( + self, + workspace: str, + suite: Suite, + host_hardware: Mapping[str, object], + ) -> dict[str, object]: + """Capture exact built image IDs and stable host runtime identities.""" + + current_hardware = self._capture_host_hardware() + if current_hardware != host_hardware: + raise RuntimeError("Remote host hardware identity changed during the build") + container_platform = str(host_hardware["container_platform"]) + + bake = self._ssh_at( + workspace, + ( + "sudo", + "docker", + "buildx", + "bake", + "-f", + "docker/docker-bake.hcl", + "--print", + "--set", + f"*.platform={container_platform}", + *suite.bake_targets, + ), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + bake_plan = json.loads(bake.stdout) + target_plan = bake_plan.get("target") + if not isinstance(target_plan, dict): + raise RuntimeError("Docker Bake did not return a target plan") + + images: dict[str, object] = {} + for target in suite.bake_targets: + raw_target = target_plan.get(target) + if not isinstance(raw_target, dict): + raise RuntimeError(f"Docker Bake omitted target {target!r}") + platforms = raw_target.get("platforms") + if platforms != [container_platform]: + raise RuntimeError( + f"Docker Bake target {target!r} resolved unexpected platforms: {platforms!r}" + ) + tags = raw_target.get("tags") + if not isinstance(tags, list) or not tags or not isinstance(tags[0], str): + raise RuntimeError(f"Docker Bake target {target!r} has no image tag") + inspected = self._ssh( + ("sudo", "docker", "image", "inspect", tags[0]), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + values = json.loads(inspected.stdout) + if ( + not isinstance(values, list) + or len(values) != 1 + or not isinstance(values[0], dict) + ): + raise RuntimeError(f"Docker returned invalid image identity for {target!r}") + value = values[0] + image_id = value.get("Id") + if ( + not isinstance(image_id, str) + or re.fullmatch(r"sha256:[0-9a-f]{64}", image_id) is None + ): + raise RuntimeError(f"Docker returned invalid image digest for {target!r}") + expected_os, expected_architecture = container_platform.split("/", maxsplit=1) + if ( + value.get("Os") != expected_os + or value.get("Architecture") != expected_architecture + ): + raise RuntimeError( + f"Built image {target!r} does not match native platform " + f"{container_platform!r}" + ) + images[target] = { + "tag": tags[0], + "id": image_id, + "repo_digests": value.get("RepoDigests") or [], + "created": value["Created"], + "os": value["Os"], + "architecture": value["Architecture"], + "resolved_platform": container_platform, + "content_digest": image_id, + } + + docker_server = self._ssh( + ("sudo", "docker", "version", "--format", "{{json .Server}}"), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + docker_buildx = self._ssh( + ("sudo", "docker", "buildx", "version"), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + try: + gpu = self._ssh( + ( + "nvidia-smi", + "--query-gpu=name,driver_version", + "--format=csv,noheader", + ), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + gpus = [line.strip() for line in gpu.stdout.splitlines() if line.strip()] + except subprocess.CalledProcessError: + gpus = [] + return { + "host_hardware": dict(host_hardware), + "container_platform": container_platform, + "host_kernel": self._ssh( + ("uname", "-srm"), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ).stdout.strip(), + "docker_server": json.loads(docker_server.stdout), + "docker_buildx": docker_buildx.stdout.strip(), + "gpus": gpus, + "images": images, + } + + def _persist_reference_container_identity( + self, + workspace: str, + execution_environment: Mapping[str, object], + ) -> dict[str, object]: + """Persist stable image identities before any native reference executes.""" + + identity = _reference_container_image_identity(execution_environment) + payload = json.dumps(identity, sort_keys=True, separators=(",", ":")) + self._ssh_at( + workspace, + ("python3", "-c", _WRITE_JSON_SCRIPT, _REFERENCE_IMAGE_IDENTITY_PATH, payload), + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ) + return identity + + def _remote_base(self) -> str: + if self.config.remote_parent is not None: + parent = PurePosixPath(self.config.remote_parent) + if not parent.is_absolute() or ".." in parent.parts: + raise ValueError("--remote-parent must be an absolute path without '..'") + return str(parent) + completed = self._ssh(("pwd",), capture=True) + home = PurePosixPath(completed.stdout.strip()) + if not home.is_absolute() or ".." in home.parts: + raise RuntimeError("Could not determine a safe remote home directory") + return str(home / "fastplms-runs") + + def run(self) -> Path: + started_at = dt.datetime.now(dt.UTC).isoformat() + _require_clean_repository(self.config.repository) + git_revision = _git_head_revision(self.config.repository) + suite = SUITES[self.config.suite] + for relative_name in suite.required_paths: + relative = PurePosixPath(relative_name) + if relative.is_absolute() or ".." in relative.parts: + raise RuntimeError(f"Suite has an unsafe required path: {relative_name!r}") + required = self.config.repository.joinpath(*relative.parts) + if ( + not required.is_file() + or required.is_symlink() + or not _is_tracked_file(self.config.repository, relative_name) + ): + capture_hint = ( + " Run --suite benchmark-capture to produce a descriptive candidate report; " + "review and commit an immutable baseline separately." + if self.config.suite == "benchmark" + else "" + ) + raise RuntimeError( + f"Suite {self.config.suite!r} requires tracked file {relative_name!r}." + + capture_hint + ) + remote_base = self._remote_base() + remote_workspace = str(PurePosixPath(remote_base) / self.run_id) + if not remote_workspace.startswith(remote_base.rstrip("/") + "/"): + raise AssertionError("Remote workspace escaped its managed parent") + output = self.config.artifacts / self.run_id + output.mkdir(parents=True, exist_ok=False) + source_archive_sha256 = "" + submodule_revisions: dict[str, str] = {} + execution_environment: dict[str, object] | None = None + phase = "initialize" + phase_started = time.monotonic() + phase_durations_seconds: dict[str, float] = {} + cache_telemetry: dict[str, object] = {} + artifact_inventory: dict[str, object] | None = None + host_hardware_preflight: dict[str, object] | None = None + kernel_capability_preflight: dict[str, object] | None = None + retrieval_returncode = -1 + cleanup_status = "retained" if self.config.keep_remote else "pending" + cleanup_failure: BaseException | None = None + inventory_failure: BaseException | None = None + remote_workspace_touched = False + remote_workspace_created = False + + def start_phase(next_phase: str) -> None: + nonlocal phase, phase_started + phase_durations_seconds[phase] = ( + phase_durations_seconds.get(phase, 0.0) + time.monotonic() - phase_started + ) + phase = next_phase + phase_started = time.monotonic() + + try: + start_phase("host-hardware-preflight") + host_hardware_preflight = self._capture_host_hardware() + start_phase("kernel-capability-preflight") + kernel_capability_preflight = _kernel_capability_preflight( + host_hardware_preflight, + suite.attention_backends, + ) + if kernel_capability_preflight["status"] != "passed": + raise RuntimeError(str(kernel_capability_preflight["reason"])) + with tempfile.TemporaryDirectory(prefix="fastplms-remote-") as temporary: + start_phase("create-source-archive") + archive = Path(temporary) / "source.tar.gz" + provenance = create_source_archive(self.config.repository, archive) + submodule_revisions = { + path: str(record["head_revision"]) for path, record in provenance.items() + } + _require_clean_repository(self.config.repository) + if _git_head_revision(self.config.repository) != git_revision: + raise RuntimeError( + "Git HEAD changed while the remote source archive was built." + ) + with archive.open("rb") as stream: + source_archive_sha256 = hashlib.file_digest(stream, "sha256").hexdigest() + + start_phase("create-remote-workspace") + remote_workspace_touched = True + self._ssh(("mkdir", "-p", remote_workspace)) + remote_workspace_created = True + + start_phase("upload-source-archive") + subprocess.run( + [ + *self.scp_prefix, + str(archive), + f"{self.config.host}:{remote_workspace}/source.tar.gz", + ], + check=True, + timeout=_TRANSFER_TIMEOUT_SECONDS, + ) + start_phase("verify-source-archive") + remote_digest_output = self._ssh( + ("sha256sum", f"{remote_workspace}/source.tar.gz"), + capture=True, + timeout_seconds=_CONTROL_TIMEOUT_SECONDS, + ).stdout + _require_matching_archive_digest( + remote_digest_output, + source_archive_sha256, + ) + start_phase("extract-source-archive") + self._ssh( + ("tar", "-xzf", f"{remote_workspace}/source.tar.gz", "-C", remote_workspace) + ) + start_phase("remove-source-archive") + self._ssh(("rm", f"{remote_workspace}/source.tar.gz")) + + start_phase("initialize-artifacts") + self._ssh(("mkdir", "-p", f"{remote_workspace}/artifacts/junit")) + start_phase("capture-cache-before-build") + cache_telemetry["before_build"] = self._docker_cache_telemetry() + start_phase("build") + self._ssh_at( + remote_workspace, + ( + "sudo", + "docker", + "buildx", + "bake", + "-f", + "docker/docker-bake.hcl", + "--set", + f"*.platform={host_hardware_preflight['container_platform']}", + *suite.bake_targets, + "--load", + ), + timeout_seconds=suite.build_timeout_seconds, + ) + start_phase("capture-environment") + execution_environment = self._execution_environment( + remote_workspace, + suite, + host_hardware_preflight, + ) + start_phase("persist-reference-container-identity") + execution_environment["reference_container_identity"] = ( + self._persist_reference_container_identity( + remote_workspace, + execution_environment, + ) + ) + for index, command in enumerate(suite.pre_commands): + start_phase(f"pre-command:{index}") + self._ssh_at( + remote_workspace, + command, + timeout_seconds=suite.pre_command_timeout_seconds, + ) + start_phase("suite") + self._ssh_at( + remote_workspace, + suite.command, + timeout_seconds=suite.command_timeout_seconds, + ) + start_phase("complete") + finally: + active_failure = sys.exception() + failure_phase = phase if active_failure is not None else None + start_phase("capture-cache-after-run") + if remote_workspace_created and "before_build" in cache_telemetry: + cache_telemetry["after_run"] = self._docker_cache_telemetry() + start_phase("artifact-retrieval") + if remote_workspace_created: + remote_artifacts = f"{self.config.host}:{remote_workspace}/artifacts/." + try: + retrieval = subprocess.run( + [*self.scp_prefix, "-r", remote_artifacts, str(output)], + check=False, + timeout=_TRANSFER_TIMEOUT_SECONDS, + ) + retrieval_returncode = retrieval.returncode + except subprocess.TimeoutExpired: + retrieval_returncode = 124 + if retrieval_returncode == 0: + start_phase("artifact-inventory") + try: + artifact_inventory = _artifact_tree_summary(output) + except BaseException as error: + inventory_failure = error + artifact_inventory = { + "status": "failed", + "error_type": type(error).__name__, + } + try: + start_phase("cleanup") + if not self.config.keep_remote and remote_workspace_touched: + self._ssh(remote_cleanup_command(remote_base, remote_workspace)) + cleanup_status = "succeeded" + elif not self.config.keep_remote: + cleanup_status = "succeeded" + except BaseException as error: + cleanup_failure = error + cleanup_status = "failed" + if active_failure is None: + raise + finally: + start_phase("report") + report_failure = active_failure or cleanup_failure or inventory_failure + report = _run_report( + run_id=self.run_id, + suite_name=self.config.suite, + suite=suite, + started_at=started_at, + finished_at=dt.datetime.now(dt.UTC).isoformat(), + source_archive_sha256=source_archive_sha256, + git_revision=git_revision, + submodule_revisions=submodule_revisions, + execution_environment=execution_environment, + failure_phase=( + failure_phase + if active_failure is not None + else ( + "cleanup" + if cleanup_failure is not None + else ("artifact-inventory" if inventory_failure is not None else None) + ) + ), + failure=report_failure, + artifact_retrieval_returncode=retrieval_returncode, + cleanup_status=cleanup_status, + phase_durations_seconds=phase_durations_seconds, + cache_telemetry=cache_telemetry, + artifact_inventory=artifact_inventory, + host_hardware_preflight=host_hardware_preflight, + kernel_capability_preflight=kernel_capability_preflight, + ) + report_path = output / "remote-run.json" + temporary_report = output / ".remote-run.json.tmp" + temporary_report.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary_report.replace(report_path) + if retrieval_returncode != 0: + raise RuntimeError(f"Remote artifacts could not be retrieved for run {self.run_id}") + if inventory_failure is not None: + raise RuntimeError( + f"Remote artifacts failed inventory validation for run {self.run_id}" + ) from inventory_failure + return output + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", required=True, help="SSH destination, for example user@gpu-host") + parser.add_argument("--identity", required=True, type=Path, help="SSH private-key path") + parser.add_argument("--repository", type=Path, default=Path.cwd()) + parser.add_argument("--suite", choices=tuple(SUITES), default="check") + parser.add_argument("--artifacts", type=Path, default=Path("artifacts/remote")) + parser.add_argument("--accept-new-host-key", action="store_true") + parser.add_argument("--keep-remote", action="store_true") + parser.add_argument( + "--remote-parent", + help="Optional absolute managed directory on the remote host", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + config = RunnerConfig( + host=arguments.host, + identity=arguments.identity, + repository=arguments.repository, + suite=arguments.suite, + artifacts=arguments.artifacts, + accept_new_host_key=arguments.accept_new_host_key, + keep_remote=arguments.keep_remote, + remote_parent=arguments.remote_parent, + ) + output = RemoteRunner(config).run() + print(output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote/runtime_import_closure.py b/tools/remote/runtime_import_closure.py new file mode 100644 index 0000000..d650b29 --- /dev/null +++ b/tools/remote/runtime_import_closure.py @@ -0,0 +1,502 @@ +"""Statically attest the declared dependency closure of FastPLMs runtime source.""" + +from __future__ import annotations + +import argparse +import ast +import json +import re +import sys +import tomllib +from collections import defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + + +_REQUIREMENT_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*") +_DISTRIBUTION_IMPORT_NAMES = { + "biopython": "Bio", + "huggingface-hub": "huggingface_hub", + "msgpack-numpy": "msgpack_numpy", + "scikit-learn": "sklearn", + "transformer-engine": "transformer_engine", + "transformer-engine-cu13": "transformer_engine", + "transformer-engine-torch": "transformer_engine", +} + + +class RuntimeImportClosureError(RuntimeError): + """Runtime source contains an undeclared import-time dependency.""" + + +@dataclass(frozen=True) +class _ImportRecord: + module: str + source: str + line: int + kind: str + guarded: bool + source_scope: str + + +def _normalized_distribution(requirement: str) -> str: + match = _REQUIREMENT_NAME.match(requirement) + if match is None: + raise RuntimeImportClosureError(f"Invalid dependency requirement: {requirement!r}") + return re.sub(r"[-_.]+", "-", match.group(0)).lower() + + +def _import_name(distribution: str) -> str: + return _DISTRIBUTION_IMPORT_NAMES.get(distribution, distribution.replace("-", "_")) + + +def _read_direct_requirements(path: Path) -> list[str]: + if not path.is_file(): + raise RuntimeImportClosureError(f"Missing direct dependency declaration: {path}") + requirements: list[str] = [] + for line_number, raw_line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + requirement = raw_line.partition("#")[0].strip() + if not requirement: + continue + if requirement.startswith(("-", "--")): + raise RuntimeImportClosureError( + f"Direct dependency file may not compose other files: {path}:{line_number}" + ) + requirements.append(requirement) + return requirements + + +def _declared_import_scopes(requirements_root: Path) -> dict[str, list[str]]: + scoped_requirements = { + "core": _read_direct_requirements(requirements_root / "core.in"), + } + features_root = requirements_root / "features" + if not features_root.is_dir(): + raise RuntimeImportClosureError( + f"Missing feature dependency declarations: {features_root}" + ) + for path in sorted(features_root.glob("*.in")): + scoped_requirements[f"extra:{path.stem}"] = _read_direct_requirements(path) + + scopes: defaultdict[str, set[str]] = defaultdict(set) + for scope, requirements in scoped_requirements.items(): + for requirement in requirements: + distribution = _normalized_distribution(requirement) + scopes[_import_name(distribution)].add(scope) + return {name: sorted(values) for name, values in sorted(scopes.items())} + + +def _literal_dynamic_import(node: ast.Call) -> str | None: + function = node.func + is_import = ( + isinstance(function, ast.Attribute) + and isinstance(function.value, ast.Name) + and function.value.id == "importlib" + and function.attr == "import_module" + ) or (isinstance(function, ast.Name) and function.id in {"__import__", "import_module"}) + if not is_import or not node.args: + return None + value = node.args[0] + return value.value if isinstance(value, ast.Constant) and isinstance(value.value, str) else None + + +def _family_runtime_scopes(source_root: Path) -> dict[str, tuple[str, ...]]: + """Return manifest-declared dependency scopes for runtime path prefixes.""" + + resolved_source_root = source_root.resolve() + manifest_path = source_root / "models.toml" + if not manifest_path.is_file(): + return {} + manifest = tomllib.loads(manifest_path.read_text(encoding="utf-8")) + families = manifest.get("families", {}) + if not isinstance(families, dict): + raise RuntimeImportClosureError("models.toml [families] must be a table") + + scopes: defaultdict[str, set[str]] = defaultdict(set) + for family_name, raw_family in families.items(): + if not isinstance(raw_family, dict): + raise RuntimeImportClosureError( + f"models.toml family {family_name!r} must be a table" + ) + extra = raw_family.get("extra") + runtime_paths = raw_family.get("runtime_paths") + if not isinstance(extra, str) or not isinstance(runtime_paths, list): + raise RuntimeImportClosureError( + f"models.toml family {family_name!r} needs extra and runtime_paths" + ) + scope = "core" if extra == "core" else f"extra:{extra}" + for raw_path in runtime_paths: + if not isinstance(raw_path, str): + raise RuntimeImportClosureError( + f"models.toml family {family_name!r} has a non-string runtime path" + ) + path = PurePosixPath(raw_path) + if ( + path.is_absolute() + or ".." in path.parts + or path.as_posix() != raw_path + ): + raise RuntimeImportClosureError( + f"models.toml family {family_name!r} has a non-portable runtime path" + ) + resolved_path = source_root.joinpath(*path.parts).resolve() + if ( + not resolved_path.is_relative_to(resolved_source_root) + or not resolved_path.exists() + ): + raise RuntimeImportClosureError( + f"models.toml family {family_name!r} has a runtime path " + "outside the runtime source root" + ) + scopes[raw_path].add(scope) + return {path: tuple(sorted(values)) for path, values in sorted(scopes.items())} + + +def _source_scope( + relative: str, + runtime_scopes: Mapping[str, tuple[str, ...]], +) -> str: + matches = [ + (len(PurePosixPath(prefix).parts), scopes) + for prefix, scopes in runtime_scopes.items() + if relative == prefix or relative.startswith(f"{prefix}/") + ] + if not matches: + return "core" + specificity = max(length for length, _ in matches) + scopes = { + scope + for length, declared in matches + if length == specificity + for scope in declared + } + if "core" in scopes: + return "core" + if len(scopes) != 1: + raise RuntimeImportClosureError( + f"Runtime source {relative!r} maps to ambiguous feature scopes: {sorted(scopes)}" + ) + return next(iter(scopes)) + + +def _caught_exception_names(node: ast.expr | None) -> set[str]: + if node is None: + return {"BaseException"} + if isinstance(node, ast.Name): + return {node.id} + if isinstance(node, ast.Attribute): + return {node.attr} + if isinstance(node, ast.Tuple): + return { + name + for element in node.elts + for name in _caught_exception_names(element) + } + return set() + + +def _catches_import_error(handler: ast.ExceptHandler) -> bool: + return bool( + _caught_exception_names(handler.type) + & {"ImportError", "ModuleNotFoundError", "Exception", "BaseException"} + ) + + +class _RuntimeImportVisitor(ast.NodeVisitor): + """Collect imports while retaining whether execution is feature guarded.""" + + def __init__(self, *, source: str, source_scope: str) -> None: + self.source = source + self.source_scope = source_scope + self.records: list[_ImportRecord] = [] + self._guard_depth = 0 + + def _record(self, module: str, line: int, kind: str) -> None: + top_level = module.partition(".")[0] + if top_level in {"fastplms"} or top_level in sys.stdlib_module_names: + return + self.records.append( + _ImportRecord( + module=top_level, + source=self.source, + line=line, + kind=kind, + guarded=self._guard_depth > 0, + source_scope=self.source_scope, + ) + ) + + def _visit_statements( + self, + statements: Sequence[ast.stmt], + *, + guarded: bool, + ) -> None: + if guarded: + self._guard_depth += 1 + try: + for statement in statements: + self.visit(statement) + finally: + if guarded: + self._guard_depth -= 1 + + def _visit_function( + self, + node: ast.FunctionDef | ast.AsyncFunctionDef, + ) -> None: + for decorator in node.decorator_list: + self.visit(decorator) + for default in (*node.args.defaults, *node.args.kw_defaults): + if default is not None: + self.visit(default) + for argument in ( + *node.args.posonlyargs, + *node.args.args, + *node.args.kwonlyargs, + ): + if argument.annotation is not None: + self.visit(argument.annotation) + if node.args.vararg is not None and node.args.vararg.annotation is not None: + self.visit(node.args.vararg.annotation) + if node.args.kwarg is not None and node.args.kwarg.annotation is not None: + self.visit(node.args.kwarg.annotation) + if node.returns is not None: + self.visit(node.returns) + for type_parameter in getattr(node, "type_params", ()): + self.visit(type_parameter) + self._visit_statements(node.body, guarded=True) + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._visit_function(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + self._visit_function(node) + + def visit_Lambda(self, node: ast.Lambda) -> None: + for default in (*node.args.defaults, *node.args.kw_defaults): + if default is not None: + self.visit(default) + self._guard_depth += 1 + try: + self.visit(node.body) + finally: + self._guard_depth -= 1 + + def _visit_try(self, node: ast.Try | ast.TryStar) -> None: + catches_import_error = any(_catches_import_error(handler) for handler in node.handlers) + self._visit_statements(node.body, guarded=catches_import_error) + for handler in node.handlers: + if handler.type is not None: + self.visit(handler.type) + self._visit_statements(handler.body, guarded=False) + self._visit_statements(node.orelse, guarded=False) + self._visit_statements(node.finalbody, guarded=False) + + def visit_Try(self, node: ast.Try) -> None: + self._visit_try(node) + + def visit_TryStar(self, node: ast.TryStar) -> None: + self._visit_try(node) + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + self._record(alias.name, node.lineno, "static") + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + if node.level == 0 and node.module: + self._record(node.module, node.lineno, "static") + + def visit_Call(self, node: ast.Call) -> None: + module = _literal_dynamic_import(node) + if module is not None: + self._record(module, node.lineno, "dynamic") + self.generic_visit(node) + + +def _record_sort_key(record: Mapping[str, object]) -> tuple[str, str, int, str]: + line = record["line"] + if not isinstance(line, int): + raise RuntimeImportClosureError("Import record line must be an integer") + return ( + str(record["module"]), + str(record["source"]), + line, + str(record["kind"]), + ) + + +def _resolved_record( + record: _ImportRecord, + declared_scopes: Sequence[str], + required_scope: str, +) -> dict[str, object]: + return { + "module": record.module, + "source": record.source, + "line": record.line, + "kind": record.kind, + "source_scope": record.source_scope, + "required_scope": required_scope, + "declared_scopes": list(declared_scopes), + } + + +def inspect_runtime_import_closure( + source_root: Path, + requirements_root: Path, +) -> dict[str, object]: + """Return deterministic, scope-aware closure evidence or fail closed.""" + + scopes = _declared_import_scopes(requirements_root) + runtime_scopes = _family_runtime_scopes(source_root) + records: list[_ImportRecord] = [] + source_files = sorted(source_root.rglob("*.py")) + if not source_files: + raise RuntimeImportClosureError(f"No Python runtime source found under {source_root}") + + for path in source_files: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + relative = path.relative_to(source_root).as_posix() + visitor = _RuntimeImportVisitor( + source=relative, + source_scope=_source_scope(relative, runtime_scopes), + ) + visitor.visit(tree) + records.extend(visitor.records) + + undeclared_static = sorted( + { + record.module + for record in records + if record.kind == "static" and record.module not in scopes + } + ) + if undeclared_static: + raise RuntimeImportClosureError( + "Runtime source has undeclared import-time dependencies: " + f"{undeclared_static}" + ) + undeclared_dynamic = sorted({ + record.module + for record in records + if record.kind == "dynamic" and record.module not in scopes + }) + if undeclared_dynamic: + raise RuntimeImportClosureError( + "Runtime source has undeclared literal dynamic dependencies: " + f"{undeclared_dynamic}" + ) + + resolved: list[tuple[_ImportRecord, dict[str, object]]] = [] + unconditional_mismatches: list[str] = [] + guarded_mismatches: list[str] = [] + ambiguous_guarded: list[str] = [] + for record in records: + declared = scopes[record.module] + if not record.guarded: + required_scope = ( + "core" if "core" in declared else record.source_scope + ) + elif record.source_scope != "core" and record.source_scope in declared: + required_scope = record.source_scope + elif "core" in declared: + required_scope = "core" + elif len(declared) == 1: + required_scope = declared[0] + else: + ambiguous_guarded.append( + f"{record.module}@{record.source}:{record.line} declared={declared}" + ) + continue + + if required_scope not in declared: + mismatch = ( + f"{record.module}@{record.source}:{record.line} " + f"requires={required_scope} declared={declared}" + ) + if record.guarded: + guarded_mismatches.append(mismatch) + else: + unconditional_mismatches.append(mismatch) + continue + resolved.append( + (record, _resolved_record(record, declared, required_scope)) + ) + + if ambiguous_guarded: + raise RuntimeImportClosureError( + "Guarded import does not map to one intended dependency scope: " + f"{sorted(ambiguous_guarded)}" + ) + if unconditional_mismatches: + raise RuntimeImportClosureError( + "Unconditional import dependency scope mismatch: " + f"{sorted(unconditional_mismatches)}" + ) + if guarded_mismatches: + raise RuntimeImportClosureError( + f"Guarded import dependency scope mismatch: {sorted(guarded_mismatches)}" + ) + + import_time = sorted( + (payload for record, payload in resolved if not record.guarded), + key=_record_sort_key, + ) + feature_gated = sorted( + ( + payload + for record, payload in resolved + if record.guarded or payload["required_scope"] != "core" + ), + key=_record_sort_key, + ) + dynamic = sorted( + (record for record in feature_gated if record["kind"] == "dynamic"), + key=_record_sort_key, + ) + return { + "schema_version": 2, + "source_files": len(source_files), + "import_time_dependencies": import_time, + "feature_gated_imports": feature_gated, + "feature_gated_dynamic_imports": dynamic, + "undeclared_import_time_dependencies": [], + "undeclared_literal_dynamic_dependencies": [], + "scope_mismatches": [], + } + + +def _atomic_json(path: Path, payload: Mapping[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-root", type=Path, default=Path("src/fastplms")) + parser.add_argument("--requirements-root", type=Path, default=Path("requirements")) + parser.add_argument( + "--output", + type=Path, + default=Path("artifacts/runtime-import-closure.json"), + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = build_parser().parse_args(argv) + payload = inspect_runtime_import_closure( + arguments.source_root.resolve(), + arguments.requirements_root.resolve(), + ) + _atomic_json(arguments.output, payload) + print(arguments.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/source_provenance.py b/tools/source_provenance.py new file mode 100644 index 0000000..bc44ed7 --- /dev/null +++ b/tools/source_provenance.py @@ -0,0 +1,462 @@ +"""Verify source-only archives without carrying Git metadata. + +The portable remote runner intentionally excludes every ``.git`` entry because +Git configuration can contain credentials or workstation-specific paths. This +module defines the small, non-secret attestations that replace those entries: +exact root tracked paths, modes, sizes, symlink targets, and content digests, +plus parent Git-link and checked-out submodule revisions. Root archive metadata +is a content attestation, not proof of a Git commit; Git-free builders therefore +use a content-addressed runtime revision. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import stat +from collections.abc import Mapping, Sequence +from pathlib import Path, PurePosixPath +from typing import Any + + +ARCHIVE_PROVENANCE_NAME = ".fastplms-source-provenance.json" +ARCHIVE_PROVENANCE_SCHEMA = 2 +_HEX_DIGEST_LENGTH = 64 +_HEX_REVISION_LENGTH = 40 +_TREE_DOMAIN = b"fastplms-tracked-submodule-tree-v1\0" +_ROOT_TREE_DOMAIN = b"fastplms-tracked-root-tree-v1\0" + + +class SourceProvenanceError(RuntimeError): + """Raised when a source archive cannot prove its submodule identity.""" + + +def _frame(digest: Any, label: bytes, payload: bytes) -> None: + digest.update(len(label).to_bytes(4, "big")) + digest.update(label) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + + +def _normalized_paths(values: Sequence[str]) -> tuple[str, ...]: + if isinstance(values, (str, bytes)): + raise SourceProvenanceError("tracked_files must be a sequence of paths") + normalized: list[str] = [] + for value in values: + if not isinstance(value, str) or not value: + raise SourceProvenanceError("tracked_files contains an invalid path") + path = PurePosixPath(value) + if ( + not path.parts + or path.is_absolute() + or ".." in path.parts + or path.as_posix() != value + or "\\" in value + or any(":" in part for part in path.parts) + ): + raise SourceProvenanceError(f"Non-portable tracked path: {value!r}") + if any(part.lower() == ".git" for part in path.parts): + raise SourceProvenanceError(f"Git metadata is forbidden in source archives: {value!r}") + normalized.append(value) + ordered = tuple(sorted(normalized)) + if len(ordered) != len(set(ordered)): + raise SourceProvenanceError("tracked_files contains duplicate paths") + return ordered + + +def _safe_symlink_target(root: Path, path: Path) -> str: + target = os.readlink(path) + if "\\" in target or any(":" in part for part in PurePosixPath(target).parts): + raise SourceProvenanceError(f"Non-portable symlink target is forbidden: {path}") + target_path = Path(target) + if target_path.is_absolute(): + raise SourceProvenanceError(f"Absolute symlink is forbidden: {path}") + resolved_root = root.resolve() + resolved_target = (path.parent / target_path).resolve(strict=False) + try: + resolved_target.relative_to(resolved_root) + except ValueError as error: + raise SourceProvenanceError(f"Symlink escapes its submodule tree: {path}") from error + return target + + +def _tracked_path(root: Path, relative_name: str) -> Path: + """Join one portable path while rejecting traversal through parent links.""" + + candidate = root + parts = PurePosixPath(relative_name).parts + for part in parts[:-1]: + candidate /= part + if candidate.is_symlink(): + raise SourceProvenanceError( + f"Tracked path traverses a symlink: {relative_name!r}" + ) + return candidate / parts[-1] + + +def tracked_tree_digest(root: Path, tracked_files: Sequence[str]) -> str: + """Hash exact tracked files, including portable in-tree symlink targets.""" + + root = root.resolve() + if not root.is_dir(): + raise SourceProvenanceError(f"Tracked tree does not exist: {root}") + paths = _normalized_paths(tracked_files) + digest = hashlib.sha256() + digest.update(_TREE_DOMAIN) + for relative_name in paths: + path = _tracked_path(root, relative_name) + try: + mode = path.lstat().st_mode + except OSError as error: + raise SourceProvenanceError(f"Tracked path is missing: {path}") from error + _frame(digest, b"path", relative_name.encode("utf-8")) + if stat.S_ISLNK(mode): + target = _safe_symlink_target(root, path) + _frame(digest, b"type", b"symlink") + _frame(digest, b"target", target.encode("utf-8")) + continue + if not stat.S_ISREG(mode): + raise SourceProvenanceError(f"Tracked path is not a regular file: {path}") + content = hashlib.sha256() + size = 0 + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + content.update(chunk) + size += len(chunk) + _frame(digest, b"type", b"file") + _frame(digest, b"size", size.to_bytes(8, "big")) + _frame(digest, b"content_sha256", content.digest()) + return digest.hexdigest() + + +def tracked_root_inventory( + root: Path, + tracked_files: Sequence[str], +) -> dict[str, dict[str, object]]: + """Describe exact root tracked bytes, portable modes, and symlink targets.""" + + root = root.resolve() + if not root.is_dir(): + raise SourceProvenanceError(f"Tracked root does not exist: {root}") + result: dict[str, dict[str, object]] = {} + for relative_name in _normalized_paths(tracked_files): + path = _tracked_path(root, relative_name) + try: + mode = path.lstat().st_mode + except OSError as error: + raise SourceProvenanceError(f"Tracked root path is missing: {path}") from error + if stat.S_ISLNK(mode): + result[relative_name] = { + "mode": "120000", + "target": _safe_symlink_target(root, path), + } + continue + if not stat.S_ISREG(mode): + raise SourceProvenanceError(f"Tracked root path is not a regular file: {path}") + content = hashlib.sha256() + size = 0 + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + content.update(chunk) + size += len(chunk) + result[relative_name] = { + "mode": "100755" if mode & stat.S_IXUSR else "100644", + "size": size, + "sha256": content.hexdigest(), + } + return result + + +def root_inventory_digest(inventory: Mapping[str, Mapping[str, object]]) -> str: + """Hash a strict, path-ordered root inventory without trusting JSON rendering.""" + + paths = _normalized_paths(tuple(inventory)) + digest = hashlib.sha256() + digest.update(_ROOT_TREE_DOMAIN) + for relative_name in paths: + record = inventory[relative_name] + _frame(digest, b"path", relative_name.encode("utf-8")) + mode = record.get("mode") + if mode == "120000" and set(record) == {"mode", "target"}: + target = record.get("target") + if not isinstance(target, str): + raise SourceProvenanceError( + f"Tracked root symlink has invalid target: {relative_name!r}" + ) + _frame(digest, b"mode", b"120000") + _frame(digest, b"target", target.encode("utf-8")) + continue + if mode not in {"100644", "100755"} or set(record) != { + "mode", + "size", + "sha256", + }: + raise SourceProvenanceError( + f"Tracked root file has invalid metadata: {relative_name!r}" + ) + size = record.get("size") + sha256 = record.get("sha256") + if ( + isinstance(size, bool) + or not isinstance(size, int) + or size < 0 + or not _valid_hex(sha256, _HEX_DIGEST_LENGTH) + ): + raise SourceProvenanceError( + f"Tracked root file has invalid size or digest: {relative_name!r}" + ) + _frame(digest, b"mode", str(mode).encode("ascii")) + _frame(digest, b"size", size.to_bytes(8, "big")) + _frame(digest, b"content_sha256", bytes.fromhex(str(sha256))) + return digest.hexdigest() + + +def archive_root_record( + root: Path, + tracked_files: Sequence[str], + *, + head_revision: str, +) -> dict[str, object]: + """Create the content attestation embedded in one Git-free source archive.""" + + if not _valid_hex(head_revision, _HEX_REVISION_LENGTH): + raise SourceProvenanceError(f"Invalid root revision: {head_revision!r}") + inventory = tracked_root_inventory(root, tracked_files) + return { + "head_revision": head_revision, + "file_count": len(inventory), + "files": inventory, + "tree_sha256": root_inventory_digest(inventory), + } + + +def actual_tree_paths(root: Path) -> tuple[str, ...]: + """Return every file or symlink present in an extracted submodule tree.""" + + root = root.resolve() + if not root.is_dir(): + raise SourceProvenanceError(f"Archived submodule does not exist: {root}") + result: list[str] = [] + for path in root.rglob("*"): + relative = path.relative_to(root) + if any(part.lower() == ".git" for part in relative.parts): + raise SourceProvenanceError(f"Archived submodule contains Git metadata: {path}") + if path.is_symlink() or path.is_file(): + result.append(relative.as_posix()) + return tuple(sorted(result)) + + +def render_archive_provenance( + submodules: Mapping[str, Mapping[str, object]], + *, + root: Mapping[str, object], +) -> bytes: + """Render a deterministic, credential-free archive provenance record.""" + + value = { + "schema_version": ARCHIVE_PROVENANCE_SCHEMA, + "root": dict(root), + "submodules": {path: dict(record) for path, record in sorted(submodules.items())}, + } + return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode("utf-8") + + +def _load_record(source_root: Path) -> dict[str, Any]: + path = source_root / ARCHIVE_PROVENANCE_NAME + try: + mode = path.lstat().st_mode + except OSError as error: + raise SourceProvenanceError( + f"Git-free source tree is missing archive provenance: {path}" + ) from error + if not stat.S_ISREG(mode): + raise SourceProvenanceError(f"Archive provenance is not a regular file: {path}") + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise SourceProvenanceError( + f"Git-free source tree is missing valid archive provenance: {path}" + ) from error + if not isinstance(value, dict) or set(value) != { + "schema_version", + "root", + "submodules", + }: + raise SourceProvenanceError("Archive provenance has an invalid top-level schema") + if value["schema_version"] != ARCHIVE_PROVENANCE_SCHEMA: + raise SourceProvenanceError("Archive provenance schema version is unsupported") + if not isinstance(value["submodules"], dict): + raise SourceProvenanceError("Archive provenance submodules must be a table") + if not isinstance(value["root"], dict): + raise SourceProvenanceError("Archive provenance root must be a table") + return value + + +def _valid_hex(value: object, length: int) -> bool: + return ( + isinstance(value, str) + and len(value) == length + and all(character in "0123456789abcdef" for character in value) + ) + + +def validate_archived_root( + source_root: Path, +) -> tuple[str, dict[str, dict[str, object]]]: + """Validate all attested root tracked bytes and return their content inventory. + + The recorded Git revision is diagnostic transport metadata. Callers must not + treat it as commit authentication because the extracted tree contains no Git + objects. Content-addressed callers should derive identity from validated + payload bytes instead. + """ + + if source_root.is_symlink(): + raise SourceProvenanceError("Archived source root may not be a symlink") + source_root = source_root.resolve() + provenance = _load_record(source_root) + raw_root = provenance["root"] + expected_fields = {"head_revision", "file_count", "files", "tree_sha256"} + if not isinstance(raw_root, dict) or set(raw_root) != expected_fields: + raise SourceProvenanceError("Archive provenance has an invalid root record") + head_revision = raw_root["head_revision"] + if not _valid_hex(head_revision, _HEX_REVISION_LENGTH): + raise SourceProvenanceError("Archive provenance has an invalid root revision") + raw_files = raw_root["files"] + if not isinstance(raw_files, dict): + raise SourceProvenanceError("Archive provenance root files must be a table") + normalized = _normalized_paths(tuple(raw_files)) + if set(normalized) != set(raw_files): + raise SourceProvenanceError("Archive provenance root file paths are invalid") + inventory: dict[str, dict[str, object]] = {} + for relative_name in normalized: + raw_record = raw_files[relative_name] + if not isinstance(raw_record, dict): + raise SourceProvenanceError( + f"Archive provenance root file record is invalid: {relative_name!r}" + ) + inventory[relative_name] = dict(raw_record) + file_count = raw_root["file_count"] + if ( + isinstance(file_count, bool) + or not isinstance(file_count, int) + or file_count != len(inventory) + ): + raise SourceProvenanceError("Archive provenance root file_count differs") + expected_tree = raw_root["tree_sha256"] + if not _valid_hex(expected_tree, _HEX_DIGEST_LENGTH): + raise SourceProvenanceError("Archive provenance root digest is invalid") + encoded_tree = root_inventory_digest(inventory) + if encoded_tree != expected_tree: + raise SourceProvenanceError("Archive provenance root inventory digest differs") + actual_inventory = tracked_root_inventory(source_root, normalized) + if actual_inventory != inventory: + differing = sorted( + name for name in normalized if actual_inventory.get(name) != inventory.get(name) + ) + raise SourceProvenanceError( + "Archived root tracked bytes or modes differ: " + ", ".join(differing[:10]) + ) + actual_tree = root_inventory_digest(actual_inventory) + if actual_tree != expected_tree: + raise SourceProvenanceError("Archived root tracked-tree digest differs") + return str(head_revision), inventory + + +def validate_archived_submodule( + source_root: Path, + *, + relative_path: str, + expected_revision: str, +) -> None: + """Validate one Git-free submodule against its archived attestation.""" + + normalized = PurePosixPath(relative_path) + if ( + not normalized.parts + or normalized.is_absolute() + or ".." in normalized.parts + or normalized.as_posix() != relative_path + or "\\" in relative_path + or any(":" in part for part in normalized.parts) + ): + raise SourceProvenanceError(f"Invalid archived submodule path: {relative_path!r}") + if not _valid_hex(expected_revision, _HEX_REVISION_LENGTH): + raise SourceProvenanceError(f"Invalid expected revision: {expected_revision!r}") + + provenance = _load_record(source_root.resolve()) + raw_record = provenance["submodules"].get(relative_path) + expected_fields = { + "file_count", + "gitlink_revision", + "head_revision", + "tracked_files", + "tree_sha256", + } + if not isinstance(raw_record, dict) or set(raw_record) != expected_fields: + raise SourceProvenanceError( + f"Archive provenance is missing a complete record for {relative_path!r}" + ) + if ( + raw_record["gitlink_revision"] != expected_revision + or raw_record["head_revision"] != expected_revision + ): + raise SourceProvenanceError( + f"Archived submodule {relative_path!r} does not match {expected_revision}" + ) + tracked_files = raw_record["tracked_files"] + if not isinstance(tracked_files, list): + raise SourceProvenanceError(f"Archived submodule {relative_path!r} has no file list") + normalized_files = _normalized_paths(tracked_files) + file_count = raw_record["file_count"] + if isinstance(file_count, bool) or not isinstance(file_count, int): + raise SourceProvenanceError(f"Archived submodule {relative_path!r} has invalid file_count") + if file_count != len(normalized_files): + raise SourceProvenanceError(f"Archived submodule {relative_path!r} file_count differs") + expected_tree = raw_record["tree_sha256"] + if not _valid_hex(expected_tree, _HEX_DIGEST_LENGTH): + raise SourceProvenanceError(f"Archived submodule {relative_path!r} has invalid digest") + + source_root = source_root.resolve() + checkout = source_root + for part in normalized.parts: + checkout /= part + if checkout.is_symlink(): + raise SourceProvenanceError( + f"Archived submodule path traverses a symlink: {relative_path!r}" + ) + try: + checkout.resolve(strict=False).relative_to(source_root) + except ValueError as error: + raise SourceProvenanceError( + f"Archived submodule escapes the source tree: {relative_path!r}" + ) from error + actual_files = actual_tree_paths(checkout) + if actual_files != normalized_files: + missing = sorted(set(normalized_files).difference(actual_files)) + extra = sorted(set(actual_files).difference(normalized_files)) + raise SourceProvenanceError( + f"Archived submodule {relative_path!r} file inventory differs; " + f"missing={missing[:10]}, extra={extra[:10]}" + ) + actual_digest = tracked_tree_digest(checkout, normalized_files) + if actual_digest != expected_tree: + raise SourceProvenanceError( + f"Archived submodule {relative_path!r} tracked-tree digest differs" + ) + + +__all__ = [ + "ARCHIVE_PROVENANCE_NAME", + "ARCHIVE_PROVENANCE_SCHEMA", + "SourceProvenanceError", + "actual_tree_paths", + "archive_root_record", + "render_archive_provenance", + "root_inventory_digest", + "tracked_root_inventory", + "tracked_tree_digest", + "validate_archived_root", + "validate_archived_submodule", +] diff --git a/tools/typing-baselines/c240d8a.json b/tools/typing-baselines/c240d8a.json new file mode 100644 index 0000000..5838c56 --- /dev/null +++ b/tools/typing-baselines/c240d8a.json @@ -0,0 +1,3184 @@ +{ + "baseline_revision": "c240d8a85eabcf5f73d7cf2618c4191295f1df5b", + "checked_source_files": 148, + "environment": { + "mypy": "1.20.2", + "python": "3.12.13" + }, + "error_count": 1064, + "error_file_count": 88, + "fingerprint_sha256": "4d2422a34fb52dcffa9d112e7203437c726d76ad8798fb14fea77ad3f6317784", + "fingerprints": [ + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "benchmarks/regression.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "benchmarks/run.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bool\"", + "path": "benchmarks/run.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"str\"", + "path": "benchmarks/run.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "benchmarks/run.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"_load_model\" has incompatible type \"SimpleNamespace\"; expected \"Namespace\"", + "path": "benchmarks/suite.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "benchmarks/suite.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"run_case\" has incompatible type \"SimpleNamespace\"; expected \"Namespace\"", + "path": "benchmarks/suite.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "benchmarks/suite.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"list[float]\", target has type \"int\")", + "path": "examples/binder_design_fastplms.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"TorchDataset\" (has type \"Any\")", + "path": "examples/fine_tuning.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"float\"", + "path": "examples/fine_tuning.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"tuple[Any, Any]\"", + "path": "examples/fine_tuning.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "examples/fine_tuning.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.attention._core\" does not explicitly export attribute \"BlockMask\"", + "path": "src/fastplms/attention/__init__.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.attention._core\" does not explicitly export attribute \"create_block_mask\"", + "path": "src/fastplms/attention/__init__.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.attention._core\" does not explicitly export attribute \"flex_attention\"", + "path": "src/fastplms/attention/__init__.py" + }, + { + "code": "attr-defined", + "count": 2, + "message": "\"object\" has no attribute \"flash_attn_func\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "attr-defined", + "count": 2, + "message": "\"object\" has no attribute \"flash_attn_varlen_func\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"object\" has no attribute \"fwd\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"object\" has no attribute \"varlen_fwd\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"Function\" (has type \"Any\")", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "no-untyped-def", + "count": 8, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[object, str | None]\", expected \"tuple[object, str]\")", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"OrderedDict\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "type-arg", + "count": 3, + "message": "Missing type arguments for generic type \"tuple\"", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"get_attention_mask\" untyped", + "path": "src/fastplms/attention/_core.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"Path\" has incompatible type \"SimplePath[Any]\"; expected \"str | PathLike[str]\"", + "path": "src/fastplms/attention/_kernel_lock.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"Path\"", + "path": "src/fastplms/attention/_kernel_lock.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/attention/_kernel_lock.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"FastPLMsAttentionMixin\" has no attribute \"config\"", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"FastPLMsAttentionMixin\" has no attribute \"modules\"", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "call-arg", + "count": 1, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "misc", + "count": 2, + "message": "\"_check_and_adjust_attn_implementation\" undefined in superclass", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/attention/interfaces.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_fasta\" has incompatible type \"Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path\"; expected \"str | Path\"", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"Iterable[str | EmbeddingInput | tuple[str, str]] | Path\", variable has type \"Iterable[str | EmbeddingInput | tuple[str, str]]\")", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "TensorValue? has no attribute \"dtype\"", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "TensorValue? has no attribute \"sha256\"", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "TensorValue? has no attribute \"shape\"", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"tuple[EmbeddingRecord, ...]\" of \"EmbeddingRecord | tuple[EmbeddingRecord, ...]\" has no attribute \"load_tensor\"", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Statement is unreachable", + "path": "src/fastplms/embeddings/runner.py" + }, + { + "code": "misc", + "count": 2, + "message": "Cannot infer type of lambda", + "path": "src/fastplms/embeddings/storage.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "src/fastplms/embeddings/storage.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "TensorValue? has no attribute \"load\"", + "path": "src/fastplms/embeddings/types.py" + }, + { + "code": "override", + "count": 1, + "message": "Signature of \"__getitem__\" incompatible with supertype \"typing.Sequence\"", + "path": "src/fastplms/embeddings/types.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Statement is unreachable", + "path": "src/fastplms/embeddings/types.py" + }, + { + "code": "valid-type", + "count": 3, + "message": "Variable \"src.fastplms.embeddings.types.TensorValue\" is not valid as a type", + "path": "src/fastplms/embeddings/types.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"Iterable[int]\"", + "path": "src/fastplms/models/_diffusion_generation.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/_esm_rotary.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "call-arg", + "count": 5, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 6, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 6, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"T5ForConditionalGeneration\" (has type \"Any\")", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 2, + "message": "Expected iterable as variadic argument", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"AnkhEncoderOutput\"", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-untyped-def", + "count": 14, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-untyped-def", + "count": 4, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "no-untyped-def", + "count": 13, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_init_weights\" untyped", + "path": "src/fastplms/models/ankh/modeling_ankh.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"sample\" in typed context", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/boltz/modeling_boltz2.py" + }, + { + "code": "var-annotated", + "count": 1, + "message": "Need type annotation for \"ref_atoms\"", + "path": "src/fastplms/models/boltz/vb_const.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_attention.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_attentionv2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_outer_product_mean.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_pair_averaging.py" + }, + { + "code": "misc", + "count": 4, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_pairformer.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_transition.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_layers_triangular_mult.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"kernel_triangular_mult\" untyped", + "path": "src/fastplms/models/boltz/vb_layers_triangular_mult.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument after ** must be a mapping, not \"dict[Any, Any] | None\"", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "no-untyped-def", + "count": 3, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "type-arg", + "count": 2, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/boltz/vb_modules_confidencev2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_diffusion_conditioning.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 3 to \"compute\" of \"Potential\" has incompatible type \"dict[str, Any] | None\"; expected \"dict[str, Any]\"", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 3 to \"compute_gradient\" of \"Potential\" has incompatible type \"dict[str, Any] | None\"; expected \"dict[str, Any]\"", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "index", + "count": 5, + "message": "Value of type \"dict[str, Any] | None\" is not indexable", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"c_in\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"c_noise\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"c_out\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"c_skip\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"loss_weight\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 2, + "message": "Call to untyped function \"noise_distribution\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"sample_schedule\" in typed context", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-def", + "count": 11, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"list\"", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"None\" of \"Any | None\" has no attribute \"float\"", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Statement is unreachable", + "path": "src/fastplms/models/boltz/vb_modules_diffusionv2.py" + }, + { + "code": "misc", + "count": 7, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_encodersv2.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"Callable\"", + "path": "src/fastplms/models/boltz/vb_modules_encodersv2.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_transformersv2.py" + }, + { + "code": "misc", + "count": 7, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_trunkv2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_modules_utils.py" + }, + { + "code": "no-untyped-def", + "count": 9, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/boltz/vb_potentials_potentials.py" + }, + { + "code": "override", + "count": 1, + "message": "Signature of \"compute_function\" incompatible with supertype \"Potential\"", + "path": "src/fastplms/models/boltz/vb_potentials_potentials.py" + }, + { + "code": "override", + "count": 1, + "message": "Signature of \"compute_variable\" incompatible with supertype \"Potential\"", + "path": "src/fastplms/models/boltz/vb_potentials_potentials.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"tuple[Any, ...]\" of \"Any | tuple[Any, Any]\" has no attribute \"dim\"", + "path": "src/fastplms/models/boltz/vb_potentials_potentials.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"tuple[Any, ...]\" of \"Any | tuple[Any, Any]\" has no attribute \"sum\"", + "path": "src/fastplms/models/boltz/vb_potentials_potentials.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_tri_attn_attention.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Linear\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_tri_attn_primitives.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/boltz/vb_tri_attn_primitives.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"kernel_triangular_attn\" untyped", + "path": "src/fastplms/models/boltz/vb_tri_attn_primitives.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"softmax_no_cast\" untyped", + "path": "src/fastplms/models/boltz/vb_tri_attn_primitives.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"list[Any] | tuple[Any, ...]\", variable has type \"dict_values[Any, Any]\")", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "misc", + "count": 1, + "message": "Incompatible redefinition (redefinition with type \"Callable[[Any, int, int], Any]\", original type \"partial[Any]\")", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bool\"", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_chunk_slice\" untyped", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_flat_idx_to_idx\" untyped", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_get_minimal_slice_set\" untyped", + "path": "src/fastplms/models/boltz/vb_tri_attn_utils.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"output_attentions\" to \"_attn\" of \"ModifiedEsmSelfAttention\" has incompatible type \"bool | None\"; expected \"bool\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"output_s_max\" to \"_attn\" of \"ModifiedEsmSelfAttention\" has incompatible type \"bool | None\"; expected \"bool\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "arg-type", + "count": 2, + "message": "Argument 5 to \"_manual_attn\" of \"ModifiedEsmSelfAttention\" has incompatible type \"bool | None\"; expected \"bool\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "assignment", + "count": 2, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"tuple[Any]\" has no attribute \"contiguous\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "call-arg", + "count": 5, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmAttention\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmConfig\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmEncoder\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmLayer\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmPreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmSelfAttention\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "misc", + "count": 2, + "message": "Tuple index out of range", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"tuple[Any] | DPLMEncoderOutput\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-call", + "count": 4, + "message": "Call to untyped function \"FAST_DPLM_ENCODER\" in typed context", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-call", + "count": 2, + "message": "Call to untyped function \"ModifiedEsmAttention\" in typed context", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmEncoder\" in typed context", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmLayer\" in typed context", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmSelfAttention\" in typed context", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-def", + "count": 10, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "no-untyped-def", + "count": 7, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "operator", + "count": 2, + "message": "Unsupported left operand type for + (\"None\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[Any, Any, Any, Any, Any]\", expected \"tuple[Any] | DPLMMaskedLMOutput\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[Any, Any, Any, Any]\", expected \"tuple[Any] | DPLMMaskedLMOutput\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"None\" of \"Any | None\" has no attribute \"device\"", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_compute_s_max\" untyped", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "var-annotated", + "count": 1, + "message": "Need type annotation for \"all_tied_weights_keys\" (hint: \"all_tied_weights_keys: dict[, ] = ...\")", + "path": "src/fastplms/models/dplm/modeling_dplm.py" + }, + { + "code": "assignment", + "count": 2, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "call-arg", + "count": 5, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "has-type", + "count": 2, + "message": "Cannot determine type of \"inv_freq\"", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmAttention\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmConfig\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmEncoder\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmLayer\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmPreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmSelfAttention\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"RotaryEmbedding\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-call", + "count": 4, + "message": "Call to untyped function \"FAST_DPLM2_ENCODER\" in typed context", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmAttention\" in typed context", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmEncoder\" in typed context", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmLayer\" in typed context", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"ModifiedEsmSelfAttention\" in typed context", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-def", + "count": 10, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "no-untyped-def", + "count": 7, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "operator", + "count": 2, + "message": "Unsupported left operand type for + (\"None\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[Any, Any, Any, Any, Any]\", expected \"tuple[Any] | DPLM2MaskedLMOutput\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[Any, Any, Any, Any]\", expected \"tuple[Any] | DPLM2MaskedLMOutput\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Right operand of \"or\" is never evaluated", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_compute_s_max\" untyped", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "var-annotated", + "count": 1, + "message": "Need type annotation for \"all_tied_weights_keys\" (hint: \"all_tied_weights_keys: dict[, ] = ...\")", + "path": "src/fastplms/models/dplm2/modeling_dplm2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmTokenizer\" (has type \"Any\")", + "path": "src/fastplms/models/dplm2/tokenization_dplm2.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"Callable\"", + "path": "src/fastplms/models/e1/attention.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"create_block_causal_mask_optimized\" untyped", + "path": "src/fastplms/models/e1/attention.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"create_within_seq_block_mask\" untyped", + "path": "src/fastplms/models/e1/attention.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"direct_block_mask\" untyped", + "path": "src/fastplms/models/e1/attention.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"doc_id_mask\" untyped", + "path": "src/fastplms/models/e1/attention.py" + }, + { + "code": "no-untyped-call", + "count": 3, + "message": "Call to untyped function \"_get_logger\" in typed context", + "path": "src/fastplms/models/e1/cache.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/e1/cache.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"forward\" of \"FAST_E1_ENCODER\" has incompatible type \"**dict[str, Any | list[str] | list[int]]\"; expected \"DynamicCache | None\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"forward\" of \"FAST_E1_ENCODER\" has incompatible type \"**dict[str, Any | list[str] | list[int]]\"; expected \"bool\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 3, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 8, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"E1BatchPreparer\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"E1ModelOutputWithPast\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"E1PreTrainedModel\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"list[float] | list[list[float]]\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-untyped-call", + "count": 7, + "message": "Call to untyped function \"_get_logger\" in typed context", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "no-untyped-def", + "count": 18, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "type-arg", + "count": 3, + "message": "Missing type arguments for generic type \"PathLike\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"None\" of \"Any | None\" has no attribute \"device\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"None\" of \"Any | None\" has no attribute \"shape\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"bool\" of \"Any | bool\" has no attribute \"long\"", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_compute_s_max\" untyped", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"embed_dataset_with_msa\" untyped", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"embed_with_msa\" untyped", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"score_ppll\" untyped", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/e1/modeling_e1.py" + }, + { + "code": "type-arg", + "count": 4, + "message": "Missing type arguments for generic type \"PathLike\"", + "path": "src/fastplms/models/e1/preparation.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/e1/preparation.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"float\"", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"_get_logger\" in typed context", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"CompletedProcess\"", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"list[int]\" of \"Any | list[str] | list[int]\" has no attribute \"max\"", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"list[str]\" of \"Any | list[str] | list[int]\" has no attribute \"max\"", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"predict\" untyped", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"predict_batch\" untyped", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"predict_batch_padded\" untyped", + "path": "src/fastplms/models/e1/retrieval.py" + }, + { + "code": "assignment", + "count": 2, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "call-arg", + "count": 5, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "\"None\" not callable", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmTokenizer\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 4, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "misc", + "count": 2, + "message": "Expected iterable as variadic argument", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"int\"", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmAttention\" in typed context", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmEncoder\" in typed context", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmLayer\" in typed context", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"feed_forward_chunk\" in typed context", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-def", + "count": 8, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-def", + "count": 10, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "no-untyped-def", + "count": 9, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_compute_s_max\" untyped", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_init_weights\" untyped", + "path": "src/fastplms/models/esm2/modeling_fastesm.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"attn_backend\" to \"ESM3Core\" has incompatible type \"str | None\"; expected \"str\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"ZipInfo\" has incompatible type \"PurePosixPath\"; expected \"str\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"encode\" of \"FastESM3Model\" has incompatible type \"str | list[str] | Any | dict[str, Any]\"; expected \"str | list[str]\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"str\", variable has type \"PurePosixPath\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 9, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedTokenizerFast\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"ESM3CoreOutput\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"FastESM3Output\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"int\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"list[int]\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "no-untyped-def", + "count": 11, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"list[str]\" of \"str | Any | list[str]\" has no attribute \"to\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"str\" of \"str | Any | list[str]\" has no attribute \"to\"", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Right operand of \"or\" is never evaluated", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Statement is unreachable", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"generate\" untyped", + "path": "src/fastplms/models/esm3/modeling_esm3.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"attentions\" to \"TransformerOutput\" has incompatible type \"tuple[()] | None\"; expected \"tuple[Any] | None\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "arg-type", + "count": 2, + "message": "Argument \"attn_backend\" to \"TransformerStack\" has incompatible type \"str | None\"; expected \"str\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"hidden_states\" to \"TransformerOutput\" has incompatible type \"tuple[()] | None\"; expected \"tuple[Any] | None\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "assignment", + "count": 3, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "call-arg", + "count": 2, + "message": "Too many arguments for \"__init__\" of \"object\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 4, + "message": "Class cannot subclass \"EmbeddingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 2, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedTokenizerFast\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmSequenceTokenizer\" in typed context", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"SwiGLU\" in typed context", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-untyped-def", + "count": 13, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-untyped-def", + "count": 8, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "no-untyped-def", + "count": 10, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "operator", + "count": 1, + "message": "Unsupported left operand type for + (\"None\")", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "override", + "count": 2, + "message": "Signature of \"forward\" incompatible with supertype \"ESMplusplusForMaskedLM\"", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"_compute_s_max\" untyped", + "path": "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"tuple[Any]\", variable has type \"tuple[()] | None\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "func-returns-value", + "count": 1, + "message": "\"fold_protein_ttt\" of \"FastEsmForProteinFolding\" does not return a value (it only ever returns None)", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 1, + "message": "\"None\" not callable", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmConfig\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"EsmForProteinFolding\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMsAttentionMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"ModelOutput\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "misc", + "count": 2, + "message": "Expected iterable as variadic argument", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmAttention\" in typed context", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"EsmLayer\" in typed context", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"FastEsmBackbone\" in typed context", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"FastEsmEncoder\" in typed context", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-def", + "count": 4, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"None\", expected \"dict[str, Any]\")", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"list[Any]\" of \"Any | list[Any]\" has no attribute \"to\"", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"infer\" untyped", + "path": "src/fastplms/models/esmfold/modeling_fast_esmfold.py" + }, + { + "code": "attr-defined", + "count": 2, + "message": "\"ESMFold2AttentionMixin\" has no attribute \"config\"", + "path": "src/fastplms/models/esmfold2/attention.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/attention.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PretrainedConfig\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/configuration_esmfold2.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"ESMFold2EmbeddingMixin\" has no attribute \"_compute_lm_hidden_states\"", + "path": "src/fastplms/models/esmfold2/embedding.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"ESMFold2EmbeddingMixin\" has no attribute \"_esmc\"", + "path": "src/fastplms/models/esmfold2/embedding.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"ESMFold2EmbeddingMixin\" has no attribute \"device\"", + "path": "src/fastplms/models/esmfold2/embedding.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"ESMFold2EmbeddingMixin\" has no attribute \"language_model\"", + "path": "src/fastplms/models/esmfold2/embedding.py" + }, + { + "code": "empty-body", + "count": 8, + "message": "Missing return statement", + "path": "src/fastplms/models/esmfold2/esmfold2_affine3d.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"bool\"", + "path": "src/fastplms/models/esmfold2/esmfold2_affine3d.py" + }, + { + "code": "no-untyped-def", + "count": 14, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_affine3d.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"float\"", + "path": "src/fastplms/models/esmfold2/esmfold2_aligner.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.models.esmfold2.esmfold2_protein_structure\" does not explicitly export attribute \"index_by_atom_name\"", + "path": "src/fastplms/models/esmfold2/esmfold2_atom_indexer.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_conformers.py" + }, + { + "code": "untyped-decorator", + "count": 2, + "message": "Untyped decorator makes function \"compute_rmsd\" untyped", + "path": "src/fastplms/models/esmfold2/esmfold2_metrics.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_metrics.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"concat\" of \"Concatable\" has incompatible type \"Sequence[Any]\"; expected \"list[Concatable]\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"Sequence[Any]\", variable has type \"list[Any]\")", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"TSequence\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"list[Any] | None\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"batched_gather\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-untyped-def", + "count": 4, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-untyped-def", + "count": 3, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-untyped-def", + "count": 3, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"Iterable\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"Sequence\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"list\"", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "unused-ignore", + "count": 2, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore[index, return-value]\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore[return-value]\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "var-annotated", + "count": 1, + "message": "Need type annotation for \"result\" (hint: \"result: list[] = ...\")", + "path": "src/fastplms/models/esmfold2/esmfold2_misc.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"_scheme_columns\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"_scheme_residue_map\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "no-untyped-def", + "count": 5, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"PathLike\"", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "type-arg", + "count": 3, + "message": "Missing type arguments for generic type \"tuple\"", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "var-annotated", + "count": 2, + "message": "Need type annotation for \"result\" (hint: \"result: dict[, ] = ...\")", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "var-annotated", + "count": 1, + "message": "Need type annotation for \"result\" (hint: \"result: set[] = ...\")", + "path": "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py" + }, + { + "code": "no-redef", + "count": 1, + "message": "Name \"structure\" already defined on line 703", + "path": "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py" + }, + { + "code": "no-untyped-call", + "count": 4, + "message": "Call to untyped function \"normalize_chain_ids_for_pdb\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"int\"", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "override", + "count": 2, + "message": "Signature of \"concat\" incompatible with supertype \"src.fastplms.models.esmfold2.esmfold2_sequential_dataclass.SequentialDataclass\"", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "var-annotated", + "count": 2, + "message": "Need type annotation for \"entries\" (hint: \"entries: list[] = ...\")", + "path": "src/fastplms/models/esmfold2/esmfold2_msa.py" + }, + { + "code": "index", + "count": 1, + "message": "Invalid index type \"str | list[str]\" for \"dict[str, int]\"; expected type \"str\"", + "path": "src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_output.py" + }, + { + "code": "no-untyped-def", + "count": 2, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_parsing.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"compute_tm\" untyped", + "path": "src/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py" + }, + { + "code": "assignment", + "count": 2, + "message": "Incompatible types in assignment (expression has type \"int\", variable has type \"TokenInfo\")", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"tuple[int, ...]\", target has type \"tuple[int, int, int]\")", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "call-overload", + "count": 1, + "message": "No overload variant of \"get\" of \"dict\" matches argument type \"TokenInfo\"", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "comparison-overlap", + "count": 1, + "message": "Non-overlapping container check (element type: \"TokenInfo\", container item type: \"int\")", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "index", + "count": 1, + "message": "Invalid index type \"TokenInfo\" for \"dict[int, tuple[int, int, int]]\"; expected type \"int\"", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "unreachable", + "count": 1, + "message": "Statement is unreachable", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "unused-ignore", + "count": 5, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_prepare_input.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 has incompatible type \"object\"; expected \"tuple[Any, ...]\"", + "path": "src/fastplms/models/esmfold2/esmfold2_processor.py" + }, + { + "code": "arg-type", + "count": 2, + "message": "Argument 1 to \"read\" of \"MmcifWrapper\" has incompatible type \"str | PathLike[str] | TextIOBase\"; expected \"str | PathLike[Any] | StringIO\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.models.esmfold2.esmfold2_protein_structure\" does not explicitly export attribute \"index_by_atom_name\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"ProteinChain\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"float\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"from_state_dict\" of \"ProteinChain\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"state_dict\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-untyped-def", + "count": 21, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-untyped-def", + "count": 3, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "no-untyped-def", + "count": 7, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "unused-ignore", + "count": 4, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_chain.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"chain_mapping\" to \"DockQResult\" has incompatible type \"float\"; expected \"dict[str, str]\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument \"native_interfaces\" to \"DockQResult\" has incompatible type \"float\"; expected \"int\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"read\" of \"MmcifWrapper\" has incompatible type \"str | PathLike[str] | TextIOBase\"; expected \"str | PathLike[Any] | StringIO\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"DockQResult\", variable has type \"dict[str, float]\")", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"dict[Any, Any]\", target has type \"float\")", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"float\" has no attribute \"items\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "Module \"src.fastplms.models.esmfold2.esmfold2_protein_chain\" does not explicitly export attribute \"index_by_atom_name\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"ProteinChain\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-any-return", + "count": 6, + "message": "Returning Any from function declared to return \"ProteinComplex\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-redef", + "count": 1, + "message": "Name \"structure\" already defined on line 1108", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"_apply_transformations_fast\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"from_state_dict\" of \"ProteinComplex\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-call", + "count": 1, + "message": "Call to untyped function \"state_dict\" in typed context", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-def", + "count": 19, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-def", + "count": 4, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "no-untyped-def", + "count": 4, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "unused-ignore", + "count": 7, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_complex.py" + }, + { + "code": "untyped-decorator", + "count": 2, + "message": "Untyped decorator makes function \"compute_affine_and_rmsd\" untyped", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_structure.py" + }, + { + "code": "untyped-decorator", + "count": 2, + "message": "Untyped decorator makes function \"compute_alignment_tensors\" untyped", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_structure.py" + }, + { + "code": "untyped-decorator", + "count": 2, + "message": "Untyped decorator makes function \"compute_rmsd_no_alignment\" untyped", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_structure.py" + }, + { + "code": "unused-ignore", + "count": 2, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/esmfold2_protein_structure.py" + }, + { + "code": "valid-type", + "count": 3, + "message": "Variable \"src.fastplms.models.esmfold2.esmfold2_sequential_dataclass.Index\" is not valid as a type", + "path": "src/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py" + }, + { + "code": "call-overload", + "count": 1, + "message": "No overload variant of \"run\" matches argument types \"tuple[Any, ...]\", \"bool\", \"dict[str, str] | None\", \"str | None\", \"bool\", \"int\", \"int | None\", \"dict[str, Any]\"", + "path": "src/fastplms/models/esmfold2/esmfold2_system.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"CompletedProcess[Any]\"", + "path": "src/fastplms/models/esmfold2/esmfold2_system.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 2 to \"_fold_protein_no_ttt\" of \"ESMFold2Model\" has incompatible type \"**dict[str, Any | str | Path | int | None]\"; expected \"int | None\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 2 to \"_fold_protein_no_ttt\" of \"ESMFold2Model\" has incompatible type \"**dict[str, Any | str | Path | int | None]\"; expected \"int\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 2 to \"_fold_protein_no_ttt\" of \"ESMFold2Model\" has incompatible type \"**dict[str, Any | str | Path | int | None]\"; expected \"str | Path | None\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 2 to \"_fold_protein_no_ttt\" of \"ESMFold2Model\" has incompatible type \"**dict[str, Any | str | Path | int | None]\"; expected \"str\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"FastPLMTestTimeTrainingMixin\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "misc", + "count": 5, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"_ESMFold2ESMplusplusAdapter\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[Any, Any]\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-untyped-call", + "count": 2, + "message": "Call to untyped function \"_ttt_select_result\" in typed context", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-untyped-def", + "count": 11, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a type annotation", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "no-untyped-def", + "count": 15, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "type-arg", + "count": 4, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"forward\" untyped", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"infer_protein\" untyped", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment, use narrower [method-assign] instead of [assignment] code", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2.py" + }, + { + "code": "misc", + "count": 29, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"tuple[Any, Any]\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "no-untyped-def", + "count": 1, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "type-arg", + "count": 1, + "message": "Missing type arguments for generic type \"dict\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "type-arg", + "count": 5, + "message": "Missing type arguments for generic type \"tuple\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "unreachable", + "count": 2, + "message": "Statement is unreachable", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"build_3d_rope\" untyped", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"sample\" untyped", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "unused-ignore", + "count": 7, + "message": "Unused \"type: ignore\" comment", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_common.py" + }, + { + "code": "method-assign", + "count": 1, + "message": "Cannot assign to a method", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "misc", + "count": 4, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"PreTrainedModel\" (has type \"Any\")", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"str\"", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "no-untyped-def", + "count": 6, + "message": "Function is missing a return type annotation", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "no-untyped-def", + "count": 10, + "message": "Function is missing a type annotation for one or more parameters", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "untyped-decorator", + "count": 1, + "message": "Untyped decorator makes function \"infer_protein\" untyped", + "path": "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"tuple[float, ...]\", target has type \"tuple[float, float, float]\")", + "path": "src/fastplms/models/esmfold2/protein_utils.py" + }, + { + "code": "return-value", + "count": 1, + "message": "Incompatible return value type (got \"tuple[int, ...]\", expected \"tuple[int, int, int, int]\")", + "path": "src/fastplms/models/esmfold2/protein_utils.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"FastPLMTestTimeTrainingMixin\" has no attribute \"config\"", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "attr-defined", + "count": 2, + "message": "\"FastPLMTestTimeTrainingMixin\" has no attribute \"modules\"", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "attr-defined", + "count": 3, + "message": "\"FastPLMTestTimeTrainingMixin\" has no attribute \"parameters\"", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "attr-defined", + "count": 5, + "message": "\"FastPLMTestTimeTrainingMixin\" has no attribute \"tokenizer\"", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"FastPLMTestTimeTrainingMixin\" has no attribute \"train\"", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "misc", + "count": 1, + "message": "Class cannot subclass \"Module\" (has type \"Any\")", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "operator", + "count": 2, + "message": "\"FastPLMTestTimeTrainingMixin\" not callable", + "path": "src/fastplms/models/ttt.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"validate_hub_license_metadata\" has incompatible type \"Any | None\"; expected \"Mapping[str, object]\"", + "path": "tools/artifacts/build.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"Path\"", + "path": "tools/artifacts/build.py" + }, + { + "code": "union-attr", + "count": 1, + "message": "Item \"None\" of \"Any | None\" has no attribute \"items\"", + "path": "tools/artifacts/build.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "tools/artifacts/generate_docs.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"str\"", + "path": "tools/artifacts/generate_docs.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "tools/artifacts/offline_probe.py" + }, + { + "code": "attr-defined", + "count": 3, + "message": "\"type\" has no attribute \"from_pretrained\"", + "path": "tools/artifacts/offline_probe.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"dict[str, Any]\"", + "path": "tools/artifacts/offline_probe.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"type\"", + "path": "tools/artifacts/offline_probe.py" + }, + { + "code": "attr-defined", + "count": 1, + "message": "\"object\" has no attribute \"detach\"", + "path": "tools/conversion/state_transforms.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "tools/debug/check_notation.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"tuple[tuple[str, int]]\", variable has type \"Iterator[tuple[str, int]]\")", + "path": "tools/debug/check_notation.py" + }, + { + "code": "assignment", + "count": 1, + "message": "Incompatible types in assignment (expression has type \"str\", variable has type \"int\")", + "path": "tools/debug/compare_structure_bundles.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"literal_eval\" has incompatible type \"expr | None\"; expected \"str | AST\"", + "path": "tools/debug/generate_boltz_conformer_patch.py" + }, + { + "code": "assignment", + "count": 2, + "message": "Incompatible types in assignment (expression has type \"expr | None\", variable has type \"Name | Attribute | Subscript\")", + "path": "tools/debug/generate_boltz_conformer_patch.py" + }, + { + "code": "index", + "count": 2, + "message": "Unsupported target for indexed assignment (\"Collection[str]\")", + "path": "tools/debug/probe_flash_kernels.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"_limitation_records\" has incompatible type \"object\"; expected \"Sequence[Mapping[str, str]] | None\"", + "path": "tools/goldens/bundle.py" + }, + { + "code": "misc", + "count": 1, + "message": "Value expression in dictionary comprehension has incompatible type \"Sequence[Any]\"; expected type \"str\"", + "path": "tools/goldens/bundle.py" + }, + { + "code": "no-any-return", + "count": 2, + "message": "Returning Any from function declared to return \"Path\"", + "path": "tools/goldens/from_native.py" + }, + { + "code": "no-any-return", + "count": 1, + "message": "Returning Any from function declared to return \"bytes\"", + "path": "tools/goldens/from_native.py" + }, + { + "code": "unused-ignore", + "count": 1, + "message": "Unused \"type: ignore\" comment", + "path": "tools/goldens/generate_e1_sampling.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "tools/remote/python_matrix.py" + }, + { + "code": "arg-type", + "count": 1, + "message": "Argument 1 to \"parse_args\" of \"ArgumentParser\" has incompatible type \"Iterable[str] | None\"; expected \"Sequence[str] | None\"", + "path": "tools/remote/run.py" + } + ], + "mypy_command": [ + "python", + "-m", + "mypy", + "--config-file=/dev/null", + "--python-version", + "3.12", + "--strict", + "--warn-unreachable", + "--ignore-missing-imports", + "--no-site-packages", + "--no-incremental", + "--explicit-package-bases", + "--follow-imports=silent", + "--show-error-codes", + "--no-color-output", + "--no-pretty", + "benchmarks", + "examples", + "src/fastplms", + "tools" + ], + "mypy_exit_code": 1, + "raw_report_sha256": "fe3e51a917b33c7650943340722258a518eae61c77bd1ebe0efd9384ad1f5a62", + "schema_version": 1, + "scope_targets": [ + "benchmarks", + "examples", + "src/fastplms", + "tools" + ], + "source_files": [ + "benchmarks/__init__.py", + "benchmarks/__main__.py", + "benchmarks/regression.py", + "benchmarks/run.py", + "benchmarks/suite.py", + "examples/__init__.py", + "examples/binder_design_fastplms.py", + "examples/fine_tuning.py", + "src/fastplms/__init__.py", + "src/fastplms/attention/__init__.py", + "src/fastplms/attention/_core.py", + "src/fastplms/attention/_kernel_lock.py", + "src/fastplms/attention/interfaces.py", + "src/fastplms/embeddings/__init__.py", + "src/fastplms/embeddings/pooling.py", + "src/fastplms/embeddings/runner.py", + "src/fastplms/embeddings/storage.py", + "src/fastplms/embeddings/types.py", + "src/fastplms/models/__init__.py", + "src/fastplms/models/_diffusion_generation.py", + "src/fastplms/models/_esm_rotary.py", + "src/fastplms/models/ankh/__init__.py", + "src/fastplms/models/ankh/modeling_ankh.py", + "src/fastplms/models/boltz/__init__.py", + "src/fastplms/models/boltz/_pair_attention.py", + "src/fastplms/models/boltz/cif_writer.py", + "src/fastplms/models/boltz/minimal_featurizer.py", + "src/fastplms/models/boltz/minimal_structures.py", + "src/fastplms/models/boltz/modeling_boltz2.py", + "src/fastplms/models/boltz/vb_const.py", + "src/fastplms/models/boltz/vb_layers_attention.py", + "src/fastplms/models/boltz/vb_layers_attentionv2.py", + "src/fastplms/models/boltz/vb_layers_confidence_utils.py", + "src/fastplms/models/boltz/vb_layers_dropout.py", + "src/fastplms/models/boltz/vb_layers_initialize.py", + "src/fastplms/models/boltz/vb_layers_outer_product_mean.py", + "src/fastplms/models/boltz/vb_layers_pair_averaging.py", + "src/fastplms/models/boltz/vb_layers_pairformer.py", + "src/fastplms/models/boltz/vb_layers_transition.py", + "src/fastplms/models/boltz/vb_layers_triangular_mult.py", + "src/fastplms/models/boltz/vb_loss_diffusionv2.py", + "src/fastplms/models/boltz/vb_modules_confidencev2.py", + "src/fastplms/models/boltz/vb_modules_diffusion_conditioning.py", + "src/fastplms/models/boltz/vb_modules_diffusionv2.py", + "src/fastplms/models/boltz/vb_modules_encodersv2.py", + "src/fastplms/models/boltz/vb_modules_transformersv2.py", + "src/fastplms/models/boltz/vb_modules_trunkv2.py", + "src/fastplms/models/boltz/vb_modules_utils.py", + "src/fastplms/models/boltz/vb_potentials_potentials.py", + "src/fastplms/models/boltz/vb_potentials_schedules.py", + "src/fastplms/models/boltz/vb_tri_attn_attention.py", + "src/fastplms/models/boltz/vb_tri_attn_primitives.py", + "src/fastplms/models/boltz/vb_tri_attn_utils.py", + "src/fastplms/models/dplm/__init__.py", + "src/fastplms/models/dplm/modeling_dplm.py", + "src/fastplms/models/dplm2/__init__.py", + "src/fastplms/models/dplm2/modeling_dplm2.py", + "src/fastplms/models/dplm2/tokenization_dplm2.py", + "src/fastplms/models/e1/__init__.py", + "src/fastplms/models/e1/attention.py", + "src/fastplms/models/e1/cache.py", + "src/fastplms/models/e1/modeling_e1.py", + "src/fastplms/models/e1/preparation.py", + "src/fastplms/models/e1/retrieval.py", + "src/fastplms/models/esm2/__init__.py", + "src/fastplms/models/esm2/modeling_fastesm.py", + "src/fastplms/models/esm3/__init__.py", + "src/fastplms/models/esm3/modeling_esm3.py", + "src/fastplms/models/esm_plusplus/__init__.py", + "src/fastplms/models/esm_plusplus/modeling_esm_plusplus.py", + "src/fastplms/models/esmfold/__init__.py", + "src/fastplms/models/esmfold/modeling_fast_esmfold.py", + "src/fastplms/models/esmfold2/__init__.py", + "src/fastplms/models/esmfold2/attention.py", + "src/fastplms/models/esmfold2/configuration_esmfold2.py", + "src/fastplms/models/esmfold2/embedding.py", + "src/fastplms/models/esmfold2/esmfold2_affine3d.py", + "src/fastplms/models/esmfold2/esmfold2_aligner.py", + "src/fastplms/models/esmfold2/esmfold2_atom_indexer.py", + "src/fastplms/models/esmfold2/esmfold2_conformers.py", + "src/fastplms/models/esmfold2/esmfold2_constants.py", + "src/fastplms/models/esmfold2/esmfold2_constants_esm3.py", + "src/fastplms/models/esmfold2/esmfold2_input_builder.py", + "src/fastplms/models/esmfold2/esmfold2_metrics.py", + "src/fastplms/models/esmfold2/esmfold2_misc.py", + "src/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py", + "src/fastplms/models/esmfold2/esmfold2_molecular_complex.py", + "src/fastplms/models/esmfold2/esmfold2_msa.py", + "src/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py", + "src/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py", + "src/fastplms/models/esmfold2/esmfold2_output.py", + "src/fastplms/models/esmfold2/esmfold2_paired_msa.py", + "src/fastplms/models/esmfold2/esmfold2_parsing.py", + "src/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py", + "src/fastplms/models/esmfold2/esmfold2_prepare_input.py", + "src/fastplms/models/esmfold2/esmfold2_processor.py", + "src/fastplms/models/esmfold2/esmfold2_protein_chain.py", + "src/fastplms/models/esmfold2/esmfold2_protein_complex.py", + "src/fastplms/models/esmfold2/esmfold2_protein_structure.py", + "src/fastplms/models/esmfold2/esmfold2_residue_constants.py", + "src/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py", + "src/fastplms/models/esmfold2/esmfold2_system.py", + "src/fastplms/models/esmfold2/esmfold2_types.py", + "src/fastplms/models/esmfold2/esmfold2_utils_types.py", + "src/fastplms/models/esmfold2/modeling_esmfold2.py", + "src/fastplms/models/esmfold2/modeling_esmfold2_common.py", + "src/fastplms/models/esmfold2/modeling_esmfold2_experimental.py", + "src/fastplms/models/esmfold2/protein_utils.py", + "src/fastplms/models/ttt.py", + "src/fastplms/registry.py", + "src/fastplms/runtime.py", + "tools/artifacts/__init__.py", + "tools/artifacts/build.py", + "tools/artifacts/build_all.py", + "tools/artifacts/generate_docs.py", + "tools/artifacts/license_metadata.py", + "tools/artifacts/offline_probe.py", + "tools/artifacts/publish.py", + "tools/artifacts/resolve_fair_esm_assets.py", + "tools/artifacts/resolve_manifest_hashes.py", + "tools/conversion/__init__.py", + "tools/conversion/extract_esmfold2_geometry.py", + "tools/conversion/state_transforms.py", + "tools/conversion/state_validation.py", + "tools/debug/analyze_boltz_opm_projection.py", + "tools/debug/check_notation.py", + "tools/debug/compare_structure_bundles.py", + "tools/debug/export_boltz_conformers.py", + "tools/debug/generate_boltz_conformer_patch.py", + "tools/debug/probe_cuda_loader.py", + "tools/debug/probe_flash_attention_forward.py", + "tools/debug/probe_flash_checkpoint_forward.py", + "tools/debug/probe_flash_kernels.py", + "tools/debug/probe_transformer_engine.py", + "tools/debug/probe_transformer_engine_fp8.py", + "tools/debug/trace_boltz2_modules.py", + "tools/goldens/__init__.py", + "tools/goldens/__main__.py", + "tools/goldens/bundle.py", + "tools/goldens/from_native.py", + "tools/goldens/generate_e1_sampling.py", + "tools/remote/__init__.py", + "tools/remote/__main__.py", + "tools/remote/prepare_references.py", + "tools/remote/python_matrix.py", + "tools/remote/python_support_smoke.py", + "tools/remote/run.py", + "tools/source_provenance.py" + ], + "source_inventory_sha256": "f4ec439a48f012353d2f4932dda6c46db09178869b428a6911d1da604d1975bf", + "source_tree_sha256": "9732188015f6fa29166428cf9c3320ee7b8045aa9cec52bf154dba671c5e65b3" +} diff --git a/tools/typing-critical-files.txt b/tools/typing-critical-files.txt new file mode 100644 index 0000000..34da434 --- /dev/null +++ b/tools/typing-critical-files.txt @@ -0,0 +1,5 @@ +src/fastplms/registry.py +tools/artifacts/build.py +tools/artifacts/publish.py +tools/conversion/state_transforms.py +tools/source_provenance.py diff --git a/tools/typing-diagnostic-files.txt b/tools/typing-diagnostic-files.txt new file mode 100644 index 0000000..9ae8747 --- /dev/null +++ b/tools/typing-diagnostic-files.txt @@ -0,0 +1,4 @@ +benchmarks +examples +src/fastplms +tools diff --git a/tools/typing_gate.py b/tools/typing_gate.py new file mode 100644 index 0000000..a4189a8 --- /dev/null +++ b/tools/typing_gate.py @@ -0,0 +1,960 @@ +"""Deterministic no-regression gate for the broad FastPLMs mypy surface.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import os +import platform +import re +import subprocess +import sys +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + + +SCHEMA_VERSION = 1 +BASELINE_REVISION = "c240d8a85eabcf5f73d7cf2618c4191295f1df5b" +BASELINE_PYTHON_VERSION = "3.12.13" +BASELINE_MYPY_VERSION = "1.20.2" +BASELINE_CHECKED_SOURCE_FILES = 148 +BASELINE_ERROR_COUNT = 1064 +BASELINE_ERROR_FILE_COUNT = 88 +BASELINE_SOURCE_INVENTORY_SHA256 = ( + "f4ec439a48f012353d2f4932dda6c46db09178869b428a6911d1da604d1975bf" +) +BASELINE_SOURCE_TREE_SHA256 = ( + "9732188015f6fa29166428cf9c3320ee7b8045aa9cec52bf154dba671c5e65b3" +) +BASELINE_RAW_REPORT_SHA256 = ( + "fe3e51a917b33c7650943340722258a518eae61c77bd1ebe0efd9384ad1f5a62" +) +BASELINE_FINGERPRINT_SHA256 = ( + "4d2422a34fb52dcffa9d112e7203437c726d76ad8798fb14fea77ad3f6317784" +) +REQUIRED_SCOPE_TARGETS = ("benchmarks", "examples", "src/fastplms", "tools") +MYPY_COMMAND = ( + "python", + "-m", + "mypy", + "--config-file=/dev/null", + "--python-version", + "3.12", + "--strict", + "--warn-unreachable", + "--ignore-missing-imports", + "--no-site-packages", + "--no-incremental", + "--explicit-package-bases", + "--follow-imports=silent", + "--show-error-codes", + "--no-color-output", + "--no-pretty", + "benchmarks", + "examples", + "src/fastplms", + "tools", +) + +_ERROR_PATTERN = re.compile( + r"^(?P.+?):(?P[0-9]+)(?::(?P[0-9]+))?: " + r"error: (?P.+?) \[(?P[^][]+)\]$" +) +_FOUND_PATTERN = re.compile( + r"^Found (?P[0-9]+) errors? in (?P[0-9]+) files? " + r"\(checked (?P[0-9]+) source files?\)$" +) +_SUCCESS_PATTERN = re.compile( + r"^Success: no issues found in (?P[0-9]+) source files?$" +) +_COMMIT_PATTERN = re.compile(r"^[0-9a-f]{40}$") +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") + + +class TypingGateError(ValueError): + """Raised when typing evidence is incomplete, malformed, or inconsistent.""" + + +@dataclass(frozen=True, order=True) +class Fingerprint: + """One line-independent mypy finding identity.""" + + path: str + code: str + message: str + + +@dataclass(frozen=True) +class MypySnapshot: + """Parsed mypy errors plus its mandatory terminal summary.""" + + fingerprints: Counter[Fingerprint] + error_count: int + error_file_count: int + checked_source_files: int + + +@dataclass(frozen=True) +class CounterComparison: + """Multiplicity-aware delta between baseline and candidate errors.""" + + retained: Counter[Fingerprint] + new: Counter[Fingerprint] + resolved: Counter[Fingerprint] + + +def load_scope_manifest(path: Path) -> tuple[str, ...]: + """Read and validate the exact broad source roots.""" + + try: + entries = tuple( + line.strip() + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() + ) + except OSError as error: + raise TypingGateError(f"Unable to read typing scope manifest: {path}") from error + if entries != REQUIRED_SCOPE_TARGETS: + raise TypingGateError( + "Broad typing scope differs from the required roots: " + + ", ".join(REQUIRED_SCOPE_TARGETS) + ) + if entries != tuple(sorted(set(entries))): + raise TypingGateError("Broad typing scope must be sorted and duplicate-free.") + for entry in entries: + value = PurePosixPath(entry) + if value.is_absolute() or value.as_posix() != entry or ".." in value.parts: + raise TypingGateError(f"Typing scope contains an unsafe path: {entry!r}") + return entries + + +def _path_is_scoped(path: PurePosixPath, scope: Sequence[str]) -> bool: + for target in scope: + target_parts = PurePosixPath(target).parts + if path.parts[: len(target_parts)] == target_parts: + return True + return False + + +def normalize_error_path(value: str, scope: Sequence[str]) -> str: + """Normalize one mypy path to a safe repository-relative POSIX path.""" + + normalized = value.replace("\\", "/") + path = PurePosixPath(normalized) + if ( + not normalized + or path.is_absolute() + or ".." in path.parts + or (path.parts and path.parts[0].endswith(":")) + or not _path_is_scoped(path, scope) + or path.suffix != ".py" + ): + raise TypingGateError(f"Mypy emitted an out-of-scope path: {value!r}") + return path.as_posix() + + +def parse_mypy_output(text: str, scope: Sequence[str]) -> MypySnapshot: + """Parse standard mypy text output while deliberately ignoring line numbers.""" + + fingerprints: Counter[Fingerprint] = Counter() + lines = [line for line in text.splitlines() if line.strip()] + summary_matches: list[tuple[int, re.Match[str], bool]] = [] + for index, line in enumerate(lines): + found_match = _FOUND_PATTERN.fullmatch(line) + if found_match is not None: + summary_matches.append((index, found_match, False)) + success_match = _SUCCESS_PATTERN.fullmatch(line) + if success_match is not None: + summary_matches.append((index, success_match, True)) + if len(summary_matches) != 1: + raise TypingGateError("Mypy output must contain exactly one terminal summary.") + summary_index, summary_match, success = summary_matches[0] + if summary_index != len(lines) - 1: + raise TypingGateError("Mypy terminal summary must be the final nonblank line.") + if success: + summary = (0, 0, int(summary_match.group("checked"))) + else: + summary = ( + int(summary_match.group("errors")), + int(summary_match.group("files")), + int(summary_match.group("checked")), + ) + for line in lines[:summary_index]: + error_match = _ERROR_PATTERN.fullmatch(line) + if error_match is not None: + fingerprint = Fingerprint( + path=normalize_error_path(error_match.group("path"), scope), + code=error_match.group("code"), + message=error_match.group("message"), + ) + fingerprints[fingerprint] += 1 + continue + if ": error:" in line: + raise TypingGateError(f"Unrecognized mypy error line: {line!r}") + error_count, error_file_count, checked_source_files = summary + parsed_count = sum(fingerprints.values()) + if parsed_count != error_count: + raise TypingGateError( + f"Mypy summary reports {error_count} errors but {parsed_count} were parsed." + ) + parsed_files = len({fingerprint.path for fingerprint in fingerprints}) + if parsed_files != error_file_count: + raise TypingGateError( + f"Mypy summary reports {error_file_count} error files but parsed {parsed_files}." + ) + return MypySnapshot( + fingerprints=fingerprints, + error_count=error_count, + error_file_count=error_file_count, + checked_source_files=checked_source_files, + ) + + +def discover_source_files(repo_root: Path, scope: Sequence[str]) -> tuple[str, ...]: + """Return every Python source under the exact diagnostic roots.""" + + root = repo_root.resolve() + discovered: set[str] = set() + for target in scope: + target_path = root.joinpath(*PurePosixPath(target).parts) + if not target_path.is_dir() or target_path.is_symlink(): + raise TypingGateError(f"Typing scope root is missing or linked: {target_path}") + for path in target_path.rglob("*"): + if "__pycache__" in path.parts: + continue + if path.is_symlink(): + raise TypingGateError(f"Typing scope contains a symlink: {path}") + if path.is_file() and path.suffix == ".py": + discovered.add(path.relative_to(root).as_posix()) + return tuple(sorted(discovered)) + + +def _source_inventory_sha256(source_files: Sequence[str]) -> str: + inventory_payload = json.dumps( + list(source_files), + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(inventory_payload).hexdigest() + + +def source_digests(repo_root: Path, source_files: Sequence[str]) -> tuple[str, str]: + """Hash the ordered path inventory and the exact source bytes.""" + + root = repo_root.resolve() + tree_records: list[dict[str, str]] = [] + for relative_name in source_files: + if ( + normalize_error_path(relative_name, REQUIRED_SCOPE_TARGETS) + != relative_name + ): + raise TypingGateError("Source digest inventory contains an unsafe path.") + path = root.joinpath(*PurePosixPath(relative_name).parts) + try: + payload = path.read_bytes() + except OSError as error: + raise TypingGateError(f"Unable to hash scoped source: {path}") from error + tree_records.append( + { + "path": relative_name, + "sha256": hashlib.sha256(payload).hexdigest(), + } + ) + tree_payload = json.dumps( + tree_records, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return ( + _source_inventory_sha256(source_files), + hashlib.sha256(tree_payload).hexdigest(), + ) + + +def verified_git_head(repo_root: Path, scope: Sequence[str]) -> str: + """Return the clean exact Git HEAD backing baseline source bytes.""" + + root = repo_root.resolve() + command = ["git", "-c", f"safe.directory={root.as_posix()}"] + try: + revision = subprocess.run( + [*command, "rev-parse", "HEAD"], + cwd=root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + status = subprocess.run( + [*command, "status", "--porcelain=v1", "--untracked-files=all", "--", *scope], + cwd=root, + check=True, + capture_output=True, + text=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as error: + raise TypingGateError("Unable to verify baseline Git source identity.") from error + if not _COMMIT_PATTERN.fullmatch(revision): + raise TypingGateError("Baseline Git HEAD is not a full lowercase commit.") + if status.strip(): + raise TypingGateError("Baseline Git source scope is not clean.") + return revision + + +def verify_git_source_identity( + repo_root: Path, + scope: Sequence[str], + source_files: Sequence[str], +) -> None: + """Match every discovered Python file byte-for-byte to the Git HEAD tree.""" + + root = repo_root.resolve() + command = ["git", "-c", f"safe.directory={root.as_posix()}"] + try: + tree_output = subprocess.run( + [*command, "ls-tree", "-r", "-z", "HEAD", "--", *scope], + cwd=root, + check=True, + capture_output=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as error: + raise TypingGateError("Unable to read baseline Git source tree.") from error + tree_blobs: dict[str, str] = {} + try: + records = (record for record in tree_output.split(b"\0") if record) + for record in records: + metadata, raw_name = record.split(b"\t", maxsplit=1) + mode, object_type, object_id = metadata.decode("ascii").split() + relative_name = raw_name.decode("utf-8") + if PurePosixPath(relative_name).suffix != ".py": + continue + normalized_name = normalize_error_path(relative_name, scope) + if ( + normalized_name != relative_name + or object_type != "blob" + or mode == "120000" + or normalized_name in tree_blobs + ): + raise TypingGateError("Baseline Git tree contains an invalid Python source.") + tree_blobs[normalized_name] = object_id + except (UnicodeDecodeError, ValueError) as error: + raise TypingGateError("Baseline Git source tree output is malformed.") from error + if tuple(sorted(tree_blobs)) != tuple(source_files): + raise TypingGateError( + "Discovered baseline Python inventory differs from its Git HEAD tree." + ) + try: + working_hashes = subprocess.run( + [*command, "hash-object", "--", *source_files], + cwd=root, + check=True, + capture_output=True, + text=True, + ).stdout.splitlines() + except (OSError, subprocess.CalledProcessError) as error: + raise TypingGateError("Unable to hash baseline working-tree sources.") from error + if len(working_hashes) != len(source_files) or any( + tree_blobs[relative_name] != object_id + for relative_name, object_id in zip(source_files, working_hashes, strict=True) + ): + raise TypingGateError("Baseline source bytes differ from the Git HEAD tree.") + + +def compare_counters( + baseline: Counter[Fingerprint], + candidate: Counter[Fingerprint], +) -> CounterComparison: + """Compare errors as multisets, retaining multiplicity.""" + + return CounterComparison( + retained=baseline & candidate, + new=candidate - baseline, + resolved=baseline - candidate, + ) + + +def _counter_records(counter: Counter[Fingerprint]) -> list[dict[str, object]]: + return [ + { + "path": fingerprint.path, + "code": fingerprint.code, + "message": fingerprint.message, + "count": count, + } + for fingerprint, count in sorted(counter.items()) + ] + + +def _fingerprint_sha256(counter: Counter[Fingerprint]) -> str: + payload = json.dumps( + _counter_records(counter), + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _write_json(path: Path, value: Mapping[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(value, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + newline="\n", + ) + temporary.replace(path) + + +def _write_bytes(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_bytes(payload) + temporary.replace(path) + + +def _run_mypy(repo_root: Path) -> tuple[int, bytes]: + """Run the immutable command through the current pinned Python interpreter.""" + + command = (sys.executable, *MYPY_COMMAND[1:]) + environment = os.environ.copy() + environment.pop("MYPYPATH", None) + environment.pop("MYPY_CONFIG_FILE", None) + environment["NO_COLOR"] = "1" + environment["PYTHONHASHSEED"] = "0" + try: + completed = subprocess.run( + command, + cwd=repo_root.resolve(), + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=environment, + ) + except OSError as error: + raise TypingGateError("Unable to execute the pinned mypy command.") from error + return completed.returncode, completed.stdout + + +def _print_report_tail(raw_report: bytes, *, line_count: int = 25) -> None: + text = raw_report.decode("utf-8", errors="replace") + print("\n".join(text.splitlines()[-line_count:])) + + +def _read_json_object(path: Path) -> dict[str, object]: + try: + value: object = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise TypingGateError(f"Unable to read typing baseline: {path}") from error + if not isinstance(value, dict) or not all(isinstance(key, str) for key in value): + raise TypingGateError("Typing baseline must be a JSON object.") + return value + + +def _require_string(value: object, context: str) -> str: + if not isinstance(value, str) or not value: + raise TypingGateError(f"{context} must be a non-empty string.") + return value + + +def _require_nonnegative_int(value: object, context: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise TypingGateError(f"{context} must be a non-negative integer.") + return value + + +def _require_string_tuple(value: object, context: str) -> tuple[str, ...]: + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise TypingGateError(f"{context} must be a string array.") + return tuple(value) + + +def _records_counter( + value: object, + *, + scope: Sequence[str], + source_files: set[str], +) -> Counter[Fingerprint]: + if not isinstance(value, list): + raise TypingGateError("Typing baseline fingerprints must be an array.") + result: Counter[Fingerprint] = Counter() + for record in value: + if not isinstance(record, dict) or set(record) != {"path", "code", "message", "count"}: + raise TypingGateError("Typing baseline contains a malformed fingerprint.") + raw_path = _require_string(record["path"], "fingerprint path") + normalized_path = normalize_error_path(raw_path, scope) + if normalized_path != raw_path or normalized_path not in source_files: + raise TypingGateError( + "Typing baseline fingerprint path is not in its source inventory." + ) + fingerprint = Fingerprint( + path=normalized_path, + code=_require_string(record["code"], "fingerprint code"), + message=_require_string(record["message"], "fingerprint message"), + ) + count = _require_nonnegative_int(record["count"], "fingerprint count") + if count == 0 or fingerprint in result: + raise TypingGateError("Typing baseline fingerprints must be positive and unique.") + result[fingerprint] = count + return result + + +def _validate_pinned_baseline_identity(value: Mapping[str, object]) -> None: + expected: dict[str, object] = { + "checked_source_files": BASELINE_CHECKED_SOURCE_FILES, + "error_count": BASELINE_ERROR_COUNT, + "error_file_count": BASELINE_ERROR_FILE_COUNT, + "source_inventory_sha256": BASELINE_SOURCE_INVENTORY_SHA256, + "source_tree_sha256": BASELINE_SOURCE_TREE_SHA256, + "raw_report_sha256": BASELINE_RAW_REPORT_SHA256, + "fingerprint_sha256": BASELINE_FINGERPRINT_SHA256, + } + mismatches = [ + field for field, expected_value in expected.items() if value.get(field) != expected_value + ] + if mismatches: + raise TypingGateError( + "Typing baseline differs from its immutable c240 identity: " + + ", ".join(mismatches) + ) + + +def baseline_payload( + snapshot: MypySnapshot, + *, + scope: Sequence[str], + source_files: Sequence[str], + revision: str, + raw_report: bytes, + source_inventory_sha256: str, + source_tree_sha256: str, + mypy_exit_code: int, +) -> dict[str, object]: + """Build one deterministic checked baseline payload.""" + + if not _COMMIT_PATTERN.fullmatch(revision): + raise TypingGateError("Baseline revision must be a full lowercase Git commit.") + if tuple(scope) != REQUIRED_SCOPE_TARGETS: + raise TypingGateError("Baseline scope differs from the required roots.") + if tuple(source_files) != tuple(sorted(set(source_files))): + raise TypingGateError("Baseline source inventory must be sorted and unique.") + if snapshot.checked_source_files != len(source_files): + raise TypingGateError( + "Mypy checked-source count differs from the discovered source inventory." + ) + expected_inventory_sha256 = _source_inventory_sha256(source_files) + if source_inventory_sha256 != expected_inventory_sha256: + raise TypingGateError("Baseline source inventory digest is inconsistent.") + if not _SHA256_PATTERN.fullmatch(source_tree_sha256): + raise TypingGateError("Baseline source tree digest is invalid.") + expected_exit_code = 0 if snapshot.error_count == 0 else 1 + if mypy_exit_code != expected_exit_code: + raise TypingGateError("Baseline mypy exit status contradicts its terminal summary.") + return { + "schema_version": SCHEMA_VERSION, + "baseline_revision": revision, + "environment": { + "python": platform.python_version(), + "mypy": importlib.metadata.version("mypy"), + }, + "mypy_command": list(MYPY_COMMAND), + "mypy_exit_code": mypy_exit_code, + "scope_targets": list(scope), + "checked_source_files": snapshot.checked_source_files, + "source_files": list(source_files), + "source_inventory_sha256": source_inventory_sha256, + "source_tree_sha256": source_tree_sha256, + "raw_report_sha256": hashlib.sha256(raw_report).hexdigest(), + "error_count": snapshot.error_count, + "error_file_count": snapshot.error_file_count, + "fingerprint_sha256": _fingerprint_sha256(snapshot.fingerprints), + "fingerprints": _counter_records(snapshot.fingerprints), + } + + +def load_baseline(path: Path) -> tuple[dict[str, object], Counter[Fingerprint]]: + """Load and validate the immutable c240 typing debt ledger.""" + + value = _read_json_object(path) + required = { + "schema_version", + "baseline_revision", + "environment", + "mypy_command", + "mypy_exit_code", + "scope_targets", + "checked_source_files", + "source_files", + "source_inventory_sha256", + "source_tree_sha256", + "raw_report_sha256", + "error_count", + "error_file_count", + "fingerprint_sha256", + "fingerprints", + } + if set(value) != required or value.get("schema_version") != SCHEMA_VERSION: + raise TypingGateError("Typing baseline schema or field inventory is invalid.") + if value.get("baseline_revision") != BASELINE_REVISION: + raise TypingGateError("Typing baseline revision differs from the pinned c240 commit.") + if ( + _require_string_tuple(value.get("scope_targets"), "baseline scope") + != REQUIRED_SCOPE_TARGETS + ): + raise TypingGateError("Typing baseline scope differs from the required roots.") + if _require_string_tuple(value.get("mypy_command"), "baseline command") != MYPY_COMMAND: + raise TypingGateError("Typing baseline command differs from the required invocation.") + source_files = _require_string_tuple(value.get("source_files"), "baseline source files") + if source_files != tuple(sorted(set(source_files))): + raise TypingGateError("Typing baseline source inventory is not sorted and unique.") + for relative_name in source_files: + if normalize_error_path(relative_name, REQUIRED_SCOPE_TARGETS) != relative_name: + raise TypingGateError("Typing baseline contains an unsafe source path.") + checked = _require_nonnegative_int( + value.get("checked_source_files"), + "baseline checked-source count", + ) + if checked != len(source_files): + raise TypingGateError("Typing baseline checked-source count differs from its inventory.") + inventory_digest = _require_string( + value.get("source_inventory_sha256"), + "baseline source inventory digest", + ) + if inventory_digest != _source_inventory_sha256(source_files): + raise TypingGateError("Typing baseline source inventory digest is invalid.") + fingerprints = _records_counter( + value.get("fingerprints"), + scope=REQUIRED_SCOPE_TARGETS, + source_files=set(source_files), + ) + if sum(fingerprints.values()) != _require_nonnegative_int( + value.get("error_count"), + "baseline error count", + ): + raise TypingGateError("Typing baseline error count differs from its fingerprints.") + if len({item.path for item in fingerprints}) != _require_nonnegative_int( + value.get("error_file_count"), + "baseline error-file count", + ): + raise TypingGateError("Typing baseline error-file count differs from its fingerprints.") + mypy_exit_code = _require_nonnegative_int( + value.get("mypy_exit_code"), + "baseline mypy exit status", + ) + expected_exit_code = 0 if not fingerprints else 1 + if mypy_exit_code != expected_exit_code: + raise TypingGateError("Typing baseline mypy exit status is inconsistent.") + digest = _require_string(value.get("fingerprint_sha256"), "baseline fingerprint digest") + if digest != _fingerprint_sha256(fingerprints): + raise TypingGateError("Typing baseline fingerprint digest is invalid.") + for field in ("source_tree_sha256", "raw_report_sha256"): + digest = _require_string(value.get(field), f"baseline {field}") + if not _SHA256_PATTERN.fullmatch(digest): + raise TypingGateError(f"Typing baseline {field} is not a SHA-256 digest.") + environment = value.get("environment") + if ( + not isinstance(environment, dict) + or set(environment) != {"python", "mypy"} + or environment.get("python") != BASELINE_PYTHON_VERSION + or environment.get("mypy") != BASELINE_MYPY_VERSION + ): + raise TypingGateError("Typing baseline environment identity is invalid.") + _validate_pinned_baseline_identity(value) + return value, fingerprints + + +def compare_payload( + *, + baseline: Mapping[str, object], + baseline_fingerprints: Counter[Fingerprint], + candidate: MypySnapshot, + scope: Sequence[str], + source_files: Sequence[str], + source_inventory_sha256: str, + source_tree_sha256: str, + mypy_exit_code: int, +) -> dict[str, object]: + """Create the reader-facing candidate comparison and fail-closed reasons.""" + + comparison = compare_counters(baseline_fingerprints, candidate.fingerprints) + reasons: list[str] = [] + if tuple(scope) != REQUIRED_SCOPE_TARGETS: + reasons.append("diagnostic scope differs from the required roots") + if tuple(source_files) != tuple(sorted(set(source_files))): + reasons.append("candidate source inventory is not sorted and unique") + source_file_set = set(source_files) + for relative_name in source_files: + try: + normalized_name = normalize_error_path(relative_name, scope) + except TypingGateError: + reasons.append("candidate source inventory contains an unsafe path") + break + if normalized_name != relative_name: + reasons.append("candidate source inventory contains a non-canonical path") + break + fingerprints_outside_inventory = sorted( + { + fingerprint.path + for fingerprint in candidate.fingerprints + if fingerprint.path not in source_file_set + } + ) + if fingerprints_outside_inventory: + reasons.append("candidate fingerprints fall outside its source inventory") + if mypy_exit_code not in {0, 1}: + reasons.append(f"mypy exited with infrastructure status {mypy_exit_code}") + expected_exit_code = 0 if candidate.error_count == 0 else 1 + if mypy_exit_code in {0, 1} and mypy_exit_code != expected_exit_code: + reasons.append("mypy exit status contradicts its terminal summary") + if candidate.checked_source_files != len(source_files): + reasons.append("mypy checked-source count differs from candidate inventory") + baseline_files = set( + _require_string_tuple(baseline.get("source_files"), "baseline source files") + ) + missing_baseline_files = sorted(baseline_files.difference(source_files)) + if source_inventory_sha256 != _source_inventory_sha256(source_files): + reasons.append("candidate source inventory digest is inconsistent") + if not _SHA256_PATTERN.fullmatch(source_tree_sha256): + reasons.append("candidate source tree digest is invalid") + if missing_baseline_files: + reasons.append("candidate source scope removed baseline files") + if comparison.new: + reasons.append("candidate introduced typing fingerprints beyond baseline debt") + return { + "schema_version": SCHEMA_VERSION, + "status": "failed" if reasons else "passed", + "failure_reasons": reasons, + "baseline_revision": baseline["baseline_revision"], + "environment": { + "python": platform.python_version(), + "mypy": importlib.metadata.version("mypy"), + }, + "mypy_command": list(MYPY_COMMAND), + "scope_targets": list(scope), + "mypy_exit_code": mypy_exit_code, + "baseline_checked_source_files": baseline["checked_source_files"], + "candidate_checked_source_files": candidate.checked_source_files, + "candidate_source_files": list(source_files), + "candidate_source_inventory_sha256": source_inventory_sha256, + "candidate_source_tree_sha256": source_tree_sha256, + "missing_baseline_source_files": missing_baseline_files, + "fingerprint_paths_outside_candidate_inventory": ( + fingerprints_outside_inventory + ), + "baseline_error_count": sum(baseline_fingerprints.values()), + "candidate_error_count": candidate.error_count, + "retained_error_count": sum(comparison.retained.values()), + "new_error_count": sum(comparison.new.values()), + "resolved_error_count": sum(comparison.resolved.values()), + "new_fingerprints": _counter_records(comparison.new), + "resolved_fingerprints": _counter_records(comparison.resolved), + "candidate_fingerprint_sha256": _fingerprint_sha256(candidate.fingerprints), + } + + +def _failed_compare_payload( + *, + baseline: Mapping[str, object], + scope: Sequence[str], + source_files: Sequence[str], + source_inventory_sha256: str, + source_tree_sha256: str, + mypy_exit_code: int, + raw_report: bytes, + reasons: Sequence[str], +) -> dict[str, object]: + """Create a durable report even when mypy output cannot be parsed.""" + + return { + "schema_version": SCHEMA_VERSION, + "status": "failed", + "failure_reasons": list(reasons), + "baseline_revision": baseline["baseline_revision"], + "environment": { + "python": platform.python_version(), + "mypy": importlib.metadata.version("mypy"), + }, + "mypy_command": list(MYPY_COMMAND), + "scope_targets": list(scope), + "mypy_exit_code": mypy_exit_code, + "candidate_source_files": list(source_files), + "candidate_source_inventory_sha256": source_inventory_sha256, + "candidate_source_tree_sha256": source_tree_sha256, + "raw_report_sha256": hashlib.sha256(raw_report).hexdigest(), + "candidate_output_parsed": False, + } + + +def _baseline_command(arguments: argparse.Namespace) -> int: + scope = load_scope_manifest(arguments.scope_manifest) + environment = { + "python": platform.python_version(), + "mypy": importlib.metadata.version("mypy"), + } + if environment != { + "python": BASELINE_PYTHON_VERSION, + "mypy": BASELINE_MYPY_VERSION, + }: + raise TypingGateError("Baseline generation requires Python 3.12.13 and mypy 1.20.2.") + revision = verified_git_head(arguments.repo_root, scope) + if revision != BASELINE_REVISION: + raise TypingGateError("Baseline Git HEAD differs from the pinned c240 commit.") + source_files = discover_source_files(arguments.repo_root, scope) + verify_git_source_identity(arguments.repo_root, scope, source_files) + source_inventory_sha256, source_tree_sha256 = source_digests( + arguments.repo_root, + source_files, + ) + mypy_exit_code, raw_report = _run_mypy(arguments.repo_root) + _write_bytes(arguments.raw_output, raw_report) + _print_report_tail(raw_report) + if verified_git_head(arguments.repo_root, scope) != revision: + raise TypingGateError("Baseline Git HEAD changed while evidence was generated.") + final_source_files = discover_source_files(arguments.repo_root, scope) + verify_git_source_identity(arguments.repo_root, scope, final_source_files) + final_digests = source_digests(arguments.repo_root, final_source_files) + if final_source_files != source_files or final_digests != ( + source_inventory_sha256, + source_tree_sha256, + ): + raise TypingGateError("Baseline source tree changed during mypy execution.") + try: + text = raw_report.decode("utf-8") + except UnicodeDecodeError as error: + raise TypingGateError("Baseline mypy output is not UTF-8.") from error + snapshot = parse_mypy_output(text, scope) + payload = baseline_payload( + snapshot, + scope=scope, + source_files=source_files, + revision=revision, + raw_report=raw_report, + source_inventory_sha256=source_inventory_sha256, + source_tree_sha256=source_tree_sha256, + mypy_exit_code=mypy_exit_code, + ) + _validate_pinned_baseline_identity(payload) + _write_json(arguments.output, payload) + return 0 + + +def _compare_command(arguments: argparse.Namespace) -> int: + baseline, baseline_fingerprints = load_baseline(arguments.baseline) + environment = { + "python": platform.python_version(), + "mypy": importlib.metadata.version("mypy"), + } + if environment != baseline["environment"]: + raise TypingGateError("Candidate typing environment differs from the baseline.") + scope = load_scope_manifest(arguments.scope_manifest) + source_files = discover_source_files(arguments.repo_root, scope) + source_inventory_sha256, source_tree_sha256 = source_digests( + arguments.repo_root, + source_files, + ) + mypy_exit_code, raw_report = _run_mypy(arguments.repo_root) + _write_bytes(arguments.raw_output, raw_report) + _print_report_tail(raw_report) + final_source_files = discover_source_files(arguments.repo_root, scope) + final_inventory_sha256, final_tree_sha256 = source_digests( + arguments.repo_root, + final_source_files, + ) + if final_source_files != source_files or ( + final_inventory_sha256, + final_tree_sha256, + ) != (source_inventory_sha256, source_tree_sha256): + payload = _failed_compare_payload( + baseline=baseline, + scope=scope, + source_files=final_source_files, + source_inventory_sha256=final_inventory_sha256, + source_tree_sha256=final_tree_sha256, + mypy_exit_code=mypy_exit_code, + raw_report=raw_report, + reasons=["candidate source tree changed during mypy execution"], + ) + _write_json(arguments.output, payload) + return 1 + try: + text = raw_report.decode("utf-8") + except UnicodeDecodeError: + text = "" + parse_error: TypingGateError | None = None + try: + snapshot = parse_mypy_output(text, scope) + except TypingGateError as error: + parse_error = error + if parse_error is not None: + reasons = [] + if mypy_exit_code not in {0, 1}: + reasons.append( + f"mypy exited with infrastructure status {mypy_exit_code}" + ) + reasons.append(f"candidate mypy output is invalid: {parse_error}") + payload = _failed_compare_payload( + baseline=baseline, + scope=scope, + source_files=source_files, + source_inventory_sha256=source_inventory_sha256, + source_tree_sha256=source_tree_sha256, + mypy_exit_code=mypy_exit_code, + raw_report=raw_report, + reasons=reasons, + ) + _write_json(arguments.output, payload) + return 1 + payload = compare_payload( + baseline=baseline, + baseline_fingerprints=baseline_fingerprints, + candidate=snapshot, + scope=scope, + source_files=source_files, + source_inventory_sha256=source_inventory_sha256, + source_tree_sha256=source_tree_sha256, + mypy_exit_code=mypy_exit_code, + ) + payload["raw_report_sha256"] = hashlib.sha256(raw_report).hexdigest() + payload["candidate_output_parsed"] = True + _write_json(arguments.output, payload) + return 0 if payload["status"] == "passed" else 1 + + +def build_parser() -> argparse.ArgumentParser: + """Build the deterministic baseline/comparison CLI.""" + + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + baseline = subparsers.add_parser("baseline") + baseline.add_argument("--raw-output", type=Path, required=True) + baseline.add_argument("--scope-manifest", type=Path, required=True) + baseline.add_argument("--repo-root", type=Path, required=True) + baseline.add_argument("--output", type=Path, required=True) + baseline.set_defaults(handler=_baseline_command) + compare = subparsers.add_parser("compare") + compare.add_argument("--baseline", type=Path, required=True) + compare.add_argument("--raw-output", type=Path, required=True) + compare.add_argument("--scope-manifest", type=Path, required=True) + compare.add_argument("--repo-root", type=Path, required=True) + compare.add_argument("--output", type=Path, required=True) + compare.set_defaults(handler=_compare_command) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run baseline generation or candidate comparison.""" + + arguments = build_parser().parse_args(argv) + try: + handler = arguments.handler + if not callable(handler): + raise TypingGateError("Typing gate command handler is invalid.") + return int(handler(arguments)) + except (OSError, TypingGateError) as error: + print(f"typing gate failed: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/update_HF.py b/update_HF.py deleted file mode 100644 index 8163504..0000000 --- a/update_HF.py +++ /dev/null @@ -1,525 +0,0 @@ -""" -Data-driven HuggingFace upload script for all FastPLMs models. - -Builds composite single-file modeling scripts by concatenating shared modules -(attention.py, embedding_mixin.py, entrypoint_setup.py) with model-specific code, -then uploads to each HF repo. - -Usage: - py -m update_HF - $env:HF_TOKEN = "..." - py -m update_HF - py -m update_HF --families esm2 dplm - py -m update_HF --skip-weights - py -m update_HF --files-only - py -m update_HF --config-only -""" - -import argparse -import os -import platform -import re -import subprocess -import tempfile -from pathlib import Path -from typing import Dict, Optional - -from huggingface_hub import HfApi, login - -_REPO_ROOT = Path(__file__).resolve().parent - -# Regex to strip the try/except import guards that reference fastplms.* -# These are only needed for local development; composites have the code inline. -_IMPORT_GUARD_PATTERN = re.compile( - r"try:\s*\n" - r"(?:\s+from fastplms\.\w+ import[^\n]*\n|\s+from fastplms\.\w+ import \(\n(?:\s+[^\n]*\n)*?\s+\)\n)+" - r"\s*except ImportError:\s*\n" - r"\s*pass[^\n]*\n", - re.MULTILINE, -) - -COMPOSITE_SHARED_MODULES = [ - "entrypoint_setup.py", - "fastplms/embedding_mixin.py", - "fastplms/attention.py", - "fastplms/test_time_training.py", -] - - -_FUTURE_IMPORT_RE = re.compile(r"^from __future__ import annotations\s*\n?", re.MULTILINE) - - -def build_composite(modeling_path: str, include_embedding_mixin: bool = True) -> str: - """Build a single self-contained modeling file for HF Hub upload. - - Concatenates shared modules + model code, stripping the try/except - import guards from the model code since shared definitions are inlined above. - Hoists `from __future__ import annotations` to the top (must be first statement). - """ - parts = [] - for shared_path in COMPOSITE_SHARED_MODULES: - if not include_embedding_mixin and "embedding_mixin" in shared_path: - continue - content = (_REPO_ROOT / shared_path).read_text(encoding="utf-8") - content = _FUTURE_IMPORT_RE.sub("", content) - parts.append(content) - - model_code = (_REPO_ROOT / modeling_path).read_text(encoding="utf-8") - model_code = _IMPORT_GUARD_PATTERN.sub("", model_code) - model_code = _FUTURE_IMPORT_RE.sub("", model_code) - parts.append(model_code) - - return "from __future__ import annotations\n\n" + "\n".join(parts) - - -def _token_from_environment() -> Optional[str]: - for key in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN"): - if key in os.environ and len(os.environ[key]) > 0: - return os.environ[key] - return None - - -def _resolve_hf_token(cli_token: Optional[str]) -> Optional[str]: - if cli_token is not None: - assert len(cli_token) > 0, "HF token cannot be empty." - return cli_token - return _token_from_environment() - - -def _login_if_token_available(token: Optional[str]) -> None: - if token is not None: - login(token=token) - - -ESMFOLD2_REPO_IDS = [ - "Synthyra/ESMFold2", - "Synthyra/ESMFold2-Fast", - "Synthyra/ESMFold2-Experimental-Fast", - "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025", - "Synthyra/ESMFold2-Experimental", - "Synthyra/ESMFold2-Experimental-Cutoff2025", -] -ESMFOLD2_REPO_IDS.extend( - f"Synthyra/ESMFold2-Experimental-Fast-base{size}-step{step}k" - for size in ("300M", "600M", "6B") - for step in ("250", "500", "750", "1000", "1500") -) -ESMFOLD2_REMOTE_CODE_IGNORE_PATTERNS = [ - "__pycache__/*", - "*.pyc", - "configuration_esmc.py", - "configuration_esmc_sae.py", - "get_weights.py", - "modeling_esmc.py", - "modeling_esmc_sae.py", -] - - -MODEL_REGISTRY = [ - { - "family": "e1", - "repo_ids": [ - "Synthyra/Profluent-E1-150M", - "Synthyra/Profluent-E1-300M", - "Synthyra/Profluent-E1-600M", - ], - "modeling_src": "fastplms/e1/modeling_e1.py", - "modeling_dest": "modeling_e1.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": { - "fastplms/e1/tokenizer.json": "tokenizer.json", - }, - "folder_uploads": [], - "readme_map": { - "Synthyra/Profluent-E1-150M": "fastplms/e1/README.md", - "Synthyra/Profluent-E1-300M": "fastplms/e1/README.md", - "Synthyra/Profluent-E1-600M": "fastplms/e1/README.md", - }, - "license_map": { - "Synthyra/Profluent-E1-150M": "fastplms/e1/LICENSE", - "Synthyra/Profluent-E1-300M": "fastplms/e1/LICENSE", - "Synthyra/Profluent-E1-600M": "fastplms/e1/LICENSE", - }, - "weight_module": "fastplms.e1.get_weights", - }, - { - "family": "esmplusplus", - "repo_ids": [ - "Synthyra/ESMplusplus_small", - "Synthyra/ESMplusplus_large", - "Synthyra/ESMplusplus_6B", - ], - "modeling_src": "fastplms/esm_plusplus/modeling_esm_plusplus.py", - "modeling_dest": "modeling_esm_plusplus.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/ESMplusplus_small": "fastplms/esm_plusplus/README_small.md", - "Synthyra/ESMplusplus_large": "fastplms/esm_plusplus/README_large.md", - "Synthyra/ESMplusplus_6B": "fastplms/esm_plusplus/README_6B.md", - }, - "license_map": { - "Synthyra/ESMplusplus_small": "fastplms/esm_plusplus/LICENSE_small", - "Synthyra/ESMplusplus_large": "fastplms/esm_plusplus/LICENSE_large", - "Synthyra/ESMplusplus_6B": "fastplms/esm_plusplus/LICENSE_6B", - }, - "weight_module": "fastplms.esm_plusplus.get_weights", - }, - { - "family": "esm3", - "repo_ids": [ - "Synthyra/ESM3_small", - ], - "modeling_src": "fastplms/esm3/modeling_esm3.py", - "modeling_dest": "modeling_esm3.py", - "composite": False, - "include_embedding_mixin": False, - "extra_files": { - "fastplms/esm3/modeling_esm3.py": "modeling_esm3.py", - }, - "folder_uploads": [], - "readme_map": { - "Synthyra/ESM3_small": "fastplms/esm3/README.md", - }, - "license_map": { - "Synthyra/ESM3_small": "fastplms/esm3/LICENSE", - }, - "weight_module": "fastplms.esm3.get_weights", - }, - { - "family": "esm2", - "repo_ids": [ - "Synthyra/ESM2-8M", - "Synthyra/ESM2-35M", - "Synthyra/ESM2-150M", - "Synthyra/ESM2-650M", - "Synthyra/ESM2-3B", - "Synthyra/FastESM2_650", - ], - "modeling_src": "fastplms/esm2/modeling_fastesm.py", - "modeling_dest": "modeling_fastesm.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/ESM2-8M": "fastplms/esm2/README.md", - "Synthyra/ESM2-35M": "fastplms/esm2/README.md", - "Synthyra/ESM2-150M": "fastplms/esm2/README.md", - "Synthyra/ESM2-650M": "fastplms/esm2/README.md", - "Synthyra/ESM2-3B": "fastplms/esm2/README.md", - "Synthyra/FastESM2_650": "fastplms/esm2/README_650.md", - }, - "license_map": { - "Synthyra/ESM2-8M": "fastplms/esm2/LICENSE", - "Synthyra/ESM2-35M": "fastplms/esm2/LICENSE", - "Synthyra/ESM2-150M": "fastplms/esm2/LICENSE", - "Synthyra/ESM2-650M": "fastplms/esm2/LICENSE", - "Synthyra/ESM2-3B": "fastplms/esm2/LICENSE", - "Synthyra/FastESM2_650": "fastplms/esm2/LICENSE", - }, - "weight_module": "fastplms.esm2.get_weights", - }, - { - "family": "dplm", - "repo_ids": [ - "Synthyra/DPLM-150M", - "Synthyra/DPLM-650M", - "Synthyra/DPLM-3B", - ], - "modeling_src": "fastplms/dplm/modeling_dplm.py", - "modeling_dest": "modeling_dplm.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/DPLM-150M": "fastplms/dplm/README.md", - "Synthyra/DPLM-650M": "fastplms/dplm/README.md", - "Synthyra/DPLM-3B": "fastplms/dplm/README.md", - }, - "license_map": {}, - "weight_module": "fastplms.dplm.get_weights", - }, - { - "family": "dplm2", - "repo_ids": [ - "Synthyra/DPLM2-150M", - "Synthyra/DPLM2-650M", - "Synthyra/DPLM2-3B", - ], - "modeling_src": "fastplms/dplm2/modeling_dplm2.py", - "modeling_dest": "modeling_dplm2.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/DPLM2-150M": "fastplms/dplm2/README.md", - "Synthyra/DPLM2-650M": "fastplms/dplm2/README.md", - "Synthyra/DPLM2-3B": "fastplms/dplm2/README.md", - }, - "license_map": {}, - "weight_module": "fastplms.dplm2.get_weights", - }, - { - "family": "ankh", - "repo_ids": [ - "Synthyra/ANKH_base", - "Synthyra/ANKH_large", - "Synthyra/ANKH2_large", - "Synthyra/ANKH3_large", - "Synthyra/ANKH3_xl", - ], - "modeling_src": "fastplms/ankh/modeling_ankh.py", - "modeling_dest": "modeling_ankh.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/ANKH_base": "fastplms/ankh/README.md", - "Synthyra/ANKH_large": "fastplms/ankh/README.md", - "Synthyra/ANKH2_large": "fastplms/ankh/README.md", - "Synthyra/ANKH3_large": "fastplms/ankh/README.md", - "Synthyra/ANKH3_xl": "fastplms/ankh/README.md", - }, - "license_map": { - "Synthyra/ANKH_base": "fastplms/ankh/ankh_license.txt", - "Synthyra/ANKH_large": "fastplms/ankh/ankh_license.txt", - "Synthyra/ANKH2_large": "fastplms/ankh/ankh_license.txt", - "Synthyra/ANKH3_large": "fastplms/ankh/ankh_license.txt", - "Synthyra/ANKH3_xl": "fastplms/ankh/ankh_license.txt", - }, - "weight_module": "fastplms.ankh.get_weights", - }, - { - "family": "esmfold", - "repo_ids": [ - "Synthyra/FastESMFold", - ], - "modeling_src": "fastplms/esmfold/modeling_fast_esmfold.py", - "modeling_dest": "modeling_fast_esmfold.py", - "composite": True, - "include_embedding_mixin": False, - "extra_files": {}, - "folder_uploads": [], - "readme_map": { - "Synthyra/FastESMFold": "fastplms/esmfold/README.md", - }, - "license_map": {}, - "weight_module": "fastplms.esmfold.get_weights", - }, - { - "family": "esmfold2", - "repo_ids": ESMFOLD2_REPO_IDS, - "modeling_src": "fastplms/esm_plusplus/modeling_esm_plusplus.py", - "modeling_dest": "modeling_esm_plusplus.py", - "composite": True, - "include_embedding_mixin": True, - "extra_files": { - "fastplms/test_time_training.py": "test_time_training.py", - }, - "folder_uploads": [ - { - "folder_path": "fastplms/esmfold2", - "ignore_patterns": ESMFOLD2_REMOTE_CODE_IGNORE_PATTERNS, - }, - ], - "readme_map": { - repo_id: "fastplms/esmfold2/README.md" - for repo_id in ESMFOLD2_REPO_IDS - }, - "license_map": { - repo_id: "fastplms/esmfold2/LICENSE" - for repo_id in ESMFOLD2_REPO_IDS - }, - "weight_module": "fastplms.esmfold2.get_weights", - }, - { - "family": "boltz", - "repo_ids": [ - "Synthyra/Boltz2", - ], - "modeling_src": None, - "modeling_dest": None, - "composite": False, - "include_embedding_mixin": False, - "extra_files": { - "fastplms/boltz/modeling_boltz2.py": "modeling_boltz2.py", - "fastplms/boltz/__init__.py": "__init__.py", - "fastplms/boltz/minimal_featurizer.py": "minimal_featurizer.py", - "fastplms/boltz/minimal_structures.py": "minimal_structures.py", - "fastplms/boltz/cif_writer.py": "cif_writer.py", - "fastplms/boltz/vb_const.py": "vb_const.py", - "fastplms/boltz/vb_layers_attention.py": "vb_layers_attention.py", - "fastplms/boltz/vb_layers_attentionv2.py": "vb_layers_attentionv2.py", - "fastplms/boltz/vb_layers_confidence_utils.py": "vb_layers_confidence_utils.py", - "fastplms/boltz/vb_layers_dropout.py": "vb_layers_dropout.py", - "fastplms/boltz/vb_layers_initialize.py": "vb_layers_initialize.py", - "fastplms/boltz/vb_layers_outer_product_mean.py": "vb_layers_outer_product_mean.py", - "fastplms/boltz/vb_layers_pair_averaging.py": "vb_layers_pair_averaging.py", - "fastplms/boltz/vb_layers_pairformer.py": "vb_layers_pairformer.py", - "fastplms/boltz/vb_layers_transition.py": "vb_layers_transition.py", - "fastplms/boltz/vb_layers_triangular_mult.py": "vb_layers_triangular_mult.py", - "fastplms/boltz/vb_loss_diffusionv2.py": "vb_loss_diffusionv2.py", - "fastplms/boltz/vb_modules_confidencev2.py": "vb_modules_confidencev2.py", - "fastplms/boltz/vb_modules_diffusion_conditioning.py": "vb_modules_diffusion_conditioning.py", - "fastplms/boltz/vb_modules_diffusionv2.py": "vb_modules_diffusionv2.py", - "fastplms/boltz/vb_modules_encodersv2.py": "vb_modules_encodersv2.py", - "fastplms/boltz/vb_modules_transformersv2.py": "vb_modules_transformersv2.py", - "fastplms/boltz/vb_modules_trunkv2.py": "vb_modules_trunkv2.py", - "fastplms/boltz/vb_modules_utils.py": "vb_modules_utils.py", - "fastplms/boltz/vb_potentials_potentials.py": "vb_potentials_potentials.py", - "fastplms/boltz/vb_potentials_schedules.py": "vb_potentials_schedules.py", - "fastplms/boltz/vb_tri_attn_attention.py": "vb_tri_attn_attention.py", - "fastplms/boltz/vb_tri_attn_primitives.py": "vb_tri_attn_primitives.py", - "fastplms/boltz/vb_tri_attn_utils.py": "vb_tri_attn_utils.py", - }, - "folder_uploads": [], - "readme_map": { - "Synthyra/Boltz2": "fastplms/boltz/README.md", - }, - "license_map": { - "Synthyra/Boltz2": "fastplms/boltz/LICENSE", - }, - "weight_module": "fastplms.boltz.get_weights", - }, -] - - -def _run_weight_scripts( - families: Optional[list], hf_token: Optional[str], skip_weights: bool -) -> None: - python_cmd = "python" if platform.system().lower() == "linux" else "py" - child_env: Optional[Dict[str, str]] = None - if hf_token is not None: - child_env = os.environ.copy() - child_env["HF_TOKEN"] = hf_token - for entry in MODEL_REGISTRY: - if families is not None and entry["family"] not in families: - continue - module = entry["weight_module"] - if module is None: - continue - command = [python_cmd, "-m", module] - if skip_weights: - command.append("--skip-weights") - print(f"Running: {' '.join(command)}") - subprocess.run(command, check=True, env=child_env) - - -def _upload_files(api: HfApi, families: Optional[list]) -> None: - for entry in MODEL_REGISTRY: - if families is not None and entry["family"] not in families: - continue - - # Build composite file if needed - composite_path = None - if entry["composite"] and entry["modeling_src"] is not None: - composite_code = build_composite( - entry["modeling_src"], - include_embedding_mixin=entry["include_embedding_mixin"], - ) - # Verify composite compiles - compile(composite_code, entry["modeling_dest"], "exec") - composite_path = os.path.join(tempfile.gettempdir(), entry["modeling_dest"]) - with open(composite_path, "w", encoding="utf-8") as f: - f.write(composite_code) - print(f"Built composite: {entry['modeling_dest']} ({len(composite_code)} chars)") - - for repo_id in entry["repo_ids"]: - print(f"\nUploading to {repo_id}") - - # Upload composite modeling file - if composite_path is not None: - api.upload_file( - path_or_fileobj=composite_path, - path_in_repo=entry["modeling_dest"], - repo_id=repo_id, - repo_type="model", - ) - - # Upload remote-code folders for multi-file AutoModel packages. - for folder_upload in entry["folder_uploads"]: - abs_folder = str(_REPO_ROOT / folder_upload["folder_path"]) - api.upload_folder( - folder_path=abs_folder, - repo_id=repo_id, - repo_type="model", - ignore_patterns=folder_upload["ignore_patterns"], - ) - - # Upload extra files (Boltz vb_* modules, E1 tokenizer, etc.) - for local_path, repo_path in entry["extra_files"].items(): - abs_local = str(_REPO_ROOT / local_path) - api.upload_file( - path_or_fileobj=abs_local, - path_in_repo=repo_path, - repo_id=repo_id, - repo_type="model", - ) - - # Upload license - license_path = None - if repo_id in entry["license_map"]: - license_path = entry["license_map"][repo_id] - if license_path is not None: - abs_license = str(_REPO_ROOT / license_path) - assert os.path.exists(abs_license), f"Missing license: {abs_license}" - api.upload_file( - path_or_fileobj=abs_license, - path_in_repo="LICENSE", - repo_id=repo_id, - repo_type="model", - ) - - # Upload readme - readme_path = None - if repo_id in entry["readme_map"]: - readme_path = entry["readme_map"][repo_id] - if readme_path is not None: - abs_readme = str(_REPO_ROOT / readme_path) - assert os.path.exists(abs_readme), f"Missing model card: {abs_readme}" - api.upload_file( - path_or_fileobj=abs_readme, - path_in_repo="README.md", - repo_id=repo_id, - repo_type="model", - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Upload FastPLMs models to HuggingFace") - parser.add_argument( - "--hf_token", - type=str, - default=None, - help="Deprecated. Prefer HF_TOKEN in the environment so tokens are not in shell history.", - ) - parser.add_argument("--families", nargs="+", default=None) - parser.add_argument( - "--skip-weights", - action="store_true", - help="Run weight scripts without downloading/pushing model weights", - ) - parser.add_argument("--files-only", action="store_true", help="Only upload files, skip weight conversion") - parser.add_argument("--config-only", action="store_true", help="Only upload config+tokenizer via --skip-weights, skip file uploads") - args = parser.parse_args() - - hf_token = _resolve_hf_token(args.hf_token) - _login_if_token_available(hf_token) - - if args.config_only: - _run_weight_scripts(args.families, hf_token, skip_weights=True) - elif not args.files_only: - _run_weight_scripts(args.families, hf_token, args.skip_weights) - - if not args.config_only: - api = HfApi() - _upload_files(api, args.families) - - print("\nDone.") diff --git a/vendor/README.md b/vendor/README.md new file mode 100644 index 0000000..99f144f --- /dev/null +++ b/vendor/README.md @@ -0,0 +1,74 @@ +# Official reference implementations + +The repositories in `vendor/upstream/` are pinned parity oracles. They define +the official configuration, tokenization, parameter, and inference behavior +against which FastPLMs is tested. Production code under `src/fastplms/` must not +import them, copy their source, or require them at runtime. + +Initialize the complete reference tree with: + +```bash +git submodule update --init --recursive +``` + +The committed Git link selects the revision. Do not follow an upstream branch +or use an unpinned source archive in a compliance run. + +## Pinned sources + +| Directory | Revision | Reference environment | License source | +|---|---|---|---| +| `ankh` | `02b4e25ce5389b9e771c9df6e546c62af1216f8e` | `reference-ankh` | `LICENSE.md` | +| `biohub-esm` | `82ee35553d39169d678f784c8d3f8712ffd7d2c4` | `reference-biohub-esm`, `reference-esmfold2` | `LICENSE.md`, `THIRD_PARTY_NOTICE.md` | +| `biohub-transformers` | `3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf` | `reference-biohub-esm`, `reference-esmfold2` | `LICENSE` | +| `boltz` | `b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc` | `reference-boltz2` | `LICENSE` | +| `dplm` | `8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d` | `reference-dplm` | `LICENSE` | +| `e1` | `bfd2620a602248499f3d2583d85a7ecddf0b6e02` | `reference-e1` | `LICENSE`, `ATTRIBUTION`, `NOTICE` | +| `fair-esm` | `2b369911bb5b4b0dda914521b9475cad1656b2ac` | `reference-esm2`, `reference-esmfold` | `LICENSE` | +| `openfold` | `4b41059694619831a7db195b7e0988fc4ff3a307` | `reference-esmfold` | `LICENSE` | +| `protein-ttt` | `fde2817cd84b936167cc76ccabf31e5c0fe49962` | `reference-protein-ttt` | `LICENSE` | + +`src/fastplms/models.toml` is the machine-readable source for these revisions, +their model-family ownership, checkpoint snapshots, and license expressions. + +## Revision updates + +Update one source at a time: + +1. Review upstream code, dependency, model, tokenizer, and license changes. +2. Check out a specific commit in the corresponding submodule. +3. Update the matching revision and legal-file digests in `models.toml`, then + refresh the verbatim copies and notices under `LICENSES/`. +4. Rebuild only the owning reference image. +5. Run exact configuration, tokenizer, state, and live inference compliance for + every affected checkpoint. +6. Record any intentional semantic difference before the Git link is changed. + +An upstream revision is not accepted because its tests pass in isolation. It is +accepted only after the FastPLMs compliance contract passes against that exact +source. + +## Container and distribution boundary + +Reference images receive only the source directories assigned to their target. +Runtime images and local Hub artifacts contain FastPLMs code, model assets, and +required notices, but never an official source checkout. Checkpoint weights are +immutable Hub snapshots identified in `models.toml`; they are not Git +submodules. + +## Reference adapter boundary + +Reference adapters may call an upstream repository's public loading and +inference APIs, then normalize the returned configuration, tokenizer, state, or +output into the shared parity protocol. They must remain independent oracles: + +- They must not import `fastplms` or any production module under + `src/fastplms/`. +- They must not patch upstream model classes or replace upstream layers. +- They must not reuse FastPLMs tokenizers, checkpoint loaders, converters, or + attention implementations. +- They must not reconstruct an upstream forward pass from copied equations. + +An adapter that cannot obtain a required value through the pinned +implementation's public API must report that limitation explicitly. It must not +manufacture an equivalent result with FastPLMs code. diff --git a/vendor/upstream/ankh b/vendor/upstream/ankh new file mode 160000 index 0000000..02b4e25 --- /dev/null +++ b/vendor/upstream/ankh @@ -0,0 +1 @@ +Subproject commit 02b4e25ce5389b9e771c9df6e546c62af1216f8e diff --git a/official/esm b/vendor/upstream/biohub-esm similarity index 100% rename from official/esm rename to vendor/upstream/biohub-esm diff --git a/vendor/upstream/biohub-transformers b/vendor/upstream/biohub-transformers new file mode 160000 index 0000000..3a8956f --- /dev/null +++ b/vendor/upstream/biohub-transformers @@ -0,0 +1 @@ +Subproject commit 3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf diff --git a/official/boltz b/vendor/upstream/boltz similarity index 100% rename from official/boltz rename to vendor/upstream/boltz diff --git a/official/dplm b/vendor/upstream/dplm similarity index 100% rename from official/dplm rename to vendor/upstream/dplm diff --git a/official/e1 b/vendor/upstream/e1 similarity index 100% rename from official/e1 rename to vendor/upstream/e1 diff --git a/vendor/upstream/fair-esm b/vendor/upstream/fair-esm new file mode 160000 index 0000000..2b36991 --- /dev/null +++ b/vendor/upstream/fair-esm @@ -0,0 +1 @@ +Subproject commit 2b369911bb5b4b0dda914521b9475cad1656b2ac diff --git a/vendor/upstream/openfold b/vendor/upstream/openfold new file mode 160000 index 0000000..4b41059 --- /dev/null +++ b/vendor/upstream/openfold @@ -0,0 +1 @@ +Subproject commit 4b41059694619831a7db195b7e0988fc4ff3a307 diff --git a/vendor/upstream/protein-ttt b/vendor/upstream/protein-ttt new file mode 160000 index 0000000..fde2817 --- /dev/null +++ b/vendor/upstream/protein-ttt @@ -0,0 +1 @@ +Subproject commit fde2817cd84b936167cc76ccabf31e5c0fe49962 diff --git a/weight_comparison.py b/weight_comparison.py deleted file mode 100644 index 80d059a..0000000 --- a/weight_comparison.py +++ /dev/null @@ -1,119 +0,0 @@ -import argparse -import os -import torch -from typing import Dict - -from safetensors.torch import load_file -from rich.console import Console -from rich.table import Table - -from transformers import AutoModelForMaskedLM, AutoConfig, AutoModel - -from fastplms.e1.modeling_e1 import E1ForMaskedLM, E1Config, E1Model - - -def load_weights(path: str, cast_fp32: bool = True) -> Dict[str, torch.Tensor]: - assert os.path.exists(path), f"File {path} not found." - if path.endswith(".safetensors"): - sd = load_file(path) - elif path.endswith(".pth") or path.endswith(".pt"): - sd = torch.load(path, map_location="cpu", weights_only=True) - if isinstance(sd, dict) and "state_dict" in sd: - sd = sd["state_dict"] - elif isinstance(sd, dict) and "model" in sd: - sd = sd["model"] - else: - try: - sd = load_file(path) - except Exception: - sd = torch.load(path, map_location="cpu", weights_only=True) - - if cast_fp32: - return {k: v.float() if isinstance(v, torch.Tensor) else v for k, v in sd.items()} - return sd - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--file1", type=str, default=None) - parser.add_argument("--files", type=str, nargs="+", default=None) - parser.add_argument("--strict", action="store_true") - parser.add_argument("--assert_exact", action="store_true") - args = parser.parse_args() - - model = E1ForMaskedLM.from_pretrained('Profluent-Bio/E1-150m', dtype=torch.float32).eval() - torch.save(model.state_dict(), 'official.pth') - - config = AutoConfig.from_pretrained('Synthyra/Profluent-E1-150M', trust_remote_code=True) - model1 = AutoModel.from_pretrained('Synthyra/Profluent-E1-150M', dtype=torch.float32, trust_remote_code=True).eval() - torch.save(model1.state_dict(), 'load_from_pretrained_1.pth') - model2 = AutoModelForMaskedLM.from_pretrained('Synthyra/Profluent-E1-150M', dtype=torch.float32, trust_remote_code=True).eval() - torch.save(model2.state_dict(), 'load_from_pretrained_2.pth') - - if args.file1 is None: - args.file1 = 'official.pth' - if args.files is None: - args.files = ['load_from_pretrained_1.pth', 'load_from_pretrained_2.pth', 'old.safetensors'] - - paths = [args.file1] + args.files - sds = [load_weights(p, cast_fp32=not args.strict) for p in paths] - all_keys = sorted(set().union(*(sd.keys() for sd in sds))) - strict_mismatches = [] - - console = Console() - table = Table(title=f"Weights Comparison (Reference: {os.path.basename(paths[0])})") - table.add_column("Tensor Name", style="cyan", no_wrap=True) - - for p in paths[1:]: - table.add_column(f"{os.path.basename(p)} == Ref", justify="center") - - sd1 = sds[0] - for k in all_keys: - row = [k] - - has_ref = k in sd1 - ref_w = sd1[k] if has_ref else None - - for sd in sds[1:]: - has_other = k in sd - other_w = sd[k] if has_other else None - - if not has_ref or not has_other: - if not has_ref and not has_other: - row.append("[dim]✔[/dim]") - else: - row.append("[red]✘[/red]") - else: - # Both present, compare shapes and MSE - assert isinstance(ref_w, torch.Tensor), f"Weight {k} in reference is not a tensor." - assert isinstance(other_w, torch.Tensor), f"Weight {k} in comparison file is not a tensor." - - if ref_w.shape != other_w.shape: - row.append("[red]✘ (Shape)[/red]") - else: - if args.strict: - if torch.equal(ref_w, other_w): - row.append("[green]✔[/green]") - else: - mse = torch.mean((ref_w.float() - other_w.float())**2).item() - row.append(f"[red]✘ (Strict, MSE: {mse:.2e})[/red]") - strict_mismatches.append(k) - else: - mse = torch.mean((ref_w - other_w)**2).item() - if mse == 0: - row.append("[green]✔[/green]") - else: - row.append(f"[red]✘ (MSE: {mse:.2e})[/red]") - - table.add_row(*row) - - console.print(table) - if args.strict and args.assert_exact: - assert len(strict_mismatches) == 0, ( - f"Found {len(strict_mismatches)} strict mismatches. " - f"First mismatches: {strict_mismatches[:10]}" - ) - - -if __name__ == "__main__": - main() \ No newline at end of file