Fix #2102 (#2105)

bertopic/_bertopic.py

@@ -453,10 +453,12 @@ def fit_transform(
             documents, embeddings, assigned_documents, assigned_embeddings = self._zeroshot_topic_modeling(
                 documents, embeddings
             )
+
             # Filter UMAP embeddings to only non-assigned embeddings to be used for clustering
-            umap_embeddings = self.umap_model.transform(embeddings)
+            if len(documents) > 0:
+                umap_embeddings = self.umap_model.transform(embeddings)
 
-        if len(documents) > 0:  # No zero-shot topics matched
+        if len(documents) > 0:
             # Cluster reduced embeddings
             documents, probabilities = self._cluster_embeddings(umap_embeddings, documents, y=y)
             if self._is_zeroshot() and len(assigned_documents) > 0:
@@ -467,7 +469,6 @@ def fit_transform(
             # All documents matches zero-shot topics
             documents = assigned_documents
             embeddings = assigned_embeddings
-        topics_before_reduction = self.topics_
 
         # Sort and Map Topic IDs by their frequency
         if not self.nr_topics:
@@ -505,17 +506,11 @@ def fit_transform(
             sim_matrix = cosine_similarity(embeddings, np.array(self.topic_embeddings_))
 
             if self.calculate_probabilities:
-                probabilities = sim_matrix
+                self.probabilities_ = sim_matrix
             else:
-                # Use `topics_before_reduction` because `self.topics_` may have already been updated from
-                # reducing topics, and the original probabilities are needed for `self._map_probabilities()`
-                probabilities = sim_matrix[
-                    np.arange(len(documents)),
-                    np.array(topics_before_reduction) + self._outliers,
-                ]
-
-        # Resulting output
-        self.probabilities_ = self._map_probabilities(probabilities, original_topics=True)
+                self.probabilities_ = np.max(sim_matrix, axis=1)
+        else:
+            self.probabilities_ = self._map_probabilities(probabilities, original_topics=True)
         predictions = documents.Topic.to_list()
 
         return predictions, self.probabilities_
@@ -2148,7 +2143,7 @@ def merge_topics(
 
         # Update topics
         documents.Topic = documents.Topic.map(mapping)
-        self.topic_mapper_.add_mappings(mapping)
+        self.topic_mapper_.add_mappings(mapping, topic_model=self)
         documents = self._sort_mappings_by_frequency(documents)
         self._extract_topics(documents, mappings=mappings)
         self._update_topic_size(documents)
@@ -3841,19 +3836,24 @@ def _zeroshot_topic_modeling(
 
         # Check that if a number of topics was specified, it exceeds the number of zeroshot topics matched
         num_zeroshot_topics = len(assigned_documents["Topic"].unique())
-        if self.nr_topics and not self.nr_topics > num_zeroshot_topics:
-            raise ValueError(
-                f"The set nr_topics ({self.nr_topics}) must exceed the number of matched zero-shot topics "
-                f"({num_zeroshot_topics}). Consider raising nr_topics or raising the "
-                f"zeroshot_min_similarity ({self.zeroshot_min_similarity})."
-            )
+        if self.nr_topics != "auto":
+            if self.nr_topics and not self.nr_topics > num_zeroshot_topics:
+                raise ValueError(
+                    f"The set nr_topics ({self.nr_topics}) must exceed the number of matched zero-shot topics "
+                    f"({num_zeroshot_topics}). Consider raising nr_topics or raising the "
+                    f"zeroshot_min_similarity ({self.zeroshot_min_similarity})."
+                )
 
         # Select non-assigned topics to be clustered
         documents = documents.iloc[non_assigned_ids]
         documents["Old_ID"] = documents["ID"].copy()
         documents["ID"] = range(len(documents))
         embeddings = embeddings[non_assigned_ids]
 
+        if len(documents) == 0:
+            self.topics_ = assigned_documents["Topic"].values.tolist()
+            self.topic_mapper_ = TopicMapper(self.topics_)
+
         logger.info("Zeroshot Step 1 - Completed \u2713")
         return documents, embeddings, assigned_documents, assigned_embeddings
 
@@ -3914,6 +3914,7 @@ def _combine_zeroshot_topics(
         # Combine the clustered documents/embeddings with assigned documents/embeddings in the original order
         documents = pd.concat([documents, assigned_documents])
         embeddings = np.vstack([embeddings, assigned_embeddings])
+        documents.ID = documents.Old_ID
         sorted_indices = documents.Old_ID.argsort()
         documents = documents.iloc[sorted_indices]
         embeddings = embeddings[sorted_indices]
@@ -4407,50 +4408,12 @@ def _reduce_to_n_topics(self, documents: pd.DataFrame, use_ctfidf: bool = False)
         # Map topics
         documents.Topic = new_topics
         self._update_topic_size(documents)
-        self.topic_mapper_.add_mappings(mapped_topics)
+        self.topic_mapper_.add_mappings(mapped_topics, topic_model=self)
 
         # Update representations
         documents = self._sort_mappings_by_frequency(documents)
         self._extract_topics(documents, mappings=mappings)
 
-        # When zero-shot topic(s) are present in the topics to merge,
-        # determine whether to take one of the zero-shot topic labels
-        # or use a calculated representation.
-        if self._is_zeroshot():
-            new_topic_id_to_zeroshot_topic_idx = {}
-            topics_to_map = {
-                topic_mapping[0]: topic_mapping[1] for topic_mapping in np.array(self.topic_mapper_.mappings_)[:, -2:]
-            }
-
-            for topic_to, topics_from in basic_mappings.items():
-                # When extracting topics, the reduced topics were reordered.
-                # Must get the updated topic_to.
-                topic_to = topics_to_map[topic_to]
-
-                # which of the original topics are zero-shot
-                zeroshot_topic_ids = [
-                    topic_id for topic_id in topics_from if topic_id in self._topic_id_to_zeroshot_topic_idx
-                ]
-                if len(zeroshot_topic_ids) == 0:
-                    continue
-
-                # If any of the original topics are zero-shot, take the best fitting zero-shot label
-                # if the cosine similarity with the new topic exceeds the zero-shot threshold
-                zeroshot_labels = [
-                    self.zeroshot_topic_list[self._topic_id_to_zeroshot_topic_idx[topic_id]]
-                    for topic_id in zeroshot_topic_ids
-                ]
-                zeroshot_embeddings = self._extract_embeddings(zeroshot_labels)
-                cosine_similarities = cosine_similarity(
-                    zeroshot_embeddings, [self.topic_embeddings_[topic_to]]
-                ).flatten()
-                best_zeroshot_topic_idx = np.argmax(cosine_similarities)
-                best_cosine_similarity = cosine_similarities[best_zeroshot_topic_idx]
-                if best_cosine_similarity >= self.zeroshot_min_similarity:
-                    new_topic_id_to_zeroshot_topic_idx[topic_to] = zeroshot_topic_ids[best_zeroshot_topic_idx]
-
-            self._topic_id_to_zeroshot_topic_idx = new_topic_id_to_zeroshot_topic_idx
-
         self._update_topic_size(documents)
         return documents
 
@@ -4503,7 +4466,7 @@ def _auto_reduce_topics(self, documents: pd.DataFrame, use_ctfidf: bool = False)
         }
 
         # Update documents and topics
-        self.topic_mapper_.add_mappings(mapped_topics)
+        self.topic_mapper_.add_mappings(mapped_topics, topic_model=self)
         documents = self._sort_mappings_by_frequency(documents)
         self._extract_topics(documents, mappings=mappings)
         self._update_topic_size(documents)
@@ -4533,13 +4496,17 @@ def _sort_mappings_by_frequency(self, documents: pd.DataFrame) -> pd.DataFrame:
             documents: Updated dataframe with documents and the mapped
                        and re-ordered topic ids
         """
-        self._update_topic_size(documents)
+        # No need to sort if it's the first pass of zero-shot topic modeling
+        nr_zeroshot = len(self._topic_id_to_zeroshot_topic_idx)
+        if self._is_zeroshot and not self.nr_topics and nr_zeroshot > 0:
+            return documents
 
         # Map topics based on frequency
+        self._update_topic_size(documents)
         df = pd.DataFrame(self.topic_sizes_.items(), columns=["Old_Topic", "Size"]).sort_values("Size", ascending=False)
         df = df[df.Old_Topic != -1]
         sorted_topics = {**{-1: -1}, **dict(zip(df.Old_Topic, range(len(df))))}
-        self.topic_mapper_.add_mappings(sorted_topics)
+        self.topic_mapper_.add_mappings(sorted_topics, topic_model=self)
 
         # Map documents
         documents.Topic = documents.Topic.map(sorted_topics).fillna(documents.Topic).astype(int)
@@ -4729,11 +4696,12 @@ def get_mappings(self, original_topics: bool = True) -> Mapping[int, int]:
             mappings = dict(zip(mappings[:, 0], mappings[:, 1]))
         return mappings
 
-    def add_mappings(self, mappings: Mapping[int, int]):
+    def add_mappings(self, mappings: Mapping[int, int], topic_model: BERTopic):
         """Add new column(s) of topic mappings.
 
         Arguments:
             mappings: The mappings to add
+            topic_model: The topic model this TopicMapper belongs to
         """
         for topics in self.mappings_:
             topic = topics[-1]
@@ -4742,6 +4710,50 @@ def add_mappings(self, mappings: Mapping[int, int]):
             else:
                 topics.append(-1)
 
+        # When zero-shot topic(s) are present in the topics to merge,
+        # determine whether to take one of the zero-shot topic labels
+        # or use a calculated representation.
+        if topic_model._is_zeroshot() and len(topic_model._topic_id_to_zeroshot_topic_idx) > 0:
+            new_topic_id_to_zeroshot_topic_idx = {}
+            topics_to_map = {
+                topic_mapping[0]: topic_mapping[1]
+                for topic_mapping in np.array(topic_model.topic_mapper_.mappings_)[:, -2:]
+            }
+
+            # Map topic_to to topics_from
+            mapping = defaultdict(list)
+            for key, value in topics_to_map.items():
+                mapping[value].append(key)
+
+            for topic_to, topics_from in mapping.items():
+                # which of the original topics are zero-shot
+                zeroshot_topic_ids = [
+                    topic_id for topic_id in topics_from if topic_id in topic_model._topic_id_to_zeroshot_topic_idx
+                ]
+                if len(zeroshot_topic_ids) == 0:
+                    continue
+
+                # If any of the original topics are zero-shot, take the best fitting zero-shot label
+                # if the cosine similarity with the new topic exceeds the zero-shot threshold
+                zeroshot_labels = [
+                    topic_model.zeroshot_topic_list[topic_model._topic_id_to_zeroshot_topic_idx[topic_id]]
+                    for topic_id in zeroshot_topic_ids
+                ]
+                zeroshot_embeddings = topic_model._extract_embeddings(zeroshot_labels)
+                cosine_similarities = cosine_similarity(
+                    zeroshot_embeddings, [topic_model.topic_embeddings_[topic_to]]
+                ).flatten()
+                best_zeroshot_topic_idx = np.argmax(cosine_similarities)
+                best_cosine_similarity = cosine_similarities[best_zeroshot_topic_idx]
+
+                if best_cosine_similarity >= topic_model.zeroshot_min_similarity:
+                    # Using the topic ID from before mapping, get the idx into the zeroshot topic list
+                    new_topic_id_to_zeroshot_topic_idx[topic_to] = topic_model._topic_id_to_zeroshot_topic_idx[
+                        zeroshot_topic_ids[best_zeroshot_topic_idx]
+                    ]
+
+            topic_model._topic_id_to_zeroshot_topic_idx = new_topic_id_to_zeroshot_topic_idx
+
     def add_new_topics(self, mappings: Mapping[int, int]):
         """Add new row(s) of topic mappings.