diff --git "a/modeling_omnigenome.py" "b/modeling_omnigenome.py" --- "a/modeling_omnigenome.py" +++ "b/modeling_omnigenome.py" @@ -1,1906 +1,1906 @@ -# coding=utf-8 -# Copyright 2022 ColaLab-UoE (https://colalab.ai/), Meta and The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" PyTorch OmniGenome model.""" - -import math -import random -import warnings -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch -import torch.utils.checkpoint -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from transformers import add_start_docstrings, PreTrainedModel - -from transformers.modeling_outputs import ( - BaseModelOutputWithPastAndCrossAttentions, - BaseModelOutputWithPoolingAndCrossAttentions, - MaskedLMOutput, - SequenceClassifierOutput, - TokenClassifierOutput, -) - -from transformers.pytorch_utils import ( - find_pruneable_heads_and_indices, - prune_linear_layer, -) - -from transformers.utils import ( - logging, - add_code_sample_docstrings, - add_start_docstrings_to_model_forward, -) - -from .configuration_omnigenome import OmniGenomeConfig - -try: - from flash_attn import flash_attn_func -except ImportError: - flash_attn_func = None - -logger = logging.get_logger(__name__) - -_CHECKPOINT_FOR_DOC = "yangheng/OmniGenome-52M" -_CONFIG_FOR_DOC = "OmniGenomeConfig" - -OmniGenome_PRETRAINED_MODEL_ARCHIVE_LIST = [ - "yangheng/OmniGenome-52M", - # This is not a complete list of all OmniGenome models! - # See all OmniGenome models at https://huggingface.co/models?filter=OmniGenome -] - - -def rotate_half(x): - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb(x, cos, sin): - cos = cos[:, :, : x.shape[-2], :] - sin = sin[:, :, : x.shape[-2], :] - - return (x * cos) + (rotate_half(x) * sin) - - -def gelu(x): - """ - This is the gelu implementation from the original OmniGenome repo. Using F.gelu yields subtly wrong results. - """ - return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) - - -def symmetrize(x): - "Make layer symmetric in final two dimensions, used for contact prediction." - return x + x.transpose(-1, -2) - - -def average_product_correct(x): - "Perform average product correct, used for contact prediction." - a1 = x.sum(-1, keepdims=True) - a2 = x.sum(-2, keepdims=True) - a12 = x.sum((-1, -2), keepdims=True) - - avg = a1 * a2 - avg.div_(a12) # in-place to reduce memory - normalized = x - avg - return normalized - - -# Copied from transformers.models.esm.modeling_esm.RotaryEmbedding -class RotaryEmbedding(torch.nn.Module): - """ - Rotary position embeddings based on those in - [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation - matrices which depend on their relative positions. - """ - - def __init__(self, dim: int): - super().__init__() - # Generate and save the inverse frequency buffer (non trainable) - inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) - inv_freq = inv_freq - self.register_buffer("inv_freq", inv_freq) - - self._seq_len_cached = None - self._cos_cached = None - self._sin_cached = None - - def _update_cos_sin_tables(self, x, seq_dimension=2): - seq_len = x.shape[seq_dimension] - - # Reset the tables if the sequence length has changed, - # or if we're on a new device (possibly due to tracing for instance) - if seq_len != self._seq_len_cached or self._cos_cached.device != x.device: - self._seq_len_cached = seq_len - t = torch.arange(x.shape[seq_dimension], device=x.device).type_as( - self.inv_freq - ) - freqs = torch.outer(t, self.inv_freq) - emb = torch.cat((freqs, freqs), dim=-1).to(x.device) - - self._cos_cached = emb.cos()[None, None, :, :] - self._sin_cached = emb.sin()[None, None, :, :] - - return self._cos_cached, self._sin_cached - - def forward( - self, q: torch.Tensor, k: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: - self._cos_cached, self._sin_cached = self._update_cos_sin_tables( - k, seq_dimension=-2 - ) - - return ( - apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), - apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), - ) - - -# Copied from transformers.models.esm.modeling_esm.EsmContactPredictionHead with Esm->OmniGenome -class OmniGenomeContactPredictionHead(nn.Module): - """Performs symmetrization, apc, and computes a logistic regression on the output features""" - - def __init__( - self, - in_features: int, - bias=True, - eos_idx: int = 2, - ): - super().__init__() - self.in_features = in_features - self.eos_idx = eos_idx - self.regression = nn.Linear(in_features, 1, bias) - self.activation = nn.Sigmoid() - - def forward(self, tokens, attentions): - # remove eos token attentions - eos_mask = tokens.ne(self.eos_idx).to(attentions) - eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2) - attentions = attentions * eos_mask[:, None, None, :, :] - attentions = attentions[..., :-1, :-1] - # remove cls token attentions - attentions = attentions[..., 1:, 1:] - batch_size, layers, heads, seqlen, _ = attentions.size() - attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen) - - # features: batch x channels x tokens x tokens (symmetric) - attentions = attentions.to( - self.regression.weight.device - ) # attentions always float32, may need to convert to float16 - attentions = average_product_correct(symmetrize(attentions)) - attentions = attentions.permute(0, 2, 3, 1) - return self.activation(self.regression(attentions).squeeze(3)) - - -# Copied from transformers.models.esm.modeling_esm.EsmEmbeddings with Esm->OmniGenome -class OmniGenomeEmbeddings(nn.Module): - """ - Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. - """ - - def __init__(self, config): - super().__init__() - self.word_embeddings = nn.Embedding( - config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id - ) - - if config.emb_layer_norm_before: - self.layer_norm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - else: - self.layer_norm = None - self.dropout = nn.Dropout(config.hidden_dropout_prob) - # position_ids (1, len position emb) is contiguous in memory and exported when serialized - self.position_embedding_type = getattr( - config, "position_embedding_type", "absolute" - ) - self.register_buffer( - "position_ids", - torch.arange(config.max_position_embeddings).expand((1, -1)), - persistent=False, - ) - - self.padding_idx = config.pad_token_id - self.position_embeddings = nn.Embedding( - config.max_position_embeddings, - config.hidden_size, - padding_idx=self.padding_idx, - ) - self.token_dropout = config.token_dropout - self.mask_token_id = config.mask_token_id - - def forward( - self, - input_ids=None, - attention_mask=None, - position_ids=None, - inputs_embeds=None, - past_key_values_length=0, - ): - if position_ids is None: - if input_ids is not None: - # Create the position ids from the input token ids. Any padded tokens remain padded. - position_ids = create_position_ids_from_input_ids( - input_ids, self.padding_idx, past_key_values_length - ) - else: - position_ids = self.create_position_ids_from_inputs_embeds( - inputs_embeds - ) - - if inputs_embeds is None: - inputs_embeds = self.word_embeddings(input_ids) - - # Note that if we want to support OmniGenome-1 (not 1b!) in future then we need to support an - # embedding_scale factor here. - embeddings = inputs_embeds - - # Matt: OmniGenome has the option to handle masking in MLM in a slightly unusual way. If the token_dropout - # flag is False then it is handled in the same was as BERT/RoBERTa. If it is set to True, however, - # masked tokens are treated as if they were selected for input dropout and zeroed out. - # This "mask-dropout" is compensated for when masked tokens are not present, by scaling embeddings by - # a factor of (fraction of unmasked tokens during training) / (fraction of unmasked tokens in sample). - # This is analogous to the way that dropout layers scale down outputs during evaluation when not - # actually dropping out values (or, equivalently, scale up their un-dropped outputs in training). - if self.token_dropout: - embeddings = embeddings.masked_fill( - (input_ids == self.mask_token_id).unsqueeze(-1), 0.0 - ) - mask_ratio_train = ( - 0.15 * 0.8 - ) # Hardcoded as the ratio used in all OmniGenome model training runs - src_lengths = attention_mask.sum(-1) - mask_ratio_observed = (input_ids == self.mask_token_id).sum( - -1 - ).float() / src_lengths - embeddings = ( - embeddings - * (1 - mask_ratio_train) - / (1 - mask_ratio_observed)[:, None, None] - ).to(embeddings.dtype) - - if self.position_embedding_type == "absolute": - position_embeddings = self.position_embeddings(position_ids) - embeddings = embeddings + position_embeddings - - if self.layer_norm is not None: - embeddings = self.layer_norm(embeddings) - if attention_mask is not None: - embeddings = (embeddings * attention_mask.unsqueeze(-1)).to( - embeddings.dtype - ) - # Matt: I think this line was copied incorrectly from BERT, disabling it for now. - # embeddings = self.dropout(embeddings) - return embeddings - - def create_position_ids_from_inputs_embeds(self, inputs_embeds): - """ - We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids. - - Args: - inputs_embeds: torch.Tensor - - Returns: torch.Tensor - """ - input_shape = inputs_embeds.size()[:-1] - sequence_length = input_shape[1] - - position_ids = torch.arange( - self.padding_idx + 1, - sequence_length + self.padding_idx + 1, - dtype=torch.long, - device=inputs_embeds.device, - ) - return position_ids.unsqueeze(0).expand(input_shape) - -# -# # Copied from transformers.models.esm.modeling_esm.EsmSelfAttention with Esm->OmniGenome -# class OmniGenomeSelfAttention(nn.Module): -# def __init__(self, config, position_embedding_type=None): -# super().__init__() -# if config.hidden_size % config.num_attention_heads != 0 and not hasattr( -# config, "embedding_size" -# ): -# raise ValueError( -# f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " -# f"heads ({config.num_attention_heads})" -# ) -# -# self.num_attention_heads = config.num_attention_heads -# self.attention_head_size = int(config.hidden_size / config.num_attention_heads) -# self.all_head_size = self.num_attention_heads * self.attention_head_size -# -# self.query = nn.Linear(config.hidden_size, self.all_head_size) -# self.key = nn.Linear(config.hidden_size, self.all_head_size) -# self.value = nn.Linear(config.hidden_size, self.all_head_size) -# -# self.dropout = nn.Dropout(config.attention_probs_dropout_prob) -# self.position_embedding_type = position_embedding_type or getattr( -# config, "position_embedding_type", "absolute" -# ) -# self.rotary_embeddings = None -# if ( -# self.position_embedding_type == "relative_key" -# or self.position_embedding_type == "relative_key_query" -# ): -# self.max_position_embeddings = config.max_position_embeddings -# self.distance_embedding = nn.Embedding( -# 2 * config.max_position_embeddings - 1, self.attention_head_size -# ) -# elif self.position_embedding_type == "rotary": -# self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) -# -# self.is_decoder = config.is_decoder -# -# def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: -# new_x_shape = x.size()[:-1] + ( -# self.num_attention_heads, -# self.attention_head_size, -# ) -# x = x.view(new_x_shape) -# return x.permute(0, 2, 1, 3) -# -# def forward( -# self, -# hidden_states: torch.Tensor, -# attention_mask: Optional[torch.FloatTensor] = None, -# head_mask: Optional[torch.FloatTensor] = None, -# encoder_hidden_states: Optional[torch.FloatTensor] = None, -# encoder_attention_mask: Optional[torch.FloatTensor] = None, -# past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, -# output_attentions: Optional[bool] = False, -# ) -> Tuple[torch.Tensor]: -# mixed_query_layer = self.query(hidden_states) -# -# # If this is instantiated as a cross-attention module, the keys -# # and values come from an encoder; the attention mask needs to be -# # such that the encoder's padding tokens are not attended to. -# is_cross_attention = encoder_hidden_states is not None -# -# if is_cross_attention and past_key_value is not None: -# # reuse k,v, cross_attentions -# key_layer = past_key_value[0] -# value_layer = past_key_value[1] -# attention_mask = encoder_attention_mask -# elif is_cross_attention: -# key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) -# value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) -# attention_mask = encoder_attention_mask -# elif past_key_value is not None: -# key_layer = self.transpose_for_scores(self.key(hidden_states)) -# value_layer = self.transpose_for_scores(self.value(hidden_states)) -# key_layer = torch.cat([past_key_value[0], key_layer], dim=2) -# value_layer = torch.cat([past_key_value[1], value_layer], dim=2) -# else: -# key_layer = self.transpose_for_scores(self.key(hidden_states)) -# value_layer = self.transpose_for_scores(self.value(hidden_states)) -# -# query_layer = self.transpose_for_scores(mixed_query_layer) -# -# # Matt: Our BERT model (which this code was derived from) scales attention logits down by sqrt(head_dim). -# # OmniGenome scales the query down by the same factor instead. Modulo numerical stability these are equivalent, -# # but not when rotary embeddings get involved. Therefore, we scale the query here to match the original -# # OmniGenome code and fix rotary embeddings. -# query_layer = query_layer * self.attention_head_size ** -0.5 -# -# if self.is_decoder: -# # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. -# # Further calls to cross_attention layer can then reuse all cross-attention -# # key/value_states (first "if" case) -# # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of -# # all previous decoder key/value_states. Further calls to uni-directional self-attention -# # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) -# # if encoder bi-directional self-attention `past_key_value` is always `None` -# past_key_value = (key_layer, value_layer) -# -# if self.position_embedding_type == "rotary": -# query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) -# -# # Take the dot product between "query" and "key" to get the raw attention scores. -# attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) -# -# if ( -# self.position_embedding_type == "relative_key" -# or self.position_embedding_type == "relative_key_query" -# ): -# seq_length = hidden_states.size()[1] -# position_ids_l = torch.arange( -# seq_length, dtype=torch.long, device=hidden_states.device -# ).view(-1, 1) -# position_ids_r = torch.arange( -# seq_length, dtype=torch.long, device=hidden_states.device -# ).view(1, -1) -# distance = position_ids_l - position_ids_r -# positional_embedding = self.distance_embedding( -# distance + self.max_position_embeddings - 1 -# ) -# positional_embedding = positional_embedding.to( -# dtype=query_layer.dtype -# ) # fp16 compatibility -# -# if self.position_embedding_type == "relative_key": -# relative_position_scores = torch.einsum( -# "bhld,lrd->bhlr", query_layer, positional_embedding -# ) -# attention_scores = attention_scores + relative_position_scores -# elif self.position_embedding_type == "relative_key_query": -# relative_position_scores_query = torch.einsum( -# "bhld,lrd->bhlr", query_layer, positional_embedding -# ) -# relative_position_scores_key = torch.einsum( -# "bhrd,lrd->bhlr", key_layer, positional_embedding -# ) -# attention_scores = ( -# attention_scores -# + relative_position_scores_query -# + relative_position_scores_key -# ) -# -# if attention_mask is not None: -# # Apply the attention mask is (precomputed for all layers in OmniGenomeModel forward() function) -# attention_scores = attention_scores + attention_mask -# -# # Normalize the attention scores to probabilities. -# attention_probs = nn.functional.softmax(attention_scores, dim=-1) -# -# # This is actually dropping out entire tokens to attend to, which might -# # seem a bit unusual, but is taken from the original Transformer paper. -# attention_probs = self.dropout(attention_probs) -# -# # Mask heads if we want to -# if head_mask is not None: -# attention_probs = attention_probs * head_mask -# -# context_layer = torch.matmul(attention_probs, value_layer) -# -# context_layer = context_layer.permute(0, 2, 1, 3).contiguous() -# new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) -# context_layer = context_layer.view(new_context_layer_shape) -# -# outputs = ( -# (context_layer, attention_probs) if output_attentions else (context_layer,) -# ) -# -# if self.is_decoder: -# outputs = outputs + (past_key_value,) -# return outputs - - -# Copied from transformers.models.esm.modeling_esm.EsmSelfAttention with Esm->OmniGenome -class OmniGenomeSelfAttention(nn.Module): - def __init__(self, config, position_embedding_type=None): - super().__init__() - if config.hidden_size % config.num_attention_heads != 0 and not hasattr( - config, "embedding_size" - ): - raise ValueError( - f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " - f"heads ({config.num_attention_heads})" - ) - - self.num_attention_heads = config.num_attention_heads - self.attention_head_size = int(config.hidden_size / config.num_attention_heads) - self.all_head_size = self.num_attention_heads * self.attention_head_size - - self.query = nn.Linear(config.hidden_size, self.all_head_size) - self.key = nn.Linear(config.hidden_size, self.all_head_size) - self.value = nn.Linear(config.hidden_size, self.all_head_size) - - self.dropout = nn.Dropout(config.attention_probs_dropout_prob) - self.position_embedding_type = position_embedding_type or getattr( - config, "position_embedding_type", "absolute" - ) - self.rotary_embeddings = None - if ( - self.position_embedding_type == "relative_key" - or self.position_embedding_type == "relative_key_query" - ): - self.max_position_embeddings = config.max_position_embeddings - self.distance_embedding = nn.Embedding( - 2 * config.max_position_embeddings - 1, self.attention_head_size - ) - elif self.position_embedding_type == "rotary": - self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) - - self.is_decoder = config.is_decoder - - # FlashAttention parameters - self.enable_flash_attn = getattr(config, "use_flash_attention", True) - if self.enable_flash_attn: - self.flash_attn_func = flash_attn_func - else: - self.flash_attn_func = None - - def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: - new_x_shape = x.size()[:-1] + ( - self.num_attention_heads, - self.attention_head_size, - ) - x = x.view(new_x_shape) - return x.permute(0, 2, 1, 3) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor]: - mixed_query_layer = self.query(hidden_states) - - is_cross_attention = encoder_hidden_states is not None - - if is_cross_attention and past_key_value is not None: - key_layer = past_key_value[0] - value_layer = past_key_value[1] - attention_mask = encoder_attention_mask - elif is_cross_attention: - key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) - value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) - attention_mask = encoder_attention_mask - elif past_key_value is not None: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - key_layer = torch.cat([past_key_value[0], key_layer], dim=2) - value_layer = torch.cat([past_key_value[1], value_layer], dim=2) - else: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - - query_layer = self.transpose_for_scores(mixed_query_layer) - - if self.is_decoder: - past_key_value = (key_layer, value_layer) - - # 使用FlashAttention的条件判断 - use_flash_attn = self.enable_flash_attn and self.position_embedding_type == "rotary" - if use_flash_attn and self.flash_attn_func is not None: - # 应用旋转位置编码 - query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) - - # 调整维度顺序为 [batch_size, seq_len, num_heads, head_dim] - q = query_layer.transpose(1, 2).half() - k = key_layer.transpose(1, 2).half() - v = value_layer.transpose(1, 2).half() - - # 使用FlashAttention计算 - context_layer = self.flash_attn_func( - q, k, v, - dropout_p=self.dropout.p if self.training else 0.0, - softmax_scale=self.attention_head_size ** -0.5, - causal=self.is_decoder - ) - - # 恢复维度顺序 [batch_size, num_heads, seq_len, head_dim] - context_layer = context_layer.transpose(1, 2).to(hidden_states.dtype) - else: - # 原始实现 - query_layer = query_layer * self.attention_head_size ** -0.5 - - if self.position_embedding_type == "rotary": - query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) - - attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) - - if self.position_embedding_type in ["relative_key", "relative_key_query"]: - seq_length = hidden_states.size()[1] - position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1) - position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1) - distance = position_ids_l - position_ids_r - positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1) - positional_embedding = positional_embedding.to(dtype=query_layer.dtype) - - if self.position_embedding_type == "relative_key": - relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) - attention_scores = attention_scores + relative_position_scores - elif self.position_embedding_type == "relative_key_query": - relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) - relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding) - attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key - - if attention_mask is not None: - attention_scores = attention_scores + attention_mask - - attention_probs = nn.functional.softmax(attention_scores, dim=-1) - attention_probs = self.dropout(attention_probs) - - if head_mask is not None: - attention_probs = attention_probs * head_mask - - context_layer = torch.matmul(attention_probs, value_layer) - - context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) - context_layer = context_layer.view(new_context_layer_shape) - - outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) - if self.is_decoder: - outputs = outputs + (past_key_value,) - return outputs - -# Copied from transformers.models.esm.modeling_esm.EsmSelfOutput with Esm->OmniGenome -class OmniGenomeSelfOutput(nn.Module): - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.hidden_size, config.hidden_size) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - - def forward(self, hidden_states, input_tensor): - hidden_states = self.dense(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = hidden_states + input_tensor - return hidden_states - - -# Copied from transformers.models.esm.modeling_esm.EsmAttention with Esm->OmniGenome -class OmniGenomeAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.self = OmniGenomeSelfAttention(config) - self.output = OmniGenomeSelfOutput(config) - self.pruned_heads = set() - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def prune_heads(self, heads): - if len(heads) == 0: - return - heads, index = find_pruneable_heads_and_indices( - heads, - self.self.num_attention_heads, - self.self.attention_head_size, - self.pruned_heads, - ) - - # Prune linear layers - self.self.query = prune_linear_layer(self.self.query, index) - self.self.key = prune_linear_layer(self.self.key, index) - self.self.value = prune_linear_layer(self.self.value, index) - self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) - - # Update hyper params and store pruned heads - self.self.num_attention_heads = self.self.num_attention_heads - len(heads) - self.self.all_head_size = ( - self.self.attention_head_size * self.self.num_attention_heads - ) - self.pruned_heads = self.pruned_heads.union(heads) - - def forward( - self, - hidden_states, - attention_mask=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_value=None, - output_attentions=False, - ): - hidden_states_ln = self.LayerNorm(hidden_states) - hidden_states_ln = hidden_states_ln.to(hidden_states.dtype) - self_outputs = self.self( - hidden_states_ln, - attention_mask, - head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - ) - attention_output = self.output(self_outputs[0], hidden_states) - outputs = (attention_output,) + self_outputs[ - 1: - ] # add attentions if we output them - return outputs - - -# Copied from transformers.models.esm.modeling_esm.EsmIntermediate with Esm->OmniGenome -class OmniGenomeIntermediate(nn.Module): - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.hidden_size, config.intermediate_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.dense(hidden_states) - hidden_states = gelu(hidden_states) - return hidden_states - - -# Copied from transformers.models.esm.modeling_esm.EsmOutput with Esm->OmniGenome -class OmniGenomeOutput(nn.Module): - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.intermediate_size, config.hidden_size) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - - def forward(self, hidden_states, input_tensor): - hidden_states = self.dense(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = hidden_states + input_tensor - return hidden_states - - -# Copied from transformers.models.esm.modeling_esm.EsmLayer with Esm->OmniGenome -class OmniGenomeLayer(nn.Module): - def __init__(self, config): - super().__init__() - self.chunk_size_feed_forward = config.chunk_size_feed_forward - self.seq_len_dim = 1 - self.attention = OmniGenomeAttention(config) - self.is_decoder = config.is_decoder - self.add_cross_attention = config.add_cross_attention - if self.add_cross_attention: - if not self.is_decoder: - raise RuntimeError( - f"{self} should be used as a decoder model if cross attention is added" - ) - self.crossattention = OmniGenomeAttention(config) - self.intermediate = OmniGenomeIntermediate(config) - self.output = OmniGenomeOutput(config) - self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states, - attention_mask=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_value=None, - output_attentions=False, - ): - # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 - self_attn_past_key_value = ( - past_key_value[:2] if past_key_value is not None else None - ) - self_attention_outputs = self.attention( - hidden_states, - attention_mask, - head_mask, - output_attentions=output_attentions, - past_key_value=self_attn_past_key_value, - ) - attention_output = self_attention_outputs[0] - - # if decoder, the last output is tuple of self-attn cache - if self.is_decoder: - outputs = self_attention_outputs[1:-1] - present_key_value = self_attention_outputs[-1] - else: - outputs = self_attention_outputs[ - 1: - ] # add self attentions if we output attention weights - - cross_attn_present_key_value = None - if self.is_decoder and encoder_hidden_states is not None: - if not hasattr(self, "crossattention"): - raise AttributeError( - f"If `encoder_hidden_states` are passed, {self} has to be instantiated" - " with cross-attention layers by setting `config.add_cross_attention=True`" - ) - - # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple - cross_attn_past_key_value = ( - past_key_value[-2:] if past_key_value is not None else None - ) - cross_attention_outputs = self.crossattention( - attention_output, - attention_mask, - head_mask, - encoder_hidden_states, - encoder_attention_mask, - cross_attn_past_key_value, - output_attentions, - ) - attention_output = cross_attention_outputs[0] - outputs = ( - outputs + cross_attention_outputs[1:-1] - ) # add cross attentions if we output attention weights - - # add cross-attn cache to positions 3,4 of present_key_value tuple - cross_attn_present_key_value = cross_attention_outputs[-1] - present_key_value = present_key_value + cross_attn_present_key_value - - layer_output = self.feed_forward_chunk(attention_output) - - outputs = (layer_output,) + outputs - - # if decoder, return the attn key/values as the last output - if self.is_decoder: - outputs = outputs + (present_key_value,) - return outputs - - def feed_forward_chunk(self, attention_output): - attention_output_ln = self.LayerNorm(attention_output) - intermediate_output = self.intermediate(attention_output_ln) - layer_output = self.output(intermediate_output, attention_output) - return layer_output - - -# Copied from transformers.models.esm.modeling_esm.EsmEncoder with Esm->OmniGenome -class OmniGenomeEncoder(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.layer = nn.ModuleList( - [OmniGenomeLayer(config) for _ in range(config.num_hidden_layers)] - ) - self.emb_layer_norm_after = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.gradient_checkpointing = False - - def forward( - self, - hidden_states, - attention_mask=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=False, - output_hidden_states=False, - return_dict=True, - ): - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting " - "`use_cache=False`..." - ) - use_cache = False - all_hidden_states = () if output_hidden_states else None - all_self_attentions = () if output_attentions else None - all_cross_attentions = ( - () if output_attentions and self.config.add_cross_attention else None - ) - - next_decoder_cache = () if use_cache else None - for i, layer_module in enumerate(self.layer): - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_head_mask = head_mask[i] if head_mask is not None else None - past_key_value = past_key_values[i] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - ) - else: - layer_outputs = layer_module( - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - ) - - hidden_states = layer_outputs[0] - if use_cache: - next_decoder_cache = next_decoder_cache + (layer_outputs[-1],) - if output_attentions: - all_self_attentions = all_self_attentions + (layer_outputs[1],) - if self.config.add_cross_attention: - all_cross_attentions = all_cross_attentions + (layer_outputs[2],) - - if self.emb_layer_norm_after: - hidden_states = self.emb_layer_norm_after(hidden_states) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_decoder_cache, - all_hidden_states, - all_self_attentions, - all_cross_attentions, - ] - if v is not None - ) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=next_decoder_cache, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - cross_attentions=all_cross_attentions, - ) - - -# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->OmniGenome -class OmniGenomePooler(nn.Module): - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.hidden_size, config.hidden_size) - self.activation = nn.Tanh() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # We "pool" the model by simply taking the hidden state corresponding - # to the first token. - first_token_tensor = hidden_states[:, 0] - pooled_output = self.dense(first_token_tensor) - pooled_output = self.activation(pooled_output) - return pooled_output - - -# Copied from transformers.models.esm.modeling_esm.EsmPreTrainedModel with Esm->OmniGenome -class OmniGenomePreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - - config_class = OmniGenomeConfig - base_model_prefix = "OmniGenome" - supports_gradient_checkpointing = True - _no_split_modules = [ - "OmniGenomeLayer", - "OmniGenomeFoldTriangularSelfAttentionBlock", - "OmniGenomeEmbeddings", - ] - - # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights - def _init_weights(self, module): - """Initialize the weights""" - if isinstance(module, nn.Linear): - # Slightly different from the TF version which uses truncated_normal for initialization - # cf https://github.com/pytorch/pytorch/pull/5617 - module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - elif isinstance(module, nn.LayerNorm): - module.bias.data.zero_() - module.weight.data.fill_(1.0) - - -OmniGenome_START_DOCSTRING = r""" - - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`OmniGenomeConfig`]): Model configuration class with all the parameters of the - model. Initializing with a config file does not load the weights associated with the model, only the - configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - -OmniGenome_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `({0})`): - Indices of input sequence tokens in the vocabulary. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are input IDs?](../glossary#input-ids) - attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - [What are attention masks?](../glossary#attention-mask) - position_ids (`torch.LongTensor` of shape `({0})`, *optional*): - Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, - config.max_position_embeddings - 1]`. - - [What are position IDs?](../glossary#position-ids) - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*): - Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This - is useful if you want more control over how to convert `input_ids` indices into associated vectors than the - model's internal embedding lookup matrix. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned - tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for - more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. -""" - - -@add_start_docstrings( - "The bare OmniGenome Model transformer outputting raw hidden-states without any specific head on top.", - OmniGenome_START_DOCSTRING, -) -# Copied from transformers.models.esm.modeling_esm.EsmModel with Esm->OmniGenome -class OmniGenomeModel(OmniGenomePreTrainedModel): - """ - - The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of - cross-attention is added between the self-attention layers, following the architecture described in [Attention is - all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, - Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. - - To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set - to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and - `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. - """ - - def __init__(self, config, add_pooling_layer=True): - super().__init__(config) - self.config = config - - self.embeddings = OmniGenomeEmbeddings(config) - self.encoder = OmniGenomeEncoder(config) - - self.pooler = OmniGenomePooler(config) if add_pooling_layer else None - - self.contact_head = OmniGenomeContactPredictionHead( - in_features=config.num_hidden_layers * config.num_attention_heads, bias=True - ) - - # Initialize weights and apply final processing - self.post_init() - - def get_input_embeddings(self): - return self.embeddings.word_embeddings - - def set_input_embeddings(self, value): - self.embeddings.word_embeddings = value - - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.layer[layer].attention.prune_heads(heads) - - @add_start_docstrings_to_model_forward( - OmniGenome_INPUTS_DOCSTRING.format("(batch_size, sequence_length)") - ) - @add_code_sample_docstrings( - checkpoint=_CHECKPOINT_FOR_DOC, - output_type=BaseModelOutputWithPoolingAndCrossAttentions, - config_class=_CONFIG_FOR_DOC, - ) - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: - r""" - encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if - the model is configured as a decoder. - encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in - the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - 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)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see - `past_key_values`). - """ - output_attentions = ( - output_attentions - if output_attentions is not None - else self.config.output_attentions - ) - output_hidden_states = ( - output_hidden_states - if output_hidden_states is not None - else self.config.output_hidden_states - ) - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - if self.config.is_decoder: - use_cache = use_cache if use_cache is not None else self.config.use_cache - else: - use_cache = False - - if input_ids is not None and inputs_embeds is not None: - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time" - ) - elif input_ids is not None: - self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) - input_shape = input_ids.size() - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - batch_size, seq_length = input_shape - device = input_ids.device if input_ids is not None else inputs_embeds.device - - # past_key_values_length - past_key_values_length = ( - past_key_values[0][0].shape[2] if past_key_values is not None else 0 - ) - - if attention_mask is None: - attention_mask = torch.ones( - ((batch_size, seq_length + past_key_values_length)), device=device - ) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( - attention_mask, input_shape - ) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.config.is_decoder and encoder_hidden_states is not None: - ( - encoder_batch_size, - encoder_sequence_length, - _, - ) = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_extended_attention_mask = self.invert_attention_mask( - encoder_attention_mask - ) - else: - encoder_extended_attention_mask = None - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x n_heads x N x N - # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] - # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] - head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - - embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - past_key_values_length=past_key_values_length, - ) - embedding_output = embedding_output.half() - encoder_outputs = self.encoder( - embedding_output, - attention_mask=extended_attention_mask, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - sequence_output = encoder_outputs[0] - pooled_output = ( - self.pooler(sequence_output) if self.pooler is not None else None - ) - - if not return_dict: - return (sequence_output, pooled_output) + encoder_outputs[1:] - - return BaseModelOutputWithPoolingAndCrossAttentions( - last_hidden_state=sequence_output, - pooler_output=pooled_output, - past_key_values=encoder_outputs.past_key_values, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - cross_attentions=encoder_outputs.cross_attentions, - ) - - def predict_contacts(self, tokens, attention_mask): - attns = self( - tokens, - attention_mask=attention_mask, - return_dict=True, - output_attentions=True, - ).attentions - attns = torch.stack(attns, dim=1) # Matches the original model layout - # In the original model, attentions for padding tokens are completely zeroed out. - # This makes no difference most of the time because the other tokens won't attend to them, - # but it does for the contact prediction task, which takes attentions as input, - # so we have to mimic that here. - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) - attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) - return self.contact_head(tokens, attns) - - -@add_start_docstrings( - """OmniGenome Model with a `language modeling` head on top.""", OmniGenome_START_DOCSTRING -) -# Copied from transformers.models.esm.modeling_esm.EsmForMaskedLM with Esm->OmniGenome -class OmniGenomeForMaskedLM(OmniGenomePreTrainedModel): - _tied_weights_keys = ["lm_head.decoder.weight"] - - def __init__(self, config): - super().__init__(config) - - if config.is_decoder: - logger.warning( - "If you want to use `OmniGenomeForMaskedLM` make sure `config.is_decoder=False` for " - "bi-directional self-attention." - ) - - self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) - self.lm_head = OmniGenomeLMHead(config) - self.init_weights() - - def get_output_embeddings(self): - return self.lm_head.decoder - - def set_output_embeddings(self, new_embeddings): - self.lm_head.decoder = new_embeddings - - @add_start_docstrings_to_model_forward( - OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") - ) - @add_code_sample_docstrings( - checkpoint=_CHECKPOINT_FOR_DOC, - output_type=MaskedLMOutput, - config_class=_CONFIG_FOR_DOC, - mask="", - ) - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, MaskedLMOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., - config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the - loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` - kwargs (`Dict[str, any]`, optional, defaults to *{}*): - Used to hide legacy arguments that have been deprecated. - """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - outputs = self.OmniGenome( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - sequence_output = outputs[0] - prediction_scores = self.lm_head(sequence_output) - - masked_lm_loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - - labels = labels.to(prediction_scores.device) - masked_lm_loss = loss_fct( - prediction_scores.view(-1, self.config.vocab_size), labels.view(-1) - ) - - if not return_dict: - output = (prediction_scores,) + outputs[2:] - return ( - ((masked_lm_loss,) + output) if masked_lm_loss is not None else output - ) - - return MaskedLMOutput( - loss=masked_lm_loss, - logits=prediction_scores, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def predict_contacts(self, tokens, attention_mask): - return self.OmniGenome.predict_contacts(tokens, attention_mask=attention_mask) - - -# Copied from transformers.models.esm.modeling_esm.EsmLMHead with Esm->OmniGenome -class OmniGenomeLMHead(nn.Module): - """OmniGenome Head for masked language modeling.""" - - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.hidden_size, config.hidden_size) - self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - - self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - self.bias = nn.Parameter(torch.zeros(config.vocab_size)) - - def forward(self, features, **kwargs): - x = self.dense(features) - x = gelu(x) - x = self.layer_norm(x) - - # project back to size of vocabulary with bias - x = self.decoder(x) + self.bias - return x - - -@add_start_docstrings( - """ - OmniGenome Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled - output) e.g. for GLUE tasks. - """, - OmniGenome_START_DOCSTRING, -) -class OmniGenomeForSequenceClassification(OmniGenomePreTrainedModel): - def __init__(self, config): - super().__init__(config) - self.num_labels = config.num_labels - self.config = config - self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) - self.classifier = OmniGenomeClassificationHead(config) - self.init_weights() - - @add_start_docstrings_to_model_forward( - OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") - ) - @add_code_sample_docstrings( - checkpoint=_CHECKPOINT_FOR_DOC, - output_type=SequenceClassifierOutput, - config_class=_CONFIG_FOR_DOC, - ) - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, SequenceClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - outputs = self.OmniGenome( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - last_hidden_state = outputs[0] - logits = self.classifier(last_hidden_state) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and ( - labels.dtype == torch.long or labels.dtype == torch.int - ): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - -@add_start_docstrings( - """ - OmniGenome Model with a token classification head on top (a linear layer on top of the hidden-states output) - Note that this model is pre-trained for RNA secondary structure prediction and can be used for zero-shot RNA - secondary structure prediction. Please find more advanced usages at https://github.com/yangheng95/OmniGenome - This model can be fine-tuned for other token classification tasks. - """, - OmniGenome_START_DOCSTRING, -) -# Copied from transformers.models.esm.modeling_esm.EsmForTokenClassification with Esm->OmniGenome -class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel): - def __init__(self, config): - super().__init__(config) - self.num_labels = config.num_labels - self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) - self.dense = torch.nn.Linear(config.hidden_size, config.hidden_size) - self.classifier = torch.nn.Linear(self.config.hidden_size, self.num_labels) - self.softmax = nn.Softmax(dim=-1) - self.init_weights() - - @add_start_docstrings_to_model_forward( - OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") - ) - @add_code_sample_docstrings( - checkpoint=_CHECKPOINT_FOR_DOC, - output_type=TokenClassifierOutput, - config_class=_CONFIG_FOR_DOC, - ) - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, TokenClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. - """ - - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - outputs = self.OmniGenome( - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - last_hidden_state = outputs[0] - last_hidden_state = self.dense(last_hidden_state) - logits = self.classifier(last_hidden_state) - logits = self.softmax(logits) - - loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return TokenClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - @staticmethod - def verify_secondary_structure(structure): - structure = list(structure) - left_brackets = [] - right_brackets = [] - for i, char in enumerate(structure): - if char == "(": - left_brackets.append(i) - elif char == ")": - if left_brackets: - left_brackets.pop() - else: - right_brackets.append(i) - - for i in left_brackets: - structure[i] = "." - for i in right_brackets: - structure[i] = "." - - structure = "".join(structure) - - return structure - - def predict_rna_structure( - self, - sequence: str, - **kwargs - ) -> List[str]: - r""" - Load the pretrained OmniGenome Model to do zero-shot prediction of the secondary structure - of a sequence given the sequence - """ - if self.tokenizer is None: - tokenizer = kwargs.get("tokenizer", None) - if tokenizer is None: - from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path) - else: - self.tokenizer = tokenizer - - inputs = self.tokenizer(sequence, return_tensors="pt", padding="max_length", truncation=True) - input_ids = inputs["input_ids"] - attention_mask = inputs["attention_mask"] - outputs = self.forward(input_ids, attention_mask, **kwargs) - - logits = torch.argmax(outputs.logits, dim=-1) - lengths = torch.sum(torch.ne(torch.tensor(0), attention_mask), dim=-1) - structures = [] - for i, length in enumerate(lengths): - structure = logits[i, :length].cpu().numpy() - structure = "".join(self.config.id2label[label] for label in structure) - if self.config.verify_ss: - structure = self.verify_secondary_structure(structure) - structures.append(structure) - return structures - - -@add_start_docstrings( - """ - This is not a standard Seq2Seq model. Instead, this model is designed for RNA design tasks. - This is the OmniGenome Model with a simple genetic algorithm based RNA design head on top. - """, - OmniGenome_START_DOCSTRING, -) -class OmniGenomeModelForSeq2SeqLM(OmniGenomePreTrainedModel): - def __init__(self, config): - super().__init__(config) - self.num_labels = config.num_labels - self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) - self.lm_head = OmniGenomeLMHead(config) - self.num_generation = config.num_generation - self.num_population = config.num_population - self.init_weights() - - self.tokenizer = None - self.predict_structure = None - - warnings.warn(f"This model {self.__class__.__name__} is not a real Seq2Seq model. " - f"Instead, this model is designed for RNA design tasks") - - @add_start_docstrings_to_model_forward( - OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") - ) - @add_code_sample_docstrings( - checkpoint=_CHECKPOINT_FOR_DOC, - output_type=TokenClassifierOutput, - config_class=_CONFIG_FOR_DOC, - ) - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = True, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, TokenClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. - """ - raise NotImplementedError("This model is not designed for standard Seq2Seq tasks. " - "Use model.rna_sequence_design() for RNA sequences design instead.") - - def rna_sequence_design( - self, - structure: str, - predict_structure_func=None, - **kwargs - ) -> List[str]: - """ - Assemble the RNA sequence given the reference sequence structure - """ - if self.tokenizer is None: - tokenizer = kwargs.get("tokenizer", None) - if tokenizer is None: - from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path) - else: - self.tokenizer = tokenizer - - candidates = self.genetic_algorithm_for_rna_design(structure, predict_structure_func=None, **kwargs) - - return candidates - - def genetic_algorithm_for_rna_design(self, structure, predict_structure_func=None, **kwargs): - if predict_structure_func is None: - import ViennaRNA - - def predict_structure(sequence): - return ViennaRNA.fold(sequence)[0] - - predict_structure_func = predict_structure - - self.predict_structure = predict_structure_func - mutation_ratio = kwargs.get("mutation_ratio", 0.5) - num_population = kwargs.get("num_population", self.num_population) - num_generation = kwargs.get("num_generation", self.num_generation) - import tqdm - population = self.init_population(structure, num_population) - population = self.mlm_mutate(population, structure, mutation_ratio=mutation_ratio) - for generation_id in tqdm.tqdm(range(num_generation), desc="Designing RNA Sequence"): - population_fitness = self.sequence_fitness(population, structure)[:num_population] - population = sorted(zip(population, population_fitness), key=lambda x: x[1])[:num_population] - population = [x[0] for x in population] - next_generation = population # Elitism - next_generation += self.crossover(population, structure) - next_generation += self.mlm_mutate(next_generation, structure, mutation_ratio) - fitness_values = self.sequence_fitness(next_generation, structure) - next_generation = sorted(zip(next_generation, fitness_values), key=lambda x: x[1]) - - candidate_sequences = [] - for sequence, fitness in next_generation: - if fitness == 0: - candidate_sequences.append(sequence) - else: - break - if candidate_sequences: - return candidate_sequences - print(f"Generation {generation_id}: {next_generation[0][0]} with fitness {next_generation[0][1]}") - population = [x[0] for x in next_generation[:num_population]] - - return [] - - def init_population(self, structure, num_population): - # Initialize lists to store population data and inputs for masked language model - population = [] - mlm_inputs = [] - # Iterate over the number of individuals in the population - for _ in range(num_population): # Changed from self.num_population to num_population - # Create a sequence by randomly choosing nucleotides or a mask token for each position in the structure - masked_sequence = [ - random.choice(["A", "G", "C", "T", ""]) - for _ in range(len(structure)) - ] - masked_sequence_str = "".join(masked_sequence) - mlm_inputs.append(f"{masked_sequence_str}{''.join(structure)}") - - # Call a function to predict outputs using the masked language model - outputs = self.mlm_predict(mlm_inputs, structure) - - # Decode the mlm outputs and construct the initial population - for i in range(len(outputs)): - sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) - fixed_sequence = [ - x if x in "AGCT" else random.choice(["G", "C"]) - for x, y in zip(sequence, list(mlm_inputs[i].replace('', '$'))) - ] - population.append("".join(fixed_sequence)) - - return population - - def mlm_mutate(self, population, structure, mutation_ratio): - def mutate(sequence, mutation_rate): - sequence = np.array(list(sequence), dtype=np.str_) - probability_matrix = np.full(sequence.shape, mutation_rate) - masked_indices = np.random.rand(*sequence.shape) < probability_matrix - sequence[masked_indices] = "$" - mut_seq = "".join(sequence.tolist()).replace("$", "") - return mut_seq - - # Initialize lists to store population data and inputs for masked language model - mlm_inputs = [] - masked_sequences = [] - - # Iterate over the number of individuals in the population - for sequence in population: - # Create a sequence by randomly choosing nucleotides or a mask token for each position in the structure - masked_sequence = mutate(sequence, mutation_ratio) - masked_sequences.append(masked_sequence) - mlm_inputs.append(f"{masked_sequence}{''.join(structure)}") - - # Call a function to predict outputs using the masked language model - outputs = self.mlm_predict(mlm_inputs, structure) - - mut_population = [] - - # Decode the mlm outputs and construct the initial population - for i in range(len(outputs)): - sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) - fixed_sequence = [ - x if x in "AGCT" else random.choice(["G", "C"]) - for x, y in zip(sequence, list(masked_sequences[i].replace('', '$'))) - ] - mut_population.append("".join(fixed_sequence)) - - return mut_population - - def crossover(self, population, structure): - crossover_population = [] - batch_crossover_inputs = [] - for i in range(len(population)): - parent1, parent2 = random.choices(population, k=2) - pos = random.randint(1, len(parent1) - 1) - child1 = parent1[:pos] + "" * len(parent2[pos:]) - child2 = "" * len(parent1[:pos]) + parent2[pos:] - batch_crossover_inputs.append(f"{child1}{structure}") - batch_crossover_inputs.append(f"{child2}{structure}") - - outputs = self.mlm_predict(batch_crossover_inputs, structure) - - for i in range(len(outputs)): - sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) - fixed_sequence = [ - x if x in "AGCT" else random.choice(["G", "C"]) - for x, y in zip(sequence, list(batch_crossover_inputs[i].replace('', '$'))) - ] - crossover_population.append("".join(fixed_sequence)) - - return crossover_population - - def sequence_fitness(self, sequences, structure): - fitness_values = [] - structures = [self.predict_structure(sequence) for sequence in sequences] - for predicted_structure in structures: - scores = [] - for i in range(len(predicted_structure)): - if predicted_structure[i] == structure[i]: - scores.append(1) - elif ( - predicted_structure[i] == ")" - and structure[i] == "(" - or predicted_structure[i] == "(" - and structure[i] == ")" - ): - scores.append(-3) - else: - scores.append(0) - score = 1 - sum(scores) / len(structure) - fitness_values.append(score) - return fitness_values - - def mlm_predict(self, mlm_inputs, structure): - batch_size = 4 - all_outputs = [] - from transformers import set_seed - set_seed(random.randint(0, 99999999), deterministic=False) - - with torch.no_grad(): - for i in range(0, len(mlm_inputs), batch_size): - batch_mlm_inputs = self.tokenizer( - mlm_inputs[i:i + batch_size], - padding=True, - max_length=len(mlm_inputs[0]) // 2, - truncation=True, - return_tensors="pt", - ) - batch_mlm_inputs = batch_mlm_inputs.to(self.device) - outputs = self.OmniGenome(**batch_mlm_inputs)[0] - outputs = self.lm_head(outputs) - outputs = outputs.argmax(dim=-1) - all_outputs.append(outputs) - outputs = torch.cat(all_outputs, dim=0) - return outputs[:, 1:1 + len(structure)] - - -# Copied from transformers.models.esm.modeling_esm.EsmClassificationHead with Esm->OmniGenome -class OmniGenomeClassificationHead(nn.Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, config): - super().__init__() - self.dense = nn.Linear(config.hidden_size, config.hidden_size) - self.dropout = nn.Dropout(config.hidden_dropout_prob) - self.out_proj = nn.Linear(config.hidden_size, config.num_labels) - - def forward(self, features, **kwargs): - x = features[:, 0, :] # take token (equiv. to [CLS]) - x = self.dropout(x) - x = self.dense(x) - x = torch.tanh(x) - x = self.dropout(x) - x = self.out_proj(x) - return x - - -def create_position_ids_from_input_ids( - input_ids, padding_idx, past_key_values_length=0 -): - """ - Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols - are ignored. This is modified from fairseq's `utils.make_positions`. - - Args: - x: torch.Tensor x: - - Returns: torch.Tensor - """ - # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA. - mask = input_ids.ne(padding_idx).int() - incremental_indices = ( - torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length - ) * mask - return incremental_indices.long() + padding_idx +# coding=utf-8 +# Copyright 2022 ColaLab-UoE (https://colalab.ai/), Meta and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" PyTorch OmniGenome model.""" + +import math +import random +import warnings +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers import add_start_docstrings, PreTrainedModel + +from transformers.modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPoolingAndCrossAttentions, + MaskedLMOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + +from transformers.pytorch_utils import ( + find_pruneable_heads_and_indices, + prune_linear_layer, +) + +from transformers.utils import ( + logging, + add_code_sample_docstrings, + add_start_docstrings_to_model_forward, +) + +from .configuration_omnigenome import OmniGenomeConfig + +try: + from flash_attn import flash_attn_func +except ImportError: + flash_attn_func = None + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "yangheng/OmniGenome-52M" +_CONFIG_FOR_DOC = "OmniGenomeConfig" + +OmniGenome_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "yangheng/OmniGenome-52M", + # This is not a complete list of all OmniGenome models! + # See all OmniGenome models at https://huggingface.co/models?filter=OmniGenome +] + + +def rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(x, cos, sin): + cos = cos[:, :, : x.shape[-2], :] + sin = sin[:, :, : x.shape[-2], :] + + return (x * cos) + (rotate_half(x) * sin) + + +def gelu(x): + """ + This is the gelu implementation from the original OmniGenome repo. Using F.gelu yields subtly wrong results. + """ + return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) + + +def symmetrize(x): + "Make layer symmetric in final two dimensions, used for contact prediction." + return x + x.transpose(-1, -2) + + +def average_product_correct(x): + "Perform average product correct, used for contact prediction." + a1 = x.sum(-1, keepdims=True) + a2 = x.sum(-2, keepdims=True) + a12 = x.sum((-1, -2), keepdims=True) + + avg = a1 * a2 + avg.div_(a12) # in-place to reduce memory + normalized = x - avg + return normalized + + +# Copied from transformers.models.esm.modeling_esm.RotaryEmbedding +class RotaryEmbedding(torch.nn.Module): + """ + Rotary position embeddings based on those in + [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation + matrices which depend on their relative positions. + """ + + def __init__(self, dim: int): + super().__init__() + # Generate and save the inverse frequency buffer (non trainable) + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) + inv_freq = inv_freq + self.register_buffer("inv_freq", inv_freq) + + self._seq_len_cached = None + self._cos_cached = None + self._sin_cached = None + + def _update_cos_sin_tables(self, x, seq_dimension=2): + seq_len = x.shape[seq_dimension] + + # Reset the tables if the sequence length has changed, + # or if we're on a new device (possibly due to tracing for instance) + if seq_len != self._seq_len_cached or self._cos_cached.device != x.device: + self._seq_len_cached = seq_len + t = torch.arange(x.shape[seq_dimension], device=x.device).type_as( + self.inv_freq + ) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1).to(x.device) + + self._cos_cached = emb.cos()[None, None, :, :] + self._sin_cached = emb.sin()[None, None, :, :] + + return self._cos_cached, self._sin_cached + + def forward( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + self._cos_cached, self._sin_cached = self._update_cos_sin_tables( + k, seq_dimension=-2 + ) + + return ( + apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), + apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), + ) + + +# Copied from transformers.models.esm.modeling_esm.EsmContactPredictionHead with Esm->OmniGenome +class OmniGenomeContactPredictionHead(nn.Module): + """Performs symmetrization, apc, and computes a logistic regression on the output features""" + + def __init__( + self, + in_features: int, + bias=True, + eos_idx: int = 2, + ): + super().__init__() + self.in_features = in_features + self.eos_idx = eos_idx + self.regression = nn.Linear(in_features, 1, bias) + self.activation = nn.Sigmoid() + + def forward(self, tokens, attentions): + # remove eos token attentions + eos_mask = tokens.ne(self.eos_idx).to(attentions) + eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2) + attentions = attentions * eos_mask[:, None, None, :, :] + attentions = attentions[..., :-1, :-1] + # remove cls token attentions + attentions = attentions[..., 1:, 1:] + batch_size, layers, heads, seqlen, _ = attentions.size() + attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen) + + # features: batch x channels x tokens x tokens (symmetric) + attentions = attentions.to( + self.regression.weight.device + ) # attentions always float32, may need to convert to float16 + attentions = average_product_correct(symmetrize(attentions)) + attentions = attentions.permute(0, 2, 3, 1) + return self.activation(self.regression(attentions).squeeze(3)) + + +# Copied from transformers.models.esm.modeling_esm.EsmEmbeddings with Esm->OmniGenome +class OmniGenomeEmbeddings(nn.Module): + """ + Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. + """ + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding( + config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id + ) + + if config.emb_layer_norm_before: + self.layer_norm = nn.LayerNorm( + config.hidden_size, eps=config.layer_norm_eps + ) + else: + self.layer_norm = None + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.position_embedding_type = getattr( + config, "position_embedding_type", "absolute" + ) + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False, + ) + + self.padding_idx = config.pad_token_id + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, + config.hidden_size, + padding_idx=self.padding_idx, + ) + self.token_dropout = config.token_dropout + self.mask_token_id = config.mask_token_id + + def forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + inputs_embeds=None, + past_key_values_length=0, + ): + if position_ids is None: + if input_ids is not None: + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = create_position_ids_from_input_ids( + input_ids, self.padding_idx, past_key_values_length + ) + else: + position_ids = self.create_position_ids_from_inputs_embeds( + inputs_embeds + ) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + # Note that if we want to support OmniGenome-1 (not 1b!) in future then we need to support an + # embedding_scale factor here. + embeddings = inputs_embeds + + # Matt: OmniGenome has the option to handle masking in MLM in a slightly unusual way. If the token_dropout + # flag is False then it is handled in the same was as BERT/RoBERTa. If it is set to True, however, + # masked tokens are treated as if they were selected for input dropout and zeroed out. + # This "mask-dropout" is compensated for when masked tokens are not present, by scaling embeddings by + # a factor of (fraction of unmasked tokens during training) / (fraction of unmasked tokens in sample). + # This is analogous to the way that dropout layers scale down outputs during evaluation when not + # actually dropping out values (or, equivalently, scale up their un-dropped outputs in training). + if self.token_dropout: + embeddings = embeddings.masked_fill( + (input_ids == self.mask_token_id).unsqueeze(-1), 0.0 + ) + mask_ratio_train = ( + 0.15 * 0.8 + ) # Hardcoded as the ratio used in all OmniGenome model training runs + src_lengths = attention_mask.sum(-1) + mask_ratio_observed = (input_ids == self.mask_token_id).sum( + -1 + ).float() / src_lengths + embeddings = ( + embeddings + * (1 - mask_ratio_train) + / (1 - mask_ratio_observed)[:, None, None] + ).to(embeddings.dtype) + + if self.position_embedding_type == "absolute": + position_embeddings = self.position_embeddings(position_ids) + embeddings = embeddings + position_embeddings + + if self.layer_norm is not None: + embeddings = self.layer_norm(embeddings) + if attention_mask is not None: + embeddings = (embeddings * attention_mask.unsqueeze(-1)).to( + embeddings.dtype + ) + # Matt: I think this line was copied incorrectly from BERT, disabling it for now. + # embeddings = self.dropout(embeddings) + return embeddings + + def create_position_ids_from_inputs_embeds(self, inputs_embeds): + """ + We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids. + + Args: + inputs_embeds: torch.Tensor + + Returns: torch.Tensor + """ + input_shape = inputs_embeds.size()[:-1] + sequence_length = input_shape[1] + + position_ids = torch.arange( + self.padding_idx + 1, + sequence_length + self.padding_idx + 1, + dtype=torch.long, + device=inputs_embeds.device, + ) + return position_ids.unsqueeze(0).expand(input_shape) + +# +# # Copied from transformers.models.esm.modeling_esm.EsmSelfAttention with Esm->OmniGenome +# class OmniGenomeSelfAttention(nn.Module): +# def __init__(self, config, position_embedding_type=None): +# super().__init__() +# if config.hidden_size % config.num_attention_heads != 0 and not hasattr( +# config, "embedding_size" +# ): +# raise ValueError( +# f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " +# f"heads ({config.num_attention_heads})" +# ) +# +# self.num_attention_heads = config.num_attention_heads +# self.attention_head_size = int(config.hidden_size / config.num_attention_heads) +# self.all_head_size = self.num_attention_heads * self.attention_head_size +# +# self.query = nn.Linear(config.hidden_size, self.all_head_size) +# self.key = nn.Linear(config.hidden_size, self.all_head_size) +# self.value = nn.Linear(config.hidden_size, self.all_head_size) +# +# self.dropout = nn.Dropout(config.attention_probs_dropout_prob) +# self.position_embedding_type = position_embedding_type or getattr( +# config, "position_embedding_type", "absolute" +# ) +# self.rotary_embeddings = None +# if ( +# self.position_embedding_type == "relative_key" +# or self.position_embedding_type == "relative_key_query" +# ): +# self.max_position_embeddings = config.max_position_embeddings +# self.distance_embedding = nn.Embedding( +# 2 * config.max_position_embeddings - 1, self.attention_head_size +# ) +# elif self.position_embedding_type == "rotary": +# self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) +# +# self.is_decoder = config.is_decoder +# +# def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: +# new_x_shape = x.size()[:-1] + ( +# self.num_attention_heads, +# self.attention_head_size, +# ) +# x = x.view(new_x_shape) +# return x.permute(0, 2, 1, 3) +# +# def forward( +# self, +# hidden_states: torch.Tensor, +# attention_mask: Optional[torch.FloatTensor] = None, +# head_mask: Optional[torch.FloatTensor] = None, +# encoder_hidden_states: Optional[torch.FloatTensor] = None, +# encoder_attention_mask: Optional[torch.FloatTensor] = None, +# past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, +# output_attentions: Optional[bool] = False, +# ) -> Tuple[torch.Tensor]: +# mixed_query_layer = self.query(hidden_states) +# +# # If this is instantiated as a cross-attention module, the keys +# # and values come from an encoder; the attention mask needs to be +# # such that the encoder's padding tokens are not attended to. +# is_cross_attention = encoder_hidden_states is not None +# +# if is_cross_attention and past_key_value is not None: +# # reuse k,v, cross_attentions +# key_layer = past_key_value[0] +# value_layer = past_key_value[1] +# attention_mask = encoder_attention_mask +# elif is_cross_attention: +# key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) +# value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) +# attention_mask = encoder_attention_mask +# elif past_key_value is not None: +# key_layer = self.transpose_for_scores(self.key(hidden_states)) +# value_layer = self.transpose_for_scores(self.value(hidden_states)) +# key_layer = torch.cat([past_key_value[0], key_layer], dim=2) +# value_layer = torch.cat([past_key_value[1], value_layer], dim=2) +# else: +# key_layer = self.transpose_for_scores(self.key(hidden_states)) +# value_layer = self.transpose_for_scores(self.value(hidden_states)) +# +# query_layer = self.transpose_for_scores(mixed_query_layer) +# +# # Matt: Our BERT model (which this code was derived from) scales attention logits down by sqrt(head_dim). +# # OmniGenome scales the query down by the same factor instead. Modulo numerical stability these are equivalent, +# # but not when rotary embeddings get involved. Therefore, we scale the query here to match the original +# # OmniGenome code and fix rotary embeddings. +# query_layer = query_layer * self.attention_head_size ** -0.5 +# +# if self.is_decoder: +# # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. +# # Further calls to cross_attention layer can then reuse all cross-attention +# # key/value_states (first "if" case) +# # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of +# # all previous decoder key/value_states. Further calls to uni-directional self-attention +# # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) +# # if encoder bi-directional self-attention `past_key_value` is always `None` +# past_key_value = (key_layer, value_layer) +# +# if self.position_embedding_type == "rotary": +# query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) +# +# # Take the dot product between "query" and "key" to get the raw attention scores. +# attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) +# +# if ( +# self.position_embedding_type == "relative_key" +# or self.position_embedding_type == "relative_key_query" +# ): +# seq_length = hidden_states.size()[1] +# position_ids_l = torch.arange( +# seq_length, dtype=torch.long, device=hidden_states.device +# ).view(-1, 1) +# position_ids_r = torch.arange( +# seq_length, dtype=torch.long, device=hidden_states.device +# ).view(1, -1) +# distance = position_ids_l - position_ids_r +# positional_embedding = self.distance_embedding( +# distance + self.max_position_embeddings - 1 +# ) +# positional_embedding = positional_embedding.to( +# dtype=query_layer.dtype +# ) # fp16 compatibility +# +# if self.position_embedding_type == "relative_key": +# relative_position_scores = torch.einsum( +# "bhld,lrd->bhlr", query_layer, positional_embedding +# ) +# attention_scores = attention_scores + relative_position_scores +# elif self.position_embedding_type == "relative_key_query": +# relative_position_scores_query = torch.einsum( +# "bhld,lrd->bhlr", query_layer, positional_embedding +# ) +# relative_position_scores_key = torch.einsum( +# "bhrd,lrd->bhlr", key_layer, positional_embedding +# ) +# attention_scores = ( +# attention_scores +# + relative_position_scores_query +# + relative_position_scores_key +# ) +# +# if attention_mask is not None: +# # Apply the attention mask is (precomputed for all layers in OmniGenomeModel forward() function) +# attention_scores = attention_scores + attention_mask +# +# # Normalize the attention scores to probabilities. +# attention_probs = nn.functional.softmax(attention_scores, dim=-1) +# +# # This is actually dropping out entire tokens to attend to, which might +# # seem a bit unusual, but is taken from the original Transformer paper. +# attention_probs = self.dropout(attention_probs) +# +# # Mask heads if we want to +# if head_mask is not None: +# attention_probs = attention_probs * head_mask +# +# context_layer = torch.matmul(attention_probs, value_layer) +# +# context_layer = context_layer.permute(0, 2, 1, 3).contiguous() +# new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) +# context_layer = context_layer.view(new_context_layer_shape) +# +# outputs = ( +# (context_layer, attention_probs) if output_attentions else (context_layer,) +# ) +# +# if self.is_decoder: +# outputs = outputs + (past_key_value,) +# return outputs + + +# Copied from transformers.models.esm.modeling_esm.EsmSelfAttention with Esm->OmniGenome +class OmniGenomeSelfAttention(nn.Module): + def __init__(self, config, position_embedding_type=None): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr( + config, "embedding_size" + ): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + self.position_embedding_type = position_embedding_type or getattr( + config, "position_embedding_type", "absolute" + ) + self.rotary_embeddings = None + if ( + self.position_embedding_type == "relative_key" + or self.position_embedding_type == "relative_key_query" + ): + self.max_position_embeddings = config.max_position_embeddings + self.distance_embedding = nn.Embedding( + 2 * config.max_position_embeddings - 1, self.attention_head_size + ) + elif self.position_embedding_type == "rotary": + self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) + + self.is_decoder = config.is_decoder + + # FlashAttention parameters + self.enable_flash_attn = getattr(config, "use_flash_attention", True) + if self.enable_flash_attn: + self.flash_attn_func = flash_attn_func + else: + self.flash_attn_func = None + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor]: + mixed_query_layer = self.query(hidden_states) + + is_cross_attention = encoder_hidden_states is not None + + if is_cross_attention and past_key_value is not None: + key_layer = past_key_value[0] + value_layer = past_key_value[1] + attention_mask = encoder_attention_mask + elif is_cross_attention: + key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) + value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) + attention_mask = encoder_attention_mask + elif past_key_value is not None: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + key_layer = torch.cat([past_key_value[0], key_layer], dim=2) + value_layer = torch.cat([past_key_value[1], value_layer], dim=2) + else: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + + query_layer = self.transpose_for_scores(mixed_query_layer) + + if self.is_decoder: + past_key_value = (key_layer, value_layer) + + # 使用FlashAttention的条件判断 + use_flash_attn = self.enable_flash_attn and self.position_embedding_type == "rotary" + if use_flash_attn and self.flash_attn_func is not None: + # 应用旋转位置编码 + query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) + + # 调整维度顺序为 [batch_size, seq_len, num_heads, head_dim] + q = query_layer.transpose(1, 2).half() + k = key_layer.transpose(1, 2).half() + v = value_layer.transpose(1, 2).half() + + # 使用FlashAttention计算 + context_layer = self.flash_attn_func( + q, k, v, + dropout_p=self.dropout.p if self.training else 0.0, + softmax_scale=self.attention_head_size ** -0.5, + causal=self.is_decoder + ) + + # 恢复维度顺序 [batch_size, num_heads, seq_len, head_dim] + context_layer = context_layer.transpose(1, 2).to(hidden_states.dtype) + else: + # 原始实现 + query_layer = query_layer * self.attention_head_size ** -0.5 + + if self.position_embedding_type == "rotary": + query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer) + + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + if self.position_embedding_type in ["relative_key", "relative_key_query"]: + seq_length = hidden_states.size()[1] + position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1) + position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1) + distance = position_ids_l - position_ids_r + positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1) + positional_embedding = positional_embedding.to(dtype=query_layer.dtype) + + if self.position_embedding_type == "relative_key": + relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) + attention_scores = attention_scores + relative_position_scores + elif self.position_embedding_type == "relative_key_query": + relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) + relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding) + attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key + + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + attention_probs = self.dropout(attention_probs) + + if head_mask is not None: + attention_probs = attention_probs * head_mask + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + if self.is_decoder: + outputs = outputs + (past_key_value,) + return outputs + +# Copied from transformers.models.esm.modeling_esm.EsmSelfOutput with Esm->OmniGenome +class OmniGenomeSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +# Copied from transformers.models.esm.modeling_esm.EsmAttention with Esm->OmniGenome +class OmniGenomeAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = OmniGenomeSelfAttention(config) + self.output = OmniGenomeSelfOutput(config) + self.pruned_heads = set() + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def prune_heads(self, heads): + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, + self.self.num_attention_heads, + self.self.attention_head_size, + self.pruned_heads, + ) + + # Prune linear layers + self.self.query = prune_linear_layer(self.self.query, index) + self.self.key = prune_linear_layer(self.self.key, index) + self.self.value = prune_linear_layer(self.self.value, index) + self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) + + # Update hyper params and store pruned heads + self.self.num_attention_heads = self.self.num_attention_heads - len(heads) + self.self.all_head_size = ( + self.self.attention_head_size * self.self.num_attention_heads + ) + self.pruned_heads = self.pruned_heads.union(heads) + + def forward( + self, + hidden_states, + attention_mask=None, + head_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_value=None, + output_attentions=False, + ): + hidden_states_ln = self.LayerNorm(hidden_states) + hidden_states_ln = hidden_states_ln.to(hidden_states.dtype) + self_outputs = self.self( + hidden_states_ln, + attention_mask, + head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + ) + attention_output = self.output(self_outputs[0], hidden_states) + outputs = (attention_output,) + self_outputs[ + 1: + ] # add attentions if we output them + return outputs + + +# Copied from transformers.models.esm.modeling_esm.EsmIntermediate with Esm->OmniGenome +class OmniGenomeIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = gelu(hidden_states) + return hidden_states + + +# Copied from transformers.models.esm.modeling_esm.EsmOutput with Esm->OmniGenome +class OmniGenomeOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +# Copied from transformers.models.esm.modeling_esm.EsmLayer with Esm->OmniGenome +class OmniGenomeLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = OmniGenomeAttention(config) + self.is_decoder = config.is_decoder + self.add_cross_attention = config.add_cross_attention + if self.add_cross_attention: + if not self.is_decoder: + raise RuntimeError( + f"{self} should be used as a decoder model if cross attention is added" + ) + self.crossattention = OmniGenomeAttention(config) + self.intermediate = OmniGenomeIntermediate(config) + self.output = OmniGenomeOutput(config) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states, + attention_mask=None, + head_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_value=None, + output_attentions=False, + ): + # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 + self_attn_past_key_value = ( + past_key_value[:2] if past_key_value is not None else None + ) + self_attention_outputs = self.attention( + hidden_states, + attention_mask, + head_mask, + output_attentions=output_attentions, + past_key_value=self_attn_past_key_value, + ) + attention_output = self_attention_outputs[0] + + # if decoder, the last output is tuple of self-attn cache + if self.is_decoder: + outputs = self_attention_outputs[1:-1] + present_key_value = self_attention_outputs[-1] + else: + outputs = self_attention_outputs[ + 1: + ] # add self attentions if we output attention weights + + cross_attn_present_key_value = None + if self.is_decoder and encoder_hidden_states is not None: + if not hasattr(self, "crossattention"): + raise AttributeError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated" + " with cross-attention layers by setting `config.add_cross_attention=True`" + ) + + # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple + cross_attn_past_key_value = ( + past_key_value[-2:] if past_key_value is not None else None + ) + cross_attention_outputs = self.crossattention( + attention_output, + attention_mask, + head_mask, + encoder_hidden_states, + encoder_attention_mask, + cross_attn_past_key_value, + output_attentions, + ) + attention_output = cross_attention_outputs[0] + outputs = ( + outputs + cross_attention_outputs[1:-1] + ) # add cross attentions if we output attention weights + + # add cross-attn cache to positions 3,4 of present_key_value tuple + cross_attn_present_key_value = cross_attention_outputs[-1] + present_key_value = present_key_value + cross_attn_present_key_value + + layer_output = self.feed_forward_chunk(attention_output) + + outputs = (layer_output,) + outputs + + # if decoder, return the attn key/values as the last output + if self.is_decoder: + outputs = outputs + (present_key_value,) + return outputs + + def feed_forward_chunk(self, attention_output): + attention_output_ln = self.LayerNorm(attention_output) + intermediate_output = self.intermediate(attention_output_ln) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +# Copied from transformers.models.esm.modeling_esm.EsmEncoder with Esm->OmniGenome +class OmniGenomeEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList( + [OmniGenomeLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.emb_layer_norm_after = nn.LayerNorm( + config.hidden_size, eps=config.layer_norm_eps + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + attention_mask=None, + head_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ): + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting " + "`use_cache=False`..." + ) + use_cache = False + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = ( + () if output_attentions and self.config.add_cross_attention else None + ) + + next_decoder_cache = () if use_cache else None + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_head_mask = head_mask[i] if head_mask is not None else None + past_key_value = past_key_values[i] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + attention_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + ) + else: + layer_outputs = layer_module( + hidden_states, + attention_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + ) + + hidden_states = layer_outputs[0] + if use_cache: + next_decoder_cache = next_decoder_cache + (layer_outputs[-1],) + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + if self.config.add_cross_attention: + all_cross_attentions = all_cross_attentions + (layer_outputs[2],) + + if self.emb_layer_norm_after: + hidden_states = self.emb_layer_norm_after(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_decoder_cache, + all_hidden_states, + all_self_attentions, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=next_decoder_cache, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + +# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->OmniGenome +class OmniGenomePooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +# Copied from transformers.models.esm.modeling_esm.EsmPreTrainedModel with Esm->OmniGenome +class OmniGenomePreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = OmniGenomeConfig + base_model_prefix = "OmniGenome" + supports_gradient_checkpointing = True + _no_split_modules = [ + "OmniGenomeLayer", + "OmniGenomeFoldTriangularSelfAttentionBlock", + "OmniGenomeEmbeddings", + ] + + # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + +OmniGenome_START_DOCSTRING = r""" + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`OmniGenomeConfig`]): Model configuration class with all the parameters of the + model. Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +OmniGenome_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + position_ids (`torch.LongTensor` of shape `({0})`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare OmniGenome Model transformer outputting raw hidden-states without any specific head on top.", + OmniGenome_START_DOCSTRING, +) +# Copied from transformers.models.esm.modeling_esm.EsmModel with Esm->OmniGenome +class OmniGenomeModel(OmniGenomePreTrainedModel): + """ + + The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of + cross-attention is added between the self-attention layers, following the architecture described in [Attention is + all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, + Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. + + To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set + to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and + `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. + """ + + def __init__(self, config, add_pooling_layer=True): + super().__init__(config) + self.config = config + + self.embeddings = OmniGenomeEmbeddings(config) + self.encoder = OmniGenomeEncoder(config) + + self.pooler = OmniGenomePooler(config) if add_pooling_layer else None + + self.contact_head = OmniGenomeContactPredictionHead( + in_features=config.num_hidden_layers * config.num_attention_heads, bias=True + ) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def _prune_heads(self, heads_to_prune): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward( + OmniGenome_INPUTS_DOCSTRING.format("(batch_size, sequence_length)") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPoolingAndCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + 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)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + + # past_key_values_length + past_key_values_length = ( + past_key_values[0][0].shape[2] if past_key_values is not None else 0 + ) + + if attention_mask is None: + attention_mask = torch.ones( + ((batch_size, seq_length + past_key_values_length)), device=device + ) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( + attention_mask, input_shape + ) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + ( + encoder_batch_size, + encoder_sequence_length, + _, + ) = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask( + encoder_attention_mask + ) + else: + encoder_extended_attention_mask = None + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + embedding_output = embedding_output.half() + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + pooled_output = ( + self.pooler(sequence_output) if self.pooler is not None else None + ) + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + def predict_contacts(self, tokens, attention_mask): + attns = self( + tokens, + attention_mask=attention_mask, + return_dict=True, + output_attentions=True, + ).attentions + attns = torch.stack(attns, dim=1) # Matches the original model layout + # In the original model, attentions for padding tokens are completely zeroed out. + # This makes no difference most of the time because the other tokens won't attend to them, + # but it does for the contact prediction task, which takes attentions as input, + # so we have to mimic that here. + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) + attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4) + return self.contact_head(tokens, attns) + + +@add_start_docstrings( + """OmniGenome Model with a `language modeling` head on top.""", OmniGenome_START_DOCSTRING +) +# Copied from transformers.models.esm.modeling_esm.EsmForMaskedLM with Esm->OmniGenome +class OmniGenomeForMaskedLM(OmniGenomePreTrainedModel): + _tied_weights_keys = ["lm_head.decoder.weight"] + + def __init__(self, config): + super().__init__(config) + + if config.is_decoder: + logger.warning( + "If you want to use `OmniGenomeForMaskedLM` make sure `config.is_decoder=False` for " + "bi-directional self-attention." + ) + + self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) + self.lm_head = OmniGenomeLMHead(config) + self.init_weights() + + def get_output_embeddings(self): + return self.lm_head.decoder + + def set_output_embeddings(self, new_embeddings): + self.lm_head.decoder = new_embeddings + + @add_start_docstrings_to_model_forward( + OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=MaskedLMOutput, + config_class=_CONFIG_FOR_DOC, + mask="", + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, MaskedLMOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + kwargs (`Dict[str, any]`, optional, defaults to *{}*): + Used to hide legacy arguments that have been deprecated. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.OmniGenome( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + prediction_scores = self.lm_head(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + + labels = labels.to(prediction_scores.device) + masked_lm_loss = loss_fct( + prediction_scores.view(-1, self.config.vocab_size), labels.view(-1) + ) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ( + ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + ) + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def predict_contacts(self, tokens, attention_mask): + return self.OmniGenome.predict_contacts(tokens, attention_mask=attention_mask) + + +# Copied from transformers.models.esm.modeling_esm.EsmLMHead with Esm->OmniGenome +class OmniGenomeLMHead(nn.Module): + """OmniGenome Head for masked language modeling.""" + + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + def forward(self, features, **kwargs): + x = self.dense(features) + x = gelu(x) + x = self.layer_norm(x) + + # project back to size of vocabulary with bias + x = self.decoder(x) + self.bias + return x + + +@add_start_docstrings( + """ + OmniGenome Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled + output) e.g. for GLUE tasks. + """, + OmniGenome_START_DOCSTRING, +) +class OmniGenomeForSequenceClassification(OmniGenomePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) + self.classifier = OmniGenomeClassificationHead(config) + self.init_weights() + + @add_start_docstrings_to_model_forward( + OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=SequenceClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.OmniGenome( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + last_hidden_state = outputs[0] + logits = self.classifier(last_hidden_state) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + OmniGenome Model with a token classification head on top (a linear layer on top of the hidden-states output) + Note that this model is pre-trained for RNA secondary structure prediction and can be used for zero-shot RNA + secondary structure prediction. Please find more advanced usages at https://github.com/yangheng95/OmniGenome + This model can be fine-tuned for other token classification tasks. + """, + OmniGenome_START_DOCSTRING, +) +# Copied from transformers.models.esm.modeling_esm.EsmForTokenClassification with Esm->OmniGenome +class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) + self.dense = torch.nn.Linear(config.hidden_size, config.hidden_size) + self.classifier = torch.nn.Linear(self.config.hidden_size, self.num_labels) + self.softmax = nn.Softmax(dim=-1) + self.init_weights() + + @add_start_docstrings_to_model_forward( + OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.OmniGenome( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = outputs[0] + last_hidden_state = self.dense(last_hidden_state) + logits = self.classifier(last_hidden_state) + logits = self.softmax(logits) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + @staticmethod + def verify_secondary_structure(structure): + structure = list(structure) + left_brackets = [] + right_brackets = [] + for i, char in enumerate(structure): + if char == "(": + left_brackets.append(i) + elif char == ")": + if left_brackets: + left_brackets.pop() + else: + right_brackets.append(i) + + for i in left_brackets: + structure[i] = "." + for i in right_brackets: + structure[i] = "." + + structure = "".join(structure) + + return structure + + def predict_rna_structure( + self, + sequence: str, + **kwargs + ) -> List[str]: + r""" + Load the pretrained OmniGenome Model to do zero-shot prediction of the secondary structure + of a sequence given the sequence + """ + if self.tokenizer is None: + tokenizer = kwargs.get("tokenizer", None) + if tokenizer is None: + from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path) + else: + self.tokenizer = tokenizer + + inputs = self.tokenizer(sequence, return_tensors="pt", padding="max_length", truncation=True) + input_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + outputs = self.forward(input_ids, attention_mask, **kwargs) + + logits = torch.argmax(outputs.logits, dim=-1) + lengths = torch.sum(torch.ne(torch.tensor(0), attention_mask), dim=-1) + structures = [] + for i, length in enumerate(lengths): + structure = logits[i, :length].cpu().numpy() + structure = "".join(self.config.id2label[label] for label in structure) + if self.config.verify_ss: + structure = self.verify_secondary_structure(structure) + structures.append(structure) + return structures + + +@add_start_docstrings( + """ + This is not a standard Seq2Seq model. Instead, this model is designed for RNA design tasks. + This is the OmniGenome Model with a simple genetic algorithm based RNA design head on top. + """, + OmniGenome_START_DOCSTRING, +) +class OmniGenomeModelForSeq2SeqLM(OmniGenomePreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False) + self.lm_head = OmniGenomeLMHead(config) + self.num_generation = config.num_generation + self.num_population = config.num_population + self.init_weights() + + self.tokenizer = None + self.predict_structure = None + + warnings.warn(f"This model {self.__class__.__name__} is not a real Seq2Seq model. " + f"Instead, this model is designed for RNA design tasks") + + @add_start_docstrings_to_model_forward( + OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = True, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + raise NotImplementedError("This model is not designed for standard Seq2Seq tasks. " + "Use model.rna_sequence_design() for RNA sequences design instead.") + + def rna_sequence_design( + self, + structure: str, + predict_structure_func=None, + **kwargs + ) -> List[str]: + """ + Assemble the RNA sequence given the reference sequence structure + """ + if self.tokenizer is None: + tokenizer = kwargs.get("tokenizer", None) + if tokenizer is None: + from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path) + else: + self.tokenizer = tokenizer + + candidates = self.genetic_algorithm_for_rna_design(structure, predict_structure_func=None, **kwargs) + + return candidates + + def genetic_algorithm_for_rna_design(self, structure, predict_structure_func=None, **kwargs): + if predict_structure_func is None: + import ViennaRNA + + def predict_structure(sequence): + return ViennaRNA.fold(sequence)[0] + + predict_structure_func = predict_structure + + self.predict_structure = predict_structure_func + mutation_ratio = kwargs.get("mutation_ratio", 0.5) + num_population = kwargs.get("num_population", self.num_population) + num_generation = kwargs.get("num_generation", self.num_generation) + import tqdm + population = self.init_population(structure, num_population) + population = self.mlm_mutate(population, structure, mutation_ratio=mutation_ratio) + for generation_id in tqdm.tqdm(range(num_generation), desc="Designing RNA Sequence"): + population_fitness = self.sequence_fitness(population, structure)[:num_population] + population = sorted(zip(population, population_fitness), key=lambda x: x[1])[:num_population] + population = [x[0] for x in population] + next_generation = population # Elitism + next_generation += self.crossover(population, structure) + next_generation += self.mlm_mutate(next_generation, structure, mutation_ratio) + fitness_values = self.sequence_fitness(next_generation, structure) + next_generation = sorted(zip(next_generation, fitness_values), key=lambda x: x[1]) + + candidate_sequences = [] + for sequence, fitness in next_generation: + if fitness == 0: + candidate_sequences.append(sequence) + else: + break + if candidate_sequences: + return candidate_sequences + print(f"Generation {generation_id}: {next_generation[0][0]} with fitness {next_generation[0][1]}") + population = [x[0] for x in next_generation[:num_population]] + + return [] + + def init_population(self, structure, num_population): + # Initialize lists to store population data and inputs for masked language model + population = [] + mlm_inputs = [] + # Iterate over the number of individuals in the population + for _ in range(num_population): # Changed from self.num_population to num_population + # Create a sequence by randomly choosing nucleotides or a mask token for each position in the structure + masked_sequence = [ + random.choice(["A", "G", "C", "T", ""]) + for _ in range(len(structure)) + ] + masked_sequence_str = "".join(masked_sequence) + mlm_inputs.append(f"{masked_sequence_str}{''.join(structure)}") + + # Call a function to predict outputs using the masked language model + outputs = self.mlm_predict(mlm_inputs, structure) + + # Decode the mlm outputs and construct the initial population + for i in range(len(outputs)): + sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) + fixed_sequence = [ + x if x in "AGCT" else random.choice(["G", "C"]) + for x, y in zip(sequence, list(mlm_inputs[i].replace('', '$'))) + ] + population.append("".join(fixed_sequence)) + + return population + + def mlm_mutate(self, population, structure, mutation_ratio): + def mutate(sequence, mutation_rate): + sequence = np.array(list(sequence), dtype=np.str_) + probability_matrix = np.full(sequence.shape, mutation_rate) + masked_indices = np.random.rand(*sequence.shape) < probability_matrix + sequence[masked_indices] = "$" + mut_seq = "".join(sequence.tolist()).replace("$", "") + return mut_seq + + # Initialize lists to store population data and inputs for masked language model + mlm_inputs = [] + masked_sequences = [] + + # Iterate over the number of individuals in the population + for sequence in population: + # Create a sequence by randomly choosing nucleotides or a mask token for each position in the structure + masked_sequence = mutate(sequence, mutation_ratio) + masked_sequences.append(masked_sequence) + mlm_inputs.append(f"{masked_sequence}{''.join(structure)}") + + # Call a function to predict outputs using the masked language model + outputs = self.mlm_predict(mlm_inputs, structure) + + mut_population = [] + + # Decode the mlm outputs and construct the initial population + for i in range(len(outputs)): + sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) + fixed_sequence = [ + x if x in "AGCT" else random.choice(["G", "C"]) + for x, y in zip(sequence, list(masked_sequences[i].replace('', '$'))) + ] + mut_population.append("".join(fixed_sequence)) + + return mut_population + + def crossover(self, population, structure): + crossover_population = [] + batch_crossover_inputs = [] + for i in range(len(population)): + parent1, parent2 = random.choices(population, k=2) + pos = random.randint(1, len(parent1) - 1) + child1 = parent1[:pos] + "" * len(parent2[pos:]) + child2 = "" * len(parent1[:pos]) + parent2[pos:] + batch_crossover_inputs.append(f"{child1}{structure}") + batch_crossover_inputs.append(f"{child2}{structure}") + + outputs = self.mlm_predict(batch_crossover_inputs, structure) + + for i in range(len(outputs)): + sequence = self.tokenizer.convert_ids_to_tokens(outputs[i].tolist()) + fixed_sequence = [ + x if x in "AGCT" else random.choice(["G", "C"]) + for x, y in zip(sequence, list(batch_crossover_inputs[i].replace('', '$'))) + ] + crossover_population.append("".join(fixed_sequence)) + + return crossover_population + + def sequence_fitness(self, sequences, structure): + fitness_values = [] + structures = [self.predict_structure(sequence) for sequence in sequences] + for predicted_structure in structures: + scores = [] + for i in range(len(predicted_structure)): + if predicted_structure[i] == structure[i]: + scores.append(1) + elif ( + predicted_structure[i] == ")" + and structure[i] == "(" + or predicted_structure[i] == "(" + and structure[i] == ")" + ): + scores.append(-3) + else: + scores.append(0) + score = 1 - sum(scores) / len(structure) + fitness_values.append(score) + return fitness_values + + def mlm_predict(self, mlm_inputs, structure): + batch_size = 4 + all_outputs = [] + from transformers import set_seed + set_seed(random.randint(0, 99999999), deterministic=False) + + with torch.no_grad(): + for i in range(0, len(mlm_inputs), batch_size): + batch_mlm_inputs = self.tokenizer( + mlm_inputs[i:i + batch_size], + padding=True, + max_length=len(mlm_inputs[0]) // 2, + truncation=True, + return_tensors="pt", + ) + batch_mlm_inputs = batch_mlm_inputs.to(self.device) + outputs = self.OmniGenome(**batch_mlm_inputs)[0] + outputs = self.lm_head(outputs) + outputs = outputs.argmax(dim=-1) + all_outputs.append(outputs) + outputs = torch.cat(all_outputs, dim=0) + return outputs[:, 1:1 + len(structure)] + + +# Copied from transformers.models.esm.modeling_esm.EsmClassificationHead with Esm->OmniGenome +class OmniGenomeClassificationHead(nn.Module): + """Head for sentence-level classification tasks.""" + + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.out_proj = nn.Linear(config.hidden_size, config.num_labels) + + def forward(self, features, **kwargs): + x = features[:, 0, :] # take token (equiv. to [CLS]) + x = self.dropout(x) + x = self.dense(x) + x = torch.tanh(x) + x = self.dropout(x) + x = self.out_proj(x) + return x + + +def create_position_ids_from_input_ids( + input_ids, padding_idx, past_key_values_length=0 +): + """ + Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols + are ignored. This is modified from fairseq's `utils.make_positions`. + + Args: + x: torch.Tensor x: + + Returns: torch.Tensor + """ + # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA. + mask = input_ids.ne(padding_idx).int() + incremental_indices = ( + torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length + ) * mask + return incremental_indices.long() + padding_idx \ No newline at end of file