Fix issue with extra token in tokenizers and pad

Signed-off-by: John St John <[email protected]>

Author: jstjohn

bionemo-recipes/recipes/evo2_megatron/tests/bionemo/evo2/data/test_tokenizer.py

@@ -17,9 +17,17 @@
 # limitations under the License.
 
 
+from pathlib import Path
+
 import pytest
+from megatron.bridge.training.tokenizers.config import TokenizerConfig
+from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer
 
-from bionemo.evo2.data.dataset_tokenizer import Evo2DatasetTokenizer
+from bionemo.evo2.data.dataset_tokenizer import (
+    DEFAULT_HF_TOKENIZER_MODEL_PATH,
+    DEFAULT_HF_TOKENIZER_MODEL_PATH_512,
+    Evo2DatasetTokenizer,
+)
 from bionemo.evo2.utils.config import Evo2PreprocessingConfig
 
 
@@ -29,6 +37,25 @@ def tokenizer() -> Evo2DatasetTokenizer:
     return Evo2DatasetTokenizer(Evo2PreprocessingConfig())
 
 
+@pytest.mark.parametrize(
+    "tokenizer_path, expected_vocab_size",
+    [
+        (DEFAULT_HF_TOKENIZER_MODEL_PATH, 256),
+        (DEFAULT_HF_TOKENIZER_MODEL_PATH_512, 512),
+    ],
+)
+def test_tokenizer_vocab_size(tokenizer_path: Path, expected_vocab_size: int) -> None:
+    """Verifies key tokenizers have the expected vocabulary size."""
+    tokenizer = build_tokenizer(
+        TokenizerConfig(
+            tokenizer_type="HuggingFaceTokenizer",
+            hf_tokenizer_kwargs={"trust_remote_code": False},
+            tokenizer_model=tokenizer_path,
+        )
+    )
+    assert tokenizer.vocab_size == expected_vocab_size
+
+
 def test_tokenizer_handles_long_dna_sequence(tokenizer: Evo2DatasetTokenizer) -> None:
     """Verifies tokenizer correctly processes a long DNA sequence into expected token IDs.
 

bionemo-recipes/recipes/evo2_megatron/tokenizers/nucleotide_fast_tokenizer_256/tokenizer.json

@@ -143,8 +143,8 @@
       "<EOS>": 0,
       "<PAD>": 1,
       "<BOS>": 2,
-      "<UNK>": 3,
-      "\u0004": 4,
+      "<SEP>": 3,
+      "<UNK>": 4,
       "\u0005": 5,
       "\u0006": 6,
       "\u0007": 7,

bionemo-recipes/recipes/evo2_megatron/tokenizers/nucleotide_fast_tokenizer_512/tokenizer.json

@@ -116,8 +116,8 @@
       "<EOS>": 0,
       "<PAD>": 1,
       "<BOS>": 2,
-      "<UNK>": 3,
-      "\u0004": 4,
+      "<SEP>": 3,
+      "<UNK>": 4,
       "\u0005": 5,
       "\u0006": 6,
       "\u0007": 7,