Suggestion: Batched cosine similarity for efficient memory handling of large documents

Feature request

To help mitigate excessive memory consumption in cosine similarity for a large set of documents, allow a batched approach (when calculate_probabilities is False)

Motivation

Oftentimes the whole similarity matrix is brought into memory when it doesn't need to be (ie when we only take the assignment, assignment_val), which can cause memory issues when the document set is large

Your contribution

Something like

@staticmethod
def batched_cosine_similarity_argmax(
        embeddings: np.ndarray, zeroshot_embeddings: np.ndarray, batch_size: int = 1000
    ):
        """
        Assigns each embedding to the zeroshot embedding with the highest cosine similarity, processing in batches.
        Args:
            embeddings (np.ndarray): Array of shape (n_samples, n_features) containing the input embeddings to be assigned.
            zeroshot_embeddings (np.ndarray): Array of shape (n_zeroshot, n_features) containing the zeroshot embeddings to compare against.
            batch_size (int, optional): Number of embeddings to process per batch. Defaults to 1000.
        Returns:
            Tuple[np.ndarray, np.ndarray]:
                - assignment: Array of shape (n_samples,) with the index of the most similar zeroshot embedding for each input embedding.
                - assignment_vals: Array of shape (n_samples,) with the corresponding maximum cosine similarity values.
        """
        logger.info("Executing batched cosine similarity")

        assignment = []
        assignment_vals = []

        for i in range(0, embeddings.shape[0], batch_size):
            batch = embeddings[i : i + batch_size]
            sim = cosine_similarity(batch, zeroshot_embeddings)
            max_indices = np.argmax(sim, axis=1)
            max_vals = np.max(sim, axis=1)
            assignment.extend(max_indices)
            assignment_vals.extend(max_vals)

        return np.array(assignment), np.array(assignment_vals)

Then for illustration, zeroshot modelling could become

def _zeroshot_topic_modeling(
        self, documents: pd.DataFrame, embeddings: np.ndarray, batch_size: int = 1000
    ) -> Tuple[pd.DataFrame, np.array, pd.DataFrame, np.array]:
        """Find documents that could be assigned to either one of the topics in self.zeroshot_topic_list.

        We transform the topics in `self.zeroshot_topic_list` to embeddings and
        compare them through cosine similarity with the document embeddings.
        If they pass the `self.zeroshot_min_similarity` threshold, they are assigned.

        Arguments:
            documents: Dataframe with documents and their corresponding IDs
            embeddings: The document embeddings

        Returns:
            documents: The leftover documents that were not assigned to any topic
            embeddings: The leftover embeddings that were not assigned to any topic
        """
        logger.info(
            "Zeroshot Step 1 - Finding documents that could be assigned to either one"
            " of the zero-shot topics"
        )
        # Similarity between document and zero-shot topic embeddings
        zeroshot_embeddings = self._extract_embeddings(self.zeroshot_topic_list)

        assignment, assignment_vals = self.batched_cosine_similarity_argmax(
            embeddings, zeroshot_embeddings, batch_size=batch_size
        )

        assigned_ids = [
            index
            for index, value in enumerate(assignment_vals)
            if value >= self.zeroshot_min_similarity
        ]
        non_assigned_ids = [
            index
            for index, value in enumerate(assignment_vals)
            if value < self.zeroshot_min_similarity
        ]

        # Assign topics
        assigned_documents = documents.iloc[assigned_ids]
        assigned_documents["Topic"] = [topic for topic in assignment[assigned_ids]]

        assigned_documents["Old_ID"] = assigned_documents["ID"].copy()
        assigned_documents["ID"] = range(len(assigned_documents))
        assigned_embeddings = embeddings[assigned_ids]

        # 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 != "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"
                    f" matched zero-shot topics ({num_zeroshot_topics}). Consider"
                    " raising nr_topics or raising the zeroshot_min_similarity"
                    f" ({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

Start by locating the existing cosine-similarity path and the _zeroshot_topic_modeling entry point. Compare the proposed batched assignment behavior with the current similarity calculation, then verify that large document sets avoid retaining the full similarity matrix while assignments and similarity values remain equivalent.