Update modeling_omnigenome.py

modeling_omnigenome.py

@@ -15,8 +15,11 @@
 """ 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
@@ -300,6 +303,178 @@ class OmniGenomeEmbeddings(nn.Module):
         )
         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):
@@ -339,6 +514,14 @@ class OmniGenomeSelfAttention(nn.Module):
 
         self.is_decoder = config.is_decoder
 
+        # FlashAttention parameters
+        self.enable_flash_attn = getattr(config, "use_flash_attention", True)
+        if self.enable_flash_attn:
+            from flash_attn import flash_attn_func
+            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,
@@ -359,13 +542,9 @@ class OmniGenomeSelfAttention(nn.Module):
     ) -> 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
@@ -384,95 +563,75 @@ class OmniGenomeSelfAttention(nn.Module):
 
         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":
+        # 使用FlashAttention的条件判断
+        use_flash_attn = self.enable_flash_attn and self.position_embedding_type == "rotary"
+        if use_flash_attn:
+            # 应用旋转位置编码
             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
+            # 调整维度顺序为 [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
             )
-            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
-                )
+            # 恢复维度顺序 [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 attention_mask is not None:
-            # Apply the attention mask is (precomputed for all layers in OmniGenomeModel forward() function)
-            attention_scores = attention_scores + attention_mask
+            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)
 
-        # Normalize the attention scores to probabilities.
-        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+                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
 
-        # 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)
+            if attention_mask is not None:
+                attention_scores = attention_scores + attention_mask
 
-        # Mask heads if we want to
-        if head_mask is not None:
-            attention_probs = attention_probs * head_mask
+            attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+            attention_probs = self.dropout(attention_probs)
 
-        context_layer = torch.matmul(attention_probs, value_layer)
+            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,)
-        )
-
+        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):
@@ -530,6 +689,7 @@ class OmniGenomeAttention(nn.Module):
             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,
@@ -1053,6 +1213,7 @@ class OmniGenomeModel(OmniGenomePreTrainedModel):
             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,
@@ -1117,7 +1278,7 @@ class OmniGenomeForMaskedLM(OmniGenomePreTrainedModel):
 
         self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False)
         self.lm_head = OmniGenomeLMHead(config)
-        # self.init_weights()
+        self.init_weights()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -1237,7 +1398,7 @@ class OmniGenomeForSequenceClassification(OmniGenomePreTrainedModel):
         self.config = config
         self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False)
         self.classifier = OmniGenomeClassificationHead(config)
-        # self.init_weights()
+        self.init_weights()
 
     @add_start_docstrings_to_model_forward(
         OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length")
@@ -1279,8 +1440,8 @@ class OmniGenomeForSequenceClassification(OmniGenomePreTrainedModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        sequence_output = outputs[0]
-        logits = self.classifier(sequence_output)
+        last_hidden_state = outputs[0]
+        logits = self.classifier(last_hidden_state)
 
         loss = None
         if labels is not None:
@@ -1336,12 +1497,10 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.OmniGenome = OmniGenomeModel(config, add_pooling_layer=False)
-        self.lm_head = OmniGenomeLMHead(config)
         self.dense = torch.nn.Linear(config.hidden_size, config.hidden_size)
         self.classifier = torch.nn.Linear(self.config.hidden_size, self.num_labels)
-        self.activation = torch.nn.Tanh()
-        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
-        # self.init_weights()
+        self.softmax = nn.Softmax(dim=-1)
+        self.init_weights()
 
     @add_start_docstrings_to_model_forward(
         OmniGenome_INPUTS_DOCSTRING.format("batch_size, sequence_length")
@@ -1367,12 +1526,12 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
         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
         )
-        
-        mlm_outputs = self.OmniGenome(
+
+        outputs = self.OmniGenome(
             input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
@@ -1382,17 +1541,11 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        try:
-            last_hidden_state = mlm_outputs[0]
-            last_hidden_state = self.dense(last_hidden_state)
-        except:
-            last_hidden_state = mlm_outputs.hidden_states[-1]
-            last_hidden_state = self.dense(last_hidden_state)
 
+        last_hidden_state = outputs[0]
+        last_hidden_state = self.dense(last_hidden_state)
         logits = self.classifier(last_hidden_state)
-        logits = torch.softmax(logits, dim=-1)
-        logits = self.activation(logits)
-        logits = self.dropout(logits)
+        logits = self.softmax(logits)
 
         loss = None
         if labels is not None:
@@ -1400,14 +1553,14 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
 
         if not return_dict:
-            output = (logits,) + mlm_outputs[2:]
+            output = (logits,) + outputs[2:]
             return ((loss,) + output) if loss is not None else output
 
         return TokenClassifierOutput(
             loss=loss,
             logits=logits,
-            hidden_states=mlm_outputs.hidden_states,
-            attentions=mlm_outputs.attentions,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
         )
 
     @staticmethod
@@ -1433,15 +1586,26 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
 
         return structure
 
-    def predict_structure(
+    def predict_rna_structure(
             self,
-            input_ids: Optional[torch.LongTensor] = None,
-            attention_mask: Optional[torch.Tensor] = None,
+            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
         """
-        Predicts the secondary structure of a sequence given the logits and attention mask.
-        """
+        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)
@@ -1458,18 +1622,26 @@ class OmniGenomeForTokenClassification(OmniGenomePreTrainedModel):
 
 @add_start_docstrings(
     """
-    OmniGenome Model with a simple genetic algorithm based RNA design head on top.
+    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 OmniGenomeMaskedLMForRNADesign(OmniGenomePreTrainedModel):
+class OmniGenomeModelForSeq2SeqLM(OmniGenomePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        self.OmniGenome = OmniGenomeForMaskedLM(config)
+        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.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")
@@ -1495,43 +1667,199 @@ class OmniGenomeMaskedLMForRNADesign(OmniGenomePreTrainedModel):
         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,
-        )
-
-        sequence_output = outputs[0]
-
-        sequence_output = self.dropout(sequence_output)
-        logits = self.classifier(sequence_output)
+        raise NotImplementedError("This model is not designed for standard Seq2Seq tasks. "
+                                  "Use model.rna_sequence_design() for RNA sequences design instead.")
 
-        loss = None
-        if labels is not None:
-            loss_fct = CrossEntropyLoss()
-
-            labels = labels.to(logits.device)
-            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,
-        )
+    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", "<mask>"])
+                for _ in range(len(structure))
+            ]
+            masked_sequence_str = "".join(masked_sequence)
+            mlm_inputs.append(f"{masked_sequence_str}<eos>{''.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('<mask>', '$')))
+            ]
+            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("$", "<mask>")
+            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}<eos>{''.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('<mask>', '$')))
+            ]
+            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] + "<mask>" * len(parent2[pos:])
+            child2 = "<mask>" * len(parent1[:pos]) + parent2[pos:]
+            batch_crossover_inputs.append(f"{child1}<eos>{structure}")
+            batch_crossover_inputs.append(f"{child2}<eos>{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('<mask>', '$')))
+            ]
+            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