Add Esm #2244

pass-lin · 2025-05-03T07:40:33Z

from #2177
Achieved a smaller error with hf.

import os
os.environ["KERAS_BACKEND"] = "torch"
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from keras import ops
from transformers.models.esm.modeling_esm import EsmAttention as hf_EsmSelfAttention
from transformers import EsmConfig
from esm2.esm2_layers import EsmSelfAttention
import numpy as np
import keras
from transformers.models.esm.modeling_esm import EsmModel
weights_path = "facebook/esm2_t6_8M_UR50D"
hf_model = EsmModel.from_pretrained(weights_path)
hf_model.cuda().eval()
hf_model.embeddings.token_dropout = False


from keras_hub.src.models.esm.esm_backbone import (
    ESMBackbone,
)


keras_model =  ESMBackbone.from_preset('hf://'+weights_path)
keras_model.summary()


x = ops.array([[1,2,3,4,5]])+1
hf_out = hf_model(x,ops.ones_like(x))[0]
keras_out = keras_model({'token_ids': x})

print(ops.all(ops.isclose(hf_out, keras_out,atol=1e-4)))

ESM Checkpoint Conversion and Numerics Verification Demo (across multiple backends): Notebook Link

Train Demo: Notebook Link

pass-lin · 2025-05-03T07:56:28Z

ruff.....................................................................Passed
ruff-format..............................................................Passed
Error: Process completed with exit code 1.

Please help me figure out how to solve this problem.

mattdangerw · 2025-05-06T18:35:36Z

Probably an issue with generating the API symbols. Looks like you need to sync with the latest changes on master, then you could try running ./shell/api_gen.sh

sachinprasadhs · 2025-05-09T17:15:23Z

ruff.....................................................................Passed
ruff-format..............................................................Passed
Error: Process completed with exit code 1.

Please help me figure out how to solve this problem.

You can rebase it to latest master code
and then run - pre-commit run --all-files
pip install -u namex

pass-lin · 2025-05-10T13:13:18Z

keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_dtype_argument_tie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_dtype_argument_untie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_int8_tie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/layers/modeling/reversible_embedding_test.py::ReversibleEmbeddingTest::test_quantize_int8_untie_weights - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/albert/albert_backbone_test.py::AlbertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bart/bart_backbone_test.py::BartBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bert/bert_backbone_test.py::BertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/bloom/bloom_backbone_test.py::BloomBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/clip/clip_backbone_test.py::CLIPBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/deberta_v3/deberta_v3_backbone_test.py::DebertaV3BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/distil_bert/distil_bert_backbone_test.py::DistilBertBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/electra/electra_backbone_test.py::ElectraBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/f_net/f_net_backbone_test.py::FNetBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/falcon/falcon_backbone_test.py::FalconBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gemma/gemma_backbone_test.py::GemmaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gemma/gemma_backbone_test.py::Gemma2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gpt2/gpt2_backbone_test.py::GPT2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/gpt_neo_x/gpt_neo_x_backbone_test.py::GPTNeoXBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/llama/llama_backbone_test.py::LlamaTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/mistral/mistral_backbone_test.py::MistralBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/opt/opt_backbone_test.py::OPTBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/pali_gemma/pali_gemma_backbone_test.py::PaliGemmaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/pali_gemma/pali_gemma_backbone_test.py::PaliGemma2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/phi3/phi3_backbone_test.py::Phi3Test::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/phi3/phi3_backbone_test.py::Phi3Test::test_backbone_basics_with_su_rotary - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/roberta/roberta_backbone_test.py::RobertaBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/siglip/siglip_backbone_test.py::SigLIPBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/siglip/siglip_backbone_test.py::SigLIP2BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/t5/t5_backbone_test.py::T5BackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/whisper/whisper_backbone_test.py::WhisperBackboneTest::test_backbone_basics - TypeError: _int8_build() takes 2 positional arguments but 3 were given
FAILED keras_hub/src/models/xlm_roberta/xlm_roberta_backbone_test.py

@mattdangerw @sachinprasadhs
Is it a problem with the test environment? Why are there so many errors that don't belong to me?

sachinprasadhs · 2025-05-12T17:44:15Z

It's not related to your code, looks like some issue with the JAX backend, we will look into it.

sachinprasadhs

