InstaDeepAI/segment_nt
9211
1# coding=utf-82# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" PyTorch ESM model."""16 17import math18from dataclasses import dataclass19from typing import List, Optional, Tuple, Union20 21import torch22import torch.utils.checkpoint23from torch import nn24from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss, SiLU25from transformers.file_utils import (26 add_code_sample_docstrings,27 add_start_docstrings,28 add_start_docstrings_to_model_forward,29)30from transformers.modeling_outputs import (31 BaseModelOutputWithPastAndCrossAttentions,32 BaseModelOutputWithPoolingAndCrossAttentions,33 MaskedLMOutput,34 SequenceClassifierOutput,35 TokenClassifierOutput,36)37from transformers.modeling_utils import (38 PreTrainedModel,39 find_pruneable_heads_and_indices,40 prune_linear_layer,41)42from transformers.utils import logging43 44from .segment_nt_config import SegmentNTConfig45 46logger = logging.get_logger(__name__)47 48_CHECKPOINT_FOR_DOC = "facebook/esm2_t6_8M_UR50D"49_CONFIG_FOR_DOC = "SegmentNTConfig"50 51ESM_PRETRAINED_MODEL_ARCHIVE_LIST = [52 "facebook/esm2_t6_8M_UR50D",53 "facebook/esm2_t12_35M_UR50D",54 # This is not a complete list of all ESM models!55 # See all ESM models at https://huggingface.co/models?filter=esm56]57 58 59def rotate_half(x):60 x1, x2 = x.chunk(2, dim=-1)61 return torch.cat((-x2, x1), dim=-1)62 63 64def apply_rotary_pos_emb(x, cos, sin):65 cos = cos[:, :, : x.shape[-2], :]66 sin = sin[:, :, : x.shape[-2], :]67 68 return (x * cos) + (rotate_half(x) * sin)69 70 71def gelu(x):72 """73 This is the gelu implementation from the original ESM repo. Using F.gelu yields subtly wrong results.74 """75 return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))76 77 78def symmetrize(x):79 "Make layer symmetric in final two dimensions, used for contact prediction."80 return x + x.transpose(-1, -2)81 82 83def average_product_correct(x):84 "Perform average product correct, used for contact prediction."85 a1 = x.sum(-1, keepdims=True)86 a2 = x.sum(-2, keepdims=True)87 a12 = x.sum((-1, -2), keepdims=True)88 89 avg = a1 * a290 avg.div_(a12) # in-place to reduce memory91 normalized = x - avg92 return normalized93 94@dataclass95class RotaryEmbeddingConfig:96 """97 Parameters to initialize the RotaryEmbedding layer. The rescaling factor allows98 to adapt the rotary embeddings to larger lengths than what was used for training.99 One of this strategy is presented in the Yarn paper: https://arxiv.org/pdf/2309.00071.pdf. # noqa100 101 Args:102 103 """104 105 rescaling_factor: Optional[float]106 107class RotaryEmbedding(torch.nn.Module):108 """109 Rotary position embeddings based on those in110 [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation111 matrices which depend on their relative positions.112 """113 114 def __init__(self, dim: int, rotary_embedding_config: RotaryEmbeddingConfig):115 super().__init__()116 117 # Extract argument from the config118 self.rescaling_factor = rotary_embedding_config.rescaling_factor119 self.upper_freq = 10000120 self.dim = dim121 122 self._seq_len_cached = None123 self._cos_cached = None124 self._sin_cached = None125 126 127 128 def _compute_cos_sin_tables(self, x, inv_freq, seq_dimension=2):129 seq_len = x.shape[seq_dimension]130 131 # Reset the tables if the sequence length has changed,132 # or if we're on a new device (possibly due to tracing for instance)133 self._seq_len_cached = seq_len134 t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(135 inv_freq136 )137 freqs = torch.outer(t, inv_freq)138 emb = torch.cat((freqs, freqs), dim=-1).to(x.device)139 140 self._cos_cached = emb.cos()[None, None, :, :]141 self._sin_cached = emb.sin()[None, None, :, :]142 143 return self._cos_cached, self._sin_cached144 145 def forward(146 self, q: torch.Tensor, k: torch.Tensor147 ) -> Tuple[torch.Tensor, torch.Tensor]:148 149 if self.rescaling_factor is None:150 inv_freq = 1.0 / (self.upper_freq ** (torch.arange(0, self.dim, 2).float() / self.dim))151 else:152 updated_base = self.upper_freq * (153 self.rescaling_factor ** (self.dim / (self.dim - 2))154 )155 inv_freq = 1.0 / (156 updated_base ** (torch.arange(0, self.dim, 2).float() / self.dim)157 )158 159 self._cos_cached, self._sin_cached = self._compute_cos_sin_tables(160 k, inv_freq, seq_dimension=-2, 161 )162 163 return (164 apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),165 apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),166 )167 168 169 170class EsmContactPredictionHead(nn.Module):171 """Performs symmetrization, apc, and computes a logistic regression on the output features"""172 173 def __init__(174 self,175 in_features: int,176 bias=True,177 eos_idx: int = 2,178 ):179 super().__init__()180 self.in_features = in_features181 self.eos_idx = eos_idx182 self.regression = nn.Linear(in_features, 1, bias)183 self.activation = nn.Sigmoid()184 185 def forward(self, tokens, attentions):186 # remove eos token attentions187 eos_mask = tokens.ne(self.eos_idx).to(attentions)188 eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2)189 attentions = attentions * eos_mask[:, None, None, :, :]190 attentions = attentions[..., :-1, :-1]191 # remove cls token attentions192 attentions = attentions[..., 1:, 1:]193 batch_size, layers, heads, seqlen, _ = attentions.size()194 attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen)195 196 # features: batch x channels x tokens x tokens (symmetric)197 attentions = attentions.to(198 self.regression.weight.device199 ) # attentions always float32, may need to convert to float16200 attentions = average_product_correct(symmetrize(attentions))201 attentions = attentions.permute(0, 2, 3, 1)202 return self.activation(self.regression(attentions).squeeze(3))203 204 205class EsmEmbeddings(nn.Module):206 """207 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.208 """209 210 def __init__(self, config):211 super().__init__()212 self.word_embeddings = nn.Embedding(213 config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id214 )215 216 if config.emb_layer_norm_before:217 self.layer_norm = nn.LayerNorm(218 config.hidden_size, eps=config.layer_norm_eps219 )220 else:221 self.layer_norm = None222 self.dropout = nn.Dropout(config.hidden_dropout_prob)223 # position_ids (1, len position emb) is contiguous in memory and exported when serialized224 self.position_embedding_type = getattr(225 config, "position_embedding_type", "absolute"226 )227 self.register_buffer(228 "position_ids",229 torch.arange(config.max_position_embeddings).expand((1, -1)),230 persistent=False,231 )232 233 self.padding_idx = config.pad_token_id234 self.position_embeddings = nn.Embedding(235 config.max_position_embeddings,236 config.hidden_size,237 padding_idx=self.padding_idx,238 )239 self.token_dropout = config.token_dropout240 self.mask_token_id = config.mask_token_id241 242 def forward(243 self,244 input_ids=None,245 attention_mask=None,246 position_ids=None,247 inputs_embeds=None,248 past_key_values_length=0,249 ):250 if position_ids is None:251 if input_ids is not None:252 # Create the position ids from the input token ids. Any padded tokens remain padded.253 position_ids = create_position_ids_from_input_ids(254 input_ids, self.padding_idx, past_key_values_length255 )256 else:257 position_ids = self.create_position_ids_from_inputs_embeds(258 inputs_embeds259 )260 261 if inputs_embeds is None:262 inputs_embeds = self.word_embeddings(input_ids)263 264 # Note that if we want to support ESM-1 (not 1b!) in future then we need to support an265 # embedding_scale factor here.266 embeddings = inputs_embeds267 268 # Matt: ESM has the option to handle masking in MLM in a slightly unusual way. If the token_dropout269 # flag is False then it is handled in the same was as BERT/RoBERTa. If it is set to True, however,270 # masked tokens are treated as if they were selected for input dropout and zeroed out.271 # This "mask-dropout" is compensated for when masked tokens are not present, by scaling embeddings by272 # a factor of (fraction of unmasked tokens during training) / (fraction of unmasked tokens in sample).273 # This is analogous to the way that dropout layers scale down outputs during evaluation when not274 # actually dropping out values (or, equivalently, scale up their un-dropped outputs in training).275 if self.token_dropout:276 embeddings.masked_fill_(277 (input_ids == self.mask_token_id).unsqueeze(-1), 0.0278 )279 mask_ratio_train = (280 0.15 * 0.8281 ) # Hardcoded as the ratio used in all ESM model training runs282 src_lengths = attention_mask.sum(-1)283 mask_ratio_observed = (input_ids == self.mask_token_id).sum(284 -1285 ).float() / src_lengths286 embeddings = (287 embeddings288 * (1 - mask_ratio_train)289 / (1 - mask_ratio_observed)[:, None, None]290 ).to(embeddings.dtype)291 292 if self.position_embedding_type == "absolute":293 position_embeddings = self.position_embeddings(position_ids)294 embeddings += position_embeddings295 296 if self.layer_norm is not None:297 embeddings = self.layer_norm(embeddings)298 if attention_mask is not None:299 embeddings = (embeddings * attention_mask.unsqueeze(-1)).to(300 embeddings.dtype301 )302 # Matt: I think this line was copied incorrectly from BERT, disabling it for now.303 # embeddings = self.dropout(embeddings)304 return embeddings305 306 def create_position_ids_from_inputs_embeds(self, inputs_embeds):307 """308 We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.309 310 Args:311 inputs_embeds: torch.Tensor312 313 Returns: torch.Tensor314 """315 input_shape = inputs_embeds.size()[:-1]316 sequence_length = input_shape[1]317 318 position_ids = torch.arange(319 self.padding_idx + 1,320 sequence_length + self.padding_idx + 1,321 dtype=torch.long,322 device=inputs_embeds.device,323 )324 return position_ids.unsqueeze(0).expand(input_shape)325 326 327class EsmSelfAttention(nn.Module):328 def __init__(self, config, position_embedding_type=None):329 super().__init__()330 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(331 config, "embedding_size"332 ):333 raise ValueError(334 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "335 f"heads ({config.num_attention_heads})"336 )337 338 self.num_attention_heads = config.num_attention_heads339 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)340 self.all_head_size = self.num_attention_heads * self.attention_head_size341 342 self.query = nn.Linear(config.hidden_size, self.all_head_size)343 self.key = nn.Linear(config.hidden_size, self.all_head_size)344 self.value = nn.Linear(config.hidden_size, self.all_head_size)345 346 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)347 self.position_embedding_type = position_embedding_type or getattr(348 config, "position_embedding_type", "absolute"349 )350 self.rotary_embeddings = None351 if (352 self.position_embedding_type == "relative_key"353 or self.position_embedding_type == "relative_key_query"354 ):355 self.max_position_embeddings = config.max_position_embeddings356 self.distance_embedding = nn.Embedding(357 2 * config.max_position_embeddings - 1, self.attention_head_size358 )359 elif self.position_embedding_type == "rotary":360 # Initiliaze rotary embedding config361 rescaling_factor = config.rescaling_factor362 rotary_embedding_config = RotaryEmbeddingConfig(rescaling_factor=rescaling_factor)363 364 self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size, rotary_embedding_config=rotary_embedding_config)365 366 self.is_decoder = config.is_decoder367 368 def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:369 new_x_shape = x.size()[:-1] + (370 self.num_attention_heads,371 self.attention_head_size,372 )373 x = x.view(new_x_shape)374 return x.permute(0, 2, 1, 3)375 376 def forward(377 self,378 hidden_states: torch.Tensor,379 attention_mask: Optional[torch.FloatTensor] = None,380 head_mask: Optional[torch.FloatTensor] = None,381 encoder_hidden_states: Optional[torch.FloatTensor] = None,382 encoder_attention_mask: Optional[torch.FloatTensor] = None,383 past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,384 output_attentions: Optional[bool] = False,385 ) -> Tuple[torch.Tensor]:386 mixed_query_layer = self.query(hidden_states)387 388 # If this is instantiated as a cross-attention module, the keys389 # and values come from an encoder; the attention mask needs to be390 # such that the encoder's padding tokens are not attended to.391 is_cross_attention = encoder_hidden_states is not None392 393 if is_cross_attention and past_key_value is not None:394 # reuse k,v, cross_attentions395 key_layer = past_key_value[0]396 value_layer = past_key_value[1]397 attention_mask = encoder_attention_mask398 elif is_cross_attention:399 key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))400 value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))401 attention_mask = encoder_attention_mask402 elif past_key_value is not None:403 key_layer = self.transpose_for_scores(self.key(hidden_states))404 value_layer = self.transpose_for_scores(self.value(hidden_states))405 key_layer = torch.cat([past_key_value[0], key_layer], dim=2)406 value_layer = torch.cat([past_key_value[1], value_layer], dim=2)407 else:408 key_layer = self.transpose_for_scores(self.key(hidden_states))409 value_layer = self.transpose_for_scores(self.value(hidden_states))410 411 query_layer = self.transpose_for_scores(mixed_query_layer)412 413 # Matt: Our BERT model (which this code was derived from) scales attention logits down by sqrt(head_dim).414 # ESM scales the query down by the same factor instead. Modulo numerical stability these are equivalent,415 # but not when rotary embeddings get involved. Therefore, we scale the query here to match the original416 # ESM code and fix rotary embeddings.417 query_layer = query_layer * self.attention_head_size**-0.5418 419 if self.is_decoder:420 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.421 # Further calls to cross_attention layer can then reuse all cross-attention422 # key/value_states (first "if" case)423 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of424 # all previous decoder key/value_states. Further calls to uni-directional self-attention425 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)426 # if encoder bi-directional self-attention `past_key_value` is always `None`427 past_key_value = (key_layer, value_layer)428 429 if self.position_embedding_type == "rotary":430 query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer)431 432 # Take the dot product between "query" and "key" to get the raw attention scores.433 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))434 435 if (436 self.position_embedding_type == "relative_key"437 or self.position_embedding_type == "relative_key_query"438 ):439 seq_length = hidden_states.size()[1]440 position_ids_l = torch.arange(441 seq_length, dtype=torch.long, device=hidden_states.device442 ).view(-1, 1)443 position_ids_r = torch.arange(444 seq_length, dtype=torch.long, device=hidden_states.device445 ).view(1, -1)446 distance = position_ids_l - position_ids_r447 positional_embedding = self.distance_embedding(448 distance + self.max_position_embeddings - 1449 )450 positional_embedding = positional_embedding.to(451 dtype=query_layer.dtype452 ) # fp16 compatibility453 454 if self.position_embedding_type == "relative_key":455 relative_position_scores = torch.einsum(456 "bhld,lrd->bhlr", query_layer, positional_embedding457 )458 attention_scores = attention_scores + relative_position_scores459 elif self.position_embedding_type == "relative_key_query":460 relative_position_scores_query = torch.einsum(461 "bhld,lrd->bhlr", query_layer, positional_embedding462 )463 relative_position_scores_key = torch.einsum(464 "bhrd,lrd->bhlr", key_layer, positional_embedding465 )466 attention_scores = (467 attention_scores468 + relative_position_scores_query469 + relative_position_scores_key470 )471 472 if attention_mask is not None:473 # Apply the attention mask is (precomputed for all layers in EsmModel forward() function)474 attention_scores = attention_scores + attention_mask475 476 # Normalize the attention scores to probabilities.477 attention_probs = nn.functional.softmax(attention_scores, dim=-1)478 479 # This is actually dropping out entire tokens to attend to, which might480 # seem a bit unusual, but is taken from the original Transformer paper.481 attention_probs = self.dropout(attention_probs)482 483 # Mask heads if we want to484 if head_mask is not None:485 attention_probs = attention_probs * head_mask486 487 context_layer = torch.matmul(attention_probs, value_layer)488 489 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()490 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)491 context_layer = context_layer.view(new_context_layer_shape)492 493 outputs = (494 (context_layer, attention_probs) if output_attentions else (context_layer,)495 )496 497 if self.is_decoder:498 outputs = outputs + (past_key_value,)499 return outputs500 501 502class EsmSelfOutput(nn.Module):503 def __init__(self, config):504 super().__init__()505 self.dense = nn.Linear(config.hidden_size, config.hidden_size)506 self.dropout = nn.Dropout(config.hidden_dropout_prob)507 508 def forward(self, hidden_states, input_tensor):509 hidden_states = self.dense(hidden_states)510 hidden_states = self.dropout(hidden_states)511 hidden_states += input_tensor512 return hidden_states513 514 515class EsmAttention(nn.Module):516 def __init__(self, config):517 super().__init__()518 self.self = EsmSelfAttention(config)519 self.output = EsmSelfOutput(config)520 self.pruned_heads = set()521 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)522 523 def prune_heads(self, heads):524 if len(heads) == 0:525 return526 heads, index = find_pruneable_heads_and_indices(527 heads,528 self.self.num_attention_heads,529 self.self.attention_head_size,530 self.pruned_heads,531 )532 533 # Prune linear layers534 self.self.query = prune_linear_layer(self.self.query, index)535 self.self.key = prune_linear_layer(self.self.key, index)536 self.self.value = prune_linear_layer(self.self.value, index)537 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)538 539 # Update hyper params and store pruned heads540 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)541 self.self.all_head_size = (542 self.self.attention_head_size * self.self.num_attention_heads543 )544 self.pruned_heads = self.pruned_heads.union(heads)545 546 def forward(547 self,548 hidden_states,549 attention_mask=None,550 head_mask=None,551 encoder_hidden_states=None,552 encoder_attention_mask=None,553 past_key_value=None,554 output_attentions=False,555 ):556 hidden_states_ln = self.LayerNorm(hidden_states)557 self_outputs = self.self(558 hidden_states_ln,559 attention_mask,560 head_mask,561 encoder_hidden_states,562 encoder_attention_mask,563 past_key_value,564 output_attentions,565 )566 attention_output = self.output(self_outputs[0], hidden_states)567 outputs = (attention_output,) + self_outputs[568 1:569 ] # add attentions if we output them570 return outputs571 572 573class EsmIntermediate(nn.Module):574 def __init__(self, config):575 super().__init__()576 577 self.dense = nn.Linear(578 config.hidden_size,579 int(config.intermediate_size * 2),580 bias=config.add_bias_fnn,581 )582 self.activation_fn = SiLU()583 584 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:585 hidden_states = self.dense(hidden_states)586 587 # GLU588 x1, x2 = hidden_states.split(int(hidden_states.size(-1) / 2), -1)589 hidden_states = self.activation_fn(x1) * x2590 591 return hidden_states592 593 594class EsmOutput(nn.Module):595 def __init__(self, config):596 super().__init__()597 self.dense = nn.Linear(598 config.intermediate_size, config.hidden_size, bias=config.add_bias_fnn599 )600 self.dropout = nn.Dropout(config.hidden_dropout_prob)601 602 def forward(self, hidden_states, input_tensor):603 hidden_states = self.dense(hidden_states)604 hidden_states = self.dropout(hidden_states)605 hidden_states += input_tensor606 return hidden_states607 608 609class EsmLayer(nn.Module):610 def __init__(self, config):611 super().__init__()612 self.chunk_size_feed_forward = config.chunk_size_feed_forward613 self.seq_len_dim = 1614 self.attention = EsmAttention(config)615 self.is_decoder = config.is_decoder616 self.add_cross_attention = config.add_cross_attention617 if self.add_cross_attention:618 if not self.is_decoder:619 raise RuntimeError(620 f"{self} should be used as a decoder model if cross attention is added"621 )622 self.crossattention = EsmAttention(config)623 self.intermediate = EsmIntermediate(config)624 self.output = EsmOutput(config)625 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)626 627 def forward(628 self,629 hidden_states,630 attention_mask=None,631 head_mask=None,632 encoder_hidden_states=None,633 encoder_attention_mask=None,634 past_key_value=None,635 output_attentions=False,636 ):637 # decoder uni-directional self-attention cached key/values tuple is at positions 1,2638 self_attn_past_key_value = (639 past_key_value[:2] if past_key_value is not None else None640 )641 self_attention_outputs = self.attention(642 hidden_states,643 attention_mask,644 head_mask,645 output_attentions=output_attentions,646 past_key_value=self_attn_past_key_value,647 )648 attention_output = self_attention_outputs[0]649 650 # if decoder, the last output is tuple of self-attn cache651 if self.is_decoder:652 outputs = self_attention_outputs[1:-1]653 present_key_value = self_attention_outputs[-1]654 else:655 outputs = self_attention_outputs[656 1:657 ] # add self attentions if we output attention weights658 659 cross_attn_present_key_value = None660 if self.is_decoder and encoder_hidden_states is not None:661 if not hasattr(self, "crossattention"):662 raise AttributeError(663 f"If `encoder_hidden_states` are passed, {self} has to be instantiated"664 " with cross-attention layers by setting `config.add_cross_attention=True`"665 )666 667 # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple668 cross_attn_past_key_value = (669 past_key_value[-2:] if past_key_value is not None else None670 )671 cross_attention_outputs = self.crossattention(672 attention_output,673 attention_mask,674 head_mask,675 encoder_hidden_states,676 encoder_attention_mask,677 cross_attn_past_key_value,678 output_attentions,679 )680 attention_output = cross_attention_outputs[0]681 outputs = (682 outputs + cross_attention_outputs[1:-1]683 ) # add cross attentions if we output attention weights684 685 # add cross-attn cache to positions 3,4 of present_key_value tuple686 cross_attn_present_key_value = cross_attention_outputs[-1]687 present_key_value = present_key_value + cross_attn_present_key_value688 689 layer_output = self.feed_forward_chunk(attention_output)690 691 outputs = (layer_output,) + outputs692 693 # if decoder, return the attn key/values as the last output694 if self.is_decoder:695 outputs = outputs + (present_key_value,)696 return outputs697 698 def feed_forward_chunk(self, attention_output):699 attention_output_ln = self.LayerNorm(attention_output)700 intermediate_output = self.intermediate(attention_output_ln)701 layer_output = self.output(intermediate_output, attention_output)702 return layer_output703 704 705class EsmEncoder(nn.Module):706 def __init__(self, config):707 super().__init__()708 self.config = config709 self.layer = nn.ModuleList(710 [EsmLayer(config) for _ in range(config.num_hidden_layers)]711 )712 self.emb_layer_norm_after = nn.LayerNorm(713 config.hidden_size, eps=config.layer_norm_eps714 )715 self.gradient_checkpointing = False716 717 def forward(718 self,719 hidden_states,720 attention_mask=None,721 head_mask=None,722 encoder_hidden_states=None,723 encoder_attention_mask=None,724 past_key_values=None,725 use_cache=None,726 output_attentions=False,727 output_hidden_states=False,728 return_dict=True,729 ):730 if self.gradient_checkpointing and self.training:731 if use_cache:732 logger.warning_once(733 "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "734 "`use_cache=False`..."735 )736 use_cache = False737 all_hidden_states = () if output_hidden_states else None738 all_self_attentions = () if output_attentions else None739 all_cross_attentions = (740 () if output_attentions and self.config.add_cross_attention else None741 )742 743 next_decoder_cache = () if use_cache else None744 for i, layer_module in enumerate(self.layer):745 if output_hidden_states:746 all_hidden_states = all_hidden_states + (hidden_states,)747 748 layer_head_mask = head_mask[i] if head_mask is not None else None749 past_key_value = past_key_values[i] if past_key_values is not None else None750 751 if self.gradient_checkpointing and self.training:752 753 def create_custom_forward(module):754 def custom_forward(*inputs):755 return module(*inputs, past_key_value, output_attentions)756 757 return custom_forward758 759 layer_outputs = torch.utils.checkpoint.checkpoint(760 create_custom_forward(layer_module),761 hidden_states,762 attention_mask,763 layer_head_mask,764 encoder_hidden_states,765 encoder_attention_mask,766 )767 else:768 layer_outputs = layer_module(769 hidden_states,770 attention_mask,771 layer_head_mask,772 encoder_hidden_states,773 encoder_attention_mask,774 past_key_value,775 output_attentions,776 )777 778 hidden_states = layer_outputs[0]779 if use_cache:780 next_decoder_cache += (layer_outputs[-1],)781 if output_attentions:782 all_self_attentions = all_self_attentions + (layer_outputs[1],)783 if self.config.add_cross_attention:784 all_cross_attentions = all_cross_attentions + (layer_outputs[2],)785 786 787 if self.emb_layer_norm_after:788 hidden_states = self.emb_layer_norm_after(hidden_states)789 790 if output_hidden_states:791 all_hidden_states = all_hidden_states + (hidden_states,)792 793 if not return_dict:794 return tuple(795 v796 for v in [797 hidden_states,798 next_decoder_cache,799 all_hidden_states,800 all_self_attentions,801 all_cross_attentions,802 ]803 if v is not None804 )805 return BaseModelOutputWithPastAndCrossAttentions(806 last_hidden_state=hidden_states,807 past_key_values=next_decoder_cache,808 hidden_states=all_hidden_states,809 attentions=all_self_attentions,810 cross_attentions=all_cross_attentions,811 )812 813 814# Copied from transformers.models.bert.modeling_bert.BertPooler815class EsmPooler(nn.Module):816 def __init__(self, config):817 super().__init__()818 self.dense = nn.Linear(config.hidden_size, config.hidden_size)819 self.activation = nn.Tanh()820 821 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:822 # We "pool" the model by simply taking the hidden state corresponding823 # to the first token.824 first_token_tensor = hidden_states[:, 0]825 pooled_output = self.dense(first_token_tensor)826 pooled_output = self.activation(pooled_output)827 return pooled_output828 829 830class EsmPreTrainedModel(PreTrainedModel):831 """832 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained833 models.834 """835 836 config_class = SegmentNTConfig837 base_model_prefix = "esm"838 _no_split_modules = ["EsmLayer", "EsmFoldTriangularSelfAttentionBlock"]839 840 # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights841 def _init_weights(self, module):842 """Initialize the weights"""843 if isinstance(module, nn.Linear):844 # Slightly different from the TF version which uses truncated_normal for initialization845 # cf https://github.com/pytorch/pytorch/pull/5617846 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)847 if module.bias is not None:848 module.bias.data.zero_()849 elif isinstance(module, nn.Embedding):850 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)851 if module.padding_idx is not None:852 module.weight.data[module.padding_idx].zero_()853 elif isinstance(module, nn.LayerNorm):854 module.bias.data.zero_()855 module.weight.data.fill_(1.0)856 857 858ESM_START_DOCSTRING = r"""859 860 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the861 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads862 etc.)863 864 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.865 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage866 and behavior.867 868 Parameters:869 config ([`EsmConfig`]): Model configuration class with all the parameters of the870 model. Initializing with a config file does not load the weights associated with the model, only the871 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.872"""873 874ESM_INPUTS_DOCSTRING = r"""875 Args:876 input_ids (`torch.LongTensor` of shape `({0})`):877 Indices of input sequence tokens in the vocabulary.878 879 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and880 [`PreTrainedTokenizer.__call__`] for details.881 882 [What are input IDs?](../glossary#input-ids)883 attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):884 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:885 886 - 1 for tokens that are **not masked**,887 - 0 for tokens that are **masked**.888 889 [What are attention masks?](../glossary#attention-mask)890 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):891 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,892 config.max_position_embeddings - 1]`.893 894 [What are position IDs?](../glossary#position-ids)895 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):896 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:897 898 - 1 indicates the head is **not masked**,899 - 0 indicates the head is **masked**.900 901 inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):902 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This903 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the904 model's internal embedding lookup matrix.905 output_attentions (`bool`, *optional*):906 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned907 tensors for more detail.908 output_hidden_states (`bool`, *optional*):909 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for910 more detail.911 return_dict (`bool`, *optional*):912 Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.913"""914 915 916@add_start_docstrings(917 "The bare ESM Model transformer outputting raw hidden-states without any specific head on top.",918 ESM_START_DOCSTRING,919)920class EsmModel(EsmPreTrainedModel):921 """922 923 The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of924 cross-attention is added between the self-attention layers, following the architecture described in [Attention is925 all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,926 Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.927 928 To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set929 to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and930 `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.931 """932 933 supports_gradient_checkpointing = False934 935 def __init__(self, config, add_pooling_layer=True):936 super().__init__(config)937 self.config = config938 939 self.embeddings = EsmEmbeddings(config)940 self.encoder = EsmEncoder(config)941 942 self.pooler = EsmPooler(config) if add_pooling_layer else None943 944 self.contact_head = EsmContactPredictionHead(945 in_features=config.num_hidden_layers * config.num_attention_heads, bias=True946 )947 948 # Initialize weights and apply final processing949 self.post_init()950 951 def _set_gradient_checkpointing(self, module, value=False):952 if isinstance(module, EsmEncoder):953 module.gradient_checkpointing = value954 955 def get_input_embeddings(self):956 return self.embeddings.word_embeddings957 958 def set_input_embeddings(self, value):959 self.embeddings.word_embeddings = value960 961 def _prune_heads(self, heads_to_prune):962 """963 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base964 class PreTrainedModel965 """966 for layer, heads in heads_to_prune.items():967 self.encoder.layer[layer].attention.prune_heads(heads)968 969 @add_start_docstrings_to_model_forward(970 ESM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)")971 )972 @add_code_sample_docstrings(973 checkpoint=_CHECKPOINT_FOR_DOC,974 output_type=BaseModelOutputWithPoolingAndCrossAttentions,975 config_class=_CONFIG_FOR_DOC,976 )977 def forward(978 self,979 input_ids: Optional[torch.Tensor] = None,980 attention_mask: Optional[torch.Tensor] = None,981 position_ids: Optional[torch.Tensor] = None,982 head_mask: Optional[torch.Tensor] = None,983 inputs_embeds: Optional[torch.Tensor] = None,984 encoder_hidden_states: Optional[torch.Tensor] = None,985 encoder_attention_mask: Optional[torch.Tensor] = None,986 past_key_values: Optional[List[torch.FloatTensor]] = None,987 use_cache: Optional[bool] = None,988 output_attentions: Optional[bool] = None,989 output_hidden_states: Optional[bool] = None,990 return_dict: Optional[bool] = None,991 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:992 r"""993 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):994 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if995 the model is configured as a decoder.996 encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):997 Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in998 the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:999 1000 - 1 for tokens that are **not masked**,1001 - 0 for tokens that are **masked**.1002 past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):1003 Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.1004 1005 If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that1006 don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all1007 `decoder_input_ids` of shape `(batch_size, sequence_length)`.1008 use_cache (`bool`, *optional*):1009 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see1010 `past_key_values`).1011 """1012 output_attentions = (1013 output_attentions1014 if output_attentions is not None1015 else self.config.output_attentions1016 )1017 output_hidden_states = (1018 output_hidden_states1019 if output_hidden_states is not None1020 else self.config.output_hidden_states1021 )1022 return_dict = (1023 return_dict if return_dict is not None else self.config.use_return_dict1024 )1025 1026 if self.config.is_decoder:1027 use_cache = use_cache if use_cache is not None else self.config.use_cache1028 else:1029 use_cache = False1030 1031 if input_ids is not None and inputs_embeds is not None:1032 raise ValueError(1033 "You cannot specify both input_ids and inputs_embeds at the same time"1034 )1035 elif input_ids is not None:1036 input_shape = input_ids.size()1037 elif inputs_embeds is not None:1038 input_shape = inputs_embeds.size()[:-1]1039 else:1040 raise ValueError("You have to specify either input_ids or inputs_embeds")1041 1042 batch_size, seq_length = input_shape1043 device = input_ids.device if input_ids is not None else inputs_embeds.device1044 1045 # past_key_values_length1046 past_key_values_length = (1047 past_key_values[0][0].shape[2] if past_key_values is not None else 01048 )1049 1050 if attention_mask is None:1051 attention_mask = torch.ones(1052 ((batch_size, seq_length + past_key_values_length)), device=device1053 )1054 1055 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]1056 # ourselves in which case we just need to make it broadcastable to all heads.1057 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(1058 attention_mask, input_shape1059 )1060 1061 # If a 2D or 3D attention mask is provided for the cross-attention1062 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]1063 if self.config.is_decoder and encoder_hidden_states is not None:1064 (1065 encoder_batch_size,1066 encoder_sequence_length,1067 _,1068 ) = encoder_hidden_states.size()1069 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)1070 if encoder_attention_mask is None:1071 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)1072 encoder_extended_attention_mask = self.invert_attention_mask(1073 encoder_attention_mask1074 )1075 else:1076 encoder_extended_attention_mask = None1077 1078 # Prepare head mask if needed1079 # 1.0 in head_mask indicate we keep the head1080 # attention_probs has shape bsz x n_heads x N x N1081 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]1082 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]1083 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)1084 1085 embedding_output = self.embeddings(1086 input_ids=input_ids,1087 position_ids=position_ids,1088 attention_mask=attention_mask,1089 inputs_embeds=inputs_embeds,1090 past_key_values_length=past_key_values_length,1091 )1092 encoder_outputs = self.encoder(1093 embedding_output,1094 attention_mask=extended_attention_mask,1095 head_mask=head_mask,1096 encoder_hidden_states=encoder_hidden_states,1097 encoder_attention_mask=encoder_extended_attention_mask,1098 past_key_values=past_key_values,1099 use_cache=use_cache,1100 output_attentions=output_attentions,1101 output_hidden_states=output_hidden_states,1102 return_dict=return_dict,1103 )1104 sequence_output = encoder_outputs[0]1105 pooled_output = (1106 self.pooler(sequence_output) if self.pooler is not None else None1107 )1108 1109 if not return_dict:1110 return (sequence_output, pooled_output) + encoder_outputs[1:]1111 1112 return BaseModelOutputWithPoolingAndCrossAttentions(1113 last_hidden_state=sequence_output,1114 pooler_output=pooled_output,1115 past_key_values=encoder_outputs.past_key_values,1116 hidden_states=encoder_outputs.hidden_states,1117 attentions=encoder_outputs.attentions,1118 cross_attentions=encoder_outputs.cross_attentions,1119 )1120 1121 def predict_contacts(self, tokens, attention_mask):1122 attns = self(1123 tokens,1124 attention_mask=attention_mask,1125 return_dict=True,1126 output_attentions=True,1127 ).attentions1128 attns = torch.stack(attns, dim=1) # Matches the original model layout1129 # In the original model, attentions for padding tokens are completely zeroed out.1130 # This makes no difference most of the time because the other tokens won't attend to them,1131 # but it does for the contact prediction task, which takes attentions as input,1132 # so we have to mimic that here.1133 attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3)1134 attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4)1135 return self.contact_head(tokens, attns)1136 1137def create_position_ids_from_input_ids(1138 input_ids, padding_idx, past_key_values_length=01139):1140 """1141 Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols1142 are ignored. This is modified from fairseq's `utils.make_positions`.1143 1144 Args:1145 x: torch.Tensor x:1146 1147 Returns: torch.Tensor1148 """1149 # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.1150 mask = input_ids.ne(padding_idx).int()1151 incremental_indices = (1152 torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length1153 ) * mask1154 return incremental_indices.long() + padding_idx1155 1156 1157 1158 1159class SegmentNT(EsmPreTrainedModel):1160 def __init__(self, config):1161 super().__init__(config)1162 self.num_labels = config.num_labels1163 self.config = config1164 self.num_features = len(config.features)1165 1166 self.esm = EsmModel(config, add_pooling_layer=False)1167 1168 embed_dim = config.hidden_size1169 num_layers = config.num_layers_head1170 self.unet = UNET1DSegmentationHead(1171 embed_dim=embed_dim,1172 num_classes=embed_dim // 2,1173 output_channels_list=tuple(1174 embed_dim * (2**i) for i in range(num_layers)1175 ),1176 )1177 self.fc = nn.Linear(in_features=embed_dim, out_features=6 * 2 * self.num_features)1178 self.activation_fn = nn.SiLU()1179 1180 self.init_weights()1181 1182 # @add_start_docstrings_to_model_forward(1183 # ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length")1184 # )1185 # @add_code_sample_docstrings(1186 # checkpoint=_CHECKPOINT_FOR_DOC,1187 # output_type=SequenceClassifierOutput,1188 # config_class=_CONFIG_FOR_DOC,1189 # )1190 def forward(1191 self,1192 input_ids: Optional[torch.LongTensor] = None,1193 attention_mask: Optional[torch.Tensor] = None,1194 position_ids: Optional[torch.LongTensor] = None,1195 head_mask: Optional[torch.Tensor] = None,1196 inputs_embeds: Optional[torch.FloatTensor] = None,1197 labels: Optional[torch.LongTensor] = None,1198 output_attentions: Optional[bool] = None,1199 output_hidden_states: Optional[bool] = None,1200 return_dict: Optional[bool] = None,