src/kbdebugger/topic_modelling/BERTopic.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from pydoc_data import topics
from typing import List, Optional, Tuple, Literal, TypedDict

from bertopic import BERTopic

from sklearn.feature_extraction.text import CountVectorizer
from .logging import save_topic_modeling_results


@dataclass(frozen=True)
class TopicModelConfig:
    """
    Configuration for topic modeling with BERTopic.

    Parameters
    ----------
    embedding_model: str or model
        The sentence transformer or embedding model used to vectorize texts.

    min_topic_size: int
        Minimum size of a topic. Used by HDBSCAN to avoid very small, noisy clusters.

    top_n_words: int
        Number of representative keywords to extract per topic.

    language: str
        Language used by CountVectorizer for stopword removal.

    seed: int
        Random seed for reproducibility.
    """
    # embedding_model: str = "all-MiniLM-L6-v2" # dimension: 384
    embedding_model: str = "all-mpnet-base-v2" # dimension: 768
    min_topic_size: int = 10
    top_n_words: int = 16
    language: str = "english"
    seed: int = 42


class TopicResult(TypedDict):
    topic_id: int
    keyword_frequency: float
    keywords: List[str]
    paragraph_indices: List[int]
    matched: Literal["exact", "synonym", "none"]


def extract_topics_from_paragraphs(
    paragraphs: List[str],
    keyword: str,
    synonyms: Optional[List[str]] = None,
    config: Optional[TopicModelConfig] = None,
) -> Tuple[
        List[TopicResult], 
        BERTopic
    ]:
    """
    Run BERTopic over the provided paragraphs and detect topics matching a given keyword.

    Parameters
    ----------
    paragraphs:
        List of clean paragraph strings to analyze.

    keyword:
        The user-selected topic of interest (e.g., "explainability").

    synonyms:
        Optional backup list of synonyms (used if exact keyword fails).

    config:
        Topic modeling hyperparameters. Defaults to sensible values.

    Returns
    -------
    - topic_matches: List of topic match metadata (see TopicResult).
    - model: Trained BERTopic model for further inspection/plotting.
    """
    cfg = config or TopicModelConfig()

    # Step 1: Instantiate BERTopic with configuration

    # Custom vectorizer with stopword removal
    vectorizer_model = CountVectorizer(stop_words=cfg.language)
    topic_model = BERTopic(
        embedding_model=cfg.embedding_model,
        vectorizer_model=vectorizer_model,
        top_n_words=cfg.top_n_words,
        language=cfg.language,
        calculate_probabilities=True,
        verbose=True,
        seed_topic_list=None,
    )

    # Step 2: Fit the model to input paragraphs
    topics, _probs = topic_model.fit_transform(paragraphs)
    # topics is List[int] of topic IDs per paragraph
    # propbs is List[List[float]] of topic probabilities per paragraph

    # cast the topics to a list of ints
    topics = list(map(int, topics))

    # Access the frequent topics that were generated:
    topic_info_df = topic_model.get_topic_info()
    # e.g.
    """
    >>> topic_info_df
    Topic   Count   Name
    -1      4630    -1_can_your_will_any
    0       693     49_windows_drive_dos_file
    1       466     32_jesus_bible_christian_faith
    2       441     2_space_launch_orbit_lunar
    3       381     22_key_encryption_keys_encrypted
    ...
    """

    # Step 3: Collect match info for all topics
    topic_matches: List[TopicResult] = []
    keyword_lower = keyword.lower()
    synonym_set = set(s.lower() for s in synonyms) if synonyms else set()

    matched_topic_ids: List[int] = [] # e.g. [3, 7]
    matched_synonyms: dict[int, set[str]] = {} # e.g. {7: {"explainable AI"}}
    match_type_by_topic: dict[int, str] = {} # e.g. {3: "exact", 7: "synonym"}


    for topic_id in topic_info_df["Topic"]:
        if topic_id == -1:
            # -1 is outlier cluster in HDBSCAN
            # Topic -1 refers to all outliers and should typically be ignored.
            continue

        # Get topic keywords and lowercased form for matching
        """
        >>> topic_model.get_topic(topic_id) # topic_id is an int
        BERTopic’s get_topic method returns the top-N highest TF-IDF keywords per topic.

        [('windows', 0.006152228076250982),
        ('drive', 0.004982897610645755),
        ('dos', 0.004845038866360651),
        ('file', 0.004140142872194834),
        ('disk', 0.004131678774810884),
        ('mac', 0.003624848635985097),
        ('memory', 0.0034840976976789903),
        ('software', 0.0034415334250699077),
        ('email', 0.0034239554442333257),
        ('pc', 0.003047105930670237)]
        """
        topic_words = [w for w, w_prob in topic_model.get_topic(topic_id)]
        topic_words_lower = [w.lower() for w in topic_words]

        match_type: Literal["exact", "synonym", "none"]
        # Matching logic
        if keyword_lower in topic_words_lower:
            # 💪 The user-chosen keyword was mentioned in the topic's keywords list
            match_type = "exact"
            matched_topic_ids.append(topic_id)
        else:
            # The exact keyword was not found, check synonyms if available
            intersecting = synonym_set.intersection(topic_words_lower)
            if intersecting:
                match_type = "synonym"
                matched_topic_ids.append(topic_id)
                # matched_synonyms[topic_id] = list(intersecting)[0] # store one matched synonym
                matched_synonyms[topic_id] = intersecting
            else:
                continue # skip the topic if no match

        match_type_by_topic[topic_id] = match_type

        # Find all paragraphs belonging to this topic
        paragraph_indices = [i for i, t in enumerate(topics) if t == topic_id]

        # Compute how often the user keyword appears in the topic's top keywords.
        # Usually 0 or 1, since BERTopic keywords are typically unique, but may repeat in rare cases.
        keyword_frequency = topic_words_lower.count(keyword_lower)

        topic_matches.append(
            TopicResult(
                topic_id=topic_id,
                keyword_frequency=keyword_frequency,
                keywords=topic_words,
                paragraph_indices=paragraph_indices,
                matched=match_type,
            )
        )

    # # Using .get_document_info, we can also extract information on a document level, 
    # # such as their corresponding topics, probabilities, whether they are representative documents for a topic, etc.   
    # """
    # >>> topic_model.get_document_info(docs)

    # Document                               Topic    Name                        Top_n_words                     Probability    ...
    # I am sure some bashers of Pens...       0       0_game_team_games_season    game - team - games...          0.200010       ...
    # My brother is in the market for...      -1     -1_can_your_will_any         can - your - will...            0.420668       ...
    # Finally you said what you dream...      -1     -1_can_your_will_any         can - your - will...            0.807259       ...
    # Think! It is the SCSI card doing...     49     49_windows_drive_dos_file    windows - drive - docs...       0.071746       ...
    # 1) I have an old Jasmine drive...       49     49_windows_drive_dos_file    windows - drive - docs...       0.038983       ...
    # """
    # topic_model.get_topic_info()

    # Save logs for this topic modeling run
    save_topic_modeling_results(
        topic_model=topic_model,
        documents=paragraphs,
        # document_chunks=...,
        keyword=keyword,
        matched_topic_ids=matched_topic_ids,
        match_type_by_topic=match_type_by_topic,
        matched_synonyms=matched_synonyms,
        generated_synonyms=synonyms,
    )

    return topic_matches, topic_model