Update FastPLMs files

README.md

@@ -14,15 +14,16 @@ This checkpoint contains 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`.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`,
+`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`.
 
 ## 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 |
+| 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: 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` |
@@ -72,6 +73,43 @@ materialize attention tensors. The configured backend does not change.
 This family declares the `compliance` tier. Release evidence identifies the
 checkpoint, backend, dtype, hardware, inputs, and reference revision.
 
+## Downstream prediction
+
+The sequence and token prediction AutoClasses use the checkpoint backbone and
+create a new, untrained `classifier`. Sequence labels have shape `(b,)`.
+Residue labels have shape `(b, l)` and use `-100` outside biological positions.
+The folding trunk is skipped. The classifier uses the checkpoint's learned pLM
+state mixture and projection, followed by one trainable transformer probe.
+
+```python
+import torch
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+model_id = "Synthyra/ESMFold2-Fast"
+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.prepare_classifier_inputs(sequences)
+biological = batch["attention_mask"].bool()
+
+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 training dependencies. Then attach LoRA to the loaded checkpoint:
@@ -81,20 +119,22 @@ python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
 ```
 
 ```python
-from peft import LoraConfig, get_peft_model
+from peft import LoraConfig, TaskType, get_peft_model
 
 peft_model = get_peft_model(
-    model,
+    sequence_model,
     LoraConfig(
+        task_type=TaskType.SEQ_CLS,
         r=8,
         lora_alpha=16,
         target_modules="all-linear",
+        modules_to_save=["classifier"],
     ),
 )
 ```
 
-This checkpoint has no advertised classifier. Supply the task objective and
-preserve any new head through `modules_to_save`.
+This checkpoint advertises a classification head. Save the separately trained
+`classifier` with the adapter.
 All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and
 can use PEFT. The ESM2-specific shipped CLI is an example, not a
 support boundary. Record the target modules, base revision, data identity, and
@@ -149,13 +189,15 @@ complex_result = model.fold(
 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.
+The typed interface also supports RNA, modifications, and covalent bonds.
+Protein MSA inputs are not supported by this Fast checkpoint; every protein
+chain must use `msa=None`. The public schema recognizes `PocketConditioning` and
+`DistogramConditioning`, but the pinned official forward consumes neither. Its
+feature builder hard-codes a zero pocket feature and constructs distogram tensors
+that the released model ignores. FastPLMs therefore rejects non-null pocket and
+distogram conditioning instead of silently ignoring scientific inputs. 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.
 
@@ -208,19 +250,15 @@ and
 [release evidence manifest](https://github.com/Synthyra/FastPLMs/blob/main/docs/generated/capability_evidence.md).
 
 
-## Hash-pinned CCD runtime asset
+## Verified 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.
+`biohub/ESMFold2`. The manifest pins its repository, revision, size, content
+identity, and MIT terms. This is a trusted-deserialization boundary. FastPLMs
+accepts only the pinned snapshot link inside the repository blob directory and
+rejects user-supplied asset and `cache_dir` symlinks. The loader verifies a
+private temporary snapshot before deserialization. Offline execution requires
+the exact cached object and never downloads a replacement.
 
 ## Optional folding TTT
 
@@ -247,8 +285,8 @@ adapter modules are excluded from checkpoint state. It is not a generic
 ## 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`
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head`
 - Attention implementations: `eager`, `sdpa`, `flex_attention`
 - Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental)
 - BF16 execution: `fp32_parameters_autocast`
@@ -262,8 +300,8 @@ adapter modules are excluded from checkpoint state. It is not a generic
 ## Release record
 
 - FastPLMs weights: `Synthyra/ESMFold2-Fast`
-- Runtime revision: recorded in the built artifact and published commit
-- Source-tree and runtime-bundle SHA-256: recorded in the source record
+- Runtime revision: recorded separately in the built artifact and published commit
+- Runtime source identities: recorded in `source-record.json`
 - Official checkpoint: `biohub/ESMFold2-Fast`
 - Artifact source: `fast`
 - State transform: `identity`

fastplms/models.toml

@@ -199,7 +199,7 @@ 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" }
+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", AutoModelForSequenceClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForTokenClassification" }
 
 [families.esm3]
 architecture = "ESM3"
@@ -223,7 +223,7 @@ 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" }
+auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model", AutoModelForSequenceClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForTokenClassification" }
 
 [families.e1]
 architecture = "E1"
@@ -370,8 +370,8 @@ conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus i
 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" }
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
+auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding", AutoModelForSequenceClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForTokenClassification" }
 
 [families.esmfold2]
 architecture = "ESMFold2"
@@ -396,8 +396,8 @@ conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its s
 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" }
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForTokenClassification" }
 
 [[models]]
 id = "esm2_8m"
@@ -1215,7 +1215,7 @@ 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" }
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
 
 [[models]]
 id = "esmfold2_experimental_fast_cutoff2025"
@@ -1236,4 +1236,4 @@ 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" }
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }

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)
+        )

fastplms/models/classification_probe.py

@@ -0,0 +1,542 @@
+"""Shared transformer probes for residue and sequence prediction tasks."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        _get_flex_attention_fn,
+        flex_attention,
+        get_attention_mask,
+        resolve_attention_backend,
+    )
+    from fastplms.embeddings.pooling import Pooler
+    from fastplms.models._esm_rotary import RotaryEmbedding
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "Pooler",
+        "RotaryEmbedding",
+        "_get_flex_attention_fn",
+        "flex_attention",
+        "get_attention_mask",
+        "resolve_attention_backend",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
+    ):
+        raise
+    # Flat Hub composites define every shared symbol above this source.
+
+
+_SUPPORTED_BACKENDS = frozenset(
+    {
+        AttentionBackend.EAGER,
+        AttentionBackend.SDPA,
+        AttentionBackend.FLEX_ATTENTION,
+    }
+)
+_SUPPORTED_PROBLEM_TYPES = frozenset(
+    {
+        "regression",
+        "single_label_classification",
+        "multi_label_classification",
+    }
+)
+_UNSUPPORTED_POOLING = frozenset({"cls", "parti"})
+
+
+def _config_value(config: Any, name: str, default: Any) -> Any:
+    value = getattr(config, name, None)
+    return default if value is None else value
+
+
+def _attention_backend(config: Any) -> AttentionBackend:
+    requested = getattr(config, "_attn_implementation", None)
+    if requested is None:
+        requested = getattr(config, "attn_backend", "sdpa")
+    backend = resolve_attention_backend(requested)
+    if backend not in _SUPPORTED_BACKENDS:
+        expected = ", ".join(sorted(item.value for item in _SUPPORTED_BACKENDS))
+        raise ValueError(
+            f"Classification probes support only {expected}; received {backend.value!r}."
+        )
+    return backend
+
+
+def resolve_problem_type(
+    config: Any,
+    labels: torch.Tensor,
+    *,
+    num_labels: int,
+) -> str:
+    """Resolve and persist the standard Transformers classification problem type."""
+
+    problem_type = getattr(config, "problem_type", None)
+    if problem_type is None:
+        if num_labels == 1:
+            problem_type = "regression"
+        elif labels.dtype in {torch.long, torch.int}:
+            problem_type = "single_label_classification"
+        else:
+            problem_type = "multi_label_classification"
+        config.problem_type = problem_type
+    if problem_type not in _SUPPORTED_PROBLEM_TYPES:
+        raise ValueError(
+            f"Unsupported problem_type {problem_type!r}; expected one of "
+            f"{sorted(_SUPPORTED_PROBLEM_TYPES)}."
+        )
+    return problem_type
+
+
+def sequence_classification_loss(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    *,
+    problem_type: str,
+    num_labels: int,
+) -> torch.Tensor:
+    """Compute a Hugging Face-compatible sequence task loss."""
+
+    labels = labels.to(logits.device)
+    if problem_type == "regression":
+        if num_labels == 1:
+            return F.mse_loss(logits.squeeze(-1), labels.squeeze(-1).to(logits.dtype))
+        return F.mse_loss(logits, labels.to(logits.dtype))
+    if problem_type == "single_label_classification":
+        return F.cross_entropy(logits.reshape(-1, num_labels), labels.reshape(-1).long())
+    if problem_type == "multi_label_classification":
+        return F.binary_cross_entropy_with_logits(logits, labels.to(logits.dtype))
+    raise ValueError(f"Unsupported problem_type {problem_type!r}.")
+
+
+def _masked_elementwise_loss(
+    losses: torch.Tensor,
+    labels: torch.Tensor,
+) -> torch.Tensor:
+    valid = labels.ne(-100)
+    if not bool(valid.any()):
+        return losses.sum() * 0
+    return losses.masked_select(valid).mean()
+
+
+def token_classification_loss(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    *,
+    problem_type: str,
+    num_labels: int,
+) -> torch.Tensor:
+    """Compute a token task loss, excluding every label element equal to ``-100``."""
+
+    labels = labels.to(logits.device)
+    if problem_type == "regression":
+        targets = labels.to(logits.dtype)
+        if num_labels == 1 and targets.ndim == logits.ndim - 1:
+            targets = targets.unsqueeze(-1)
+        if targets.shape != logits.shape:
+            raise ValueError(
+                "Token regression labels must match logits, except that the final "
+                "singleton dimension may be omitted when num_labels=1."
+            )
+        return _masked_elementwise_loss(F.mse_loss(logits, targets, reduction="none"), targets)
+    if problem_type == "single_label_classification":
+        if not bool(labels.ne(-100).any()):
+            return logits.sum() * 0
+        return F.cross_entropy(
+            logits.reshape(-1, num_labels),
+            labels.reshape(-1).long(),
+            ignore_index=-100,
+        )
+    if problem_type == "multi_label_classification":
+        if labels.shape != logits.shape:
+            raise ValueError("Multilabel token labels must have the same shape as logits.")
+        losses = F.binary_cross_entropy_with_logits(
+            logits,
+            labels.to(logits.dtype),
+            reduction="none",
+        )
+        return _masked_elementwise_loss(losses, labels)
+    raise ValueError(f"Unsupported problem_type {problem_type!r}.")
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation used by the Protify-aligned feed-forward layer."""
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        gate, values = inputs.chunk(2, dim=-1)
+        return F.silu(gate) * values
+
+
+class ProbeSelfAttention(nn.Module):
+    """Four-head RoPE self-attention with explicit, fail-closed dispatch."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        dropout: float,
+        backend: AttentionBackend,
+        use_bias: bool,
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads:
+            raise ValueError("classifier_probe_hidden_size must be divisible by its head count.")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_size = hidden_size // num_heads
+        self.dropout = dropout
+        self.backend = backend
+        self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=use_bias)
+        self.output = nn.Linear(hidden_size, hidden_size, bias=use_bias)
+        self.rotary = RotaryEmbedding(self.head_size)
+
+    def _reshape(self, tensor: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, _ = tensor.shape
+        return tensor.view(
+            batch_size,
+            sequence_length,
+            self.num_heads,
+            self.head_size,
+        ).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        attention_mask: torch.Tensor | None,
+        output_attentions: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query, key, value = self.qkv(hidden_states).chunk(3, dim=-1)
+        query = self._reshape(query)
+        key = self._reshape(key)
+        value = self._reshape(value)
+        query, key = self.rotary(query, key)
+        if output_attentions and self.backend != AttentionBackend.EAGER:
+            raise ValueError(
+                f"output_attentions=True is unavailable for {self.backend.value!r}; "
+                "select 'eager' explicitly."
+            )
+        _, attention_mask_4d, flex_block_mask = get_attention_mask(
+            self.backend,
+            batch_size,
+            sequence_length,
+            hidden_states.device,
+            attention_mask,
+            hidden_states.dtype,
+        )
+        dropout = self.dropout if self.training else 0.0
+        attention_weights = None
+        if self.backend == AttentionBackend.EAGER:
+            scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_size)
+            if attention_mask_4d is not None:
+                scores = scores.masked_fill(~attention_mask_4d, float("-inf"))
+            attention_weights = scores.softmax(dim=-1)
+            context = F.dropout(attention_weights, p=dropout, training=self.training) @ value
+        elif self.backend == AttentionBackend.SDPA:
+            context = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask_4d,
+                dropout_p=dropout,
+            )
+        elif self.backend == AttentionBackend.FLEX_ATTENTION:
+            if flex_attention is None:
+                raise RuntimeError("'flex_attention' was requested but is unavailable.")
+            flex_fn = _get_flex_attention_fn(
+                device=query.device,
+                dtype=query.dtype,
+                shape=tuple(query.shape),
+                mask_semantics="padding",
+            )
+            if flex_fn is None:
+                raise RuntimeError("'flex_attention' was requested but is unavailable.")
+            context = flex_fn(
+                query,
+                key,
+                value,
+                block_mask=flex_block_mask,
+                scale=1.0 / math.sqrt(self.head_size),
+                kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32},
+            )
+        else:
+            raise AssertionError(f"Unhandled attention backend {self.backend.value!r}.")
+        context = context.transpose(1, 2).contiguous().view(
+            batch_size,
+            sequence_length,
+            self.hidden_size,
+        )
+        return self.output(context), attention_weights
+
+
+class ProteinTransformerProbe(nn.Module):
+    """Project residue embeddings and refine them with exactly one pre-LN block."""
+
+    def __init__(self, config: Any, input_size: int) -> None:
+        super().__init__()
+        hidden_size = int(_config_value(config, "classifier_probe_hidden_size", 512))
+        num_heads = int(_config_value(config, "classifier_probe_num_heads", 4))
+        dropout = float(_config_value(config, "classifier_probe_dropout", 0.1))
+        use_bias = bool(
+            _config_value(
+                config,
+                "classifier_use_bias",
+                _config_value(config, "use_bias", False),
+            )
+        )
+        if hidden_size != 512 or num_heads != 4 or hidden_size // num_heads != 128:
+            raise ValueError(
+                "The folding classification probe requires a 512-wide projection with "
+                "four 128-wide attention heads."
+            )
+        self.hidden_size = hidden_size
+        self.input_norm = nn.LayerNorm(input_size)
+        self.input_projection = nn.Linear(input_size, hidden_size, bias=use_bias)
+        self.attention_norm = nn.LayerNorm(hidden_size)
+        self.attention = ProbeSelfAttention(
+            hidden_size,
+            num_heads,
+            dropout,
+            _attention_backend(config),
+            use_bias,
+        )
+        intermediate_size = int(math.ceil((8 / 3) * hidden_size / 256) * 256)
+        self.feed_forward_norm = nn.LayerNorm(hidden_size)
+        self.feed_forward = nn.Sequential(
+            nn.Linear(hidden_size, 2 * intermediate_size, bias=use_bias),
+            SwiGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(intermediate_size, hidden_size, bias=use_bias),
+        )
+        self.residual_dropout = nn.Dropout(dropout)
+
+    @property
+    def attn_backend(self) -> str:
+        return self.attention.backend.value
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        *,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> BaseModelOutput | tuple[torch.Tensor, ...]:
+        if embeddings.ndim != 3:
+            raise ValueError("embeddings must have shape (batch, residue, channel).")
+        embeddings = embeddings.to(dtype=self.input_projection.weight.dtype)
+        hidden_states = self.input_projection(self.input_norm(embeddings))
+        attention_output, attention_weights = self.attention(
+            self.attention_norm(hidden_states),
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = hidden_states + self.residual_dropout(attention_output)
+        hidden_states = hidden_states + self.residual_dropout(
+            self.feed_forward(self.feed_forward_norm(hidden_states))
+        )
+        output = BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=(hidden_states,) if output_hidden_states else None,
+            attentions=(attention_weights,) if output_attentions else None,
+        )
+        return output if return_dict else output.to_tuple()
+
+
+class _ClassificationProbe(nn.Module):
+    def __init__(self, config: Any, input_size: int, *, sequence_task: bool) -> None:
+        super().__init__()
+        self.config = config
+        self.num_labels = int(_config_value(config, "num_labels", 2))
+        self.transformer = ProteinTransformerProbe(config, input_size)
+        self.sequence_task = sequence_task
+        pooling_types = _config_value(config, "classifier_pooling_types", ["mean"])
+        self.pooler = Pooler(pooling_types) if sequence_task else None
+        if self.pooler is not None:
+            unsupported = sorted(set(self.pooler.names) & _UNSUPPORTED_POOLING)
+            if unsupported:
+                raise ValueError(
+                    "Classification probes consume residue-only representations and do not "
+                    f"support pooling operation(s) {unsupported}."
+                )
+        hidden_size = self.transformer.hidden_size
+        classifier_input = hidden_size * (len(self.pooler.names) if self.pooler else 1)
+        classifier_hidden = int(_config_value(config, "classifier_hidden_size", 4096))
+        classifier_dropout = float(_config_value(config, "classifier_dropout", 0.2))
+        use_bias = bool(
+            _config_value(
+                config,
+                "classifier_use_bias",
+                _config_value(config, "use_bias", False),
+            )
+        )
+        projection_size = int(math.ceil((2 * self.num_labels) / 256) * 256)
+        classifier_layers: list[nn.Module] = [
+            nn.LayerNorm(classifier_input),
+            nn.Linear(classifier_input, classifier_hidden, bias=use_bias),
+            nn.ReLU(),
+            nn.Dropout(classifier_dropout),
+            nn.Linear(classifier_hidden, projection_size, bias=use_bias),
+            nn.ReLU(),
+            nn.Dropout(classifier_dropout),
+        ]
+        if not sequence_task:
+            classifier_layers.extend(
+                [
+                    nn.Linear(projection_size, projection_size, bias=use_bias),
+                    nn.ReLU(),
+                ]
+            )
+        classifier_layers.append(nn.Linear(projection_size, self.num_labels, bias=use_bias))
+        self.classifier = nn.Sequential(*classifier_layers)
+
+    def _forward_transformer(
+        self,
+        embeddings: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        output_attentions: bool | None,
+        output_hidden_states: bool | None,
+    ) -> BaseModelOutput:
+        output_attentions = (
+            bool(output_attentions)
+            if output_attentions is not None
+            else bool(getattr(self.config, "output_attentions", False))
+        )
+        output_hidden_states = (
+            bool(output_hidden_states)
+            if output_hidden_states is not None
+            else bool(getattr(self.config, "output_hidden_states", False))
+        )
+        return self.transformer(
+            embeddings,
+            attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+
+
+class SequenceClassificationProbe(_ClassificationProbe):
+    """Protify-style sequence classifier over externally supplied residue embeddings."""
+
+    def __init__(self, config: Any, input_size: int) -> None:
+        super().__init__(config, input_size, sequence_task=True)
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        attention_mask: 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,
+    ) -> SequenceClassifierOutput | tuple[torch.Tensor, ...]:
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                embeddings.shape[:2],
+                device=embeddings.device,
+                dtype=torch.bool,
+            )
+        outputs = self._forward_transformer(
+            embeddings,
+            attention_mask,
+            output_attentions,
+            output_hidden_states,
+        )
+        if self.pooler is None:
+            raise AssertionError("Sequence classification requires a configured pooler.")
+        pooled = self.pooler(outputs.last_hidden_state, attention_mask)
+        logits = self.classifier(pooled)
+        loss = None
+        if labels is not None:
+            problem_type = resolve_problem_type(self.config, labels, num_labels=self.num_labels)
+            loss = sequence_classification_loss(
+                logits,
+                labels,
+                problem_type=problem_type,
+                num_labels=self.num_labels,
+            )
+        result = SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+        use_return_dict = (
+            bool(return_dict)
+            if return_dict is not None
+            else bool(getattr(self.config, "use_return_dict", True))
+        )
+        return result if use_return_dict else result.to_tuple()
+
+
+class TokenClassificationProbe(_ClassificationProbe):
+    """Protify-style residue classifier or regressor over supplied embeddings."""
+
+    def __init__(self, config: Any, input_size: int) -> None:
+        super().__init__(config, input_size, sequence_task=False)
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        attention_mask: 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,
+    ) -> TokenClassifierOutput | tuple[torch.Tensor, ...]:
+        outputs = self._forward_transformer(
+            embeddings,
+            attention_mask,
+            output_attentions,
+            output_hidden_states,
+        )
+        logits = self.classifier(outputs.last_hidden_state)
+        loss = None
+        if labels is not None:
+            problem_type = resolve_problem_type(self.config, labels, num_labels=self.num_labels)
+            loss = token_classification_loss(
+                logits,
+                labels,
+                problem_type=problem_type,
+                num_labels=self.num_labels,
+            )
+        result = TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+        use_return_dict = (
+            bool(return_dict)
+            if return_dict is not None
+            else bool(getattr(self.config, "use_return_dict", True))
+        )
+        return result if use_return_dict else result.to_tuple()
+
+
+__all__ = [
+    "ProbeSelfAttention",
+    "ProteinTransformerProbe",
+    "SequenceClassificationProbe",
+    "SwiGLU",
+    "TokenClassificationProbe",
+    "resolve_problem_type",
+    "sequence_classification_loss",
+    "token_classification_loss",
+]

fastplms/models/esmfold2/__init__.py

@@ -9,6 +9,12 @@ 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_classification import (
+        ESMFold2ExperimentalForSequenceClassification as ESMFold2ExperimentalForSequenceClassification,
+        ESMFold2ExperimentalForTokenClassification as ESMFold2ExperimentalForTokenClassification,
+        ESMFold2ForSequenceClassification as ESMFold2ForSequenceClassification,
+        ESMFold2ForTokenClassification as ESMFold2ForTokenClassification,
+    )
     from .modeling_esmfold2_experimental import (
         ESMFold2ExperimentalModel as ESMFold2ExperimentalModel,
     )
@@ -17,6 +23,10 @@ if TYPE_CHECKING:
 _EXPORT_MODULES = {
     "ESMFold2Config": ".configuration_esmfold2",
     "ESMFold2ExperimentalModel": ".modeling_esmfold2_experimental",
+    "ESMFold2ExperimentalForSequenceClassification": ".modeling_esmfold2_classification",
+    "ESMFold2ExperimentalForTokenClassification": ".modeling_esmfold2_classification",
+    "ESMFold2ForSequenceClassification": ".modeling_esmfold2_classification",
+    "ESMFold2ForTokenClassification": ".modeling_esmfold2_classification",
     "ESMFold2Model": ".modeling_esmfold2",
     "ESMFold2Output": ".modeling_esmfold2",
     "seed_context": ".reproducibility",

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -289,6 +289,23 @@ class ESMFold2Config(PretrainedConfig):
             )
         self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
 
+        self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe"))
+        if self.classifier_train_scope not in {"probe", "projection"}:
+            raise ValueError(
+                "classifier_train_scope must be 'probe' or 'projection', "
+                f"got {self.classifier_train_scope!r}."
+            )
+        self.classifier_probe_hidden_size = int(
+            kwargs.get("classifier_probe_hidden_size", 512)
+        )
+        self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4))
+        self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1))
+        self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096))
+        self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2))
+        self.classifier_pooling_types = list(
+            kwargs.get("classifier_pooling_types", ["mean"])
+        )
+
     def to_dict(self) -> dict[str, Any]:
         output = cast(dict[str, Any], super().to_dict())
         for name, _config_type in _NESTED_CONFIGS:

fastplms/models/esmfold2/esmfold2_processor.py

@@ -123,6 +123,12 @@ def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePrediction
             "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this "
             "input instead of silently emitting an all-zero pocket feature."
         )
+    if input.distogram_conditioning is not None:
+        raise NotImplementedError(
+            "ESMFold2 distogram conditioning is present in the upstream input schema but "
+            "the published ESMFold2 forward does not consume it. FastPLMs refuses this "
+            "input instead of silently ignoring the supplied distogram."
+        )
 
     cleaned: list[Any] = []
     for item in input.sequences:

fastplms/models/esmfold2/modeling_esmfold2.py

@@ -74,6 +74,7 @@ from .modeling_esmfold2_common import (
     maybe_subsample_msa,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
+    validate_prepared_auxiliary_inputs,
 )
 
 _ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj"
@@ -1417,6 +1418,12 @@ class ESMFold2Model(
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        pocket_feature: Tensor | None = None,
+        gt_coords: Tensor | None = None,
+        is_resolved: Tensor | None = None,
+        frames_idx: Tensor | None = None,
+        disto_cond: Tensor | None = None,
+        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -1432,6 +1439,12 @@ class ESMFold2Model(
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
+        validate_prepared_auxiliary_inputs(
+            pocket_feature=pocket_feature,
+            disto_cond=disto_cond,
+            disto_cond_mask=disto_cond_mask,
+        )
+        del gt_coords, is_resolved, frames_idx
         tok_mask = token_attention_mask
         atm_mask = atom_attention_mask
         disto_idx = distogram_atom_idx

fastplms/models/esmfold2/modeling_esmfold2_classification.py

@@ -0,0 +1,230 @@
+"""Sequence and residue prediction heads for ESMFold2 checkpoints."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from ..classification_probe import SequenceClassificationProbe, TokenClassificationProbe
+from .configuration_esmfold2 import ESMFold2Config
+from .embedding import _TOKEN_TO_ID, _VALID_RESIDUES, _encode_single_chain
+from .esmfold2_constants_esm3 import SEQUENCE_PAD_TOKEN
+from .modeling_esmfold2 import ESMFold2Model
+from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel
+
+
+ClassifierTrainScope = Literal["probe", "projection"]
+_VALID_RESIDUE_IDS = frozenset(_TOKEN_TO_ID[residue] for residue in _VALID_RESIDUES)
+
+
+class _ESMFold2ClassificationMixin:
+    """Run a task probe on the checkpoint-owned ESMC sequence projection."""
+
+    _classifier_config_type = ""
+    _keys_to_ignore_on_load_unexpected: list[str] = [r"\._extra_state$"]
+
+    def _initialize_classifier(self, classifier: nn.Module) -> None:
+        self.requires_grad_(False)
+        self.classifier = classifier
+        self.set_classifier_train_scope(self.config.classifier_train_scope)
+
+    def set_classifier_train_scope(self, scope: ClassifierTrainScope) -> None:
+        """Select whether fine-tuning updates only the probe or its input projection too."""
+
+        if scope not in {"probe", "projection"}:
+            raise ValueError(
+                "classifier_train_scope must be 'probe' or 'projection', "
+                f"got {scope!r}."
+            )
+        self.config.classifier_train_scope = scope
+        self.requires_grad_(False)
+        self.classifier.requires_grad_(True)
+        if scope == "projection":
+            self.language_model.base_z_combine.requires_grad_(True)
+            self.language_model.base_z_linear.requires_grad_(True)
+        if self._esmc is not None:
+            self._esmc.requires_grad_(False)
+
+    def load_esmc(self, *args: Any, **kwargs: Any) -> None:
+        super().load_esmc(*args, **kwargs)
+        if self._esmc is None:
+            raise RuntimeError("ESMFold2 ESMC loading completed without a backbone.")
+        self._esmc.requires_grad_(False)
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        if self._esmc is not None:
+            self._esmc.eval()
+        return self
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args: Any, **kwargs: Any):
+        if "config" not in kwargs:
+            kwargs["config"] = ESMFold2Config.from_pretrained(
+                pretrained_model_name_or_path, **kwargs
+            )
+        config = kwargs["config"]
+        if not isinstance(config, ESMFold2Config):
+            raise TypeError("ESMFold2 classifiers require an ESMFold2Config.")
+        if config.type != cls._classifier_config_type:
+            raise ValueError(
+                f"{cls.__name__} requires config.type={cls._classifier_config_type!r}, "
+                f"got {config.type!r}."
+            )
+        loaded = super().from_pretrained(
+            pretrained_model_name_or_path, *model_args, **kwargs
+        )
+        model = loaded[0] if isinstance(loaded, tuple) else loaded
+        model.set_classifier_train_scope(model.config.classifier_train_scope)
+        return loaded
+
+    def prepare_classifier_inputs(
+        self, sequence_or_sequences: str | list[str] | tuple[str, ...]
+    ) -> dict[str, Tensor]:
+        """Encode one or more ungapped single-chain proteins without special tokens."""
+
+        sequences = (
+            [sequence_or_sequences]
+            if isinstance(sequence_or_sequences, str)
+            else list(sequence_or_sequences)
+        )
+        if not sequences:
+            raise ValueError("prepare_classifier_inputs requires at least one sequence.")
+        encoded = [_encode_single_chain(sequence) for sequence in sequences]
+        sequence_length = max(map(len, encoded))
+        input_ids = torch.full(
+            (len(encoded), sequence_length),
+            SEQUENCE_PAD_TOKEN,
+            dtype=torch.long,
+            device=self.device,
+        )
+        attention_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        for batch_index, token_ids in enumerate(encoded):
+            length = len(token_ids)
+            input_ids[batch_index, :length] = torch.tensor(
+                token_ids, dtype=torch.long, device=self.device
+            )
+            attention_mask[batch_index, :length] = True
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def _classifier_embeddings(
+        self, input_ids: Tensor, attention_mask: Tensor | None
+    ) -> tuple[Tensor, Tensor]:
+        if input_ids.ndim != 2:
+            raise ValueError(
+                "ESMFold2 classifier input_ids must have shape (batch, residue), "
+                f"got {tuple(input_ids.shape)}."
+            )
+        if attention_mask is None:
+            attention_mask = input_ids.ne(SEQUENCE_PAD_TOKEN)
+        elif attention_mask.shape != input_ids.shape:
+            raise ValueError(
+                "ESMFold2 classifier attention_mask must match input_ids, got "
+                f"{tuple(attention_mask.shape)} and {tuple(input_ids.shape)}."
+            )
+        residue_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+        if not residue_mask.any(dim=1).all():
+            raise ValueError("Every ESMFold2 classifier input must contain a protein residue.")
+        if input_ids.masked_select(residue_mask).eq(SEQUENCE_PAD_TOKEN).any():
+            raise ValueError("ESMFold2 classifier padding tokens cannot be attended residues.")
+        residue_ids = input_ids.masked_select(residue_mask)
+        valid_residue_ids = torch.tensor(
+            sorted(_VALID_RESIDUE_IDS), dtype=input_ids.dtype, device=input_ids.device
+        )
+        if not torch.isin(residue_ids, valid_residue_ids).all():
+            raise ValueError(
+                "ESMFold2 classifiers accept residue-only single-chain protein inputs."
+            )
+
+        batch_size, sequence_length = input_ids.shape
+        residue_index = torch.arange(sequence_length, device=input_ids.device).expand(
+            batch_size, -1
+        )
+        asym_id = torch.zeros_like(input_ids)
+        mol_type = torch.zeros_like(input_ids)
+        with torch.no_grad():
+            hidden_states = self._compute_lm_hidden_states(
+                input_ids,
+                asym_id,
+                residue_index,
+                mol_type,
+                residue_mask,
+            )
+        embeddings = self.project_esmc_hidden_states(hidden_states, residue_mask)
+        return embeddings, residue_mask
+
+    def _classifier_forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        labels: Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ):
+        embeddings, residue_mask = self._classifier_embeddings(input_ids, attention_mask)
+        return self.classifier(
+            embeddings,
+            attention_mask=residue_mask,
+            labels=labels,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+class _ESMFold2SequenceClassificationMixin(_ESMFold2ClassificationMixin):
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__(config)
+        self._initialize_classifier(SequenceClassificationProbe(config, config.d_pair))
+
+    forward = _ESMFold2ClassificationMixin._classifier_forward
+
+
+class _ESMFold2TokenClassificationMixin(_ESMFold2ClassificationMixin):
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__(config)
+        self._initialize_classifier(TokenClassificationProbe(config, config.d_pair))
+
+    forward = _ESMFold2ClassificationMixin._classifier_forward
+
+
+class ESMFold2ForSequenceClassification(
+    _ESMFold2SequenceClassificationMixin, ESMFold2Model
+):
+    """Released ESMFold2 with a sequence classification or regression probe."""
+
+    _classifier_config_type = "release"
+
+
+class ESMFold2ForTokenClassification(_ESMFold2TokenClassificationMixin, ESMFold2Model):
+    """Released ESMFold2 with a residue classification or regression probe."""
+
+    _classifier_config_type = "release"
+
+
+class ESMFold2ExperimentalForSequenceClassification(
+    _ESMFold2SequenceClassificationMixin, ESMFold2ExperimentalModel
+):
+    """Experimental ESMFold2 with a sequence classification or regression probe."""
+
+    _classifier_config_type = "experimental"
+
+
+class ESMFold2ExperimentalForTokenClassification(
+    _ESMFold2TokenClassificationMixin, ESMFold2ExperimentalModel
+):
+    """Experimental ESMFold2 with a residue classification or regression probe."""
+
+    _classifier_config_type = "experimental"
+
+
+__all__ = [
+    "ESMFold2ExperimentalForSequenceClassification",
+    "ESMFold2ExperimentalForTokenClassification",
+    "ESMFold2ForSequenceClassification",
+    "ESMFold2ForTokenClassification",
+]

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -57,6 +57,14 @@ MSA_CONDITIONING_INPUT_NAMES = (
     "deletion_value",
     "deletion_mean",
 )
+PREPARED_AUXILIARY_INPUT_NAMES = (
+    "pocket_feature",
+    "gt_coords",
+    "is_resolved",
+    "frames_idx",
+    "disto_cond",
+    "disto_cond_mask",
+)
 
 
 def validate_kernel_backend(backend: str | None) -> None:
@@ -105,6 +113,29 @@ def validate_msa_conditioning_inputs(
         )
 
 
+def validate_prepared_auxiliary_inputs(
+    *,
+    pocket_feature: Tensor | None,
+    disto_cond: Tensor | None,
+    disto_cond_mask: Tensor | None,
+) -> None:
+    """Accept inert prepared fields and reject unsupported conditioning."""
+
+    if pocket_feature is not None and torch.any(pocket_feature != 0).item():
+        raise NotImplementedError(
+            "The published ESMFold2 forward does not consume pocket conditioning; "
+            "nonzero pocket_feature values are unsupported."
+        )
+    distogram_is_active = (disto_cond is not None and torch.any(disto_cond != 0).item()) or (
+        disto_cond_mask is not None and torch.any(disto_cond_mask).item()
+    )
+    if distogram_is_active:
+        raise NotImplementedError(
+            "The published ESMFold2 forward does not consume distogram conditioning; "
+            "nonzero disto_cond or disto_cond_mask values are unsupported."
+        )
+
+
 def _fused_active(module: nn.Module, tensor: Tensor) -> bool:
     """Return whether an optional fused implementation can handle this call."""
     return (

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -59,6 +59,7 @@ from .modeling_esmfold2_common import (
     gather_token_to_atom,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
+    validate_prepared_auxiliary_inputs,
 )
 
 _EPS = 1e-5
@@ -689,6 +690,12 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        pocket_feature: Tensor | None = None,
+        gt_coords: Tensor | None = None,
+        is_resolved: Tensor | None = None,
+        frames_idx: Tensor | None = None,
+        disto_cond: Tensor | None = None,
+        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -704,6 +711,12 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
+        validate_prepared_auxiliary_inputs(
+            pocket_feature=pocket_feature,
+            disto_cond=disto_cond,
+            disto_cond_mask=disto_cond_mask,
+        )
+        del gt_coords, is_resolved, frames_idx
         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

modeling_fastplms.py

@@ -8,11 +8,12 @@ import sys
 import tempfile
 from io import BytesIO
 from pathlib import Path
+from typing import ClassVar
 from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "5f557b57da6c13e8b26b2a144ab0983b9621d853a3141a3a52e912d2de26289a":
+if RUNTIME_HASH != "dacab314c14443a90eb43e0dfaeca56233336bb0f0ca617756b8353e4f5e3fc2":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []
@@ -179,9 +180,14 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_181 = _import_without_bytecode("fastplms.models.esmfold2.configuration_esmfold2")
-ESMFold2Config = _module_181.ESMFold2Config
+_module_182 = _import_without_bytecode("fastplms.models.esmfold2.configuration_esmfold2")
+ESMFold2Config = _module_182.ESMFold2Config
 ESMFold2Config.__module__ = __name__
-_module_184 = _import_without_bytecode("fastplms.models.esmfold2.modeling_esmfold2")
-ESMFold2Model = _module_184.ESMFold2Model
+_module_185 = _import_without_bytecode("fastplms.models.esmfold2.modeling_esmfold2")
+ESMFold2Model = _module_185.ESMFold2Model
 ESMFold2Model.__module__ = __name__
+_module_188 = _import_without_bytecode("fastplms.models.esmfold2.modeling_esmfold2_classification")
+ESMFold2ForSequenceClassification = _module_188.ESMFold2ForSequenceClassification
+ESMFold2ForSequenceClassification.__module__ = __name__
+ESMFold2ForTokenClassification = _module_188.ESMFold2ForTokenClassification
+ESMFold2ForTokenClassification.__module__ = __name__