diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..4cb52978 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1429,7 +1429,7 @@ def approximate_distribution( return topic_distributions, topic_token_distributions def find_topics( - self, search_term: str | None = None, image: str | None = None, top_n: int = 5 + self, search_term: str | List[str] | None = None, image: str | None = None, top_n: int = 5 ) -> Tuple[List[int], List[float]]: """Find topics most similar to a search_term. @@ -1444,7 +1444,10 @@ def find_topics( below 5 words. Arguments: - search_term: the term you want to use to search for topics. + search_term: the term you want to use to search for topics. Either a + single string or a list of strings. When a list is passed, + the embeddings of the individual terms are averaged into a + single query embedding. image: path to the image you want to use to search for topics. top_n: the number of topics to return @@ -1460,6 +1463,13 @@ def find_topics( topics, similarity = topic_model.find_topics("sports", top_n=5) ``` + Multiple search terms can be combined into a single query by + passing a list of strings: + + ```python + topics, similarity = topic_model.find_topics(["sports", "football"], top_n=5) + ``` + Note that the search query is typically more accurate if the search_term consists of a phrase or multiple words. """ @@ -1471,7 +1481,12 @@ def find_topics( # Extract search_term embeddings and compare with topic embeddings if search_term is not None: - search_embedding = self._extract_embeddings([search_term], method="word", verbose=False).flatten() + search_terms = [search_term] if isinstance(search_term, str) else list(search_term) + if not search_terms: + raise ValueError("Make sure to pass at least one search term to `search_term`.") + if not all(isinstance(term, str) for term in search_terms): + raise TypeError("`search_term` should either be a string or a list of strings.") + search_embedding = self._extract_embeddings(search_terms, method="word", verbose=False).mean(axis=0) elif image is not None: search_embedding = self._extract_embeddings( [None], images=[image], method="document", verbose=False diff --git a/tests/test_representation/test_representations.py b/tests/test_representation/test_representations.py index fa756625..00481be7 100644 --- a/tests/test_representation/test_representations.py +++ b/tests/test_representation/test_representations.py @@ -182,3 +182,58 @@ def test_find_topics(model, request): assert np.mean(similarity) > 0.1 assert len(similar_topics) > 0 + + +@pytest.mark.parametrize( + "model", + [ + ("kmeans_pca_topic_model"), + ("base_topic_model"), + ], +) +def test_find_topics_single_element_list(model, request): + """A one-element list should behave the same as passing the string itself.""" + topic_model = copy.deepcopy(request.getfixturevalue(model)) + + topics_str, similarity_str = topic_model.find_topics("car") + topics_list, similarity_list = topic_model.find_topics(["car"]) + + assert topics_list == topics_str + assert np.allclose(similarity_list, similarity_str) + + +@pytest.mark.parametrize( + "model", + [ + ("kmeans_pca_topic_model"), + ("base_topic_model"), + ], +) +def test_find_topics_multiple_search_terms(model, request): + """All search terms should contribute, regardless of the order they are passed in.""" + topic_model = copy.deepcopy(request.getfixturevalue(model)) + + topics, similarity = topic_model.find_topics(["car", "computer"]) + reversed_topics, reversed_similarity = topic_model.find_topics(["computer", "car"]) + + assert len(topics) > 0 + assert topics == reversed_topics + assert np.allclose(similarity, reversed_similarity) + + +@pytest.mark.parametrize( + "model", + [ + ("kmeans_pca_topic_model"), + ("base_topic_model"), + ], +) +def test_find_topics_invalid_search_term(model, request): + """Invalid search terms should raise an informative error.""" + topic_model = copy.deepcopy(request.getfixturevalue(model)) + + with pytest.raises(ValueError): + topic_model.find_topics([]) + + with pytest.raises(TypeError): + topic_model.find_topics(["car", 1])