File size: 9,096 Bytes
77a9c3b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 | """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",
]
|