chore: moved model2vec as in internal package

src/distiller/model2vec/__init__.py

@@ -0,0 +1,4 @@
+from .model import StaticModel
+from .version import __version__
+
+__all__ = ["StaticModel", "__version__"]

src/distiller/model2vec/distill/__init__.py

@@ -0,0 +1,10 @@
+from distiller.model2vec.utils import get_package_extras, importable
+
+_REQUIRED_EXTRA = "distill"
+
+for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA):
+    importable(extra_dependency, _REQUIRED_EXTRA)
+
+from distiller.model2vec.distill.distillation import distill, distill_from_model
+
+__all__ = ["distill", "distill_from_model"]

src/distiller/model2vec/distill/distillation.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import logging
+import os
+import re
+from typing import cast
+
+import numpy as np
+from huggingface_hub import model_info
+from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerFast
+
+from distiller.model2vec.distill.inference import PCADimType, create_embeddings, post_process_embeddings
+from distiller.model2vec.distill.utils import select_optimal_device
+from distiller.model2vec.model import StaticModel
+from distiller.model2vec.quantization import DType, quantize_embeddings
+from distiller.model2vec.tokenizer import clean_and_create_vocabulary, replace_vocabulary, turn_tokens_into_ids
+
+logger = logging.getLogger(__name__)
+
+
+def distill_from_model(
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizerFast,
+    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,
+) -> StaticModel:
+    """
+    Distill a staticmodel from a sentence transformer.
+
+    This function creates a set of embeddings from a sentence transformer. It does this by doing either
+    a forward pass for all subword tokens in the tokenizer, or by doing a forward pass for all tokens in a passed vocabulary.
+
+    If you pass through a vocabulary, we create a custom word tokenizer for that vocabulary.
+    If you don't pass a vocabulary, we use the model's tokenizer directly.
+
+    :param model: The model to use.
+    :param tokenizer: The tokenizer to use.
+    :param vocabulary: The vocabulary to use. If this is None, we use the model's vocabulary.
+    :param device: The device to use.
+    :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.
+    :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)
+
+    if vocabulary is None:
+        vocabulary = []
+
+    device = select_optimal_device(device)
+
+    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
+    )
+    n_tokens_after = len([token for token in all_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:
+        msg = "The vocabulary is empty after preprocessing. Please check your token_remove_pattern."
+        raise ValueError(msg)
+
+    unk_token = cast("str | None", tokenizer.special_tokens_map.get("unk_token"))
+    pad_token = cast("str | None", 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)
+
+    logger.info(f"Creating embeddings for {len(all_tokens)} tokens")
+    # Convert tokens to IDs
+    token_ids = turn_tokens_into_ids(all_tokens, tokenizer, unk_token)
+
+    # Create the embeddings
+    embeddings = create_embeddings(
+        tokenized=token_ids, model=model, device=device, pad_token_id=tokenizer.get_vocab()[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)
+
+    model_name = getattr(model, "name_or_path", "")
+
+    config = {
+        "model_type": "model2vec",
+        "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.
+        "normalize": True,
+    }
+
+    if os.path.exists(model_name):
+        # Using a local model. Get the model name from the path.
+        model_name = os.path.basename(model_name)
+        language = None
+    else:
+        # Get the language from the model card.
+        try:
+            info = model_info(model_name)
+            language = info.cardData.get("language", None) if info.cardData is not None else None
+        except Exception as e:
+            # NOTE: bare except because there's many reasons this can fail.
+            logger.warning(f"Couldn't get the model info from the Hugging Face Hub: {e}. Setting language to None.")
+            language = None
+
+    return StaticModel(
+        vectors=embeddings,
+        tokenizer=backend_tokenizer,
+        config=config,
+        base_model_name=model_name,
+        language=language,
+        normalize=True,
+    )
+
+
+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 and not 0 < sif_coefficient < 1.0:
+        msg = "SIF coefficient must be a value > 0 and < 1.0."
+        raise ValueError(msg)
+
+    token_remove_regex: re.Pattern | None = None
+    if token_remove_pattern is not None:
+        try:
+            token_remove_regex = re.compile(token_remove_pattern)
+        except re.error as e:
+            msg = f"Couldn't compile the regex pattern: {e}"
+            raise ValueError(msg)
+
+    return sif_coefficient, token_remove_regex
+
+
+def distill(
+    model_name: str,
+    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,
+) -> StaticModel:
+    """
+    Distill a staticmodel from a sentence transformer.
+
+    This function creates a set of embeddings from a sentence transformer. It does this by doing either
+    a forward pass for all subword tokens in the tokenizer, or by doing a forward pass for all tokens in a passed vocabulary.
+
+    If you pass through a vocabulary, we create a custom word tokenizer for that vocabulary.
+    If you don't pass a vocabulary, we use the model's tokenizer directly.
+
+    :param model_name: The model name to use. Any sentencetransformer compatible model works.
+    :param vocabulary: The vocabulary to use. If this is None, we use the model's vocabulary.
+    :param device: The device to use.
+    :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.
+    :return: A StaticModel
+
+    """
+    model: PreTrainedModel = AutoModel.from_pretrained(model_name, trust_remote_code=trust_remote_code)
+    tokenizer = cast(
+        "PreTrainedTokenizerFast",
+        AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code, use_fast=True),
+    )
+
+    return distill_from_model(
+        model=model,
+        tokenizer=tokenizer,
+        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,
+    )

src/distiller/model2vec/distill/inference.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import inspect
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, Protocol, Union
+
+import numpy as np
+import torch
+from sklearn.decomposition import PCA
+from torch.nn.utils.rnn import pad_sequence
+from tqdm import tqdm
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+    from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
+
+logger = logging.getLogger(__name__)
+
+
+PathLike = Union[Path, str]
+PCADimType = Union[int, None, float, Literal["auto"]]
+
+
+_DEFAULT_BATCH_SIZE = 256
+
+
+class ModulewithWeights(Protocol):
+    weight: torch.nn.Parameter
+
+
+def create_embeddings(
+    model: PreTrainedModel,
+    tokenized: list[list[int]],
+    device: str,
+    pad_token_id: int,
+) -> np.ndarray:
+    """
+    Create output embeddings for a bunch of tokens using a pretrained model.
+
+    It does a forward pass for all tokens passed in `tokens`.
+
+    :param model: The model to use.
+        This should be a transformers model.
+    :param tokenized: All tokenized tokens.
+    :param device: The torch device to use.
+    :param pad_token_id: The pad token id. Used to pad sequences.
+    :return: The output embeddings.
+    """
+    model = model.to(device)
+
+    out_weights: np.ndarray
+    intermediate_weights: list[np.ndarray] = []
+
+    # Add token_type_ids only if the model supports it
+    add_token_type_ids = "token_type_ids" in inspect.getfullargspec(model.forward).args
+
+    lengths = np.asarray([len(sequence) for sequence in tokenized])
+    sort_order = np.argsort(lengths)
+
+    sorted_tokenized = [tokenized[i] for i in sort_order]
+
+    pbar = tqdm(total=len(sorted_tokenized), desc="Encoding tokens", unit=" tokens")
+
+    for batch_idx in range(0, len(sorted_tokenized), _DEFAULT_BATCH_SIZE):
+        batch = [torch.Tensor(x).long() for x in sorted_tokenized[batch_idx : batch_idx + _DEFAULT_BATCH_SIZE]]
+
+        encoded = {}
+        encoded["input_ids"] = pad_sequence(batch, batch_first=True, padding_value=pad_token_id)
+        encoded["attention_mask"] = encoded["input_ids"] != pad_token_id
+
+        if add_token_type_ids:
+            encoded["token_type_ids"] = torch.zeros_like(encoded["input_ids"])
+
+        out = _encode_mean_using_model(model, encoded)
+        intermediate_weights.extend(out.numpy())
+        pbar.update(len(batch))
+
+    # Sort the output back to the original order
+    intermediate_weights = [intermediate_weights[i] for i in np.argsort(sort_order)]
+    out_weights = np.stack(intermediate_weights)
+
+    out_weights = np.nan_to_num(out_weights)
+
+    return out_weights
+
+
+@torch.no_grad()
+def _encode_mean_using_model(model: PreTrainedModel, encodings: dict[str, torch.Tensor]) -> torch.Tensor:
+    """
+    Encode a batch of tokens using a model.
+
+    Note that if a token in the input batch does not have any embeddings, it will be output as a vector of zeros.
+    So detection of these is necessary.
+
+    :param model: The model to use.
+    :param encodings: The encoded tokens to turn into features.
+    :return: The mean of the output for each token.
+    """
+    encodings = {k: v.to(model.device) for k, v in encodings.items()}
+    encoded: BaseModelOutputWithPoolingAndCrossAttentions = model(**encodings)
+    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
+    # See here: https://github.com/numpy/numpy/issues/19808
+    if out.dtype == torch.bfloat16:
+        out = out.float()
+
+    # Take the mean by averaging over the attention mask.
+    mask = encodings["attention_mask"].cpu().float()
+    mask /= mask.sum(1)[:, None]
+
+    return torch.bmm(mask[:, None, :].float(), out).squeeze(1)
+
+
+
+def post_process_embeddings(
+    embeddings: np.ndarray, pca_dims: PCADimType, sif_coefficient: float | None = 1e-4
+) -> np.ndarray:
+    """Post process embeddings by applying PCA and SIF weighting by estimating the frequencies through Zipf's law."""
+    if pca_dims is not None:
+        if pca_dims == "auto":
+            pca_dims = embeddings.shape[1]
+        if pca_dims > embeddings.shape[1]:
+            logger.warning(
+                f"PCA dimension ({pca_dims}) is larger than the number of dimensions in the embeddings ({embeddings.shape[1]}). "
+                "Applying PCA, but not reducing dimensionality. Is this is not desired, please set `pca_dims` to None. "
+                "Applying PCA will probably improve performance, so consider just leaving it."
+            )
+            pca_dims = embeddings.shape[1]
+        if pca_dims >= embeddings.shape[0]:
+            logger.warning(
+                f"PCA dimension ({pca_dims}) is larger than the number of tokens in the vocabulary ({embeddings.shape[0]}). Not applying PCA."
+            )
+        elif pca_dims <= embeddings.shape[1]:
+            if isinstance(pca_dims, float):
+                logger.info(f"Applying PCA with {pca_dims} explained variance.")
+            else:
+                logger.info(f"Applying PCA with n_components {pca_dims}")
+
+            orig_dims = embeddings.shape[1]
+            p = PCA(n_components=pca_dims, svd_solver="full")
+            embeddings = p.fit_transform(embeddings)
+
+            if embeddings.shape[1] < orig_dims:
+                explained_variance_ratio = np.sum(p.explained_variance_ratio_)
+                explained_variance = np.sum(p.explained_variance_)
+                logger.info(f"Reduced dimensionality from {orig_dims} to {embeddings.shape[1]}.")
+                logger.info(f"Explained variance ratio: {explained_variance_ratio:.3f}.")
+                logger.info(f"Explained variance: {explained_variance:.3f}.")
+
+    if sif_coefficient is not None:
+        logger.info("Estimating word frequencies using Zipf's law, and then applying SIF.")
+        inv_rank = 1 / (np.arange(2, embeddings.shape[0] + 2))
+        proba = inv_rank / np.sum(inv_rank)
+        embeddings *= (sif_coefficient / (sif_coefficient + proba))[:, None]
+
+    return embeddings

src/distiller/model2vec/distill/utils.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from logging import getLogger
+
+import torch
+
+logger = getLogger(__name__)
+
+
+def select_optimal_device(device: str | None) -> str:
+    """
+    Guess what your optimal device should be based on backend availability.
+
+    If you pass a device, we just pass it through.
+
+    :param device: The device to use. If this is not None you get back what you passed.
+    :return: The selected device.
+    """
+    if device is None:
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+        logger.info(f"Automatically selected device: {device}")
+
+    return device

