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-07-22T02:20:15Z

Hi, Just one small comment, rest everything looks good, please replace head_size argument in both backbone and classifier example with a modified argument name which i believe num_heads?

Thanks for the reminder, I'll make changes now.

pass-lin · 2025-07-22T03:11:39Z

Thank you for the PR @pass-lin. Looks great already. Added a few NIT comments. Thanks for the train demo, do you have a demo that shows output for classifier and PLM with presets loaded?

Thanks for your review, relevant changes have been made.

keras_hub/src/models/esm/esm_backbone.py

pass-lin · 2025-07-22T15:38:22Z

I forgot to switch branches when I was developing other tasks, which led to this error. It has been fixed now @sachinprasadhs

pass-lin · 2025-07-25T02:09:14Z

@sachinprasadhs This test error should have nothing to do with me

sachinprasadhs · 2025-07-25T05:14:05Z

Yes, you can ignore the GPU Test failures

pass-lin · 2025-07-25T05:15:16Z

Yes, you can ignore the GPU Test failures

So, can we merge now?

divyashreepathihalli · 2025-07-25T20:13:17Z

The GPU tests are fixed, please rebase with master. Will run the GPU tests again

pass-lin · 2025-07-26T01:53:19Z

The GPU tests are fixed, please rebase with master. Will run the GPU tests again

It seems that there are no problems with the test

sachinprasadhs

LGTM

pass-lin · 2025-08-02T08:53:12Z

@divyashreepathihalli Can this pr merge now?

divyashreepathihalli · 2025-08-11T02:39:38Z

Thank you for the contribution!!

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

divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Jul 21, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 21, 2025

update doc and test

c6bee00

sachinprasadhs reviewed Jul 22, 2025

View reviewed changes

keras_hub/src/models/esm/esm_backbone.py Outdated Show resolved Hide resolved

pass-lin force-pushed the esm branch from 7cb83f2 to f5e6868 Compare July 22, 2025 15:04

fix commit error

59464f2

pass-lin force-pushed the esm branch from f5e6868 to 59464f2 Compare July 22, 2025 15:36

sachinprasadhs added the kokoro:force-run Runs Tests on GPU label Jul 22, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 22, 2025

Merge branch 'keras-team:master' into esm

241e7ae

divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Jul 25, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 25, 2025

divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Jul 25, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 25, 2025

Merge branch 'keras-team:master' into esm

b01eda3

sachinprasadhs added the kokoro:force-run Runs Tests on GPU label Jul 28, 2025

kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 28, 2025

sachinprasadhs approved these changes Jul 29, 2025

View reviewed changes

divyashreepathihalli merged commit 3f8d3b7 into keras-team:master Aug 11, 2025
10 checks passed

github-project-automation bot moved this from In Progress to Done in KerasHub Aug 11, 2025

		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):

Add Esm #2244

Add Esm #2244

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!

sachinprasadhs May 15, 2025

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!

sachinprasadhs Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sachinprasadhs Jun 2, 2025

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 Jul 22, 2025

Uh oh!

pass-lin commented Jul 22, 2025

Uh oh!

Uh oh!

pass-lin commented Jul 22, 2025

Uh oh!

pass-lin commented Jul 25, 2025

Uh oh!

sachinprasadhs commented Jul 25, 2025

Uh oh!

pass-lin commented Jul 25, 2025

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