Flatten esm-2 model package

bionemo-recipes/models/esm2/.ruff.toml

@@ -0,0 +1 @@
+extend = "../.ruff.toml"

bionemo-recipes/models/esm2/Dockerfile

@@ -2,4 +2,4 @@ FROM nvcr.io/nvidia/pytorch:26.01-py3
 WORKDIR /workspace/bionemo
 COPY . .
 RUN --mount=type=cache,target=/root/.cache/pip \
-    PIP_CONSTRAINT= pip install -e .
+    PIP_CONSTRAINT= pip install -r requirements.txt

bionemo-recipes/models/esm2/README.md

@@ -80,7 +80,7 @@ Hugging Face Transformers format for sharing and deployment. The workflow involv
 ```python
 from transformers import AutoModelForMaskedLM
 
-from esm.convert import convert_esm_hf_to_te
+from convert import convert_esm_hf_to_te
 
 hf_model = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
 te_model = convert_esm_hf_to_te(hf_model)
@@ -92,8 +92,8 @@ This loads the pre-trained ESM2 model that will serve as our reference for compa
 ### Converting from TE back to HF Transformers
 
 ```python
-from esm.convert import convert_esm_te_to_hf
-from esm.modeling_esm_te import NVEsmForMaskedLM
+from convert import convert_esm_te_to_hf
+from modeling_esm_te import NVEsmForMaskedLM
 
 te_model = NVEsmForMaskedLM.from_pretrained("/path/to/te_checkpoint")
 hf_model = convert_esm_te_to_hf(te_model)
@@ -130,8 +130,8 @@ To run tests locally, run `recipes_local_test.py` from the repository root with
 ### Development container
 
 To use the provided devcontainer, use "Dev Containers: Reopen in Container" from the VSCode menu, and choose the
-"BioNeMo Recipes Dev Container" option. To run the tests inside the container, first install the model package in
-editable mode with `pip install -e .`, then run `pytest -v .` in the model directory.
+"BioNeMo Recipes Dev Container" option. To run the tests inside the container, first install the dependencies with
+`pip install -r requirements.txt`, then run `pytest -v .` in the model directory.
 
 ### Deploying converted checkpoints to HuggingFace Hub
 

bionemo-recipes/models/esm2/src/esm/convert.py → bionemo-recipes/models/esm2/convert.py

@@ -18,8 +18,8 @@
 from torch import nn
 from transformers import EsmConfig, EsmForMaskedLM
 
