Change SaAMPLIFY API to be compatible with AMPLIFY v1

src --> input_ids
pad_mask --> attention_mask

amplify.py

@@ -16,7 +16,6 @@ from .tokenizer import ProteinTokenizer
 from transformers import PreTrainedModel, PretrainedConfig
 from transformers.modeling_outputs import MaskedLMOutput
 
-
 class DotDict(dict):
     """Dictionary that supports the dot notation to access attributes (similarly to HuggingFace)."""
 
@@ -24,10 +23,8 @@ class DotDict(dict):
     __setattr__ = dict.__setitem__
     __delattr__ = dict.__delitem__
 
-
 class AMPLIFYConfig(PretrainedConfig):
     model_type = "AMPLIFY"
-
     # All config parameters must have a default value.
     def __init__(
         self,
@@ -51,7 +48,7 @@ class AMPLIFYConfig(PretrainedConfig):
         **kwargs,
     ):
         super().__init__(**kwargs)
-
+        
         self.hidden_size = hidden_size
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
@@ -69,7 +66,7 @@ class AMPLIFYConfig(PretrainedConfig):
         self.att_bias = att_bias
         self.pad_token_id = pad_token_id
         self.max_length = max_length
-
+        
 
 class EncoderBlock(nn.Module):
     """Transformer encoder block."""
@@ -111,7 +108,12 @@ class EncoderBlock(nn.Module):
             multiple_of = 8
             intermediate_size = int(2 * config.intermediate_size / 3)
             intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