src/distiller/model2vec/hf_utils.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import huggingface_hub
+import safetensors
+from huggingface_hub import ModelCard, ModelCardData
+from safetensors.numpy import save_file
+from tokenizers import Tokenizer
+
+if TYPE_CHECKING:
+    import numpy as np
+
+    from distiller.model2vec.utils import SafeOpenProtocol
+
+logger = logging.getLogger(__name__)
+
+
+def save_pretrained(
+    folder_path: Path,
+    embeddings: np.ndarray,
+    tokenizer: Tokenizer,
+    config: dict[str, Any],
+    create_model_card: bool = True,
+    subfolder: str | None = None,
+    **kwargs: Any,
+) -> None:
+    """
+    Save a model to a folder.
+
+    :param folder_path: The path to the folder.
+    :param embeddings: The embeddings.
+    :param tokenizer: The tokenizer.
+    :param config: A metadata config.
+    :param create_model_card: Whether to create a model card.
+    :param subfolder: The subfolder to save the model in.
+    :param **kwargs: Any additional arguments.
+    """
+    folder_path = folder_path / subfolder if subfolder else folder_path
+    folder_path.mkdir(exist_ok=True, parents=True)
+    save_file({"embeddings": embeddings}, folder_path / "model.safetensors")
+    tokenizer.save(str(folder_path / "tokenizer.json"), pretty=False)
+    json.dump(config, open(folder_path / "config.json", "w"), indent=4)
+
+    # Create modules.json
+    modules = [{"idx": 0, "name": "0", "path": ".", "type": "sentence_transformers.models.StaticEmbedding"}]
+    if config.get("normalize"):
+        # If normalize=True, add sentence_transformers.models.Normalize
+        modules.append({"idx": 1, "name": "1", "path": "1_Normalize", "type": "sentence_transformers.models.Normalize"})
+    json.dump(modules, open(folder_path / "modules.json", "w"), indent=4)
+
+    logger.info(f"Saved model to {folder_path}")
+
+    # Optionally create the model card
+    if create_model_card:
+        _create_model_card(folder_path, **kwargs)
+
+
+def _create_model_card(
+    folder_path: Path,
+    base_model_name: str = "unknown",
+    license: str = "mit",
+    language: list[str] | None = None,
+    model_name: str | None = None,
+    template_path: str = "modelcards/model_card_template.md",
+    **kwargs: Any,
+) -> None:
+    """
+    Create a model card and store it in the specified path.
+
+    :param folder_path: The path where the model card will be stored.
+    :param base_model_name: The name of the base model.
+    :param license: The license to use.
+    :param language: The language of the model.
+    :param model_name: The name of the model to use in the Model Card.
+    :param template_path: The path to the template.
+    :param **kwargs: Additional metadata for the model card (e.g., model_name, base_model, etc.).
+    """
+    folder_path = Path(folder_path)
+    model_name = model_name or folder_path.name
+    full_path = Path(__file__).parent / template_path
+
+    model_card_data = ModelCardData(
+        model_name=model_name,
+        base_model=base_model_name,
+        license=license,
+        language=language,
+        tags=["embeddings", "static-embeddings", "sentence-transformers"],
+        library_name="model2vec",
+        **kwargs,
+    )
+    model_card = ModelCard.from_template(model_card_data, template_path=str(full_path))
+    model_card.save(folder_path / "README.md")
+
+
+def load_pretrained(
+    folder_or_repo_path: str | Path,
+    subfolder: str | None = None,
+    token: str | None = None,
+    from_sentence_transformers: bool = False,
+) -> tuple[np.ndarray, Tokenizer, dict[str, Any], dict[str, Any]]:
+    """
+    Loads a pretrained model from a folder.
+
+    :param folder_or_repo_path: The folder or repo path to load from.
+        - If this is a local path, we will load from the local path.
+        - If the local path is not found, we will attempt to load from the huggingface hub.
+    :param subfolder: The subfolder to load from.
+    :param token: The huggingface token to use.
+    :param from_sentence_transformers: Whether to load the model from a sentence transformers model.
+    :raises: FileNotFoundError if the folder exists, but the file does not exist locally.
+    :return: The embeddings, tokenizer, config, and metadata.
+
+    """
+    if from_sentence_transformers:
+        model_file = "0_StaticEmbedding/model.safetensors"
+        tokenizer_file = "0_StaticEmbedding/tokenizer.json"
+        config_name = "config_sentence_transformers.json"
+    else:
+        model_file = "model.safetensors"
+        tokenizer_file = "tokenizer.json"
+        config_name = "config.json"
+
+    folder_or_repo_path = Path(folder_or_repo_path)
+
+    local_folder = folder_or_repo_path / subfolder if subfolder else folder_or_repo_path
+
+    if local_folder.exists():
+        embeddings_path = local_folder / model_file
+        if not embeddings_path.exists():
+            msg = f"Embeddings file does not exist in {local_folder}"
+            raise FileNotFoundError(msg)
+
+        config_path = local_folder / config_name
+        if not config_path.exists():
+            msg = f"Config file does not exist in {local_folder}"
+            raise FileNotFoundError(msg)
+
+        tokenizer_path = local_folder / tokenizer_file
+        if not tokenizer_path.exists():
+            msg = f"Tokenizer file does not exist in {local_folder}"
+            raise FileNotFoundError(msg)
+
+        # README is optional, so this is a bit finicky.
+        readme_path = local_folder / "README.md"
+        metadata = _get_metadata_from_readme(readme_path)
+
+    else:
+        logger.info("Folder does not exist locally, attempting to use huggingface hub.")
+        embeddings_path = Path(
+            huggingface_hub.hf_hub_download(
+                folder_or_repo_path.as_posix(), model_file, token=token, subfolder=subfolder
+            )
+        )
+
+        try:
+            readme_path = Path(
+                huggingface_hub.hf_hub_download(
+                    folder_or_repo_path.as_posix(), "README.md", token=token, subfolder=subfolder
+                )
+            )
+            metadata = _get_metadata_from_readme(Path(readme_path))
+        except Exception as e:
+            # NOTE: we don't want to raise an error here, since the README is optional.
+            logger.info(f"No README found in the model folder: {e} No model card loaded.")
+            metadata = {}
+
+        config_path = Path(
+            huggingface_hub.hf_hub_download(
+                folder_or_repo_path.as_posix(), config_name, token=token, subfolder=subfolder
+            )
+        )
+        tokenizer_path = Path(
+            huggingface_hub.hf_hub_download(
+                folder_or_repo_path.as_posix(), tokenizer_file, token=token, subfolder=subfolder
+            )
+        )
+
+    opened_tensor_file = cast("SafeOpenProtocol", safetensors.safe_open(embeddings_path, framework="numpy"))
+    if from_sentence_transformers:
+        embeddings = opened_tensor_file.get_tensor("embedding.weight")
+    else:
+        embeddings = opened_tensor_file.get_tensor("embeddings")
+
+    tokenizer: Tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    config = json.load(open(config_path))
+
+    if len(tokenizer.get_vocab()) != len(embeddings):
+        logger.warning(
+            f"Number of tokens does not match number of embeddings: `{len(tokenizer.get_vocab())}` vs `{len(embeddings)}`"
+        )
+
+    return embeddings, tokenizer, config, metadata
+
+
+def _get_metadata_from_readme(readme_path: Path) -> dict[str, Any]:
+    """Get metadata from a README file."""
+    if not readme_path.exists():
+        logger.info(f"README file not found in {readme_path}. No model card loaded.")
+        return {}
+    model_card = ModelCard.load(readme_path)
+    data: dict[str, Any] = model_card.data.to_dict()
+    if not data:
+        logger.info("File README.md exists, but was empty. No model card loaded.")
+    return data
+
+
+def push_folder_to_hub(
+    folder_path: Path, subfolder: str | None, repo_id: str, private: bool, token: str | None
+) -> None:
+    """
+    Push a model folder to the huggingface hub, including model card.
+
+    :param folder_path: The path to the folder.
+    :param subfolder: The subfolder to push to.
+        If None, the folder will be pushed to the root of the repo.
+    :param repo_id: The repo name.
+    :param private: Whether the repo is private.
+    :param token: The huggingface token.
+    """
+    if not huggingface_hub.repo_exists(repo_id=repo_id, token=token):
+        huggingface_hub.create_repo(repo_id, token=token, private=private)
+
+    # Push model card and all model files to the Hugging Face hub
+    huggingface_hub.upload_folder(repo_id=repo_id, folder_path=folder_path, token=token, path_in_repo=subfolder)
+
+    logger.info(f"Pushed model to {repo_id}")

src/distiller/model2vec/inference/README.md

@@ -0,0 +1,18 @@
+# Inference
+
+This subpackage mainly contains helper functions for inference with trained models that have been exported to `scikit-learn` compatible pipelines.
+
+If you're looking for information on how to train a model, see [here](../train/README.md).
+
+# Usage
+
+Let's assume you're using our [potion-edu classifier](https://huggingface.co/minishlab/potion-8m-edu-classifier).
+
+```python
+from model2vec.inference import StaticModelPipeline
+
+classifier = StaticModelPipeline.from_pretrained("minishlab/potion-8m-edu-classifier")
+label = classifier.predict("Attitudes towards cattle in the Alps: a study in letting go.")
+```
+
+This should just work.

src/distiller/model2vec/inference/__init__.py

@@ -0,0 +1,10 @@
+from distiller.model2vec.utils import get_package_extras, importable
+
+_REQUIRED_EXTRA = "inference"
+
+for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA):
+    importable(extra_dependency, _REQUIRED_EXTRA)
+
+from distiller.model2vec.inference.model import StaticModelPipeline, evaluate_single_or_multi_label
+
+__all__ = ["StaticModelPipeline", "evaluate_single_or_multi_label"]

src/distiller/model2vec/inference/model.py

@@ -0,0 +1,312 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, TypeVar
+
+import huggingface_hub
+import numpy as np
+import skops.io
+from sklearn.metrics import classification_report
+from sklearn.neural_network import MLPClassifier
+from sklearn.preprocessing import MultiLabelBinarizer
+
+from distiller.model2vec.hf_utils import _create_model_card
+from distiller.model2vec.model import PathLike, StaticModel
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sklearn.pipeline import Pipeline
+
+_DEFAULT_TRUST_PATTERN = re.compile(r"sklearn\..+")
+_DEFAULT_MODEL_FILENAME = "pipeline.skops"
+
+LabelType = TypeVar("LabelType", list[str], list[list[str]])
+
+
+class StaticModelPipeline:
+    def __init__(self, model: StaticModel, head: Pipeline) -> None:
+        """Create a pipeline with a StaticModel encoder."""
+        self.model = model
+        self.head = head
+        classifier = self.head[-1]
+        # Check if the classifier is a multilabel classifier.
+        # NOTE: this doesn't look robust, but it is.
+        # Different classifiers, such as OVR wrappers, support multilabel output natively, so we
+        # can just use predict.
+        self.multilabel = False
+        if isinstance(classifier, MLPClassifier) and classifier.out_activation_ == "logistic":
+            self.multilabel = True
+
+    @property
+    def classes_(self) -> np.ndarray:
+        """The classes of the classifier."""
+        return self.head.classes_
+
+    @classmethod
+    def from_pretrained(
+        cls: type[StaticModelPipeline], path: PathLike, token: str | None = None, trust_remote_code: bool = False
+    ) -> StaticModelPipeline:
+        """
+        Load a StaticModel from a local path or huggingface hub path.
+
+        NOTE: if you load a private model from the huggingface hub, you need to pass a token.
+
+        :param path: The path to the folder containing the pipeline, or a repository on the Hugging Face Hub
+        :param token: The token to use to download the pipeline from the hub.
+        :param trust_remote_code: Whether to trust the remote code. If this is False, we will only load components coming from `sklearn`.
+        :return: The loaded pipeline.
+        """
+        model, head = _load_pipeline(path, token, trust_remote_code)
+        model.embedding = np.nan_to_num(model.embedding)
+
+        return cls(model, head)
+
+    def save_pretrained(self, path: str) -> None:
+        """Save the model to a folder."""
+        save_pipeline(self, path)
+
+    def push_to_hub(
+        self, repo_id: str, subfolder: str | None = None, token: str | None = None, private: bool = False
+    ) -> None:
+        """
+        Save a model to a folder, and then push that folder to the hf hub.
+
+        :param repo_id: The id of the repository to push to.
+        :param subfolder: The subfolder to push to.
+        :param token: The token to use to push to the hub.
+        :param private: Whether the repository should be private.
+        """
+        from distiller.model2vec.hf_utils import push_folder_to_hub
+
+        with TemporaryDirectory() as temp_dir:
+            save_pipeline(self, temp_dir)
+            self.model.save_pretrained(temp_dir)
+            push_folder_to_hub(Path(temp_dir), subfolder, repo_id, private, token)
+
+    def _encode_and_coerce_to_2d(
+        self,
+        X: Sequence[str],
+        show_progress_bar: bool,
+        max_length: int | None,
+        batch_size: int,
+        use_multiprocessing: bool,
+        multiprocessing_threshold: int,
+    ) -> np.ndarray:
+        """Encode the instances and coerce the output to a matrix."""
+        encoded = self.model.encode(
+            X,
+            show_progress_bar=show_progress_bar,
+            max_length=max_length,
+            batch_size=batch_size,
+            use_multiprocessing=use_multiprocessing,
+            multiprocessing_threshold=multiprocessing_threshold,
+        )
+        if np.ndim(encoded) == 1:
+            encoded = encoded[None, :]
+
+        return encoded
+
+    def predict(
+        self,
+        X: Sequence[str],
+        show_progress_bar: bool = False,
+        max_length: int | None = 512,
+        batch_size: int = 1024,
+        use_multiprocessing: bool = True,
+        multiprocessing_threshold: int = 10_000,
+        threshold: float = 0.5,
+    ) -> np.ndarray:
+        """
+        Predict the labels of the input.
+
+        :param X: The input data to predict. Can be a list of strings or a single string.
+        :param show_progress_bar: Whether to display a progress bar during prediction. Defaults to False.
+        :param max_length: The maximum length of the input sequences. Defaults to 512.
+        :param batch_size: The batch size for prediction. Defaults to 1024.
+        :param use_multiprocessing: Whether to use multiprocessing for encoding. Defaults to True.
+        :param multiprocessing_threshold: The threshold for the number of samples to use multiprocessing. Defaults to 10,000.
+        :param threshold: The threshold for multilabel classification. Defaults to 0.5. Ignored if not multilabel.
+        :return: The predicted labels or probabilities.
+        """
+        encoded = self._encode_and_coerce_to_2d(
+            X,
+            show_progress_bar=show_progress_bar,
+            max_length=max_length,
+            batch_size=batch_size,
+            use_multiprocessing=use_multiprocessing,
+            multiprocessing_threshold=multiprocessing_threshold,
+        )
+
+        if self.multilabel:
+            out_labels = []
+            proba = self.head.predict_proba(encoded)
+            for vector in proba:
+                out_labels.append(self.classes_[vector > threshold])
+            return np.asarray(out_labels, dtype=object)
+
+        return self.head.predict(encoded)
+
+    def predict_proba(
+        self,
+        X: Sequence[str],
+        show_progress_bar: bool = False,
+        max_length: int | None = 512,
+        batch_size: int = 1024,
+        use_multiprocessing: bool = True,
+        multiprocessing_threshold: int = 10_000,
+    ) -> np.ndarray:
+        """
+        Predict the labels of the input.
+
+        :param X: The input data to predict. Can be a list of strings or a single string.
+        :param show_progress_bar: Whether to display a progress bar during prediction. Defaults to False.
+        :param max_length: The maximum length of the input sequences. Defaults to 512.
+        :param batch_size: The batch size for prediction. Defaults to 1024.
+        :param use_multiprocessing: Whether to use multiprocessing for encoding. Defaults to True.
+        :param multiprocessing_threshold: The threshold for the number of samples to use multiprocessing. Defaults to 10,000.
+        :return: The predicted labels or probabilities.
+        """
+        encoded = self._encode_and_coerce_to_2d(
+            X,
+            show_progress_bar=show_progress_bar,
+            max_length=max_length,
+            batch_size=batch_size,
+            use_multiprocessing=use_multiprocessing,
+            multiprocessing_threshold=multiprocessing_threshold,
+        )
+
+        return self.head.predict_proba(encoded)
+
+    def evaluate(
+        self, X: Sequence[str], y: LabelType, batch_size: int = 1024, threshold: float = 0.5, output_dict: bool = False
+    ) -> str | dict[str, dict[str, float]]:
+        """
+        Evaluate the classifier on a given dataset using scikit-learn's classification report.
+
+        :param X: The texts to predict on.
+        :param y: The ground truth labels.
+        :param batch_size: The batch size.
+        :param threshold: The threshold for multilabel classification.
+        :param output_dict: Whether to output the classification report as a dictionary.
+        :return: A classification report.
+        """
+        predictions = self.predict(X, show_progress_bar=True, batch_size=batch_size, threshold=threshold)
+        return evaluate_single_or_multi_label(predictions=predictions, y=y, output_dict=output_dict)
+
+
+
+def _load_pipeline(
+    folder_or_repo_path: PathLike, token: str | None = None, trust_remote_code: bool = False
+) -> tuple[StaticModel, Pipeline]:
+    """
+    Load a model and an sklearn pipeline.
+
+    This assumes the following files are present in the repo:
+    - `pipeline.skops`: The head of the pipeline.
+    - `config.json`: The configuration of the model.
+    - `model.safetensors`: The weights of the model.
+    - `tokenizer.json`: The tokenizer of the model.
+
+    :param folder_or_repo_path: The path to the folder containing the pipeline.
+    :param token: The token to use to download the pipeline from the hub. If this is None, you will only
+        be able to load the pipeline from a local folder, public repository, or a repository that you have access to
+        because you are logged in.
+    :param trust_remote_code: Whether to trust the remote code. If this is False,
+        we will only load components coming from `sklearn`. If this is True, we will load all components.
+        If you set this to True, you are responsible for whatever happens.
+    :return: The encoder model and the loaded head
+    :raises FileNotFoundError: If the pipeline file does not exist in the folder.
+    :raises ValueError: If an untrusted type is found in the pipeline, and `trust_remote_code` is False.
+    """
+    folder_or_repo_path = Path(folder_or_repo_path)
+    model_filename = _DEFAULT_MODEL_FILENAME
+    head_pipeline_path: str | Path
+    if folder_or_repo_path.exists():
+        head_pipeline_path = folder_or_repo_path / model_filename
+        if not head_pipeline_path.exists():
+            msg = f"Pipeline file does not exist in {folder_or_repo_path}"
+            raise FileNotFoundError(msg)
+    else:
+        head_pipeline_path = huggingface_hub.hf_hub_download(
+            folder_or_repo_path.as_posix(), model_filename, token=token
+        )
+
+    model = StaticModel.from_pretrained(folder_or_repo_path)
+
+    unknown_types = skops.io.get_untrusted_types(file=head_pipeline_path)
+    # If the user does not trust remote code, we should check that the unknown types are trusted.
+    # By default, we trust everything coming from scikit-learn.
+    if not trust_remote_code:
+        for t in unknown_types:
+            if not _DEFAULT_TRUST_PATTERN.match(t):
+                msg = f"Untrusted type {t}."
+                raise ValueError(msg)
+    head = skops.io.load(head_pipeline_path, trusted=unknown_types)
+
+    return model, head
+
+
+def save_pipeline(pipeline: StaticModelPipeline, folder_path: str | Path) -> None:
+    """
+    Save a pipeline to a folder.
+
+    :param pipeline: The pipeline to save.
+    :param folder_path: The path to the folder to save the pipeline to.
+    """
+    folder_path = Path(folder_path)
+    folder_path.mkdir(parents=True, exist_ok=True)
+    model_filename = _DEFAULT_MODEL_FILENAME
+    head_pipeline_path = folder_path / model_filename
+    skops.io.dump(pipeline.head, head_pipeline_path)
+    pipeline.model.save_pretrained(folder_path)
+    base_model_name = pipeline.model.base_model_name
+    if isinstance(base_model_name, list) and base_model_name:
+        name = base_model_name[0]
+    elif isinstance(base_model_name, str):
+        name = base_model_name
+    else:
+        name = "unknown"
+    _create_model_card(
+        folder_path,
+        base_model_name=name,
+        language=pipeline.model.language,
+        template_path="modelcards/classifier_template.md",
+    )
+
+
+def _is_multi_label_shaped(y: LabelType) -> bool:
+    """Check if the labels are in a multi-label shape."""
+    return isinstance(y, (list, tuple)) and len(y) > 0 and isinstance(y[0], (list, tuple, set))
+
+
+def evaluate_single_or_multi_label(
+    predictions: np.ndarray,
+    y: LabelType,
+    output_dict: bool = False,
+) -> str | dict[str, dict[str, float]]:
+    """
+    Evaluate the classifier on a given dataset using scikit-learn's classification report.
+
+    :param predictions: The predictions.
+    :param y: The ground truth labels.
+    :param output_dict: Whether to output the classification report as a dictionary.
+    :return: A classification report.
+    """
+    if _is_multi_label_shaped(y):
+        classes = sorted({label for labels in y for label in labels})
+        mlb = MultiLabelBinarizer(classes=classes)
+        y = mlb.fit_transform(y)
+        predictions = mlb.transform(predictions)
+    elif isinstance(y[0], (str, int)):
+        classes = sorted(set(y))
+
+    return classification_report(
+        y,
+        predictions,
+        output_dict=output_dict,
+        zero_division=0,
+    )
+

