diff --git a/model2vec/distill/__init__.py b/model2vec/distill/__init__.py index fbc59f1..b2e23af 100644 --- a/model2vec/distill/__init__.py +++ b/model2vec/distill/__init__.py @@ -5,6 +5,6 @@ for extra_dependency in get_package_extras("model2vec", _REQUIRED_EXTRA): importable(extra_dependency, _REQUIRED_EXTRA) -from model2vec.distill.distillation import distill, distill_from_model +from model2vec.distill.distillation import distill, distill_from_model, init_from_tokenizer, random_init -__all__ = ["distill", "distill_from_model"] +__all__ = ["distill", "distill_from_model", "init_from_tokenizer", "random_init"] diff --git a/model2vec/distill/distillation.py b/model2vec/distill/distillation.py index 3cf5c0b..511513a 100644 --- a/model2vec/distill/distillation.py +++ b/model2vec/distill/distillation.py @@ -5,6 +5,7 @@ import re from typing import Sequence, cast +import numpy as np from huggingface_hub.hf_api import model_info from skeletoken import TokenizerModel from skeletoken.external.transformers import reshape_embeddings @@ -65,35 +66,16 @@ def distill_from_model( 'first': use the first token's hidden state ([CLS] token in BERT-style models). 'pooler': use the pooler output (if available). This is often a non-linear projection of the [CLS] token. :return: A StaticModel. - :raises ValueError: if the vocabulary is empty after preprocessing. """ quantize_to = DType(quantize_to) sif_coefficient, token_remove_regex = _validate_parameters(sif_coefficient, token_remove_pattern) - if vocabulary is None: - vocabulary = [] - device = select_optimal_device(device) - original_tokenizer_model = TokenizerModel.from_transformers_tokenizer(tokenizer) - original_tokenizer_model = original_tokenizer_model.prune_added_tokens() - - # Clean the vocabulary by removing duplicate tokens and tokens that are in the internal vocabulary. - # Copy the original tokenizer model. - tokenizer_model = original_tokenizer_model.deep_copy() - if tokenizer_model.adds_prefix_space is not None: - tokenizer_model.adds_prefix_space = True - - # Create the vocabulary in the new tokenizer. - tokenizer_model = clean_and_create_vocabulary(tokenizer_model, vocabulary, token_remove_regex=token_remove_regex) - - # All tokens in a single list. - all_tokens = tokenizer_model.sorted_vocabulary - if not all_tokens: - raise ValueError("The vocabulary is empty after preprocessing. Please check your token_remove_pattern.") + original_tokenizer_model, tokenizer_model = _prepare_tokenizer(tokenizer, vocabulary, token_remove_regex) # Turn all _new_ tokens into ids using the original tokenizer - token_ids = turn_tokens_into_ids(all_tokens, original_tokenizer_model) + token_ids = turn_tokens_into_ids(tokenizer_model.sorted_vocabulary, original_tokenizer_model) # Reshape the transformer model = reshape_embeddings(model, original_tokenizer_model) @@ -138,7 +120,176 @@ def distill_from_model( "pooling": pooling, } - if os.path.exists(model_name): + return _create_static_model( + embeddings=embeddings, + weights=weights, + token_mapping=token_mapping, + tokenizer_model=tokenizer_model, + config=config, + model_name=model_name, + ) + + +def random_init( + model_name: str, + vocabulary: Sequence[str] | None = None, + dim: int = 256, + sif_coefficient: float | None = 1e-4, + token_remove_pattern: str | None = r"\[unused\d+\]", + quantize_to: DType | str = DType.Float16, + seed: int = 42, + trust_remote_code: bool = False, +) -> StaticModel: + """Create a staticmodel with random embeddings from a tokenizer. + + The tokenizer and vocabulary are processed in the same way as in `distill_from_model`, but instead of + computing embeddings with a model, each token gets an embedding drawn from N(0, 1). + + :param model_name: The name or path of the model whose tokenizer to load. + :param vocabulary: The vocabulary to use. If this is None, we use the tokenizer's vocabulary. + :param dim: The dimensionality of the embeddings. + :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 quantize_to: The data type to quantize to. Can be any of the DType enum members or their string equivalents. + :param seed: The random seed used to draw the embeddings. + :param trust_remote_code: Whether to trust remote code when loading the tokenizer. + :return: A StaticModel. + + """ + tokenizer = cast( + PreTrainedTokenizerFast, + AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code, use_fast=True), + ) + return init_from_tokenizer( + tokenizer=tokenizer, + vocabulary=vocabulary, + dim=dim, + sif_coefficient=sif_coefficient, + token_remove_pattern=token_remove_pattern, + quantize_to=quantize_to, + seed=seed, + ) + + +def init_from_tokenizer( + tokenizer: PreTrainedTokenizerFast, + vocabulary: Sequence[str] | None = None, + dim: int = 256, + sif_coefficient: float | None = 1e-4, + token_remove_pattern: str | None = r"\[unused\d+\]", + quantize_to: DType | str = DType.Float16, + seed: int = 42, +) -> StaticModel: + """Create a staticmodel with random embeddings from a tokenizer. + + The tokenizer and vocabulary are processed in the same way as in `distill_from_model`, but instead of + computing embeddings with a model, each token gets an embedding drawn from N(0, 1). + + :param tokenizer: The tokenizer to use. + :param vocabulary: The vocabulary to use. If this is None, we use the tokenizer's vocabulary. + :param dim: The dimensionality of the embeddings. + :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 quantize_to: The data type to quantize to. Can be any of the DType enum members or their string equivalents. + :param seed: The random seed used to draw the embeddings. + :return: A StaticModel. + :raises ValueError: if the parameters are invalid or the vocabulary is empty after preprocessing. + + """ + if dim <= 0: + raise ValueError(f"dim must be a positive integer, got {dim}.") + quantize_to = DType(quantize_to) + sif_coefficient, token_remove_regex = _validate_parameters(sif_coefficient, token_remove_pattern) + + _, tokenizer_model = _prepare_tokenizer(tokenizer, vocabulary, token_remove_regex) + + embeddings = np.random.default_rng(seed).standard_normal((tokenizer_model.vocabulary_size, dim)) + weights = compute_weights(len(embeddings), sif_coefficient=sif_coefficient) + embeddings = quantize_embeddings(embeddings * weights[:, None], quantize_to) + + tokenizer_name = tokenizer.name_or_path + + config: StaticModelConfig = { + "model_type": "model2vec", + "architectures": ["StaticModel"], + "tokenizer_name": tokenizer_name, + "apply_pca": None, + "sif_coefficient": sif_coefficient, + "hidden_dim": embeddings.shape[1], + "seq_length": 1000000, + "normalize": True, + "pooling": PoolingMode.RANDOM, + } + + return _create_static_model( + embeddings=embeddings, + weights=None, + token_mapping=None, + tokenizer_model=tokenizer_model, + config=config, + model_name=tokenizer_name, + ) + + +def _prepare_tokenizer( + tokenizer: PreTrainedTokenizerFast, + vocabulary: Sequence[str] | None, + token_remove_regex: re.Pattern[str] | None, +) -> tuple[TokenizerModel, TokenizerModel]: + """Create the original tokenizer model and the new tokenizer model. + + :param tokenizer: The original tokenizer. + :param vocabulary: The vocabulary to add. If this is None, no tokens are added. + :param token_remove_regex: Tokens matching this regex are removed from the vocabulary. + :return: The original tokenizer model and the new tokenizer model. + :raises ValueError: if the vocabulary is empty after preprocessing. + + """ + original_tokenizer_model = TokenizerModel.from_transformers_tokenizer(tokenizer) + original_tokenizer_model = original_tokenizer_model.prune_added_tokens() + + # Clean the vocabulary by removing duplicate tokens and tokens that are in the internal vocabulary. + # Copy the original tokenizer model. + tokenizer_model = original_tokenizer_model.deep_copy() + if tokenizer_model.adds_prefix_space is not None: + tokenizer_model.adds_prefix_space = True + + vocabulary = vocabulary if vocabulary is not None else [] + # Create the vocabulary in the new tokenizer. + tokenizer_model = clean_and_create_vocabulary(tokenizer_model, vocabulary, token_remove_regex=token_remove_regex) + + if tokenizer_model.vocabulary_size == 0: + raise ValueError("The vocabulary is empty after preprocessing. Please check your token_remove_pattern.") + + return original_tokenizer_model, tokenizer_model + + +def _create_static_model( + embeddings: np.ndarray, + weights: np.ndarray | None, + token_mapping: np.ndarray | None, + tokenizer_model: TokenizerModel, + config: StaticModelConfig, + model_name: str, +) -> StaticModel: + """Create a StaticModel, looking up the language of the base model on the Hugging Face Hub. + + :param embeddings: The embeddings. + :param weights: The token weights, if any. + :param token_mapping: The mapping from tokens to embeddings, if any. + :param tokenizer_model: The tokenizer model. + :param config: The config of the model. + :param model_name: The name or path of the base model. + :return: A StaticModel. + + """ + if not model_name: + language = None + elif 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 @@ -171,7 +322,7 @@ def _validate_parameters( """Validate the parameters passed to the distillation function. :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. + 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. @@ -220,7 +371,7 @@ def distill( If this is None, we don't apply PCA. If this is 'auto', we don't reduce dimenionality, but still apply PCA. :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. + 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 diff --git a/model2vec/distill/inference.py b/model2vec/distill/inference.py index e580e05..b69eab4 100644 --- a/model2vec/distill/inference.py +++ b/model2vec/distill/inference.py @@ -31,12 +31,14 @@ class PoolingMode(str, Enum): - POOLER: use the model's `pooler_output`. In BERT-like models this is computed as the hidden state at [CLS], passed through a learned dense layer + activation. Not all models provide this. + - RANDOM: no model is used; each token gets a randomly drawn embedding. """ MEAN = "mean" LAST = "last" FIRST = "first" POOLER = "pooler" + RANDOM = "random" def create_embeddings( diff --git a/tests/test_distillation.py b/tests/test_distillation.py index 2b0522d..ef5b1cd 100644 --- a/tests/test_distillation.py +++ b/tests/test_distillation.py @@ -3,16 +3,21 @@ import json from importlib import import_module +from pathlib import Path from unittest.mock import MagicMock, patch import numpy as np import pytest from pytest import LogCaptureFixture from skeletoken import TokenizerModel +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.pre_tokenizers import ByteLevel from transformers import BertTokenizer from transformers.modeling_utils import PreTrainedModel from transformers.tokenization_utils_tokenizers import PreTrainedTokenizerFast +from model2vec.distill import init_from_tokenizer, random_init from model2vec.distill.distillation import distill, distill_from_model from model2vec.distill.inference import PoolingMode, apply_pca, compute_weights, create_embeddings from model2vec.model import StaticModel @@ -198,6 +203,14 @@ def test_distill_removal_pattern( token_remove_pattern="[...papapa", ) + # Random init from the tokenizer only + full = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=8, token_remove_pattern=None) + pruned = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=8) + assert len(full.embedding) == expected_vocab_size + assert any(token.startswith("[unused") for token in full.tokens) + assert not any(token.startswith("[unused") for token in pruned.tokens) + assert len(pruned.embedding) == len(pruned.tokens) + @pytest.mark.parametrize( "vocabulary, pca_dims, sif_coefficient, expected_shape", @@ -270,6 +283,149 @@ def test_missing_modelinfo( assert static_model.language is None +@pytest.mark.parametrize( + "vocabulary, dim, sif_coefficient, token_remove_pattern, quantize_to, expected_shape", + [ + (None, 256, 1e-4, None, "float16", (30519, 256)), # SIF on + (None, 16, None, None, "float32", (30519, 16)), # SIF off + (None, 16, None, None, "float64", (30519, 16)), + (["wordA", "wordB"], 8, None, None, "int8", (30521, 8)), # Custom vocab + (None, 0, None, None, "float16", None), # invalid dim -> raises + (None, -1, None, None, "float16", None), # invalid dim -> raises + (None, 8, 0, None, "float16", None), # invalid SIF (too low) -> raises + (None, 8, 1, None, "float16", None), # invalid SIF (too high) -> raises + (None, 8, None, "[", "float16", None), # invalid regex -> raises + ], +) +@patch.object(import_module("model2vec.distill.distillation"), "model_info") +def test_init_from_tokenizer( + mock_model_info: MagicMock, + mock_berttokenizer: PreTrainedTokenizerFast, + tmp_path: Path, + vocabulary: list[str] | None, + dim: int, + sif_coefficient: float | None, + token_remove_pattern: str | None, + quantize_to: str, + expected_shape: tuple[int, int] | None, +) -> None: + """Test creating a random model from a tokenizer only.""" + if expected_shape is None: + with pytest.raises(ValueError): + init_from_tokenizer( + tokenizer=mock_berttokenizer, + dim=dim, + sif_coefficient=sif_coefficient, + token_remove_pattern=token_remove_pattern, + ) + return + + static_model = init_from_tokenizer( + tokenizer=mock_berttokenizer, + vocabulary=vocabulary, + dim=dim, + sif_coefficient=sif_coefficient, + token_remove_pattern=token_remove_pattern, + quantize_to=quantize_to, + ) + assert static_model.embedding.shape == expected_shape + assert static_model.embedding.dtype == np.dtype(quantize_to) + assert static_model.weights is None + assert static_model.token_mapping is None + assert static_model.config["hidden_dim"] == dim + assert static_model.config["apply_pca"] is None + assert static_model.config["sif_coefficient"] == sif_coefficient + assert static_model.config["pooling"] == PoolingMode.RANDOM + assert static_model.base_model_name == "test_tokenizer" + mock_model_info.assert_not_called() + + for token in vocabulary or []: + assert static_model.tokenizer.normalizer.normalize_str(token) in static_model.tokens + assert static_model.encode([" ".join(vocabulary or ["hello"])]).shape == (1, dim) + + static_model.save_pretrained(tmp_path) + loaded = StaticModel.from_pretrained(tmp_path) + assert np.array_equal(loaded.embedding, static_model.embedding) + assert loaded.tokens == static_model.tokens + + seeded = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=dim, seed=1) + seeded_again = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=dim, seed=1) + other_seed = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=dim, seed=2) + assert np.array_equal(seeded.embedding, seeded_again.embedding) + assert not np.array_equal(seeded.embedding, other_seed.embedding) + + +@pytest.fixture +def byte_level_tokenizer() -> PreTrainedTokenizerFast: + """Create a byte-level BPE tokenizer without a prefix space.""" + vocab = ["word1", "word2", "Ġword3", "[UNK]"] + tokenizer = Tokenizer(BPE(vocab={token: idx for idx, token in enumerate(vocab)}, merges=[], unk_token="[UNK]")) + tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False) # type: ignore + return PreTrainedTokenizerFast(tokenizer_object=tokenizer) + + +def test_init_from_tokenizer_byte_level(byte_level_tokenizer: PreTrainedTokenizerFast) -> None: + """Test that byte-level tokenizers get a prefix space, and that removing all tokens raises.""" + static_model = init_from_tokenizer(tokenizer=byte_level_tokenizer, dim=4, token_remove_pattern=None) + assert static_model.tokenizer.pre_tokenizer.add_prefix_space # type: ignore + assert not byte_level_tokenizer.backend_tokenizer.pre_tokenizer.add_prefix_space # type: ignore + assert static_model.embedding.shape == (4, 4) + + with pytest.raises(ValueError, match="The vocabulary is empty after preprocessing"): + init_from_tokenizer(tokenizer=byte_level_tokenizer, dim=4, token_remove_pattern=r".*") + + +def test_init_from_tokenizer_sif_weighting(mock_berttokenizer: PreTrainedTokenizerFast) -> None: + """Test that SIF weights are folded into the embeddings.""" + unweighted = init_from_tokenizer( + tokenizer=mock_berttokenizer, dim=8, sif_coefficient=None, token_remove_pattern=None, quantize_to="float64" + ) + weighted = init_from_tokenizer( + tokenizer=mock_berttokenizer, dim=8, sif_coefficient=1e-4, token_remove_pattern=None, quantize_to="float64" + ) + weights = compute_weights(len(unweighted.embedding), sif_coefficient=1e-4) + assert np.allclose(weighted.embedding, unweighted.embedding * weights[:, None]) + + +@pytest.mark.parametrize( + "name_or_path, card_data, expected_language", + [ + ("", None, None), # No name -> no hub call + ("some-org/some-model", {"language": "en"}, "en"), # Hub name -> language from model card + ], +) +@patch.object(import_module("model2vec.distill.distillation"), "model_info") +def test_init_from_tokenizer_model_info( + mock_model_info: MagicMock, + mock_berttokenizer: PreTrainedTokenizerFast, + name_or_path: str, + card_data: dict[str, str] | None, + expected_language: str | None, +) -> None: + """Test that the hub is only queried for tokenizers with a hub name.""" + mock_model_info.return_value = type("ModelInfo", (object,), {"cardData": card_data}) + mock_berttokenizer.name_or_path = name_or_path + static_model = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=8) + assert static_model.language == expected_language + assert static_model.config["tokenizer_name"] == name_or_path + if card_data is None: + mock_model_info.assert_not_called() + else: + mock_model_info.assert_called_once_with(name_or_path) + assert static_model.base_model_name == name_or_path + + +@patch("model2vec.distill.distillation.AutoTokenizer.from_pretrained") +def test_random_init(mock_from_pretrained: MagicMock, mock_berttokenizer: PreTrainedTokenizerFast) -> None: + """Test that random_init loads the tokenizer and matches init_from_tokenizer.""" + mock_from_pretrained.return_value = mock_berttokenizer + static_model = random_init("tests/data/test_tokenizer", dim=8, seed=3, trust_remote_code=True) + mock_from_pretrained.assert_called_once_with("tests/data/test_tokenizer", trust_remote_code=True, use_fast=True) + expected = init_from_tokenizer(tokenizer=mock_berttokenizer, dim=8, seed=3) + assert np.array_equal(static_model.embedding, expected.embedding) + assert static_model.tokens == expected.tokens + + @pytest.mark.parametrize( "embeddings, pca_dims, sif_coefficient, expected_shape", [