Fix geneformer tests

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

Author: jstjohn

sub-packages/bionemo-geneformer/examples/.gitignore

@@ -1 +1,2 @@
 **.png
+pretrain-recipe-short.yaml

sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/celltype_classification_bench/bench.py

@@ -153,8 +153,9 @@ def load_data_run_benchmark(result_path, adata_path, write_results=True):
     import torch
 
     adata = read_h5ad(adata_path)
-
-    infer_Xs = torch.load(result_path / "predictions__rank_0__dp_rank_0.pt")["embeddings"].float().cpu().numpy()
+    # TODO: update the prediction writer to support model and data parallelism, and modify the
+    #   path accordingly.
+    infer_Xs = torch.load(result_path / "predictions__rank_0.pt")["embeddings"].float().cpu().numpy()
     assert len(adata) == len(infer_Xs), (len(adata), len(infer_Xs))
 
     infer_metadata = adata.obs

sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_celltype_bench.py

@@ -88,7 +88,8 @@ def test_celltype_embeddings_golden_values():
         config_class=GeneformerConfig,
         include_unrecognized_vocab_in_dataset=False,
     )
-
+    # TODO: update the prediction writer to support model and data parallelism, and modify the
+    #   path accordingly.
     result = torch.load(results_path / "predictions__rank_0.pt")["embeddings"]
     expected_vals = torch.load(golden_values_path / "predictions__rank_0.pt")["embeddings"]