src/distiller/model2vec/model.py

@@ -0,0 +1,480 @@
+from __future__ import annotations
+
+import math
+import os
+from logging import getLogger
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any, Union, overload
+
+import numpy as np
+from joblib import delayed
+from tqdm import tqdm
+
+from .quantization import DType, quantize_and_reduce_dim
+from .utils import ProgressParallel, load_local_model
+
+if TYPE_CHECKING:
+	from collections.abc import Iterator, Sequence
+
+	from tokenizers import Encoding, Tokenizer
+
+PathLike = Union[Path, str]
+
+logger = getLogger(__name__)
+
+
+class StaticModel:
+	def __init__(
+		self,
+		vectors: np.ndarray,
+		tokenizer: Tokenizer,
+		config: dict[str, Any] | None = None,
+		normalize: bool | None = None,
+		base_model_name: str | None = None,
+		language: list[str] | None = None,
+	) -> None:
+		"""
+		Initialize the StaticModel.
+
+		:param vectors: The vectors to use.
+		:param tokenizer: The Transformers tokenizer to use.
+		:param config: Any metadata config.
+		:param normalize: Whether to normalize the embeddings.
+		:param base_model_name: The used base model name. Used for creating a model card.
+		:param language: The language of the model. Used for creating a model card.
+		:raises: ValueError if the number of tokens does not match the number of vectors.
+		"""
+		super().__init__()
+		tokens, _ = zip(*sorted(tokenizer.get_vocab().items(), key=lambda x: x[1]), strict=False)
+		self.tokens = tokens
+
+		self.embedding = vectors
+
+		if len(tokens) != vectors.shape[0]:
+			msg = f"Number of tokens ({len(tokens)}) does not match number of vectors ({vectors.shape[0]})"
+			raise ValueError(msg)
+
+		self.tokenizer = tokenizer
+		self.unk_token_id: int | None
+		if hasattr(self.tokenizer.model, "unk_token") and self.tokenizer.model.unk_token is not None:
+			self.unk_token_id = tokenizer.get_vocab()[self.tokenizer.model.unk_token]
+		else:
+			self.unk_token_id = None  # pragma: no cover  # Doesn't actually happen, but can happen.
+
+		self.median_token_length = int(np.median([len(token) for token in self.tokens]))
+		self.config = config or {}
+		self.base_model_name = base_model_name
+		self.language = language
+		if hasattr(self.tokenizer, "encode_batch_fast"):
+			self._can_encode_fast = True
+		else:
+			self._can_encode_fast = False
+
+		if normalize is not None:
+			self.normalize = normalize
+		else:
+			self.normalize = self.config.get("normalize", False)
+
+	@property
+	def dim(self) -> int:
+		"""Get the dimension of the model."""
+		return self.embedding.shape[1]
+
+	@property
+	def normalize(self) -> bool:
+		"""
+		Get the normalize value.
+
+		:return: The normalize value.
+		"""
+		return self._normalize
+
+	@normalize.setter
+	def normalize(self, value: bool) -> None:
+		"""Update the config if the value of normalize changes."""
+		config_normalize = self.config.get("normalize")
+		self._normalize = value
+		if config_normalize is not None and value != config_normalize:
+			logger.warning(
+				f"Set normalization to `{value}`, which does not match config value `{config_normalize}`. Updating config."
+			)
+		self.config["normalize"] = value
+
+	def save_pretrained(self, path: PathLike, model_name: str | None = None, subfolder: str | None = None) -> None:
+		"""
+		Save the pretrained model.
+
+		:param path: The path to save to.
+		:param model_name: The model name to use in the Model Card.
+		:param subfolder: The subfolder to save to.
+		"""
+		from .hf_utils import save_pretrained
+
+		save_pretrained(
+			folder_path=Path(path),
+			embeddings=self.embedding,
+			tokenizer=self.tokenizer,
+			config=self.config,
+			base_model_name=self.base_model_name,
+			language=self.language,
+			model_name=model_name,
+			subfolder=subfolder,
+		)
+
+	def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> list[list[int]]:
+		"""
+		Tokenize a list of sentences.
+
+		:param sentences: The sentences to tokenize.
+		:param max_length: The maximum length of the sentences in tokens. If this is None, sequences
+		    are not truncated.
+		:return: A list of list of tokens.
+		"""
+		if max_length is not None:
+			m = max_length * self.median_token_length
+			sentences = [sentence[:m] for sentence in sentences]
+
+		if self._can_encode_fast:
+			encodings: list[Encoding] = self.tokenizer.encode_batch_fast(sentences, add_special_tokens=False)
+		else:
+			encodings = self.tokenizer.encode_batch(sentences, add_special_tokens=False)
+
+		encodings_ids = [encoding.ids for encoding in encodings]
+
+		if self.unk_token_id is not None:
+			# NOTE: Remove the unknown token: necessary for word-level models.
+			encodings_ids = [
+				[token_id for token_id in token_ids if token_id != self.unk_token_id] for token_ids in encodings_ids
+			]
+		if max_length is not None:
+			encodings_ids = [token_ids[:max_length] for token_ids in encodings_ids]
+
+		return encodings_ids
+
+	@classmethod
+	def from_pretrained(
+		cls: type[StaticModel],
+		path: PathLike,
+		token: str | None = None,
+		normalize: bool | None = None,
+		subfolder: str | None = None,
+		quantize_to: str | DType | None = None,
+		dimensionality: int | None = None,
+	) -> StaticModel:
+		"""
+		Load a StaticModel from a local path or huggingface hub path.
+
+		NOTE: if you load a private model from the huggingface hub, you need to pass a token.
+
+		:param path: The path to load your static model from.
+		:param token: The huggingface token to use.
+		:param normalize: Whether to normalize the embeddings.
+		:param subfolder: The subfolder to load from.
+		:param quantize_to: The dtype to quantize the model to. If None, no quantization is done.
+		    If a string is passed, it is converted to a DType.
+		:param dimensionality: The dimensionality of the model. If this is None, use the dimensionality of the model.
+		    This is useful if you want to load a model with a lower dimensionality.
+		    Note that this only applies if you have trained your model using mrl or PCA.
+		:return: A StaticModel.
+		"""
+		from .hf_utils import load_pretrained
+
+		embeddings, tokenizer, config, metadata = load_pretrained(
+			folder_or_repo_path=path,
+			token=token,
+			from_sentence_transformers=False,
+			subfolder=subfolder,
+		)
+
+		embeddings = quantize_and_reduce_dim(
+			embeddings=embeddings,
+			quantize_to=quantize_to,
+			dimensionality=dimensionality,
+		)
+
+		return cls(
+			embeddings,
+			tokenizer,
+			config,
+			normalize=normalize,
+			base_model_name=metadata.get("base_model"),
+			language=metadata.get("language"),
+		)
+
+	@classmethod
+	def from_sentence_transformers(
+		cls: type[StaticModel],
+		path: PathLike,
+		token: str | None = None,
+		normalize: bool | None = None,
+		quantize_to: str | DType | None = None,
+		dimensionality: int | None = None,
+	) -> StaticModel:
+		"""
+		Load a StaticModel trained with sentence transformers from a local path or huggingface hub path.
+
+		NOTE: if you load a private model from the huggingface hub, you need to pass a token.
+
+		:param path: The path to load your static model from.
+		:param token: The huggingface token to use.
+		:param normalize: Whether to normalize the embeddings.
+		:param quantize_to: The dtype to quantize the model to. If None, no quantization is done.
+		    If a string is passed, it is converted to a DType.
+		:param dimensionality: The dimensionality of the model. If this is None, use the dimensionality of the model.
+		    This is useful if you want to load a model with a lower dimensionality.
+		    Note that this only applies if you have trained your model using mrl or PCA.
+		:return: A StaticModel.
+		"""
+		from .hf_utils import load_pretrained
+
+		embeddings, tokenizer, config, metadata = load_pretrained(
+			folder_or_repo_path=path,
+			token=token,
+			from_sentence_transformers=True,
+			subfolder=None,
+		)
+
+		embeddings = quantize_and_reduce_dim(
+			embeddings=embeddings,
+			quantize_to=quantize_to,
+			dimensionality=dimensionality,
+		)
+
+		return cls(
+			embeddings,
+			tokenizer,
+			config,
+			normalize=normalize,
+			base_model_name=metadata.get("base_model"),
+			language=metadata.get("language"),
+		)
+
+	@overload
+	def encode_as_sequence(
+		self,
+		sentences: str,
+		max_length: int | None = None,
+		batch_size: int = 1024,
+		show_progress_bar: bool = False,
+		use_multiprocessing: bool = True,
+		multiprocessing_threshold: int = 10_000,
+	) -> np.ndarray: ...
+
+	@overload
+	def encode_as_sequence(
+		self,
+		sentences: list[str],
+		max_length: int | None = None,
+		batch_size: int = 1024,
+		show_progress_bar: bool = False,
+		use_multiprocessing: bool = True,
+		multiprocessing_threshold: int = 10_000,
+	) -> list[np.ndarray]: ...
+
+	def encode_as_sequence(
+		self,
+		sentences: str | list[str],
+		max_length: int | None = None,
+		batch_size: int = 1024,
+		show_progress_bar: bool = False,
+		use_multiprocessing: bool = True,
+		multiprocessing_threshold: int = 10_000,
+	) -> list[np.ndarray] | np.ndarray:
+		"""
+		Encode a list of sentences as a list of numpy arrays of tokens.
+
+		This is useful if you want to use the tokens for further processing, or if you want to do sequence
+		modeling.
+		Note that if you just want the mean, you should use the `encode` method.
+		This is about twice as slow.
+		Sentences that do not contain any tokens will be turned into an empty array.
+
+		NOTE: the input type is currently underspecified. The actual input type is `Sequence[str] | str`, but this
+		    is not possible to implement in python typing currently.
+
+		:param sentences: The list of sentences to encode.
+		:param max_length: The maximum length of the sentences. Any tokens beyond this length will be truncated.
+		    If this is None, no truncation is done.
+		:param batch_size: The batch size to use.
+		:param show_progress_bar: Whether to show the progress bar.
+		:param use_multiprocessing: Whether to use multiprocessing.
+		    By default, this is enabled for inputs > multiprocessing_threshold sentences and disabled otherwise.
+		:param multiprocessing_threshold: The threshold in number of sentences for using multiprocessing.
+		:return: The encoded sentences with an embedding per token.
+		"""
+		was_single = False
+		if isinstance(sentences, str):
+			sentences = [sentences]
+			was_single = True
+
+		# Prepare all batches
+		sentence_batches = list(self._batch(sentences, batch_size))
+		total_batches = math.ceil(len(sentences) / batch_size)
+
+		# Use joblib for multiprocessing if requested, and if we have enough sentences
+		if use_multiprocessing and len(sentences) > multiprocessing_threshold:
+			# Disable parallelism for tokenizers
+			os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+			results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)(
+				delayed(self._encode_batch_as_sequence)(batch, max_length) for batch in sentence_batches
+			)
+			out_array: list[np.ndarray] = []
+			for r in results:
+				out_array.extend(r)
+		else:
+			out_array = []
+			for batch in tqdm(
+				sentence_batches,
+				total=total_batches,
+				disable=not show_progress_bar,
+			):
+				out_array.extend(self._encode_batch_as_sequence(batch, max_length))
+
+		if was_single:
+			return out_array[0]
+		return out_array
+
+	def _encode_batch_as_sequence(self, sentences: Sequence[str], max_length: int | None) -> list[np.ndarray]:
+		"""Encode a batch of sentences as a sequence."""
+		ids = self.tokenize(sentences=sentences, max_length=max_length)
+		out: list[np.ndarray] = []
+		for id_list in ids:
+			if id_list:
+				out.append(self.embedding[id_list])
+			else:
+				out.append(np.zeros((0, self.dim)))
+
+		return out
+
+	def encode(
+		self,
+		sentences: Sequence[str],
+		show_progress_bar: bool = False,
+		max_length: int | None = 512,
+		batch_size: int = 1024,
+		use_multiprocessing: bool = True,
+		multiprocessing_threshold: int = 10_000,
+		**kwargs: Any,
+	) -> np.ndarray:
+		"""
+		Encode a list of sentences.
+
+		This function encodes a list of sentences by averaging the word embeddings of the tokens in the sentence.
+		For ease of use, we don't batch sentences together.
+
+		NOTE: the return type is currently underspecified. In the case of a single string, this returns a 1D array,
+		    but in the case of a list of strings, this returns a 2D array. Not possible to implement in numpy currently.
+
+		:param sentences: The list of sentences to encode. You can also pass a single sentence.
+		:param show_progress_bar: Whether to show the progress bar.
+		:param max_length: The maximum length of the sentences. Any tokens beyond this length will be truncated.
+		    If this is None, no truncation is done.
+		:param batch_size: The batch size to use.
+		:param use_multiprocessing: Whether to use multiprocessing.
+		    By default, this is enabled for inputs > multiprocessing_threshold sentences and disabled otherwise.
+		:param multiprocessing_threshold: The threshold in number of sentences for using multiprocessing.
+		:param **kwargs: Any additional arguments. These are ignored.
+		:return: The encoded sentences. If a single sentence was passed, a vector is returned.
+		"""
+		was_single = False
+		if isinstance(sentences, str):
+			sentences = [sentences]
+			was_single = True
+
+		# Prepare all batches
+		sentence_batches = list(self._batch(sentences, batch_size))
+		total_batches = math.ceil(len(sentences) / batch_size)
+
+		# Use joblib for multiprocessing if requested, and if we have enough sentences
+		if use_multiprocessing and len(sentences) > multiprocessing_threshold:
+			# Disable parallelism for tokenizers
+			os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+			results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)(
+				delayed(self._encode_batch)(batch, max_length) for batch in sentence_batches
+			)
+			out_array = np.concatenate(results, axis=0)
+		else:
+			# Don't use multiprocessing
+			out_arrays: list[np.ndarray] = []
+			for batch in tqdm(
+				sentence_batches,
+				total=total_batches,
+				disable=not show_progress_bar,
+			):
+				out_arrays.append(self._encode_batch(batch, max_length))
+			out_array = np.concatenate(out_arrays, axis=0)
+
+		if was_single:
+			return out_array[0]
+		return out_array
+
+	def _encode_batch(self, sentences: Sequence[str], max_length: int | None) -> np.ndarray:
+		"""Encode a batch of sentences."""
+		ids = self.tokenize(sentences=sentences, max_length=max_length)
+		out: list[np.ndarray] = []
+		for id_list in ids:
+			if id_list:
+				out.append(self.embedding[id_list].mean(0))
+			else:
+				out.append(np.zeros(self.dim))
+
+		out_array = np.stack(out)
+		if self.normalize:
+			norm = np.linalg.norm(out_array, axis=1, keepdims=True) + 1e-32
+			out_array = out_array / norm
+
+		return out_array
+
+	@staticmethod
+	def _batch(sentences: Sequence[str], batch_size: int) -> Iterator[Sequence[str]]:
+		"""Batch the sentences into equal-sized."""
+		return (sentences[i : i + batch_size] for i in range(0, len(sentences), batch_size))
+
+	def push_to_hub(
+		self, repo_id: str, private: bool = False, token: str | None = None, subfolder: str | None = None
+	) -> None:
+		"""
+		Push the model to the huggingface hub.
+
+		NOTE: you need to pass a token if you are pushing a private model.
+
+		:param repo_id: The repo id to push to.
+		:param private: Whether the repo, if created is set to private.
+		    If the repo already exists, this doesn't change the visibility.
+		:param token: The huggingface token to use.
+		:param subfolder: The subfolder to push to.
+		"""
+		from .hf_utils import push_folder_to_hub
+
+		with TemporaryDirectory() as temp_dir:
+			self.save_pretrained(temp_dir, model_name=repo_id)
+			push_folder_to_hub(Path(temp_dir), subfolder=subfolder, repo_id=repo_id, private=private, token=token)
+
+	@classmethod
+	def load_local(cls: type[StaticModel], path: PathLike) -> StaticModel:
+		"""
+		Loads a model from a local path.
+
+		You should only use this code path if you are concerned with start-up time.
+		Loading via the `from_pretrained` method is safer, and auto-downloads, but
+		also means we import a whole bunch of huggingface code that we don't need.
+
+		Additionally, huggingface will check the most recent version of the model,
+		which can be slow.
+
+		:param path: The path to load the model from. The path is a directory saved by the
+		    `save_pretrained` method.
+		:return: A StaticModel
+		:raises: ValueError if the path is not a directory.
+		"""
+		path = Path(path)
+		if not path.is_dir():
+			msg = f"Path {path} is not a directory."
+			raise ValueError(msg)
+
+		embeddings, tokenizer, config = load_local_model(path)
+
+		return StaticModel(embeddings, tokenizer, config)

