Update FastPLMs files

README.md

@@ -14,15 +14,16 @@ This checkpoint contains 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`.
+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 | Unavailable for this structure-only checkpoint |
 | Test-time training | Unavailable: the checkpoint has no trained MLM head |
 | 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/FastESMFold"
+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
@@ -132,8 +172,8 @@ folding requests raise.
 ## 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`
+- 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: `default`
 - BF16 execution: `fp32_parameters_autocast`
@@ -147,8 +187,8 @@ folding requests raise.
 ## Release record
 
 - FastPLMs weights: `Synthyra/FastESMFold`
-- 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: `facebook/esmfold_v1`
 - Artifact source: `fast`
 - State transform: `esmfold_meta_to_fastplms_v1`

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/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/esmfold/modeling_fast_esmfold.py

@@ -37,6 +37,22 @@ from transformers.models.esm.openfold_utils import residue_constants
 
 from fastplms.models._esm_rotary import RotaryEmbedding
 
+try:
+    from fastplms.models.classification_probe import (
+        SequenceClassificationProbe,
+        TokenClassificationProbe,
+    )
+except ModuleNotFoundError as error:
+    _COMPOSITE_CLASSIFIER_NAMES = (
+        "SequenceClassificationProbe",
+        "TokenClassificationProbe",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_CLASSIFIER_NAMES
+    ):
+        raise
+    # Hub composites define the shared classifier probes before this source.
+
 
 # Hub composite artifacts define these shared names earlier in the assembled file.
 try:
@@ -531,13 +547,42 @@ class FastEsmBackbone(nn.Module):
 class FastEsmFoldConfig(EsmConfig):
     model_type = "fast_esmfold"
 
