model.py

import json

from bertopic import BERTopic
from bertopic.representation import MaximalMarginalRelevance, KeyBERTInspired
from bertopic.vectorizers import ClassTfidfTransformer

from umap import UMAP
from hdbscan import HDBSCAN
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import CountVectorizer

import streamlit as st
import pandas as pd

import plotly.graph_objects as go
import plotly.express as px
from wordcloud import WordCloud
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle
import matplotlib.pyplot as plt

import numpy as np

class BERTmodel:
    def __init__(self, fname:str)->None:
        self.topic_model = None
        self.docs = None
        self.embedding_model = None
        self.umap_model = None
        self.hdbscan_model = None
        self.vectorizer_model = None
        self.ctfidf_model = None
        self.representation_model = None
        self.embeddings = None
        self.topics = None
        self.num_topics = None
        self.probs = None

        self.vis_method_map = {
            "topic map": self.__topic_map,
            "topic similarity": self.__topic_similarity,
            "topic barchart": self.__topic_barchart,
            "topic clouds": self.__word_clouds,
            "topic sunburst": self.__sunburst,
            "topic treemap": self.__treemap,
            "document topics": self.__document_topics,
            "documents": self.__documents,
            "3D document topics": self.__3d_topic_map,
            "cluster map": self.__cluster_map  
        }

        self.__get_spans(fname)

