feat: add skeletoken

Author: stephantul

model2vec/distill/distillation.py

@@ -3,11 +3,13 @@
 import logging
 import os
 import re
-from typing import Optional, cast
+from typing import cast
 
 import numpy as np
 from huggingface_hub import model_info
-from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerFast
+from transformers import AutoModel, AutoTokenizer
+from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
 from model2vec.distill.inference import PCADimType, create_embeddings, post_process_embeddings
 from model2vec.distill.utils import select_optimal_device
@@ -24,11 +26,10 @@ def distill_from_model(
     vocabulary: list[str] | None = None,
     device: str | None = None,
     pca_dims: PCADimType = 256,
-    apply_zipf: bool | None = None,
     sif_coefficient: float | None = 1e-4,
     token_remove_pattern: str | None = r"\[unused\d+\]",
     quantize_to: DType | str = DType.Float16,
-    use_subword: bool | None = None,
+    lower_case: bool = True,
 ) -> StaticModel:
     """
     Distill a staticmodel from a sentence transformer.
@@ -46,25 +47,20 @@ def distill_from_model(
     :param pca_dims: The number of components to use for PCA.
         If this is None, we don't apply PCA.
         If this is 'auto', we don't reduce dimensionality, but still apply PCA.
-    :param apply_zipf: DEPRECATED: This parameter used to control whether Zipf is applied.
-        Zipf weighting is now controlled by the sif_coefficient parameter. If this is set to None, no weighting is applied.
     :param sif_coefficient: The SIF coefficient to use. If this is None, no weighting is applied.
         Should be a value > 0 and < 1.0. A value of 1e-4 is a good default.
     :param token_remove_pattern: If this is set to a string, we compile this into a regex. Any tokens that conform to this regex pattern will be removed from the vocabulary.
         If the pattern is so general that it removes all tokens, we throw an error. If the pattern can't be compiled into a valid regex, we also throw an error.
     :param quantize_to: The data type to quantize to. Can be any of the DType enum members or their string equivalents.
-    :param use_subword: DEPRECATED: If this is not set to None, we show a warning. It doesn't do anything.
+    :param lower_case: If this is set, all tokens in the model vocabulary will be converted to lowercase, and
+        a lowercase normalizer will be inserted. This almost always improves performance.
     :return: A StaticModel
     :raises: ValueError if the vocabulary is empty after preprocessing.
 
     """
-    if use_subword is not None:
-        logger.warning(
-            "The `use_subword` parameter is deprecated and will be removed in the next release. It doesn't do anything."
-        )
     quantize_to = DType(quantize_to)
     backend_tokenizer = tokenizer.backend_tokenizer
-    sif_coefficient, token_remove_regex = _validate_parameters(apply_zipf, sif_coefficient, token_remove_pattern)
+    sif_coefficient, token_remove_regex = _validate_parameters(sif_coefficient, token_remove_pattern)
 
     if vocabulary is None:
         vocabulary = []
