Skip to content

chore: remove quant from distill #404

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stephantul merged 1 commit into from
Oct 19, 2026
Merged

chore: remove quant from distill #404

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 4 additions & 32 deletions model2vec/distill/distillation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from model2vec.quantization import DType, quantize_embeddings
from model2vec.tokenizer import clean_and_create_vocabulary, turn_tokens_into_ids
from model2vec.types import StaticModelConfig
from model2vec.vocabulary_quantization import quantize_vocabulary

logger = logging.getLogger(__name__)

Expand All @@ -32,7 +31,6 @@ def distill_from_model(
sif_coefficient: float | None = 1e-4,
token_remove_pattern: str | None = r"\[unused\d+\]",
quantize_to: DType | str = DType.Float16,
vocabulary_quantization: int | None = None,
pooling: PoolingMode | str = PoolingMode.MEAN,
) -> StaticModel:
"""Distill a staticmodel from a sentence transformer.
Expand All @@ -58,8 +56,6 @@ def distill_from_model(
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 vocabulary_quantization: The number of clusters to use for vocabulary quantization. If this is None, no
quantization is performed.
:param pooling: The pooling mode to use for creating embeddings. Can be one of:
'mean' (default): mean over all tokens. Robust and works well in most cases.
'last': use the last token's hidden state (often the [EOS] token). Common for decoder-style models.
Expand Down Expand Up @@ -89,20 +85,10 @@ def distill_from_model(
pooling=pooling,
)

# Apply quantization
if vocabulary_quantization is not None:
weights = compute_weights(len(embeddings), sif_coefficient=sif_coefficient)
embeddings, token_mapping, weights = quantize_vocabulary(
n_clusters=vocabulary_quantization, weights=weights, embeddings=embeddings
)
embeddings = apply_pca(embeddings, pca_dims)
else:
# Post-process the embeddings.
weights = compute_weights(len(embeddings), sif_coefficient=sif_coefficient)
embeddings = apply_pca(embeddings, pca_dims)
embeddings = embeddings * weights[:, None]
weights = None
token_mapping = None
# Post-process the embeddings.
weights = compute_weights(len(embeddings), sif_coefficient=sif_coefficient)
embeddings = apply_pca(embeddings, pca_dims)
embeddings = embeddings * weights[:, None]
# Quantize the embeddings.
embeddings = quantize_embeddings(embeddings, quantize_to)

Expand All @@ -122,8 +108,6 @@ def distill_from_model(

return _create_static_model(
embeddings=embeddings,
weights=weights,
token_mapping=token_mapping,
tokenizer_model=tokenizer_model,
config=config,
model_name=model_name,
Expand Down Expand Up @@ -227,8 +211,6 @@ def init_from_tokenizer(

return _create_static_model(
embeddings=embeddings,
weights=None,
token_mapping=None,
tokenizer_model=tokenizer_model,
config=config,
model_name=tokenizer_name,
Expand Down Expand Up @@ -270,17 +252,13 @@ def _prepare_tokenizer(

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.
Expand All @@ -305,8 +283,6 @@ def _create_static_model(

return StaticModel(
vectors=embeddings,
weights=weights,
token_mapping=token_mapping,
tokenizer=tokenizer_model.to_tokenizer(),
config=config,
base_model_name=model_name,
Expand Down Expand Up @@ -352,7 +328,6 @@ def distill(
token_remove_pattern: str | None = r"\[unused\d+\]",
trust_remote_code: bool = False,
quantize_to: DType | str = DType.Float16,
vocabulary_quantization: int | None = None,
pooling: PoolingMode | str = PoolingMode.MEAN,
) -> StaticModel:
"""Distill a staticmodel from a sentence transformer.
Expand All @@ -377,8 +352,6 @@ def distill(
: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 vocabulary_quantization: The number of clusters to use for vocabulary quantization. If this is None, no
quantization is performed.
:param pooling: The pooling mode to use for creating embeddings. Can be one of:
'mean' (default): mean over all tokens. Robust and works well in most cases.
'last': use the last token's hidden state (often the [EOS] token). Common for decoder-style models.
Expand All @@ -402,6 +375,5 @@ def distill(
token_remove_pattern=token_remove_pattern,
sif_coefficient=sif_coefficient,
quantize_to=quantize_to,
vocabulary_quantization=vocabulary_quantization,
pooling=pooling,
)
30 changes: 0 additions & 30 deletions tests/test_distillation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,36 +93,6 @@ def test_distill_from_model(
assert token in static_model.tokens or normalized in static_model.tokens


@patch.object(import_module("model2vec.distill.distillation"), "model_info")
@patch("transformers.AutoModel.from_pretrained")
def test_distill_quantization(
mock_auto_model: MagicMock,
mock_model_info: MagicMock,
mock_berttokenizer: PreTrainedTokenizerFast,
mock_transformer: PreTrainedModel,
) -> None:
"""Test distill function with different parameters."""
# Mock the return value of model_info to avoid calling the Hugging Face API
mock_model_info.return_value = type("ModelInfo", (object,), {"cardData": {"language": "en"}})
mock_auto_model.return_value = mock_transformer

static_model = distill_from_model(
model=mock_transformer,
tokenizer=mock_berttokenizer,
vocabulary=None,
device="cpu",
pca_dims="auto",
sif_coefficient=1e-4,
token_remove_pattern=None,
vocabulary_quantization=3,
)

assert static_model.embedding.shape == (3, 768)
assert static_model.weights is not None
assert static_model.token_mapping is not None
assert len(static_model.weights) == static_model.tokenizer.get_vocab_size()


@patch.object(import_module("model2vec.distill.distillation"), "model_info")
@patch("transformers.AutoModel.from_pretrained")
def test_distill_removal_pattern_all_tokens(
Expand Down
Loading