feat: remove flag argument

Author: stephantul

model2vec/model.py

@@ -12,7 +12,7 @@
 from tokenizers import Encoding, Tokenizer
 from tqdm import tqdm
 
-from model2vec.quantization import DType, quantize_embeddings
+from model2vec.quantization import DType, quantize_and_reduce_dim
 from model2vec.utils import ProgressParallel, load_local_model
 
 PathLike = Union[Path, str]
@@ -147,51 +147,6 @@ def tokenize(self, sentences: list[str], max_length: int | None = None) -> list[
 
         return encodings_ids
 
-    @classmethod
-    def _load_model(
-        cls: type[StaticModel],
-        path: PathLike,
-        token: str | None,
-        normalize: bool | None,
-        quantize_to: str | DType | None,
-        dimensionality: int | None,
-        from_sentence_transformers: bool,
-        subfolder: str | None = None,
-    ) -> StaticModel:
-        """Helper function to load a model from a path and optionally quantize it/reduce the dimensionality."""
-        from model2vec.hf_utils import load_pretrained
-
-        embeddings, tokenizer, config, metadata = load_pretrained(
-            folder_or_repo_path=path,
-            token=token,
-            from_sentence_transformers=from_sentence_transformers,
-            subfolder=subfolder,
-        )
-
-        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]:
-                raise ValueError(
-                    f"Dimensionality {dimensionality} is greater than the model dimensionality {embeddings.shape[1]}"
-                )
-            embeddings = embeddings[:, :dimensionality]
-            if config.get("apply_pca", None) is None:
-                logger.warning(
-                    "You are reducing the dimensionality of the model, but we can't find a pca key in the model config. This might not work as expected."
-                )
-
-        return cls(
-            embeddings,
-            tokenizer,
-            config,
-            normalize=normalize,
-            base_model_name=metadata.get("base_model"),
-            language=metadata.get("language"),
-        )
-
     @classmethod
     def from_pretrained(
         cls: type[StaticModel],
@@ -218,16 +173,30 @@ def from_pretrained(
             Note that this only applies if you have trained your model using mrl or PCA.
         :return: A StaticModel.
         """
-        return cls._load_model(
-            path=path,
+        from model2vec.hf_utils import load_pretrained
+
+        embeddings, tokenizer, config, metadata = load_pretrained(
+            folder_or_repo_path=path,
             token=token,
-            normalize=normalize,
-            quantize_to=quantize_to,
-            dimensionality=dimensionality,
             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],
@@ -252,13 +221,28 @@ def from_sentence_transformers(
             Note that this only applies if you have trained your model using mrl or PCA.
         :return: A StaticModel.
         """
-        return cls._load_model(
-            path=path,
+        from model2vec.hf_utils import load_pretrained
+
+        embeddings, tokenizer, config, metadata = load_pretrained(
+            folder_or_repo_path=path,
             token=token,
-            normalize=normalize,
+            from_sentence_transformers=True,
+            subfolder=None,
+        )
+
+        embeddings = quantize_and_reduce_dim(
+            embeddings=embeddings,
             quantize_to=quantize_to,
             dimensionality=dimensionality,
-            from_sentence_transformers=True,
+        )
+
+        return cls(
+            embeddings,
+            tokenizer,
+            config,
+            normalize=normalize,
+            base_model_name=metadata.get("base_model"),
+            language=metadata.get("language"),
         )
 
     def encode_as_sequence(

model2vec/quantization.py

@@ -33,3 +33,29 @@ def quantize_embeddings(embeddings: np.ndarray, quantize_to: DType) -> np.ndarra
         return quantized
     else:
         raise ValueError("Not a valid enum member of DType.")
+
+
+def quantize_and_reduce_dim(
+    embeddings: np.ndarray, quantize_to: DType | str | 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]:
+            raise ValueError(
+                f"Dimensionality {dimensionality} is greater than the model dimensionality {embeddings.shape[1]}"
+            )
+        embeddings = embeddings[:, :dimensionality]
+
+    return embeddings