## PRIVATE METHODS
    # Topic map in 2-D
    def __topic_map(self)->go.Figure:
        st.session_state.fig_plotly = True
        return self.topic_model.visualize_topics(top_n_topics=self.num_topics, custom_labels=True)
    
    # Topic map in 3-D
    def __3d_topic_map(self)->go.Figure:
        st.session_state.fig_plotly = True
        # Reduce 384-dimensional embeddings to 3 dimensions
        reduced_embeddings_3d = UMAP(
            n_neighbors=10, 
            n_components=3, 
            min_dist=0.07, 
            metric='cosine',
            random_state=42
            ).fit_transform(self.embeddings)   

        # Create dataframe with 3D coordinates
        df_3d = pd.DataFrame(reduced_embeddings_3d, columns=["x", "y", "z"])
        # Add title and cluster columns to the dataframe
        df_3d["title"] = [self.topic_info[self.topic_info['Topic'] == self.topics[i]]['CustomName'].iloc[0] for i, _ in enumerate(self.topics)]
        df_3d["cluster"] = [str(c) for c in self.topics]   

        # Create 3D scatter plot 
        fig = px.scatter_3d(
            df_3d, 
            x='x', 
            y='y', 
            z='z', 
            color='cluster', 
            title='3D Document Clusters', 
            color_continuous_scale='viridis', 
            size_max=0.5, 
            hover_data=['title']
            )
        # Update layout
        fig.update_layout(
            dict(width=900,
            height=800,
            legend_title_text="Topic Clusters",
            showlegend=True,
            uirevision='foo'    # Setting to a constant prevents resizing of the 3-D canvas when a legend item is toggled
            )
            )

        return fig
  
    # Generate topic heatmap
    def __topic_similarity(self)->go.Figure:
        st.session_state.fig_plotly = True
        return self.topic_model.visualize_heatmap(top_n_topics=self.num_topics, custom_labels=True)
    
    # Generate topic barchart
    def __topic_barchart(self)->go.Figure:
        st.session_state.fig_plotly = True
        return self.topic_model.visualize_barchart(top_n_topics=self.num_topics, autoscale=True)
    
    # Generate topic word clouds as subplots
    def __word_clouds(self)->Figure:
        st.session_state.fig_plotly = False
        # dynamically allocate subplots
        n_cols = 4
        n_rows = int((self.num_topics + n_cols - 1)/n_cols)
        
        cloud = WordCloud(width=2200, height=2000, background_color="white")   

        fig, axes = plt.subplots(n_rows, n_cols, figsize=(10,10), sharex=True, sharey=True)
        for i, ax in enumerate(axes.flatten()):
            if (i < self.num_topics):
                fig.add_subplot(ax)    
                text = {word: value for word, value in self.topic_model.get_topic(i)}
                cloud.generate_from_frequencies(text)
                plt.imshow(cloud, interpolation="bilinear")
                plt.gca().set_title('Topic ' + str(i), fontdict=dict(size=12))
                plt.axis("off")
                rect = Rectangle((0, 0), 1, 1, transform=ax.transAxes, edgecolor='gray', facecolor='none', lw=1.5)
                # Add the rectangle to the axes
                ax.add_patch(rect)
            else:
                fig.delaxes(ax)   

        plt.subplots_adjust(wspace=0, hspace=0)
        plt.axis('off')
        plt.margins(x=0, y=0)
        plt.tight_layout()
        
        return fig
    
    # Create a sunburst plot showing topics and their most representative words, colour coded as to importance
    def __sunburst(self)->go.Figure:
        st.session_state.fig_plotly = True

        topic_word_imp = self.__get_vis_df()
        topic_word_imp['hole'] = 'Topics'

        fig = px.sunburst(topic_word_imp, 
                        path=['hole', 'topic', 'word'], 
                        values='importance', 
                        color='importance', 
                        color_continuous_scale='viridis', 
                        labels={'importance': 'Importance', 'name': 'Topic Terms'},
                        hover_data=['name']
                        )
        fig.update_layout(width=900, 
                          height=900, 
                          margin = dict(t=0, l=0, r=0, b=0),
                          coloraxis_colorbar = dict(len=0.8, thickness=20)
                          )

        return fig
    
    # Create a treemap plot showing topics and their most representative words, colour coded as to importance
    def __treemap(self)->go.Figure:
        st.session_state.fig_plotly = True

        treemap_data = self.__get_vis_df()

        fig = px.treemap(treemap_data, path=['topic', 'word'],
                        values='importance',
                        color='importance', 
                        hover_data=['name'],
                        color_continuous_scale='viridis',
                        color_continuous_midpoint=np.average(treemap_data['importance'])
                        )
        fig.update_layout(margin = dict(t=0, l=0, r=0, b=0),
                          coloraxis_colorbar = dict(len=0.8, thickness=20)                          
                          )        

        return fig

    # Display the document corpus grouped by the dominant topic in each document
    def __document_topics(self)->go.Figure:
        st.session_state.fig_plotly = True

        probs_df=pd.DataFrame(self.probs)
        probs_df['main percentage'] = pd.DataFrame({'max': probs_df.max(axis=1)})
        column_index_of_max = probs_df.idxmax(axis=1)

        counts = column_index_of_max.value_counts()

        chart_df = pd.DataFrame(columns = ["topic", "value", "keywords"])

        for i in range(len(counts)):
            chart_df.loc[len(chart_df)] = [i, counts[i], self.topic_info.iloc[i+1, self.topic_info.columns.get_loc('Representation')]] 

        fig = px.bar(chart_df, 
                     x='topic', 
                     y='value', 
                     title='Documents Grouped by Dominant Topic', 
                     hover_data=['keywords'], 
                     labels={'value': 'No. of Documents', 'topic': 'Topic Number'}, 
                     width=750, 
                     height=500)
        fig.update_traces(width=.5)
        fig.update_layout(bargap=.1)
        fig.update_xaxes(type='category')

        return fig
    
    def __documents(self)->go.Figure:
        st.session_state.fig_plotly = True

        # Reduce dimensionality of embeddings
        reduced = UMAP(n_neighbors=5, n_components=2, min_dist=0.0, metric='cosine', random_state=42).fit_transform(self.embeddings)

        # Create the documents dataframe
        df_2d = pd.DataFrame(reduced, columns=["x", "y"])
        # Add title and cluster columns to the dataframe
        df_2d["title"] = [self.topic_info[self.topic_info['Topic'] == self.topics[i]]['CustomName'].iloc[0] for i, _ in enumerate(self.topics)]
        df_2d["cluster"] = [str(c) for c in self.topics]   

        # Create the scatter plot of document clusters
        fig = px.scatter(
            df_2d, 
            x='x', 
            y='y', 
            color='cluster', 
            title='Document Clusters by Topic', 
            color_continuous_scale='viridis', 
            size_max=0.5, 
            hover_data=['title']
            )
        
        # Update layout
        fig.update_layout(
            dict(width=900,
            height=800,
            legend_title_text="Topics",
            showlegend=True)
            )  

        return fig      
    
    def __cluster_map(self)->go.Figure:
        st.session_state.fig_plotly = True
        hierarchical_topics = self.topic_model.hierarchical_topics(self.docs, use_ctfidf=True)
        return self.topic_model.visualize_hierarchy(height=1000, 
                                                    custom_labels=True, 
                                                    orientation="bottom", 
                                                    use_ctfidf=False, 
                                                    hierarchical_topics=hierarchical_topics)

    # Load the text spans from a json file
    def __get_spans(self, fname:str)->None:
        with open(fname, "r", encoding="utf-8") as f:
            self.docs = json.load(f)["paragraphs"]

    # Helper method to create the Pandas dataframe used by the sunburst and treemap methods
    def __get_vis_df(self)->pd.DataFrame:
        topic_words = self.topic_model.get_topics()

        topic_data = pd.DataFrame(columns=['topic', 'word', 'importance', 'name'])

        for topic_id, words in topic_words.items():
            if topic_id != -1:  # Filter out outliers
                for word, importance in words:
                    topic_data.loc[len(topic_data)] = [topic_id, 
                                                       word, 
                                                       importance, 
                                                       self.topic_info.iloc[topic_id+1, self.topic_info.columns.get_loc('CustomName')]
                                                       ] 
        return topic_data        

