Upload modeling_fastesm.py with huggingface_hub

modeling_fastesm.py

@@ -4,9 +4,10 @@ from torch.nn import functional as F
 from torch.utils.data import Dataset, DataLoader
 from typing import Optional, Tuple, Union
 from einops import rearrange
+from dataclasses import dataclass
 from transformers import PreTrainedModel, PretrainedConfig, EsmTokenizer
 from transformers.modeling_outputs import (
-    MaskedLMOutput,
+    ModelOutput,
     BaseModelOutputWithPastAndCrossAttentions,
     BaseModelOutputWithPoolingAndCrossAttentions,
     SequenceClassifierOutput,
@@ -23,6 +24,15 @@ from transformers.models.esm.modeling_esm import (
 from tqdm.auto import tqdm
 
 
+@dataclass
+class EsmMaskedLMOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
 class FastEsmConfig(PretrainedConfig):
     model_type = "fast_esm"
     def __init__(
@@ -656,9 +666,7 @@ class FAST_ESM_ENCODER(FastEsmPreTrainedModel):
             Model outputs including hidden states and optionally attention weights
         """
         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
-        )
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 
         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")
@@ -739,9 +747,7 @@ class FastEsmModel(FastEsmPreTrainedModel):
             Model outputs including hidden states and optionally attention weights
         """
         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
-        )
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 
         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")
@@ -798,7 +804,7 @@ class FastEsmForMaskedLM(FastEsmPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None, # to play nice with HF adjacent packages
-    ) -> Union[Tuple, MaskedLMOutput]:
+    ) -> Union[Tuple, EsmMaskedLMOutput]:
         outputs = self.esm(
             input_ids,
             attention_mask=attention_mask,
@@ -815,7 +821,7 @@ class FastEsmForMaskedLM(FastEsmPreTrainedModel):
             labels = labels.to(prediction_scores.device)
             loss = self.loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
 
-        return MaskedLMOutput(
+        return EsmMaskedLMOutput(
             loss=loss,
             logits=prediction_scores,
             hidden_states=outputs.hidden_states,