-    def __init__(self, attn_backend: str | None = None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        attn_backend: str | None = None,
+        classifier_train_scope: str = "probe",
+        classifier_pooling_types: list[str] | None = None,
+        classifier_probe_hidden_size: int = 512,
+        classifier_probe_num_heads: int = 4,
+        classifier_probe_dropout: float = 0.1,
+        classifier_hidden_size: int = 4096,
+        classifier_dropout: float = 0.2,
+        classifier_use_bias: bool = False,
+        **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
+        if classifier_train_scope not in {"probe", "projection"}:
+            raise ValueError(
+                "classifier_train_scope must be 'probe' or 'projection', got "
+                f"{classifier_train_scope!r}."
+            )
+        self.classifier_train_scope = classifier_train_scope
+        self.classifier_pooling_types = (
+            ["mean"]
+            if classifier_pooling_types is None
+            else list(classifier_pooling_types)
+        )
+        self.classifier_probe_hidden_size = classifier_probe_hidden_size
+        self.classifier_probe_num_heads = classifier_probe_num_heads
+        self.classifier_probe_dropout = classifier_probe_dropout
+        self.classifier_hidden_size = classifier_hidden_size
+        self.classifier_dropout = classifier_dropout
+        self.classifier_use_bias = classifier_use_bias
 
 
 class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
@@ -873,3 +918,139 @@ class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
 
         del sequence, return_pdb_string
         self._ttt_unavailable()
+
+
+class _FastEsmFoldClassificationMixin:
+    """Expose ESMFold's checkpoint-trained residue projection to task probes."""
+
+    _classifier_probe_class: type[nn.Module]
+
+    def __init__(self, config: FastEsmFoldConfig) -> None:
+        super().__init__(config)
+        sequence_state_dim = config.esmfold_config.trunk.sequence_state_dim
+        self.classifier = self._classifier_probe_class(config, sequence_state_dim)
+        self.classifier.apply(self._init_weights)
+        self._configure_classifier_train_scope()
+
+    def _configure_classifier_train_scope(self) -> None:
+        """Apply the serialized fine-tuning boundary deterministically."""
+
+        self.requires_grad_(False)
+        self.classifier.requires_grad_(True)
+        if self.config.classifier_train_scope == "projection":
+            self.esm_s_combine.requires_grad_(True)
+            self.esm_s_mlp.requires_grad_(True)
+
+    @staticmethod
+    def prepare_classifier_inputs(
+        sequences: str | list[str],
+    ) -> dict[str, torch.Tensor]:
+        """Encode single-chain proteins as residue-only ESMFold inputs."""
+
+        sequence_batch = [sequences] if isinstance(sequences, str) else list(sequences)
+        if not sequence_batch:
+            raise ValueError("At least one protein sequence is required.")
+
+        supported_residues = residue_constants.restype_order_with_x
+        encoded_sequences: list[torch.Tensor] = []
+        for sequence in sequence_batch:
+            if not isinstance(sequence, str):
+                raise TypeError("Each protein sequence must be a string.")
+            normalized = sequence.upper()
+            if not normalized:
+                raise ValueError("Protein sequences must not be empty.")
+            invalid = sorted(set(normalized).difference(supported_residues))
+            if invalid:
+                raise ValueError(
+                    "ESMFold classifiers accept single-chain proteins containing "
+                    "the 20 standard amino acids or X; unsupported residues: "
+                    f"{', '.join(invalid)}."
+                )
+            encoded_sequences.append(
+                torch.tensor(
+                    [supported_residues[residue] for residue in normalized],
+                    dtype=torch.int64,
+                )
+            )
+
+        input_ids = collate_dense_tensors(encoded_sequences, pad_v=0)
+        attention_mask = collate_dense_tensors(
+            [torch.ones_like(sequence) for sequence in encoded_sequences],
+            pad_v=0,
+        )
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    def _classifier_features(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_ids.ndim != 2:
+            raise ValueError(
+                "input_ids must have shape (batch, residue), got "
+                f"{tuple(input_ids.shape)}."
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        elif attention_mask.shape != input_ids.shape:
+            raise ValueError(
+                "attention_mask must match input_ids, got "
+                f"{tuple(attention_mask.shape)} and {tuple(input_ids.shape)}."
+            )
+        attention_mask = attention_mask.to(device=input_ids.device)
+        if torch.any(attention_mask.sum(dim=-1) == 0):
+            raise ValueError("Every classifier input must contain at least one residue.")
+
+        esmaa = self.af2_idx_to_esm_idx(input_ids, attention_mask)
+        layer_states = self.compute_language_model_representations(esmaa)
+        layer_states = layer_states.to(self.esm_s_combine.dtype).detach()
+        if self.config.esmfold_config.esm_ablate_sequence:
+            layer_states = layer_states * 0
+        mixed_states = (
+            self.esm_s_combine.softmax(0).unsqueeze(0) @ layer_states
+        ).squeeze(2)
+        residue_embeddings = self.esm_s_mlp(mixed_states)
+        residue_embeddings = residue_embeddings * attention_mask.unsqueeze(-1).to(
+            residue_embeddings.dtype
+        )
+        return residue_embeddings, attention_mask
+
+    def forward(
+        self,
+        input_ids: 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,
+    ) -> Any:
+        residue_embeddings, attention_mask = self._classifier_features(
+            input_ids,
+            attention_mask,
+        )
+        return self.classifier(
+            residue_embeddings,
+            attention_mask=attention_mask,
+            labels=labels,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+class FastEsmForSequenceClassification(
+    _FastEsmFoldClassificationMixin,
+    FastEsmForProteinFolding,
+):
+    """Sequence classification or regression over ESMFold residue features."""
+
+    _classifier_probe_class = SequenceClassificationProbe
+
+
+class FastEsmForTokenClassification(
+    _FastEsmFoldClassificationMixin,
+    FastEsmForProteinFolding,
+):
+    """Residue classification or regression over ESMFold residue features."""
+
+    _classifier_probe_class = TokenClassificationProbe

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 != "9a0627bd8e6f68ac389ee97690f802284cffc7e7a7a537880aacd8a7ae440cc3":
+if RUNTIME_HASH != "027d8ff36110255e68e1050fa838f94dc3e5d6255611e73b934d543993cd823a":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []
@@ -179,8 +180,12 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_181 = _import_without_bytecode("fastplms.models.esmfold.modeling_fast_esmfold")
-FastEsmFoldConfig = _module_181.FastEsmFoldConfig
+_module_182 = _import_without_bytecode("fastplms.models.esmfold.modeling_fast_esmfold")
+FastEsmFoldConfig = _module_182.FastEsmFoldConfig
 FastEsmFoldConfig.__module__ = __name__
-FastEsmForProteinFolding = _module_181.FastEsmForProteinFolding
+FastEsmForProteinFolding = _module_182.FastEsmForProteinFolding
 FastEsmForProteinFolding.__module__ = __name__
+FastEsmForSequenceClassification = _module_182.FastEsmForSequenceClassification
+FastEsmForSequenceClassification.__module__ = __name__
+FastEsmForTokenClassification = _module_182.FastEsmForTokenClassification
+FastEsmForTokenClassification.__module__ = __name__