Smarter saving

docs/encoders.md

@@ -1,20 +1,32 @@
 # Encoders
 
+!!! warning
+    We now recommend that you do NOT pre-load `SentenceTransformer` encoder models, but rather pass them by name to the topic models.
+    This allows the model not to store the encoder model on disk, and load it when loading the model.
+    This can lead to extreme reductions in disk space usage.
+    For example, instead of writing:
+    ```python
+    model = SensTopic(
+        encoder=SentenceTransformer("intfloat/multilingual-e5-large-instruct", default_prompt_name="query")
+    )
+    ```
+    Write:
+    ```python
+    model = SensTopic(
+        encoder="intfloat/multilingual-e5-large-instruct",
+        trf_kwargs=dict(default_prompt_name="query")
+    )
+    ```
+
+
 Turftopic by default encodes documents using sentence transformers.
 You can always change the encoder model either by passing the name of a sentence transformer from the Huggingface Hub to a model, or by passing a `SentenceTransformer` instance.
 
 Here's an example of building a multilingual topic model by using multilingual embeddings:
 
 ```python
-from sentence_transformers import SentenceTransformer
 from turftopic import GMM
 
-trf = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
-
-model = GMM(10, encoder=trf)
-
-# or
-
 model = GMM(10, encoder="paraphrase-multilingual-MiniLM-L12-v2")
 ```
 
@@ -35,31 +47,36 @@ In this case, documents will serve as the queries and words as the passages:
 from turftopic import KeyNMF
 from sentence_transformers import SentenceTransformer
 
