Update modeling_fastesm.py

modeling_fastesm.py

@@ -233,10 +233,6 @@ class EsmSelfAttention(nn.Module):
         if self.position_embedding_type == "rotary":
             query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer)
 
-        # Ensure all tensors have the same dtype before calling scaled_dot_product_attention
-        #query_layer = query_layer.to(value_layer.dtype)
-        #key_layer = key_layer.to(value_layer.dtype)
-
         context_layer = F.scaled_dot_product_attention(
             query_layer,
             key_layer,
@@ -422,10 +418,7 @@ class FastEsmModel(FastEsmPreTrainedModel):
             # Expand to (batch_size, 1, seq_length, seq_length)
             extended_attention_mask = attention_mask[:, None, None, :].expand(
                 batch_size, 1, seq_length, seq_length
-            )
-            # Convert mask to float with 0.0 for positions to keep and -inf for masked positions
-            attention_mask = attention_mask.to(dtype=embedding_output.dtype)  # fp16 compatibility
-            attention_mask = (1.0 - attention_mask) * torch.finfo(embedding_output.dtype).min
+            ).bool()
         else:
             extended_attention_mask = None
 
@@ -608,13 +601,13 @@ if __name__ == "__main__":
     In Pytorch 2.5+ (and linux kernel), this implementation is very fast and uses less memory than the HF implementation.
     """
     import random
-    from transformers import EsmModel as TransformersEsmModel, EsmTokenizer
+    from transformers import EsmForMaskedLM as TransformersEsmModel, EsmTokenizer
 
     model_paths = [
         "facebook/esm2_t6_8M_UR50D",
         "facebook/esm2_t12_35M_UR50D",
-        "facebook/esm2_t30_150M_UR50D",
-        "facebook/esm2_t33_650M_UR50D",
+        #"facebook/esm2_t30_150M_UR50D",
+        #"facebook/esm2_t33_650M_UR50D",
     ]
     canonical_amino_acids = "ACDEFGHIKLMNPQRSTVWY"
     length = 64
@@ -630,17 +623,17 @@ if __name__ == "__main__":
         print(f"Testing {model_path}...")
         tokenizer = EsmTokenizer.from_pretrained(model_path)
         config = FastEsmConfig.from_pretrained(model_path)
-        fast_model = FastEsmModel(config).from_pretrained(model_path, torch_dtype=torch.float16).to(device)
-        model = TransformersEsmModel.from_pretrained(model_path, token_dropout=False, torch_dtype=torch.float16).to(device)
+        fast_model = FastEsmForMaskedLM(config).from_pretrained(model_path).to(device)
+        model = TransformersEsmModel.from_pretrained(model_path, token_dropout=False).to(device)
 
         counts = [0] * len(tolerances)
         for _ in range(seq_count):
             example_seq = generate_random_sequence(length)
             fast_tokens = tokenizer(example_seq, return_tensors="pt").input_ids.to(device)
-            fast_output = fast_model(fast_tokens).last_hidden_state.detach().cpu()
+            fast_output = fast_model(fast_tokens, output_hidden_states=True).hidden_states[-1].detach().cpu()
 
             model_tokens = tokenizer(example_seq, return_tensors="pt").input_ids.to(device)
-            model_output = model(model_tokens).last_hidden_state.detach().cpu()
+            model_output = model(model_tokens, output_hidden_states=True).hidden_states[-1].detach().cpu()
 
             for i, atol in enumerate(tolerances):
                 if torch.allclose(fast_output, model_output, atol=atol):