Method run

src/ml_models/topic_modeling/topic_model.py:71–125

Run the topic modeling.

(self, data, gnn_outputs: Tuple[np.ndarray, np.ndarray], key="query")

        return hierarchical_topics, tree

    def run(self, data, gnn_outputs: Tuple[np.ndarray, np.ndarray], key="query") -> dict:
        """
        Run the topic modeling.
        """
        # Prepare the data
        if self.config["auto_cluster"]:
            if self.config["representation_model"].split("(")[0] == "pipeline":
                generator = self.config["representation_model"]
                representation_model = [MaximalMarginalRelevance(diversity=0.3), TextGeneration(eval(generator))]
            elif self.config["representation_model"].split("(")[0] == "OpenAI":
                representation_model = [
                    MaximalMarginalRelevance(diversity=0.3),
                    eval(self.config["representation_model"]),
                ]
            else:
                lg.info("Using MaximalMarginalRelevance as representation model")
                representation_model = MaximalMarginalRelevance(diversity=0.3)

            topic_model = BERTopic(
                nr_topics="auto",
                vectorizer_model=eval(self.config["vectorizer_model"]),
                embedding_model=self.config["embedding_model"],
                representation_model=representation_model,
                n_gram_range=(1, 2),
                min_topic_size=10,
                top_n_words=10,
                calculate_probabilities=False,
            )
        else:
            # Fit the BERTopic model
            topic_model = BERTopic(
                nr_topics="auto",
                vectorizer_model=eval(self.config["vectorizer_model"]),
                umap_model=eval(self.config["dimensionality_reduction"]),
                hdbscan_model=eval(self.config["clustering_model"]),
                embedding_model=eval(self.config["embedding_model"]),
                representation_model=eval(self.config["representation_model"]),
                n_gram_range=(1, 2),
                min_topic_size=10,
                top_n_words=10,
                calculate_probabilities=False,
            )
        topics, probabilities = topic_model.fit_transform(data[key].to_list(), gnn_outputs[1])
        freq = self.get_topic_info(topic_model)
        rep_docs = self.get_representative_docs(topic_model)
        hr, tree = self.compute_hierarchical_topic_tree(topic_model=topic_model, data=data)
        # add dict with topic info
        docs = []
        for k, v in rep_docs.items():
            docs.append((k, v))
        rep = pd.DataFrame(docs, columns=["Topic", "Rep_docs"])
        topic_info_freq = pd.merge(freq, rep, on="Topic")
        data["topic"] = topics
        data["prob"] = probabilities
        return {"data": data, "topic_info": topic_info_freq, "tree": tree}