Thanks fro the PR, I have added my comments, also add checkpoints conversion under: keras-hub/tools/checkpoint_conversion

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-05-15T22:02:55Z

keras_hub/src/models/esm/esm_backbone.py

+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        layer_norm_eps:bool.Should we use ln after embedding?


Didn't get the point here, are you asking our input or it's the arg detail, if it is the arg details, it needs to be repharsed, avoid question marks and the argument name is emb_layer_norm_before

layer_norm_eps discription needs to be updated.

keras_hub/src/models/esm/esm_backbone.py

keras_hub/src/models/esm/esm_classifier_test.py

keras_hub/src/models/esm/esm_masked_plm_test.py

keras_hub/src/models/esm/esm_masked_plm.py

keras_hub/src/models/esm/esm_classifier.py

keras_hub/src/utils/transformers/convert_esm.py

pass-lin · 2025-05-17T18:13:05Z

@sachinprasadhs @mattdangerw
Can anybody review my code?

pass-lin · 2025-06-02T18:11:23Z

@mattdangerw @sachinprasadhs
Please check my code, thank you.

sachinprasadhs

Added few more comments and few of the previous review comments still needs to be addressed

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-06-02T18:37:34Z

keras_hub/src/models/esm/esm_backbone.py

+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:


Still activation and max_wavelength description is missing!

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-06-02T18:42:21Z

keras_hub/src/models/esm/esm_backbone.py

+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:


add arg description for pad_token_id as well

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-06-02T18:46:33Z

keras_hub/src/models/esm/esm_backbone.py

+        position_embedding_type:esm1 use abs position embeding,esm2 use rope.
+            so this parameter is only except for absolute and rotary.


This still needs to be changed to:

position_embedding_type: str. The position embedding type to use. One of "absolute" and "rotary". Use "absolute" for ESM1. Use "rotary" for ESM2. Defaults to "rotary".

keras_hub/src/models/esm/esm_backbone_test.py

sachinprasadhs · 2025-06-02T18:50:08Z

keras_hub/src/models/esm/esm_classifier_preprocessor.py

+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifierPreprocessor")
+class ESMProteinClassifierPreprocessor(BertTextClassifierPreprocessor):


Pending change here which should be subclassed from TextClassifierPreprocessor instead of BertTextClassifierPreprocessor

sachinprasadhs · 2025-06-02T18:55:23Z

keras_hub/src/models/esm/esm_backbone.py

+        max_sequence_length=1024,
+        max_wavelength=10000,
+        layer_norm_eps=1e-12,
+        emb_layer_norm_before=False,


pending change, instead emb_layer_norm_before --> use_pre_layer_norm

sachinprasadhs · 2025-06-02T18:56:17Z

keras_hub/src/models/esm/esm_classifier.py

+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifier")
+class ESMProteinClassifier(RobertaTextClassifier):


pending change.
You can subclass TextClassifier and make the same changes as RobertaTextClassifier instead of subclassing from another model.

sachinprasadhs · 2025-06-02T19:13:43Z

Once you address all the comments, add end to end working colab along with the checkpoints conversion under: keras-hub/tools/checkpoint_conversion

pass-lin · 2025-06-03T10:53:20Z

Once you address all the comments, add end to end working colab along with the checkpoints conversion under: keras-hub/tools/checkpoint_conversion

Ok, please check the new code.

sachinprasadhs

Thanks, few minor comments.

Also, need more details specific to Keras 3.6 older version issue.

Finally in the PR description, add the colab notebook to show end to end working of the model, numerics verification.
you can follow the PR description template from the recent PR.

sachinprasadhs · 2025-06-05T21:07:01Z

keras_hub/src/models/esm/esm_backbone.py

+        position_embedding_type: str. The position embedding type to use.
+            One of "absolute" and "rotary".
+            Use "absolute" for ESM1. Use "rotary" for ESM2. Defaults to "rotary"
+        max_wavelength : The maximum angular wavelength of


change to --> max_wavelength : int. The maximum ...

sachinprasadhs · 2025-06-05T21:12:42Z

keras_hub/src/models/esm/esm_backbone.py

+            the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
+        activation :string or keras.activations. The activation to
+            use for the transformer. Defaults to `"gelu"`.
+        pad_token_id: int .padding token id,at esm2 it's 1.default 0.


Does it accepts only 1 for ESM2 when you say at esm2 it's 1
Also add clear description like: The index of the padding token in the vocabulary

