Method transform

bertopic/_bertopic.py:545–647

After having fit a model, use transform to predict new instances. Arguments: documents: A single document or a list of documents to predict on embeddings: Pre-trained document embeddings. These can be used instead of the sentence-transformer m

(
        self,
        documents: Union[str, List[str]],
        embeddings: np.ndarray = None,
        images: List[str] | None = None,
    )

        return predictions, self.probabilities_

    def transform(
        self,
        documents: Union[str, List[str]],
        embeddings: np.ndarray = None,
        images: List[str] | None = None,
    ) -> Tuple[List[int], np.ndarray]:
        """After having fit a model, use transform to predict new instances.

        Arguments:
            documents: A single document or a list of documents to predict on
            embeddings: Pre-trained document embeddings. These can be used
                        instead of the sentence-transformer model.
            images: A list of paths to the images to predict on or the images themselves

        Returns:
            predictions: Topic predictions for each documents
            probabilities: The topic probability distribution which is returned by default.
                           If `calculate_probabilities` in BERTopic is set to False, then the
                           probabilities are not calculated to speed up computation and
                           decrease memory usage.

        Examples:
        ```python
        from bertopic import BERTopic
        from sklearn.datasets import fetch_20newsgroups

        docs = fetch_20newsgroups(subset='all')['data']
        topic_model = BERTopic().fit(docs)
        topics, probs = topic_model.transform(docs)
        ```

        If you want to use your own embeddings:

        ```python
        from bertopic import BERTopic
        from sklearn.datasets import fetch_20newsgroups
        from sentence_transformers import SentenceTransformer

        # Create embeddings
        docs = fetch_20newsgroups(subset='all')['data']
        sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
        embeddings = sentence_model.encode(docs, show_progress_bar=True)

        # Create topic model
        topic_model = BERTopic().fit(docs, embeddings)
        topics, probs = topic_model.transform(docs, embeddings)
        ```
        """
        check_is_fitted(self)
        check_embeddings_shape(embeddings, documents)

        if isinstance(documents, str) or documents is None:
            documents = [documents]

        if embeddings is None:
            embeddings = self._extract_embeddings(documents, images=images, method="document", verbose=self.verbose)

        # Check if an embedding model was found