fix tests, make tokenizer changes better

Author: stephantul

model2vec/distill/distillation.py

@@ -73,7 +73,9 @@ def distill_from_model(
 
     n_tokens_before = len(vocabulary)
     # Clean the vocabulary by removing duplicate tokens and tokens that are in the internal vocabulary.
-    all_tokens = clean_and_create_vocabulary(tokenizer, vocabulary, token_remove_regex=token_remove_regex)
+    all_tokens, backend_tokenizer = clean_and_create_vocabulary(
+        tokenizer, vocabulary, token_remove_regex=token_remove_regex
+    )
     n_tokens_after = len([token for token in all_tokens if not token.is_internal])
     if n_tokens_before:
         logger.info(

model2vec/tokenizer/normalizer.py

@@ -1,28 +1,34 @@
 from string import punctuation
 
-from tokenizers import Regex
+from tokenizers import Regex, Tokenizer
 from tokenizers.normalizers import Normalizer, Replace, Sequence, Strip
 
 
-def prepare_normalizer(
-    normalizer: Normalizer,
-) -> Normalizer:
+def replace_normalizer(
+    tokenizer: Tokenizer,
+) -> Tokenizer:
     """
-    Prepare the normalizer for the tokenizer.
+    Replace the normalizer for the tokenizer.
 
-    This function sets the normalizer for the tokenizer based on the provided normalizer type.
-    If no normalizer is provided, it uses the default one.
+    The new normalizer will replace punctuation with a space before and after the punctuation.
+    It will also replace multiple spaces with a single space and strip the right side of the string.
+    If the tokenizer already has a normalizer, it will be added to the new normalizer.
+    If the tokenizer does not have a normalizer, a new normalizer will be created.
 
-    :param normalizer: The tokenizer to prepare.
-    :return: The prepared tokenizer.
+    :param tokenizer: The tokenizer to change.
+    :return: The tokenizer with a replaced normalizer.
     """
+    normalizer = tokenizer.normalizer
     new_normalizers = []
     for char in punctuation:
         new_normalizers.append(Replace(char, f" {char} "))
 
     new_normalizers.append(Replace(Regex(r"\s+"), " "))
     new_normalizers.append(Strip(right=True))
     if normalizer is None:
-        return Sequence(new_normalizers)
+        normalizer = Sequence(new_normalizers)
+    else:
+        normalizer = Sequence([normalizer] + new_normalizers)  # type: ignore
+    tokenizer.normalizer = normalizer
 
-    return Sequence([normalizer] + new_normalizers)  # type: ignore
+    return tokenizer

model2vec/tokenizer/pretokenizer.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
+import json
 from typing import Any
 
+from tokenizers import Tokenizer
+
 _FORBIDDEN_PRETOKENIZERS = (
     "WhiteSpace",
     "WhitespaceSplit",
@@ -28,26 +31,27 @@ def _fix_single_pretokenizer(pre_tokenizer: dict[str, Any]) -> dict[str, Any] |
     return pre_tokenizer
 
 
-def fix_pretokenizer(pretokenizer: dict[str, Any] | None) -> dict[str, Any]:
+def replace_pretokenizer(tokenizer: Tokenizer) -> Tokenizer:
     """Fixes a single pretokenizer to allow multiword units."""
-    if pretokenizer is None:
-        return _BASIC_METASPACE
+    tokenizer_json = json.loads(tokenizer.to_str())
+    pre_tokenizer_json = tokenizer_json.get("pre_tokenizer", None)
+
+    if pre_tokenizer_json is None:
+        pre_tokenizer_json = _BASIC_METASPACE
 
-    if pretokenizer["type"] == "Sequence":
+    elif pre_tokenizer_json["type"] == "Sequence":
         new_pretokenizers = []
-        for single_pretokenizer in pretokenizer["pretokenizers"]:
+        for single_pretokenizer in pre_tokenizer_json["pretokenizers"]:
             new_pretokenizer = _fix_single_pretokenizer(single_pretokenizer)
             if new_pretokenizer is not None:
                 new_pretokenizers.append(new_pretokenizer)
-        pretokenizer["pretokenizers"] = new_pretokenizers
-
-        if not pretokenizer:
-            return _BASIC_METASPACE
 
-        return pretokenizer
+        if new_pretokenizers:
+            pre_tokenizer_json["pretokenizers"] = new_pretokenizers
+        else:
+            pre_tokenizer_json = _BASIC_METASPACE
 
-    single_pretokenizer = _fix_single_pretokenizer(pretokenizer)
-    if single_pretokenizer is None:
-        return _BASIC_METASPACE
+    pre_tokenizer_json = _fix_single_pretokenizer(pre_tokenizer_json) or _BASIC_METASPACE
+    tokenizer_json["pre_tokenizer"] = pre_tokenizer_json
 
-    return single_pretokenizer
+    return tokenizer.from_str(json.dumps(tokenizer_json))

model2vec/tokenizer/tokenizer.py

@@ -14,8 +14,8 @@
 
 from model2vec.tokenizer.datamodels import Token
 from model2vec.tokenizer.model import process_tokenizer
-from model2vec.tokenizer.normalizer import prepare_normalizer
-from model2vec.tokenizer.pretokenizer import fix_pretokenizer
+from model2vec.tokenizer.normalizer import replace_normalizer
+from model2vec.tokenizer.pretokenizer import replace_pretokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -54,11 +54,7 @@ def replace_vocabulary(
     tokenizer: Tokenizer, new_vocabulary: list[Token], unk_token: str | None, pad_token: str | None
 ) -> Tokenizer:
     """Replace the vocabulary of a tokenizer with a new one."""
-    tokenizer = tokenizer.from_str(tokenizer.to_str())
-    tokenizer.normalizer = prepare_normalizer(tokenizer.normalizer)  # type: ignore[assignment]  # Is just wrong
     tokenizer_json: dict[str, Any] = json.loads(tokenizer.to_str())
-    tokenizer_json["pre_tokenizer"] = fix_pretokenizer(tokenizer_json["pre_tokenizer"])
-
     added_tokens: list[dict[str, Any]] = tokenizer_json["added_tokens"]
 
     pre_tokenized_tokens = [x.normalized_form for x in new_vocabulary]
@@ -102,7 +98,7 @@ def clean_and_create_vocabulary(
     tokenizer: PreTrainedTokenizerFast,
     vocabulary: list[str],
     token_remove_regex: re.Pattern | None,
-) -> list[Token]:
+) -> tuple[list[Token], Tokenizer]:
     """Cleans a vocabulary by removing duplicates and tokens that were already in the vocabulary."""
     seen_tokens = set()
     post_normalize_seen_tokens = set()
@@ -115,15 +111,12 @@ def clean_and_create_vocabulary(
     internal_vocab: dict[str, int] = tokenizer.get_vocab()
     internal_tokens: list[str] = [k for k, _ in sorted(internal_vocab.items(), key=lambda x: x[1])]
 
-    cleaned_vocabulary = _process_internal_tokens(tokenizer, internal_tokens, token_remove_regex)
-    internal_tokens_set = {token.form for token in cleaned_vocabulary}
-
-    # Change the backend tokenizer to the new one.
+    # Copy the backend tokenizer to avoid modifying the original.
     backend_tokenizer = backend_tokenizer.from_str(backend_tokenizer.to_str())
-    backend_tokenizer.normalizer = prepare_normalizer(backend_tokenizer.normalizer)  # type: ignore[assignment]  # Is just wrong
-    tokenizer_json: dict[str, Any] = json.loads(backend_tokenizer.to_str())
-    tokenizer_json["pre_tokenizer"] = fix_pretokenizer(tokenizer_json["pre_tokenizer"])
-    backend_tokenizer = Tokenizer.from_str(json.dumps(tokenizer_json))
+    backend_tokenizer = replace_normalizer(backend_tokenizer)
+
+    cleaned_vocabulary = _process_internal_tokens(tokenizer, backend_tokenizer, internal_tokens, token_remove_regex)
+    internal_tokens_set = {token.form for token in cleaned_vocabulary}
 
     normalizer: Normalizer | None = backend_tokenizer.normalizer
     for token in vocabulary:
@@ -178,11 +171,14 @@ def clean_and_create_vocabulary(
     if n_empty:
         logger.warning(f"Removed {n_empty} empty tokens.")
 
-    return cleaned_vocabulary
+    return cleaned_vocabulary, replace_pretokenizer(backend_tokenizer)
 
 
 def _process_internal_tokens(
-    tokenizer: PreTrainedTokenizerFast, internal_tokens: list[str], token_remove_regex: re.Pattern | None
+    tokenizer: PreTrainedTokenizerFast,
+    backend_tokenizer: Tokenizer,
+    internal_tokens: list[str],
+    token_remove_regex: re.Pattern | None,
 ) -> list[Token]:
     """Clean internal tokens."""
     # Get the pad and unk token from the tokenizer.
@@ -193,7 +189,6 @@ def _process_internal_tokens(
     added_tokens_to_remove = set(tokenizer.added_tokens_encoder) - added_tokens_to_keep
     cleaned_internal_tokens: list[Token] = []
 
-    backend_tokenizer = tokenizer.backend_tokenizer
     # Figure out whether token is a subword or not.
     encoded = backend_tokenizer.encode(f" {'a' * 25}", add_special_tokens=False)
     first_token, second_token, *_ = encoded.tokens
@@ -378,7 +373,7 @@ def create_tokenizer(
     """
     unk_token = cast(str | None, tokenizer.special_tokens_map.get("unk_token"))
     pad_token = cast(str | None, tokenizer.special_tokens_map.get("pad_token"))
-    cleaned_vocabulary = clean_and_create_vocabulary(tokenizer, vocabulary, token_remove_regex)
-    new_tokenizer = replace_vocabulary(tokenizer.backend_tokenizer, cleaned_vocabulary, unk_token, pad_token)
+    cleaned_vocabulary, backend_tokenizer = clean_and_create_vocabulary(tokenizer, vocabulary, token_remove_regex)
+    new_tokenizer = replace_vocabulary(backend_tokenizer, cleaned_vocabulary, unk_token, pad_token)
 
     return PreTrainedTokenizerFast(tokenizer_object=new_tokenizer)

tests/test_distillation.py

@@ -22,7 +22,7 @@
     from huggingface_hub.errors import RepositoryNotFoundError
 except ImportError:
     # For huggingface_hub<0.25.0
-    from huggingface_hub.utils._errors import RepositoryNotFoundError
+    from huggingface_hub.utils._errors import RepositoryNotFoundError  # type: ignore
 
 rng = np.random.default_rng()
 
@@ -275,7 +275,7 @@ def test_clean_and_create_vocabulary(
 ) -> None:
     """Test the _clean_vocabulary function."""
     with caplog.at_level("WARNING"):
-        tokens = clean_and_create_vocabulary(mock_berttokenizer, added_tokens, None)
+        tokens, _ = clean_and_create_vocabulary(mock_berttokenizer, added_tokens, None)
 
         cleaned_vocab = [token.form for token in tokens if not token.is_internal]
         # Check the cleaned vocabulary matches the expected output