biohub/ESMFold2-Fast

1"""Load this checkpoint with the ``esm`` package instead of the transformers port.2 3Shipped inside each ``biohub/ESMFold2*`` repo and named by its ``config.json``4``auto_map``, so ``AutoModel.from_pretrained(repo, trust_remote_code=True)``5returns an ``esm.models.esmfold2`` model. Without the flag transformers keeps6using its own port.7 8Runs from the Hub module cache in the user's environment, so it imports nothing9but ``transformers``, ``packaging`` and ``esm``.10"""11 12from typing import Any13 14from packaging.version import Version15from transformers.configuration_utils import PretrainedConfig16from transformers.modeling_utils import PreTrainedModel17 18#: Inclusive floor. Older esm cannot read the bundled single-checkpoint layout.19MIN_ESM_VERSION = "3.4.1"20 21#: Everything esm's from_pretrained consumes, directly or via resolve_model_dir.22_ESM_KWARGS = frozenset(23    {24        "load_esmc",25        "esmc_precision",26        "device",27        "dtype",28        "revision",29        "cache_dir",30        "token",31        "local_files_only",32        "force_download",33    }34)35 36 37def esmfold2_class() -> Any:38    """Return esm's ``EsmFold2Model``, or raise ``ImportError`` naming the pip command."""39    try:40        import esm41        from esm.models.esmfold2 import EsmFold2Model42    except ImportError as exc:43        raise ImportError(44            f"trust_remote_code=True runs the esm package, which failed to "45            f"import ({exc}). Install it with:\n\n"46            f"    pip install 'esm>={MIN_ESM_VERSION}'\n\n"47            "Or drop the flag to use the ESMFold2 port in transformers."48        ) from exc49    if Version(esm.__version__) < Version(MIN_ESM_VERSION):50        raise ImportError(51            f"esm {esm.__version__} is installed, but this checkpoint needs "52            f"{MIN_ESM_VERSION} or newer. Upgrade with:\n\n"53            f"    pip install --upgrade 'esm>={MIN_ESM_VERSION}'"54        )55    return EsmFold2Model56 57 58class EsmFold2RemoteConfig(PretrainedConfig):59    """Holds config.json verbatim. esm re-reads the file and builds its own config."""60 61    model_type = "esmfold2"62 63 64class EsmFold2RemoteModel(PreTrainedModel):65    """Loader shim; ``from_pretrained`` returns an esm model, not one of these.66 67    A ``PreTrainedModel`` subclass only because the Auto classes call68    ``register_for_auto_class`` on what they load and check its ``config_class``.69    """70 71    config_class = EsmFold2RemoteConfig72 73    @classmethod74    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):  # type: ignore[override]75        # esm forwards leftover kwargs to resolve_model_dir, whose signature is76        # fixed, so transformers' trust_remote_code / _from_auto would TypeError.77        forwarded = {k: v for k, v in kwargs.items() if k in _ESM_KWARGS}78        return esmfold2_class().from_pretrained(79            pretrained_model_name_or_path, *args, **forwarded80        )81