-import esm.state as io
-from esm.modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
+import state as io
+from modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 
 
 mapping = {

bionemo-recipes/models/esm2/export.py

@@ -13,9 +13,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
+import json
+import shutil
 from pathlib import Path
 
-from esm.export import export_hf_checkpoint
+import torch
+from jinja2 import Template
+from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer
+
+from convert import convert_esm_hf_to_te
+from modeling_esm_te import AUTO_MAP
+
+
+BENCHMARK_RESULTS = {
+    "esm2_t6_8M_UR50D": {"CAMEO": 0.48, "CASP14": 0.37},
+    "esm2_t12_35M_UR50D": {"CAMEO": 0.56, "CASP14": 0.41},
+    "esm2_t30_150M_UR50D": {"CAMEO": 0.65, "CASP14": 0.49},
+    "esm2_t33_650M_UR50D": {"CAMEO": 0.70, "CASP14": 0.51},
+    "esm2_t36_3B_UR50D": {"CAMEO": 0.72, "CASP14": 0.52},
+    "esm2_t48_15B_UR50D": {"CAMEO": 0.72, "CASP14": 0.55},
+}
 
 
 ESM_TAGS = [
@@ -28,6 +46,85 @@
 ]
 
 
+def format_parameter_count(num_params: int, sig: int = 1) -> str:
+    """Format parameter count in scientific notation (e.g., 6.5 x 10^8).
+
+    Args:
+        num_params: Total number of parameters
+        sig: Number of digits to include after the decimal point
+
+    Returns:
+        Formatted string in scientific notation
+    """
+    s = f"{num_params:.{sig}e}"
+    base, exp = s.split("e")
+    return f"{base} x 10^{int(exp)}"
+
+
+def export_hf_checkpoint(tag: str, export_path: Path):
+    """Export a Hugging Face checkpoint to a Transformer Engine checkpoint.
+
+    Args:
+        tag: The tag of the checkpoint to export.
+        export_path: The parent path to export the checkpoint to.
+    """
+    model_hf_masked_lm = AutoModelForMaskedLM.from_pretrained(f"facebook/{tag}")
+    model_hf = AutoModel.from_pretrained(f"facebook/{tag}")
+    model_hf_masked_lm.esm.pooler = model_hf.pooler
+    model_te = convert_esm_hf_to_te(model_hf_masked_lm)
+    model_te.save_pretrained(export_path / tag)
+
+    tokenizer = AutoTokenizer.from_pretrained("esm_fast_tokenizer")  # Use our PreTrainedTokenizerFast implementation.
+    tokenizer.save_pretrained(export_path / tag)
+
+    # Patch the config
+    with open(export_path / tag / "config.json", "r") as f:
+        config = json.load(f)
+
+    config["auto_map"] = AUTO_MAP
+
+    with open(export_path / tag / "config.json", "w") as f:
+        json.dump(config, f, indent=2, sort_keys=True)
+
+    shutil.copy("modeling_esm_te.py", export_path / tag / "esm_nv.py")
+
+    # Calculate model parameters and render README template
+    num_params = sum(p.numel() for p in model_te.parameters())
+    formatted_params = format_parameter_count(num_params)
+
+    # Read and render the template
+    with open("model_readme.template", "r", encoding="utf-8") as f:
+        template_content = f.read()
+
+    template = Template(template_content)
+    rendered_readme = template.render(
+        num_params=formatted_params,
+        model_tag=tag,
+        cameo_score=BENCHMARK_RESULTS[tag]["CAMEO"],
+        casp14_score=BENCHMARK_RESULTS[tag]["CASP14"],
+    )
+
+    # Write the rendered README
+    with open(export_path / tag / "README.md", "w") as f:
+        f.write(rendered_readme)
+
+    shutil.copy("LICENSE", export_path / tag / "LICENSE")
+
+    del model_hf, model_te, model_hf_masked_lm
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # Smoke test that the model can be loaded.
+    model_te = AutoModelForMaskedLM.from_pretrained(
+        export_path / tag,
+        dtype=torch.bfloat16,
+        trust_remote_code=True,
+    )
+    del model_te
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
 def main():
     """Export the ESM2 models from Hugging Face to the Transformer Engine format."""
     # TODO (peter): maybe add a way to specify the model to export or option to export all models?

bionemo-recipes/models/esm2/requirements.txt

@@ -0,0 +1,11 @@
+accelerate
+datasets
+hydra-core
+jinja2
+megatron-fsdp
+omegaconf
+peft
+torch
+torchao!=0.14.0
+transformer_engine[pytorch]
+transformers

bionemo-recipes/models/esm2/tests/conftest.py

@@ -25,12 +25,11 @@
 from transformer_engine.pytorch import fp8
 from transformers import AutoModelForMaskedLM, AutoTokenizer, DataCollatorForLanguageModeling
 
-from esm.convert import convert_esm_hf_to_te
-
 
 sys.path.append(Path(__file__).parent.parent.as_posix())
 sys.path.append(Path(__file__).parent.as_posix())
 
+
 pytest_plugins = ["tests.common.fixtures"]
 
 
@@ -105,6 +104,8 @@ def input_data(tokenizer, tokenized_proteins):
 
 @pytest.fixture
 def te_model_checkpoint(tmp_path):
+    from convert import convert_esm_hf_to_te
+
     model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_te.save_pretrained(tmp_path / "te_model_checkpoint")

bionemo-recipes/models/esm2/tests/test_collator.py

@@ -20,7 +20,7 @@
 import torch
 from transformers import DataCollatorForLanguageModeling
 
-from esm.collator import (
+from collator import (
     DataCollatorWithFlattening,
     TokenPackingDataset,
     _split_sample_by_num_tokens,