diff --git a/docs/encoders.md b/docs/encoders.md index ef791281..3ea87a7d 100644 --- a/docs/encoders.md +++ b/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 diff --git a/docs/persistence.md b/docs/persistence.md index 94825570..12dc0315 100644 --- a/docs/persistence.md +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index eded9c19..db3cc10c 100644 --- a/pyproject.toml +++ b/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 ", email = "martonkardos@cas.au.dk" } diff --git a/turftopic/base.py b/turftopic/base.py index 6606c650..d74f6987 100644 --- a/turftopic/base.py +++ b/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 diff --git a/turftopic/models/cluster.py b/turftopic/models/cluster.py index feb59988..ca828fba 100644 --- a/turftopic/models/cluster.py +++ b/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, diff --git a/turftopic/models/ctm.py b/turftopic/models/ctm.py index e432f4b5..046569e7 100644 --- a/turftopic/models/ctm.py +++ b/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() diff --git a/turftopic/models/cvp.py b/turftopic/models/cvp.py index 8b9b64c1..37c985d2 100644 --- a/turftopic/models/cvp.py +++ b/turftopic/models/cvp.py @@ -10,14 +10,14 @@ from sentence_transformers import SentenceTransformer from sklearn.base import BaseEstimator, TransformerMixin -from turftopic.base import Encoder +from turftopic.base import ContextualModel, Encoder from turftopic.encoders.multimodal import MultimodalEncoder from turftopic.serialization import create_readme, get_package_versions Seeds = tuple[list[str], list[str]] -class ConceptVectorProjection(BaseEstimator, TransformerMixin): +class ConceptVectorProjection(ContextualModel): """Concept Vector Projection model from [Lyngbæk et al. (2025)](https://doi.org/10.63744/nVu1Zq5gRkuD) Can be used to project document embeddings onto a difference projection vector between positive and negative seed phrases. The primary use case is sentiment analysis, and continuous sentiment scores, @@ -34,6 +34,10 @@ class ConceptVectorProjection(BaseEstimator, TransformerMixin): encoder: str or SentenceTransformer Model to produce document representations, paraphrase-multilingual-mpnet-base-v2 is the default per Lyngbæk et al. (2025). + 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__( @@ -42,6 +46,8 @@ def __init__( encoder: Union[ Encoder, str, MultimodalEncoder ] = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2", + trf_kwargs=None, + encode_kwargs=None, ): self.seeds = seeds if isinstance(seeds, OrderedDict): @@ -55,15 +61,14 @@ def __init__( else: self._seeds = OrderedDict(seeds) 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.classes_ = np.array([name for name in self._seeds]) self.concept_matrix_ = [] for _, (positive, negative) in self._seeds.items(): - positive_emb = self.encoder_.encode(list(positive)) - negative_emb = self.encoder_.encode(list(negative)) + positive_emb = self.encode_documents(list(positive)) + negative_emb = self.encode_documents(list(negative)) cv = np.mean(positive_emb, axis=0) - np.mean(negative_emb, axis=0) self.concept_matrix_.append(cv / np.linalg.norm(cv)) self.concept_matrix_ = np.stack(self.concept_matrix_) @@ -92,7 +97,7 @@ def fit_transform(self, raw_documents=None, y=None, embeddings=None): "Either embeddings or raw_documents has to be passed, both are None." ) if embeddings is None: - embeddings = self.encoder_.encode(list(raw_documents)) + embeddings = self.encode_documents(list(raw_documents)) return embeddings @ self.concept_matrix_.T def transform(self, raw_documents=None, embeddings=None): @@ -111,39 +116,3 @@ def transform(self, raw_documents=None, embeddings=None): Prevalance of each concept in each document. """ return self.fit_transform(raw_documents, embeddings=embeddings) - - def to_disk(self, out_dir: Union[Path, str]): - """Persists model to directory on your machine. - - Parameters - ---------- - out_dir: Path | str - Directory to save the model to. - """ - out_dir = Path(out_dir) - out_dir.mkdir(exist_ok=True) - 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")) - - def push_to_hub(self, repo_id: str): - """Uploads model to HuggingFace Hub - - Parameters - ---------- - repo_id: str - Repository to upload the model to. - """ - api = HfApi() - api.create_repo(repo_id, exist_ok=True) - with tempfile.TemporaryDirectory() as tmp_dir: - readme_path = Path(tmp_dir).joinpath("README.md") - with readme_path.open("w") as readme_file: - readme_file.write(create_readme(self, repo_id)) - self.to_disk(tmp_dir) - api.upload_folder( - folder_path=tmp_dir, - repo_id=repo_id, - repo_type="model", - ) diff --git a/turftopic/models/decomp.py b/turftopic/models/decomp.py index d2a3facd..00d150a4 100644 --- a/turftopic/models/decomp.py +++ b/turftopic/models/decomp.py @@ -67,6 +67,10 @@ class SemanticSignalSeparation( or their combination ('combined') should determine the word's importance for a topic. 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__( @@ -82,14 +86,15 @@ def __init__( "axial", "angular", "combined" ] = "combined", random_state: Optional[int] = None, + trf_kwargs=None, + encode_kwargs=None, ): self.n_components = n_components self.encoder = encoder self.feature_importance = feature_importance - 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() diff --git a/turftopic/models/fastopic.py b/turftopic/models/fastopic.py index 1642964b..bec0bc3f 100644 --- a/turftopic/models/fastopic.py +++ b/turftopic/models/fastopic.py @@ -59,6 +59,10 @@ class FASTopic(ContextualModel): Learning rate for the ADAM optimizer. device: str, default "cpu" Device to run the model on. Defaults to CPU. + 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__( @@ -76,14 +80,15 @@ def __init__( n_epochs: int = 200, learning_rate: float = 0.002, device: str = "cpu", + trf_kwargs=None, + encode_kwargs=None, ): self.n_components = n_components self.encoder = encoder + self.trf_kwargs = trf_kwargs + self.encode_kwargs = encode_kwargs self.random_state = random_state - if isinstance(encoder, str): - self.encoder_ = SentenceTransformer(encoder) - else: - self.encoder_ = encoder + self.load_encoder() if vectorizer is None: self.vectorizer = default_vectorizer() else: @@ -153,7 +158,7 @@ def fit_transform( with console.status("Fitting model") as status: if embeddings is None: status.update("Encoding documents") - embeddings = self.encoder_.encode(raw_documents) + embeddings = self.encode_documents(raw_documents) console.log("Documents encoded.") self.train_doc_embeddings = embeddings status.update("Extracting terms.") @@ -188,7 +193,7 @@ def transform( Document-topic matrix. """ if embeddings is None: - embeddings = self.encoder_.encode(raw_documents) + embeddings = self.encode_documents(raw_documents) with torch.no_grad(): self.model.eval() theta = self.model.get_theta( diff --git a/turftopic/models/gmm.py b/turftopic/models/gmm.py index cf2b7c53..04d5210b 100644 --- a/turftopic/models/gmm.py +++ b/turftopic/models/gmm.py @@ -91,6 +91,10 @@ class GMM(ContextualModel, DynamicTopicModel, MultimodalModel): not embedding-based ones.* 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. Attributes ---------- @@ -110,16 +114,17 @@ def __init__( weight_prior: Literal["dirichlet", "dirichlet_process", None] = None, gamma: Optional[float] = None, random_state: Optional[int] = None, + trf_kwargs=None, + encode_kwargs=None, ): self.n_components = n_components self.encoder = encoder + self.encode_kwargs = encode_kwargs + self.trf_kwargs = trf_kwargs self.weight_prior = weight_prior self.gamma = gamma self.random_state = random_state - 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() diff --git a/turftopic/models/keynmf.py b/turftopic/models/keynmf.py index c9618a7b..f00562a6 100644 --- a/turftopic/models/keynmf.py +++ b/turftopic/models/keynmf.py @@ -69,6 +69,10 @@ class KeyNMF(ContextualModel, DynamicTopicModel, MultimodalModel): This is useful when you have a corpus containing multiple languages. term_match_threshold: float, default 0.9 Cosine similarity threshold for matching terms across languages. + 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__( @@ -85,6 +89,8 @@ def __init__( seed_exponent: float = 2.0, cross_lingual: bool = False, term_match_threshold: float = 0.9, + trf_kwargs=None, + encode_kwargs=None, ): self.random_state = random_state self.n_components = n_components @@ -92,6 +98,8 @@ def __init__( self.seed_exponent = seed_exponent self.metric = metric self.encoder = encoder + self.trf_kwargs = trf_kwargs + self.encode_kwargs = encode_kwargs self._has_custom_vectorizer = vectorizer is not None if isinstance(encoder, str): if ( @@ -101,9 +109,7 @@ def __init__( "all-MiniLM is incompatible with cross-lingual transfer, using paraphrase-multilingual-MiniLM-L12-v2." ) encoder = "paraphrase-multilingual-MiniLM-L12-v2" - self.encoder_ = SentenceTransformer(encoder) - else: - self.encoder_ = encoder + self.load_encoder() self.validate_encoder() if vectorizer is None: self.vectorizer = default_vectorizer() @@ -121,7 +127,7 @@ def __init__( self.seed_phrase = seed_phrase self.seed_embedding = None if self.seed_phrase is not None: - self.seed_embedding = self.encoder_.encode([self.seed_phrase])[0] + self.seed_embedding = self.encode_documents([self.seed_phrase])[0] self.cross_lingual = cross_lingual self.term_match_threshold = term_match_threshold diff --git a/turftopic/models/senstopic.py b/turftopic/models/senstopic.py index 3a992464..016543dc 100644 --- a/turftopic/models/senstopic.py +++ b/turftopic/models/senstopic.py @@ -89,6 +89,10 @@ class SensTopic(ContextualModel, DynamicTopicModel, MultimodalModel): L1 penalty applied to document-topic proportions. Higher values push the model to assign fewer topics to a single document, while lower values will distribute topics across documents. + 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__( @@ -104,14 +108,15 @@ def __init__( ] = "combined", random_state: Optional[int] = None, sparsity: float = 1, + trf_kwargs=None, + encode_kwargs=None, ): self.n_components = n_components self.encoder = encoder + self.trf_kwargs = trf_kwargs + self.encode_kwargs = encode_kwargs self.feature_importance = feature_importance - 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() @@ -346,7 +351,7 @@ def partial_fit( def transform(self, raw_documents, embeddings=None): if embeddings is None: - embeddings = self.encoder_.encode(raw_documents) + embeddings = self.encode_documents(raw_documents) return self.decomposition.transform(embeddings) def fit_transform_multimodal( diff --git a/turftopic/models/topeax.py b/turftopic/models/topeax.py index 52f05742..e3c7cc00 100644 --- a/turftopic/models/topeax.py +++ b/turftopic/models/topeax.py @@ -141,6 +141,10 @@ class Topeax(GMM): Number of neighbours to take into account when running TSNE. 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. """ @@ -152,6 +156,8 @@ def __init__( vectorizer: Optional[CountVectorizer] = None, perplexity: int = 50, random_state: Optional[int] = None, + trf_kwargs=None, + encode_kwargs=None, ): dimensionality_reduction = TSNE( 2, @@ -166,6 +172,8 @@ def __init__( vectorizer=vectorizer, dimensionality_reduction=dimensionality_reduction, random_state=random_state, + trf_kwargs=trf_kwargs, + encode_kwargs=encode_kwargs, ) def estimate_components( diff --git a/turftopic/serialization.py b/turftopic/serialization.py index 8336a4ed..5ffccf90 100644 --- a/turftopic/serialization.py +++ b/turftopic/serialization.py @@ -110,7 +110,12 @@ def load_model(repo_id_or_path: Union[str, Path]): remote_versions = json.loads(ver_file.read()) validate_package_versions(remote_versions) model = joblib.load(path.joinpath("model.joblib")) - return model else: in_dir = snapshot_download(repo_id=repo_id_or_path) - return load_model(in_dir) + model = load_model(in_dir) + if getattr(model, "encoder_", None) is None: + print( + "Model does not come with prepackaged encoder_ attribute, loading it." + ) + model.load_encoder() + return model