src/distiller/model2vec/modelcards/classifier_template.md

@@ -0,0 +1,50 @@
+---
+{{ card_data }}
+---
+
+# {{ model_name }} Model Card
+
+This [Model2Vec](https://github.com/MinishLab/model2vec) model is a fine-tuned version of {% if base_model %}the [{{ base_model }}](https://huggingface.co/{{ base_model }}){% else %}a{% endif %} Model2Vec model. It also includes a classifier head on top.
+
+## Installation
+
+Install model2vec using pip:
+```
+pip install model2vec[inference]
+```
+
+## Usage
+Load this model using the `from_pretrained` method:
+```python
+from model2vec.inference import StaticModelPipeline
+
+# Load a pretrained Model2Vec model
+model = StaticModelPipeline.from_pretrained("{{ model_name }}")
+
+# Predict labels
+predicted = model.predict(["Example sentence"])
+```
+
+## Additional Resources
+
+- [Model2Vec Repo](https://github.com/MinishLab/model2vec)
+- [Model2Vec Base Models](https://huggingface.co/collections/minishlab/model2vec-base-models-66fd9dd9b7c3b3c0f25ca90e)
+- [Model2Vec Results](https://github.com/MinishLab/model2vec/tree/main/results)
+- [Model2Vec Tutorials](https://github.com/MinishLab/model2vec/tree/main/tutorials)
+- [Website](https://minishlab.github.io/)
+
+## Library Authors
+
+Model2Vec was developed by the [Minish Lab](https://github.com/MinishLab) team consisting of [Stephan Tulkens](https://github.com/stephantul) and [Thomas van Dongen](https://github.com/Pringled).
+
+## Citation
+
+Please cite the [Model2Vec repository](https://github.com/MinishLab/model2vec) if you use this model in your work.
+```
+@article{minishlab2024model2vec,
+  author = {Tulkens, Stephan and {van Dongen}, Thomas},
+  title = {Model2Vec: Fast State-of-the-Art Static Embeddings},
+  year = {2024},
+  url = {https://github.com/MinishLab/model2vec}
+}
+```

src/distiller/model2vec/modelcards/model_card_template.md

@@ -0,0 +1,91 @@
+---
+{{ card_data }}
+---
+
+# {{ model_name }} Model Card
+
+This [Model2Vec](https://github.com/MinishLab/model2vec) model is a distilled version of {% if base_model %}the {{ base_model }}(https://huggingface.co/{{ base_model }}){% else %}a{% endif %} Sentence Transformer. It uses static embeddings, allowing text embeddings to be computed orders of magnitude faster on both GPU and CPU. It is designed for applications where computational resources are limited or where real-time performance is critical. Model2Vec models are the smallest, fastest, and most performant static embedders available. The distilled models are up to 50 times smaller and 500 times faster than traditional Sentence Transformers.
+
+
+## Installation
+
+Install model2vec using pip:
+```
+pip install model2vec
+```
+
+## Usage
+
+### Using Model2Vec
+
+The [Model2Vec library](https://github.com/MinishLab/model2vec) is the fastest and most lightweight way to run Model2Vec models.
+
+Load this model using the `from_pretrained` method:
+```python
+from model2vec import StaticModel
+
+# Load a pretrained Model2Vec model
+model = StaticModel.from_pretrained("{{ model_name }}")
+
+# Compute text embeddings
+embeddings = model.encode(["Example sentence"])
+```
+
+### Using Sentence Transformers
+
+You can also use the [Sentence Transformers library](https://github.com/UKPLab/sentence-transformers) to load and use the model:
+
+```python
+from sentence_transformers import SentenceTransformer
+
+# Load a pretrained Sentence Transformer model
+model = SentenceTransformer("{{ model_name }}")
+
+# Compute text embeddings
+embeddings = model.encode(["Example sentence"])
+```
+
+### Distilling a Model2Vec model
+
+You can distill a Model2Vec model from a Sentence Transformer model using the `distill` method. First, install the `distill` extra with `pip install model2vec[distill]`. Then, run the following code:
+
+```python
+from model2vec.distill import distill
+
+# Distill a Sentence Transformer model, in this case the BAAI/bge-base-en-v1.5 model
+m2v_model = distill(model_name="BAAI/bge-base-en-v1.5", pca_dims=256)
+
+# Save the model
+m2v_model.save_pretrained("m2v_model")
+```
+
+## How it works
+
+Model2vec creates a small, fast, and powerful model that outperforms other static embedding models by a large margin on all tasks we could find, while being much faster to create than traditional static embedding models such as GloVe. Best of all, you don't need any data to distill a model using Model2Vec.
+
+It works by passing a vocabulary through a sentence transformer model, then reducing the dimensionality of the resulting embeddings using PCA, and finally weighting the embeddings using [SIF weighting](https://openreview.net/pdf?id=SyK00v5xx). During inference, we simply take the mean of all token embeddings occurring in a sentence.
+
+## Additional Resources
+
+- [Model2Vec Repo](https://github.com/MinishLab/model2vec)
+- [Model2Vec Base Models](https://huggingface.co/collections/minishlab/model2vec-base-models-66fd9dd9b7c3b3c0f25ca90e)
+- [Model2Vec Results](https://github.com/MinishLab/model2vec/tree/main/results)
+- [Model2Vec Tutorials](https://github.com/MinishLab/model2vec/tree/main/tutorials)
+- [Website](https://minishlab.github.io/)
+
+
+## Library Authors
+
+Model2Vec was developed by the [Minish Lab](https://github.com/MinishLab) team consisting of [Stephan Tulkens](https://github.com/stephantul) and [Thomas van Dongen](https://github.com/Pringled).
+
+## Citation
+
+Please cite the [Model2Vec repository](https://github.com/MinishLab/model2vec) if you use this model in your work.
+```
+@article{minishlab2024model2vec,
+  author = {Tulkens, Stephan and {van Dongen}, Thomas},
+  title = {Model2Vec: Fast State-of-the-Art Static Embeddings},
+  year = {2024},
+  url = {https://github.com/MinishLab/model2vec}
+}
+```

src/distiller/model2vec/quantization.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from enum import Enum
+
+import numpy as np
+
+
+class DType(str, Enum):
+    Float16 = "float16"
+    Float32 = "float32"
+    Float64 = "float64"
+    Int8 = "int8"
+
+
+def quantize_embeddings(embeddings: np.ndarray, quantize_to: DType) -> np.ndarray:
+    """
+    Quantize embeddings to a specified data type to reduce memory usage.
+
+    :param embeddings: The embeddings to quantize, as a numpy array.
+    :param quantize_to: The data type to quantize to.
+    :return: The quantized embeddings.
+    :raises ValueError: If the quantization type is not valid.
+    """
+    if quantize_to == DType.Float16:
+        return embeddings.astype(np.float16)
+    if quantize_to == DType.Float32:
+        return embeddings.astype(np.float32)
+    if quantize_to == DType.Float64:
+        return embeddings.astype(np.float64)
+    if quantize_to == DType.Int8:
+        # Normalize to [-128, 127] range for int8
+        # We normalize to -127 to 127 to keep symmetry.
+        scale = np.max(np.abs(embeddings)) / 127.0
+        return np.round(embeddings / scale).astype(np.int8)
+    msg = "Not a valid enum member of DType."
+    raise ValueError(msg)
+
+
+def quantize_and_reduce_dim(
+    embeddings: np.ndarray, quantize_to: str | DType | None, dimensionality: int | None
+) -> np.ndarray:
+    """
+    Quantize embeddings to a datatype and reduce dimensionality.
+
+    :param embeddings: The embeddings to quantize and reduce, as a numpy array.
+    :param quantize_to: The data type to quantize to. If None, no quantization is performed.
+    :param dimensionality: The number of dimensions to keep. If None, no dimensionality reduction is performed.
+    :return: The quantized and reduced embeddings.
+    :raises ValueError: If the passed dimensionality is not None and greater than the model dimensionality.
+    """
+    if quantize_to is not None:
+        quantize_to = DType(quantize_to)
+        embeddings = quantize_embeddings(embeddings, quantize_to)
+
+    if dimensionality is not None:
+        if dimensionality > embeddings.shape[1]:
+            msg = f"Dimensionality {dimensionality} is greater than the model dimensionality {embeddings.shape[1]}"
+            raise ValueError(
+                msg
+            )
+        embeddings = embeddings[:, :dimensionality]
+
+    return embeddings

src/distiller/model2vec/tokenizer/__init__.py

@@ -0,0 +1,12 @@
+from distiller.model2vec.utils import importable
+
+importable("transformers", "tokenizer")
+
+from distiller.model2vec.tokenizer.tokenizer import (
+    clean_and_create_vocabulary,
+    create_tokenizer,
+    replace_vocabulary,
+    turn_tokens_into_ids,
+)
+
+__all__ = ["clean_and_create_vocabulary", "create_tokenizer", "replace_vocabulary", "turn_tokens_into_ids"]

src/distiller/model2vec/tokenizer/datamodels.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Token:
+    """A class to represent a token."""
+
+    form: str
+    # The normalized and pretokenized form of the token
+    normalized_form: str
+    # Whether the word is a continuing subword.
+    is_subword: bool
+    # Whether the token is internal to the model.
+    is_internal: bool

src/distiller/model2vec/tokenizer/model.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+def process_tokenizer(
+    tokenizer_json: dict[str, Any], pre_tokenized_tokens: list[str], unk_token: str | None
+) -> dict[str, Any]:
+    """Process the WordPiece tokenizer JSON."""
+    if tokenizer_json["model"]["type"] == "Unigram":
+        return _process_unigram(tokenizer_json, pre_tokenized_tokens, unk_token)
+    tokenizer_json["model"]["type"] = "Unigram"
+    tokenizer_json["model"]["unk_id"] = pre_tokenized_tokens.index(unk_token) if unk_token else None
+
+    token_weights = np.asarray([_calculate_token_weight_for_unigram(token) for token in pre_tokenized_tokens])
+    proba = (token_weights / np.sum(token_weights)).tolist()
+    tokenizer_json["model"]["vocab"] = [(token, np.log(p)) for token, p in zip(pre_tokenized_tokens, proba, strict=False)]
+
+    return tokenizer_json
+
+
+def _process_unigram(
+    tokenizer_json: dict[str, Any], pre_tokenized_tokens: list[str], unk_token: str | None
+) -> dict[str, Any]:
+    """Process the Unigram tokenizer JSON."""
+    current_probas = dict(tokenizer_json["model"]["vocab"])
+    avg_proba = sum(current_probas.values()) / len(current_probas)
+    new_probas = [[word, current_probas.get(word, avg_proba)] for word in pre_tokenized_tokens]
+    tokenizer_json["model"]["vocab"] = new_probas
+
+    tokens, _ = zip(*tokenizer_json["model"]["vocab"], strict=False)
+    if unk_token is not None:
+        tokenizer_json["model"]["unk_id"] = list(tokens).index(unk_token)
+
+    return tokenizer_json
+
+
+def _calculate_token_weight_for_unigram(token: str) -> float:
+    """Calculate the token weight for Unigram."""
+    # Always prefer longer tokens.
+    return len(token) + token.count("▁") + token.count("Ġ")

src/distiller/model2vec/tokenizer/normalizer.py

@@ -0,0 +1,34 @@
+from string import punctuation
+
+from tokenizers import Regex, Tokenizer
+from tokenizers.normalizers import Replace, Sequence, Strip
+
+
+def replace_normalizer(
+    tokenizer: Tokenizer,
+) -> Tokenizer:
+    """
+    Replace the normalizer for the tokenizer.
+
+    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 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:
+        normalizer = Sequence(new_normalizers)  # type: ignore
+    else:
+        normalizer = Sequence([normalizer, *new_normalizers])  # type: ignore
+    tokenizer.normalizer = normalizer  # type: ignore
+
+    return tokenizer

src/distiller/model2vec/tokenizer/pretokenizer.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from tokenizers import Tokenizer
+
+_FORBIDDEN_PRETOKENIZERS = (
+    "WhiteSpace",
+    "WhitespaceSplit",
+    "BertPreTokenizer",
+    "CharDelimiterSplit",
+    "Punctuation",
+    "Split",
+    "UnicodeScripts",
+)
+_BASIC_METASPACE = {"type": "Metaspace", "replacement": "▁", "prepend_scheme": "always", "split": False}
+
+
+def _fix_single_pretokenizer(pre_tokenizer: dict[str, Any]) -> dict[str, Any] | None:
+    """Fixes a single pretokenizer to allow multiword units."""
+    if pre_tokenizer["type"] in _FORBIDDEN_PRETOKENIZERS:
+        return None
+    if pre_tokenizer["type"] == "ByteLevel":
+        pre_tokenizer["add_prefix_space"] = True
+        pre_tokenizer["use_regex"] = False
+    if pre_tokenizer["type"] == "Metaspace":
+        pre_tokenizer["split"] = False
+        pre_tokenizer["prepend_scheme"] = "always"
+
+    return pre_tokenizer
+
+
+def replace_pretokenizer(tokenizer: Tokenizer) -> Tokenizer:
+    """Fixes a single pretokenizer to allow multiword units."""
+    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
+
+    elif pre_tokenizer_json["type"] == "Sequence":
+        new_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)
+
+        if new_pretokenizers:
+            pre_tokenizer_json["pretokenizers"] = new_pretokenizers
+        else:
+            pre_tokenizer_json = _BASIC_METASPACE
+
+    pre_tokenizer_json = _fix_single_pretokenizer(pre_tokenizer_json) or _BASIC_METASPACE
+    tokenizer_json["pre_tokenizer"] = pre_tokenizer_json
+
+    return tokenizer.from_str(json.dumps(tokenizer_json))

src/distiller/model2vec/tokenizer/tokenizer.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any, cast
+
+from tokenizers import Tokenizer
+from transformers import PreTrainedTokenizerFast
+
+from distiller.model2vec.tokenizer.datamodels import Token
+from distiller.model2vec.tokenizer.model import process_tokenizer
+from distiller.model2vec.tokenizer.normalizer import replace_normalizer
+from distiller.model2vec.tokenizer.pretokenizer import replace_pretokenizer
+
+if TYPE_CHECKING:
+    import re
+
+    from tokenizers.normalizers import Normalizer
+    from tokenizers.pre_tokenizers import (
+        PreTokenizer,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+_DEFAULT_POST_PROCESSOR_TEMPLATE = {
+    "type": "TemplateProcessing",
+    "single": [{"Sequence": {"id": "A", "type_id": 0}}],
+    "pair": [{"Sequence": {"id": "A", "type_id": 0}}, {"Sequence": {"id": "B", "type_id": 0}}],
+    "special_tokens": {},
+}
+
+
+def _remap_added_tokens(
+    special_tokens: list[dict[str, Any]],
+    vocabulary: list[str],
+) -> list[dict[str, Any]]:
+    """
+    Remap special tokens in the tokenizer.
+
+    This function updates the special tokens in the tokenizer based on a mapping provided.
+    It also ensures that the special tokens are present in the vocabulary.
+
+    :param special_tokens: The special tokens to remap.
+    :param vocabulary: The vocabulary as a list of tokens.
+    :return: The updated special tokens.
+    """
+    # Deepcopy
+    special_tokens = [{**x} for x in special_tokens]
+    for token in special_tokens:
+        token["id"] = vocabulary.index(token["content"])
+
+    return special_tokens
+
+
+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_json: dict[str, Any] = json.loads(tokenizer.to_str())
+    added_tokens: list[dict[str, Any]] = tokenizer_json["added_tokens"]
+
+    pre_tokenized_tokens = [x.normalized_form for x in new_vocabulary]
+
+    # We need to remove the added tokens but keep [UNK] and [PAD] tokens.
+    added_tokens = _rename_added_token(unk_token, "[UNK]", added_tokens, pre_tokenized_tokens)
+    added_tokens = _rename_added_token(pad_token, "[PAD]", added_tokens, pre_tokenized_tokens)
+
+    # Remove old added tokens from added tokens
+    tokenizer_json["added_tokens"] = [x for x in added_tokens if x["content"] in {"[UNK]", "[PAD]"}]
+    tokenizer_json = process_tokenizer(
+        tokenizer_json, pre_tokenized_tokens, "[UNK]" if "[UNK]" in pre_tokenized_tokens else None
+    )
+
+    # Remap special tokens
+    tokenizer_json["added_tokens"] = _remap_added_tokens(
+        special_tokens=tokenizer_json["added_tokens"],
+        vocabulary=pre_tokenized_tokens,
+    )
+    tokenizer_json["post_processor"] = _DEFAULT_POST_PROCESSOR_TEMPLATE
+
+    return Tokenizer.from_str(json.dumps(tokenizer_json))
+
+
+def _rename_added_token(
+    form: str | None, new_form: str, added_tokens: list[dict[str, Any]], vocabulary: list[str]
+) -> list[dict[str, Any]]:
+    """Rename added tokens in the tokenizer."""
+    if form is None:
+        return added_tokens
+
+    idx = vocabulary.index(form)
+    added_token = [x for x in added_tokens if x["content"] == form]
+    if added_token:
+        added_token[0]["id"] = idx
+        added_token[0]["content"] = new_form
+        vocabulary[idx] = new_form
+
+    return added_tokens
+
+
+def clean_and_create_vocabulary(
+    tokenizer: PreTrainedTokenizerFast,
+    vocabulary: list[str],
+    token_remove_regex: re.Pattern | None,
+) -> 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()
+    n_empty = 0
+    n_duplicates = 0
+
+    backend_tokenizer = tokenizer.backend_tokenizer
+
+    # Make a base list of tokens.
+    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, backend_tokenizer, internal_tokens, token_remove_regex)
+    # Copy the backend tokenizer to avoid modifying the original.
+    backend_tokenizer = backend_tokenizer.from_str(backend_tokenizer.to_str())
+    backend_tokenizer = replace_normalizer(backend_tokenizer)
+
+    internal_tokens_set = {token.form for token in cleaned_vocabulary}
+
+    normalizer: Normalizer | None = backend_tokenizer.normalizer
+    for token in vocabulary:
+        if normalizer is not None:
+            token = cast("str", normalizer.normalize_str(token))
+
+        if not token:
+            n_empty += 1
+            continue
+
+        pre_tokenizer: PreTokenizer | None = backend_tokenizer.pre_tokenizer
+        normalized_token = token
+        if pre_tokenizer is not None:
+            normalized_token = _normalize_vocabulary_token(
+                token=token,
+                pre_tokenizer=pre_tokenizer,
+            )
+
+        # We need to check whether the pretokenized token is in the vocabulary.
+        # But we need to return the original token, because that will be tokenized
+        # again by the tokenizer during featurization.
+        if normalized_token in seen_tokens or normalized_token in internal_tokens_set:
+            n_duplicates += 1
+            continue
+
+        # Add the possibly pretokenized token to seen
+        seen_tokens.add(normalized_token)
+
+        # After checking the token exists, we need to normalize it into the token
+        # it will become. For byte tokens, this means we don't do anything. For
+        # other types of tokens, we will insert a metaspace.
+        # In the case of multiword tokens, we replace any spaces with the metaspace
+        # or byte prefix token.
+        if not normalized_token.startswith(("▁", "Ġ")):
+            normalized_token = normalized_token.replace(" ", "▁")
+            normalized_token = f"▁{normalized_token}"
+        else:
+            normalized_token = normalized_token.replace(" ", normalized_token[0])
+
+        if normalized_token in post_normalize_seen_tokens:
+            n_duplicates += 1
+            continue
+
+        post_normalize_seen_tokens.add(normalized_token)
+        # Add the original string to the vocabulary.
+        cleaned_vocabulary.append(
+            Token(form=token, normalized_form=normalized_token, is_subword=False, is_internal=False)
+        )
+
+    if n_duplicates:
+        logger.warning(f"Removed {n_duplicates} duplicate tokens.")
+    if n_empty:
+        logger.warning(f"Removed {n_empty} empty tokens.")
+
+    return cleaned_vocabulary, replace_pretokenizer(backend_tokenizer)
+
+
+def _process_internal_tokens(
+    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.
+    pad_token: str | None = tokenizer.special_tokens_map.get("pad_token")  # type: ignore[assignment]
+    unk_token: str | None = tokenizer.special_tokens_map.get("unk_token")  # type: ignore[assignment]
+    # Empty set if no pad or unk token is set.
+    added_tokens_to_keep: set[str] = {x for x in (pad_token, unk_token) if x is not None}
+    added_tokens_to_remove = set(tokenizer.added_tokens_encoder) - added_tokens_to_keep
+    cleaned_internal_tokens: list[Token] = []
+
+    # 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
+    # Isolate the prefix. We can't do first_token[0] because we don't know
+    # how long the prefix is.
+    # e.g., "Ġaaaa" -> "Ġ"
+    a_index = None if "a" not in first_token else first_token.index("a")
+    word_prefix = first_token[:a_index]
+    is_byte_prefix = word_prefix == "Ġ"
+    second_token = encoded.tokens[1]
+    # The second token is the first subword token.
+    # If a tokenizer uses subwords, this token will have been prefixed.
+    # We don't know how long the prefix is.
+    a_index = None if "a" not in second_token else second_token.index("a")
+    subword_prefix = second_token[:a_index]
+
+    pre_tokenizer: PreTokenizer | None = backend_tokenizer.pre_tokenizer
+
+    for token in internal_tokens:
+        # Create the token objects. If this returns None, it was unsucessful for some reason.
+        if token_object := _create_single_internal_token(
+            token=token,
+            subword_prefix=subword_prefix,
+            word_prefix=word_prefix,
+            pre_tokenizer=pre_tokenizer,
+            is_byte_prefix=is_byte_prefix,
+            token_remove_regex=token_remove_regex,
+            added_tokens_to_keep=added_tokens_to_keep,
+            added_tokens_to_remove=added_tokens_to_remove,
+        ):
+            cleaned_internal_tokens.append(token_object)
+
+    if len(cleaned_internal_tokens) != len(internal_tokens):
+        logger.info(
+            f"Removed {len(internal_tokens) - len(cleaned_internal_tokens)} internal tokens from the vocabulary."
+        )
+
+    return cleaned_internal_tokens
+
+
+def _create_single_internal_token(
+    token: str,
+    subword_prefix: str,
+    word_prefix: str,
+    pre_tokenizer: PreTokenizer | None,
+    is_byte_prefix: bool,
+    token_remove_regex: re.Pattern | None,
+    added_tokens_to_keep: set[str],
+    added_tokens_to_remove: set[str],
+) -> Token | None:
+    """Create a token object from a string."""
+    if token in added_tokens_to_remove:
+        # We remove any tokens that are added tokens that aren't [UNK] or [PAD].
+        return None
+    if token in added_tokens_to_keep:
+        # Don't put added tokens through the regular motions.
+        return Token(form=token, normalized_form=token, is_subword=False, is_internal=True)
+    if token_remove_regex and token_remove_regex.match(token):
+        # If the regex matches, remove the token.
+        return None
+
+    # A token is a subword if there is a subword prefix and the word
+    # starts with a subword prefix, or if there is a WORD prefix, and the word
+    # does not start with this prefix. For metaspace tokenizers, for example:
+    # "doghouse" -> ["_dog", "house"]
+    # So we can only tell that "house" is a subword by knowing that it is not prefixed
+    # and word-initial tokens are.
+    is_subword = False
+    if subword_prefix:
+        is_subword = bool(token.startswith(subword_prefix))
+    if word_prefix:
+        is_subword = not bool(token.startswith(word_prefix))
+
+    # Byte prefixed tokenizers don't need to be checked.
+    if pre_tokenizer is not None and not is_byte_prefix:
+        # We need to check the thing without prefixes. If we have a word prefix,
+        # we need to check tokens that have are subwords. Other way around for subword
+        # prefixes.
+        if (subword_prefix and not is_subword) or (word_prefix and is_subword):
+            # If this is True, the token is unreachable, even though it is a subword token.
+            if len(pre_tokenizer.pre_tokenize_str(token)) > 1:
+                return None
+
+    # Turn a token into a normalized form for later processing.
+    normalized_form = _create_normalized_form(token, subword_prefix, word_prefix, is_byte_prefix, is_subword)
+
+    return Token(form=token, normalized_form=normalized_form, is_subword=is_subword, is_internal=True)
+
+
+def _create_normalized_form(
+    token: str, subword_prefix: str, word_prefix: str, is_byte_prefix: bool, is_subword: bool
+) -> str:
+    """Turn an internal token string into a normalized form."""
+    # We don't need to check byte prefixed strings.
+    if is_byte_prefix:
+        return token
+    # We need to check if the token is a subword or not and remove the prefix.
+    if is_subword:
+        return token.removeprefix(subword_prefix)
+    # If the token is not a subword, we need to remove the word prefix, and add metaspace.
+    return f"▁{token.removeprefix(word_prefix)}"
+
+
+def turn_tokens_into_ids(
+    tokens: list[Token], tokenizer: PreTrainedTokenizerFast, unk_token: str | None
+) -> list[list[int]]:
+    """
+    Convert a list of Token objects to their corresponding token ID sequences.
+
+    :param tokens: List of Token objects to convert
+    :param tokenizer: The tokenizer to use for converting tokens to IDs
+    :param unk_token: The string form of the unk token.
+    :return: List of token IDs corresponding to the input tokens
+    """
+    unk_id = None if unk_token is None else tokenizer.convert_tokens_to_ids(unk_token)
+    prefix, suffix = find_eos_bos(tokenizer)
+
+    token_ids: list[list[int]] = []
+    for token in tokens:
+        if token.is_internal:
+            # Careful. Any incorrect tokens will just get `[UNK]``, so this could go horribly wrong
+            # Cast because return type is wrong.
+            token_id: int = cast("int", tokenizer.convert_tokens_to_ids(token.form)) or 0
+            # Explicitly check and warn if `unk_id` appears, but don't crash.
+            if unk_id is not None and token_id == unk_id and token.form != unk_token:
+                logger.warning(f"Token {token.form} was set to unk. This is wrong.")
+            token_ids.append([*prefix, token_id, *suffix])
+        else:
+            token_ids.append(tokenizer.encode(token.form))
+
+    return token_ids
+
+
+def find_eos_bos(tokenizer: PreTrainedTokenizerFast) -> tuple[list[int], list[int]]:
+    """Finds the eos and bos tokens for a tokenizer."""
+    # Little bit complicated, because not all tokenizers have eos and bos tokens.
+    encoding = tokenizer.encode("a", add_special_tokens=True)
+    if len(encoding) != 3:
+        a_encoded = tokenizer.encode("a", add_special_tokens=False)
+        if len(a_encoded) != 1:
+            msg = f"Error while encoding, couldn't determine eos and bos tokens. The model tokenizes 'a' to '{a_encoded}'"
+            raise ValueError(
+                msg
+            )
+        a_idx = encoding.index(a_encoded[0])
+        prefix, suffix = encoding[:a_idx], encoding[a_idx + 1 :]
+    else:
+        prefix, suffix = encoding[:1], encoding[2:]
+    return prefix, suffix
+
+
+def _normalize_vocabulary_token(token: str, pre_tokenizer: PreTokenizer) -> str:
+    """Normalize a token that is not in the initial token vocabulary."""
+    # Add prefix space for byte tokenizers.
+    prefixed_token = f" {token}"
+    pretokenized_tokens: tuple[str, ...]
+    pretokenized_tokens, offsets = zip(*pre_tokenizer.pre_tokenize_str(prefixed_token), strict=False)
+    # The first item is always the start of the token.
+    new_token = [pretokenized_tokens[0]]
+    # Loop over the subtokens and offsets.
+    for t, (s, _) in zip(pretokenized_tokens[1:], offsets[1:], strict=False):
+        # Do not prefix the token with a space if it starts with a metaspace.
+        if t.startswith("▁"):
+            new_token.append(t)
+        # If the character before the subtoken is a space, we have a
+        # multiword token. e.g., "room for the moon", which is split into
+        # ["room", "for", "the", "moon"].
+        # If it doesn't have a space, it is part of a complex multiword token,
+        # e.g., "chat-gpt", which is split into ["chat", "-", "gpt"].
+        elif prefixed_token[s - 1] == " ":
+            new_token.append(f" {t}")
+        else:
+            new_token.append(t)
+    return "".join(new_token)
+
+
+
+def create_tokenizer(
+    tokenizer: PreTrainedTokenizerFast,
+    vocabulary: list[str],
+    token_remove_regex: re.Pattern | None = None,
+) -> PreTrainedTokenizerFast:
+    """
+    Create a tokenizer by adding tokens to the vocabulary.
+
+    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.
+    :param token_remove_regex: The regex to use to remove tokens from the vocabulary.
+    :return: The created 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, 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)

src/distiller/model2vec/train/README.md

@@ -0,0 +1,151 @@
+# Training
+
+Aside from [distillation](../../README.md#distillation), `model2vec` also supports training simple classifiers on top of static models, using [pytorch](https://pytorch.org/), [lightning](https://lightning.ai/) and [scikit-learn](https://scikit-learn.org/stable/index.html).
+
+We support both single and multi-label classification, which work seamlessly based on the labels you provide.
+
+# Installation
+
+To train, make sure you install the training extra:
+
+```
+pip install model2vec[training]
+```
+
+# Quickstart
+
+To train a model, simply initialize it using a `StaticModel`, or from a pre-trained model, as follows:
+
+```python
+from model2vec.distill import distill
+from model2vec.train import StaticModelForClassification
+
+# From a distilled model
+distilled_model = distill("baai/bge-base-en-v1.5")
+classifier = StaticModelForClassification.from_static_model(model=distilled_model)
+
+# From a pre-trained model: potion is the default
+classifier = StaticModelForClassification.from_pretrained(model_name="minishlab/potion-base-32m")
+```
+
+This creates a very simple classifier: a StaticModel with a single 512-unit hidden layer on top. You can adjust the number of hidden layers and the number units through some parameters on both functions. Note that the default for `from_pretrained` is [potion-base-32m](https://huggingface.co/minishlab/potion-base-32M), our best model to date. This is our recommended path if you're working with general English data.
+
+Now that you have created the classifier, let's just train a model. The example below assumes you have the [`datasets`](https://github.com/huggingface/datasets) library installed.
+
+```python
+import numpy as np
+from datasets import load_dataset
+
+# Load the subj dataset
+ds = load_dataset("setfit/subj")
+train = ds["train"]
+test = ds["test"]
+
+s = perf_counter()
+classifier = classifier.fit(train["text"], train["label"])
+
+print(f"Training took {int(perf_counter() - s)} seconds.")
+# Training took 81 seconds
+classification_report = classifier.evaluate(ds["test"]["text"], ds["test"]["label"])
+print(classification_report)
+# Achieved 91.0 test accuracy
+```
+
+As you can see, we got a pretty nice 91% accuracy, with only 81 seconds of training.
+
+The training loop is handled by [`lightning`](https://pypi.org/project/lightning/). By default the training loop splits the data into a train and validation split, with 90% of the data being used for training and 10% for validation. By default, it runs with early stopping on the validation set accuracy, with a patience of 5.
+
+Note that this model is as fast as you're used to from us:
+
+```python
+from time import perf_counter
+
+s = perf_counter()
+classifier.predict(test["text"])
+print(f"Took {int((perf_counter() - s) * 1000)} milliseconds for {len(test)} instances on CPU.")
+# Took 67 milliseconds for 2000 instances on CPU.
+```
+
+## Multi-label classification
+
+Multi-label classification is supported out of the box. Just pass a list of lists to the `fit` function (e.g. `[[label1, label2], [label1, label3]]`), and a multi-label classifier will be trained. For example, the following code trains a multi-label classifier on the [go_emotions](https://huggingface.co/datasets/google-research-datasets/go_emotions) dataset:
+
+```python
+from datasets import load_dataset
+from model2vec.train import StaticModelForClassification
+
+# Initialize a classifier from a pre-trained model
+classifier = StaticModelForClassification.from_pretrained(model_name="minishlab/potion-base-32M")
+
+# Load a multi-label dataset
+ds = load_dataset("google-research-datasets/go_emotions")
+
+# Inspect some of the labels
+print(ds["train"]["labels"][40:50])
+# [[0, 15], [15, 18], [16, 27], [27], [7, 13], [10], [20], [27], [27], [27]]
+
+# Train the classifier on text (X) and labels (y)
+classifier.fit(ds["train"]["text"], ds["train"]["labels"])
+```
+
+Then, we can evaluate the classifier:
+
+```python
+from sklearn import metrics
+from sklearn.preprocessing import MultiLabelBinarizer
+
+classification_report = classifier.evaluate(ds["test"]["text"], ds["test"]["labels"], threshold=0.3)
+print(classification_report)
+# Accuracy: 0.410
+# Precision: 0.527
+# Recall: 0.410
+# F1: 0.439
+```
+
+The scores are competitive with the popular [roberta-base-go_emotions](https://huggingface.co/SamLowe/roberta-base-go_emotions) model, while our model is orders of magnitude faster.
+
+# Persistence
+
+You can turn a classifier into a scikit-learn compatible pipeline, as follows:
+
+```python
+pipeline = classifier.to_pipeline()
+```
+
+This pipeline object can be persisted using standard pickle-based methods, such as [joblib](https://joblib.readthedocs.io/en/stable/). This makes it easy to use your model in inferene pipelines (no installing torch!), although `joblib` and `pickle` should not be used to share models outside of your organization.
+
+If you want to persist your pipeline to the Hugging Face hub, you can use our built-in functions:
+
+```python
+pipeline.save_pretrained(path)
+pipeline.push_to_hub("my_cool/project")
+```
+
+Later, you can load these as follows:
+
+```python
+from model2vec.inference import StaticModelPipeline
+
+pipeline = StaticModelPipeline.from_pretrained("my_cool/project")
+```
+
+Loading pipelines in this way is _extremely_ fast. It takes only 30ms to load a pipeline from disk.
+
+
+# Bring your own architecture
+
+Our training architecture is set up to be extensible, with each task having a specific class. Right now, we only offer `StaticModelForClassification`, but in the future we'll also offer regression, etc.
+
+The core functionality of the `StaticModelForClassification` is contained in a couple of functions:
+
+* `construct_head`: This function constructs the classifier on top of the staticmodel. For example, if you want to create a model that has LayerNorm, just subclass, and replace this function. This should be the main function to update if you want to change model behavior.
+* `train_test_split`: governs the train test split before classification.
+* `prepare_dataset`: Selects the `torch.Dataset` that will be used in the `Dataloader` during training.
+* `_encode`: The encoding function used in the model.
+* `fit`: contains all the lightning-related fitting logic.
+
+The training of the model is done in a `lighting.LightningModule`, which can be modified but is very basic.
+
+# Results
+
+We ran extensive benchmarks where we compared our model to several well known architectures. The results can be found in the [training results](https://github.com/MinishLab/model2vec/tree/main/results#training-results) documentation.

src/distiller/model2vec/train/__init__.py

@@ -0,0 +1,10 @@
+from distiller.model2vec.utils import get_package_extras, importable
+
+_REQUIRED_EXTRA = "train"
+
+for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA):
+    importable(extra_dependency, _REQUIRED_EXTRA)
+
+from distiller.model2vec.train.classifier import StaticModelForClassification
+
+__all__ = ["StaticModelForClassification"]

src/distiller/model2vec/train/base.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self, TypeVar
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset
+
+from distiller.model2vec import StaticModel
+
+if TYPE_CHECKING:
+    from tokenizers import Encoding, Tokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class FinetunableStaticModel(nn.Module):
+    def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int = 2, pad_id: int = 0) -> None:
+        """
+        Initialize a trainable StaticModel from a StaticModel.
+
+        :param vectors: The embeddings of the staticmodel.
+        :param tokenizer: The tokenizer.
+        :param out_dim: The output dimension of the head.
+        :param pad_id: The padding id. This is set to 0 in almost all model2vec models
+        """
+        super().__init__()
+        self.pad_id = pad_id
+        self.out_dim = out_dim
+        self.embed_dim = vectors.shape[1]
+
+        self.vectors = vectors
+        if self.vectors.dtype != torch.float32:
+            dtype = str(self.vectors.dtype)
+            logger.warning(
+                f"Your vectors are {dtype} precision, converting to to torch.float32 to avoid compatibility issues."
+            )
+            self.vectors = vectors.float()
+
+        self.embeddings = nn.Embedding.from_pretrained(self.vectors.clone(), freeze=False, padding_idx=pad_id)
+        self.head = self.construct_head()
+        self.w = self.construct_weights()
+        self.tokenizer = tokenizer
+
+    def construct_weights(self) -> nn.Parameter:
+        """Construct the weights for the model."""
+        weights = torch.zeros(len(self.vectors))
+        weights[self.pad_id] = -10_000
+        return nn.Parameter(weights)
+
+    def construct_head(self) -> nn.Sequential:
+        """Method should be overridden for various other classes."""
+        return nn.Sequential(nn.Linear(self.embed_dim, self.out_dim))
+
+    @classmethod
+    def from_pretrained(
+        cls, *, out_dim: int = 2, model_name: str = "minishlab/potion-base-32m", **kwargs: Any
+    ) -> Self:
+        """Load the model from a pretrained model2vec model."""
+        model = StaticModel.from_pretrained(model_name)
+        return cls.from_static_model(model=model, out_dim=out_dim, **kwargs)
+
+    @classmethod
+    def from_static_model(cls, *, model: StaticModel, out_dim: int = 2, **kwargs: Any) -> Self:
+        """Load the model from a static model."""
+        model.embedding = np.nan_to_num(model.embedding)
+        embeddings_converted = torch.from_numpy(model.embedding)
+        return cls(
+            vectors=embeddings_converted,
+            pad_id=model.tokenizer.token_to_id("[PAD]"),
+            out_dim=out_dim,
+            tokenizer=model.tokenizer,
+            **kwargs,
+        )
+
+    def _encode(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        A forward pass and mean pooling.
+
+        This function is analogous to `StaticModel.encode`, but reimplemented to allow gradients
+        to pass through.
+
+        :param input_ids: A 2D tensor of input ids. All input ids are have to be within bounds.
+        :return: The mean over the input ids, weighted by token weights.
+        """
+        w = self.w[input_ids]
+        w = torch.sigmoid(w)
+        zeros = (input_ids != self.pad_id).float()
+        w = w * zeros
+        # Add a small epsilon to avoid division by zero
+        length = zeros.sum(1) + 1e-16
+        embedded = self.embeddings(input_ids)
+        # Weigh each token
+        embedded = torch.bmm(w[:, None, :], embedded).squeeze(1)
+        # Mean pooling by dividing by the length
+        embedded = embedded / length[:, None]
+
+        return nn.functional.normalize(embedded)
+
+    def forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass through the mean, and a classifier layer after."""
+        encoded = self._encode(input_ids)
+        return self.head(encoded), encoded
+
+    def tokenize(self, texts: list[str], max_length: int | None = 512) -> torch.Tensor:
+        """
+        Tokenize a bunch of strings into a single padded 2D tensor.
+
+        Note that this is not used during training.
+
+        :param texts: The texts to tokenize.
+        :param max_length: If this is None, the sequence lengths are truncated to 512.
+        :return: A 2D padded tensor
+        """
+        encoded: list[Encoding] = self.tokenizer.encode_batch_fast(texts, add_special_tokens=False)
+        encoded_ids: list[torch.Tensor] = [torch.Tensor(encoding.ids[:max_length]).long() for encoding in encoded]
+        return pad_sequence(encoded_ids, batch_first=True, padding_value=self.pad_id)
+
+    @property
+    def device(self) -> str:
+        """Get the device of the model."""
+        return self.embeddings.weight.device
+
+    def to_static_model(self) -> StaticModel:
+        """Convert the model to a static model."""
+        emb = self.embeddings.weight.detach().cpu().numpy()
+        w = torch.sigmoid(self.w).detach().cpu().numpy()
+
+        return StaticModel(emb * w[:, None], self.tokenizer, normalize=True)
+
+
+class TextDataset(Dataset):
+    def __init__(self, tokenized_texts: list[list[int]], targets: torch.Tensor) -> None:
+        """
+        A dataset of texts.
+
+        :param tokenized_texts: The tokenized texts. Each text is a list of token ids.
+        :param targets: The targets.
+        :raises ValueError: If the number of labels does not match the number of texts.
+        """
+        if len(targets) != len(tokenized_texts):
+            msg = "Number of labels does not match number of texts."
+            raise ValueError(msg)
+        self.tokenized_texts = tokenized_texts
+        self.targets = targets
+
+    def __len__(self) -> int:
+        """Return the length of the dataset."""
+        return len(self.tokenized_texts)
+
+    def __getitem__(self, index: int) -> tuple[list[int], torch.Tensor]:
+        """Gets an item."""
+        return self.tokenized_texts[index], self.targets[index]
+
+    @staticmethod
+    def collate_fn(batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Collate function."""
+        texts, targets = zip(*batch, strict=False)
+
+        tensors = [torch.LongTensor(x) for x in texts]
+        padded = pad_sequence(tensors, batch_first=True, padding_value=0)
+
+        return padded, torch.stack(targets)
+
+    def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader:
+        """Convert the dataset to a DataLoader."""
+        return DataLoader(self, collate_fn=self.collate_fn, shuffle=shuffle, batch_size=batch_size)
+
+
+ModelType = TypeVar("ModelType", bound=FinetunableStaticModel)

src/distiller/model2vec/train/classifier.py

@@ -0,0 +1,426 @@
+from __future__ import annotations
+
+import logging
+from collections import Counter
+from itertools import chain
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, TypeVar, cast
+
+import lightning as pl
+import numpy as np
+import torch
+from lightning.pytorch.callbacks import Callback, EarlyStopping
+from sklearn.metrics import jaccard_score
+from sklearn.model_selection import train_test_split
+from sklearn.neural_network import MLPClassifier
+from sklearn.pipeline import make_pipeline
+from torch import nn
+from tqdm import trange
+
+from distiller.model2vec.inference import StaticModelPipeline, evaluate_single_or_multi_label
+from distiller.model2vec.train.base import FinetunableStaticModel, TextDataset
+
+if TYPE_CHECKING:
+    from lightning.pytorch.utilities.types import OptimizerLRScheduler
+    from tokenizers import Tokenizer
+
+logger = logging.getLogger(__name__)
+_RANDOM_SEED = 42
+
+LabelType = TypeVar("LabelType", list[str], list[list[str]])
+
+
+class StaticModelForClassification(FinetunableStaticModel):
+    def __init__(
+        self,
+        *,
+        vectors: torch.Tensor,
+        tokenizer: Tokenizer,
+        n_layers: int = 1,
+        hidden_dim: int = 512,
+        out_dim: int = 2,
+        pad_id: int = 0,
+    ) -> None:
+        """Initialize a standard classifier model."""
+        self.n_layers = n_layers
+        self.hidden_dim = hidden_dim
+        # Alias: Follows scikit-learn. Set to dummy classes
+        self.classes_: list[str] = [str(x) for x in range(out_dim)]
+        # multilabel flag will be set based on the type of `y` passed to fit.
+        self.multilabel: bool = False
+        super().__init__(vectors=vectors, out_dim=out_dim, pad_id=pad_id, tokenizer=tokenizer)
+
+    @property
+    def classes(self) -> np.ndarray:
+        """Return all clasess in the correct order."""
+        return np.array(self.classes_)
+
+    def construct_head(self) -> nn.Sequential:
+        """Constructs a simple classifier head."""
+        if self.n_layers == 0:
+            return nn.Sequential(nn.Linear(self.embed_dim, self.out_dim))
+        modules = [
+            nn.Linear(self.embed_dim, self.hidden_dim),
+            nn.ReLU(),
+        ]
+        for _ in range(self.n_layers - 1):
+            modules.extend([nn.Linear(self.hidden_dim, self.hidden_dim), nn.ReLU()])
+        modules.extend([nn.Linear(self.hidden_dim, self.out_dim)])
+
+        for module in modules:
+            if isinstance(module, nn.Linear):
+                nn.init.kaiming_uniform_(module.weight)
+                nn.init.zeros_(module.bias)
+
+        return nn.Sequential(*modules)
+
+    def predict(
+        self, X: list[str], show_progress_bar: bool = False, batch_size: int = 1024, threshold: float = 0.5
+    ) -> np.ndarray:
+        """
+        Predict labels for a set of texts.
+
+        In single-label mode, each prediction is a single class.
+        In multilabel mode, each prediction is a list of classes.
+
+        :param X: The texts to predict on.
+        :param show_progress_bar: Whether to show a progress bar.
+        :param batch_size: The batch size.
+        :param threshold: The threshold for multilabel classification.
+        :return: The predictions.
+        """
+        pred = []
+        for batch in trange(0, len(X), batch_size, disable=not show_progress_bar):
+            logits = self._predict_single_batch(X[batch : batch + batch_size])
+            if self.multilabel:
+                probs = torch.sigmoid(logits)
+                mask = (probs > threshold).cpu().numpy()
+                pred.extend([self.classes[np.flatnonzero(row)] for row in mask])
+            else:
+                pred.extend([self.classes[idx] for idx in logits.argmax(dim=1).tolist()])
+        if self.multilabel:
+            # Return as object array to allow for lists of varying lengths.
+            return np.array(pred, dtype=object)
+        return np.array(pred)
+
+    @torch.no_grad()
+    def _predict_single_batch(self, X: list[str]) -> torch.Tensor:
+        input_ids = self.tokenize(X)
+        vectors, _ = self.forward(input_ids)
+        return vectors
+
+    def predict_proba(self, X: list[str], show_progress_bar: bool = False, batch_size: int = 1024) -> np.ndarray:
+        """
+        Predict probabilities for each class.
+
+        In single-label mode, returns softmax probabilities.
+        In multilabel mode, returns sigmoid probabilities.
+        """
+        pred = []
+        for batch in trange(0, len(X), batch_size, disable=not show_progress_bar):
+            logits = self._predict_single_batch(X[batch : batch + batch_size])
+            if self.multilabel:
+                pred.append(torch.sigmoid(logits).cpu().numpy())
+            else:
+                pred.append(torch.softmax(logits, dim=1).cpu().numpy())
+        return np.concatenate(pred, axis=0)
+
+    def fit(
+        self,
+        X: list[str],
+        y: LabelType,
+        learning_rate: float = 1e-3,
+        batch_size: int | None = None,
+        min_epochs: int | None = None,
+        max_epochs: int | None = -1,
+        early_stopping_patience: int | None = 5,
+        test_size: float = 0.1,
+        device: str = "auto",
+        X_val: list[str] | None = None,
+        y_val: LabelType | None = None,
+    ) -> StaticModelForClassification:
+        """
+        Fit a model.
+
+        This function creates a Lightning Trainer object and fits the model to the data.
+        It supports both single-label and multi-label classification.
+        We use early stopping. After training, the weights of the best model are loaded back into the model.
+
+        This function seeds everything with a seed of 42, so the results are reproducible.
+        It also splits the data into a train and validation set, again with a random seed.
+
+        If `X_val` and `y_val` are not provided, the function will automatically
+        split the training data into a train and validation set using `test_size`.
+
+        :param X: The texts to train on.
+        :param y: The labels to train on. If the first element is a list, multi-label classification is assumed.
+        :param learning_rate: The learning rate.
+        :param batch_size: The batch size. If None, a good batch size is chosen automatically.
+        :param min_epochs: The minimum number of epochs to train for.
+        :param max_epochs: The maximum number of epochs to train for.
+            If this is -1, the model trains until early stopping is triggered.
+        :param early_stopping_patience: The patience for early stopping.
+            If this is None, early stopping is disabled.
+        :param test_size: The test size for the train-test split.
+        :param device: The device to train on. If this is "auto", the device is chosen automatically.
+        :param X_val: The texts to be used for validation.
+        :param y_val: The labels to be used for validation.
+        :return: The fitted model.
+        :raises ValueError: If either X_val or y_val are provided, but not both.
+        """
+        pl.seed_everything(_RANDOM_SEED)
+        logger.info("Re-initializing model.")
+
+        # Determine whether the task is multilabel based on the type of y.
+
+        self._initialize(y)
+
+        if (X_val is not None) != (y_val is not None):
+            msg = "Both X_val and y_val must be provided together, or neither."
+            raise ValueError(msg)
+
+        if X_val is not None and y_val is not None:
+            # Additional check to ensure y_val is of the same type as y
+            if type(y_val[0]) != type(y[0]):
+                msg = "X_val and y_val must be of the same type as X and y."
+                raise ValueError(msg)
+
+            train_texts = X
+            train_labels = y
+            validation_texts = X_val
+            validation_labels = y_val
+        else:
+            train_texts, validation_texts, train_labels, validation_labels = self._train_test_split(
+                X,
+                y,
+                test_size=test_size,
+            )
+
+        if batch_size is None:
+            # Set to a multiple of 32
+            base_number = int(min(max(1, (len(train_texts) / 30) // 32), 16))
+            batch_size = int(base_number * 32)
+            logger.info("Batch size automatically set to %d.", batch_size)
+
+        logger.info("Preparing train dataset.")
+        train_dataset = self._prepare_dataset(train_texts, train_labels)
+        logger.info("Preparing validation dataset.")
+        val_dataset = self._prepare_dataset(validation_texts, validation_labels)
+
+        c = _ClassifierLightningModule(self, learning_rate=learning_rate)
+
+        n_train_batches = len(train_dataset) // batch_size
+        callbacks: list[Callback] = []
+        if early_stopping_patience is not None:
+            callback = EarlyStopping(monitor="val_accuracy", mode="max", patience=early_stopping_patience)
+            callbacks.append(callback)
+
+        # If the dataset is small, we check the validation set every epoch.
+        # If the dataset is large, we check the validation set every 250 batches.
+        if n_train_batches < 250:
+            val_check_interval = None
+            check_val_every_epoch = 1
+        else:
+            val_check_interval = max(250, 2 * len(val_dataset) // batch_size)
+            check_val_every_epoch = None
+
+        with TemporaryDirectory() as tempdir:
+            trainer = pl.Trainer(
+                min_epochs=min_epochs,
+                max_epochs=max_epochs,
+                callbacks=callbacks,
+                val_check_interval=val_check_interval,
+                check_val_every_n_epoch=check_val_every_epoch,
+                accelerator=device,
+                default_root_dir=tempdir,
+            )
+
+            trainer.fit(
+                c,
+                train_dataloaders=train_dataset.to_dataloader(shuffle=True, batch_size=batch_size),
+                val_dataloaders=val_dataset.to_dataloader(shuffle=False, batch_size=batch_size),
+            )
+            best_model_path = trainer.checkpoint_callback.best_model_path  # type: ignore
+            best_model_weights = torch.load(best_model_path, weights_only=True)
+
+        state_dict = {}
+        for weight_name, weight in best_model_weights["state_dict"].items():
+            state_dict[weight_name.removeprefix("model.")] = weight
+
+        self.load_state_dict(state_dict)
+        self.eval()
+        return self
+
+    def evaluate(
+        self, X: list[str], y: LabelType, batch_size: int = 1024, threshold: float = 0.5, output_dict: bool = False
+    ) -> str | dict[str, dict[str, float]]:
+        """
+        Evaluate the classifier on a given dataset using scikit-learn's classification report.
+
+        :param X: The texts to predict on.
+        :param y: The ground truth labels.
+        :param batch_size: The batch size.
+        :param threshold: The threshold for multilabel classification.
+        :param output_dict: Whether to output the classification report as a dictionary.
+        :return: A classification report.
+        """
+        self.eval()
+        predictions = self.predict(X, show_progress_bar=True, batch_size=batch_size, threshold=threshold)
+        return evaluate_single_or_multi_label(predictions=predictions, y=y, output_dict=output_dict)
+
+
+    def _initialize(self, y: LabelType) -> None:
+        """
+        Sets the output dimensionality, the classes, and initializes the head.
+
+        :param y: The labels.
+        :raises ValueError: If the labels are inconsistent.
+        """
+        if isinstance(y[0], (str, int)):
+            # Check if all labels are strings or integers.
+            if not all(isinstance(label, (str, int)) for label in y):
+                msg = "Inconsistent label types in y. All labels must be strings or integers."
+                raise ValueError(msg)
+            self.multilabel = False
+            classes = sorted(set(y))
+        else:
+            # Check if all labels are lists or tuples.
+            if not all(isinstance(label, (list, tuple)) for label in y):
+                msg = "Inconsistent label types in y. All labels must be lists or tuples."
+                raise ValueError(msg)
+            self.multilabel = True
+            classes = sorted(set(chain.from_iterable(y)))
+
+        self.classes_ = classes
+        self.out_dim = len(self.classes_)  # Update output dimension
+        self.head = self.construct_head()
+        self.embeddings = nn.Embedding.from_pretrained(self.vectors.clone(), freeze=False, padding_idx=self.pad_id)
+        self.w = self.construct_weights()
+        self.train()
+
+    def _prepare_dataset(self, X: list[str], y: LabelType, max_length: int = 512) -> TextDataset:
+        """
+        Prepare a dataset. For multilabel classification, each target is converted into a multi-hot vector.
+
+        :param X: The texts.
+        :param y: The labels.
+        :param max_length: The maximum length of the input.
+        :return: A TextDataset.
+        """
+        # This is a speed optimization.
+        # assumes a mean token length of 10, which is really high, so safe.
+        truncate_length = max_length * 10
+        X = [x[:truncate_length] for x in X]
+        tokenized: list[list[int]] = [
+            encoding.ids[:max_length] for encoding in self.tokenizer.encode_batch_fast(X, add_special_tokens=False)
+        ]
+        if self.multilabel:
+            # Convert labels to multi-hot vectors
+            num_classes = len(self.classes_)
+            labels_tensor = torch.zeros(len(y), num_classes, dtype=torch.float)
+            mapping = {label: idx for idx, label in enumerate(self.classes_)}
+            for i, sample_labels in enumerate(y):
+                indices = [mapping[label] for label in sample_labels]
+                labels_tensor[i, indices] = 1.0
+        else:
+            labels_tensor = torch.tensor([self.classes_.index(label) for label in cast("list[str]", y)], dtype=torch.long)
+        return TextDataset(tokenized, labels_tensor)
+
+    def _train_test_split(
+        self,
+        X: list[str],
+        y: list[str] | list[list[str]],
+        test_size: float,
+    ) -> tuple[list[str], list[str], LabelType, LabelType]:
+        """
+        Split the data.
+
+        For single-label classification, stratification is attempted (if possible).
+        For multilabel classification, a random split is performed.
+        """
+        if not self.multilabel:
+            label_counts = Counter(y)
+            if min(label_counts.values()) < 2:
+                logger.info("Some classes have less than 2 samples. Stratification is disabled.")
+                return train_test_split(X, y, test_size=test_size, random_state=42, shuffle=True)
+            return train_test_split(X, y, test_size=test_size, random_state=42, shuffle=True, stratify=y)
+        # Multilabel classification does not support stratification.
+        return train_test_split(X, y, test_size=test_size, random_state=42, shuffle=True)
+
+    def to_pipeline(self) -> StaticModelPipeline:
+        """Convert the model to an sklearn pipeline."""
+        static_model = self.to_static_model()
+
+        random_state = np.random.RandomState(_RANDOM_SEED)
+        n_items = len(self.classes)
+        X = random_state.randn(n_items, static_model.dim)
+        y = self.classes
+
+        converted = make_pipeline(MLPClassifier(hidden_layer_sizes=(self.hidden_dim,) * self.n_layers))
+        converted.fit(X, y)
+        mlp_head: MLPClassifier = converted[-1]
+
+        for index, layer in enumerate([module for module in self.head if isinstance(module, nn.Linear)]):
+            mlp_head.coefs_[index] = layer.weight.detach().cpu().numpy().T
+            mlp_head.intercepts_[index] = layer.bias.detach().cpu().numpy()
+        # Below is necessary to ensure that the converted model works correctly.
+        # In scikit-learn, a binary classifier only has a single vector of output coefficients
+        # and a single intercept. We use two output vectors.
+        # To convert correctly, we need to set the outputs correctly, and fix the activation function.
+        # Make sure n_outputs is set to > 1.
+        mlp_head.n_outputs_ = self.out_dim
+        # Set to softmax or sigmoid
+        mlp_head.out_activation_ = "logistic" if self.multilabel else "softmax"
+
+        return StaticModelPipeline(static_model, converted)
+
+
+class _ClassifierLightningModule(pl.LightningModule):
+    def __init__(self, model: StaticModelForClassification, learning_rate: float) -> None:
+        """Initialize the LightningModule."""
+        super().__init__()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.loss_function = nn.CrossEntropyLoss() if not model.multilabel else nn.BCEWithLogitsLoss()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Simple forward pass."""
+        return self.model(x)
+
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Training step using cross-entropy loss for single-label and binary cross-entropy for multilabel training."""
+        x, y = batch
+        head_out, _ = self.model(x)
+        loss = self.loss_function(head_out, y)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Validation step computing loss and accuracy."""
+        x, y = batch
+        head_out, _ = self.model(x)
+        loss = self.loss_function(head_out, y)
+        if self.model.multilabel:
+            preds = (torch.sigmoid(head_out) > 0.5).float()
+            # Multilabel accuracy is defined as the Jaccard score averaged over samples.
+            accuracy = jaccard_score(y.cpu(), preds.cpu(), average="samples")
+        else:
+            accuracy = (head_out.argmax(dim=1) == y).float().mean()
+        self.log("val_loss", loss)
+        self.log("val_accuracy", accuracy, prog_bar=True)
+
+        return loss
+
+    def configure_optimizers(self) -> OptimizerLRScheduler:
+        """Configure optimizer and learning rate scheduler."""
+        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer,
+            mode="min",
+            factor=0.5,
+            patience=3,
+            min_lr=1e-6,
+            threshold=0.03,
+            threshold_mode="rel",
+        )
+        return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "monitor": "val_loss"}}

src/distiller/model2vec/utils.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+from importlib import import_module
+from importlib.metadata import metadata
+from typing import TYPE_CHECKING, Any, Protocol, cast
+
+import safetensors
+from joblib import Parallel
+from tokenizers import Tokenizer
+from tqdm import tqdm
+
+if TYPE_CHECKING:
+	from collections.abc import Iterator
+	from pathlib import Path
+
+	import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressParallel(Parallel):
+	"""A drop-in replacement for joblib.Parallel that shows a tqdm progress bar."""
+
+	def __init__(self, use_tqdm: bool = True, total: int | None = None, *args: Any, **kwargs: Any) -> None:
+		"""
+		Initialize the ProgressParallel object.
+
+		:param use_tqdm: Whether to show the progress bar.
+		:param total: Total number of tasks (batches) you expect to process. If None,
+		            it updates the total dynamically to the number of dispatched tasks.
+		:param *args: Additional arguments to pass to `Parallel.__init__`.
+		:param **kwargs: Additional keyword arguments to pass to `Parallel.__init__`.
+		"""
+		self._use_tqdm = use_tqdm
+		self._total = total
+		super().__init__(*args, **kwargs)
+
+	def __call__(self, *args: Any, **kwargs: Any) -> Any:
+		"""Create a tqdm context."""
+		with tqdm(disable=not self._use_tqdm, total=self._total) as self._pbar:
+			self._pbar = self._pbar
+			return super().__call__(*args, **kwargs)
+
+	def print_progress(self) -> None:
+		"""Hook called by joblib as tasks complete. We update the tqdm bar here."""
+		if self._total is None:
+			# If no fixed total was given, we dynamically set the total
+			self._pbar.total = self.n_dispatched_tasks
+		# Move the bar to the number of completed tasks
+		self._pbar.n = self.n_completed_tasks
+		self._pbar.refresh()
+
+
+class SafeOpenProtocol(Protocol):
+	"""Protocol to fix safetensors safe open."""
+
+	def get_tensor(self, key: str) -> np.ndarray:
+		"""Get a tensor."""
+		...  # pragma: no cover
+
+
+_MODULE_MAP = (("scikit-learn", "sklearn"),)
+_DIVIDERS = re.compile(r"[=<>!]+")
+
+
+def get_package_extras(package: str, extra: str) -> Iterator[str]:
+	"""Get the extras of the package."""
+	try:
+		message = metadata(package)
+	except Exception as e:
+		# For local packages without metadata, return empty iterator
+		# This allows the package to work without installed metadata
+		logger.debug(f"Could not retrieve metadata for package '{package}': {e}")
+		return iter([])
+
+	all_packages = message.get_all("Requires-Dist") or []
+	for package in all_packages:
+		name, *rest = package.split(";", maxsplit=1)
+		if rest:
+			# Extract and clean the extra requirement
+			found_extra = rest[0].split("==")[-1].strip(" \"'")
+			if found_extra == extra:
+				prefix, *_ = _DIVIDERS.split(name)
+				yield prefix.strip()
+
+
+def importable(module: str, extra: str) -> None:
+	"""Check if a module is importable."""
+	module = dict(_MODULE_MAP).get(module, module)
+	try:
+		import_module(module)
+	except ImportError:
+		msg = f"`{module}`, is required. Please reinstall model2vec with the `{extra}` extra. `pip install model2vec[{extra}]`"
+		raise ImportError(msg)
+
+
+def setup_logging() -> None:
+	"""Simple logging setup."""
+	from rich.logging import RichHandler
+
+	logging.basicConfig(
+		level="INFO",
+		format="%(name)s - %(message)s",
+		datefmt="%Y-%m-%d %H:%M:%S",
+		handlers=[RichHandler(rich_tracebacks=True)],
+	)
+
+
+def load_local_model(folder: Path) -> tuple[np.ndarray, Tokenizer, dict[str, str]]:
+	"""Load a local model."""
+	embeddings_path = folder / "model.safetensors"
+	tokenizer_path = folder / "tokenizer.json"
+	config_path = folder / "config.json"
+
+	opened_tensor_file = cast("SafeOpenProtocol", safetensors.safe_open(embeddings_path, framework="numpy"))
+	embeddings = opened_tensor_file.get_tensor("embeddings")
+
+	config = json.load(open(config_path)) if config_path.exists() else {}
+
+	tokenizer: Tokenizer = Tokenizer.from_file(str(tokenizer_path))
+
+	if len(tokenizer.get_vocab()) != len(embeddings):
+		logger.warning(
+			f"Number of tokens does not match number of embeddings: `{len(tokenizer.get_vocab())}` vs `{len(embeddings)}`"
+		)
+
+	return embeddings, tokenizer, config

src/distiller/model2vec/version.py

@@ -0,0 +1,2 @@
+__version_triple__ = (0, 5, 0)
+__version__ = ".".join(map(str, __version_triple__))