tests starting to run

Signed-off-by: Peter St. John <[email protected]>
NVIDIA · Jan 23, 2025 · 507b35f · 507b35f
1 parent e13bbee
commit 507b35f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/sub-packages/bionemo-amplify/pyproject.toml b/sub-packages/bionemo-amplify/pyproject.toml
@@ -14,6 +14,7 @@ dynamic = ["version"]
 dependencies = [
     # internal
     'bionemo-core',
+    'bionemo-esm2',
     'bionemo-llm',
     # external
 ]

diff --git a/sub-packages/bionemo-amplify/tests/bionemo/amplify/test_tokenizer.py b/sub-packages/bionemo-amplify/tests/bionemo/amplify/test_tokenizer.py
@@ -30,6 +30,6 @@ def test_tokenizer_serialization(tokenizer, tmp_path):
     tokenizer.io_dump(tmp_path / "tokenizer", yaml_attrs=[])  # BioNeMoESMTokenizer takes no __init__ arguments
     deserialized_tokenizer = io.load(tmp_path / "tokenizer", tokenizer.__class__)
 
-    our_tokens = deserialized_tokenizer.encode("K A <mask> I S Q", add_special_tokens=False)
-    esm_tokens = torch.tensor([15, 5, 32, 12, 8, 16])
-    torch.testing.assert_close(torch.tensor(our_tokens), esm_tokens)
+    our_tokens = deserialized_tokenizer.encode("KA<mask>ISQ", add_special_tokens=False)
+    amplify_tokens = torch.tensor([17, 7, 2, 14, 10, 18])
+    torch.testing.assert_close(torch.tensor(our_tokens), amplify_tokens)