chore: typing on loading should return a T, not a StaticModel

model2vec/model.py

@@ -3,11 +3,11 @@
 import json
 import math
 import os
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Iterator, Sequence
 from logging import getLogger
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Any, cast, overload
+from typing import Any, Mapping, TypeVar, cast, overload
 
 import numpy as np
 from joblib import delayed
@@ -191,7 +191,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,
@@ -201,7 +201,7 @@ def from_pretrained(
         vocabulary_quantization: int | None = None,
         max_length: int | None | _UnsetType = _UNSET,
         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.
@@ -226,6 +226,7 @@ def from_pretrained(
         """
         return _loading_helper(
             cls=cls,
+            max_length=max_length,
             path=path,
             token=token,
             vocabulary_quantization=vocabulary_quantization,
@@ -234,7 +235,6 @@ def from_pretrained(
             normalize=normalize,
             subfolder=subfolder,
             force_download=force_download,
-            max_length=max_length,
         )
 
     @overload
@@ -468,11 +468,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.
@@ -505,7 +505,7 @@ def quantize_model(
             dimensionality=dimensionality,
         )
 
-    return StaticModel(
+    return type(model)(
         vectors=embeddings,
         tokenizer=model.tokenizer,
         config=model.config,
@@ -519,7 +519,7 @@ def quantize_model(
 
 
 def _loading_helper(
-    cls: type[StaticModel],
+    cls: type[T],
     path: PathLike,
     token: str | None,
     vocabulary_quantization: int | None,
@@ -529,7 +529,7 @@ def _loading_helper(
     subfolder: str | None,
     force_download: bool,
     max_length: int | None | _UnsetType,
-) -> StaticModel:
+) -> T:
     """Helper function to load a model from a directory."""
     from model2vec.persistence import load_pretrained
 
@@ -581,3 +581,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)

model2vec/train/base.py

@@ -139,12 +139,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.")
@@ -154,13 +154,13 @@ def from_pretrained(
 
     @classmethod
     def from_static_model(
-        cls: type[ModelType],
+        cls: type[T],
         *,
         model: StaticModel,
         pad_token: str | None = None,
         max_length: int | None = None,
         **kwargs: Any,
-    ) -> ModelType:
+    ) -> T:
         """Load the model from a static model.
 
         :param model: The static model to load from.
@@ -430,4 +430,4 @@ def _create_datasets(
         return train_dataset, val_dataset
 
 
-ModelType = TypeVar("ModelType", bound=BaseFinetuneable)
+T = TypeVar("T", bound=BaseFinetuneable)

model2vec/train/similarity.py

@@ -64,7 +64,7 @@ def __init__(
         )
 
     def fit(
-        self: _T,
+        self: T,
         X: list[str],
         y: torch.Tensor,
         learning_rate: float = 1e-3,
@@ -78,7 +78,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.
@@ -132,4 +132,4 @@ def fit(
         return self
 
 
-_T = TypeVar("_T", bound=StaticModelForSimilarity)
+T = TypeVar("T", bound=StaticModelForSimilarity)