From b681cf9067aeaadb75e5241d987a35a917443624 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sun, 30 Aug 2026 17:07:09 +0200 Subject: [PATCH 1/2] fix: loading --- model2vec/model.py | 49 +++++++++++----------------------------------- 1 file changed, 11 insertions(+), 38 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 3a63206..d454660 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -3,12 +3,11 @@ import json import math import os -import warnings from collections.abc import Iterator, Sequence from logging import getLogger from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, overload +from typing import Any, TypeVar, overload import numpy as np from joblib import delayed @@ -172,7 +171,7 @@ def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> l @classmethod def from_pretrained( - cls: type[StaticModel], + cls: type[T], path: PathLike, token: str | None = None, normalize: bool | None = None, @@ -181,7 +180,7 @@ def from_pretrained( dimensionality: int | None = None, vocabulary_quantization: int | None = None, force_download: bool = True, - ) -> StaticModel: + ) -> T: """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. @@ -212,35 +211,6 @@ def from_pretrained( force_download=force_download, ) - @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, - vocabulary_quantization: int | None = None, - force_download: bool = True, - ) -> StaticModel: - """Deprecated: use from_pretrained.""" - warnings.warn( - "StaticModel.from_sentence_transformers() is deprecated; use from_pretrained() instead.", - DeprecationWarning, - stacklevel=2, - ) - return _loading_helper( - cls=cls, - path=path, - token=token, - vocabulary_quantization=vocabulary_quantization, - quantize_to=quantize_to, - dimensionality=dimensionality, - normalize=normalize, - subfolder=None, - force_download=force_download, - ) - @overload def encode_as_sequence( self, @@ -464,11 +434,11 @@ def push_to_hub( def quantize_model( - model: StaticModel, + model: T, vocabulary_quantization: int | None = None, quantize_to: str | DType | None = None, dimensionality: int | None = None, -) -> StaticModel: +) -> T: """Quantize the model to a lower precision and possibly lower dimensionality. :param model: The model to quantize. @@ -501,7 +471,7 @@ def quantize_model( dimensionality=dimensionality, ) - return StaticModel( + return type(model)( vectors=embeddings, tokenizer=model.tokenizer, config=dict(model.config), @@ -514,7 +484,7 @@ def quantize_model( def _loading_helper( - cls: type[StaticModel], + cls: type[T], path: PathLike, token: str | None, vocabulary_quantization: int | None, @@ -523,7 +493,7 @@ def _loading_helper( normalize: bool | None, subfolder: str | None, force_download: bool, -) -> StaticModel: +) -> T: """Helper function to load a model from a directory.""" from model2vec.persistence import load_pretrained @@ -569,3 +539,6 @@ def _get_unk_token_id(tokenizer: Tokenizer) -> int | None: return tokenizer.token_to_id(token) # Unigram return json.loads(tokenizer.to_str())["model"].get("unk_id") + + +T = TypeVar("T", bound=StaticModel) From 5d63298d0245d805e5ac86087fae489255162086 Mon Sep 17 00:00:00 2001 From: stephantul Date: Mon, 31 Aug 2026 09:15:14 +0200 Subject: [PATCH 2/2] everything to 'T' for naming consistency --- model2vec/train/base.py | 10 +++++----- model2vec/train/similarity.py | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 16a032c..3ec4da2 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -135,12 +135,12 @@ def _initialize(self) -> None: @classmethod def from_pretrained( - cls: type[ModelType], + cls: type[T], path: PathLike = "minishlab/potion-base-32m", *, token: str | None = None, **kwargs: Any, - ) -> ModelType: + ) -> T: """Load the model from a pretrained model2vec model.""" if model_name := kwargs.pop("model_name", None): logger.warning("The 'model_name' argument is deprecated. Use 'path' instead.") @@ -150,12 +150,12 @@ def from_pretrained( @classmethod def from_static_model( - cls: type[ModelType], + cls: type[T], *, model: StaticModel, pad_token: str | None = None, **kwargs: Any, - ) -> ModelType: + ) -> T: """Load the model from a static model.""" model.embedding = np.nan_to_num(model.embedding) weights = torch.from_numpy(model.weights) if model.weights is not None else None @@ -412,4 +412,4 @@ def _create_datasets( return train_dataset, val_dataset -ModelType = TypeVar("ModelType", bound=BaseFinetuneable) +T = TypeVar("T", bound=BaseFinetuneable) diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index d2e7068..7f73bce 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -61,7 +61,7 @@ def __init__( ) def fit( - self: _T, + self: T, X: list[str], y: torch.Tensor, learning_rate: float = 1e-3, @@ -75,7 +75,7 @@ def fit( y_val: torch.Tensor | None = None, validation_steps: int | None = None, random_seed: int = DEFAULT_RANDOM_SEED, - ) -> _T: + ) -> T: """Fit a model. This function trains the model with a plain torch training loop. @@ -129,4 +129,4 @@ def fit( return self -_T = TypeVar("_T", bound=StaticModelForSimilarity) +T = TypeVar("T", bound=StaticModelForSimilarity)