@@ -73,45 +69,37 @@ 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, backend_tokenizer = clean_and_create_vocabulary(
-        tokenizer, vocabulary, token_remove_regex=token_remove_regex
+    tokens, backend_tokenizer = clean_and_create_vocabulary(
+        tokenizer, vocabulary, token_remove_regex=token_remove_regex, lower_case=lower_case
     )
-    n_tokens_after = len([token for token in all_tokens if not token.is_internal])
+    n_tokens_after = len([token for token in tokens if not token.is_internal])
     if n_tokens_before:
         logger.info(
             f"Adding {n_tokens_after} tokens to the vocabulary. Removed {n_tokens_before - n_tokens_after} tokens during preprocessing."
         )
 
-    if not all_tokens:
+    if not tokens:
         raise ValueError("The vocabulary is empty after preprocessing. Please check your token_remove_pattern.")
 
-    unk_token = cast(Optional[str], tokenizer.special_tokens_map.get("unk_token"))
-    pad_token = cast(Optional[str], tokenizer.special_tokens_map.get("pad_token"))
-
-    # Weird if to satsify mypy
-    if pad_token is None:
-        if unk_token is not None:
-            pad_token = unk_token
-            logger.warning(
-                "The pad token is not set. Setting it to the unk token. This is a workaround for models that don't have a pad token."
-            )
-        else:
-            pad_token = unk_token or all_tokens[0].form
-            logger.warning(
-                "The pad token is not set. Setting it to the first token in the vocabulary. This is a workaround for models that don't have a pad token."
-            )
-
     # Replace the vocabulary in the tokenizer with the new vocabulary.
-    backend_tokenizer = replace_vocabulary(backend_tokenizer, all_tokens, unk_token=unk_token, pad_token=pad_token)
+    backend_tokenizer = replace_vocabulary(backend_tokenizer, tokens)
 
-    logger.info(f"Creating embeddings for {len(all_tokens)} tokens")
+    logger.info(f"Creating embeddings for {len(tokens)} tokens")
     # Convert tokens to IDs
-    token_ids = turn_tokens_into_ids(all_tokens, tokenizer, unk_token)
+    token_ids = turn_tokens_into_ids(tokens, tokenizer.backend_tokenizer)
 
     # Create the embeddings
-    embeddings = create_embeddings(
-        tokenized=token_ids, model=model, device=device, pad_token_id=tokenizer.get_vocab()[pad_token]
-    )
+    pad_token = cast(str | None, tokenizer.special_tokens_map.get("pad_token", None))
+    vocab = tokenizer.get_vocab()
+    if pad_token is None:
+        sep_token = cast(str | None, tokenizer.special_tokens_map.get("sep_token", None))
+        if sep_token is None:
+            pad_token_id = 0
+        else:
+            pad_token_id = vocab[sep_token]
+    else:
+        pad_token_id = vocab[pad_token]
+    embeddings = create_embeddings(tokenized=token_ids, model=model, device=device, pad_token_id=pad_token_id)
 
     # Post process the embeddings by applying PCA and Zipf weighting.
     embeddings = post_process_embeddings(np.asarray(embeddings), pca_dims, sif_coefficient=sif_coefficient)
@@ -125,7 +113,6 @@ def distill_from_model(
         "architectures": ["StaticModel"],
         "tokenizer_name": model_name,
         "apply_pca": pca_dims,
-        "apply_zipf": apply_zipf,
         "sif_coefficient": sif_coefficient,
         "hidden_dim": embeddings.shape[1],
         "seq_length": 1000000,  # Set this to a high value since we don't have a sequence length limit.
@@ -157,35 +144,19 @@ def distill_from_model(
 
 
 def _validate_parameters(
-    apply_zipf: bool | None,
     sif_coefficient: float | None,
     token_remove_pattern: str | None,
 ) -> tuple[float | None, re.Pattern | None]:
     """
     Validate the parameters passed to the distillation function.
 
-    :param apply_zipf: DEPRECATED: This parameter used to control whether Zipf is applied.
-        Zipf weighting is now controlled by the sif_coefficient parameter. If this is set to None, no weighting is applied.
     :param sif_coefficient: The SIF coefficient to use. If this is None, no weighting is applied.
         Should be a value >= 0 and < 1.0. A value of 1e-4 is a good default.
     :param token_remove_pattern: If this is set to a string, we compile this into a regex. Any tokens that conform to this regex pattern will be removed from the vocabulary.
     :return: The SIF coefficient to use.
     :raises: ValueError if the regex can't be compiled.
 
     """
-    if apply_zipf is not None:
-        logger.warning(
-            "The `apply_zipf` parameter is deprecated and will be removed in the next release. "
-            "Zipf weighting is applied based on the sif_coefficient parameter. If this is set to None, "
-            "no weighting is applied."
-        )
-        if apply_zipf and sif_coefficient is None:
-            logger.warning("You set apply_zipf to True, but sif_coefficient is None. Setting sif_coefficient to 1e-4.")
-            sif_coefficient = 1e-4
-        elif not apply_zipf:
-            logger.warning("Because you set apply_zipf to False, we ignore the sif_coefficient parameter.")
-            sif_coefficient = None
-
     if sif_coefficient is not None:
         if not 0 < sif_coefficient < 1.0:
             raise ValueError("SIF coefficient must be a value > 0 and < 1.0.")
@@ -205,12 +176,11 @@ def distill(
     vocabulary: list[str] | None = None,
     device: str | None = None,
     pca_dims: PCADimType = 256,
-    apply_zipf: bool | None = None,
     sif_coefficient: float | None = 1e-4,
     token_remove_pattern: str | None = r"\[unused\d+\]",
     trust_remote_code: bool = False,
     quantize_to: DType | str = DType.Float16,
-    use_subword: bool | None = None,
+    lower_case: bool = True,
 ) -> StaticModel:
     """
     Distill a staticmodel from a sentence transformer.
@@ -227,14 +197,13 @@ def distill(
     :param pca_dims: The number of components to use for PCA.
         If this is None, we don't apply PCA.
         If this is 'auto', we don't reduce dimenionality, but still apply PCA.
-    :param apply_zipf: DEPRECATED: This parameter used to control whether Zipf is applied.
-        Zipf weighting is now controlled by the sif_coefficient parameter. If this is set to None, no weighting is applied.
     :param sif_coefficient: The SIF coefficient to use. If this is None, no weighting is applied.
         Should be a value >= 0 and < 1.0. A value of 1e-4 is a good default.
     :param token_remove_pattern: If this is set to a string, we compile this into a regex. Any tokens that conform to this regex pattern will be removed from the vocabulary.
     :param trust_remote_code: Whether to trust the remote code. If this is False, we will only load components coming from `transformers`. If this is True, we will load all components.
     :param quantize_to: The data type to quantize to. Can be any of the DType enum members or their string equivalents.
-    :param use_subword: DEPRECATED: If this is not set to None, we show a warning. It doesn't do anything.
+    :param lower_case: If this is set, all tokens in the model vocabulary will be converted to lowercase, and
+        a lowercase normalizer will be inserted. This almost always improves performance.
     :return: A StaticModel
 
     """
@@ -250,9 +219,8 @@ def distill(
         vocabulary=vocabulary,
         device=device,
         pca_dims=pca_dims,
-        apply_zipf=apply_zipf,
         token_remove_pattern=token_remove_pattern,
         sif_coefficient=sif_coefficient,
         quantize_to=quantize_to,
-        use_subword=use_subword,
+        lower_case=lower_case,
     )

model2vec/distill/inference.py

@@ -46,7 +46,7 @@ def create_embeddings(
     :param pad_token_id: The pad token id. Used to pad sequences.
     :return: The output embeddings.
     """
-    model = model.to(device)
+    model = model.to(device)  # type: ignore
 
     out_weights: np.ndarray
     intermediate_weights: list[np.ndarray] = []
@@ -98,6 +98,7 @@ def _encode_mean_using_model(model: PreTrainedModel, encodings: dict[str, torch.
     """
     encodings = {k: v.to(model.device) for k, v in encodings.items()}
     encoded: BaseModelOutputWithPoolingAndCrossAttentions = model(**encodings)
+    assert encoded.last_hidden_state is not None
     out: torch.Tensor = encoded.last_hidden_state.cpu()
     # NOTE: If the dtype is bfloat 16, we convert to float32,
     # because numpy does not suport bfloat16

model2vec/tokenizer/__init__.py

@@ -4,9 +4,8 @@
 
 from model2vec.tokenizer.tokenizer import (
     clean_and_create_vocabulary,
-    create_tokenizer,
     replace_vocabulary,
     turn_tokens_into_ids,
 )
 
-__all__ = ["clean_and_create_vocabulary", "create_tokenizer", "turn_tokens_into_ids", "replace_vocabulary"]
+__all__ = ["clean_and_create_vocabulary", "turn_tokens_into_ids", "replace_vocabulary"]

model2vec/tokenizer/datamodels.py

@@ -5,10 +5,13 @@
 class Token:
     """A class to represent a token."""
 
+    # The surface form: used for featurizing
     form: str
-    # The normalized and pretokenized form of the token
+    # The normalized form: preprocessed by the new tokenizer
     normalized_form: str
     # Whether the word is a continuing subword.
     is_subword: bool
     # Whether the token is internal to the model.
     is_internal: bool
+    # Whether the token is a multiword token
+    is_multiword: bool = False