keras_hub/src/models/esm/esm_backbone.py

sachinprasadhs · 2025-06-05T21:31:58Z

keras_hub/src/models/esm/esm_classifier.py

+
+    # Pretrained classifier.
+    classifier = keras_hub.models.ESMProteinClassifier.from_preset(
+        "roformer_v2_base_zh",


chnage preset name.

更改预设名称。

Keras doesn't add presets, so I can only write random ones

Write the preset name, like you're mapping in convert_esm_checkpoints.py

Write the preset name, like you're mapping in convert_esm_checkpoints.py

OK，i have renewed preset name

sachinprasadhs · 2025-06-05T21:34:29Z

keras_hub/src/models/esm/esm_classifier_test.py

+        if version.parse(keras.__version__) < version.parse("3.6"):
+            self.skipTest("Failing on keras lower version")


Why does it fail in Keras 3.5? and does it fail for all backends?

Because ops. dot_product_attention only available above 3.6.

keras_hub/src/models/esm/esm_classifier_preprocessor.py

keras_hub/src/models/esm/esm_classifier.py

pass-lin · 2025-06-11T05:00:14Z

Thanks, few minor comments.

Also, need more details specific to Keras 3.6 older version issue.

Finally in the PR description, add the colab notebook to show end to end working of the model, numerics verification. you can follow the PR description template from the recent PR.

How to add a Colab notebook? Can you give me give a demo?

sachinprasadhs · 2025-06-11T18:14:09Z

Thanks, few minor comments.
Also, need more details specific to Keras 3.6 older version issue.
Finally in the PR description, add the colab notebook to show end to end working of the model, numerics verification. you can follow the PR description template from the recent PR.

How to add a Colab notebook? Can you give me give a demo?

Adding from one of the recent PR which got merged, you can do something like this

DeiT Checkpoint Conversion and Numerics Verification Demo (across multiple backends): Notebook Link
DeiT End-to-End Demo (zero-shot and finetuning): Notebook Link
Here are the converted DeiT presets from Hugging Face checkpoints for reference.

pass-lin · 2025-06-12T07:20:20Z

Thanks, few minor comments.
Also, need more details specific to Keras 3.6 older version issue.
Finally in the PR description, add the colab notebook to show end to end working of the model, numerics verification. you can follow the PR description template from the recent PR.

How to add a Colab notebook? Can you give me give a demo?

Adding from one of the recent PR which got merged, you can do something like this

DeiT Checkpoint Conversion and Numerics Verification Demo (across multiple backends): Notebook Link

DeiT End-to-End Demo (zero-shot and finetuning): Notebook Link

Here are the converted DeiT presets from Hugging Face checkpoints for reference.

Hello, I've already added the Colab demo of tools/checkpoint_conversion/convert_esm_checkpoints.py in the PR description. I think this is enough, and we can refer to BERT for the rest.
Can we merge now?

sachinprasadhs · 2025-06-12T17:58:24Z

We don't have access to view the notebook, can you make it public. Thanks

pass-lin · 2025-06-12T18:36:32Z

We don't have access to view the notebook, can you make it public. Thanks

OK,It has been enable sharing

sachinprasadhs · 2025-06-13T20:05:11Z

Hi, The intention of the notebook is to verify the correctness of the model including, backbone, tasks with the usage details and the expected outcome and to verify the numerics stablity after weights transfer to the Keras architecture, with wither forward pass.

pass-lin · 2025-06-14T05:37:10Z

Hi, The intention of the notebook is to verify the correctness of the model including, backbone, tasks with the usage details and the expected outcome and to verify the numerics stablity after weights transfer to the Keras architecture, with wither forward pass.

Okay, I've added another notebook, which is a demo for predicting the suitable pH of enzymes using ESM.

sachinprasadhs · 2025-06-16T21:20:01Z

You can remove the esm2_t6_8M directory, that will be generated using the conversion script you have provided and will be uploaded to Kaggle.

The notebook which you have provided doesn't have predict method,
take any sample suitable input and display the output with predict.

Also in your conversion script, you have mentioned atol=1e-3, what would be the error percentage when the atol=1e-04 and we need following things in your notebook

Numerics verification, load the original ESM model and do forward pass, and do the same forward pass to Keras-Hub ESM implementation and compare the numerics layer by layer to show if numerics are matching(preferably to the 1e-4 precision)
Demonstrating usage of proprocessor, Tokenizer and other functionalities of ESM

I have provided the reference notebooks, please refer those.

You can keep only ESM changes in this PR, you can create a new PR for roformer which also needs checkpoint conversion script, so that we can maintain the latest weight in Kaggle by generating the new weights with the script with any future changes to Keras Hub model specific.

pass-lin · 2025-06-17T06:24:17Z

You can remove the esm2_t6_8M directory, that will be generated using the conversion script you have provided and will be uploaded to Kaggle.

The notebook which you have provided doesn't have predict method, take any sample suitable input and display the output with predict.

Also in your conversion script, you have mentioned atol=1e-3, what would be the error percentage when the atol=1e-04 and we need following things in your notebook

Numerics verification, load the original ESM model and do forward pass, and do the same forward pass to Keras-Hub ESM implementation and compare the numerics layer by layer to show if numerics are matching(preferably to the 1e-4 precision)

Demonstrating usage of proprocessor, Tokenizer and other functionalities of ESM

I have provided the reference notebooks, please refer those.

You can keep only ESM changes in this PR, you can create a new PR for roformer which also needs checkpoint conversion script, so that we can maintain the latest weight in Kaggle by generating the new weights with the script with any future changes to Keras Hub model specific.

OK, I have modified the notebook, please check. In addition, roformerV2 does not need to convert scripts, it is a native keras model. I just modified the keras2 api

pass-lin added 2 commits May 3, 2025 01:28

add esm

f9ff098

add esm2

cc4123b

pass-lin added 4 commits May 3, 2025 17:35

fix

d3f598d

fix

737a147

format

140207b

fix test

cc9a11c

divyashreepathihalli requested a review from sachinprasadhs May 5, 2025 17:09

format

f8da784

pass-lin and others added 2 commits May 10, 2025 19:26

renew

72e9829

Merge branch 'keras-team:master' into esm

16bb9f2

pass-lin force-pushed the esm branch from a66ee78 to 19f4b1f Compare May 10, 2025 11:33

format

5cbf577

pass-lin force-pushed the esm branch from 19f4b1f to 5cbf577 Compare May 10, 2025 12:05

format

6e9f817

pass-lin mentioned this pull request May 10, 2025

_int8_build() bug from keras-nightly keras-team/keras#21272

Closed

sachinprasadhs reviewed May 16, 2025

View reviewed changes

pass-lin added 3 commits May 17, 2025 12:18

update

2815e9c

update

79e738c

update

20d5051

sachinprasadhs reviewed Jun 2, 2025

View reviewed changes

pass-lin added 2 commits June 3, 2025 17:59

update

fb18c98

add new tool

7609ab4

sachinprasadhs reviewed Jun 5, 2025

View reviewed changes

pass-lin added 3 commits June 6, 2025 16:25

update

e7b9b76

update

11f56fd

update

bd7380a

pass-lin added 3 commits June 12, 2025 14:44

fix

0f6261c

fix

95bb704

fix

c2a2787

fix tool

cc69d43

		position_embedding_type:esm1 use abs position embeding,esm2 use rope.
		so this parameter is only except for absolute and rotary.



		@keras_hub_export("keras_hub.models.ESMProteinClassifierPreprocessor")
		class ESMProteinClassifierPreprocessor(BertTextClassifierPreprocessor):



		@keras_hub_export("keras_hub.models.ESMProteinClassifier")
		class ESMProteinClassifier(RobertaTextClassifier):

		if version.parse(keras.__version__) < version.parse("3.6"):
		self.skipTest("Failing on keras lower version")

Add Esm #2244

Are you sure you want to change the base?

Add Esm #2244

Uh oh!

Conversation

pass-lin commented May 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pass-lin commented May 3, 2025

Uh oh!

mattdangerw commented May 6, 2025

Uh oh!

sachinprasadhs commented May 9, 2025

Uh oh!

pass-lin commented May 10, 2025

Uh oh!

sachinprasadhs commented May 12, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pass-lin commented May 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pass-lin commented Jun 2, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs commented Jun 2, 2025

Uh oh!

pass-lin commented Jun 3, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pass-lin commented May 3, 2025 •

edited

Loading

pass-lin commented May 17, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading

sachinprasadhs Jun 2, 2025 •

edited

Loading