review fix

Signed-off-by: Ohad Mosafi <[email protected]>

Author: ohadmo

bionemo-recipes/models/esm2/README.md

@@ -16,7 +16,7 @@ The ESM-2 implementation natively supports the following TransformerEngine-provi
 | **Sequence Packing / THD input format** | ✅ Supported                                                                     |
 | **FP8 with THD input format**           | ✅ Supported where FP8 is supported                                              |
 | **Import from HuggingFace checkpoints** | ✅ Supported                                                                     |
-| **Export to HuggingFace checkpoints**   | ✅ Under development                                                             |
+| **Export to HuggingFace checkpoints**   | ✅ Supported                                                           |
 
 See [BioNemo Recipes](../../recipes/README.md) for more details on how to use these features to accelerate model
 training and inference.
@@ -137,7 +137,7 @@ from esm.export import export_te_checkpoint
 
 hf_export_path = Path("hf_export")
 exported_model_path = te_checkpoint_path / "esm2_t6_8M_UR50D"
-export_te_checkpoint(str(exported_model_path), str(hf_export_path))
+export_te_checkpoint(str(exported_model_path), hf_export_path)
 ```
 
 This step creates a new Hugging Face model that should be functionally equivalent to the original.

bionemo-recipes/models/esm2/src/esm/export.py

@@ -119,7 +119,7 @@ def export_hf_checkpoint(tag: str, export_path: Path):
     torch.cuda.empty_cache()
 
 
-def export_te_checkpoint(te_checkpoint_path: str, output_path: str):
+def export_te_checkpoint(te_checkpoint_path: str, output_path: Path):
     """Export a Transformer Engine checkpoint back to the original HuggingFace Facebook ESM-2 format.
 
     This function converts from the NVIDIA Transformer Engine (TE) format back to the
@@ -129,7 +129,7 @@ def export_te_checkpoint(te_checkpoint_path: str, output_path: str):
 
     Args:
         te_checkpoint_path (str): Path to the TE checkpoint
-        output_path (str): Output path for the converted Facebook ESM-2 format model
+        output_path (Path): Output path for the converted Facebook ESM-2 format model
     """
     if not Path(te_checkpoint_path).exists():
         raise FileNotFoundError(f"TE checkpoint {te_checkpoint_path} not found")
@@ -143,15 +143,15 @@ def export_te_checkpoint(te_checkpoint_path: str, output_path: str):
 
     tokenizer_config_path = Path(te_checkpoint_path) / "tokenizer_config.json"
     if tokenizer_config_path.exists():
-        shutil.copy(tokenizer_config_path, Path(output_path) / "tokenizer_config.json")
+        shutil.copy(tokenizer_config_path, output_path / "tokenizer_config.json")
 
     vocab_path = Path(te_checkpoint_path) / "vocab.txt"
     if vocab_path.exists():
-        shutil.copy(vocab_path, Path(output_path) / "vocab.txt")
+        shutil.copy(vocab_path, output_path / "vocab.txt")
 
     special_tokens_path = Path(te_checkpoint_path) / "special_tokens_map.json"
     if special_tokens_path.exists():
-        shutil.copy(special_tokens_path, Path(output_path) / "special_tokens_map.json")
+        shutil.copy(special_tokens_path, output_path / "special_tokens_map.json")
 
     model_hf = AutoModelForMaskedLM.from_pretrained(
         output_path,

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

@@ -87,7 +87,7 @@ def test_export_te_checkpoint_to_hf(model_name):
         te_model_path = te_checkpoint_path / model_name
 
         hf_export_path = temp_path / "hf_export"
-        export_te_checkpoint(str(te_model_path), str(hf_export_path))
+        export_te_checkpoint(str(te_model_path), hf_export_path)
 
         model_hf_exported = AutoModelForMaskedLM.from_pretrained(str(hf_export_path))