loading directly from hugginface

[pengzhangzhi]

Hi,
In the current codebase, we have to download the ckpt to the local and load it using the following method:

# download "{model}.safetensors" to the local 
# and load it like below
model = ESM2.from_pretrained("{model}.safetensors", device=0)

I wonder if we can directly load the ckpt from Hugginface?
such as

model = ESM2.from_pretrained("facebook/esm2_t30_150M_UR50D", device=0)

That way, it's more straightforward to replace existing codebase with a flash-attention version of esm2.

It seems doable to me bc eesm shares the same model architecture with ESM2 except for the use of flash attention?

Would love to hear ur thoughts!

[pengzhangzhi]

I'm working on it. I guess the work is converting of the names defined in esm-efficient to be what's defined in esm2, which is the standrad hugginface names? Let me know if you are interested! We can talk more about that!!

[MuhammedHasan]

Just so you know, pull requests are welcome. Please make sure any change passes the test cases. I renamed and created safetensors from the checkpoints because pickles are not reliable in my experience.

Given we have the weights in the huggingface https://huggingface.co/mhcelik/esm-efficient/tree/main, it need to fetched to implement something like:

model = ESM2.from_pretrained("esm-efficient/esm2_8M", device=0)

I plan to support this at some point, but pull requests are appreciated and make it sooner.

[MuhammedHasan]

Fixed by #9