[canary] Test for CanaryTokenizer + refactoring

nemo/collections/common/tokenizers/canary_tokenizer.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 from functools import cached_property
+from pathlib import Path
 from typing import Dict
 
 from nemo.collections.common.tokenizers.aggregate_tokenizer import AggregateTokenizer
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
 
 __all__ = ['CanaryTokenizer']
 
@@ -32,6 +34,7 @@
 }
 
 SPECIAL_TOKENS = [
+    "<pad>",
     "<|endoftext|>",
     "<|startoftranscript|>",
     *[f"<|{lang}|>" for lang in list(LANGUAGES.keys())],
@@ -40,9 +43,10 @@
     "<|nopnc|>",
     "<|pnc|>",
     "<|nospeech|>",
-    "<pad>",
 ]
 
+UNUSED_SPECIAL_TOKENS = [f"<|spltoken{i}|>" for i in range(18)]
+
 
 class CanaryTokenizer(AggregateTokenizer):
     """
@@ -96,3 +100,22 @@ def to_language_id(self, language):
             return token_id
 
         raise KeyError(f"Language {language} not found in tokenizer.")
+
+    @staticmethod
+    def build_special_tokenizer(output_dir: str | Path) -> SentencePieceTokenizer:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(exist_ok=True, parents=True)
+        text_path = output_dir / "train_text.txt"
+        all_tokens = SPECIAL_TOKENS + UNUSED_SPECIAL_TOKENS
+        train_text = "\n".join(all_tokens)
+        text_path.write_text(train_text)
+        model_path = output_dir / "tokenizer.model"
+        create_spt_model(
+            str(text_path),
+            vocab_size=32,
+            sample_size=-1,
+            do_lower_case=False,
+            output_dir=str(output_dir),
+            user_defined_symbols=all_tokens,
+        )
+        return SentencePieceTokenizer(str(model_path))

tests/collections/asr/test_custom_tokenizer.py

@@ -0,0 +1,65 @@
+from unittest.mock import Mock
+
+import pytest
+import sentencepiece as spm
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.parts.mixins import ASRBPEMixin
+from nemo.collections.common.tokenizers.canary_tokenizer import SPECIAL_TOKENS, UNUSED_SPECIAL_TOKENS, CanaryTokenizer
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+
+
+@pytest.fixture(scope="session")
+def special_tokenizer_path(tmp_path_factory) -> str:
+    tmpdir = tmp_path_factory.mktemp("spl_tokens")
+    CanaryTokenizer.build_special_tokenizer(tmpdir)
+    return str(tmpdir)
+
+
+@pytest.fixture(scope="session")
+def lang_tokenizer_path(tmp_path_factory) -> str:
+    tmpdir = tmp_path_factory.mktemp("klingon_tokens")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("a\nb\nc\nd\n")
+    create_spt_model(text_path, vocab_size=8, sample_size=-1, do_lower_case=False, output_dir=str(tmpdir))
+    return str(tmpdir)
+
+
+def test_canary_tokenizer_build_special_tokenizer(tmp_path):
+    tokenizer = CanaryTokenizer.build_special_tokenizer(tmp_path)
+    expected_tokens = ["<unk>"] + SPECIAL_TOKENS + UNUSED_SPECIAL_TOKENS + ["▁"]
+    tokens = []
+    for i in range(tokenizer.tokenizer.vocab_size()):
+        tokens.append(tokenizer.tokenizer.IdToPiece(i))
+    assert expected_tokens == tokens
+
+
+def test_canary_tokenizer_init_from_cfg(special_tokenizer_path, lang_tokenizer_path):
+    bpe_mixin = ASRBPEMixin()
+    bpe_mixin.register_artifact = Mock(side_effect=lambda self, x: x)
+    config = OmegaConf.create(
+        {
+            "type": "agg",
+            "is_canary": True,
+            "dir": None,
+            "langs": {
+                "spl_tokens": {"dir": special_tokenizer_path, "type": "bpe",},
+                "en": {"dir": lang_tokenizer_path, "type": "bpe",},
+            },
+        }
+    )
+    bpe_mixin._setup_aggregate_tokenizer(config)
+    tokenizer = bpe_mixin.tokenizer
+
+    assert isinstance(tokenizer, CanaryTokenizer)
+    assert len(tokenizer.tokenizers_dict) == 2
+    assert set(tokenizer.tokenizers_dict.keys()) == {"spl_tokens", "en"}
+
+    assert isinstance(tokenizer.tokenizers_dict["spl_tokens"], SentencePieceTokenizer)
+    assert tokenizer.tokenizers_dict["spl_tokens"].vocab_size == 32
+
+    assert isinstance(tokenizer.tokenizers_dict["en"], SentencePieceTokenizer)
+    assert tokenizer.tokenizers_dict["en"].vocab_size == 6
+
+    assert tokenizer.text_to_ids("<|startoftranscript|>", lang_id="spl_tokens") == [31, 3]  # "_" comes first
+    assert tokenizer.text_to_ids("a", lang_id="en") == [32 + 1, 32 + 2]