working version

Author: stephantul

model2vec/distill/distillation.py

@@ -107,10 +107,8 @@ def distill_from_model(
     pad_token = tokenizer.special_tokens_map.get("pad_token")
     # Add the cleaned vocabulary to the tokenizer.
     backend_tokenizer = replace_vocabulary(backend_tokenizer, all_tokens, unk_token=unk_token, pad_token=pad_token)
-
     # Post process the embeddings by applying PCA and Zipf weighting.
     embeddings = _post_process_embeddings(np.asarray(embeddings), pca_dims, sif_coefficient=sif_coefficient)
-
     # Quantize the embeddings.
     embeddings = quantize_embeddings(embeddings, quantize_to)
 

model2vec/distill/inference.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import torch
+from tokenizers.models import BPE, Unigram, WordPiece
 from torch.nn.utils.rnn import pad_sequence
 from tqdm import tqdm
 from transformers import PreTrainedModel, PreTrainedTokenizerFast
@@ -71,6 +72,23 @@ def create_embeddings(
         # If the token remove regex is None, just use all tokens.
         id_list = list(range(len(tokenizer.get_vocab())))
 
+    id_set = set(id_list)
+    new_id_list = []
+    for token, idx in tokenizer.get_vocab().items():
+        if idx not in id_set:
+            continue
+
+        if (
+            tokenizer.backend_tokenizer.pre_tokenizer is not None
+            and not token.startswith("##")
+            and not token in tokens_to_keep
+        ):
+            pre_token = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(token)
+            if len(pre_token) > 1:
+                continue
+        new_id_list.append(idx)
+    id_list = new_id_list
+
     added_tokens_ids = [id for token, id in tokenizer.added_tokens_encoder.items() if token not in tokens_to_keep]
     ids = torch.Tensor(sorted(set(id_list) - set(added_tokens_ids))).long()
 
@@ -82,7 +100,28 @@ def create_embeddings(
         eos = torch.full([len(ids)], fill_value=eos_token_id)
 
         tokenized.extend(torch.stack([bos, ids, eos], dim=1))
-        subword_tokens = [Token(x, True) for x in tokenizer.convert_ids_to_tokens(ids.tolist())]
+
+        subword_tokens = []
+        for token in tokenizer.convert_ids_to_tokens(ids.tolist()):
+            is_subword = False
+            should_be_pretokenized = True
+            if token == unk_token or token == pad_token:
+                is_subword = True
+            elif isinstance(tokenizer.backend_tokenizer.model, WordPiece):
+                prefix_char = tokenizer.backend_tokenizer.model.continuing_subword_prefix
+                if token.startswith(prefix_char):
+                    is_subword = True
+            elif isinstance(tokenizer.backend_tokenizer.model, Unigram):
+                if not token.startswith("▁"):
+                    is_subword = True
+            elif isinstance(tokenizer.backend_tokenizer.model, BPE):
+                if not token.startswith("Ġ"):
+                    is_subword = True
+                    should_be_pretokenized = False
+                else:
+                    should_be_pretokenized = False
+            subword_tokens.append(Token(token, is_subword, should_be_pretokenized))
+
         out_tokens.extend(subword_tokens)
 
     tokenized.extend([tokenizer.encode_plus(token, return_tensors="pt")["input_ids"][0] for token in tokens])
@@ -113,7 +152,7 @@ def create_embeddings(
 
     # Sort the output back to the original order
     intermediate_weights = [intermediate_weights[i] for i in np.argsort(sort_order)]
-    out_tokens.extend([Token(x, False) for x in tokens])
+    out_tokens.extend([Token(x, False, True) for x in tokens])
     out_weights = np.stack(intermediate_weights)
 
     return out_tokens, out_weights

model2vec/distill/tokenizer.py

@@ -2,14 +2,16 @@
 
 import json
 import logging
+from string import punctuation
 from typing import Any
 
-from tokenizers import Tokenizer
+from tokenizers import Regex, Tokenizer
+from tokenizers.normalizers import Lowercase, Normalizer, Replace, Strip
+from tokenizers.normalizers import Sequence as NormalizerSequence
 from tokenizers.pre_tokenizers import (
     BertPreTokenizer,
     ByteLevel,
     CharDelimiterSplit,
-    Digits,
     Metaspace,
     PreTokenizer,
     Punctuation,
@@ -45,7 +47,7 @@
 }
 
 
-def _pre_tokenize_vocabulary(tokenizer: Tokenizer, tokens: list[Token]) -> list[str]:
+def _pre_tokenize_vocabulary(tokenizer: Tokenizer, tokens: list[Token], subword_prefix: str) -> list[str]:
     """
     Apply pre-tokenization to vocabulary tokens if a pre-tokenizer is present.
 
@@ -54,19 +56,28 @@ def _pre_tokenize_vocabulary(tokenizer: Tokenizer, tokens: list[Token]) -> list[
 
     :param tokenizer: The tokenizer to use.
     :param tokens: The tokens to pre-tokenize.
+    :param subword_prefix: The prefix for subwords.
     :return: The pre-tokenized tokens.
     """
     pre_tokenized_tokens = []
 
     if tokenizer.pre_tokenizer is not None:
         for token in tokens:
-            if token.is_original:
+            if token.is_subword:
                 # Original tokens do not need to be pre-tokenized.
-                pre_tokenized_tokens.append(token.form)
-            else:
+                form = token.form
+                if subword_prefix is not None:
+                    form = token.form.removeprefix(subword_prefix)
+                pre_tokenized_tokens.append(form)
+            elif token.should_be_pretokenized:
                 # Join tokens just to be sure.
+                token.form = tokenizer.normalizer.normalize_str(token.form).rstrip()
                 pretokenized_tokens, _ = zip(*tokenizer.pre_tokenizer.pre_tokenize_str(token.form))
-                pre_tokenized_tokens.append(" ".join(pretokenized_tokens))
+                form = " ".join(pretokenized_tokens)
+                pre_tokenized_tokens.append(form)
+            else:
+                token.form = tokenizer.normalizer.normalize_str(token.form).rstrip()
+                pre_tokenized_tokens.append(token.form)
     else:
         pre_tokenized_tokens = [token.form for token in tokens]
 
@@ -95,12 +106,38 @@ def _remap_added_tokens(
     return special_tokens
 
 
+def _prepare_normalizer(
+    normalizer: Normalizer,
+) -> Normalizer:
+    """
+    Prepare 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.
+
+    :param normalizer: The tokenizer to prepare.
+    :return: The prepared tokenizer.
+    """
+    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 NormalizerSequence(new_normalizers)
+
+    return NormalizerSequence([normalizer] + new_normalizers)
+
+
 def _fix_single_pretokenizer(pretokenizer: PreTokenizer) -> PreTokenizer | None:
     """Fixes a single pretokenizer to allow multiword units."""
+    if isinstance(pretokenizer, Metaspace):
+        return Metaspace(split=False, replacement=pretokenizer.replacement, prepend_scheme=pretokenizer.prepend_scheme)
     if isinstance(pretokenizer, _FORBIDDEN_PRETOKENIZERS):
-        return Metaspace(split=False, replacement="Ġ")
+        return Metaspace(split=False, replacement="▁")
     elif isinstance(pretokenizer, ByteLevel):
         pretokenizer.use_regex = False
+        pretokenizer.add_prefix_space = True
 
     return pretokenizer
 
@@ -111,68 +148,29 @@ def _fix_pretokenizer_for_super(pre: PreTokenizer | None) -> Tokenizer:
         return pre
 
     if isinstance(pre, Sequence):
-        new_pretokenizers = []
-        for pretokenizer in pre:
-            new_pretokenizers.append(_fix_single_pretokenizer(pretokenizer))
-        return Sequence(new_pretokenizers)
+        return Metaspace(split=False)
 
     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]]:
-    """
-    Generate new merges from a vocabulary.
-
-    This function creates new merge pairs from a given vocabulary of tokens.
-    The merges are used to build or extend a tokenizer's merge table.
-
-    :param merges: The list of existing merges in the form (first, second) where first and second are tokens.
-    :param tokens: The list of tokens (vocabulary) from which to generate new merges.
-    :param special_tokens: Tokens that should not be merged.
-    :return: The list of new merges in the form (first, second) where first and second are tokens.
-    """
-    new_merges = merges.copy()
-    current_vocab = set(tokens) - special_tokens
-    already_merged = set("".join(merge) for merge in merges)
-
-    for token in tokens:
-        if token in special_tokens:
-            continue
-        if token in already_merged:
-            continue
-        if len(token) == 1:
-            continue
-        merges = []
-        for index in range(1, len(token)):
-            first, second = token[:index], token[index:]
-            if first in current_vocab and second in current_vocab:
-                merges.append((first, second))
-        if not merges:
-            logger.warning(f"Token {token} has no merges.")
-            continue
-        new_merges.extend(merges)
-
-    return new_merges
-
-
 def _process_wordpiece(
     tokenizer_json: dict[str, Any], pre_tokenized_tokens: list[str], unk_token: str | None
 ) -> dict[str, Any]:
     """Process the WordPiece tokenizer JSON."""
-    tokenizer_json["model"]["unk_token"] = unk_token
-    tokenizer_json["model"]["vocab"] = {token: idx for idx, token in enumerate(pre_tokenized_tokens)}
+    tokenizer_json["model"]["type"] = "Unigram"
+    tokenizer_json["model"]["unk_id"] = pre_tokenized_tokens.index(unk_token) if unk_token else None
+    tokenizer_json["model"]["vocab"] = [(token, 0.0) for token in pre_tokenized_tokens]
 
     return tokenizer_json
 
 
-def _process_bpe(tokenizer_json: dict[str, Any], pre_tokenized_tokens: list[str]) -> dict[str, Any]:
+def _process_bpe(
+    tokenizer_json: dict[str, Any], pre_tokenized_tokens: list[str], unk_token: str | None
+) -> dict[str, Any]:
     """Process the BPE tokenizer JSON."""
-    tokenizer_json = _process_wordpiece(tokenizer_json, pre_tokenized_tokens, None)
-    merges = tokenizer_json["model"]["merges"]
-    merges = _make_new_merges_from_vocab(merges, pre_tokenized_tokens, {"[UNK]", "[PAD]"})
-    tokenizer_json["model"]["merges"] = merges
+    tokenizer_json["model"]["type"] = "Unigram"
+    tokenizer_json["model"]["unk_id"] = pre_tokenized_tokens.index(unk_token) if unk_token else None
+    tokenizer_json["model"]["vocab"] = [(token, 0.0) for token in pre_tokenized_tokens]
 
     return tokenizer_json
 
@@ -194,13 +192,16 @@ 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.normalizer = _prepare_normalizer(tokenizer.normalizer)
     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.
     # Very careful, we need to pretokenize words before adding them to the vocabulary.
     # But only if they are not part of the original vocabulary.
-    pre_tokenized_tokens = _pre_tokenize_vocabulary(tokenizer, new_vocabulary)
+    subword_prefix = tokenizer_json["model"].get("continuing_subword_prefix", "")
+
+    pre_tokenized_tokens = _pre_tokenize_vocabulary(tokenizer, new_vocabulary, subword_prefix=subword_prefix)
 
     model_type = tokenizer_json["model"]["type"]
     added_tokens: list[dict[str, Any]] = tokenizer_json["added_tokens"]
@@ -215,7 +216,7 @@ def replace_vocabulary(
     if model_type == "WordPiece":
         tokenizer_json = _process_wordpiece(tokenizer_json, pre_tokenized_tokens, "[UNK]")
     elif model_type == "BPE":
-        tokenizer_json = _process_bpe(tokenizer_json, pre_tokenized_tokens)
+        tokenizer_json = _process_bpe(tokenizer_json, pre_tokenized_tokens, "[UNK]")
     elif model_type == "Unigram":
         tokenizer_json = _process_unigram(tokenizer_json, pre_tokenized_tokens, "[UNK]")
     else:

model2vec/distill/utils.py

@@ -14,7 +14,13 @@ class Token:
     """A class to represent a token."""
 
     form: str
-    is_original: bool
+    # Whether the word is a continuing subword.
+    is_subword: bool
+    # Whether it should be pretokenized.
+    # This is independent of is_subword, because some
+    # tokenizer models like BPE and Unigram do not have a
+    # continuing subword prefix, but instead prefix nonsubwords.
+    should_be_pretokenized: bool
 
 
 def select_optimal_device(device: str | None) -> str: