Class TopicModel

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

class TopicModel:
    def __init__(self, config):
        self.config = config

    def __repr__(self) -> str:
        return f"TopicModel(config={self.config})"

    def fit_topic_model(self, topic_model, data, embeddings, key="query"):
        topics, probs = topic_model.fit_transform(data[key].to_list(), embeddings)
        return topics, probs

    def get_topic_info(self, topic_model):
        return topic_model.get_topic_info()

    def reduce_topics(self, topic_model, data, nr_topics, key="query"):
        topic_model.reduce_topics(data[key].to_list(), nr_topics)
        return topic_model

    def get_representative_docs(self, topic_model):
        return topic_model.get_representative_docs()

    def reduce_outliers(self, topic_model, data, topics, probs, key="query", strategy="c-tf-idf"):
        if strategy == "c-tf-idf":
            new_topics = topic_model.reduce_outliers(data[key].to_list(), topics, strategy, threshold=0.1)
        elif strategy == "embeddings":
            new_topics = topic_model.reduce_outliers(data[key].to_list(), topics, strategy)
        elif strategy == "distributions":
            new_topics = topic_model.reduce_outliers(
                data[key].to_list(), topics, probabilities=probs, strategy=strategy
            )
        else:
            raise ValueError("Invalid strategy")
        topic_model.update_topics(
            data[key].to_list(),
            topics=new_topics,
            representation_model=self.representation_model,
            vectorizer_model=self.vectorizer_model,
            ctfidf_model=self.ctfidf_model,
        )
        return topic_model, new_topics

    def compute_hierarchical_topic_tree(self, topic_model, data, key="query"):
        hierarchical_topics = topic_model.hierarchical_topics(data[key].to_list())
        tree = topic_model.get_topic_tree(hierarchical_topics)
        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),