Flatten esm-2 model package (#1464)

As we're adding the ESM-C model, we're going to need to install the
`esm` package in the devcontainer for running esm-c model tests. This
will conflict with our `esm` pseudo-package I created for the original
ESM-2 model folder.

This PR just flattens the model into a single requirements.txt folder
like llama3 and the other, newer models

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added comprehensive export workflow for converting HuggingFace ESM-2
checkpoints to Transformer Engine format, including parameter count
formatting and verification.

* **Chores**
* Reorganized project structure and dependencies from pyproject.toml to
modular configuration files.
* Updated import paths and development setup instructions throughout
codebase.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

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
 

…-recipes/models/esm2/src/esm/collator.py bionemo-recipes/models/esm2/collator.pybionemo-recipes/models/esm2/src/esm/collator.py renamed to bionemo-recipes/models/esm2/collator.py bionemo-recipes/models/esm2/src/esm/collator.py renamed to bionemo-recipes/models/esm2/collator.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 = {

…o-recipes/models/esm2/src/esm/convert.py bionemo-recipes/models/esm2/convert.pybionemo-recipes/models/esm2/src/esm/convert.py renamed to bionemo-recipes/models/esm2/convert.py bionemo-recipes/models/esm2/src/esm/convert.py renamed to bionemo-recipes/models/esm2/convert.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/export.py

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