# PUBLIC METHODS
    def BERT_topic_model(self,
            top_n_words=10, 
            n_gram_range=3,
            min_topic_size=10,
            n_neighbors=15,
            min_dist=0.01,
            n_components=5,
            min_cluster_size=10,
            min_samples=10,
            diversity=0.75):
        
        ngram_start = 2
        ngram_end = n_gram_range

        ngram = (ngram_start, ngram_end)

        # Extract embeddings
        self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
        # Reduce dimensionality
        self.umap_model = UMAP(n_neighbors=n_neighbors,
                               n_components=n_components, 
                               min_dist=min_dist, 
                               metric='cosine',
                               random_state=42)
        # Cluster reduced embeddings
        self.hdbscan_model = HDBSCAN(min_cluster_size=min_cluster_size, 
                                     min_samples=min_samples, 
                                     metric='euclidean', 
                                     cluster_selection_method='eom', 
                                     prediction_data=True, 
                                     approx_min_span_tree=False)
        # Tokenize topics
        self.vectorizer_model = CountVectorizer(stop_words="english", ngram_range=ngram)
        # Create topic representation
        self.ctfidf_model = ClassTfidfTransformer(reduce_frequent_words=True)
        # Fine-tune topic representations with a bertopic representation model
        # Using KeyBERTInspired and MMR
        keyword = KeyBERTInspired(top_n_words=top_n_words)
        mmr = MaximalMarginalRelevance(diversity=diversity)
        self.representation_model = [keyword, mmr]

        # Generate the model
        self.topic_model = BERTopic( 
            embedding_model = self.embedding_model,          
            umap_model = self.umap_model,                     
            hdbscan_model = self.hdbscan_model,               
            vectorizer_model = self.vectorizer_model,        
            ctfidf_model = self.ctfidf_model,                 
            representation_model = self.representation_model, 
            top_n_words = top_n_words,
            min_topic_size=min_topic_size,
            calculate_probabilities=True
            ) 
        # Encode the sentence model
        self.embeddings = self.embedding_model.encode(self.docs, show_progress_bar=False)  
        # Train the model     
        self.topics, self.probs = self.topic_model.fit_transform(self.docs)

        # Create custom labels, these are saved under the heading "CustomName" in the topic info
        # Each custom label is composed of the top 5 words for each cluster. The label does not 
        # have the default cluster number prefix and components of the label are separated by a 
        # comma instead of the default underscore
        labels = self.topic_model.generate_topic_labels(nr_words=3,topic_prefix=False, separator=',')
        self.topic_model.set_topic_labels(labels)

        self.topic_info = self.topic_model.get_topic_info()

        # Filter out the outlier topic (-1) and get the count of remaining topics
        self.num_topics = len(self.topic_info[self.topic_info['Topic'] != -1])

        return self.topic_model
    
    def vis_data(self, key:str)->go.Figure:
        return self.vis_method_map[key]()