MaartenGr / MaartenGr/BERTopic

KeyBERTInspired re computes representative_docs embeddings even when precomputed embeddings are available

Open
#2,367 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
7.8k
Forks
920
Avg merge
22h 24m
Merged PRs (30d)
5

Description

Have you searched existing issues? 🔎
  • I have searched and found no existing issues

There are a few closely existing issue, but they dont explicitly mention why recomputing embeddings for documents is unnecessary.
Eg- in zero shot modelling, we would need it for word embeddings, but not for document embeddings.
As mentioned in https://github.com/MaartenGr/BERTopic/issues/2011

There is also an issue that closely mentions this, but the code works differently today

You can use pre-calculated embeddings but only for the documents themselves. KeyBERTInspired extracts, based on a number of representative documents, certain words/tokens, embeds them and compare them with the topic embeddings. As such, word/token embeddings will need to be calculated #alongside document embeddings#.

Originally posted by @MaartenGr in #1484

Desribe the bug

Hi,

Essentially, we shouldnt be computing the document embeddings if they are precomputed

    def _extract_embeddings(
        self,
        topic_model,
        topics: Mapping[str, List[Tuple[str, float]]],
        representative_docs: List[str],
        repr_doc_indices: List[List[int]],
    ) -> Union[np.ndarray, List[str]]:
        """Extract the representative document embeddings and create topic embeddings.
        Then extract word embeddings and calculate the cosine similarity between topic
        embeddings and the word embeddings. Topic embeddings are the average of
        representative document embeddings.

        Arguments:
            topic_model: A BERTopic model
            topics: The top words per topic
            representative_docs: A flat list of representative documents
            repr_doc_indices: The indices of representative documents
                              that belong to each topic

        Returns:
            sim: The similarity matrix between word and topic embeddings
            vocab: The complete vocabulary of input documents
        """
        # Calculate representative docs embeddings and create topic embeddings
        repr_embeddings = topic_model._extract_embeddings(representative_docs, method="document", verbose=False)

When running BERTopic for precomputed embeddings, we would have embeddings of documents and also have ids of representative_docs, to extract the repr_embeddings from "embeddings"

    def extract_topics(
        self,
        topic_model,
        documents: pd.DataFrame,
        c_tf_idf: csr_matrix,
        topics: Mapping[str, List[Tuple[str, float]]],
    ) -> Mapping[str, List[Tuple[str, float]]]:
        """Extract topics.

        Arguments:
            topic_model: A BERTopic model
            documents: All input documents
            c_tf_idf: The topic c-TF-IDF representation
            topics: The candidate topics as calculated with c-TF-IDF

        Returns:
            updated_topics: Updated topic representations
        """
        # We extract the top n representative documents per class
        _, representative_docs, repr_doc_indices, _ = topic_model._extract_representative_docs(
            c_tf_idf, documents, topics, self.nr_samples, self.nr_repr_docs
        )

        # We extract the top n words per class
        topics = self._extract_candidate_words(topic_model, c_tf_idf, topics)

        # We calculate the similarity between word and document embeddings and create
        # topic embeddings from the representative document embeddings
        sim_matrix, words = self._extract_embeddings(topic_model, topics, representative_docs, repr_doc_indices)

nor does the class have "embedding" arguments

class KeyBERTInspired(BaseRepresentation):
    def __init__(
        self,
        top_n_words: int = 10,
        nr_repr_docs: int = 5,
        nr_samples: int = 500,
        nr_candidate_words: int = 100,
        random_state: int = 42,
    ):

Proposed fix would be to pass "embeddings" parameter when executing "extract_topics" of KeyBERTInspired, and create embeddings only when its none. This will speed up the pipeline and so will reduce OpenAI consumption.
If this is good, I have added a PR for the same. Kindly advice, thank you.

Reproduction

Here is a sample code that doesnt fail, but is unnecessarily consumes OpenAI tokens and delays the model fit.

client = AzureOpenAI(
api_version = os.environ["OPENAI_API_VERSION"],
api_key =os.environ["AZURE_OPENAI_API_KEY"].strip(),
azure_endpoint =os.environ["AZURE_OPENAI_ENDPOINT"] ,
)
#openAI as embedding tool
embedding_model = OpenAIBackend(client, os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"])

# We define a number of topics that we know are in the documents
zeroshot_topic_list = ["Human resources ", "Information Technology ","Finance"]

topic_model = BERTopic(
    embedding_model=embedding_model,
    min_topic_size=15,
    zeroshot_topic_list=zeroshot_topic_list,
    zeroshot_min_similarity=.85,
    representation_model=KeyBERTInspired()
)
topics, _ = topic_model.fit_transform(df["content_detail"], embeddings)
BERTopic Version

v0.17.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with KeyBERTInspired.extract_topics and its _extract_embeddings helper, then trace how BERTopic.fit_transform passes precomputed embeddings into representation models. Verify whether representative document indices can select existing embeddings instead of recomputing them, while preserving word-embedding behavior. Done means the reproduction no longer consumes unnecessary document-embedding tokens and the existing topic extraction behavior remains intact.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.