-            self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=config.ffn_bias)
+            self.ffn = SwiGLU(
+                config.hidden_size,
+                intermediate_size,
+                config.hidden_size,
+                bias=config.ffn_bias
+            )
         elif act == "relu":
             self.ffn = nn.Sequential(
                 nn.Linear(config.hidden_size, config.intermediate_size, bias=config.ffn_bias),
@@ -127,26 +129,18 @@ class EncoderBlock(nn.Module):
         else:
             raise ValueError(f"Unsupported hidden_act: {config.hidden_act}")
 
-        self.attention_norm = (
-            RMSNorm(config.hidden_size, config.norm_eps)
-            if config.rms_norm
-            else nn.LayerNorm(config.hidden_size, config.norm_eps)
-        )
-        self.ffn_norm = (
-            RMSNorm(config.hidden_size, config.norm_eps)
-            if config.rms_norm
-            else nn.LayerNorm(config.hidden_size, config.norm_eps)
-        )
+        self.attention_norm = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
+        self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
 
         self.ffn_dropout = nn.Dropout(config.dropout_prob)
 
-    def forward(self, x: torch.Tensor, pad_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
-        attn, contact = self._att_block(self.attention_norm(x), pad_mask, freqs_cis, output_attentions)
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
+        attn, contact = self._att_block(self.attention_norm(x), attention_mask, freqs_cis, output_attentions)
         x = x + attn
         x = x + self._ff_block(self.ffn_norm(x))
         return x, contact
 
-    def _att_block(self, x: torch.Tensor, pad_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
+    def _att_block(self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
         batch_size, seq_len, _ = x.shape
         xq, xk, xv = self.q(x), self.k(x), self.v(x)
 
@@ -160,8 +154,8 @@ class EncoderBlock(nn.Module):
         attn_weights = None
         if output_attentions:
             attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
-            if pad_mask is not None:
-                attn_weights = attn_weights + pad_mask
+            if attention_mask is not None:
+                attn_weights = attn_weights + attention_mask
             attn_weights = attn_weights.softmax(-1)
 
         # Compute the attention using xformers if the tensors are on GPU
@@ -172,7 +166,7 @@ class EncoderBlock(nn.Module):
                 query=xq,
                 key=xk,
                 value=xv,
-                attn_bias=pad_mask,
+                attn_bias=attention_mask,
                 p=self.config.dropout_prob if self.training else 0,
             )
         else:
@@ -181,13 +175,13 @@ class EncoderBlock(nn.Module):
                 query=xq.transpose(1, 2),
                 key=xk.transpose(1, 2),
                 value=xv.transpose(1, 2),
-                attn_mask=pad_mask,
+                attn_mask=attention_mask,
                 dropout_p=self.config.dropout_prob if self.training else 0,
             ).transpose(1, 2)
 
         attn_scores = self.wo(attn.reshape(batch_size, seq_len, self.config.num_attention_heads * self.d_head))
         return (self.resid_dropout(attn_scores), attn_weights)
-
+    
     def _ff_block(self, x: torch.Tensor):
         return self.ffn_dropout(self.ffn(x))
 
@@ -207,10 +201,9 @@ class AMPLIFYPreTrainedModel(PreTrainedModel):
 class AMPLIFY(AMPLIFYPreTrainedModel):
     """The main model class.
 
-    Args:
-       config (amplify.model.amplify.AMPLIFYConfig): model configuration, usually defined from the Hydra configuration.
+       Args:
+          config (amplify.model.amplify.AMPLIFYConfig): model configuration, usually defined from the Hydra configuration.
     """
-
     def __init__(self, config: AMPLIFYConfig, **kwargs):
         super().__init__(config)
 
@@ -219,30 +212,23 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
 
         if config.layer_norm_after_embedding:
-            self.layer_norm_1 = (
-                RMSNorm(config.hidden_size, config.norm_eps)
-                if config.rms_norm
-                else nn.LayerNorm(config.hidden_size, config.norm_eps)
-            )
+            self.layer_norm_1 = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
 
         self.transformer_encoder = nn.ModuleList()
         for _ in range(config.num_hidden_layers):
             self.transformer_encoder.append(EncoderBlock(config))
 
         if config.layer_norm_before_last_layer:
-            self.layer_norm_2 = (
-                RMSNorm(config.hidden_size, config.norm_eps)
-                if config.rms_norm
-                else nn.LayerNorm(config.hidden_size, config.norm_eps)
-            )
+            self.layer_norm_2 = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
 
         self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
 
         self.freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
-
+        
         # Initialize weights and apply final processing
         self.post_init()
 
+
     @classmethod
     def load(cls, checkpoint_path: str, config_path: str):
 
@@ -254,7 +240,7 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         if checkpoint_path.endswith(".safetensors"):
             state_dict = safetensors.torch.load_file(checkpoint_path)
         elif checkpoint_path.endswith(".pt"):
-            state_dict = torch.load(checkpoint_path)
+            state_dict = torch.load(checkpoint_path, map_location="cpu")
         else:
             raise ValueError(f"Expected checkpoint to be a `.pt` or `.safetensors` file.")
 
@@ -262,29 +248,28 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         tokenizer = ProteinTokenizer(**cfg["tokenizer"])
         return model, tokenizer
 
-    def forward(self, src, pad_mask=None, output_hidden_states=False, output_attentions=False):
+
+    def forward(self, input_ids, attention_mask=None, output_hidden_states=False, output_attentions=False):
         # Initialize
         hidden_states, attentions = [], []
 
         # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
-        if pad_mask is not None:
-            assert pad_mask.dtype != torch.bool and 1.0 not in pad_mask, "AMPLIFY expects an additive pad_mask"
-            pad_mask = (
-                pad_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, pad_mask.size(-1), 1)
-            )
+        if attention_mask is not None:
+            assert attention_mask.dtype != torch.bool and 1.0 not in attention_mask, "AMPLIFY expects an additive attention_mask"
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
 
         # RoPE
-        self.freqs_cis = self.freqs_cis.to(src.device, non_blocking=True)
-        freqs_cis = self.freqs_cis[: src.shape[1]]
+        self.freqs_cis = self.freqs_cis.to(input_ids.device, non_blocking=True)
+        freqs_cis = self.freqs_cis[: input_ids.shape[1]]
 
         # Embedding
-        x = self.encoder(src)
+        x = self.encoder(input_ids)
         if self.config.layer_norm_after_embedding:
             x = self.layer_norm_1(x)
 
         # Transformer encoder
         for layer in self.transformer_encoder:
-            x, attn = layer(x, pad_mask, freqs_cis, output_attentions)
+            x, attn = layer(x, attention_mask, freqs_cis, output_attentions)
             if output_hidden_states:
                 hidden_states.append(x)
             if output_attentions:
@@ -295,3 +280,4 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
 
         # Return logits or the output of the last hidden layer
         return MaskedLMOutput(logits=logits, hidden_states=hidden_states, attentions=attentions)
+