-encoder = SentenceTransformer(
-    "intfloat/multilingual-e5-large-instruct",
-    prompts={
-        "query": "Instruct: Retrieve relevant keywords from the given document. Query: "
-        "passage": "Passage: "
-    },
-    # Make sure to set default prompt to query!
-    default_prompt_name="query",
+model = KeyNMF(
+    10,
+    encoder="intfloat/multilingual-e5-large-instruct",
+    # These are the arguments that get applied when loading the encoder.
+    trf_kwargs=dict(
+        prompts={
+            "query": "Instruct: Retrieve relevant keywords from the given document. Query: "
+            "passage": "Passage: "
+        },
+        # Make sure to set default prompt to query!
+        default_prompt_name="query",
+    )
 )
-model = KeyNMF(10, encoder=encoder)
 ```
 
 And a regular, asymmetric example:
 
 ```python
-encoder = SentenceTransformer(
-    "intfloat/e5-large-v2",
-    prompts={
-        "query": "query: "
-        "passage": "passage: "
-    },
-    # Make sure to set default prompt to query!
-    default_prompt_name="query",
+model = KeyNMF(
+    10,
+    encoder="intfloat/e5-large-v2",
+    trf_kwargs=dict(
+        prompts={
+            "query": "query: "
+            "passage": "passage: "
+        },
+        # Make sure to set default prompt to query!
+        default_prompt_name="query",
+    )
 )
-model = KeyNMF(10, encoder=encoder)
 ```
 
 ## Performance tips
@@ -75,9 +92,8 @@ pip install sentence-transformers[onnx, onnx-gpu]
 from turftopic import SemanticSignalSeparation
 from sentence_transformers import SentenceTransformer
 
-encoder = SentenceTransformer("all-MiniLM-L6-v2", backend="onnx")
 
-model = SemanticSignalSeparation(10, encoder=encoder)
+model = SemanticSignalSeparation(10, encoder="all-MiniLM-L6-v2", trf_kwargs=dict(backend="onnx"))
 ```
 
 ## External Embeddings

docs/persistence.md

@@ -3,6 +3,24 @@
 ## Model persistence
 All models in Turftopic can be serialized and saved to disk, or published to the HuggingFace Hub.
 
+!!! warning
+    We now recommend that you do NOT pre-load `SentenceTransformer` encoder models, but rather pass them by name to the topic models.
+    This allows the model not to store the encoder model on disk, and load it when loading the model.
+    This can lead to extreme reductions in disk space usage.
+    For example, instead of writing:
+    ```python
+    model = SensTopic(
+        encoder=SentenceTransformer("intfloat/multilingual-e5-large-instruct", default_prompt_name="query")
+    )
+    ```
+    Write:
+    ```python
+    model = SensTopic(
+        encoder="intfloat/multilingual-e5-large-instruct",
+        trf_kwargs=dict(default_prompt_name="query")
+    )
+    ```
+
 ### Saving locally
 
 Turftopic models can now be saved to disk using the `to_disk()` method of models:

pyproject.toml

@@ -9,7 +9,7 @@ profile = "black"
 
 [project]
 name = "turftopic"
-version = "0.26.1"
+version = "0.27.0"
 description = "Topic modeling with contextual representations from sentence transformers."
 authors = [
    { name = "Márton Kardos <power.up1163@gmail.com>", email = "martonkardos@cas.au.dk" }

turftopic/base.py

@@ -22,6 +22,18 @@
 class ContextualModel(BaseEstimator, TransformerMixin, TopicContainer):
     """Base class for contextual topic models in Turftopic."""
 
+    def load_encoder(self):
+        if isinstance(self.encoder, str):
+            if self.trf_kwargs is None:
+                trf_kwargs = dict()
+            else:
+                trf_kwargs = self.trf_kwargs
+            self.encoder_ = SentenceTransformer(self.encoder, **trf_kwargs)
+            self._encoder_preloaded = False
+        else:
+            self.encoder_ = self.encoder
+            self._encoder_preloaded = True
+
     @property
     def has_negative_side(self) -> bool:
         return False
@@ -41,7 +53,11 @@ def encode_documents(self, raw_documents: Iterable[str]) -> np.ndarray:
         """
         if not hasattr(self.encoder_, "encode"):
             return self.encoder.get_text_embeddings(list(raw_documents))
-        return self.encoder_.encode(list(raw_documents))
+        if getattr(self, "encode_kwargs", None) is None:
+            encode_kwargs = dict()
+        else:
+            encode_kwargs = self.encode_kwargs
+        return self.encoder_.encode(list(raw_documents), **encode_kwargs)
 
     @abstractmethod
     def fit_transform(
@@ -173,7 +189,13 @@ def to_disk(self, out_dir: Union[Path, str]):
         package_versions = get_package_versions()
         with out_dir.joinpath("package_versions.json").open("w") as ver_file:
             ver_file.write(json.dumps(package_versions))
-        joblib.dump(self, out_dir.joinpath("model.joblib"))
+        if getattr(self, "_encoder_preloaded", True):
+            joblib.dump(self, out_dir.joinpath("model.joblib"))
+        else:
+            encoder_ = self.encoder_
+            delattr(self, "encoder_")
+            joblib.dump(self, out_dir.joinpath("model.joblib"))
+            self.encoder_ = encoder_
 
     def push_to_hub(self, repo_id: str):
         """Uploads model to HuggingFace Hub

turftopic/models/cluster.py

@@ -180,6 +180,10 @@ class ClusteringTopicModel(
         if 'centroid' the centroid vectors of clusters will be used as topic vectors (Top2Vec).
     random_state: int, default None
         Random state to use so that results are exactly reproducible.
+    trf_kwargs: dict, default None
+        Keyword arguments to apply when loading the Encoder model.
+    encode_kwargs: dict, default None
+        Keyword arguments to apply encoding documents with the encoder.
     """
 
     def __init__(
@@ -196,9 +200,13 @@ def __init__(
         reduction_distance_metric: DistanceMetric = "cosine",
         reduction_topic_representation: TopicRepresentation = "component",
         random_state: Optional[int] = None,
+        trf_kwargs=None,
+        encode_kwargs=None,
     ):
         self.encoder = encoder
         self.random_state = random_state
+        self.trf_kwargs = trf_kwargs
+        self.encode_kwargs = encode_kwargs
         if feature_importance not in VALID_WORD_IMPORTANCE:
             raise ValueError(
                 f"feature_importance must be one of {VALID_WORD_IMPORTANCE} got {feature_importance} instead."
@@ -217,10 +225,7 @@ def __init__(
             )
         if isinstance(encoder, int):
             raise TypeError(integer_message)
-        if isinstance(encoder, str):
-            self.encoder_ = SentenceTransformer(encoder)
-        else:
-            self.encoder_ = encoder
+        self.load_encoder()
         self.validate_encoder()
         if vectorizer is None:
             self.vectorizer = default_vectorizer()
@@ -766,6 +771,8 @@ def __init__(
         reduction_distance_metric: DistanceMetric = "cosine",
         reduction_topic_representation: TopicRepresentation = "component",
         random_state: Optional[int] = None,
+        trf_kwargs=None,
+        encode_kwargs=None,
     ):
         if dimensionality_reduction is None:
             try:
@@ -798,6 +805,8 @@ def __init__(
             reduction_method=reduction_method,
             reduction_distance_metric=reduction_distance_metric,
             reduction_topic_representation=reduction_topic_representation,
+            trf_kwargs=trf_kwargs,
+            encode_kwargs=encode_kwargs,
         )
 
 
@@ -834,6 +843,8 @@ def __init__(
         reduction_distance_metric: DistanceMetric = "cosine",
         reduction_topic_representation: TopicRepresentation = "centroid",
         random_state: Optional[int] = None,
+        trf_kwargs=None,
+        encode_kwargs=None,
     ):
         if dimensionality_reduction is None:
             try:
@@ -866,6 +877,8 @@ def __init__(
             reduction_method=reduction_method,
             reduction_distance_metric=reduction_distance_metric,
             reduction_topic_representation=reduction_topic_representation,
+            trf_kwargs=trf_kwargs,
+            encode_kwargs=encode_kwargs,
         )
 
 
@@ -911,6 +924,8 @@ def __init__(
         step_size: Optional[int] = 40,
         pooling: Optional[Callable] = np.nanmean,
         random_state: Optional[int] = None,
+        trf_kwargs=None,
+        encode_kwargs=None,
     ):
         if dimensionality_reduction is None:
             try:
@@ -956,6 +971,8 @@ def __init__(
             reduction_method=reduction_method,
             reduction_distance_metric=reduction_distance_metric,
             reduction_topic_representation=reduction_topic_representation,
+            trf_kwargs=trf_kwargs,
+            encode_kwargs=encode_kwargs,
         )
         super().__init__(
             model,

turftopic/models/ctm.py

@@ -139,6 +139,10 @@ class AutoEncodingTopicModel(ContextualModel, MultimodalModel):
         Number of epochs to run during training.
     random_state: int, default None
         Random state to use so that results are exactly reproducible.
+    trf_kwargs: dict, default None
+        Keyword arguments to apply when loading the Encoder model.
+    encode_kwargs: dict, default None
+        Keyword arguments to apply encoding documents with the encoder.
     """
 
     def __init__(
@@ -155,14 +159,15 @@ def __init__(
         learning_rate: float = 1e-2,
         n_epochs: int = 50,
         random_state: Optional[int] = None,
+        trf_kwargs=None,
+        encode_kwargs=None,
     ):
         self.n_components = n_components
         self.random_state = random_state
         self.encoder = encoder
-        if isinstance(encoder, str):
-            self.encoder_ = SentenceTransformer(encoder)
-        else:
-            self.encoder_ = encoder
+        self.trf_kwargs = trf_kwargs
+        self.encode_kwargs = encode_kwargs
+        self.load_encoder()
         self.validate_encoder()
         if vectorizer is None:
             self.vectorizer = default_vectorizer()