add superbpe tokenizers

Author: stephantul

model2vec/distill/distillation.py

@@ -302,7 +302,6 @@ def _clean_vocabulary(tokenizer: Tokenizer, vocabulary: list[str], added_tokens:
     cleaned_vocabulary = []
     n_empty = 0
     n_duplicates = 0
-    n_multiword = 0
     for token in vocabulary:
         if tokenizer.normalizer is not None:
             token = tokenizer.normalizer.normalize_str(token)
@@ -321,7 +320,5 @@ def _clean_vocabulary(tokenizer: Tokenizer, vocabulary: list[str], added_tokens:
         logger.warning(f"Removed {n_duplicates} duplicate tokens.")
     if n_empty:
         logger.warning(f"Removed {n_empty} empty tokens.")
-    if n_multiword:
-        logger.warning(f"Removed {n_multiword} multiword tokens.")
 
     return cleaned_vocabulary

model2vec/distill/tokenizer.py

@@ -5,6 +5,32 @@
 from typing import Any
 
 from tokenizers import Tokenizer
+from tokenizers.pre_tokenizers import (
+    BertPreTokenizer,
+    ByteLevel,
+    CharDelimiterSplit,
+    Digits,
+    Metaspace,
+    PreTokenizer,
+    Punctuation,
+    Sequence,
+    Split,
+    UnicodeScripts,
+    Whitespace,
+    WhitespaceSplit,
+)
+
+_FORBIDDEN_PRETOKENIZERS = (
+    BertPreTokenizer,
+    CharDelimiterSplit,
+    Metaspace,
+    Punctuation,
+    Split,
+    UnicodeScripts,
+    Whitespace,
+    WhitespaceSplit,
+)
+
 
 logger = logging.getLogger(__name__)
 
@@ -36,9 +62,9 @@ def _pre_tokenize_vocabulary(tokenizer: Tokenizer, tokens: list[str]) -> list[st
             if token in current_tokenizer_vocab:
                 pre_tokenized_tokens.append(token)
             else:
-                # We know 100% sure that all pretokenized tokens will have length 1.
-                pretokenized_tokens, _ = zip(*tokenizer.pre_tokenizer.pre_tokenize_str(f" {token}"))
-                pre_tokenized_tokens.append(pretokenized_tokens[-1])
+                # Join tokens just to be sure.
+                pretokenized_tokens, _ = zip(*tokenizer.pre_tokenizer.pre_tokenize_str(token))
+                pre_tokenized_tokens.append(" ".join(pretokenized_tokens))
     else:
         pre_tokenized_tokens = tokens
 
@@ -67,6 +93,30 @@ def _remap_added_tokens(
     return special_tokens
 
 
+def _fix_single_pretokenizer(pretokenizer: PreTokenizer) -> PreTokenizer | None:
+    """Fixes a single pretokenizer to allow multiword units."""
+    if isinstance(pretokenizer, _FORBIDDEN_PRETOKENIZERS):
+        return Metaspace(split=False, replacement="Ġ")
+    elif isinstance(pretokenizer, ByteLevel):
+        pretokenizer.use_regex = False
+
+    return pretokenizer
+
+
+def _fix_pretokenizer_for_super(pre: PreTokenizer | None) -> Tokenizer:
+    """Fixes the pretokenizer to allow multiword units."""
+    if pre is None:
+        return pre
+
+    if isinstance(pre, Sequence):
+        new_pretokenizers = []
+        for pretokenizer in pre:
+            new_pretokenizers.append(_fix_single_pretokenizer(pretokenizer))
+        return Sequence(new_pretokenizers)
+
+    return _fix_single_pretokenizer(pre)
+
+
 def _make_new_merges_from_vocab(
     merges: list[tuple[str, str]], tokens: list[str], special_tokens: set[str | None]
 ) -> list[tuple[str, str]]:
@@ -109,6 +159,7 @@ def replace_vocabulary(
     tokenizer: Tokenizer, new_vocabulary: list[str], unk_token: str | None, pad_token: str | None
 ) -> Tokenizer:
     """Replace the vocabulary of a tokenizer with a new one."""
+    tokenizer.pre_tokenizer = _fix_pretokenizer_for_super(tokenizer.pre_tokenizer)
     tokenizer_json: dict[str, Any] = json.loads(tokenizer.to_str())
 
     # NOTE: all tokens have been normalized before.
@@ -124,10 +175,30 @@ def replace_vocabulary(
         unk_token = unk_token or tokenizer_json["model"]["unk_token"]
         tokenizer_json["model"]["unk_token"] = unk_token
         tokenizer_json["added_tokens"] = [x for x in tokenizer_json["added_tokens"] if x["content"] in special_tokens]
-        tokenizer_json["model"]["vocab"] = {token: idx for idx, token in enumerate(pre_tokenized_tokens)}
+
+        if model_type == "WordPiece":
+            subword_prefix = tokenizer_json["model"]["continuing_subword_prefix"]
+            new_vocab = {}
+            for idx, token in enumerate(pre_tokenized_tokens):
+                if token in special_tokens:
+                    # We need to remove the prefix from the token
+                    pass
+                elif token.startswith(subword_prefix):
+                    # We need to remove the prefix from the token
+                    token = token.removeprefix(subword_prefix)
+                elif token.startswith("Ġ"):
+                    pass
+                else:
+                    # We need to add the prefix to the token
+                    token = f"Ġ{token}"
+                new_vocab[token] = idx
+            tokenizer_json["model"]["continuing_subword_prefix"] = ""
+            tokenizer_json["model"]["max_input_chars_per_word"] = 10_000
+            tokenizer_json["model"]["vocab"] = new_vocab
 
         if model_type == "BPE":
             # Bit more difficult, we need to take into account merges.
+            tokenizer_json["model"]["vocab"] = {token: idx for idx, token in enumerate(pre_tokenized_tokens)}
             merges = tokenizer_json["model"]["merges"]
             merges = _make_new_merges_from_vocab(merges, pre_tokenized_tokens, special_tokens)
             tokenizer_json["model"]["merges"] = merges
@@ -154,4 +225,5 @@ def replace_vocabulary(
     tokenizer_json["added_tokens"] = _remap_added_tokens(added_tokens, pre_tokenized_tokens)
     tokenizer_json["post_processor"] = _DEFAULT_POST_PROCESSOR_TEMPLATE
 
-    return Tokenizer.from_str(json.dumps(tokenizer_json))
+    tokenizer = Tokenizer.from_str(json.dumps(tokenizer_json))
+    return tokenizer