tests: add a lot of tests

Author: stephantul

model2vec/tokenizer/normalizer.py

@@ -26,9 +26,9 @@ def replace_normalizer(
     new_normalizers.append(Replace(Regex(r"\s+"), " "))
     new_normalizers.append(Strip(right=True))
     if normalizer is None:
-        normalizer = Sequence(new_normalizers)
+        normalizer = Sequence(new_normalizers)  # type: ignore
     else:
         normalizer = Sequence([normalizer] + new_normalizers)  # type: ignore
-    tokenizer.normalizer = normalizer
+    tokenizer.normalizer = normalizer  # type: ignore
 
     return tokenizer

model2vec/tokenizer/tokenizer.py

@@ -361,10 +361,13 @@ def create_tokenizer(
     token_remove_regex: re.Pattern | None = None,
 ) -> PreTrainedTokenizerFast:
     """
-    Create a tokenizer from a vocabulary.
+    Create a tokenizer by adding tokens to the vocabulary.
 
-    This function creates a tokenizer from a vocabulary and a tokenizer.
-    It also sets the normalizer and pre-tokenizer for the tokenizer.
+    This function turns any tokenizer into a supertoken tokenizer. It does the following:
+    1. Turns the tokenizer model into a unigram model.
+    2. Adds a new pretokenizer, splitting on punctuation.
+    3. Adds all tokens in vocabulary to the model.
+    4. Removes any internal tokens that conform to the regex.
 
     :param tokenizer: The tokenizer to use.
     :param vocabulary: The vocabulary to use.

tests/conftest.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import pytest
 import torch
 from tokenizers import Tokenizer
 from tokenizers.models import BPE, Unigram, WordPiece
 from tokenizers.pre_tokenizers import Whitespace
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizerFast
 
 from model2vec.inference import StaticModelPipeline
 from model2vec.train import StaticModelForClassification
@@ -25,7 +25,9 @@ def mock_tokenizer(request: pytest.FixtureRequest) -> Tokenizer:
     tokenizer_type = request.param
 
     if tokenizer_type == "wordpiece":
-        model = WordPiece(vocab={token: idx for idx, token in enumerate(vocab)}, unk_token=unk_token)
+        model = WordPiece(
+            vocab={token: idx for idx, token in enumerate(vocab)}, unk_token=unk_token, max_input_chars_per_word=100
+        )
     elif tokenizer_type == "bpe":
         model = BPE(
             vocab={token: idx for idx, token in enumerate(vocab)},
@@ -35,17 +37,19 @@ def mock_tokenizer(request: pytest.FixtureRequest) -> Tokenizer:
             ignore_merges=True,
         )
     elif tokenizer_type == "unigram":
-        model = Unigram(vocab=[(token, 0.0) for token in vocab], unk_id=0)
+        model = Unigram(vocab=[(token, 0.0) for token in vocab], unk_id=0, byte_fallback=False)
+    else:
+        raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")
     tokenizer = Tokenizer(model)
-    tokenizer.pre_tokenizer = Whitespace()
+    tokenizer.pre_tokenizer = Whitespace()  # type: ignore  # Tokenizer issue
 
     return tokenizer
 
 
 @pytest.fixture(scope="function")
-def mock_berttokenizer() -> AutoTokenizer:
+def mock_berttokenizer() -> PreTrainedTokenizerFast:
     """Load the real BertTokenizerFast from the provided tokenizer.json file."""
-    return AutoTokenizer.from_pretrained("tests/data/test_tokenizer")
+    return cast(PreTrainedTokenizerFast, AutoTokenizer.from_pretrained("tests/data/test_tokenizer"))
 
 
 @pytest.fixture

tests/test_tokenizer.py

@@ -0,0 +1,123 @@
+import json
+
+import pytest
+from transformers import PreTrainedTokenizerFast
+
+from model2vec.tokenizer.model import _calculate_token_weight_for_unigram, _process_unigram, process_tokenizer
+from model2vec.tokenizer.normalizer import replace_normalizer
+from model2vec.tokenizer.pretokenizer import _FORBIDDEN_PRETOKENIZERS, _fix_single_pretokenizer, replace_pretokenizer
+from model2vec.tokenizer.tokenizer import _rename_added_token, create_tokenizer
+
+
+def test_fix_single_pretokenizer() -> None:
+    """Test the _fix_single_pretokenizer function."""
+    result = _fix_single_pretokenizer({"type": "ByteLevel", "add_prefix_space": False, "use_regex": True})
+    assert result == {"type": "ByteLevel", "add_prefix_space": True, "use_regex": False}
+
+    for tokenizer_type in _FORBIDDEN_PRETOKENIZERS:
+        result = _fix_single_pretokenizer({"type": tokenizer_type})
+        assert result is None
+
+    result = _fix_single_pretokenizer(
+        {"type": "Metaspace", "split": True, "prepend_scheme": "never", "replacement": "▁"}
+    )
+    assert result == {"type": "Metaspace", "replacement": "▁", "prepend_scheme": "always", "split": False}
+
+
+def test_replace_pretokenizer(mock_berttokenizer: PreTrainedTokenizerFast) -> None:
+    """Test the replace_pretokenizer function."""
+    tokenizer = replace_pretokenizer(mock_berttokenizer.backend_tokenizer)
+    assert tokenizer.pre_tokenizer is not None
+    assert tokenizer.pre_tokenizer.__class__.__name__ == "Metaspace"
+    assert tokenizer.pre_tokenizer.replacement == "▁"
+    assert tokenizer.pre_tokenizer.prepend_scheme == "always"
+    assert not tokenizer.pre_tokenizer.split
+
+    tokenizer.pre_tokenizer = None  # type: ignore
+    tokenizer = replace_pretokenizer(tokenizer)
+    assert tokenizer.pre_tokenizer is not None
+    assert tokenizer.pre_tokenizer.__class__.__name__ == "Metaspace"
+    assert tokenizer.pre_tokenizer.replacement == "▁"
+    assert tokenizer.pre_tokenizer.prepend_scheme == "always"
+    assert tokenizer.pre_tokenizer.split is False
+
+
+def test_replace_normalizer(mock_berttokenizer: PreTrainedTokenizerFast) -> None:
+    """Test the replace_normalizer function."""
+    tokenizer = replace_normalizer(mock_berttokenizer.backend_tokenizer)
+    assert tokenizer.normalizer is not None
+    assert tokenizer.normalizer.__class__.__name__ == "Sequence"
+
+    assert tokenizer.normalizer.normalize_str("Hello, World!") == "hello , world !"
+
+    tokenizer.normalizer = None  # type: ignore
+    tokenizer = replace_normalizer(tokenizer)
+    assert tokenizer.normalizer.normalize_str("Hello, World!") == "Hello , World !"
+
+
+@pytest.mark.parametrize(
+    "word,weight",
+    [
+        ("dog", 3),
+        ("cat", 3),
+        ("▁longer▁word", 14),
+        ("▁word", 6),
+        ("▁", 2),  # Single underscore
+        ("", 0),  # Empty string
+        ("▁a" * 100, 300),  # Long word with underscores
+    ],
+)
+def test_calculate_token_weight_for_unigram(word: str, weight: int) -> None:
+    """Test the _calculate_token_weight_for_unigram function."""
+    assert _calculate_token_weight_for_unigram(word) == weight
+
+
+def test_process_tokenizer(mock_berttokenizer: PreTrainedTokenizerFast) -> None:
+    """Test the process_tokenizer function."""
+    vocab = ["dog", "cat", "longer_word", "word", "a" * 100, "[UNK]"]
+    tokenizer_json = json.loads(mock_berttokenizer.backend_tokenizer.to_str())
+    tokenizer_json = process_tokenizer(tokenizer_json=tokenizer_json, pre_tokenized_tokens=vocab, unk_token="[UNK]")
+
+    assert tokenizer_json["model"]["type"] == "Unigram"
+    assert tokenizer_json["model"]["unk_id"] == 5  # Index of "[UNK]"
+    assert len(tokenizer_json["model"]["vocab"]) == 6
+    assert all(isinstance(token, tuple) and len(token) == 2 for token in tokenizer_json["model"]["vocab"])
+    for (x, _), y in zip(tokenizer_json["model"]["vocab"], vocab):
+        assert x == y, f"Expected {y}, but got {x}"
+
+
+def test_process_unigram() -> None:
+    """Test the _process_unigram function."""
+    vocab = ["dog", "cat", "longer_word", "word", "a" * 100, "[UNK]"]
+    orig_vocab = [("dog", 0), ("cat", 0)]
+    model = {"model": {"type": "Unigram", "vocab": orig_vocab}}
+    processed_model = _process_unigram(model, vocab, "[UNK]")
+    assert processed_model["model"]["type"] == "Unigram"
+    assert processed_model["model"]["unk_id"] == 5  # Index of "[UNK]"
+    assert len(processed_model["model"]["vocab"]) == 6
+    assert all(isinstance(token, list) and len(token) == 2 for token in processed_model["model"]["vocab"])
+
+    for (x, score), y in zip(processed_model["model"]["vocab"], vocab):
+        assert x == y, f"Expected {y}, but got {x}"
+        if x in orig_vocab:
+            assert score == 0
+
+    assert process_tokenizer(model, vocab, "[UNK]") == processed_model
+
+
+def test_rename_added_token() -> None:
+    """Test the _rename_added_token function."""
+    # Invalid input
+    result = _rename_added_token(None, "a", [{"content": "a", "id": 0}], ["a"])
+    assert result == [{"content": "a", "id": 0}]
+
+    # Rename 'a' to 'c'
+    result = _rename_added_token("a", "c", [{"content": "a"}], ["a"])
+    assert result == [{"content": "c", "id": 0}]
+
+
+def test_create_tokenizer(mock_berttokenizer: PreTrainedTokenizerFast) -> None:
+    """Test the create_tokenizer function."""
+    tokenizer = create_tokenizer(tokenizer=mock_berttokenizer, vocabulary=["dog", "catssssss"], token_remove_regex=None)
+    assert tokenizer.backend_tokenizer.get_vocab_size() == 29525
+    assert tokenizer.encode("catssssss") == [29524]