Skip to content

Backends

BaseEmbedder

The Base Embedder used for creating embedding models.

Parameters:

Name Type Description Default
embedding_model

The main embedding model to be used for extracting document and word embedding

None
Source code in bertopic\backend\_base.py
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class BaseEmbedder:
    """The Base Embedder used for creating embedding models.

    Arguments:
        embedding_model: The main embedding model to be used for extracting
                         document and word embedding
    """

    def __init__(self, embedding_model=None):
        self.embedding_model = embedding_model

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        pass

    def embed_words(self, words: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n words into an n-dimensional
        matrix of embeddings.

        Words are the short texts BERTopic compares with documents or topics, such as
        a search term, a zero-shot label or a keyword. They are embedded as documents,
        unless a backend embeds them differently, such as without a prompt meant for documents.

        Arguments:
            words: A list of words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Word embeddings with shape (n, m) with `n` words
            that each have an embeddings size of `m`
        """
        return self.embed_documents(words, verbose)

    def embed_documents(self, document: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents into an n-dimensional
        matrix of embeddings.

        Arguments:
            document: A list of documents to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document embeddings with shape (n, m) with `n` documents
            that each have an embeddings size of `m`
        """
        return self.embed(document, verbose)

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_base.py
16
17
18
19
20
21
22
23
24
25
26
27
28
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    pass

embed_documents(document, verbose=False)

Embed a list of n documents into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
document List[str]

A list of documents to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document embeddings with shape (n, m) with n documents

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_base.py
48
49
50
51
52
53
54
55
56
57
58
59
60
def embed_documents(self, document: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents into an n-dimensional
    matrix of embeddings.

    Arguments:
        document: A list of documents to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document embeddings with shape (n, m) with `n` documents
        that each have an embeddings size of `m`
    """
    return self.embed(document, verbose)

embed_words(words, verbose=False)

Embed a list of n words into an n-dimensional matrix of embeddings.

Words are the short texts BERTopic compares with documents or topics, such as a search term, a zero-shot label or a keyword. They are embedded as documents, unless a backend embeds them differently, such as without a prompt meant for documents.

Parameters:

Name Type Description Default
words List[str]

A list of words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Word embeddings with shape (n, m) with n words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_base.py
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def embed_words(self, words: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n words into an n-dimensional
    matrix of embeddings.

    Words are the short texts BERTopic compares with documents or topics, such as
    a search term, a zero-shot label or a keyword. They are embedded as documents,
    unless a backend embeds them differently, such as without a prompt meant for documents.

    Arguments:
        words: A list of words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Word embeddings with shape (n, m) with `n` words
        that each have an embeddings size of `m`
    """
    return self.embed_documents(words, verbose)

CohereBackend

Bases: BaseEmbedder

Cohere Embedding Model.

Parameters:

Name Type Description Default
client

A cohere client.

required
embedding_model str

A Cohere model. Default is "large". For an overview of models see: https://docs.cohere.ai/docs/generation-card

'large'
delay_in_seconds float | None

If a batch_size is given, use this set the delay in seconds between batches.

None
batch_size int | None

The size of each batch.

None
embed_kwargs Mapping[str, Any] | None

Kwargs passed to cohere.Client.embed. Can be used to define additional parameters such as input_type

None

Examples:

import cohere
from bertopic.backend import CohereBackend

client = cohere.Client("APIKEY")
cohere_model = CohereBackend(client)

If you want to specify input_type:

cohere_model = CohereBackend(
    client,
    embedding_model="embed-english-v3.0",
    embed_kwargs={"input_type": "clustering"}
)
Source code in bertopic\backend\_cohere.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
class CohereBackend(BaseEmbedder):
    """Cohere Embedding Model.

    Arguments:
        client: A `cohere` client.
        embedding_model: A Cohere model. Default is "large".
                         For an overview of models see:
                         https://docs.cohere.ai/docs/generation-card
        delay_in_seconds: If a `batch_size` is given, use this set
                          the delay in seconds between batches.
        batch_size: The size of each batch.
        embed_kwargs: Kwargs passed to `cohere.Client.embed`.
                            Can be used to define additional parameters
                            such as `input_type`

    Examples:
    ```python
    import cohere
    from bertopic.backend import CohereBackend

    client = cohere.Client("APIKEY")
    cohere_model = CohereBackend(client)
    ```

    If you want to specify `input_type`:

    ```python
    cohere_model = CohereBackend(
        client,
        embedding_model="embed-english-v3.0",
        embed_kwargs={"input_type": "clustering"}
    )
    ```
    """

    def __init__(
        self,
        client,
        embedding_model: str = "large",
        delay_in_seconds: float | None = None,
        batch_size: int | None = None,
        embed_kwargs: Mapping[str, Any] | None = None,
    ):
        super().__init__()
        self.client = client
        self.embedding_model = embedding_model
        self.delay_in_seconds = delay_in_seconds
        self.batch_size = batch_size
        self.embed_kwargs = dict(embed_kwargs or {})

        if self.embed_kwargs.get("model"):
            self.embedding_model = self.embed_kwargs.get("model")
        else:
            self.embed_kwargs["model"] = self.embedding_model

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        # Batch-wise embedding extraction
        if self.batch_size is not None:
            embeddings = []
            for batch in tqdm(self._chunks(documents), disable=not verbose):
                response = self.client.embed(texts=batch, **self.embed_kwargs)
                embeddings.extend(response.embeddings)

                # Delay subsequent calls
                if self.delay_in_seconds:
                    time.sleep(self.delay_in_seconds)

        # Extract embeddings all at once
        else:
            response = self.client.embed(texts=documents, **self.embed_kwargs)
            embeddings = response.embeddings
        return np.array(embeddings)

    def _chunks(self, documents):
        for i in range(0, len(documents), self.batch_size):
            yield documents[i : i + self.batch_size]

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_cohere.py
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    # Batch-wise embedding extraction
    if self.batch_size is not None:
        embeddings = []
        for batch in tqdm(self._chunks(documents), disable=not verbose):
            response = self.client.embed(texts=batch, **self.embed_kwargs)
            embeddings.extend(response.embeddings)

            # Delay subsequent calls
            if self.delay_in_seconds:
                time.sleep(self.delay_in_seconds)

    # Extract embeddings all at once
    else:
        response = self.client.embed(texts=documents, **self.embed_kwargs)
        embeddings = response.embeddings
    return np.array(embeddings)

FastEmbedBackend

Bases: BaseEmbedder

FastEmbed embedding model.

The FastEmbed embedding model used for generating sentence embeddings.

Parameters:

Name Type Description Default
embedding_model str

A FastEmbed embedding model

'BAAI/bge-small-en-v1.5'

Examples: To create a model, you can load in a string pointing to a supported FastEmbed model:

from bertopic.backend import FastEmbedBackend

sentence_model = FastEmbedBackend("BAAI/bge-small-en-v1.5")
Source code in bertopic\backend\_fastembed.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class FastEmbedBackend(BaseEmbedder):
    """FastEmbed embedding model.

    The FastEmbed embedding model used for generating sentence embeddings.

    Arguments:
        embedding_model: A FastEmbed embedding model

    Examples:
    To create a model, you can load in a string pointing to a supported
    FastEmbed model:

    ```python
    from bertopic.backend import FastEmbedBackend

    sentence_model = FastEmbedBackend("BAAI/bge-small-en-v1.5")
    ```
    """

    def __init__(self, embedding_model: str = "BAAI/bge-small-en-v1.5"):
        super().__init__()

        supported_models = [m["model"] for m in TextEmbedding.list_supported_models()]

        if isinstance(embedding_model, str) and embedding_model in supported_models:
            self.embedding_model = TextEmbedding(model_name=embedding_model)
        else:
            raise ValueError(
                "Please select a correct FasteEmbed model: \n"
                "the model must be a string and must be supported. \n"
                "The supported TextEmbedding model list is here: https://qdrant.github.io/fastembed/examples/Supported_Models/"
            )

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        embeddings = np.array(list(self.embedding_model.embed(documents, show_progress_bar=verbose)))
        return embeddings

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_fastembed.py
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    embeddings = np.array(list(self.embedding_model.embed(documents, show_progress_bar=verbose)))
    return embeddings

LangChainBackend

Bases: BaseEmbedder

LangChain Embedding Model.

This class uses the LangChain Embedding class to embed the documents. Argument: embedding_model: A LangChain Embedding Instance.

Examples:

from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from bertopic.backend import LangChainBackend

hf_embedding = HuggingFaceInstructEmbeddings()
langchain_embedder = LangChainBackend(hf_embedding)

Source code in bertopic\backend\_langchain.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class LangChainBackend(BaseEmbedder):
    """LangChain Embedding Model.

    This class uses the LangChain Embedding class to embed the documents.
    Argument:
        embedding_model: A LangChain Embedding Instance.

    Examples:
    ```python
    from langchain_community.embeddings import HuggingFaceInstructEmbeddings
    from bertopic.backend import LangChainBackend

    hf_embedding = HuggingFaceInstructEmbeddings()
    langchain_embedder = LangChainBackend(hf_embedding)
    ```
    """

    def __init__(self, embedding_model: Embeddings):
        self.embedding_model = embedding_model

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        # Prepare documents, replacing empty strings with a single space
        prepared_documents = [" " if doc == "" else doc for doc in documents]
        response = self.embedding_model.embed_documents(prepared_documents)
        return np.array(response)

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_langchain.py
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    # Prepare documents, replacing empty strings with a single space
    prepared_documents = [" " if doc == "" else doc for doc in documents]
    response = self.embedding_model.embed_documents(prepared_documents)
    return np.array(response)

Model2VecBackend

Bases: BaseEmbedder

Model2Vec embedding model.

Parameters:

Name Type Description Default
embedding_model Union[str, StaticModel]

Either a model2vec model or a string pointing to a model2vec model

required
distill bool

Indicates whether to distill a sentence-transformers compatible model. The distillation will happen during fitting of the topic model. NOTE: Only works if embedding_model is a string.

False
distill_kwargs dict | None

Keyword arguments to pass to the distillation process of model2vec.distill.distill

None
distill_vectorizer str | None

A CountVectorizer used for creating a custom vocabulary based on the same documents used for topic modeling. NOTE: If "vocabulary" is in distill_kwargs, this will be ignored.

None

Examples: To create a model, you can load in a string pointing to a model2vec model:

from bertopic.backend import Model2VecBackend

sentence_model = Model2VecBackend("minishlab/potion-base-8M")

or you can instantiate a model yourself:

from bertopic.backend import Model2VecBackend
from model2vec import StaticModel

embedding_model = StaticModel.from_pretrained("minishlab/potion-base-8M")
sentence_model = Model2VecBackend(embedding_model)

If you want to distill a sentence-transformers model with the vocabulary of the documents, run the following:

from bertopic.backend import Model2VecBackend

sentence_model = Model2VecBackend("sentence-transformers/all-MiniLM-L6-v2", distill=True)
Source code in bertopic\backend\_model2vec.py
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
class Model2VecBackend(BaseEmbedder):
    """Model2Vec embedding model.

    Arguments:
        embedding_model: Either a model2vec model or a
                         string pointing to a model2vec model
        distill: Indicates whether to distill a sentence-transformers compatible model.
                 The distillation will happen during fitting of the topic model.
                 NOTE: Only works if `embedding_model` is a string.
        distill_kwargs: Keyword arguments to pass to the distillation process
                        of `model2vec.distill.distill`
        distill_vectorizer: A CountVectorizer used for creating a custom vocabulary
                            based on the same documents used for topic modeling.
                            NOTE: If "vocabulary" is in `distill_kwargs`, this will be ignored.

    Examples:
    To create a model, you can load in a string pointing to a
    model2vec model:

    ```python
    from bertopic.backend import Model2VecBackend

    sentence_model = Model2VecBackend("minishlab/potion-base-8M")
    ```

    or  you can instantiate a model yourself:

    ```python
    from bertopic.backend import Model2VecBackend
    from model2vec import StaticModel

    embedding_model = StaticModel.from_pretrained("minishlab/potion-base-8M")
    sentence_model = Model2VecBackend(embedding_model)
    ```

    If you want to distill a sentence-transformers model with the vocabulary of the documents,
    run the following:

    ```python
    from bertopic.backend import Model2VecBackend

    sentence_model = Model2VecBackend("sentence-transformers/all-MiniLM-L6-v2", distill=True)
    ```
    """

    def __init__(
        self,
        embedding_model: Union[str, StaticModel],
        distill: bool = False,
        distill_kwargs: dict | None = None,
        distill_vectorizer: str | None = None,
    ):
        super().__init__()

        self.distill = distill
        self.distill_kwargs = dict(distill_kwargs or {})
        self.distill_vectorizer = distill_vectorizer
        self._has_distilled = False

        # When we distill, we need a string pointing to a sentence-transformer model
        if self.distill:
            self._check_model2vec_installation()
            if not self.distill_vectorizer:
                self.distill_vectorizer = CountVectorizer()
            if isinstance(embedding_model, str):
                self.embedding_model = embedding_model
            else:
                raise ValueError(
                    "Please pass a string pointing to a sentence-transformer model when distilling."
                )

        # If we don't distill, we can pass a model2vec model directly or load from a string
        elif isinstance(embedding_model, StaticModel):
            self.embedding_model = embedding_model
        elif isinstance(embedding_model, str):
            self.embedding_model = StaticModel.from_pretrained(embedding_model)
        else:
            raise ValueError(
                "Please select a correct Model2Vec model: \n"
                "`from model2vec import StaticModel` \n"
                "`model = StaticModel.from_pretrained('minishlab/potion-base-8M')`"
            )

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        # Distill the model
        if self.distill and not self._has_distilled:
            from model2vec.distill import distill

            # Distill with the vocabulary of the documents
            if not self.distill_kwargs.get("vocabulary"):
                X = self.distill_vectorizer.fit_transform(documents)
                word_counts = np.array(X.sum(axis=0)).flatten()
                words = self.distill_vectorizer.get_feature_names_out()
                vocabulary = [
                    word for word, _ in sorted(zip(words, word_counts), key=lambda x: x[1], reverse=True)
                ]
                self.distill_kwargs["vocabulary"] = vocabulary

            # Distill the model
            self.embedding_model = distill(self.embedding_model, **self.distill_kwargs)

            # Distillation should happen only once and not for every embed call
            # The distillation should only happen the first time on the entire vocabulary
            self._has_distilled = True

        # Embed the documents
        embeddings = self.embedding_model.encode(documents, show_progress_bar=verbose)
        return embeddings

    def _check_model2vec_installation(self):
        try:
            from model2vec.distill import distill  # noqa: F401
        except ImportError:
            raise ImportError(
                "To distill a model using model2vec, you need to run `pip install model2vec[distill]`"
            )

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_model2vec.py
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    # Distill the model
    if self.distill and not self._has_distilled:
        from model2vec.distill import distill

        # Distill with the vocabulary of the documents
        if not self.distill_kwargs.get("vocabulary"):
            X = self.distill_vectorizer.fit_transform(documents)
            word_counts = np.array(X.sum(axis=0)).flatten()
            words = self.distill_vectorizer.get_feature_names_out()
            vocabulary = [
                word for word, _ in sorted(zip(words, word_counts), key=lambda x: x[1], reverse=True)
            ]
            self.distill_kwargs["vocabulary"] = vocabulary

        # Distill the model
        self.embedding_model = distill(self.embedding_model, **self.distill_kwargs)

        # Distillation should happen only once and not for every embed call
        # The distillation should only happen the first time on the entire vocabulary
        self._has_distilled = True

    # Embed the documents
    embeddings = self.embedding_model.encode(documents, show_progress_bar=verbose)
    return embeddings

MultiModalBackend

Bases: BaseEmbedder

Multimodal backend using Sentence-transformers.

One model may cover every modality, or each may be given its own. A modality without its own model falls back to embedding_model.

Note that rows of different modalities are only comparable when a single model embeds them all, since two separately trained encoders do not share a vector space. Per-modality models are therefore for corpora of one modality; mixed corpora want a joint model.

Parameters:

Name Type Description Default
embedding_model Union[str, SentenceTransformer]

A sentence-transformers embedding model covering every modality that has no model of its own, text included.

required
image_model Union[str, SentenceTransformer]

A sentence-transformers embedding model used to embed only images.

None
audio_model Union[str, SentenceTransformer]

A sentence-transformers embedding model used to embed only audio.

None
video_model Union[str, SentenceTransformer]

A sentence-transformers embedding model used to embed only video.

None
code_model Union[str, SentenceTransformer]

A sentence-transformers embedding model used to embed only code.

None
batch_size int

The number of items to embed at a time. Images, audio and video are capped at smaller batches of their own to prevent using too much memory.

32

Examples: To create a model, you can load in a string pointing to a sentence-transformers model:

from bertopic.backend import MultiModalBackend

sentence_model = MultiModalBackend("google/embeddinggemma-2")

or you can instantiate a model yourself:

from bertopic.backend import MultiModalBackend
from sentence_transformers import SentenceTransformer

embedding_model = SentenceTransformer("google/embeddinggemma-2", default_prompt_name="Clustering")
embedding_model.max_seq_length = 1024
sentence_model = MultiModalBackend(embedding_model)

Source code in bertopic\backend\_multimodal.py
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
class MultiModalBackend(BaseEmbedder):
    """Multimodal backend using Sentence-transformers.

    One model may cover every modality, or each may be given its own. A modality without
    its own model falls back to `embedding_model`.

    Note that rows of different modalities are only comparable when a single model embeds
    them all, since two separately trained encoders do not share a vector space. Per-modality
    models are therefore for corpora of one modality; mixed corpora want a joint model.

    Arguments:
        embedding_model: A sentence-transformers embedding model covering every modality
                         that has no model of its own, text included.
        image_model: A sentence-transformers embedding model used to embed only images.
        audio_model: A sentence-transformers embedding model used to embed only audio.
        video_model: A sentence-transformers embedding model used to embed only video.
        code_model: A sentence-transformers embedding model used to embed only code.
        batch_size: The number of items to embed at a time. Images, audio and video are capped at
                    smaller batches of their own to prevent using too much memory.

    Examples:
    To create a model, you can load in a string pointing to a
    sentence-transformers model:

    ```python
    from bertopic.backend import MultiModalBackend

    sentence_model = MultiModalBackend("google/embeddinggemma-2")
    ```

    or  you can instantiate a model yourself:
    ```python
    from bertopic.backend import MultiModalBackend
    from sentence_transformers import SentenceTransformer

    embedding_model = SentenceTransformer("google/embeddinggemma-2", default_prompt_name="Clustering")
    embedding_model.max_seq_length = 1024
    sentence_model = MultiModalBackend(embedding_model)
    ```
    """

    def __init__(
        self,
        embedding_model: Union[str, SentenceTransformer],
        image_model: Union[str, SentenceTransformer] = None,
        batch_size: int = 32,
        *,
        audio_model: Union[str, SentenceTransformer] = None,
        video_model: Union[str, SentenceTransformer] = None,
        code_model: Union[str, SentenceTransformer] = None,
    ):
        super().__init__()
        self.batch_size = batch_size
        self.models = {}

        # Text is resolved first, since it is what every other modality falls back to
        for modality, model in (
            (Modality.TEXT, embedding_model),
            (Modality.IMAGE, image_model),
            (Modality.AUDIO, audio_model),
            (Modality.VIDEO, video_model),
            (Modality.CODE, code_model),
        ):
            if isinstance(model, SentenceTransformer):
                self.models[modality] = model
            elif isinstance(model, str):
                self.models[modality] = load_sentence_transformer(model)
            elif model is None:
                self.models[modality] = self.models[Modality.TEXT]
            else:
                raise ValueError(
                    "Please select a correct SentenceTransformers model: \n"
                    "`from sentence_transformers import SentenceTransformer` \n"
                    "`model = SentenceTransformer('google/embeddinggemma-2')`"
                )

        self.embedding_model = self.models[Modality.TEXT]

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        return self.embed_media(documents, Modality.TEXT, verbose)

    def embed_words(self, words: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n words into an n-dimensional matrix of embeddings.

        A clustering prompt is meant for texts that are clustered together, and on a few
        words it outweighs what they say, so words are embedded without the model's prompt.

        Arguments:
            words: A list of words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Word embeddings with shape (n, m) with `n` words
            that each have an embeddings size of `m`
        """
        return self.models[Modality.TEXT].encode(
            words, batch_size=self.batch_size, show_progress_bar=verbose, prompt=""
        )

    def embed_media(
        self, media: List, modality: Modality = Modality.TEXT, verbose: bool = False
    ) -> np.ndarray:
        """Embed a list of n items sharing one modality, with the model that handles it.

        Arguments:
            media: A list of items, all of the same modality
            modality: Which modality the items are, choosing the model that embeds them
            verbose: Controls the verbosity of the process

        Returns:
            Embeddings with shape (n, m) with `n` items that each have an
            embeddings size of `m`
        """
        modality = Modality(modality)
        model = self.models[modality]
        batch_size = min(self.batch_size, MEDIA_BATCH_SIZES.get(modality, self.batch_size))

        # A model that takes this kind of media is told what each item is, so a path is never read as text.
        # The clustering prompt and the length limit are for text, so media goes without both
        if modality != Modality.TEXT and modality in model.modalities:
            items = [{modality.value: item} for item in media]
            return model.encode(
                items,
                batch_size=batch_size,
                show_progress_bar=verbose,
                prompt="",
                processing_kwargs={"text": {"truncation": False}},
            )
        return model.encode(media, batch_size=batch_size, show_progress_bar=verbose)

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_multimodal.py
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    return self.embed_media(documents, Modality.TEXT, verbose)

embed_media(media, modality=Modality.TEXT, verbose=False)

Embed a list of n items sharing one modality, with the model that handles it.

Parameters:

Name Type Description Default
media List

A list of items, all of the same modality

required
modality Modality

Which modality the items are, choosing the model that embeds them

TEXT
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Embeddings with shape (n, m) with n items that each have an

ndarray

embeddings size of m

Source code in bertopic\backend\_multimodal.py
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def embed_media(
    self, media: List, modality: Modality = Modality.TEXT, verbose: bool = False
) -> np.ndarray:
    """Embed a list of n items sharing one modality, with the model that handles it.

    Arguments:
        media: A list of items, all of the same modality
        modality: Which modality the items are, choosing the model that embeds them
        verbose: Controls the verbosity of the process

    Returns:
        Embeddings with shape (n, m) with `n` items that each have an
        embeddings size of `m`
    """
    modality = Modality(modality)
    model = self.models[modality]
    batch_size = min(self.batch_size, MEDIA_BATCH_SIZES.get(modality, self.batch_size))

    # A model that takes this kind of media is told what each item is, so a path is never read as text.
    # The clustering prompt and the length limit are for text, so media goes without both
    if modality != Modality.TEXT and modality in model.modalities:
        items = [{modality.value: item} for item in media]
        return model.encode(
            items,
            batch_size=batch_size,
            show_progress_bar=verbose,
            prompt="",
            processing_kwargs={"text": {"truncation": False}},
        )
    return model.encode(media, batch_size=batch_size, show_progress_bar=verbose)

embed_words(words, verbose=False)

Embed a list of n words into an n-dimensional matrix of embeddings.

A clustering prompt is meant for texts that are clustered together, and on a few words it outweighs what they say, so words are embedded without the model's prompt.

Parameters:

Name Type Description Default
words List[str]

A list of words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Word embeddings with shape (n, m) with n words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_multimodal.py
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def embed_words(self, words: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n words into an n-dimensional matrix of embeddings.

    A clustering prompt is meant for texts that are clustered together, and on a few
    words it outweighs what they say, so words are embedded without the model's prompt.

    Arguments:
        words: A list of words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Word embeddings with shape (n, m) with `n` words
        that each have an embeddings size of `m`
    """
    return self.models[Modality.TEXT].encode(
        words, batch_size=self.batch_size, show_progress_bar=verbose, prompt=""
    )

OpenAIBackend

Bases: BaseEmbedder

OpenAI Embedding Model.

Parameters:

Name Type Description Default
client OpenAI

A openai.OpenAI client.

required
embedding_model str

An OpenAI model. Default is For an overview of models see: https://platform.openai.com/docs/models/embeddings

'text-embedding-ada-002'
delay_in_seconds float | None

If a batch_size is given, use this set the delay in seconds between batches.

None
batch_size int | None

The size of each batch.

None
generator_kwargs Mapping[str, Any] | None

Kwargs passed to openai.Embedding.create. Can be used to define custom engines or deployment_ids.

None

Examples:

import openai
from bertopic.backend import OpenAIBackend

client = openai.OpenAI(api_key="sk-...")
openai_embedder = OpenAIBackend(client, "text-embedding-ada-002")

Source code in bertopic\backend\_openai.py
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class OpenAIBackend(BaseEmbedder):
    """OpenAI Embedding Model.

    Arguments:
        client: A `openai.OpenAI` client.
        embedding_model: An OpenAI model. Default is
                         For an overview of models see:
                         https://platform.openai.com/docs/models/embeddings
        delay_in_seconds: If a `batch_size` is given, use this set
                          the delay in seconds between batches.
        batch_size: The size of each batch.
        generator_kwargs: Kwargs passed to `openai.Embedding.create`.
                          Can be used to define custom engines or
                          deployment_ids.

    Examples:
    ```python
    import openai
    from bertopic.backend import OpenAIBackend

    client = openai.OpenAI(api_key="sk-...")
    openai_embedder = OpenAIBackend(client, "text-embedding-ada-002")
    ```
    """

    def __init__(
        self,
        client: openai.OpenAI,
        embedding_model: str = "text-embedding-ada-002",
        delay_in_seconds: float | None = None,
        batch_size: int | None = None,
        generator_kwargs: Mapping[str, Any] | None = None,
    ):
        super().__init__()
        self.client = client
        self.embedding_model = embedding_model
        self.delay_in_seconds = delay_in_seconds
        self.batch_size = batch_size
        self.generator_kwargs = dict(generator_kwargs or {})

        if self.generator_kwargs.get("model"):
            self.embedding_model = self.generator_kwargs.get("model")
        elif not self.generator_kwargs.get("engine"):
            self.generator_kwargs["model"] = self.embedding_model

    def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
        """Embed a list of n documents/words into an n-dimensional
        matrix of embeddings.

        Arguments:
            documents: A list of documents or words to be embedded
            verbose: Controls the verbosity of the process

        Returns:
            Document/words embeddings with shape (n, m) with `n` documents/words
            that each have an embeddings size of `m`
        """
        # Prepare documents, replacing empty strings with a single space
        prepared_documents = [" " if doc == "" else doc for doc in documents]

        # Batch-wise embedding extraction
        if self.batch_size is not None:
            embeddings = []
            for batch in tqdm(self._chunks(prepared_documents), disable=not verbose):
                response = self.client.embeddings.create(input=batch, **self.generator_kwargs)
                embeddings.extend([r.embedding for r in response.data])

                # Delay subsequent calls
                if self.delay_in_seconds:
                    time.sleep(self.delay_in_seconds)

        # Extract embeddings all at once
        else:
            response = self.client.embeddings.create(input=prepared_documents, **self.generator_kwargs)
            embeddings = [r.embedding for r in response.data]
        return np.array(embeddings)

    def _chunks(self, documents):
        for i in range(0, len(documents), self.batch_size):
            yield documents[i : i + self.batch_size]

embed(documents, verbose=False)

Embed a list of n documents/words into an n-dimensional matrix of embeddings.

Parameters:

Name Type Description Default
documents List[str]

A list of documents or words to be embedded

required
verbose bool

Controls the verbosity of the process

False

Returns:

Type Description
ndarray

Document/words embeddings with shape (n, m) with n documents/words

ndarray

that each have an embeddings size of m

Source code in bertopic\backend\_openai.py
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def embed(self, documents: List[str], verbose: bool = False) -> np.ndarray:
    """Embed a list of n documents/words into an n-dimensional
    matrix of embeddings.

    Arguments:
        documents: A list of documents or words to be embedded
        verbose: Controls the verbosity of the process

    Returns:
        Document/words embeddings with shape (n, m) with `n` documents/words
        that each have an embeddings size of `m`
    """
    # Prepare documents, replacing empty strings with a single space
    prepared_documents = [" " if doc == "" else doc for doc in documents]

    # Batch-wise embedding extraction
    if self.batch_size is not None:
        embeddings = []
        for batch in tqdm(self._chunks(prepared_documents), disable=not verbose):
            response = self.client.embeddings.create(input=batch, **self.generator_kwargs)
            embeddings.extend([r.embedding for r in response.data])

            # Delay subsequent calls
            if self.delay_in_seconds:
                time.sleep(self.delay_in_seconds)

    # Extract embeddings all at once
    else:
        response = self.client.embeddings.create(input=prepared_documents, **self.generator_kwargs)
        embeddings = [r.embedding for r in response.data]
    return np.array(embeddings)