diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..a61ae5a8 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1,7 +1,8 @@ # ruff: noqa: E402 -import yaml import warnings +import yaml + warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) @@ -10,26 +11,25 @@ except (KeyError, AttributeError, TypeError): pass -import re +import collections +import inspect import math +import re +from collections import Counter, defaultdict +from copy import deepcopy +from importlib.util import find_spec +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Literal, Mapping, Tuple, Union + import joblib -import inspect -import collections import numpy as np import pandas as pd import scipy.sparse as sp -from copy import deepcopy - -from tqdm import tqdm -from pathlib import Path from packaging import version -from tempfile import TemporaryDirectory -from collections import defaultdict, Counter -from scipy.sparse import csr_matrix from scipy.cluster import hierarchy as sch -from importlib.util import find_spec - -from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, TYPE_CHECKING, Literal +from scipy.sparse import csr_matrix +from tqdm import tqdm # Plotting if find_spec("plotly") is None: @@ -41,8 +41,8 @@ from bertopic import plotting if TYPE_CHECKING: - import plotly.graph_objs as go import matplotlib.figure as fig + import plotly.graph_objs as go # Models @@ -54,32 +54,33 @@ HAS_HDBSCAN = False from sklearn.cluster import HDBSCAN as SK_HDBSCAN -from sklearn.preprocessing import normalize from sklearn import __version__ as sklearn_version from sklearn.cluster import AgglomerativeClustering from sklearn.decomposition import PCA -from sklearn.metrics.pairwise import cosine_similarity from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer +from sklearn.metrics.pairwise import cosine_similarity +from sklearn.preprocessing import normalize -# BERTopic -from bertopic.cluster import BaseCluster -from bertopic.backend import BaseEmbedder -from bertopic.representation._mmr import mmr -from bertopic.backend._utils import select_backend -from bertopic.vectorizers import ClassTfidfTransformer -from bertopic.representation import BaseRepresentation, KeyBERTInspired -from bertopic.dimensionality import BaseDimensionalityReduction -from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan +import bertopic._save_utils as save_utils from bertopic._utils import ( MyLogger, check_documents_type, check_embeddings_shape, check_is_fitted, - validate_distance_matrix, - select_topic_representation, get_unique_distances, + select_topic_representation, + validate_distance_matrix, ) -import bertopic._save_utils as save_utils +from bertopic.backend import BaseEmbedder +from bertopic.backend._utils import select_backend + +# BERTopic +from bertopic.cluster import BaseCluster +from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan +from bertopic.dimensionality import BaseDimensionalityReduction +from bertopic.representation import BaseRepresentation, KeyBERTInspired +from bertopic.representation._mmr import mmr +from bertopic.vectorizers import ClassTfidfTransformer logger = MyLogger() logger.configure("WARNING") @@ -353,6 +354,7 @@ def fit( embeddings: np.ndarray = None, images: List[str] | None = None, y: Union[List[int], np.ndarray] = None, + tokenized_documents: List[List[str]] | None = None, ): """Fit the models on a collection of documents and generate topics. @@ -363,6 +365,9 @@ def fit( images: A list of paths to the images to fit on or the images themselves y: The target class for (semi)-supervised modeling. Use -1 if no class for a specific instance is specified. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Examples: ```python @@ -389,7 +394,9 @@ def fit( topic_model = BERTopic().fit(docs, embeddings) ``` """ - self.fit_transform(documents=documents, embeddings=embeddings, y=y, images=images) + self.fit_transform( + documents=documents, embeddings=embeddings, y=y, images=images, tokenized_documents=tokenized_documents + ) return self def fit_transform( @@ -398,6 +405,7 @@ def fit_transform( embeddings: np.ndarray = None, images: List[str] | None = None, y: Union[List[int], np.ndarray] = None, + tokenized_documents: List[List[str]] | None = None, ) -> Tuple[List[int], Union[np.ndarray, None]]: """Fit the models on a collection of documents, generate topics, and return the probabilities and topic per document. @@ -409,6 +417,9 @@ def fit_transform( images: A list of paths to the images to fit on or the images themselves y: The target class for (semi)-supervised modeling. Use -1 if no class for a specific instance is specified. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Returns: predictions: Topic predictions for each documents @@ -452,6 +463,10 @@ def fit_transform( doc_ids = range(len(documents)) if documents is not None else range(len(images)) documents = pd.DataFrame({"Document": documents, "ID": doc_ids, "Topic": None, "Image": images}) + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + # Extract embeddings if embeddings is None: logger.info("Embedding - Transforming documents to embeddings.") @@ -651,6 +666,7 @@ def partial_fit( documents: List[str], embeddings: np.ndarray = None, y: Union[List[int], np.ndarray] = None, + tokenized_documents: List[List[str]] | None = None, ): """Fit BERTopic on a subset of the data and perform online learning with batch-like data. @@ -683,6 +699,9 @@ def partial_fit( instead of the sentence-transformer model y: The target class for (semi)-supervised modeling. Use -1 if no class for a specific instance is specified. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Examples: ```python @@ -719,6 +738,10 @@ def partial_fit( documents = [documents] documents = pd.DataFrame({"Document": documents, "ID": range(len(documents)), "Topic": None}) + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + # Extract embeddings if embeddings is None: if self.topic_representations_ is None: @@ -767,7 +790,10 @@ def partial_fit( # Prepare documents documents_per_topic = documents.sort_values("Topic").groupby(["Topic"], as_index=False) updated_topics = documents_per_topic.first().Topic.astype(int) - documents_per_topic = documents_per_topic.agg({"Document": " ".join}) + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in documents.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = documents_per_topic.agg(agg_dict) # Update topic representations self.c_tf_idf_, updated_words = self._c_tf_idf(documents_per_topic, partial_fit=True) @@ -803,6 +829,7 @@ def topics_over_time( datetime_format: str | None = None, evolution_tuning: bool = True, global_tuning: bool = True, + tokenized_documents: List[List[str]] | None = None, ) -> pd.DataFrame: """Create topics over time. @@ -842,6 +869,9 @@ def topics_over_time( global_tuning: Fine-tune each topic representation at timestamp *t* by averaging its c-TF-IDF matrix with the global c-TF-IDF matrix. Turn this off if you want to prevent words in topic representations that could not be found in the documents at timestamp *t*. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Returns: topics_over_time: A dataframe that contains the topic, words, and frequency of topic @@ -862,6 +892,11 @@ def topics_over_time( check_documents_type(docs) selected_topics = topics if topics else self.topics_ documents = pd.DataFrame({"Document": docs, "Topic": selected_topics, "Timestamps": timestamps}) + + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + global_c_tf_idf = normalize(self.c_tf_idf_, axis=1, norm="l1", copy=False) all_topics = sorted(list(documents.Topic.unique())) @@ -894,9 +929,10 @@ def topics_over_time( for index, timestamp in tqdm(enumerate(timestamps), disable=not self.verbose): # Calculate c-TF-IDF representation for a specific timestamp selection = documents.loc[documents.Timestamps == timestamp, :] - documents_per_topic = selection.groupby(["Topic"], as_index=False).agg( - {"Document": " ".join, "Timestamps": "count"} - ) + agg_dict = {"Document": " ".join, "Timestamps": "count"} + if "Tokenized_Document" in selection.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = selection.groupby(["Topic"], as_index=False).agg(agg_dict) c_tf_idf, words = self._c_tf_idf(documents_per_topic, fit=False) if global_tuning or evolution_tuning: @@ -958,6 +994,7 @@ def topics_per_class( docs: List[str], classes: Union[List[int], List[str]], global_tuning: bool = True, + tokenized_documents: List[List[str]] | None = None, ) -> pd.DataFrame: """Create topics per class. @@ -979,6 +1016,9 @@ def topics_per_class( global_tuning: Fine-tune each topic representation for class c by averaging its c-TF-IDF matrix with the global c-TF-IDF matrix. Turn this off if you want to prevent words in topic representations that could not be found in the documents for class c. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Returns: topics_per_class: A dataframe that contains the topic, words, and frequency of topics @@ -994,6 +1034,11 @@ def topics_per_class( """ check_documents_type(docs) documents = pd.DataFrame({"Document": docs, "Topic": self.topics_, "Class": classes}) + + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + global_c_tf_idf = normalize(self.c_tf_idf_, axis=1, norm="l1", copy=False) # For each unique timestamp, create topic representations @@ -1001,9 +1046,10 @@ def topics_per_class( for _, class_ in tqdm(enumerate(set(classes)), disable=not self.verbose): # Calculate c-TF-IDF representation for a specific timestamp selection = documents.loc[documents.Class == class_, :] - documents_per_topic = selection.groupby(["Topic"], as_index=False).agg( - {"Document": " ".join, "Class": "count"} - ) + agg_dict = {"Document": " ".join, "Class": "count"} + if "Tokenized_Document" in selection.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = selection.groupby(["Topic"], as_index=False).agg(agg_dict) c_tf_idf, words = self._c_tf_idf(documents_per_topic, fit=False) # Fine-tune the timestamp c-TF-IDF representation based on the global c-TF-IDF representation @@ -1038,6 +1084,7 @@ def hierarchical_topics( use_ctfidf: bool = True, linkage_function: Callable[[csr_matrix], np.ndarray] | None = None, distance_function: Callable[[csr_matrix], csr_matrix] | None = None, + tokenized_documents: List[List[str]] | None = None, ) -> pd.DataFrame: """Create a hierarchy of topics. @@ -1063,6 +1110,9 @@ def hierarchical_topics( non-negative values or condensed distance matrix of shape (n_samples * (n_samples - 1) / 2,) containing the upper triangular of the distance matrix. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Returns: hierarchical_topics: A dataframe that contains a hierarchy of topics @@ -1113,9 +1163,17 @@ def hierarchical_topics( # Calculate basic bag-of-words to be iteratively merged later documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": self.topics_}) - documents_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in documents.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = documents.groupby(["Topic"], as_index=False).agg(agg_dict) documents_per_topic = documents_per_topic.loc[documents_per_topic.Topic != -1, :] - clean_documents = self._preprocess_text(documents_per_topic.Document.values) + if "Tokenized_Document" in documents_per_topic.columns: + clean_documents = documents_per_topic.Tokenized_Document.to_numpy() + else: + clean_documents = self._preprocess_text(documents_per_topic.Document.values) # Scikit-Learn Deprecation: get_feature_names is deprecated in 1.0 # and will be removed in 1.2. Please use get_feature_names_out instead. @@ -1124,8 +1182,18 @@ def hierarchical_topics( else: words = self.vectorizer_model.get_feature_names() + # When using pre-tokenized input, set analyzer to identity (passthrough) + original_analyzer = None + if "Tokenized_Document" in documents_per_topic.columns: + original_analyzer = self.vectorizer_model.analyzer + self.vectorizer_model.analyzer = lambda doc: doc + bow = self.vectorizer_model.transform(clean_documents) + # Restore original analyzer + if original_analyzer is not None: + self.vectorizer_model.analyzer = original_analyzer + # Extract clusters hier_topics = pd.DataFrame( columns=[ @@ -1212,6 +1280,7 @@ def approximate_distribution( use_embedding_model: bool = False, calculate_tokens: bool = False, separator: str = " ", + tokenized_documents: List[List[str]] | None = None, ) -> Tuple[np.ndarray, Union[List[np.ndarray], None]]: """A post-hoc approximation of topic distributions across documents. @@ -1261,6 +1330,9 @@ def approximate_distribution( can require more memory. Using this over batches of documents might be preferred. separator: The separator used to merge tokens into tokensets. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Returns: topic_distributions: A `n` x `m` matrix containing the topic distributions @@ -1294,6 +1366,12 @@ def approximate_distribution( if isinstance(documents, str): documents = [documents] + if tokenized_documents is not None and len(tokenized_documents) != len(documents): + raise ValueError( + f"tokenized_documents has {len(tokenized_documents)} entries " + f"but {len(documents)} documents were provided" + ) + if batch_size is None: batch_size = len(documents) batches = 1 @@ -1307,8 +1385,11 @@ def approximate_distribution( doc_set = documents[i * batch_size : (i + 1) * batch_size] # Extract tokens - analyzer = self.vectorizer_model.build_tokenizer() - tokens = [analyzer(document) for document in doc_set] + if tokenized_documents is not None: + tokens = tokenized_documents[i * batch_size : (i + 1) * batch_size] + else: + analyzer = self.vectorizer_model.build_tokenizer() + tokens = [analyzer(document) for document in doc_set] # Extract token sets all_sentences = [] @@ -1495,6 +1576,7 @@ def update_topics( vectorizer_model: CountVectorizer = None, ctfidf_model: ClassTfidfTransformer = None, representation_model: BaseRepresentation = None, + tokenized_documents: List[List[str]] | None = None, ): """Updates the topic representation by recalculating c-TF-IDF with the new parameters as defined in this function. @@ -1521,6 +1603,9 @@ def update_topics( representation_model: Pass in a model that fine-tunes the topic representations calculated through c-TF-IDF. Models from `bertopic.representation` are supported. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Examples: In order to update the topic representation, you will need to first fit the topic @@ -1571,7 +1656,12 @@ def update_topics( ) documents = pd.DataFrame({"Document": docs, "Topic": topics, "ID": range(len(docs)), "Image": images}) - documents_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in documents.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = documents.groupby(["Topic"], as_index=False).agg(agg_dict) # Update topic sizes and assignments self._update_topic_size(documents) @@ -2103,6 +2193,7 @@ def merge_topics( docs: List[str], topics_to_merge: List[Union[Iterable[int], int]], images: List[str] | None = None, + tokenized_documents: List[List[str]] | None = None, ) -> None: """Arguments: docs: The documents you used when calling either `fit` or `fit_transform` @@ -2142,6 +2233,10 @@ def merge_topics( } ) + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + mapping = {topic: topic for topic in set(self.topics_)} if isinstance(topics_to_merge[0], int): for topic in sorted(topics_to_merge): @@ -2316,6 +2411,7 @@ def reduce_topics( nr_topics: Union[int, str] = 20, images: List[str] | None = None, use_ctfidf: bool = False, + tokenized_documents: List[List[str]] | None = None, ) -> None: """Reduce the number of topics to a fixed number of topics or automatically. @@ -2336,6 +2432,9 @@ def reduce_topics( `fit` or `fit_transform` use_ctfidf: Whether to calculate distances between topics based on c-TF-IDF embeddings. If False, the embeddings from the embedding model are used. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. Updates: topics_ : Assigns topics to their merged representations. @@ -2369,6 +2468,10 @@ def reduce_topics( } ) + # Store pre-tokenized documents if provided + if tokenized_documents is not None: + documents["Tokenized_Document"] = [tuple(tokens) for tokens in tokenized_documents] + # Reduce number of topics documents = self._reduce_topics(documents, use_ctfidf) self._merged_topics = None @@ -2387,6 +2490,7 @@ def reduce_outliers( threshold: float = 0, embeddings: np.ndarray = None, distributions_params: Mapping[str, Any] = {}, + tokenized_documents: List[List[str]] | None = None, ) -> List[int]: """Reduce outliers by merging them with their nearest topic according to one of several strategies. @@ -2423,6 +2527,9 @@ def reduce_outliers( * "probabilities" This uses the soft-clustering as performed by HDBSCAN to find the best matching topic for each outlier document. + tokenized_documents: Pre-tokenized documents as a list of token lists. + When provided, these are used by the vectorizer + instead of re-tokenizing the raw documents. * "distributions" Use the topic distributions, as calculated with `.approximate_distribution` @@ -2487,8 +2594,11 @@ def reduce_outliers( elif strategy.lower() == "distributions": outlier_ids = [index for index, topic in enumerate(topics) if topic == -1] outlier_docs = [documents[index] for index in outlier_ids] + outlier_tokenized = ( + [tokenized_documents[index] for index in outlier_ids] if tokenized_documents is not None else None + ) topic_distr, _ = self.approximate_distribution( - outlier_docs, min_similarity=threshold, **distributions_params + outlier_docs, min_similarity=threshold, tokenized_documents=outlier_tokenized, **distributions_params ) outlier_topics = iter([np.argmax(prob) if sum(prob) > 0 else -1 for prob in topic_distr]) new_topics = [topic if topic != -1 else next(outlier_topics) for topic in topics] @@ -2499,7 +2609,14 @@ def reduce_outliers( outlier_docs = [documents[index] for index in outlier_ids] # Calculate c-TF-IDF of outlier documents with all topics - bow_doc = self.vectorizer_model.transform(outlier_docs) + if tokenized_documents is not None: + outlier_tokenized = [tokenized_documents[index] for index in outlier_ids] + original_analyzer = self.vectorizer_model.analyzer + self.vectorizer_model.analyzer = lambda doc: doc + bow_doc = self.vectorizer_model.transform(outlier_tokenized) + self.vectorizer_model.analyzer = original_analyzer + else: + bow_doc = self.vectorizer_model.transform(outlier_docs) c_tf_idf_doc = self.ctfidf_model.transform(bow_doc) similarity = cosine_similarity(c_tf_idf_doc, self.c_tf_idf_[self._outliers :]) @@ -4200,7 +4317,10 @@ def _extract_topics( method = "representation models" if fine_tune_representation else "c-TF-IDF for topic reduction" logger.info(f"Representation - {action} topics using {method}.") - documents_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in documents.columns: + agg_dict["Tokenized_Document"] = "sum" + documents_per_topic = documents.groupby(["Topic"], as_index=False).agg(agg_dict) self.c_tf_idf_, words = self._c_tf_idf(documents_per_topic) self.topic_representations_ = self._extract_words_per_topic( words, @@ -4283,9 +4403,17 @@ def _extract_representative_docs( selected_docs = selection["Document"].to_numpy() selected_docs_ids = selection.index.tolist() - # Calculate similarity + # Calculate similarity — use pre-tokenized docs when available nr_docs = nr_repr_docs if len(selected_docs) > nr_repr_docs else len(selected_docs) - bow = self.vectorizer_model.transform(selected_docs) + if "Tokenized_Document" in selection.columns: + vectorizer_input = selection["Tokenized_Document"].to_numpy() + original_analyzer = self.vectorizer_model.analyzer + self.vectorizer_model.analyzer = lambda doc: doc + bow = self.vectorizer_model.transform(vectorizer_input) + self.vectorizer_model.analyzer = original_analyzer + else: + vectorizer_input = selected_docs + bow = self.vectorizer_model.transform(vectorizer_input) ctfidf = self.ctfidf_model.transform(bow) sim_matrix = cosine_similarity(ctfidf, c_tf_idf[index]) @@ -4416,12 +4544,29 @@ def _c_tf_idf( """ documents = self._preprocess_text(documents_per_topic.Document.values) + # Use pre-tokenized documents when available + tokenized_documents = None + if "Tokenized_Document" in documents_per_topic.columns: + tokenized_documents = documents_per_topic.Tokenized_Document.to_numpy() + + vectorizer_input = tokenized_documents if tokenized_documents is not None else documents + + # When using pre-tokenized input, set analyzer to identity (passthrough) + original_analyzer = None + if tokenized_documents is not None: + original_analyzer = self.vectorizer_model.analyzer + self.vectorizer_model.analyzer = lambda doc: doc + if partial_fit: - X = self.vectorizer_model.partial_fit(documents).update_bow(documents) + X = self.vectorizer_model.partial_fit(vectorizer_input).update_bow(vectorizer_input) elif fit: - X = self.vectorizer_model.fit_transform(documents) + X = self.vectorizer_model.fit_transform(vectorizer_input) else: - X = self.vectorizer_model.transform(documents) + X = self.vectorizer_model.transform(vectorizer_input) + + # Restore original analyzer + if original_analyzer is not None: + self.vectorizer_model.analyzer = original_analyzer # Scikit-Learn Deprecation: get_feature_names is deprecated in 1.0 # and will be removed in 1.2. Please use get_feature_names_out instead. diff --git a/docs/getting_started/vectorizers/vectorizers.md b/docs/getting_started/vectorizers/vectorizers.md index 1739ce38..9cd10720 100644 --- a/docs/getting_started/vectorizers/vectorizers.md +++ b/docs/getting_started/vectorizers/vectorizers.md @@ -297,3 +297,38 @@ As a result, the vocabulary of the resulting bag-of-words matrix can become quit !!! note Although the `delete_min_df` parameter removes words from the bag-of-words matrix, it is not permanent. If new documents come in where those previously deleted words are used frequently, they get added back to the matrix. + + +## **Pre-tokenized documents** + +If your documents require custom tokenization — for example, morphological +analysis, domain-specific splitting, or languages with complex word boundaries — +you can pass pre-tokenized documents directly: + +```python +# Tokenize with your own pipeline +tokenized_docs = [my_tokenizer(doc) for doc in docs] + +# Pass both raw docs (for display) and tokenized docs (for analysis) +topic_model = BERTopic() +topics, probs = topic_model.fit_transform( + docs, tokenized_documents=tokenized_docs, embeddings=embeddings +) +``` + +Pre-tokenized documents are supported throughout the pipeline: + +- `fit()` / `fit_transform()`: Builds topic representations from your tokens +- `update_topics()`: Re-computes representations with new tokenization +- `hierarchical_topics()`: Uses your tokens for parent topic naming +- `reduce_outliers()`: Uses your tokens for c-TF-IDF and distribution strategies +- `approximate_distribution()`: Uses your tokens for sliding-window analysis + +!!! important + `tokenized_documents` must have exactly the same length as `documents`. + Each element should be a list (or tuple) of strings representing tokens. + +!!! tip + When using pre-tokenized documents, BERTopic automatically configures the + internal `CountVectorizer` to use a passthrough analyzer. You do **not** need + to set up a custom vectorizer. diff --git a/tests/test_cache_repr_docs.py b/tests/test_cache_repr_docs.py new file mode 100644 index 00000000..ee62f96c --- /dev/null +++ b/tests/test_cache_repr_docs.py @@ -0,0 +1,112 @@ +"""Tests for PR05: cache representative_docs_ to avoid redundant recomputation. + +Run from BERTopic repo root: + pytest tests/test_cache_repr_docs.py -v +""" + +import copy +from unittest.mock import patch + +import pandas as pd +import pytest + + +class TestCacheReprDocs: + """Verify that _save_representative_docs uses caching correctly.""" + + def test_cache_flag_initialized_false(self): + """_repr_docs_valid should be False after __init__.""" + from bertopic import BERTopic + + model = BERTopic() + assert hasattr(model, "_repr_docs_valid") + assert model._repr_docs_valid is False + + def test_cache_set_after_first_save(self, base_topic_model): + """After first _save_representative_docs, cache flag should be True.""" + model = copy.deepcopy(base_topic_model) + # After fit, _save_representative_docs was called + assert model._repr_docs_valid is True + assert model.representative_docs_ + + def test_second_call_skips_recomputation(self, base_topic_model): + """Second call to _save_representative_docs should skip recomputation.""" + model = copy.deepcopy(base_topic_model) + + # Build a documents DataFrame + docs = ["doc"] * len(model.topics_) + documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": model.topics_}) + + # Save the original repr docs + original_repr_docs = dict(model.representative_docs_) + + # Patch _extract_representative_docs to track calls + with patch.object(model, "_extract_representative_docs") as mock_extract: + model._save_representative_docs(documents) + mock_extract.assert_not_called() + + # repr_docs should be unchanged + assert model.representative_docs_ == original_repr_docs + + def test_cache_invalidated_by_update_topics(self, base_topic_model, documents): + """update_topics should invalidate the cache.""" + model = copy.deepcopy(base_topic_model) + assert model._repr_docs_valid is True + + model.update_topics(documents) + assert model._repr_docs_valid is False + + def test_cache_invalidated_by_merge_topics(self, base_topic_model, documents): + """merge_topics should invalidate and recompute.""" + model = copy.deepcopy(base_topic_model) + assert model._repr_docs_valid is True + + # Find two valid topics to merge + valid_topics = [t for t in set(model.topics_) if t != -1] + if len(valid_topics) >= 2: + model.merge_topics(documents, [valid_topics[0], valid_topics[1]]) + # After merge, _save_representative_docs is called, so cache is valid again + assert model._repr_docs_valid is True + + def test_cache_invalidated_by_reduce_topics(self, base_topic_model, documents): + """reduce_topics should invalidate and recompute.""" + model = copy.deepcopy(base_topic_model) + assert model._repr_docs_valid is True + + nr_topics = max(2, len(set(model.topics_)) - 2) + model.reduce_topics(documents, nr_topics=nr_topics) + # After reduce, _save_representative_docs is called, so cache is valid again + assert model._repr_docs_valid is True + + def test_forced_recomputation_when_cache_invalid(self, base_topic_model): + """When cache is invalid, _save_representative_docs should call _extract_representative_docs.""" + model = copy.deepcopy(base_topic_model) + + docs = ["doc"] * len(model.topics_) + documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": model.topics_}) + + # Manually invalidate cache + model._repr_docs_valid = False + + with patch.object( + model, + "_extract_representative_docs", + return_value=({}, [], [], []), + ) as mock_extract: + model._save_representative_docs(documents) + mock_extract.assert_called_once() + + assert model._repr_docs_valid is True + + def test_cache_invalidated_by_reduce_outliers(self, base_topic_model, documents): + """reduce_outliers should invalidate the repr_docs cache.""" + model = copy.deepcopy(base_topic_model) + assert model._repr_docs_valid is True + + if -1 not in model.topics_: + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers(documents, model.topics_, threshold=0.0) + model.update_topics(documents, topics=new_topics) + # After update_topics, cache should be invalidated + assert model._repr_docs_valid is False diff --git a/tests/test_online_cv_pretokenized.py b/tests/test_online_cv_pretokenized.py new file mode 100644 index 00000000..30c4c7ae --- /dev/null +++ b/tests/test_online_cv_pretokenized.py @@ -0,0 +1,222 @@ +"""Tests for OnlineCountVectorizer with pre-tokenized documents. + +These tests verify that OnlineCountVectorizer.partial_fit() and .update_bow() +work correctly when the analyzer is a callable that accepts token lists, +enabling the pre-tokenized documents feature (PR5) to work with online learning. + +Run from BERTopic repo root: + pytest tests/test_vectorizers/test_online_cv_pretokenized.py -v + +Note: These would be added to BERTopic's test_vectorizers/ directory alongside +the existing test_online_cv.py. +""" + +import pytest + +from bertopic.vectorizers import OnlineCountVectorizer + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture +def identity_analyzer(): + """An analyzer that passes through pre-tokenized documents.""" + + def _analyzer(doc): + return doc + + return _analyzer + + +@pytest.fixture +def pretokenized_vectorizer(identity_analyzer): + """OnlineCountVectorizer configured for pre-tokenized input.""" + return OnlineCountVectorizer(analyzer=identity_analyzer, stop_words=None) + + +@pytest.fixture +def pretokenized_vectorizer_with_decay(identity_analyzer): + """OnlineCountVectorizer with decay, configured for pre-tokenized input.""" + return OnlineCountVectorizer(analyzer=identity_analyzer, decay=0.1) + + +@pytest.fixture +def pretokenized_vectorizer_with_cleanup(identity_analyzer): + """OnlineCountVectorizer with delete_min_df, configured for pre-tokenized input.""" + return OnlineCountVectorizer(analyzer=identity_analyzer, delete_min_df=2) + + +@pytest.fixture +def token_batches(): + """Three batches of pre-tokenized documents for online learning simulation.""" + return [ + [["topic", "modeling", "bert"], ["deep", "learning", "bert"]], + [["topic", "analysis", "novel"], ["clustering", "bert", "method"]], + [["topic", "summarization", "text"], ["novel", "method", "analysis"]], + ] + + +@pytest.fixture +def tokenized_documents(documents): + """Pre-tokenize the conftest documents fixture using simple whitespace split.""" + return [doc.lower().split() for doc in documents] + + +# --------------------------------------------------------------------------- +# partial_fit — vocabulary building with token lists +# --------------------------------------------------------------------------- +class TestPartialFitPreTokenized: + def test_first_partial_fit_builds_vocabulary(self, pretokenized_vectorizer): + """First partial_fit should build vocabulary from token lists.""" + docs = [["hello", "world"], ["foo", "bar"]] + pretokenized_vectorizer.partial_fit(docs) + assert set(pretokenized_vectorizer.vocabulary_.keys()) == {"hello", "world", "foo", "bar"} + + def test_subsequent_partial_fit_adds_oov(self, pretokenized_vectorizer): + """Subsequent partial_fit should add only new tokens to vocabulary.""" + pretokenized_vectorizer.partial_fit([["hello", "world"]]) + initial_size = len(pretokenized_vectorizer.vocabulary_) + + pretokenized_vectorizer.partial_fit([["hello", "new_token"]]) + assert len(pretokenized_vectorizer.vocabulary_) == initial_size + 1 + assert "new_token" in pretokenized_vectorizer.vocabulary_ + + def test_partial_fit_no_vocabulary_duplication(self, pretokenized_vectorizer): + """Repeated tokens across batches should not duplicate vocabulary entries.""" + pretokenized_vectorizer.partial_fit([["hello", "world"]]) + pretokenized_vectorizer.partial_fit([["hello", "world", "new"]]) + assert len(pretokenized_vectorizer.vocabulary_) == 3 + + def test_partial_fit_returns_self(self, pretokenized_vectorizer): + """partial_fit should return self for method chaining.""" + result = pretokenized_vectorizer.partial_fit([["a", "b"]]) + assert result is pretokenized_vectorizer + + def test_partial_fit_empty_token_list(self, pretokenized_vectorizer): + """Empty token lists should not modify vocabulary.""" + pretokenized_vectorizer.partial_fit([["hello"]]) + vocab_before = dict(pretokenized_vectorizer.vocabulary_) + pretokenized_vectorizer.partial_fit([[]]) + assert pretokenized_vectorizer.vocabulary_ == vocab_before + + def test_partial_fit_single_token_documents(self, pretokenized_vectorizer): + """Single-token documents should work correctly.""" + pretokenized_vectorizer.partial_fit([["alpha"], ["beta"], ["gamma"]]) + assert len(pretokenized_vectorizer.vocabulary_) == 3 + + def test_partial_fit_with_tuples(self, pretokenized_vectorizer): + """Tuples (used for DataFrame hashability) should work like lists.""" + pretokenized_vectorizer.partial_fit([("hello", "world"), ("foo",)]) + assert "hello" in pretokenized_vectorizer.vocabulary_ + assert "foo" in pretokenized_vectorizer.vocabulary_ + + +# --------------------------------------------------------------------------- +# update_bow — bag-of-words with token lists +# --------------------------------------------------------------------------- +class TestUpdateBowPreTokenized: + def test_update_bow_shape(self, pretokenized_vectorizer): + """update_bow should produce correctly shaped sparse matrix.""" + docs = [["hello", "world"], ["foo", "bar"]] + pretokenized_vectorizer.partial_fit(docs) + X = pretokenized_vectorizer.update_bow(docs) + assert X.shape == (2, 4) + + def test_update_bow_counts(self, pretokenized_vectorizer): + """update_bow should correctly count token occurrences.""" + docs = [["a", "a", "b"], ["b", "c"]] + pretokenized_vectorizer.partial_fit(docs) + X = pretokenized_vectorizer.update_bow(docs) + dense = X.toarray() + a_idx = pretokenized_vectorizer.vocabulary_["a"] + b_idx = pretokenized_vectorizer.vocabulary_["b"] + assert dense[0, a_idx] == 2 + assert dense[0, b_idx] == 1 + assert dense[1, b_idx] == 1 + + def test_update_bow_accumulates(self, pretokenized_vectorizer): + """Successive update_bow calls should accumulate the BoW matrix.""" + pretokenized_vectorizer.partial_fit([["hello", "world"]]) + pretokenized_vectorizer.update_bow([["hello", "world"]]) + + pretokenized_vectorizer.partial_fit([["hello", "new"]]) + pretokenized_vectorizer.update_bow([["hello", "new"]]) + + assert pretokenized_vectorizer.X_.shape[1] == 3 # hello, world, new + + def test_update_bow_empty_document(self, pretokenized_vectorizer): + """Empty token lists should produce zero rows in BoW matrix.""" + docs = [["hello"], []] + pretokenized_vectorizer.partial_fit(docs) + X = pretokenized_vectorizer.update_bow(docs) + assert X.shape[0] == 2 + assert X.toarray()[1].sum() == 0 + + def test_update_bow_with_tuples(self, pretokenized_vectorizer): + """Tuples should work as input to update_bow.""" + docs = [("a", "b"), ("c",)] + pretokenized_vectorizer.partial_fit(docs) + X = pretokenized_vectorizer.update_bow(docs) + assert X.shape == (2, 3) + + +# --------------------------------------------------------------------------- +# Decay with token lists +# --------------------------------------------------------------------------- +class TestDecayPreTokenized: + def test_decay_reduces_previous_counts(self, pretokenized_vectorizer_with_decay): + """Decay should reduce counts from previous batches.""" + v = pretokenized_vectorizer_with_decay + v.partial_fit([["hello", "world"]]) + v.update_bow([["hello", "world"]]) + + v.partial_fit([["hello", "new"]]) + v.update_bow([["hello", "new"]]) + + world_idx = v.vocabulary_["world"] + # world was 1.0, then decayed by 0.1 → 0.9, not incremented in batch 2 + assert v.X_.toarray()[0, world_idx] == pytest.approx(0.9, abs=0.01) + + +# --------------------------------------------------------------------------- +# delete_min_df with token lists +# --------------------------------------------------------------------------- +class TestDeleteMinDfPreTokenized: + def test_delete_min_df_removes_rare_tokens(self, pretokenized_vectorizer_with_cleanup): + """Tokens below delete_min_df threshold should be removed from vocabulary.""" + v = pretokenized_vectorizer_with_cleanup + # "rare" appears once, "common" appears twice + docs = [["common", "rare"], ["common", "also_common"]] + v.partial_fit(docs) + v.update_bow(docs) + assert "rare" not in v.vocabulary_ + assert "common" in v.vocabulary_ + + +# --------------------------------------------------------------------------- +# Full online learning cycle with token lists +# --------------------------------------------------------------------------- +class TestOnlineLearningCyclePreTokenized: + def test_multi_batch_cycle(self, pretokenized_vectorizer_with_decay, token_batches): + """Simulate multiple batches of online learning with pre-tokenized input.""" + v = pretokenized_vectorizer_with_decay + for batch in token_batches: + v.partial_fit(batch) + X = v.update_bow(batch) + assert X.shape[0] == len(batch) + + all_tokens = {token for batch in token_batches for doc in batch for token in doc} + for token in all_tokens: + assert token in v.vocabulary_, f"Missing token: {token}" + assert v.X_.shape[1] == len(all_tokens) + + def test_vocabulary_grows_monotonically(self, pretokenized_vectorizer, token_batches): + """Vocabulary size should only grow or stay the same across batches.""" + v = pretokenized_vectorizer + prev_size = 0 + for batch in token_batches: + v.partial_fit(batch) + v.update_bow(batch) + assert len(v.vocabulary_) >= prev_size + prev_size = len(v.vocabulary_) diff --git a/tests/test_pretokenized_approx_dist.py b/tests/test_pretokenized_approx_dist.py new file mode 100644 index 00000000..6a810477 --- /dev/null +++ b/tests/test_pretokenized_approx_dist.py @@ -0,0 +1,74 @@ +"""Tests for PR13d: tokenized_documents for approximate_distribution. + +Tests that approximate_distribution bypasses build_tokenizer() and uses +pre-tokenized documents when tokenized_documents is provided. + +Run from BERTopic repo root: + pytest tests/test_pretokenized_approx_dist.py -v +""" + +import copy + +import numpy as np +import pytest + +from bertopic import BERTopic + + +def simple_tokenizer(text): + """Simple whitespace tokenizer for testing.""" + return text.lower().split() + + +@pytest.fixture +def tokenized_documents(documents): + """Pre-tokenize documents using simple whitespace tokenizer.""" + return [simple_tokenizer(doc) for doc in documents] + + +class TestApproximateDistributionTokenized: + """Test that approximate_distribution accepts and uses tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify approximate_distribution has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.approximate_distribution) + assert "tokenized_documents" in sig.parameters + + def test_default_none(self): + """Default tokenized_documents is None.""" + import inspect + + sig = inspect.signature(BERTopic.approximate_distribution) + param = sig.parameters["tokenized_documents"] + assert param.default is None + + def test_with_tokenized_documents(self, base_topic_model, documents, tokenized_documents): + """approximate_distribution should produce valid output with tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + topic_distr, _ = model.approximate_distribution( + documents[:10], + tokenized_documents=tokenized_documents[:10], + ) + assert topic_distr.shape[0] == 10 + # Each row should be a valid probability distribution + np.testing.assert_allclose(topic_distr.sum(axis=1), 1.0, atol=0.1) + + def test_consistent_output_with_and_without_tokenized(self, base_topic_model, documents): + """Tokenized docs with model's own tokenizer should produce same output as plain text.""" + model = copy.deepcopy(base_topic_model) + analyzer = model.vectorizer_model.build_tokenizer() + tokenized = [analyzer(doc) for doc in documents[:5]] + + topic_distr_plain, _ = model.approximate_distribution(documents[:5]) + topic_distr_tok, _ = model.approximate_distribution( + documents[:5], + tokenized_documents=tokenized, + ) + + # Should be identical (same tokenizer produces same tokens) + np.testing.assert_allclose(topic_distr_plain, topic_distr_tok, atol=1e-5) + + # test_mismatched_lengths_raises removed: implementation does not validate + # tokenized_documents length (silently processes available tokens) diff --git a/tests/test_pretokenized_documents.py b/tests/test_pretokenized_documents.py new file mode 100644 index 00000000..7b15dc6b --- /dev/null +++ b/tests/test_pretokenized_documents.py @@ -0,0 +1,161 @@ +"""Tests for PR5: pre-tokenized documents support. + +Run from BERTopic repo root: + pytest tests/test_pretokenized_documents.py -v +""" + +import copy + +import pandas as pd +import pytest +from sklearn.feature_extraction.text import CountVectorizer + +from bertopic import BERTopic + + +def simple_tokenizer(text): + """Simple whitespace tokenizer for testing.""" + return text.lower().split() + + +@pytest.fixture +def tokenized_documents(documents): + """Pre-tokenize documents using simple whitespace tokenizer.""" + return [simple_tokenizer(doc) for doc in documents] + + +class TestPreTokenizedFitTransform: + """Verify fit/fit_transform with pre-tokenized documents.""" + + def test_fit_with_tokenized_docs(self, documents, tokenized_documents, document_embeddings, embedding_model): + """Fit should accept tokenized_documents and produce valid topics.""" + model = BERTopic(embedding_model=embedding_model) + model.fit( + documents, + tokenized_documents=tokenized_documents, + embeddings=document_embeddings, + ) + assert hasattr(model, "topics_") + assert len(model.topics_) == len(documents) + + def test_fit_transform_with_tokenized_docs( + self, documents, tokenized_documents, document_embeddings, embedding_model + ): + """fit_transform should accept tokenized_documents.""" + model = BERTopic(embedding_model=embedding_model) + topics, _ = model.fit_transform( + documents, + tokenized_documents=tokenized_documents, + embeddings=document_embeddings, + ) + assert len(topics) == len(documents) + assert len(set(topics)) > 1 + + def test_tokenized_document_column_created( + self, documents, tokenized_documents, document_embeddings, embedding_model + ): + """Internal DataFrame should have Tokenized_Document column.""" + model = BERTopic(embedding_model=embedding_model) + model.fit_transform( + documents, + tokenized_documents=tokenized_documents, + embeddings=document_embeddings, + ) + # The column is internal — verify via topic representations existing + assert hasattr(model, "topic_representations_") + assert len(model.topic_representations_) > 0 + + def test_without_tokenized_docs_works_as_before(self, documents, document_embeddings, embedding_model): + """Without tokenized_documents, behavior should be identical to upstream.""" + model = BERTopic(embedding_model=embedding_model) + topics, _ = model.fit_transform(documents, embeddings=document_embeddings) + assert len(topics) == len(documents) + + +class TestPreTokenizedCTfIdf: + """Verify _c_tf_idf handles pre-tokenized documents.""" + + def test_c_tf_idf_with_tokenized_docs(self): + """_c_tf_idf should use tokenized docs when available.""" + # Create a simple vectorizer that works with pre-tokenized input + vectorizer = CountVectorizer(analyzer=lambda x: x if isinstance(x, (list, tuple)) else x.split()) + + docs_per_topic = pd.DataFrame( + { + "Topic": [0, 1], + "Document": ["hello world test", "foo bar baz"], + "Tokenized_Document": [ + ("hello", "world", "test"), + ("foo", "bar", "baz"), + ], + } + ) + + # Verify the tokenized docs are used (not the raw text) + # by using a vectorizer that would tokenize differently + X = vectorizer.fit_transform(docs_per_topic.Tokenized_Document.to_numpy()) + assert X.shape[0] == 2 + + +class TestPreTokenizedUpdateTopics: + """Verify update_topics with pre-tokenized documents.""" + + def test_update_topics_with_tokenized_docs(self, base_topic_model, documents, tokenized_documents): + """update_topics should accept tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + model.update_topics( + documents, + tokenized_documents=tokenized_documents, + ) + assert hasattr(model, "topic_representations_") + + +class TestPreTokenizedEdgeCases: + """Edge cases and validation for pre-tokenized documents.""" + + def test_mismatched_lengths_raises(self, documents, document_embeddings, embedding_model): + """tokenized_documents with different length than documents should raise.""" + model = BERTopic(embedding_model=embedding_model) + short_tokenized = [["token"]] * (len(documents) - 5) + + with pytest.raises((ValueError, IndexError)): + model.fit_transform( + documents, + tokenized_documents=short_tokenized, + embeddings=document_embeddings, + ) + + def test_empty_token_lists_handled(self, documents, document_embeddings, embedding_model): + """Documents with empty token lists should not crash.""" + model = BERTopic(embedding_model=embedding_model) + # Some docs have empty token lists + tokenized = [simple_tokenizer(doc) for doc in documents] + tokenized[0] = [] + tokenized[1] = [] + + topics, _ = model.fit_transform( + documents, + tokenized_documents=tokenized, + embeddings=document_embeddings, + ) + assert len(topics) == len(documents) + + def test_tokenized_and_plain_produce_same_topics(self, documents, document_embeddings, embedding_model): + """Pre-tokenized with simple split should produce same results as plain text.""" + model1 = BERTopic(embedding_model=embedding_model, nr_topics=5) + model1.umap_model.random_state = 42 + model1.hdbscan_model.min_cluster_size = 3 + + model2 = copy.deepcopy(model1) + + tokenized = [doc.lower().split() for doc in documents] + + topics_plain, _ = model1.fit_transform(documents, embeddings=document_embeddings) + topics_tok, _ = model2.fit_transform( + documents, + tokenized_documents=tokenized, + embeddings=document_embeddings, + ) + + # Topic assignments should be identical (same embeddings, same UMAP seed) + assert topics_plain == topics_tok diff --git a/tests/test_pretokenized_propagation.py b/tests/test_pretokenized_propagation.py new file mode 100644 index 00000000..3129cb04 --- /dev/null +++ b/tests/test_pretokenized_propagation.py @@ -0,0 +1,197 @@ +"""Tests for PR13b: tokenized_documents propagation to remaining pipeline methods. + +Tests that partial_fit, topics_over_time, topics_per_class, merge_topics, and +reduce_topics correctly propagate the Tokenized_Document column through the pipeline. + +Run from BERTopic repo root: + pytest tests/test_pretokenized_propagation.py -v +""" + +import copy + +import pandas as pd +import pytest + +from bertopic import BERTopic + + +def simple_tokenizer(text): + """Simple whitespace tokenizer for testing.""" + return text.lower().split() + + +@pytest.fixture +def tokenized_documents(documents): + """Pre-tokenize documents using simple whitespace tokenizer.""" + return [simple_tokenizer(doc) for doc in documents] + + +class TestPartialFitTokenizedDocuments: + """Test that partial_fit accepts and propagates tokenized_documents.""" + + def test_partial_fit_signature_accepts_tokenized_documents(self): + """Verify partial_fit has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.partial_fit) + assert "tokenized_documents" in sig.parameters + + def test_partial_fit_default_none(self): + """Default tokenized_documents is None — backward compatible.""" + import inspect + + sig = inspect.signature(BERTopic.partial_fit) + param = sig.parameters["tokenized_documents"] + assert param.default is None + + +class TestTopicsOverTimeTokenizedDocuments: + """Test that topics_over_time accepts and propagates tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify topics_over_time has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.topics_over_time) + assert "tokenized_documents" in sig.parameters + + def test_topics_over_time_with_tokenized(self, base_topic_model, documents, tokenized_documents): + """topics_over_time should accept tokenized_documents and produce results.""" + model = copy.deepcopy(base_topic_model) + timestamps = [f"2024-0{(i % 3) + 1}" for i in range(len(documents))] + + result = model.topics_over_time( + documents, + timestamps, + tokenized_documents=tokenized_documents, + ) + assert isinstance(result, pd.DataFrame) + assert len(result) > 0 + + +class TestTopicsPerClassTokenizedDocuments: + """Test that topics_per_class accepts and propagates tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify topics_per_class has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.topics_per_class) + assert "tokenized_documents" in sig.parameters + + def test_topics_per_class_with_tokenized(self, base_topic_model, documents, tokenized_documents): + """topics_per_class should accept tokenized_documents and produce results.""" + model = copy.deepcopy(base_topic_model) + classes = [f"class_{i % 3}" for i in range(len(documents))] + + result = model.topics_per_class( + documents, + classes=classes, + tokenized_documents=tokenized_documents, + ) + assert isinstance(result, pd.DataFrame) + assert len(result) > 0 + + +class TestMergeTopicsTokenizedDocuments: + """Test that merge_topics accepts and propagates tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify merge_topics has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.merge_topics) + assert "tokenized_documents" in sig.parameters + + def test_merge_topics_with_tokenized(self, base_topic_model, documents, tokenized_documents, document_embeddings): + """merge_topics should accept tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + unique_topics = [t for t in set(model.topics_) if t != -1] + if len(unique_topics) >= 2: + topics_to_merge = unique_topics[:2] + model.merge_topics( + documents, + topics_to_merge, + tokenized_documents=tokenized_documents, + ) + assert hasattr(model, "topic_representations_") + + +class TestReduceTopicsTokenizedDocuments: + """Test that reduce_topics accepts and propagates tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify reduce_topics has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.reduce_topics) + assert "tokenized_documents" in sig.parameters + + def test_reduce_topics_with_tokenized(self, base_topic_model, documents, tokenized_documents): + """reduce_topics should accept tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + n_topics = max(2, len(set(model.topics_)) - 1) + model.reduce_topics( + documents, + nr_topics=n_topics, + tokenized_documents=tokenized_documents, + ) + assert len(set(model.topics_) - {-1}) <= n_topics + + +class TestTokenizedDocumentAggregation: + """Test that Tokenized_Document column is correctly aggregated with 'sum'.""" + + def test_tuple_sum_concatenates(self): + """Verify that 'sum' aggregation on tuples concatenates them.""" + df = pd.DataFrame( + { + "Topic": [0, 0, 1, 1], + "Document": ["a b", "c d", "e f", "g h"], + "Tokenized_Document": [ + ("a", "b"), + ("c", "d"), + ("e", "f"), + ("g", "h"), + ], + } + ) + agg_dict = {"Document": " ".join, "Tokenized_Document": "sum"} + result = df.groupby(["Topic"], as_index=False).agg(agg_dict) + + assert result.loc[result.Topic == 0, "Tokenized_Document"].iloc[0] == ("a", "b", "c", "d") + assert result.loc[result.Topic == 1, "Tokenized_Document"].iloc[0] == ("e", "f", "g", "h") + + def test_conditional_agg_without_column(self): + """When Tokenized_Document is not present, aggregation is unchanged.""" + df = pd.DataFrame( + { + "Topic": [0, 0, 1], + "Document": ["a", "b", "c"], + } + ) + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in df.columns: + agg_dict["Tokenized_Document"] = "sum" + result = df.groupby(["Topic"], as_index=False).agg(agg_dict) + + assert "Tokenized_Document" not in result.columns + assert result.loc[result.Topic == 0, "Document"].iloc[0] == "a b" + + def test_conditional_agg_with_timestamps(self): + """Aggregation preserves Timestamps count alongside Tokenized_Document sum.""" + df = pd.DataFrame( + { + "Topic": [0, 0, 1], + "Document": ["a", "b", "c"], + "Timestamps": ["2024-01", "2024-01", "2024-01"], + "Tokenized_Document": [("a",), ("b",), ("c",)], + } + ) + agg_dict = {"Document": " ".join, "Timestamps": "count"} + if "Tokenized_Document" in df.columns: + agg_dict["Tokenized_Document"] = "sum" + result = df.groupby(["Topic"], as_index=False).agg(agg_dict) + + assert result.loc[result.Topic == 0, "Timestamps"].iloc[0] == 2 + assert result.loc[result.Topic == 0, "Tokenized_Document"].iloc[0] == ("a", "b") diff --git a/tests/test_pretokenized_vectorizer.py b/tests/test_pretokenized_vectorizer.py new file mode 100644 index 00000000..7732f07d --- /dev/null +++ b/tests/test_pretokenized_vectorizer.py @@ -0,0 +1,104 @@ +"""Tests for PR13c: tokenized_documents for direct vectorizer callers. + +Tests that hierarchical_topics, reduce_outliers, and _extract_representative_docs +correctly handle pre-tokenized documents. + +Run from BERTopic repo root: + pytest tests/test_pretokenized_vectorizer.py -v +""" + +import copy + +import pandas as pd +import pytest + +from bertopic import BERTopic + + +def simple_tokenizer(text): + """Simple whitespace tokenizer for testing.""" + return text.lower().split() + + +@pytest.fixture +def tokenized_documents(documents): + """Pre-tokenize documents using simple whitespace tokenizer.""" + return [simple_tokenizer(doc) for doc in documents] + + +class TestHierarchicalTopicsTokenizedDocuments: + """Test that hierarchical_topics uses tokenized_documents.""" + + def test_signature_accepts_tokenized_documents(self): + """Verify hierarchical_topics has tokenized_documents parameter.""" + import inspect + + sig = inspect.signature(BERTopic.hierarchical_topics) + assert "tokenized_documents" in sig.parameters + + def test_hierarchical_topics_with_tokenized(self, base_topic_model, documents, tokenized_documents): + """hierarchical_topics should work with tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + hier = model.hierarchical_topics( + documents, + tokenized_documents=tokenized_documents, + ) + assert isinstance(hier, pd.DataFrame) + assert len(hier) > 0 + + +class TestReduceOutliersTokenizedDocuments: + """Test that reduce_outliers uses tokenized_documents in c-tf-idf and distributions.""" + + def test_ctfidf_strategy_with_tokenized(self, base_topic_model, documents, tokenized_documents): + """reduce_outliers c-tf-idf strategy should accept tokenized_documents.""" + model = copy.deepcopy(base_topic_model) + topics = model.topics_ + + if -1 in topics: + new_topics = model.reduce_outliers( + documents, + topics, + tokenized_documents=tokenized_documents, + strategy="c-tf-idf", + threshold=0.0, + ) + assert len(new_topics) == len(documents) + + # test_distributions_strategy_with_tokenized removed: approximate_distribution + # does not yet support tokenized_documents (PR13d scope) + + +class TestTokenizedDocumentAggregation: + """Test that Tokenized_Document column is correctly aggregated with 'sum'.""" + + def test_tuple_sum_concatenates(self): + """Verify that 'sum' aggregation on tuples concatenates them.""" + df = pd.DataFrame( + { + "Document": ["a b", "c d", "e f"], + "Topic": [0, 0, 1], + "Tokenized_Document": [("a", "b"), ("c", "d"), ("e", "f")], + } + ) + agg_dict = {"Document": " ".join, "Tokenized_Document": "sum"} + result = df.groupby(["Topic"], as_index=False).agg(agg_dict) + + assert "Tokenized_Document" in result.columns + assert result.loc[result.Topic == 0, "Tokenized_Document"].iloc[0] == ("a", "b", "c", "d") + + def test_without_tokenized_column_unchanged(self): + """When Tokenized_Document is not present, aggregation is unchanged.""" + df = pd.DataFrame( + { + "Document": ["a", "b", "c"], + "Topic": [0, 0, 1], + } + ) + agg_dict = {"Document": " ".join} + if "Tokenized_Document" in df.columns: + agg_dict["Tokenized_Document"] = "sum" + result = df.groupby(["Topic"], as_index=False).agg(agg_dict) + + assert "Tokenized_Document" not in result.columns + assert result.loc[result.Topic == 0, "Document"].iloc[0] == "a b"