diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..3039a968 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -765,9 +765,9 @@ def partial_fit( missing_topics = {} # 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}) + documents_sorted = documents.sort_values("Topic") + updated_topics = documents_sorted.groupby(["Topic"], as_index=False).first().Topic.astype(int) + documents_per_topic = self._aggregate_documents(documents_sorted) # Update topic representations self.c_tf_idf_, updated_words = self._c_tf_idf(documents_per_topic, partial_fit=True) @@ -1113,16 +1113,11 @@ 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}) + documents_per_topic = self._aggregate_documents(documents) documents_per_topic = documents_per_topic.loc[documents_per_topic.Topic != -1, :] 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. - if version.parse(sklearn_version) >= version.parse("1.0.0"): - words = self.vectorizer_model.get_feature_names_out() - else: - words = self.vectorizer_model.get_feature_names() + words = self._get_feature_names() bow = self.vectorizer_model.transform(clean_documents) @@ -1162,14 +1157,14 @@ def hierarchical_topics( # Extract parent's name and ID parent_id = index + len(clusters) - parent_name = "_".join([x[0] for x in words_per_topic[0]][:5]) + parent_name = self._topic_name_from_words(words_per_topic[0]) # Extract child's name and ID Z_id = Z[index][0] child_left_id = Z_id if Z_id - nr_clusters < 0 else Z_id - nr_clusters if Z_id - nr_clusters < 0: - child_left_name = "_".join([x[0] for x in self.get_topic(Z_id)][:5]) + child_left_name = self._topic_name_from_words(self.get_topic(Z_id)) else: child_left_name = hier_topics.iloc[int(child_left_id)].Parent_Name @@ -1178,7 +1173,7 @@ def hierarchical_topics( child_right_id = Z_id if Z_id - nr_clusters < 0 else Z_id - nr_clusters if Z_id - nr_clusters < 0: - child_right_name = "_".join([x[0] for x in self.get_topic(Z_id)][:5]) + child_right_name = self._topic_name_from_words(self.get_topic(Z_id)) else: child_right_name = hier_topics.iloc[int(child_right_id)].Parent_Name @@ -1571,7 +1566,7 @@ 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}) + documents_per_topic = self._aggregate_documents(documents) # Update topic sizes and assignments self._update_topic_size(documents) @@ -4200,7 +4195,7 @@ 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}) + documents_per_topic = self._aggregate_documents(documents) self.c_tf_idf_, words = self._c_tf_idf(documents_per_topic) self.topic_representations_ = self._extract_words_per_topic( words, @@ -4423,12 +4418,7 @@ def _c_tf_idf( else: X = self.vectorizer_model.transform(documents) - # 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. - if version.parse(sklearn_version) >= version.parse("1.0.0"): - words = self.vectorizer_model.get_feature_names_out() - else: - words = self.vectorizer_model.get_feature_names() + words = self._get_feature_names() multiplier = None if self.ctfidf_model.seed_words and self.seed_topic_list: @@ -4461,6 +4451,44 @@ def _update_topic_size(self, documents: pd.DataFrame): self.topic_sizes_ = collections.Counter(documents.Topic.to_numpy().tolist()) self.topics_ = documents.Topic.astype(int).tolist() + @staticmethod + def _aggregate_documents(documents: pd.DataFrame) -> pd.DataFrame: + """Group documents by topic, joining text. + + Arguments: + documents: DataFrame with at least ``Document`` and ``Topic`` columns. + + Returns: + DataFrame with one row per topic, text joined per topic. + """ + return documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + + def _get_feature_names(self) -> List[str]: + """Return feature names from the fitted vectorizer (sklearn-version safe). + + Returns: + List of feature name strings from the fitted vectorizer. + """ + if version.parse(sklearn_version) >= version.parse("1.0.0"): + return list(self.vectorizer_model.get_feature_names_out()) + return list(self.vectorizer_model.get_feature_names()) + + @staticmethod + def _topic_name_from_words( + words_per_topic: List[Tuple[str, float]], + n: int = 5, + ) -> str: + """Build a topic name by joining the top *n* words with underscores. + + Arguments: + words_per_topic: List of (word, score) tuples for a single topic. + n: Number of top words to include in the name. + + Returns: + A string like ``"word1_word2_word3_word4_word5"``. + """ + return "_".join([w for w, _ in words_per_topic[:n]]) + def _extract_words_per_topic( self, words: List[str], diff --git a/tests/test_shared_helpers.py b/tests/test_shared_helpers.py new file mode 100644 index 00000000..c381e849 --- /dev/null +++ b/tests/test_shared_helpers.py @@ -0,0 +1,134 @@ +"""Tests for PR02: shared helper methods extraction. + +Run from BERTopic repo root: + pytest tests/test_shared_helpers.py -v + +These tests validate the three helpers: +- _aggregate_documents (static) +- _get_feature_names (instance) +- _topic_name_from_words (static) +""" + +import copy + +import pandas as pd +from sklearn.feature_extraction.text import CountVectorizer + +from bertopic import BERTopic + + +class TestAggregateDocuments: + """Tests for BERTopic._aggregate_documents.""" + + def test_basic_groupby(self): + df = pd.DataFrame( + { + "Document": ["hello world", "foo bar", "baz qux"], + "Topic": [0, 0, 1], + "ID": [0, 1, 2], + } + ) + result = BERTopic._aggregate_documents(df) + assert len(result) == 2 + assert result.loc[result.Topic == 0, "Document"].to_numpy()[0] == "hello world foo bar" + assert result.loc[result.Topic == 1, "Document"].to_numpy()[0] == "baz qux" + + def test_empty_dataframe(self): + df = pd.DataFrame({"Document": [], "Topic": []}) + result = BERTopic._aggregate_documents(df) + assert len(result) == 0 + + def test_single_doc_per_topic(self): + df = pd.DataFrame( + { + "Document": ["only doc"], + "Topic": [0], + } + ) + result = BERTopic._aggregate_documents(df) + assert len(result) == 1 + assert result.loc[result.Topic == 0, "Document"].to_numpy()[0] == "only doc" + + def test_drops_extra_columns(self): + """Extra columns should not appear in output unless aggregated.""" + df = pd.DataFrame( + { + "Document": ["a", "b"], + "Topic": [0, 0], + "ID": [0, 1], + "Image": [None, None], + } + ) + result = BERTopic._aggregate_documents(df) + assert "Document" in result.columns + assert "Topic" in result.columns + + +class TestGetFeatureNames: + """Tests for BERTopic._get_feature_names.""" + + def test_returns_feature_names(self): + model = BERTopic() + model.vectorizer_model = CountVectorizer() + model.vectorizer_model.fit(["hello world", "foo bar baz"]) + names = model._get_feature_names() + assert "hello" in names + assert "world" in names + assert "foo" in names + + +class TestTopicNameFromWords: + """Tests for BERTopic._topic_name_from_words.""" + + def test_default_n(self): + """Default n=5 should join the first 5 words.""" + words = [(f"word{i}", 0.5 - i * 0.1) for i in range(10)] + name = BERTopic._topic_name_from_words(words) + assert name == "word0_word1_word2_word3_word4" + + def test_custom_n(self): + words = [(f"w{i}", 0.5) for i in range(10)] + name = BERTopic._topic_name_from_words(words, n=3) + assert name == "w0_w1_w2" + + def test_fewer_words_than_n(self): + words = [("only", 0.9)] + name = BERTopic._topic_name_from_words(words, n=5) + assert name == "only" + + def test_empty_words(self): + name = BERTopic._topic_name_from_words([]) + assert name == "" + + +class TestHelpersIntegration: + """Verify refactored methods produce identical output to upstream.""" + + def test_hierarchical_topics_unchanged(self, base_topic_model, documents): + """hierarchical_topics should produce valid output after refactoring.""" + model = copy.deepcopy(base_topic_model) + hier = model.hierarchical_topics(documents) + + assert len(hier) > 0 + assert "Parent_Name" in hier.columns + # Parent names should be underscore-joined keywords + for name in hier["Parent_Name"]: + assert "_" in name or len(name.split("_")) >= 1 + + def test_extract_topics_unchanged(self, base_topic_model, documents, document_embeddings): + """_extract_topics should produce identical topic representations.""" + model = copy.deepcopy(base_topic_model) + original_topics = dict(model.topic_representations_) + + # Re-run _extract_topics + docs_df = pd.DataFrame( + { + "Document": documents, + "ID": range(len(documents)), + "Topic": model.topics_, + } + ) + model._extract_topics(docs_df, embeddings=document_embeddings) + + # Topic representations should be identical + assert set(model.topic_representations_.keys()) == set(original_topics.keys())