import gradio as gr
from bertopic import BERTopic
import numpy as np
import pickle

with open('C1_msg.pkl', 'rb') as f:
  C1_msg = pickle.load(f)
    
reduced_embeddings = np.load('C1_reduced_embeddings.npy')

model = BERTopic.load("AlexanderHolmes0/C1-Topic-Model-100")

#model.set_topic_labels({8: "8_shimmer_religion"})

model2 = BERTopic.load("AlexanderHolmes0/C2-Topic-Model-100")

with gr.Blocks(theme=gr.themes.Soft()) as demo:
   with gr.Tab("C1 Topic Model"):
        gr.load("models/AlexanderHolmes0/C1-Topic-Model-100")
        with gr.Row():
            gr.Plot(model.visualize_topics(width=550),
                    min_width=500)
            gr.Plot(model.visualize_documents(C1_msg, 
                                              reduced_embeddings=reduced_embeddings,
                                              hide_annotations = True,
                                              #topics = range(0,7), 
                                              width = 900,
                                              custom_labels=False),
                    min_width=900)
   with gr.Tab("C2 Topic Model"): 
            gr.load("models/AlexanderHolmes0/C2-Topic-Model-100",alias = "C2 Topic Model")
            with gr.Row():
                gr.Plot(model2.visualize_topics(width=650),
                        min_width=650)
                gr.Plot(model2.visualize_hierarchy(width=650,
                                                   height=650),
                        min_width=650)

demo.launch()